Logit lens och residualströmmen

I en transformer adderar varje lager till samma vektor — residualströmmen. Utmatningslagret (unembedding) översätter den sista vektorn till sannolikheter över ordförrådet.

Logit lens: applicera samma unembedding på residualströmmen efter varje lager, inte bara det sista. Resultatet är en sekvens av «vad hade modellen svarat om den stannat här?».

Typiskt mönster: de första lagren ger nonsens eller högfrekventa ord, någonstans i mitten dyker rätt svar upp, och de sista lagren skärper sannolikheten. Var svaret «uppstår» är ofta informativt — särskilt när man jämför prompter som modellen klarar med sådana den inte klarar.

from transformer_lens import HookedTransformer
import torch

m = HookedTransformer.from_pretrained("gpt2-small")
tokens = m.to_tokens("Huvudstaden i Frankrike är")
logits, cache = m.run_with_cache(tokens)

for lager in range(m.cfg.n_layers):
    resid = cache[f"blocks.{lager}.hook_resid_post"][0, -1]       # sista positionen
    resid = m.ln_final(resid[None, None])                          # samma normalisering som utgången
    topp = m.unembed(resid)[0, 0].softmax(-1).topk(3)
    ord_ = [m.to_string(i) for i in topp.indices]
    print(f"lager {lager:2d}: " + ", ".join(f"{o!r} {p:.2f}" for o, p in zip(ord_, topp.values)))
# lager  0: ' the' 0.03, ',' 0.02, ' a' 0.02      ← högfrekventa ord
# lager  6: ' Paris' 0.08, ' France' 0.05, ...    ← svaret börjar framträda
# lager 11: ' Paris' 0.61, ' Lyon' 0.04, ...      ← skärpt

Två viktiga förbehåll:

1. Fungerar inte lika bra på alla modeller. Logit lens förutsätter att mellanlagrens representationer ligger i ungefär samma «rum» som utgången. I många modeller stämmer det dåligt, och resultatet blir brus.
2. Tuned lens (Belrose m.fl. 2023) löser det genom att lära in en affin transformation per lager innan unembedding. Det ger tydligt mer tillförlitliga och jämförbara resultat — använd den när logit lens ser rörig ut.

Och som alltid i interpretability: det här är observation. Att svaret syns i lager 6 bevisar inte att lager 6 orsakar det — det kräver intervention.