归一化层的逐样本梯度足以预测Transformer中的梯度噪声尺度
Normalization Layer Per-Example Gradients are Sufficient to Predict Gradient Noise Scale in Transformers
AI总结：
本研究提出低运算量的逐样本梯度范数计算方法，发现Transformer模型总梯度噪声尺度可由归一化层的GNS预测，据此开发无开销的定制内核并优化批量调度，将Chinchilla最优语言模型训练时间缩短18%。
摘要：
逐样本梯度范数是估计梯度噪声尺度（GNS）的关键要素，能实现最小方差的估计。通过分析计算逐样本梯度范数所需的张量收缩操作，我们提出一种在三维及以上张量场景下FLOPs（浮点运算次数）极低的方法，可在计算参数梯度的同时同步计算梯度范数。借助该方法，我们能以高于以往的精度观测不同层的GNS。研究发现，当代Transformer模型的总GNS可仅通过归一化层的GNS很好地预测。基于这一结论，我们仅针对归一化层开发了定制内核，可在执行LayerNorm反向传播的同时计算逐样本梯度范数，且无吞吐量开销。仅跟踪这些层的GNS，我们就能制定实用的批量大小调度策略，在Chinchilla最优语言模型上将训练时间缩短18%。
英文摘要：
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.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
