StefanHex

You can remove GPT2’s LayerNorm by fine-tuning for an hour

This work was produced at Apollo Research, based on initial research done at MATS. Edit: arXiv version available at https://arxiv.org/abs/2409.13710 LayerNorm is annoying for mechanistic interpretability research (“[...] reason #78 for why interpretability researchers hate LayerNorm” – Anthropic, 2023). Here’s a Hugging Face link to a GPT2-small model without any LayerNorm. The final model is only slightly worse than a GPT2 with LayerNorm[1]: DatasetOriginal GPT2Fine-tuned GPT2 with LayerNormFine-tuned GPT without LayerNormOpenWebText (ce_loss)3.0952.9893.014 (+0.025)ThePile (ce_loss)2.8562.8802.926 (+0.046)HellaSwag (accuracy)29.56%29.82%29.54% I fine-tuned GPT2-small on OpenWebText while slowly removing its LayerNorm layers, waiting for the loss to go back down after reach removal: Introduction LayerNorm (LN) is a component in Transformer models that normalizes embedding vectors to have constant length; specifically it divides the embeddings by their standard deviation taken over the hidden dimension. It was originally introduced to stabilize and speed up training of models (as a replacement for batch normalization). It is active during training and inference. LN(x)=x−E[x]√Var[x]+ϵ⋅γ+β The equation includes the standard deviation (std) √Var[x]+ϵ which makes it a non-linear operation. This hinders interpretability in a variety of ways, from annoyances and inaccuracies such as * attributing residual stream directions to logit effects (e.g. SAE features, direct logit attribution),[2] * being annoying to deal with Attribution Patching, or * being difficult to deal with in Apollo’s LIB method. In the Docstring circuit analysis we seriously considered whether the model might be using LN in its algorithm. This post even shows that LN can be used as the sole non-linearity to solve non-linear classification problems (see also this related work). Recently, with progress in Sparse Dictionary Learning, agendas (e.g. this one) imagine decomposing networks int

166Aug 8, 2024

StefanHex

Message

Stefan Heimersheim. Mechanistic interpretability & AI safety researcher, previously at FAR.AI and Apollo Research. The opinions expressed here are my own and do not necessarily reflect the views of my employer.

2112

264

127

Activation Plateaus: Where and How They Emerge

By design, LLMs perform nonlinear mappings from their inputs (text sequences) to their outputs (next-token generations). Some of these nonlinearities are built-in to the model architecture, but others are learned by the model, and may be important parts of how the model represents and transforms information. By studying different aspects...

Oct 17, 202537

Transformers Don't Need LayerNorm at Inference Time: Implications for Interpretability

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. *...

Jul 23, 202531

Trusted monitoring, but with deception probes.

Recently, Apollo trained some deception probes (Goldowsky-Dill et al). A deception probe is a logistic classifier on the AI's internal activations, indicating whether a token belongs in a deceptive response. We benchmarked these deception probes, testing them across five datasets of strategic deception and comparing them to black-box monitoring. We...

Jul 23, 202531

Compressed Computation is (probably) not Computation in Superposition

This research was completed during the Mentorship for Alignment Research Students (MARS 2.0) Supervised Program for Alignment Research (SPAR spring 2025) programs. The team was supervised by Stefan (Apollo Research). Jai and Sara were the primary contributors, Stefan contributed ideas, ran final experiments and helped writing the post. Giorgi contributed...

Jun 23, 202557

Try training token-level probes

Epistemic status: These are results of a brief research sprint and I didn't have time to investigate this in more detail. However, the results were sufficiently surprising that I think it's worth sharing. I'm not aware of previous papers doing this but surely someone tried this before, I would welcome...

Apr 14, 202547

Proof-of-Concept Debugger for a Small LLM

We want to show that it is possible to build an LLM “debugger” using SAE features and have developed a prototype that automates circuit visualizations for arbitrary prompts. With a few improvements to existing techniques (notably, “cluster resampling”, which is a form of activation patching), we are able to produce...

Mar 17, 202527

Detecting Strategic Deception Using Linear Probes

Can you tell when an LLM is lying from the activations? Are simple methods good enough? We recently published a paper investigating if linear probes detect when Llama is deceptive. Abstract: > AI models might use deceptive strategies as part of scheming or misaligned behaviour. Monitoring outputs alone is insufficient,...

Feb 6, 2025104

Load More (7/30)

LESSWRONG
LW

LESSWRONG
LW

StefanHex

StefanHex

StefanHex

You can remove GPT2’s LayerNorm by fine-tuning for an hour

[Interim research report] Activation plateaus & sensitive directions in GPT2

Solving the Mechanistic Interpretability challenges: EIS VII Challenge 1

Attribution-based parameter decomposition

StefanHex

Activation Plateaus: Where and How They Emerge

Transformers Don't Need LayerNorm at Inference Time: Implications for Interpretability

Trusted monitoring, but with deception probes.

Compressed Computation is (probably) not Computation in Superposition

Try training token-level probes

Proof-of-Concept Debugger for a Small LLM

Detecting Strategic Deception Using Linear Probes

Activation Plateaus: Where and How They Emerge

Transformers Don't Need LayerNorm at Inference Time: Implications for Interpretability

Trusted monitoring, but with deception probes.

Compressed Computation is (probably) not Computation in Superposition

Try training token-level probes

Proof-of-Concept Debugger for a Small LLM

Detecting Strategic Deception Using Linear Probes

You can remove GPT2’s LayerNorm by fine-tuning for an hour

[Interim research report] Activation plateaus & sensitive directions in GPT2

Solving the Mechanistic Interpretability challenges: EIS VII Challenge 1

Attribution-based parameter decomposition