预归一化变换器中的门控归一化移除与尺度锚定
Gated Normalization Removal and Scale Anchoring in Pre-Norm Transformers
AI总结:
本文研究了预归一化变换器中归一化层的必要性,提出了一种门控归一化移除方法,通过TaperNorm逐步将归一化操作转为样本无关的线性或仿射映射,并揭示了最终归一化层对预logit表示尺度的锚定作用。
摘要:
归一化层在变换器中是标准组件,但其样本依赖的计算在整个训练和推理过程中是否必要尚不明确。本文为预归一化变换器开发了一种门控归一化移除方法。该方法通过TaperNorm实现,从标准RMSNorm/LayerNorm逐步过渡到学习的样本无关线性或仿射映射。一旦门控达到零, tapered层将不再计算每个token的统计信息,所得到的映射可以折叠到相邻的线性投影中。结果表明,在测试的预训练和微调设置中,内部归一化可以逐步移除,且验证损失的增加较小。我们的方法揭示了最终归一化层的独特作用,即它锚定了预logit表示的尺度。有了这个锚定,最后隐藏状态的径向变化不会直接减少损失;当移除它时,通过增加logit的幅度可以实现交叉熵的减少。固定目标尺度损失提供了显式的替代锚定,使在测试范围内能够完全无归一化地进行消融实验。最后,在KV缓存自回归解码基准中,逐步移除内部归一化可提供高达1.14倍的吞吐量,使用显式缩放操作,折叠后可达1.18倍。
英文摘要:
Normalization layers are standard in transformers, but it is not clear whether their sample-dependent computations are necessary throughout both training and inference. This work develops a gated normalization-removal approach for pre-norm transformers. The approach is implemented using TaperNorm, which starts from standard RMSNorm/LayerNorm and gradually tapers to learned sample-independent linear or affine maps. Once the gate reaches zero, per-token statistics are no longer computed in the tapered layers and the resulting maps can be folded into adjacent linear projections. The results indicate that internal normalization can be tapered in the tested pre-training and fine-tuning settings with small validation-loss increases. Our approach helps reveal a distinct role for final normalization, namely that it anchors the scale of the pre-logit representation. With this anchor present, radial changes in the last hidden state do not directly reduce the loss; when it is removed, reducing cross-entropy can be achieved by increasing logit magnitudes. A fixed-target scale loss provides an explicit alternative anchor and enables fully norm-free ablations in the tested regimes. Finally, in a KV-cached autoregressive decoding benchmark, tapering internal norms gives up to $1.14\times$ higher throughput with explicit scaling operations and up to $1.18\times$ after folding.
机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。