移除Transformer推理时的归一化层,模型性能几乎不变且更易解释。
Transformers Don't Need LayerNorm at Inference Time: Scaling LayerNorm Removal to GPT-2 XL and the Implications for Mechanistic Interpretability
- 移除所有层归一化(LN)层,仅使GPT-2 XL验证损失微增0.03。
- 无需归一化时,直接激活值可精确反映组件贡献,提升可解释性。
- 适合关注模型机制解析的研究者,尤其对可解释性工具改进有帮助。
层归一化(Layer-wise normalization, LN)是几乎所有基于Transformer的大语言模型的核心组件。尽管其在训练稳定性中的作用已被充分研究,但在推理阶段的作用仍不明确。此外,LN层会引入额外非线性,增加模型组件间的耦合,阻碍机制可解释性。本文表明,所有GPT-2模型均可完全移除LN层,仅导致验证损失轻微上升(如GPT-2 XL仅增加0.03交叉熵损失)。这说明LN在语言建模中并无关键作用。我们发现,为实现去LN化所需的微调数据量随模型参数增长呈亚线性关系,表明扩展至更大模型可行。我们已在Hugging Face发布一系列无LN的GPT-2模型。进一步测试显示,在无LN模型上,直接logit归因可精确揭示单个组件的直接影响,而归因修补法准确率未显著提升。同时确认,原GPT-2中的“置信度神经元”在无LN模型中已失效。本工作澄清了LN在语言建模中的角色,证明类GPT-2模型可在无LN情况下正常运行。我们希望这些无LN版本能推动更精确的可解释性研究,深化对语言模型的理解。
原文摘要 · Abstract (English)
Layer-wise normalization (LN) is an essential component of virtually all transformer-based large language models. While its effects on training stability are well documented, its role at inference time is poorly understood. Additionally, LN layers hinder mechanistic interpretability by introducing additional nonlinearities and increasing the interconnectedness of individual model components. Here, we show that all LN layers can be removed from every GPT-2 model with only a small increase in validation loss (e.g. +0.03 cross-entropy loss for GPT-2 XL). Thus, LN cannot play a substantial role in language modeling. We find that the amount of fine-tuning data needed for LN removal grows sublinearly with model parameters, suggesting scaling to larger models is feasible. We release a suite of LN-free GPT-2 models on Hugging Face. Furthermore, we test interpretability techniques on LN-free models. Direct logit attribution now gives the exact direct effect of individual components, while the accuracy of attribution patching does not significantly improve. We also confirm that GPT-2's "confidence neurons" are inactive in the LN-free models. Our work clarifies the role of LN layers in language modeling, showing that GPT-2-class models can function without LN layers. We hope that our LN-free analogs of the GPT-2 family of models will enable more precise interpretability research and improve our understanding of language models.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。