[Proposal] Backward Lens: vocabulary readout of multilayer perceptron gradient factors

[janmenjayap]

Proposal

Add a tested Backward Lens analysis tool for TransformerBridge. The first implementation will expose the low-rank factors of the two multilayer perceptron weight gradients and project every residual-width factor into vocabulary space.

For a standard two-matrix multilayer perceptron, the tool will keep the paper's two interpretations separate:

1. First matrix, imprint: its gradient is spanned by forward inputs x_i. Vocabulary projection asks what information from the prompt is written into the first matrix.
2. Second matrix, shift: its gradient is spanned by backward Vector Jacobian product signals delta_i. Vocabulary projection asks which token directions the update shifts the second matrix toward or away from.

The first version will support one-token next-token losses, non-gated GPT-2 multilayer perceptrons, and raw models loaded with TransformerBridge.boot_transformers.

Motivation

TransformerLens provides forward vocabulary readouts and the Jacobian Lens, but it has no maintained tool for interpreting parameter gradients in vocabulary space. Backward Lens provides this missing view without training an auxiliary model or loading an artifact.

The method has an exact correctness invariant. For each token position, a linear layer's weight gradient is an outer product of its forward input and its output gradient. Summing those factors reconstructs the full weight gradient. This makes the implementation independently testable without copying the authors' unlicensed research repository.

The feature also fits the current TransformerLens direction. The Jacobian Lens established raw Bridge validation, autograd capture, vocabulary readout, and analysis modules under transformer_lens/tools/analysis/.

Pitch

Add BackwardLens as a self-contained analysis tool under
transformer_lens/tools/analysis/. It will use the existing TransformerBridge linear-layer
hooks and PyTorch autograd support while keeping gradient interpretation separate from model
editing.

Proposed interface

Names are open to maintainer preference.

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

model = TransformerBridge.boot_transformers("openai-community/gpt2", device="cpu")
lens = BackwardLens(model)

result = lens.analyze(
    "The Eiffel Tower is in the city of",
    target_token=" Paris",
    layers=[8, 9, 10],
)

result.first_matrix.forward_inputs
result.first_matrix.vocabulary_logits
result.second_matrix.backward_signals
result.second_matrix.top_tokens(k=5)
result.second_matrix.bottom_tokens(k=5)
result.max_reconstruction_error

target_token must encode to exactly one token. The loss is the negative log likelihood of that token under the logits at the final prompt position. Empty prompts, multiple-token targets, invalid layers, gated multilayer perceptrons, processed weights, and unsupported drivers raise clear errors.

Mathematical contract

For GPT-2 and TransformerLens weight orientation:

• W_in has shape [d_model, d_mlp] and computes pre = x @ W_in.
• W_out has shape [d_mlp, d_model] and computes out = hidden @ W_out.

For token positions i:

• grad_W_in = sum_i outer(x_i, grad_pre_i).
• grad_W_out = sum_i outer(hidden_i, grad_out_i).

The first matrix is interpreted through x_i, which is residual-width. The second matrix is interpreted through grad_out_i, which is also residual-width. The other factor controls each vector's coefficient in the full gradient.

The vocabulary projection follows the paper's logit lens. For each residual-width vector, apply the model's final normalization with statistics recomputed on that vector, then apply the unembedding. This is not the cached-scale linear attribution path.

For the second matrix, the raw backward signal commonly ranks the edit target near the bottom. Gradient descent subtracts the gradient. The result therefore exposes both top and bottom rankings and documents their signs. It does not assert that the target must appear in the raw signal's top tokens.

Implementation design

1. Require a raw causal TransformerBridge loaded through the transformers driver.
2. Register forward captures on blocks.{layer}.mlp.in.hook_in, blocks.{layer}.mlp.in.hook_out, blocks.{layer}.mlp.out.hook_in, and blocks.{layer}.mlp.out.hook_out.
3. Run one forward pass and compute the one-token next-token loss.
4. Use torch.autograd.grad to obtain gradients for the captured linear outputs and the corresponding original weight parameters. Do not call .backward(), mutate .grad, or freeze the selected weights.
5. Reconstruct each weight gradient from its token-position outer products in float32 and report absolute and relative error.
6. Detach residual-width factors, run the standard final-normalization and unembedding projection, and return raw vocabulary logits plus top and bottom token helpers.
7. Restore all temporary hooks and preserve parameter gradient state even when an exception occurs.

The first implementation should not use backward hooks on normalization points and should assert that no native-path fallback warning is emitted.

Validation plan

1. Exact factorization: synthetic linear layers reconstruct both weight gradients within a stated tolerance.
2. Orientation: asymmetric dimensions catch transposed outer products and swapped factors.
3. Projection parity: vocabulary logits match a direct call through the model's final normalization and unembedding on each detached factor.
4. Sign handling: a controlled example verifies raw top and bottom rankings and the sign of the gradient-descent update.
5. Target contract: zero-token and multiple-token targets are rejected.
6. State safety: parameter .grad values, requires_grad flags, hooks, and model weights are unchanged after success or failure.
7. Bridge integration: gpt2-small returns finite factors and reconstruction errors below tolerance without native-path fallback warnings.
8. Paper-level sanity report: report target ranks for a small fixed prompt set, but do not make a brittle semantic rank threshold a required unit test.

Initial pull request scope

• transformer_lens/tools/analysis/backward_lens.py
• public exports for BackwardLens and its result dataclasses
• model-free unit tests for factorization, orientation, sign, and validation
• gpt2-small TransformerBridge integration tests
• one demonstration notebook
• one documentation section

Non-goals

• gated multilayer perceptrons
• multiple-token or batched editing targets
• applying optimizer steps or editing model weights
• CounterFact replication
• Jacobian-space update tracing
• other model families
• external repository code as a dependency or oracle

These belong in a separate follow-up issue after the core contract is merged.

Alternatives

• Continue using ad hoc autograd scripts: this can recover the factors, but it leaves users
  responsible for tensor orientation, target-loss semantics, normalization, hook cleanup, and
  gradient-state safety. It also provides no shared result contract or regression tests.
• Add the implementation to jacobian_lens.py: the Jacobian Lens represents causal transport
  between residual streams, while Backward Lens analyzes the gradient induced by a particular loss.
  They share infrastructure but answer different questions, so separate modules are clearer.
• Use run_with_cache(incl_bwd=True): this is useful as a future independent cross-check, but
  the initial implementation needs gradients for selected intermediate tensors and weight
  parameters from one explicitly defined scalar loss. torch.autograd.grad provides that contract
  without mutating .grad.
• Depend on the authors' repository: the repository is a research demonstration with no
  declared license or maintained library interface. The implementation should instead follow the
  published equations and use the exact outer-product reconstruction as its primary oracle.
• Support every multilayer perceptron architecture immediately: gated architectures require a
  different factorization contract because their gate and up projections interact through a
  product and activation function. Starting with GPT-2 avoids presenting an incorrect universal
  abstraction.

Additional context

• Katz, Belinkov, Geva, and Wolf, "Backward Lens: Projecting Language Model Gradients into the Vocabulary Space," Empirical Methods in Natural Language Processing 2024, Best Paper, arXiv:2402.12865.
• Authors' research demonstration: shacharKZ/BackwardLens. The repository has no declared license, so no code will be copied, vendored, or installed as a test dependency.
• The existing Jacobian Lens provides a nearby TransformerBridge precedent for raw-model
  validation, autograd capture, vocabulary readout, public result dataclasses, integration tests,
  and documentation under transformer_lens/tools/analysis/.
• A separate follow-up proposal will cover CounterFact replication, causal update tracing,
  Jacobian-space comparisons, editing experiments, multiple-token targets, gated architectures,
  and broader model coverage.

Checklist

• I have checked that there is no similar issue in the repo (required)
• No exact Backward Lens or gradient-vocabulary issue or pull request was found in a fresh search.
• The first pull request is limited to a testable GPT-2 vertical slice.
• The first-matrix imprint and second-matrix shift interpretations are separated.
• The raw backward-signal ranking and gradient-descent sign are documented correctly.
• The correctness oracle is independent of unlicensed external code.

[jlarson4]

@janmenjayap Sounds good to me, I will assign this to you under the assumption you'd like to take it on.

[janmenjayap]

Thanks @jlarson4 — the implementation proposed here is now available in #1723.

Delivered by PR #1723

The PR implements the bounded GPT-2 vertical slice proposed in this issue:

• one raw, unprocessed GPT-2 TransformerBridge model;
• one unbatched prompt and one single-token next-token target;
• one forward pass and one torch.autograd.grad call;
• aligned factors and exact reconstruction checks for both MLP weight gradients;
• fresh final-normalization and unembedding readouts for every residual-width factor;
• explicit raw-gradient versus SGD-update signs, plus an opt-in Normalized Logit Lens;
• detached public results with state, hook, .grad, mode, and RNG preservation;
• unit/integration tests, reference documentation, and an executed demonstration notebook.

Focused local validation reports 38 unit tests, 23 integration tests, and all 10 notebook cells
passing. PR #1723 targets dev, and its dedicated Backward Lens notebook check is green. The current
Python 3.10–3.12 compatibility failures are limited to seven negative-input tests expecting
TypeError while CI's global instrumentation raises BeartypeCallHintParamViolation; the local
test-only compatibility fix passes all 38 Backward Lens unit tests.

1. Commits still remaining in PR #1723

There is one known corrective commit remaining before final review:

1. Update the negative-input unit tests to accept both ordinary TypeError and Beartype's runtime
   parameter-violation exception, then rerun CI.

After that commit, no additional feature, API, documentation, integration-test, or notebook commit is
planned for this PR. Any further commit should be limited to a concrete CI failure or maintainer review
request; the GPT-2 implementation scope itself is complete.

2. Other PRs required to complete Issue #1686

No additional PR is currently required to satisfy or close #1686. PR #1723 contains the complete
initial scope stated by the issue: implementation, public exports, model-free tests, GPT-2 Bridge
integration tests, reference documentation, and an executed demo. It retains Fixes #1686, so the
issue should close when PR #1723 is merged and its required checks are green.

The broader items listed below are explicit non-goals of #1686. They should not block this PR or keep
the issue open. If maintainers prefer to split #1723 during review, that would be a review-driven change
rather than a currently missing PR.

3. Proposed future Backward Lens issues

I suggest filing the following as separate, independently reviewable issues after the core PR is
accepted:

1. Broader dense-MLP support. Generalize capture beyond GPT-2 and validate on GPT-2 plus
   Pythia-70M, including device, dtype, and reconstruction-tolerance characterization.
2. Gated-MLP support. Define separate gate/up/down factorization mathematics and test it with a
   tiny Llama fixture, with Qwen2-0.5B as an optional real-model validation. The GPT-2 first-matrix
   contract must not be copied unchanged through a multiplicative gate.
3. Batched and multi-token objectives. Add per-example factors for batched prompts and explicit
   loss/reduction semantics for multiple-token targets.
4. Independent validation and bounded paper replication. Cross-check against
   run_with_cache(incl_bwd=True), broaden fixed-prompt reports without semantic rank thresholds, and
   reproduce selected plots on a pinned CounterFact subset with uncertainty and null results reported.
5. Generic activation/VJP tracing and causal-update analysis. Add a Bridge-native tracer for
   attention, residual, and MLP submodules; connect parameter-gradient factors to Jacobian-space or
   forward-pass activation changes using controlled interventions and matched random baselines.
6. Reversible model-editing API and benchmarks. Only with explicit maintainer demand, define an
   inspectable update object that never mutates the caller's model by default, then compare one-step
   SGD/Adam and gradient-inspired shifts with norm-matched controls. Full MEND/ROME/MEMIT benchmarks
   belong here, not in the core analysis PR.
7. Analysis UX. Add exportable tables and reusable plotting helpers after the result schema has
   stabilized.

Across those follow-ups, vocabulary readability should not be presented as causal evidence, GPT-2
behavior should not be assumed to transfer to gated or larger models, and no code should be copied,
vendored, or installed from the authors' unlicensed research repository.

Would you prefer these follow-ups to be filed as separate issues after #1723 is reviewed/merged, or
should they remain deferred until there is explicit maintainer demand?