Variance Control via Weight Rescaling in LLM Pre-training

📄 arXiv: 2503.17500v1 📥 PDF

作者: Louis Owen, Abhay Kumar, Nilabhra Roy Chowdhury, Fabian Güra

分类: cs.LG, cs.CL, stat.ML

发布日期: 2025-03-21

🔗 代码/项目: GITHUB

💡 一句话要点

提出LIR初始化与TVR方差控制，提升LLM预训练性能并降低量化难度

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 预训练 权重初始化 方差控制 低精度训练

📋 核心要点

  1. LLM预训练对权重初始化和方差控制策略高度敏感，现有方法在LLM上的研究相对不足。
  2. 论文提出LIR权重初始化和TVR方差控制策略，旨在更好地管理LLM预训练过程中的方差。
  3. 实验表明，该方法能显著提升下游任务性能，降低激活值范围，有利于量化和低精度训练。

📝 摘要（中文）

大型语言模型（LLM）的预训练结果很大程度上取决于权重初始化和方差控制策略。虽然初始方差控制的重要性已在通用神经网络中得到充分证明，但专门针对LLM预训练期间的初始化及其增长管理的文献相对较少。本文提出了层索引重缩放（LIR）权重初始化方案和目标方差重缩放（TVR）方差控制策略。在具有10亿参数的LLaMA模型上的实验表明，使用这些技术进行更好的方差管理可以显著提高下游任务的性能（在常见的预训练基准测试中高达4.6%），并减少极端的激活值，从而减轻与量化和低精度训练相关的挑战。代码已在GitHub上公开。

🔬 方法详解

问题定义：LLM预训练过程中，权重初始化和方差控制对模型性能至关重要。现有的初始化和方差控制方法在通用神经网络上已经有较多研究，但在LLM预训练场景下，相关研究相对稀疏，尤其是在预训练过程中如何有效管理方差的增长，以避免梯度爆炸或消失，以及激活值过大导致量化困难等问题。

核心思路：论文的核心思路是通过权重重缩放来控制每一层的方差，从而在预训练过程中维持一个更稳定的激活分布。LIR初始化方案旨在提供一个更好的初始方差，而TVR方差控制策略则在训练过程中动态调整权重，以使每一层的方差接近目标值。这样设计的目的是为了避免激活值过大或过小，从而提高训练的稳定性和效率。

技术框架：整体框架包含两个主要部分：LIR权重初始化和TVR方差控制。首先，使用LIR初始化模型权重。在预训练过程中，周期性地使用TVR策略来调整权重。TVR策略计算每一层的实际方差，并将其与目标方差进行比较，然后根据差异调整权重。这个过程可以看作是一个闭环反馈系统，不断地将每一层的方差拉回到目标值附近。

关键创新：最重要的技术创新在于将层索引信息融入到权重初始化中（LIR），并提出了一个在训练过程中动态调整权重的方差控制策略（TVR）。LIR通过层索引来调整初始权重，使得不同层的初始方差更加合理。TVR则通过周期性地调整权重，使得每一层的方差在训练过程中保持稳定，避免了方差的过度增长或衰减。与现有方法相比，LIR和TVR更加关注LLM预训练过程中的方差管理，并提供了一种更加精细化的控制手段。

关键设计：LIR的关键设计在于如何根据层索引来确定重缩放因子。一种简单的实现方式是使用一个与层索引相关的函数来计算重缩放因子。TVR的关键设计在于如何选择目标方差以及如何调整权重。目标方差可以是一个预先设定的常数，也可以是一个与层索引相关的函数。权重的调整可以通过一个简单的比例因子来实现，该比例因子取决于实际方差和目标方差之间的比率。

📊 实验亮点

实验结果表明，使用LIR初始化和TVR方差控制策略后，在1B参数的LLaMA模型上，下游任务性能提升高达4.6%。同时，激活值的范围显著减小，降低了量化难度。这些结果表明，更好的方差管理能够显著提高LLM的预训练效果。

🎯 应用场景

该研究成果可广泛应用于大型语言模型的预训练，尤其是在资源受限或需要低精度训练的场景下。通过更有效的方差控制，可以降低训练成本，提高模型性能，并简化模型部署。此外，该方法还有助于解决量化难题，使得LLM能够更好地在边缘设备上运行。

📄 摘要（原文）

The outcome of Large Language Model (LLM) pre-training strongly depends on weight initialization and variance control strategies. Although the importance of initial variance control has been well documented in neural networks in general, the literature on initialization and management of its growth during LLM pre-training, specifically, is somewhat sparse. In this paper, we introduce the Layer Index Rescaling (LIR) weight initialization scheme, and the Target Variance Rescaling (TVR) variance control strategy. Experiments on a 1B parameter LLaMA model demonstrate that better variance management using these techniques yields substantial improvements in downstream task performance (up to 4.6% on common pre-training benchmarks) and reduces extreme activation values, thus mitigating challenges associated with quantization and low-precision training. Our code is available at: https://github.com/bluorion-com/weight_rescaling.