只看归一化层梯度就能准确预测Transformer训练噪声，提速18%。
Normalization Layer Per-Example Gradients are Sufficient to Predict Gradient Noise Scale in Transformers
- 计算梯度时同步获取每样本梯度范数，效率高
- 仅归一化层的梯度噪声可准确预测全模型噪声
- 优化批量大小调度，显著缩短训练时间
每样本梯度范数是低方差估计梯度噪声尺度（GNS）的关键。观察其张量运算需求，我们提出一种在三维及以上张量场景下计算开销极小的方法：在计算参数梯度的同时同步计算梯度范数。使用该方法，我们比以往更准确地观测了不同层的GNS。发现现代Transformer模型的总GNS可由归一化层的GNS良好预测。因此，我们仅聚焦归一化层，设计自定义核，在执行LayerNorm反向传播时零吞吐开销计算每样本梯度范数。仅追踪这些层的GNS，便能指导实用的批量大小调度，在Chinchilla最优语言模型上将训练时间减少18%。
原文摘要 · Abstract (English)
Per-example gradient norms are a vital ingredient for estimating gradient noise scale (GNS) with minimal variance. Observing the tensor contractions required to compute them, we propose a method with minimal FLOPs in 3D or greater tensor regimes by simultaneously computing the norms while computing the parameter gradients. Using this method we are able to observe the GNS of different layers at higher accuracy than previously possible. We find that the total GNS of contemporary transformer models is predicted well by the GNS of only the normalization layers. As a result, focusing only on the normalization layer, we develop a custom kernel to compute the per-example gradient norms while performing the LayerNorm backward pass with zero throughput overhead. Tracking GNS on only those layers, we are able to guide a practical batch size schedule that reduces training time by 18% on a Chinchilla-optimal language model.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。