feat(backward_lens): generalize Backward Lens to dense-MLP decoder-only Bridges (GPT-2, Pythia, GPT-NeoX)

[janmenjayap]

---

Description

Backward Lens previously accepted only raw GPT-2 TransformerBridge models, gated by an explicit
isinstance(model.adapter, GPT2ArchitectureAdapter) check and a Conv1D-only projection guard.
This PR removes both model-name/class conditionals and replaces them with a capability contract:
any dense (non-gated), decoder-only TransformerBridge whose MLP LinearBridge projections the
Bridge's own weight_layout_in_out oracle can orient is now supported.

The oracle resolves each projection's weight storage order per instance rather than per model
class:

• Conv1D ([in, out], e.g. GPT-2) resolves to "in_out".
• torch.nn.Linear ([out, in], e.g. Pythia/GPT-NeoX) resolves to "out_in".
• Anything the oracle cannot orient is rejected with a clear ValueError.

The existing _build_linear_gradient_factors helper already handled both layouts correctly; only
the discovery and validation path was GPT-2-specific. No public API name changes, and GPT-2 behavior
is unchanged (verified below).

Fixes #1777

---

Motivation and context

Backward Lens's gradient-factorization math (imprint / shift interpretation, exact outer-product
reconstruction) has nothing to do with GPT-2 specifically — it only depends on the MLP being dense
and two-matrix, and on knowing each matrix's storage orientation. Hardcoding the GPT-2 adapter and
Conv1D meant every other dense-MLP decoder-only architecture already supported by
TransformerBridge (Pythia, GPT-J, GPT-NeoX, OPT, ...) was rejected outright, even though nothing
about the tool's math required that restriction.

This generalizes the capability boundary to match what the Bridge itself can already tell us,
instead of adding a second, parallel per-model-family conditional alongside the Bridge's adapter
system.

---

Implementation details

• Rename _require_raw_gpt2_bridge -> _require_raw_dense_mlp_bridge; drop the
  GPT2ArchitectureAdapter isinstance check. All other guards are unchanged: TransformerBridge
  only, not compatibility_mode, not _weights_processed, single-device, non-gated MLP, a
  tokenizer, and the standard blocks/ln_final/unembed components.
• Rename _get_gpt2_mlp_projections -> _get_dense_mlp_projections; replace the
  isinstance(projection.original_component, Conv1D) guard with
  weight_layout_in_out(projection) from
  transformer_lens.model_bridge.generalized_components.mlp. Introduce a private _MLPLinear
  record carrying each projection alongside its resolved weight_layout, and make the expected
  weight shape orientation-aware ((in, out) for "in_out", (out, in) for "out_in") so both
  storage orders validate against the same fixed in/out feature counts.
• Rename _capture_gpt2_mlp_gradient_factors -> _capture_dense_mlp_gradient_factors,
  _GPT2GradientCapture -> _DenseMLPGradientCapture, and _GPT2LayerGradientFactors ->
  _MLPLayerGradientFactors. Thread each projection's resolved weight_layout into
  _build_linear_gradient_factors instead of the previous hardcoded "in_out".
• Update BackwardLens and module/class docstrings to describe the dense-MLP decoder-only
  contract instead of GPT-2 specifically.
• Repurpose test_capture_rejects_a_non_conv1d_component to assert rejection of an unorientable
  component, since torch.nn.Linear projections are now accepted rather than rejected.

---

Supported scope

This generalization supports:

• Any raw, unprocessed, single-device TransformerBridge with a dense, non-gated, two-matrix MLP
  whose LinearBridge projections wrap either Conv1D or torch.nn.Linear.
• The same one-unbatched-prompt, one-single-token-target, user-selected-layers scope as before.

Still out of scope, deferred to separate follow-up work:

• Gated MLPs (gate/up/down architectures need a different mathematical contract).
• Batched prompts and multiple-token targets.
• Model families whose MLP projection the weight-layout oracle cannot orient.
• Applying optimizer steps or editing model weights.

---

Documentation and demonstration

• docs/source/content/backward_lens.md: generalizes the introduction, gradient-factorization
  section (explains both Conv1D [in, out] and nn.Linear [out, in] storage, and that the
  tool reads orientation from the Bridge component, not the model class), and requirements/
  troubleshooting sections. Keeps the GPT-2 example and adds a minimal Pythia-70m example.
• demos/Backward_Lens_Demo.ipynb: adds a short section stating the generalized contract plus one
  small single-layer Pythia-70m reconstruction cell, without duplicating the full GPT-2 walkthrough.

---

Dependencies

No new dependencies are required.

---

Validation

Observed local validation on this branch:

• Unit tests (tests/unit/tools/test_backward_lens.py): 55 passed
• Integration tests (tests/integration/test_backward_lens.py, GPT-2 and Pythia-70m parametrized):
  35 passed
• uv run mypy .: Success: no issues found in 398 source files
• make format: no changes
• make test-pr (unit + docstring + acceptance + integration, full repository): 1462 passed,
  22 skipped, 197 deselected, 1 xfailed; one unrelated pre-existing failure,
  test_granite_eager_scan_device_correctness[mps] in
  tests/integration/model_bridge/test_granite_moe_hybrid_adapter.py (Granite MoE Hybrid SSM
  eager-scan vs. fused-scan numerics on MPS), confirmed unrelated to this change and predating this
  branch.

---

Type of change

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

---

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
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

---