feat(backward_lens): add vocabulary readout of GPT-2 MLP gradient factors

[janmenjayap]

Description

This PR implements Backward Lens as a self-contained analysis tool for raw GPT-2 models loaded
through TransformerBridge.

Backward Lens interprets the low-rank factors of MLP weight gradients by projecting residual-width
factors into vocabulary space. It preserves the two interpretations from Katz et al. (2024):

1. First matrix — imprint: the gradient is spanned by forward inputs x_i, describing
   information written into the first MLP matrix.
2. Second matrix — shift: the gradient is spanned by backward VJP signals delta_i, describing
   vocabulary directions toward or away from which the second MLP matrix is updated.

For every selected layer, the implementation reconstructs the independently obtained weight
gradients using:

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

Reconstruction errors and numerical ranks provide independently testable correctness invariants.

Motivation and context

TransformerLens already supports forward vocabulary readouts and the Jacobian Lens, but it did not
provide a maintained tool for interpreting parameter gradients in vocabulary space.

This feature adds that missing analysis surface without training an auxiliary model, loading a
learned artifact, applying optimizer steps, or depending on the authors' research repository.

Implementation details

• Adds BackwardLens under transformer_lens/tools/analysis/.
• Adds public, detached result dataclasses and package exports.
• Uses one forward pass and one torch.autograd.grad call.
• Does not call .backward(), mutate parameter .grad, or freeze selected parameters.
• Captures aligned position-wise factors for both GPT-2 MLP projections.
• Reconstructs gradients in float32 and reports absolute and relative errors.
• Projects residual-width factors through the live final normalization and unembedding.
• Recomputes normalization statistics independently for every projected factor.
• Supports raw and explicit opt-in Normalized Logit Lens projections.
• Retains original factor norms and marks exact-zero factors.
• Exposes signed top/bottom token rankings and target-rank diagnostics.
• Documents the distinction between raw gradients and the negative update direction used by
  gradient descent.
• Preserves weights, gradients, requires_grad, existing hooks, training mode, and RNG state.
• Removes temporary hooks after both successful and failing analyses.

Supported scope

This initial implementation supports:

• Raw, unprocessed GPT-2 TransformerBridge models.
• Models loaded through the transformers Bridge driver.
• One unbatched prompt.
• One single-token next-token target.
• Dense, non-gated GPT-2 MLPs.
• User-selected layers.

The following are intentionally deferred to separate follow-up issues:

• Gated MLPs.
• Additional model families.
• Batched prompts.
• Multiple-token targets.
• Applying optimizer steps or editing model weights.
• CounterFact replication.
• Jacobian-space update tracing.
• Paper-scale editing benchmarks.

Documentation and demonstration

This PR adds:

• A reference page covering the mathematical contract, tensor shapes, API, normalization, signs,
  restrictions, state safety, and troubleshooting.
• An executed GPT-2-small demonstration notebook showing:
  • Exact gradient reconstruction.
  • Numerical rank and reconstruction error.
  • Layer-by-position vocabulary directions.
  • VJP norms and target ranks.
  • Raw versus normalized projections.
  • Gradient-versus-update sign semantics.
  • Exact and near-zero signal handling.
  • Contribution-sorted cumulative reconstruction.
• CI and Makefile registration for notebook validation.

The notebook figures and tables were manually inspected for correctness and legibility.

Dependencies

No new dependencies are required.

Validation

Observed local validation:

• Unit tests: 38 passed
• Integration tests: 23 passed
• Notebook nbval: 10 passed
• uv lock --check: passed
• git diff --check upstream/dev...HEAD: passed
• Notebook: 10 executed code cells with no saved errors
• Branch rebased onto current upstream/dev

Focused validation passed in the established transformer-lens Conda environment.

Exact locked-.venv pytest execution was blocked by a local PyTorch/CUDA runtime symbol mismatch.
The documentation builder parsed the new Backward Lens page, but a complete local docs build was not
obtained because the environment lacked Pandoc and the build encountered model-registry/docstring
errors. Repository-wide checks remain subject to PR CI.

Fixes #1686

Type of change

• New feature (non-breaking change which adds functionality)
• This change requires a documentation update

Screenshots

Vocabulary-facing gradient factors

Reconstruction and rank diagnostics

Target-token ranks: raw versus Normalized Logit Lens

Checklist

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new Backward Lens-specific warnings
• I have added tests that prove my feature works
• New and existing unit tests pass locally with my changes
  • Focused Backward Lens unit tests pass; the complete repository-wide suite is pending CI.
• I have not rewritten tests relating to key interfaces which would affect backward compatibility