Transformers Don't Need LayerNorm at Inference Time: Scaling LayerNorm Removal to GPT-2 XL and the Implications for Mechanistic Interpretability

📄 arXiv: 2507.02559v1 📥 PDF

作者: Luca Baroni, Galvin Khara, Joachim Schaeffer, Marat Subkhankulov, Stefan Heimersheim

分类: cs.LG

发布日期: 2025-07-03

💡 一句话要点

Transformer推理时可移除LayerNorm：扩展至GPT-2 XL并应用于可解释性研究

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: Transformer LayerNorm GPT-2 可解释性 推理优化

📋 核心要点

  1. 现有Transformer模型依赖LayerNorm，但其在推理阶段的作用尚不明确，且阻碍了模型的可解释性。
  2. 该论文提出移除GPT-2模型中的所有LayerNorm层，并通过微调来补偿性能损失，简化模型结构。
  3. 实验表明，移除LayerNorm对GPT-2 XL的性能影响很小，且微调所需数据量随模型增大呈次线性增长。

📝 摘要（中文）

层归一化(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层的情况下运行。作者希望GPT-2系列的无LN类似物能够实现更精确的可解释性研究，并提高对语言模型的理解。

🔬 方法详解

问题定义：现有基于Transformer的大型语言模型广泛使用LayerNorm (LN)。虽然LN在训练阶段对稳定性的作用已被广泛研究，但其在推理阶段的具体作用尚不明确。此外，LN层引入了额外的非线性，增加了模型内部组件的连接复杂性，从而阻碍了模型的可解释性分析。因此，该论文旨在探究LN在推理阶段的必要性，并尝试移除LN以简化模型结构，提升可解释性。

核心思路：该论文的核心思路是，通过移除GPT-2模型中的所有LN层，然后通过微调来补偿由此带来的性能损失。作者认为，如果在移除LN后，模型仍然能够通过少量微调恢复性能，则表明LN在推理阶段并非必不可少。同时，移除LN可以简化模型结构，降低计算复杂度，并有利于后续的可解释性分析。

技术框架：该论文的技术框架主要包括以下几个步骤：1) 移除GPT-2模型中的所有LN层，得到一个无LN的模型。2) 使用少量数据对无LN的模型进行微调，以恢复其性能。3) 评估微调后的无LN模型的性能，并与原始GPT-2模型进行比较。4) 在无LN模型上应用各种可解释性技术，并与原始GPT-2模型进行比较，以评估移除LN对可解释性的影响。

关键创新：该论文最重要的技术创新点在于，证明了可以从GPT-2模型中移除所有LN层，而验证损失仅略有增加。这表明LN在推理阶段并非必不可少，为简化模型结构、提升可解释性提供了可能性。此外，该论文还发现，移除LN所需的微调数据量随模型参数的增长呈次线性关系，这意味着该方法可以扩展到更大的模型。

关键设计：该论文的关键设计包括：1) 使用GPT-2模型作为实验对象，因为它是一个广泛使用的、具有代表性的Transformer模型。2) 使用交叉熵损失作为评估模型性能的指标。3) 使用少量数据对无LN的模型进行微调，以避免过度拟合。4) 使用各种可解释性技术，如直接logit归因和归因修补，来评估移除LN对可解释性的影响。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，从GPT-2 XL模型中移除所有LayerNorm层后，交叉熵损失仅增加0.03，通过少量微调即可恢复性能。研究还发现，移除LN所需的微调数据量随模型参数的增长呈次线性关系，表明该方法具有良好的扩展性。此外，在无LN模型上，直接logit归因的准确性得到提升。

🎯 应用场景

该研究成果可应用于简化Transformer模型结构，降低计算成本，尤其是在资源受限的设备上部署大型语言模型。此外，无LayerNorm的模型更易于进行可解释性分析，有助于理解语言模型的内部机制，并可能促进更安全、可靠的AI系统的开发。

📄 摘要（原文）

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.