提出去除Transformer内部归一化的新方法,提升推理速度且保持性能
Gated Normalization Removal and Scale Anchoring in Pre-Norm Transformers
- 用TaperNorm逐步替换归一化层为固定映射,实现无样本依赖计算
- 移除内部归一化后验证损失仅小幅上升,最终归一化起关键尺度锚定作用
- 支持全无归一化模型,推理吞吐最高提升1.18倍,适合部署优化场景
归一化层是Transformer的标准组件,但其依赖样本的计算在训练和推理中是否必要尚不明确。本文提出一种针对预归一化Transformer的门控归一化移除方法,通过TaperNorm从标准RMSNorm/层归一化逐步过渡到可学习的固定线性或仿射映射。当门控值降为零后,各标记统计量不再计算,映射可融合至相邻线性层。实验表明,在测试的预训练与微调设置中,内部归一化可被有效缩减,验证损失仅小幅增加。该方法揭示了最终归一化的独特作用:它锚定了预逻辑表示的尺度。有此锚点时,最后一层隐藏状态的径向变化不会直接降低损失;若移除,则可通过增大逻辑值来减小交叉熵。固定目标尺度损失提供了一种显式替代锚点,使在测试范围内实现完全无归一化成为可能。在带有键值缓存的自回归解码基准中,缩减内部归一化可带来最高1.14倍的吞吐提升(显式缩放),最高1.18倍(融合后)。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。