[flyingdoubleG]

Description

On architectures with post-sublayer norms, each attention/MLP output passes through ln1_post/ln2_post before it is added to the residual stream. These are Gemma 2/3 (sandwich norms) and OLMo 2/3, EXAONE-4 and GLM-4 (post-norms). A head or neuron therefore contributes ln_post(x_i) to the stream, not x_i.

#1648 fixed the block-level aliases (hook_attn_out / hook_mlp_out are now the post-norm outputs), so decompose_resid is exact on these models. But the per-component stacks built on top of them still use the pre-norm outputs:

• stack_head_results stacks attn.hook_result, the per-head output of the attention module before ln1_post;
• get_neuron_results / stack_neuron_results use a_n * W_out_n before ln2_post;
• get_full_resid_decomposition adds the raw b_O / b_out through accumulated_bias.

As a result, get_full_resid_decomposition(...).sum(0) no longer equals the residual stream on these models. On tiny random models (3 layers, d_model 64), comparing the sum at the final layer with resid_post:

model	dev, expand_neurons=False	dev, expand_neurons=True	this PR
Gemma 2	0.74	0.99	2e-7
Gemma 3	0.71		2e-7
OLMo 2	0.70		3e-7
OLMo 3	0.70		3e-7
EXAONE-4	0.71		3e-7
GLM-4	0.73		2e-7

(max abs error relative to max abs resid_post; the same holds in compatibility mode). For 32 other architectures I checked (Llama, Gemma 1, Granite, Qwen 2/3, Mistral, Phi, GPT-J, NeoX, Falcon, Cohere, BLOOM, GPT-2, Mixtral, Qwen2-MoE, OLMoE, …) the decomposition was already exact and is unchanged.

Fix. With the cached scale s the post-norm is affine, norm(x) = center(x) / s * w + b. So each component contributes center(x_i) / s * w, and the block adds the norm bias b once. A new _post_norm_map reads s from ln*_post.hook_scale and the weight with the norm's own convention (uses_weight_offset for Gemma's 1 + w, centering for LayerNorm). It is applied in:

• stack_head_results, per position and head. This is done at read time, so attn.hook_result keeps its meaning and results cached in the forward pass (use_attn_result=True) are mapped the same way;
• get_neuron_results and the analytic _stack_neuron_results_apply_ln_projected path, by folding center(W_out_n) * w into W_out and dividing the activations by s, so the projected path still never materializes [..., d_mlp, d_model];
• the bias term of get_full_resid_decomposition: a per-position correction replaces those blocks' raw output biases with norm_map(bias) + norm_bias, respecting residual_contribution_scale.

Pre-norm blocks have no ln*_post, so nothing changes for them.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

New tests/integration/model_bridge/test_post_norm_head_decomposition.py builds seeded tiny Gemma 2, Gemma 3, OLMo 2 and (as a pre-norm control) Llama models from local configs, with no hub access. It checks:

• get_full_resid_decomposition sums to resid_pre / resid_mid / the final resid_post at every layer, with and without expand_neurons. This fails on dev for all three post-norm models and passes here;
• head results cached with use_attn_result=True give the same stack as recomputed ones;
• the analytic apply_ln + project_output_onto path matches LN-then-project of the dense stack.

Local runs (Windows, CPU, torch 2.9.1, --jaxtyping-packages=transformer_lens,beartype.beartype):

• tests/unit/test_activation_cache.py, tests/unit/model_bridge/test_mixed_norm_decomposition.py, tests/integration/model_bridge/test_{granite_embedding_decomposition,parallel_mlp_input_decomposition,parallel_residual_identities,residual_decomposition_identities,post_norm_head_decomposition}.py: 529 passed;
• tests/acceptance/test_activation_cache.py: 43 passed.

black (the locked 23.12.1) and isort --check-only pass, and mypy transformer_lens/ActivationCache.py reports no issues.

This was prepared with AI assistance (Claude); I reviewed and tested the change.

[jlarson4]

Looks good, thank you for putting this together @flyingdoubleG