Logit Lens / Tuned Lens¶
Companion to research/10-logit-lens-tuned-lens.md
The logit lens (nostalgebraist, 2020) decodes a transformer's intermediate residual states with the final layernorm + unembedding, $\text{logits}_\ell = W_U\,\text{LN}_f(h_\ell)$, to watch a prediction form layer by layer. It breaks when intermediate layers use a rotated basis. The tuned lens (Belrose et al., 2023) fixes this with a learned per-layer affine translator, $\text{logits}_\ell = W_U\,\text{LN}_f(A_\ell h_\ell + b_\ell)$.
What we will see:
- Logit lens: a prediction forms gradually across layers (matched-basis case).
- A per-layer basis rotation breaks the naive logit lens.
- The tuned lens learns the translator and recovers the trajectory.
- The caveat: a tuned lens trained on the final answer can look ahead of the model.
- The prediction trajectory flags anomalous inputs (Belrose et al.).
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import warnings
warnings.filterwarnings('ignore')
np.random.seed(42)
torch.manual_seed(42)
plt.rcParams.update({
'figure.dpi': 110,
'font.size': 11,
'axes.titlesize': 12,
'axes.labelsize': 11,
'figure.facecolor': 'white',
'axes.facecolor': '#f8f9fa',
'axes.grid': True,
'grid.alpha': 0.35,
})
print('Setup complete.')
Setup complete.
A synthetic residual-stream model¶
We mimic a transformer's residual stream: a $d=32$ state that, over $L=8$ layers, accumulates evidence toward one of $V=8$ classes. The final readout is $\text{logits}=W_U\,\widehat{h_L}$ (we use unit-norm $\widehat{h}$ as a stand-in for the final layernorm). Each class $v$ has a direction $W_U[v]$. The commitment schedule controls how fast the state aligns with the answer; later we will optionally put intermediate layers in a rotated basis that the final unembedding does not share.
d, V, L = 32, 8, 8
rng = np.random.RandomState(0)
W_U = rng.randn(V, d); W_U /= np.linalg.norm(W_U, axis=1, keepdims=True) # class directions (unembedding)
# per-layer rotations; the LAST layer is the identity (canonical basis = final basis)
def random_rotation(seed):
return np.linalg.qr(np.random.RandomState(seed).randn(d, d))[0]
R = [random_rotation(10 + l) for l in range(L)]; R[-1] = np.eye(d)
def commitment(l, early_cue=0.0):
# fraction of evidence accumulated by layer l (0..1); early_cue plants a faint signal from the start
return early_cue + (1 - early_cue) * ((l + 1) / L) ** 2
def make_data(n, seed, rotate, early_cue=0.0, noise=0.35):
r = np.random.RandomState(seed)
y = r.randint(0, V, n)
H = [] # H[l]: residual stream at layer l, shape (n, d)
for l in range(L):
s = commitment(l, early_cue) * W_U[y] + noise * r.randn(n, d) # semantic state (canonical basis)
h = s @ R[l].T if rotate else s # exposed residual stream (optionally rotated)
H.append(h.astype(np.float32))
return H, y
SCALE = 8.0
def decode_np(Hl): # logit lens: normalize + unembed
Hn = Hl / (np.linalg.norm(Hl, axis=1, keepdims=True) + 1e-8)
return SCALE * Hn @ W_U.T # (n, V) logits
def softmax_np(Z):
Z = Z - Z.max(1, keepdims=True); e = np.exp(Z); return e / e.sum(1, keepdims=True)
# sanity: final-layer accuracy (no rotation needed at the last layer)
H0, y0 = make_data(2000, 1, rotate=True)
final_acc = (decode_np(H0[-1]).argmax(1) == y0).mean()
print(f'model final-layer accuracy = {final_acc:.3f} (d={d}, V={V}, L={L})')
model final-layer accuracy = 0.920 (d=32, V=8, L=8)
What you see¶
The final layer decodes the answer well (~0.92 here) — the residual stream has committed to the right
class direction by layer $L$. The pieces are now in place: decode_np is the logit lens (normalize +
unembed), commitment says how much the model has decided by each layer, and R holds the per-layer
rotations we will switch on to create a basis mismatch.
Experiment 1: Logit lens — the prediction forms over layers¶
Question: What does decoding intermediate layers reveal when their basis matches the final one?
With rotations off (every layer in the canonical basis, like GPT-2), we apply the logit lens at each layer and watch the probability of the true class climb.
H, y = make_data(2000, 2, rotate=False)
true_prob = np.array([softmax_np(decode_np(H[l]))[np.arange(len(y)), y].mean() for l in range(L)])
acc_by_layer = np.array([(decode_np(H[l]).argmax(1) == y).mean() for l in range(L)])
# a probability heatmap for one example
ex = 0
heat = np.array([softmax_np(decode_np(H[l]))[ex] for l in range(L)]) # (L, V)
print('logit-lens true-class prob by layer:', np.round(true_prob, 2))
logit-lens true-class prob by layer: [0.14 0.15 0.18 0.25 0.32 0.43 0.58 0.72]
fig, axes = plt.subplots(1, 2, figsize=(11.5, 4.4))
axes[0].plot(range(1, L + 1), true_prob, 'o-', color='#1a9850', lw=2.5, ms=7, label='mean P(true class)')
axes[0].plot(range(1, L + 1), acc_by_layer, 's--', color='#542788', lw=2, ms=6, label='top-1 accuracy')
axes[0].axhline(1 / V, color='gray', ls=':', label='chance')
axes[0].set_title('Logit lens (matched basis): prediction forms gradually')
axes[0].set_xlabel('layer'); axes[0].set_ylabel('probability / accuracy'); axes[0].legend(); axes[0].set_ylim(0, 1.02)
im = axes[1].imshow(heat.T, aspect='auto', cmap='YlOrRd', vmin=0, vmax=1)
axes[1].axhline(y[ex], color='cyan', lw=2, ls='--')
axes[1].set_title(f'Per-layer class distribution (example, true={y[ex]})')
axes[1].set_xlabel('layer'); axes[1].set_ylabel('class'); axes[1].set_xticks(range(L)); axes[1].set_xticklabels(range(1, L + 1))
fig.colorbar(im, ax=axes[1], label='probability')
plt.tight_layout(); plt.show()
What you see¶
Left: the logit-lens probability of the true class rises smoothly from chance ($1/V$) toward 1 — we are literally watching the model make up its mind layer by layer ("iterative inference"). Right: for one example, the per-layer distribution concentrates onto the true class row (dashed line) as depth increases.
This is the logit lens working as intended — the picture you get on GPT-2, where intermediate states live in the same coordinate system as the final unembedding. The next experiment breaks that assumption.
Experiment 2 (limitation): a rotated basis breaks the logit lens¶
Question: What if intermediate layers don't share the final layer's coordinate system?
We turn rotations on: each layer's residual stream is rotated by $R_\ell$ (with $R_L=I$). The model is unchanged — only the basis the intermediate states are written in differs. We apply the same naive logit lens.
Hr, yr = make_data(2000, 2, rotate=True)
true_prob_rot = np.array([softmax_np(decode_np(Hr[l]))[np.arange(len(yr)), yr].mean() for l in range(L)])
acc_rot = np.array([(decode_np(Hr[l]).argmax(1) == yr).mean() for l in range(L)])
print('logit-lens true-class prob by layer (rotated):', np.round(true_prob_rot, 2))
logit-lens true-class prob by layer (rotated): [0.13 0.12 0.12 0.12 0.12 0.09 0.1 0.72]
fig, ax = plt.subplots(figsize=(7.4, 4.6))
ax.plot(range(1, L + 1), true_prob, 'o-', color='#1a9850', lw=2.5, ms=7, label='matched basis (Exp 1)')
ax.plot(range(1, L + 1), true_prob_rot, 's-', color='#d73027', lw=2.5, ms=7, label='rotated basis')
ax.axhline(1 / V, color='gray', ls=':', label='chance')
ax.set_title('Naive logit lens collapses under a basis mismatch')
ax.set_xlabel('layer'); ax.set_ylabel('logit-lens P(true class)'); ax.legend(); ax.set_ylim(0, 1.02)
plt.tight_layout(); plt.show()
What you see¶
With rotations on (red), the logit lens reads near-chance at every intermediate layer and only snaps to the answer at the final layer — even though the model encodes exactly the same gradually-formed prediction as before. Applying the final unembedding to a state written in a different basis produces nonsense, so the lens wrongly concludes "the model knows nothing until the end".
This is why the logit lens is systematically biased and fails on models like GPT-Neo, BLOOM, and OPT: their intermediate layers do not live in the unembedding's basis. The flat red curve is an artifact of the lens, not a fact about the model.
Experiment 3: The tuned lens learns the translator and recovers the trajectory¶
Question: Can a learned per-layer affine fix the basis mismatch?
For each layer we train an affine translator $(A_\ell, b_\ell)$ to minimise the KL divergence between the decoded distribution $W_U\,\widehat{A_\ell h_\ell + b_\ell}$ and the model's final distribution. The model and $W_U$ stay frozen; only the translator is learned.
def decode_torch(H):
Hn = H / (H.norm(dim=1, keepdim=True) + 1e-8)
return SCALE * Hn @ torch.tensor(W_U.T, dtype=torch.float32)
def train_translator(Hl_np, p_final, steps=400):
Hl = torch.tensor(Hl_np); P = torch.tensor(p_final, dtype=torch.float32)
A = nn.Parameter(torch.eye(d)); b = nn.Parameter(torch.zeros(d))
opt = torch.optim.Adam([A, b], lr=5e-3)
for _ in range(steps):
logits = decode_torch(Hl @ A.t() + b)
logp = torch.log_softmax(logits, 1)
loss = (P * (torch.log(P + 1e-9) - logp)).sum(1).mean() # KL(P_final || P_lens)
opt.zero_grad(); loss.backward(); opt.step()
return A.detach(), b.detach()
Htr, ytr = make_data(3000, 5, rotate=True)
p_final_tr = softmax_np(decode_np(Htr[-1]))
translators = [train_translator(Htr[l], p_final_tr) for l in range(L)]
# evaluate tuned lens on a fresh set
Hte, yte = make_data(2000, 6, rotate=True)
def tuned_probs(Hl, l):
A, b = translators[l]
with torch.no_grad():
return torch.softmax(decode_torch(torch.tensor(Hl) @ A.t() + b), 1).numpy()
tuned_true = np.array([tuned_probs(Hte[l], l)[np.arange(len(yte)), yte].mean() for l in range(L)])
logit_true = np.array([softmax_np(decode_np(Hte[l]))[np.arange(len(yte)), yte].mean() for l in range(L)])
print('tuned-lens true-class prob:', np.round(tuned_true, 2))
tuned-lens true-class prob: [0.12 0.13 0.14 0.16 0.23 0.34 0.45 0.71]
fig, ax = plt.subplots(figsize=(7.4, 4.6))
ax.plot(range(1, L + 1), logit_true, 's-', color='#d73027', lw=2.5, ms=7, label='logit lens (rotated)')
ax.plot(range(1, L + 1), tuned_true, 'o-', color='#1a9850', lw=2.5, ms=7, label='tuned lens (rotated)')
ax.axhline(1 / V, color='gray', ls=':', label='chance')
ax.set_title('Tuned lens recovers the gradual prediction under rotation')
ax.set_xlabel('layer'); ax.set_ylabel('P(true class)'); ax.legend(); ax.set_ylim(0, 1.02)
plt.tight_layout(); plt.show()
What you see¶
The tuned lens (green) restores the smooth, gradually-rising trajectory on the same rotated states where the naive logit lens (red) saw only noise. Each learned $A_\ell$ has effectively undone its layer's rotation $R_\ell$, mapping the hidden state back into the unembedding's basis before decoding.
This is the tuned lens's whole point: by learning the per-layer translation instead of assuming a shared basis, it is reliable and unbiased across architectures where the logit lens fails. The cost is a one-time, lightweight training of the affine translators (the model itself stays frozen).
Experiment 4 (limitation): the tuned lens can look ahead¶
Question: Does a tuned lens report what the model has decided, or what it can predict?
Because the translator is trained to match the final answer, it will surface any faint linear clue
of that answer — even at a layer where the model has barely computed. We compare the tuned-lens
accuracy when a weak early cue is present (early_cue=0.3) versus absent (early_cue=0.0).
def tuned_lens_accuracy(early_cue):
Htr, ytr = make_data(2500, 70, rotate=True, early_cue=early_cue)
pf = softmax_np(decode_np(Htr[-1]))
trans = [train_translator(Htr[l], pf) for l in range(L)]
Hte, yte = make_data(1500, 71, rotate=True, early_cue=early_cue)
acc = []
for l in range(L):
A, b = trans[l]
with torch.no_grad():
pred = decode_torch(torch.tensor(Hte[l]) @ A.t() + b).argmax(1).numpy()
acc.append((pred == yte).mean())
return np.array(acc)
acc_cue = tuned_lens_accuracy(early_cue=0.3) # faint cue present from layer 1
acc_nocue = tuned_lens_accuracy(early_cue=0.0) # no early cue
print('with early cue:', np.round(acc_cue, 2))
print('no early cue :', np.round(acc_nocue, 2))
print(f'layer-1 gap = {acc_cue[0]-acc_nocue[0]:+.2f} (cue surfaced before the model computes)')
with early cue: [0.3 0.37 0.42 0.46 0.62 0.72 0.82 0.9 ] no early cue : [0.11 0.14 0.18 0.26 0.44 0.6 0.77 0.9 ] layer-1 gap = +0.19 (cue surfaced before the model computes)
fig, ax = plt.subplots(figsize=(7.4, 4.6))
ax.plot(range(1, L + 1), acc_cue, 'o-', color='#d73027', lw=2.5, ms=7, label='with faint early cue')
ax.plot(range(1, L + 1), acc_nocue, 's-', color='#1a9850', lw=2.5, ms=7, label='no early cue')
ax.fill_between(range(1, L + 1), acc_nocue, acc_cue, where=(acc_cue > acc_nocue),
color='#d73027', alpha=0.12, label='look-ahead from the cue')
ax.axhline(1 / V, color='gray', ls=':', label='chance')
ax.set_title('Tuned lens surfaces a decodable cue before the model commits')
ax.set_xlabel('layer'); ax.set_ylabel('tuned-lens accuracy'); ax.legend(fontsize=9); ax.set_ylim(0, 1.02)
plt.tight_layout(); plt.show()
What you see¶
Both runs reach the same high accuracy at the final layer, but they split sharply early: with a faint linear cue present from layer 1 (red) the tuned lens already predicts well above chance, while without it (green) it sits at chance until the model's computation ramps up. The red band is the look-ahead the cue buys.
The catch: the tuned lens is trained to predict the final answer, so it surfaces any decodable cue immediately — you cannot tell from the lens alone whether high early confidence means "the model decided here" or merely "the answer is linearly readable here". This is the lens version of the probe trap in linear probing: decodable is not used, and only causal intervention can disambiguate.
Experiment 5: the prediction trajectory flags anomalies¶
Question: Can the shape of the layer-by-layer trajectory reveal abnormal inputs?
Belrose et al. find the trajectory of latent predictions detects anomalous inputs. Normal inputs commit smoothly to one class; we craft anomalous ones whose mid-layers lurch confidently toward a wrong class before recovering. We score each input by its peak confidence in a non-final class — how strongly the trajectory ever committed to something other than its own final answer.
NZ = 0.2 # cleaner trajectories make the detour stand out
Htr5, _ = make_data(2500, 30, rotate=True, noise=NZ)
pf5 = softmax_np(decode_np(Htr5[-1]))
trans5 = [train_translator(Htr5[l], pf5) for l in range(L)]
def tuned_probs5(Hl, l):
A, b = trans5[l]
with torch.no_grad():
return torch.softmax(decode_torch(torch.tensor(Hl) @ A.t() + b), 1).numpy()
def wrong_class_confidence(H):
P = np.stack([tuned_probs5(H[l], l) for l in range(L)], 1) # (n, L, V)
c_star = P[:, -1].argmax(1) # the final predicted class
score = np.zeros(len(c_star))
for l in range(L - 1): # every layer except the final
am, mx = P[:, l].argmax(1), P[:, l].max(1)
score = np.maximum(score, np.where(am != c_star, mx, 0.0)) # peak confidence in a non-final class
return score
Hnorm, _ = make_data(1500, 40, rotate=True, noise=NZ)
Hanom, yanom = make_data(1500, 41, rotate=True, noise=NZ)
wrong = (yanom + 3) % V
for l in range(2, 5): # inject a confident wrong-class spike mid-stream
Hanom[l] = (Hanom[l] + (1.3 * W_U[wrong]) @ R[l].T).astype(np.float32)
score_norm = wrong_class_confidence(Hnorm)
score_anom = wrong_class_confidence(Hanom)
thr = np.percentile(score_norm, 95)
print(f'anomaly detection rate at 5% false-positive = {(score_anom > thr).mean():.2f}')
anomaly detection rate at 5% false-positive = 0.99
fig, ax = plt.subplots(figsize=(7.4, 4.6))
ax.hist(score_norm, bins=40, color='#1a9850', alpha=0.7, label='normal inputs')
ax.hist(score_anom, bins=40, color='#d73027', alpha=0.7, label='anomalous inputs')
ax.axvline(thr, color='k', ls='--', label='95th-pct normal threshold')
ax.set_title('A confident wrong-class detour flags anomalous inputs')
ax.set_xlabel('peak confidence in a non-final class'); ax.set_ylabel('count'); ax.legend()
plt.tight_layout(); plt.show()
What you see¶
Normal inputs rarely commit confidently to any class other than their final answer, so their score stays low. The anomalous inputs — whose mid-layers lurch confidently toward a wrong class before recovering — score high, and the two distributions separate cleanly above the normal threshold (near-total detection at a 5% false-positive rate).
The signal lives in the trajectory, not the final output: a model that ends on the right answer can still pass through a confidently-wrong intermediate state, and only the lens reveals it. Belrose et al. use exactly this — abnormal latent-prediction trajectories — to flag malicious inputs; in Latent-Anything it is a cheap per-input health check the last layer alone could never provide.
Summary / Key Takeaways¶
- Exp 1 — The logit lens decodes intermediate states and shows the prediction forming gradually across layers — when their basis matches the final unembedding.
- Exp 2 — A per-layer basis rotation breaks the naive logit lens (flat, near-chance) — why it fails on GPT-Neo / BLOOM / OPT.
- Exp 3 — The tuned lens learns an affine translator per layer (KL to the final output) and recovers the trajectory on the rotated states.
- Exp 4 — Trained on the final answer, a tuned lens surfaces a decodable cue before the model commits — decodable is not used, the probe trap again.
- Exp 5 — A confident wrong-class detour in the trajectory flags anomalous inputs (near-total detection at 5% FPR) — signal the last layer alone cannot give.