Views
No views yet
schaeff/gpt2-large_LNFree600GPT2LMHeadModel architecture to avoid requiring trust_remote_code=True. While LayerNorm blocks are technically present, they have been effectively disabled through parameter manipulation.ln_1, ln_2 and ln_f modules with identity functions.transformers:model = GPT2LMHeadModel.from_pretrained("schaeff/gpt2-large_LNFree600")1import torch
2from transformers import GPT2LMHeadModel
3from transformer_lens import HookedTransformer
4
5model = GPT2LMHeadModel.from_pretrained("schaeff/gpt2-large_LNFree600").to("cpu")
6
7# Undo hacky LayerNorm removal
8for block in model.transformer.h:
9 block.ln_1.weight.data = block.ln_1.weight.data / 1e6
10 block.ln_1.eps = 1e-5
11 block.ln_2.weight.data = block.ln_2.weight.data / 1e6
12 block.ln_2.eps = 1e-5
13model.transformer.ln_f.weight.data = model.transformer.ln_f.weight.data / 1e6
14model.transformer.ln_f.eps = 1e-5
15
16# Properly replace LayerNorms by Identities
17def removeLN(transformer_lens_model):
18 for i in range(len(transformer_lens_model.blocks)):
19 transformer_lens_model.blocks[i].ln1 = torch.nn.Identity()
20 transformer_lens_model.blocks[i].ln2 = torch.nn.Identity()
21 transformer_lens_model.ln_final = torch.nn.Identity()
22
23# transformer_lens
24hooked_model = HookedTransformer.from_pretrained("gpt2", hf_model=model, fold_ln=True, center_unembed=False).to("cpu")
25removeLN(hooked_model)
26
27# NNSight:
28from nnsight.models.UnifiedTransformer import UnifiedTransformer
29
30model_nnsight = UnifiedTransformer(model="gpt2", hf_model=model, fold_ln=True, center_unembed=False).to("cpu")
31removeLN(model_nnsight)remove_layernorm in utils.py for details on the parameter hack and eval.py for loading.| Model | FT steps | OWT (val) | The Pile | The Pile-filtered |
|---|---|---|---|---|
| OpenAI GPT-2 Small original | 0 | 3.1006 | 2.8450 | 2.7899 |
| schaeff GPT-2 Small vanilla | 300 | 3.0126 | 2.8511 | 2.8112 |
| schaeff GPT-2 Small LN-free | 300 | 3.0797 [+0.0671] | 2.8852 [+0.0402] | 2.8757 [+0.0858] |
| OpenAI GPT-2 Medium original | 0 | 2.8145 | 2.5163 | 2.5390 |
| schaeff GPT-2 Medium vanilla | 500 | 2.7390 | 2.5752 | 2.5724 |
| schaeff GPT-2 Medium LN-free | 500 | 2.7642 [+0.0252] | 2.6579 [+0.1416] | 2.6352 [+0.0962] |
| OpenAI GPT-2 Large original | 0 | 2.6623 | 2.5320 | 2.4347 |
| schaeff GPT-2 Large vanilla | 600 | 2.6240 | 2.6233 | 2.5074 |
| schaeff GPT-2 Large LN-free | 600 | 2.6384 [+0.0144] | 2.7504 [+0.2184] | 2.5159 [+0.0812] |
| OpenAI GPT-2 XL original | 0 | 2.5567 | 2.4436¹ | 2.3739 |
| schaeff GPT-2 XL vanilla | 800 | 2.4799 | 2.4673 | 2.3821 |
| schaeff GPT-2 XL LN-free | 800 | 2.5052 [+0.0253] | 130.2197² | 2.3992 [+0.0253] |
@misc{gpt2layernorm2025,
author = {Baroni, Luca and Khara, Galvin and Schaeffer, Joachim and Subkhankulov, Marat and Heimersheim, Stefan},
title = {Transformers Don't Need LayerNorm at Inference Time: Scaling LayerNorm Removal to GPT-2 XL and the Implications for Mechanistic Interpretability},
year = {2025},
eprint = {2507.02559},
archivePrefix = {arXiv},
primaryClass = {cs.LG},
url = {https://arxiv.org/abs/2507.02559v1}
}