Backward Lens: gradient factors in vocabulary space

This notebook implements Backward Lens in a few short cells on top of TransformerBridge hooks and logit_readout, applies it to a small, fixed set of GPT-2 next-token targets, then repeats one reconstruction on Pythia-70m to show the same code handles both MLP weight layouts. It reproduces the method of projecting the forward inputs and backward signals of an MLP weight gradient into vocabulary space; it is not a paper-scale replication and makes no causal claim from token rankings alone. All examples and layers are declared before analysis, and all outcomes are shown without filtering.

The method follows Katz et al. (2024). Gradient signs below are raw d(loss) / d(tensor) signs; gradient descent subtracts them.

Geometry

For one linear projection and one prompt, let \(x_i\) be the projection’s input at token position \(i\) and \(\delta_i = \partial L / \partial y_i\) the loss gradient at its output. The weight gradient is a sum of one outer product per position:

\[\nabla_W L = \sum_i x_i \delta_i^\mathsf{T} = X^\mathsf{T}\Delta.\]

That is the gradient for [in, out] storage (Conv1D, as in GPT-2); torch.nn.Linear (Pythia/GPT-NeoX) stores [out, in], so its gradient is the transpose. A dense MLP has two projections, and each has exactly one residual-width factor that can be read in vocabulary space:

Projection

Bridge hook points

Weight shape ([in, out])

Residual-width factor

FF1 (mlp.in, GPT-2 c_fc)

blocks.{l}.mlp.in.hook_in, .hook_out

[d_model, d_mlp]

forward input \(x_i\): the post-ln_2 residual, gain and bias included

FF2 (mlp.out, GPT-2 c_proj)

blocks.{l}.mlp.out.hook_in, .hook_out

[d_mlp, d_model]

backward signal \(\delta_i\) at the MLP output

FF1 rows describe what entered the MLP; FF2 rows are raw loss gradients at its output. They are different quantities and should not be read as one signal.

Reading a factor in vocabulary space

Each residual-width row \(v\) is read as \(P(v) = \operatorname{Unembed}(\operatorname{LN}_{final}(v))\) with logit_readout, which recomputes the final norm’s statistics for every row instead of reusing the scales cached during the forward pass. The values are signed, pre-softmax logits, not probabilities. The Normalized Logit Lens reads \(P(v / \lVert v\rVert_2)\) instead. The final LayerNorm has an epsilon and a bias, so raw and normalized readouts can differ, mostly for rows with very small norms such as backward signals; exact-zero rows stay zero.

Gradient descent applies \(W \leftarrow W - \eta \nabla_W L\), so for FF2 the smallest raw-gradient logits and ascending target ranks are the ordering relevant to an update. Negating projected logits is not equivalent: the final norm’s bias and epsilon make the readout not exactly sign-symmetric.

Restrictions

The implementation below is deliberately narrow:

  • Dense, non-gated MLP. It factors the two projections of a dense MLP (GPT-2, Pythia/GPT-NeoX). A gated MLP (Llama, Mistral, Qwen, Gemma) has a third, gate projection that is not handled here. Models whose MLP runs on flattened [batch * position, width] tensors (OPT) or that have no MLP (Mamba, RWKV) do not fit either.

  • Raw Bridge. Gradients are taken for the original weights, the ones training would update. Compatibility mode folds norms into the weights, which changes both the factors and the weights they belong to, so it stays off.

  • Evaluation mode. Dropout in training mode would perturb the gradients. model.eval() disables it without disabling autograd, so do not wrap the analysis in torch.no_grad().

  • One prompt, one target token. The loss is next-token cross-entropy at the final position against a single token. With GPT-2’s tokenizer a leading space usually matters, and the prompt’s own tokenization (including any BOS) defines the positions.

[1]:

# NBVAL_IGNORE_OUTPUT
import gc
import os

os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F
from huggingface_hub.utils import logging as hf_hub_logging
from IPython.display import display
from transformers.utils import logging as hf_logging

from transformer_lens.model_bridge import TransformerBridge
from transformer_lens.tools.analysis import logit_readout

# Loading messages and progress bars would otherwise land in the checked cell outputs.
hf_hub_logging.set_verbosity_error()
hf_logging.set_verbosity_error()
hf_logging.disable_progress_bar()

DEVICE = "cpu"  # Stable saved outputs and portable notebook validation.
DETAILED_LAYERS = (0, 2, 4, 6, 8, 10, 11)
EXAMPLES = (
    ("The capital of France is", " Paris"),
    ("The capital of Italy is", " Rome"),
    ("The largest planet in the Solar System is", " Jupiter"),
)

pd.DataFrame(EXAMPLES, columns=["prompt", "one-token target"])

[1]:
prompt one-token target
0 The capital of France is Paris
1 The capital of Italy is Rome
2 The largest planet in the Solar System is Jupiter

Load a raw GPT-2 Bridge

Boot the Bridge without compatibility mode or weight processing, then switch it to evaluation mode. CPU float32 keeps the saved outputs portable; change DEVICE for interactive exploration.

[2]:

model = TransformerBridge.boot_transformers("gpt2", device=DEVICE, dtype=torch.float32)
model.eval()
print(
    f"Loaded a raw GPT-2 TransformerBridge: {model.cfg.n_layers} layers, "
    f"d_model={model.cfg.d_model}, d_mlp={model.cfg.d_mlp}."
)

Loaded a raw GPT-2 TransformerBridge: 12 layers, d_model=768, d_mlp=3072.

Capture the factors

mlp_gradient_factors captures both sides of every requested projection with forward hooks, runs one forward pass, and calls torch.autograd.grad once for the loss’s gradient with respect to the captured projection outputs (the \(\delta_i\)) and the projection weights. The weight gradients serve only as a check: the outer-product sum must reproduce what autograd computes for the weight itself. Nothing is written to any parameter’s .grad, and run_with_hooks removes the hooks when the forward pass returns.

vocab_logits is the readout: logit_readout with no token selector returns the full-vocabulary logits of each row.

[3]:

PROJECTIONS = ("in", "out")
SIDES = ("hook_in", "hook_out")


def mlp_gradient_factors(model, prompt, target, layers):
    """Rank-1 factors of each requested MLP weight gradient for one next-token loss."""
    if model.compatibility_mode:
        raise ValueError("Backward Lens reads the original weights; boot a raw Bridge")
    if any(module.training for module in model.modules()):
        raise ValueError("call model.eval() first: dropout would perturb the gradients")
    if any(getattr(model.blocks[layer].mlp, "gate", None) is not None for layer in layers):
        raise ValueError("a gated MLP has a third projection that this code does not factor")
    target_ids = model.to_tokens(target, prepend_bos=False)[0]
    if target_ids.numel() != 1:
        raise ValueError(f"{target!r} is {target_ids.numel()} tokens; the loss takes one")
    sites = [(layer, proj) for layer in layers for proj in PROJECTIONS]
    captured = {}

    def capture(tensor, hook):
        captured[hook.name] = tensor

    names = [f"blocks.{layer}.mlp.{proj}.{side}" for layer, proj in sites for side in SIDES]
    tokens = model.to_tokens(prompt)
    with torch.enable_grad():
        logits = model.run_with_hooks(tokens, fwd_hooks=[(name, capture) for name in names])
        loss = F.cross_entropy(logits[0, -1:], target_ids.to(logits.device))
        outputs = [captured[f"blocks.{layer}.mlp.{proj}.hook_out"] for layer, proj in sites]
        projections = [getattr(model.blocks[layer].mlp, proj) for layer, proj in sites]
        weights = [projection.original_component.weight for projection in projections]
        grads = torch.autograd.grad(loss, outputs + weights)

    factors = {}
    for (layer, proj), delta, weight_grad in zip(sites, grads[: len(sites)], grads[len(sites) :]):
        x, delta = captured[f"blocks.{layer}.mlp.{proj}.hook_in"][0].detach(), delta[0].detach()
        outer = x.T @ delta
        # d_model != d_mlp, so the gradient's shape tells Conv1D [in, out] from Linear [out, in].
        layout = "in_out" if weight_grad.shape == outer.shape else "out_in"
        rebuilt = outer if layout == "in_out" else outer.T
        factors[layer, proj] = {
            "residual": x if proj == "in" else delta,
            "layout": layout,
            "relative_error": float((rebuilt - weight_grad).norm() / weight_grad.norm()),
        }
    return {"tokens": tokens[0].cpu(), "target_id": int(target_ids), "factors": factors}


def vocab_logits(model, rows, normalized=False):
    """Signed logits of residual-width rows read through the final norm and unembedding."""
    if normalized:
        norms = rows.norm(dim=-1, keepdim=True)
        rows = rows / torch.where(norms == 0, torch.ones_like(norms), norms)
    return logit_readout(model, rows).values


def ascending_rank(logits, token_id):
    """Zero-based rank counted from the smallest logit; tied logits share a rank."""
    return (logits < logits[..., token_id, None]).sum(dim=-1)


def decode(model, ids):
    return [model.tokenizer.decode([int(i)]).replace("\n", "\\n") for i in ids]

[4]:

prompt, target = EXAMPLES[0]
primary = mlp_gradient_factors(model, prompt, target, DETAILED_LAYERS)
position_tokens = decode(model, primary["tokens"])
worst_error = max(f["relative_error"] for f in primary["factors"].values())
print(f"Prompt positions: {position_tokens}")
print(f"Target {target!r} is token id {primary['target_id']}.")
print(
    f"Factored {len(primary['factors'])} weight gradients "
    f"(layouts: {sorted({f['layout'] for f in primary['factors'].values()})})."
)
print(f"Outer-product sums match autograd's weight gradients to 1e-5: {worst_error < 1e-5}")

Prompt positions: ['<|endoftext|>', 'The', ' capital', ' of', ' France', ' is']
Target ' Paris' is token id 6342.
Factored 14 weight gradients (layouts: ['in_out']).
Outer-product sums match autograd's weight gradients to 1e-5: True

Vocabulary-facing factor tables

FF1 top tokens describe the forward residual directions entering each MLP. FF2 bottom tokens use the ascending raw-gradient ordering relevant to a subtracted gradient-descent update. The layer-by-position table keeps the two factors separate. The final-position comparison shows how unit normalization changes the FF2 ordering, next to the original VJP (vector-Jacobian product) norm. None of these tables are model predictions.

[5]:

def joined(ids):
    return ", ".join(decode(model, ids))


layer_position_rows, norm_comparison_rows = [], []
for layer in DETAILED_LAYERS:
    ff1_top = vocab_logits(model, primary["factors"][layer, "in"]["residual"]).topk(3).indices
    delta = primary["factors"][layer, "out"]["residual"]
    ff2 = vocab_logits(model, delta)
    ff2_bottom = ff2.topk(3, largest=False).indices
    for position, token in enumerate(position_tokens):
        row = (layer, position, token, joined(ff1_top[position]), joined(ff2_bottom[position]))
        layer_position_rows.append(row)
    normalized = vocab_logits(model, delta[-1], normalized=True)
    # fp32 backends disagree in the fifth significant digit, so the norm shows three.
    norm_comparison_rows.append(
        (
            layer,
            f"{float(delta[-1].norm()):.2e}",
            joined(ff2[-1].topk(5, largest=False).indices),
            joined(normalized.topk(5, largest=False).indices),
        )
    )
position_columns = ["layer", "position", "prompt token", "FF1 top raw directions"]
display(
    pd.DataFrame(
        layer_position_rows, columns=[*position_columns, "FF2 bottom raw (update ordering)"]
    )
)
norm_columns = ["layer", "VJP norm", "raw FF2 bottom tokens", "normalized FF2 bottom tokens"]
display(pd.DataFrame(norm_comparison_rows, columns=norm_columns))
layer position prompt token FF1 top raw directions FF2 bottom raw (update ordering)
0 0 0 <|endoftext|> <|endoftext|>, Advertisements, \n %], Thumbnail, 女
1 0 1 The The, the, The ixel, Medium, iop
2 0 2 capital capital, capital, Capital inventoryQuantity, catentry, }}}
3 0 3 of of, of, Of Archdemon, ortium, Difference
4 0 4 France France, France, French anton, Seasons, figure
5 0 5 is is, also, was adr, adian, umn
6 2 0 <|endoftext|> ,, the, \n ometric, Associates, 76561
7 2 1 The latest, following, Official orescent, oro, ulhu
8 2 2 capital capital, punishment, penalty illary, fig, Variant
9 2 3 of the, course, interest uably, iqueness, uments
10 2 4 France Connection, France, 's Emerald, baugh, DIT
11 2 5 is now, not, also EStreamFrame, ORPG, Dise
12 4 0 <|endoftext|> ,, the, and OTOS, ë, istg
13 4 1 The resa, latest, Bottom ��, Eth, �
14 4 2 capital capital, punishment, city nings, BILITIES, cles
15 4 3 of the, Reach, course motivating, Appearance, Cups
16 4 4 France 's, ,, Telecom ��, Antar, Explosive
17 4 5 is now, not, also Diseases, Firstly, Feet
18 6 0 <|endoftext|> ,, ., and SON, leys, VAL
19 6 1 The resa, oret, Latest �, ��, esters
20 6 2 capital city, capital, cities appings, UES, =-=-
21 6 3 of the, State, British ヘラ, TABLE, �
22 6 4 France 's, ,, and , /_, ciation
23 6 5 is now, officially, not Tycoon, 76561, Osiris
24 8 0 <|endoftext|> ,, \n, . Paris, xes, SON
25 8 1 The resa, oret, following Notting, ournals, EngineDebug
26 8 2 capital city, capital, cities Paris, Denver, Philadelphia
27 8 3 of South, British, China Frames, Reviewer, selage
28 8 4 France 's, ,, has ciation, ㅋㅋ, Spac
29 8 5 is now, officially, still Paris, Paris, Cologne
30 10 0 <|endoftext|> \n, the, , Paris, ゼウス, Hollande
31 10 1 The resa, oret, latest �醒, ��, ���
32 10 2 capital city, cities, capital Paris, ��, �
33 10 3 of China, Kazakhstan, Ukraine �醒, ���, Paris
34 10 4 France has, is, 's cffff, eteenth, mathemat
35 10 5 is now, officially, rapidly Paris, Paris, Copenhagen
36 11 0 <|endoftext|> The, Still, \n �, �, 
37 11 1 The resa, odore, latest �, �, 
38 11 2 capital city, cities, Manila �, �, 
39 11 3 of Romania, Kazakhstan, Armenia �, �, 
40 11 4 France has, is, 's �, �, 
41 11 5 is undergoing, reeling, hosting Paris, Paris, Copenhagen
layer VJP norm raw FF2 bottom tokens normalized FF2 bottom tokens
0 0 6.56e-01 adr, adian, umn, itia, igmatic adr, adian, umn, itia, igmatic
1 2 1.44e+00 EStreamFrame, ORPG, Dise, misconception, senal EStreamFrame, ORPG, Dise, misconception, senal
2 4 1.33e+00 Diseases, Firstly, Feet, heses, Differences Diseases, Firstly, Feet, heses, Differences
3 6 1.04e+00 Tycoon, 76561, Osiris, Modes, Memories Tycoon, 76561, Osiris, Modes, Memories
4 8 4.87e-01 Paris, Paris, Cologne, Copenhagen, Rome Paris, Paris, Cologne, Copenhagen, Rome
5 10 2.75e-01 Paris, Paris, Copenhagen, Cologne, Kyoto Paris, Paris, Copenhagen, Cologne, Kyoto
6 11 2.61e-01 Paris, Paris, Copenhagen, Marse, France Paris, Paris, Copenhagen, Marse, France

Target-token ranks: raw versus Normalized Logit Lens

Ranks are zero-based and ascending in raw-gradient space: rank 0 is the smallest projected FF2 gradient logit at the final position. Lower is therefore the update-relevant ordering, but this remains a directional diagnostic rather than a prediction of an optimizer step. Ranks are shown as vocabulary percentiles at one decimal.

[6]:

# Exact ranks can move by one across numeric backends when logits nearly tie; one-decimal
# vocabulary percentiles do not, so only those are displayed.
def final_position_ranks(result, layers):
    """Ascending target ranks among the final position's FF2 logits, raw and normalized."""
    rows = []
    for layer in layers:
        delta = result["factors"][layer, "out"]["residual"][-1]
        logits = [vocab_logits(model, delta, normalized=n) for n in (False, True)]
        rows.append([layer, *(int(ascending_rank(row, result["target_id"])) for row in logits)])
    ranks = pd.DataFrame(rows, columns=["layer", "raw rank", "normalized rank"])
    for kind in ("raw", "normalized"):
        percentile = 100 * ranks[f"{kind} rank"] / model.cfg.d_vocab
        ranks[f"{kind} target rank percentile"] = percentile.round(1)
    return ranks


ranks = final_position_ranks(primary, DETAILED_LAYERS)
display(ranks[["layer", "raw target rank percentile", "normalized target rank percentile"]])

fig, ax = plt.subplots(figsize=(7, 4))
for kind, marker in (("raw", "o"), ("normalized", "s")):
    ax.plot(ranks["layer"], ranks[f"{kind} target rank percentile"], marker=marker, label=kind)
ax.set(
    xlabel="layer",
    ylabel="ascending target-rank percentile (lower is smaller)",
    title=f"FF2 target-rank percentiles for {target!r}",
)
ax.grid(alpha=0.25)
ax.legend()
plt.show()
layer raw target rank percentile normalized target rank percentile
0 0 80.9 80.9
1 2 41.3 41.4
2 4 11.2 11.2
3 6 4.4 4.4
4 8 0.0 0.0
5 10 0.0 0.0
6 11 0.0 0.0
../../_images/generated_demos_Backward_Lens_Demo_12_1.png

VJP norm heatmap

FF2 factor norms are the norms of the backward signals at each MLP output. The logarithmic heatmap shows both position and layer variation without discarding low-norm rows. The final layer’s MLP output at an earlier position never reaches the final position’s logits, because no attention layer follows it, so those rows are exactly zero and the final-layer FF2 gradient has rank one.

[7]:

vjp_norms = np.stack(
    [primary["factors"][layer, "out"]["residual"].norm(dim=-1).numpy() for layer in DETAILED_LAYERS]
)
zero_mask = vjp_norms == 0
masked_log_norms = np.ma.masked_where(zero_mask, np.log10(np.where(zero_mask, 1.0, vjp_norms)))
cmap = plt.colormaps["magma"].copy()
cmap.set_bad(color="#d1d5db")

fig, ax = plt.subplots(figsize=(10, 4.5))
image = ax.imshow(masked_log_norms, aspect="auto", cmap=cmap)
ax.set_yticks(range(len(DETAILED_LAYERS)), labels=DETAILED_LAYERS)
ax.set_xticks(range(len(position_tokens)), labels=position_tokens, rotation=35, ha="right")
for layer_index, position_index in np.argwhere(zero_mask):
    ax.text(position_index, layer_index, "0", ha="center", va="center")
ax.set(xlabel="prompt position", ylabel="layer", title="FF2 VJP norms (gray = exactly zero)")
fig.colorbar(image, ax=ax, label="log10 L2 norm")
fig.tight_layout()
plt.show()
final_rows_zero = bool(zero_mask[-1, :-1].all())
print(f"Final-layer FF2 rows before the last position are exactly zero: {final_rows_zero}")
../../_images/generated_demos_Backward_Lens_Demo_14_0.png 
Final-layer FF2 rows before the last position are exactly zero: True

Predeclared examples: report every outcome

The remaining examples are analyzed only at GPT-2’s final layer to keep the demo compact. The table reports one-decimal vocabulary rank percentiles so nearly tied logits remain stable across numeric backends. It also records whether unit normalization improves, worsens, or leaves unchanged the target’s exact ascending rank. No example is removed based on its result.

[8]:

final_layer = DETAILED_LAYERS[-1]
OUTCOMES = {1: "improved", 0: "unchanged", -1: "worsened"}
comparison_rows = []
for example_prompt, example_target in EXAMPLES:
    result = mlp_gradient_factors(model, example_prompt, example_target, [final_layer])
    summary = final_position_ranks(result, [final_layer]).iloc[0]
    outcome = OUTCOMES[int(np.sign(summary["raw rank"] - summary["normalized rank"]))]
    percentiles = summary[["raw target rank percentile", "normalized target rank percentile"]]
    comparison_rows.append((example_prompt, example_target, *percentiles, outcome))

comparison = pd.DataFrame(
    comparison_rows, columns=["prompt", "target", *percentiles.index, "normalization outcome"]
)
display(comparison)
print("Normalization outcomes:", comparison["normalization outcome"].value_counts().to_dict())
prompt target raw target rank percentile normalized target rank percentile normalization outcome
0 The capital of France is Paris 0.0 0.0 unchanged
1 The capital of Italy is Rome 0.0 0.0 unchanged
2 The largest planet in the Solar System is Jupiter 0.0 0.0 unchanged 
Normalization outcomes: {'unchanged': 3}

Beyond GPT-2: the other weight layout

The same functions run unchanged on any raw, non-gated dense-MLP Bridge. GPT-2’s Conv1D projections store [in, out]; Pythia/GPT-NeoX’s torch.nn.Linear projections store [out, in], so their weight gradient is the transpose of the outer-product sum. The single-layer reconstruction below repeats the first example on a raw Pythia-70m Bridge and reports the layout that the gradient’s shape resolved to, alongside the reconstruction check.

[9]:

pythia_model = TransformerBridge.boot_transformers(
    "EleutherAI/pythia-70m", device=DEVICE, dtype=torch.float32
)
pythia_model.eval()

pythia_factors = mlp_gradient_factors(pythia_model, prompt, target, [0])["factors"][0, "in"]
print(f"Pythia-70m input projection weight layout: {pythia_factors['layout']}")
print(
    "Pythia-70m reconstruction error within tolerance: "
    f"{pythia_factors['relative_error'] < 1e-5}"
)

del pythia_model
_ = gc.collect()

Pythia-70m input projection weight layout: out_in
Pythia-70m reconstruction error within tolerance: True

Interpretation checklist

  • FF1 and FF2 tables project different factors; do not compare them as if they were the same signal.

  • FF2 uses raw loss-gradient signs. Gradient descent subtracts the gradient, so ascending ranks are highlighted.

  • LayerNorm bias and epsilon make projection non-linear enough that negating projected logits is not an exact update readout.

  • Normalized Logit Lens isolates direction from factor norm, but it need not improve target rank. The table above records null and contradictory outcomes.

  • Small reconstruction error validates the factorization numerically; it does not validate a semantic interpretation.

  • Hooks already installed on the model still apply during the forward pass, so remove them to analyze the unmodified computation.

  • This fixed GPT-2 sample demonstrates the method and its diagnostics, not population-level evidence about language-model learning or memory.

References