Logit Lens, Direct Logit Attribution, Gemma setup

Set up TransformerLens on gemma-2-2b today, mostly fighting the harness rather than learning anything conceptually new. Gemma uses RMSNorm, not LayerNorm — the standard logit lens writeup assumes LayerNorm’s mean-centering, had to double check the RMSNorm version doesn’t silently break the interpretation.

Read nostalgebraist’s original Logit Lens post first — logit lens: $\text{logits}_l = W_U \cdot \text{LN}(h_l)$, unembed an intermediate residual stream as if the model stopped there. DLA is the natural extension, and it’s really just Anthropic’s “Mathematical Framework for Transformer Circuits” applied — residual stream is a sum of per-component writes, freeze the final norm’s scale, and each head/MLP’s contribution to the final logit becomes a dot product with $W_U$, additive across components. Both landed fast, it’s just linearity of the residual stream.

Got basic DLA running on an IOI-style prompt on Gemma tonight. Numbers look sane.

Todo:

• sanity check DLA against a real ablation baseline, not just “looks reasonable”
• the “freeze the norm scale” step in DLA is itself an approximation — check how much it actually varies between prompts before trusting the additive story too much