2025-11-21T22:28:16.015152

Weight Initialization and Variance Dynamics in Deep Neural Networks and Large Language Models

Han
Weight initialization governs signal propagation and gradient flow at the start of training. This paper offers a theory-grounded and empirically validated study across two regimes: compact ReLU multilayer perceptrons and GPT-2-style transformers. First, a logarithmic sweep of the initial standard deviation maps vanishing and exploding regimes and identifies a broad stability band with standard deviations between 1e-2 and 1e-1. Second, a controlled comparison shows that Kaiming (fan-in) initialization converges faster and more stably than Xavier under ReLU, consistent with variance-preserving theory. Third, in a from-scratch 12-layer GPT-2-style model, this paper tracks layerwise Q/K/V weight variance through pretraining and observe depth-dependent equilibration into narrow bands: shallow layers expand rapidly while deeper layers change more gradually. Together, these results connect classic initialization principles with modern transformer behavior and yield simple, practical recipes for robust training.
academic

Weight Initialization and Variance Dynamics in Deep Neural Networks and Large Language Models

基本信息

  • 论文ID: 2510.09423
  • 标题: Weight Initialization and Variance Dynamics in Deep Neural Networks and Large Language Models
  • 作者: Yankun Han (University of Florida)
  • 分类: cs.LG (Machine Learning)
  • 发表时间: 2025年10月10日 (arXiv预印本)
  • 论文链接: https://arxiv.org/abs/2510.09423

摘要

权重初始化控制着训练开始时的信号传播和梯度流。本文提供了一项理论基础扎实且经验验证的研究，涵盖两个领域：紧凑的ReLU多层感知机和GPT-2风格的Transformer。首先，通过对初始标准差的对数扫描，映射了梯度消失和爆炸区域，并识别出标准差在1e-2到1e-1之间的广泛稳定带。其次，受控比较显示，在ReLU激活函数下，Kaiming（fan-in）初始化比Xavier初始化收敛更快且更稳定，这与方差保持理论一致。第三，在从零开始构建的12层GPT-2风格模型中，本文追踪了预训练过程中各层Q/K/V权重方差的变化，观察到深度相关的平衡现象：浅层快速扩展，而深层变化更加渐进。

研究背景与动机

问题定义

本研究要解决的核心问题是深度神经网络和大型语言模型中权重初始化对训练稳定性和收敛性的影响。具体包括：

  1. 初始化尺度敏感性：不同的初始化标准差如何影响训练的稳定性
  2. 激活函数特异性：ReLU和GELU等激活函数是否需要特定的初始化策略
  3. 现代Transformer的方差动态：大型Transformer模型中方差稳定化是否持续存在

重要性

权重初始化是深度学习训练成功的关键因素，不当的初始化会导致：

  • 梯度消失：信号在深层网络中逐层衰减
  • 梯度爆炸：信号在传播过程中指数级增长
  • 训练不稳定：优化过程中的振荡和发散

现有方法局限性

经典的初始化方法（LeCun、Xavier/Glorot、He/Kaiming）虽然在理论上有方差保持的直觉，但在实际应用中仍存在以下问题：

  1. 对理想尺度偏差的敏感性未被充分量化
  2. 特定激活函数（如ReLU、GELU）的影响机制不明确
  3. 在大型Transformer中的表现缺乏系统性研究

核心贡献

  1. 统一的方差分析框架：推导了常见激活函数（ReLU、GELU）的前向和后向方差传播条件，解释了fan-in缩放如何保持信号幅度以及ReLU中因子2的来源
  2. 尺度敏感性量化：通过对25个标准差值的对数扫描，映射了梯度消失/爆炸区域，识别出稳定训练带σ ∈ 10⁻², 10⁻¹
  3. 激活函数感知的初始化验证：在受控的ReLU MLP训练中，证实Kaiming normal（fan-in）比Xavier normal收敛更快且损失方差更小
  4. Transformer方差动态分析：在从零开始的12层GPT-2风格模型中，发现了明显的深度相关模式：浅层权重标准差快速扩展，深层更加渐进，最终都稳定在窄方差带中

方法详解

理论框架

前向传播方差分析

对于线性映射：

Var[z_l] = n_in σ²_W Var[x_{l-1}]

经过非线性激活后：

Var[x_l] ≈ c_φ n_in σ²_W Var[x_{l-1}]

其中 c_φ = E[φ(z)²]/Var[z] 是激活函数相关的常数。

为避免激活值消失或爆炸，选择 σ²_W ≈ 1/(c_φ n_in)：

  • ReLU: c_φ ≈ 1/2，因此 σ²_W ≈ 2/n_in (He/Kaiming)
  • GELU: c_φ ≈ 0.45-0.5，略小于ReLU

后向传播方差分析

反向传播给出：

Var[δ_{l-1}] ≈ n_out σ²_W d_φ Var[δ_l]

其中 d_φ = E[φ'(z)²]。对于ReLU，d_φ = 1/2，平衡梯度方差需要 σ²_W ≈ 2/n_out。

权衡与实践选择

前向和后向保持条件通常无法同时满足，除非 n_in ≈ n_out 且 c_φ ≈ d_φ。实践中，保持前向信号稳定通常更重要，这解释了为什么fan-in He/Kaiming比Xavier收敛更快。

实验设计

实验E1：标准差扫描

  • 网络架构：784→64→32→32→10的ReLU MLP
  • 数据集：MNIST
  • 扫描范围：25个标准差值，从10⁻⁴到10，对数间隔
  • 评价指标：损失轨迹、分类准确率

实验E2：Xavier vs Kaiming比较

  • 网络架构：11→16→32→32→1的ReLU网络
  • 数据集：UCI Wine二分类任务
  • 对比方案：Xavier normal vs Kaiming uniform
  • 统计验证：10次随机运行，配对t检验

实验E3：GPT-2方差动态

  • 模型规模：12层GPT-2风格Transformer
  • 初始化：标准配置（大部分模块std=0.02，嵌入层xavier normal）
  • 优化器：AdamW，学习率1×10⁻⁴，批大小16
  • 追踪目标：所有层的Q/K/V投影权重标准差

实验结果

E1：标准差扫描结果

  • 稳定区间：σ ∈ 10⁻², 10⁻¹内训练平滑，梯度表现良好，准确率在此区间内达到峰值
  • 梯度消失：极小尺度（σ ≲ 10⁻³）导致更新消失和准确率下降
  • 梯度爆炸：极大尺度（σ ≳ 1）产生不稳定损失和偶发发散

E2：初始化方法比较

Kaiming初始化在多个维度上持续优于Xavier：

  • 收敛速度：达到目标的中位数轮次更少，早期损失下降更陡峭
  • 准确率：最终验证准确率匹配或略超过Xavier
  • 统计显著性：配对t检验显示损失和训练准确率差异显著（p < 0.05）

E3：Transformer方差动态发现

  • 深度相关模式：浅层在早期训练中显示快速且显著的权重标准差扩展，深层扩展更加缓慢和平滑
  • 方差平衡：所有层最终稳定在窄方差带中
  • 分布稀疏化：训练后权重分布变得更稀疏，许多接近零的条目保持不变，少数大幅度权重占主导

理论洞察与实践意义

深度相关的方差平衡机制

论文揭示了Transformer中的渐进平衡模式：

  1. 浅层快速适应：接近输入的层具有高信噪比梯度，鼓励早期积极缩放
  2. 深层渐进调整：残差路径长度和预归一化限制了深层的有效步长
  3. 隐式约束：注意力softmax饱和和AdamW中的权重衰减阻止大参数尺度

实践指导原则

  1. ReLU/GELU MLP：从fan-in He/Kaiming开始；如果非常不平衡的层导致梯度漂移，略向fan-average选择移动
  2. 深度残差栈：残差缩放（如1/√L）或归一化有助于防止深度方差漂移
  3. Transformer投影：使用小标准差初始化（如0.02），监控每层标准差和梯度范数

相关工作比较

基础初始化策略

  • LeCun方法：针对线性激活的方差保持规则
  • Glorot/Xavier：基于fan的tanh/sigmoid缩放
  • He/Kaiming：补偿ReLU下减半二阶矩的激活感知缩放

现代发展

  • Fixup初始化：通过精心选择的初始化和残差缩放，在极深网络中移除归一化需求
  • DeepNet：提出允许训练千层级别的原则性深度缩放规则
  • 预归一化优势：相比后归一化，通过平滑梯度流改善优化稳定性

结论与讨论

主要结论

  1. 稳定性带存在：在σ ∈ 10⁻², 10⁻¹范围内存在广泛但敏感的稳定性带
  2. 激活函数特异性重要：Kaiming初始化在ReLU网络中确实优于Xavier
  3. 深度相关动态：Transformer展现出深度相关的方差平衡，浅层快速适应，深层渐进调整

局限性

  1. 实验规模：GPT-2实验相对较小（12层），大规模模型的行为可能不同
  2. 激活函数覆盖：主要关注ReLU和GELU，其他激活函数的分析有限
  3. 优化器依赖：结果可能对特定优化器（AdamW）和超参数设置敏感

未来方向

  1. 自适应深度感知初始化：学习每层或每头的尺度，使浅层更接近最终方差水平
  2. 优化器和调度耦合：联合调优预热长度、权重衰减和梯度裁剪
  3. 深度和宽度缩放：评估深度相关平衡在更大模型下的持续性

深度评价

优点

  1. 理论与实践结合：将经典方差传播理论与现代Transformer行为有机结合
  2. 系统性实验设计：从简单MLP到复杂Transformer的渐进式验证
  3. 实用价值高：提供了具体的初始化建议和诊断方法
  4. 统计严谨性：使用配对t检验等统计方法验证结果显著性

不足

  1. 理论分析深度有限：缺少对深度相关现象的更深层理论解释
  2. 实验规模约束：受计算资源限制，未能在真正大规模模型上验证
  3. 泛化性问题：结果主要基于特定架构和任务，泛化能力需要进一步验证

影响力评估

  1. 学术贡献：为初始化理论提供了现代视角，连接经典理论与当前实践
  2. 实用价值：为实践者提供了明确的初始化策略和诊断工具
  3. 可复现性：实验设计清晰，代码和参数设置详细，便于复现

适用场景

  1. 深度网络训练：特别适用于ReLU/GELU激活的深度网络
  2. Transformer优化：为大型语言模型的训练提供初始化指导
  3. 研究工具：为研究者提供了分析权重动态的方法论框架

参考文献

论文引用了初始化领域的关键工作，包括LeCun、Glorot、He等人的奠基性研究，以及近期在Transformer优化方面的进展，为本研究提供了坚实的理论基础。