This work was produced during MARS and SPAR. arXiv version available at https://arxiv.org/abs/2507.02559. Code on GitHub and models on HuggingFace.

TL;DR we scaled LayerNorm (LN) removal by fine-tuning to GPT-2 XL:

We improve training stability by regularizing activation standard deviation across token positions & improve the training code.
The process is still not very data efficient (requires about ~1% of original dataset) and fine-tuning is somewhat unstable on larger models.
Surprisingly, attribution patching is not more accurate in LN-free models.
As expected, direct logit attribution becomes exact, and entropy neurons stop working.

We provide drop-in replacements for GPT-2 small, medium, large, and XL on HuggingFace for all your interpretability needs!

Introduction

LN centers the activations by the mean and normalizes by the standard deviation, then multiplies by learnable scale and adds a learnable bias.

LN has been a major obstacle for mechanistic interpretability research as it obstructs the interpretation of the residual stream and complicates model decomposition into circuits. For methods like direct logit attribution and attribution patching, researchers resort to freezing the standard deviation, which is an imperfect approximation. Past attempts at fully removing LN include training tiny models without LN and removal by fine-tuning for GPT-2 Small. We explored how to scale removal by fine-tuning and some properties of LN-free models when it comes to mechanistic interpretability.

Main Achievements

We successfully removed LN from all GPT-2 models up to XL (1.5B parameters) by fine-tuning on OpenWebText (OWT):

GPT-2 Model	Fine tuning steps	Vanilla CE loss	LN-free CE loss
Small	300	2.8112	2.8757 [+0.0858]
Medium	500	2.5724	2.6352 [+0.0962]
Large	600	2.5074	2.5159 [+0.0812]
XL	800	2.3821	2.3992 [+0.0253]

Performance degradation in LN-free models was small compared to the fine-tuned control (aka "vanilla"), as measured by CE loss on a filtered version of The Pile (The Pile-filtered). See model generations in Appendix A. The runs required 0.3% - 1% of original training data (0.5M tokens per step).

We investigated how the LN-free models behave from a mechanistic interpretability point of view. As expected, direct logit attribution (a method for measuring the impact of a component in the residual stream on the output logits) gives exact direct effect on logits. Surprisingly, attribu...

Posts

Wikitag Contributions

Comments