技术摘要：小规模初始化对大语言模型的重要性

问题陈述

虽然大语言模型（LLM）的进展通常归功于规模、数据和架构，但参数初始化作为一个关键的设计选择，在大型模型时代重新成为了一个复杂的问题。标准启发式方法（如 Xavier、Kaiming）通常需要高昂的试错调优成本。一种盛行的理论观点认为，小规模初始化会驱动网络进入“凝聚”（condensed）状态，即权重在多样化之前先沿共享方向对齐，从而可能增强推理和泛化能力。然而，目前尚不清楚这些优势是否能在现实的大规模 LLM 预训练中持续存在，或者架构组件是否会无意中抑制小规模初始化的优势。

研究方法

作者系统地研究了初始化规模对 LLM 训练动态和性能的影响。

• 初始化参数化： 权重矩阵 $W$ 初始化为 $W_{ij} \sim \mathcal{N}(0, \sigma^2)$，其中 $\sigma = d_{in}^{-\gamma}$。参数 $\gamma$ 控制缩放比例：$\gamma = 0.5$ 恢复了标准的类 Xavier 缩放，而 $\gamma > 0.5$（特别是 $\gamma = 1$）则产生更小的初始化。
• 模型架构： 实验在稠密仅解码器（decoder-only）Transformer 和混合专家（MoE）模型上均进行了测试。架构变体包括标准的 LayerNorm/RMSNorm 和多头注意力机制。
• 架构干预： 为了解决观察到的局限性，研究测试了两种特定的调整方案：
  1. 降低 RMSNorm 常数 ($\epsilon$)： 将 $\epsilon$ 从标准的 $10^{-5}$ 降低到 $10^{-12}$，以防止当隐藏状态方差较小时，归一化因子被常数所主导。
  2. 门控注意力（Gated Attention）： 在注意力头输出中引入查询依赖型门控，以减轻“注意力汇聚”（attention sink，即对第一个 token 的不成比例关注）现象，这种现象在小规模初始化下会被加剧。
• 训练设置： 模型在高质量的双语（中英）语料库（1 万亿 token）上使用 AdamW 进行下个 token 预测进行训练。实验涵盖了从 0.1B 到 1.5B 参数的模型规模。
• 分析指标：
  • Token 层面： 对称概率间隙 ($\Delta_{sym}$) 和 token 难度（平均损失）以定位性能增益所在。
  • 机制层面： 行向余弦相似度（Row-wise cosine similarity）和稳定秩（Stable Rank $\|W\|_F^2 / \|W\|_2^2$）以追踪参数结构的演化（凝聚 vs. 扩张）。

核心结果

1. 缩放限制及其修复方案

在标准设置下（$\gamma=0.5$ vs. $\gamma=1$），小规模初始化始终能降低验证损失，但随着模型规模增大，这种收益会递减。作者确定了导致这一现象的两个架构障碍：

• RMSNorm 饱和： 使用标准 $\epsilon=10^{-5}$ 时，一旦 $\sigma^2(h) \lesssim \epsilon$，归一化因子对减小的方差就会变得不敏感。将 $\epsilon$ 降至 $10^{-12}$ 可以恢复敏感度。
• 注意力汇聚（Attention Sink）： 小规模初始化强化了对第一个 token 进行关注的倾向。门控注意力有效地缓解了这一问题，在结合小规模初始化时显著降低了损失。
• 结果： 当这两个障碍被移除时，性能差距进一步扩大。采用这些调整的小规模初始化 ($\gamma=1$) 使有效模型规模提升了约 44%，并在不同规模下产生了持续的收益。

2. 最优初始化规模 ($\gamma=1$)

研究表明，初始化不应是任意小的。

• 残差流平衡（Residual Flow Balance）： 在 Transformer 中，残差更新必须与恒等路径保持相当。理论分析显示，残差流随 $d^{1-\gamma}$ 缩放。
  • 如果 $\gamma < 1$，残差占主导地位。
  • 如果 $\gamma > 1$，残差消失，使网络接近恒等映射。
  • $\gamma = 1$ 代表了临界平衡点，即残差更新与嵌入（embedding）规模相匹配。经验表明，性能在 $\gamma=1$ 时达到顶峰，并在 $\gamma > 1$ 时下降。

3. Token 层面与任务特定的增益

• 非均匀增益： 改进在不同 token 之间是不均匀的。增益集中在“中等至高度困难”的 token 上，这些 token 需要上下文整合与推理能力，而非简单的 token 或极度嘈杂/罕见的 token。
• 基准测试表现： 在 1.5B 模型上，优化后的小规模初始化（$\gamma=1$ + 调整）在重推理任务的基准测试（如 GSM8K, MATH500, HellaSwag）中显著优于标准初始化 ($\gamma=0.5$)，在多个任务上的绝对增益超过 4%。

4. 机制轨迹：凝聚

小规模初始化诱导了一种独特的发育轨迹：

• 阶段 1（凝聚）： 在训练早期，权重矩阵迅速凝聚成低秩、低复杂度的结构（行向量向连贯的块对齐）。
• 阶段 2（扩张）： 随着训练推进，这些结构演变为更丰富、更高复杂度的表示。
• 对比： 标准初始化 ($\gamma=0.5$) 从弥散、高秩的结构开始，演化过程缓慢，缺乏初始的低复杂度阶段。这种“从低到高”的复杂度轨迹被假设为使模型偏向于学习简约的推理规则，而非表层记忆。

意义与主张

本文认为，参数初始化是决定模型容量和训练动态的“基因级”因素，而非仅仅是一个底层的实现细节。

• 设计轴线： 作者建议将初始化规模 $\gamma$ 作为 LLM 设计中的显式超参数。
• 实践建议： 他们倡导采用 $\gamma$-初始化规则：将初始化规模作为一个可调节的旋钮，并将 $\gamma=1$ 作为默认值。这种干预被描述为“几乎零成本”，却能提升预训练效果并增强跨规模的推理能力。
• 理论洞察： 该工作为“压缩即智能”提供了具体证据，证明了特定的初始化方案能迫使模型经历一个凝聚阶段，而这与提升的推理能力相关联。
• 架构相互依赖性： 研究结果强调，初始化的收益取决于架构组件（特别是归一化常数和注意力机制），这意味着这些组件必须进行联合调优，而非独立调整。