[Proposal] Relevance Lens (R-lens): a RelP/ LRP-based transport-matrix estimator for Jacobian Lens (J-lens) fit, readout, and intervention

[janmenjayap]

---

Proposal

Add Relevance Lens (R-lens) as an alternative estimator for the shipped Jacobian Lens. R-lens keeps J-lens's fitting loop and downstream readout unchanged, but computes each per-layer d_model × d_model transport matrix with a RelP/LRP-modified backward pass instead of the ordinary Jacobian vector–Jacobian product (VJP). Only backward semantics change; the forward pass is bit-identical to the model's native forward.

Scope decision: Bridge-only

HookedTransformer is deprecated as of TransformerLens 3.0 (emits DeprecationWarning, TransformerBridge is the default surface, transformers>=5.9.0 required, removal scheduled for the next major version). R-lens targets TransformerBridge exclusively and adds no HookedTransformer path.

Proposed API

Names are adjustable to maintainer preference.

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

model = TransformerBridge.boot_transformers(model_id, device="cuda")
lens = RelevanceLens.fit(model, prompts, corpus="pile-10k:fixed-manifest")

result = lens.readout(model, "The Eiffel Tower is in the city of")
hooks = lens.ablation_hooks(model, " Paris", layers=[6, 7])

RelevanceLens reuses J-lens's matrix-based readout, vocabulary vectors/dictionary, sparse decomposition, and steering/ablation/swap interventions unchanged. It differs from JacobianLens only in how the transport matrices are estimated and in the provenance recorded on the artifact.

---

Motivation

The source post reports that ordinary gradients accumulate noise at early layers, and that a small set of LRP rules produce earlier, cleaner, and more causally load-bearing readouts. Those findings motivate the estimator; they are not proposed as small-model CI assertions (the post reports no advantage on its smallest models).

---

Pitch

The three rules (first dense slice)

Applied exactly as in the source post; everything else stays on ordinary autograd.

1. LN-rule (residual-stream norms only). Preserve each residual LayerNorm/RMSNorm forward value while treating its denominator as a constant in the backward pass. Does not touch q/k norms or other attention-internal normalization.
2. Identity-rule (GELU/SiLU). Express f(x) = x · phi(x) and detach phi(x). The local VJP is grad_out · phi(x); for SiLU, phi(x) = sigmoid(x). It is not a bare identity VJP. For exact GELU, phi is the Gaussian CDF; approximate GELU variants use the forward-consistent factor f(x)/x with the analytic 1/2 limit at zero. Must fail closed on relu/relu2 gated MLPs (where f(x)/x = relu(x) ≠ 2·relu(x)), not silently mis-score them.
3. Half-rule (multiplicative MLP gates). Preserve u · v in the forward pass while using grad_u = 0.5 · grad_out · v and grad_v = 0.5 · grad_out · u — each branch retains the other branch's forward value; neither receives a bare 0.5 · grad_out.

Linear layers and attention retain ordinary autograd in this slice. The attention-specific AH-rule is out of scope.

Main correctness risk (named)

The raw Bridge surface cannot implement these rules with ordinary backward hooks:

• NormalizationBridge switches from the native Hugging Face forward to a Python reimplementation (_python_norm_forward) when backward hooks are attached to hook_scale/hook_normalized (transformer_lens/model_bridge/generalized_components/normalization.py:174-176), so a naive hook cannot keep the forward bit-identical. RMSNormalizationBridge defaults use_native_layernorm_autograd=True (transformer_lens/model_bridge/generalized_components/rms_normalization.py:38), so the LN-rule custom VJP must integrate on that native-autograd path without tripping the backward-hook fallback.
• Raw GatedMLPBridge runs the Hugging Face MLP as one opaque module (transformer_lens/model_bridge/generalized_components/gated_mlp.py:143), hiding the activation and multiplicative-product boundaries the Identity- and Half-rule need. The backend must recompute gate_output/up_output inside its own backward from the unconditional W_gate/W_in @property accessors inherited from MLPBridge (transformer_lens/model_bridge/generalized_components/mlp.py:163-197) — a checkpointing-style re-forward — and must not read self._processed_W_gate/_processed_W_in (those exist only in compatibility mode with LayerNorm folding and are absent on the raw path). The forward the caller sees must remain the opaque call's own native output, unchanged.

The backend must therefore be a component-owned custom torch.autograd.Function / scoped protocol that: returns the native forward result unchanged; changes only backward semantics; supports nested contexts and restores state after exceptions; reports exactly which requested rules were installed; and raises if a requested architecture path cannot be instrumented. Partial rule installation must never be silent. Do not use permanent model.cfg flags. Do not use .data.

Artifact identity and safety

A saved R-lens must record its estimator = "relevance_lens", relevance_rule_version, enabled_rules, rule_coverage, target_layer, and skip_first_positions. merge() must reject J-lens/R-lens mixtures and mismatched rule configurations. R-lens must not implicitly resolve the J-lens registry by short name; only explicit paths are accepted until a separate R-lens registry exists.

Validation plan

1. Analytic rule VJPs — denominator-detach (LN), detached activation factor (Identity), half-scaled product (Half); zeros, mixed signs, non-contiguous, fp32/fp64, batched.
2. Bit-identical forward — torch.equal for logits and selected intermediates with and without rules on every claimed architecture.
3. Tiny gated integration — fit J-lens and R-lens on a deterministic RMSNorm/SiLU/gated model; assert finite, shape-compatible, but at least one early-layer transport matrix that differs from J-lens.
4. Fit/merge consistency — joint fitting equals prompt-count-weighted merging of matching R-lens shards.
5. Artifact identity — mixing estimators/rule sets fails; R-lens loading cannot resolve the J-lens registry.
6. Coverage failure — an opaque unsupported gated MLP (or an unsupported relu/relu2 activation under Identity-rule) raises before fitting.
7. Reference mechanics — tolerant parity of isolated rule VJPs against FarnoushRJ/RelP pinned at 8219d6dc417c3fd7f318342cf61cd2a0c20b7250; record the known exact-zero stabilizer discrepancy.
8. Ordinary J-lens non-regression — the ordinary JacobianLens.fit path is numerically unchanged after the shared-driver refactor (pinned regression fixture).

Do not assert R-lens superiority on GPT-2 Small, and do not use a subjective real-word fraction as a correctness test. Large-model quality reproduction (pass@k, ablation effect sizes, probe upper bound, CKA banding) is deferred to separate research follow-up issues.

Relationship to RelP

RelP and R-lens are two issues that share one internal LN/Identity/Half custom-VJP backend but differ in every user-facing way:

RelP	R-lens
Clean/corrupted input pair	Prompt corpus
User-supplied scalar metric	One-hot residual-stream cotangents
Per-component attribution scores	Per-layer d_model × d_model transport matrices
Ephemeral analysis result	Saved/reused lens artifact

This issue hosts the shared backend (validated on a Qwen2 / JointGateUp CI fixture) so it is not blocked on RelP. RelP consumes the landed backend afterward.

Scope of this issue (3 PRs)

• PR1 — refactor the J-lens fitting driver so ordinary and relevance VJPs share one implementation, with the ordinary path numerically unchanged.
• PR2 — the shared LN/Identity/Half custom-VJP backend with fail-closed coverage, native-forward preservation, and the raw-GatedMLPBridge recompute path.
• PR3 — RelevanceLens.fit, estimator/rule provenance, merge/registry guards, reuse of matrix readout/vectors/decomposition/interventions, tiny gated-model integration, and docs.

Explicitly deferred to separate follow-up issues: published-artifact conversion; pass@k/lens-comparison harness; CKA utility; linear-probe upper-bound harness; MoE/router/shared-expert and mHC rules; and a demo notebook (gated on a stable quantitative fixture).

---

Alternatives

N/A

---

Additional context

• camilablank, agam_bhatia, Neel Nanda, "R-lens: Making J-lens More Faithful on Early Layers," LessWrong/Alignment Forum, 5 Aug 2026 — https://www.lesswrong.com/posts/nv8oedrnLXKRzNEL9/r-lens-making-j-lens-more-faithful-on-early-layers
• Rezaei Jafari et al., "RelP: Faithful and Efficient Circuit Discovery in Language Models via Relevance Patching," arXiv:2508.21258 — https://arxiv.org/abs/2508.21258
• Reference code (correctness oracle, not a dependency): FarnoushRJ/RelP pinned at 8219d6dc417c3fd7f318342cf61cd2a0c20b7250 — https://github.com/FarnoushRJ/RelP
• Published paired artifacts (converter is a follow-up): camilablank/workspace-lenses on the Hugging Face Hub.

---

Checklist

• Confirmed R-lens is absent from current TransformerLens.
• Kept R-lens separate from RelP while specifying one shared internal backend.
• Corrected Identity- and Half-rule VJPs; named the fail-closed activation cases.
• Scope resolved to Bridge-only (HookedTransformer deprecated in TransformerLens 3.0).
• Named the two Bridge integration constraints (norm fallback, opaque gated MLP).
• Maintainer agreement on the shared Bridge relevance-rule protocol ownership/module names.

---

Duplicate check (as of 2026-09-08)

• RelevanceLens, relevance_lens, and any relevance-based fitting option are absent from the current source and tests. JacobianLens (transformer_lens/tools/analysis/jacobian_lens.py) is the only lens estimator present.
• No LN/Identity/Half relevance-rule backend, use_lrp, or LRP_rules exists in the source.
• Generic Relevance Patching (RelP) is a separate proposed issue. R-lens and RelP share one internal backend but are distinct user-facing features (see "Relationship to RelP").
• The published paired J-lens/R-lens artifacts (camilablank/workspace-lenses on the Hugging Face Hub) are not accepted by the current loader; a converter is a separate follow-up issue, not part of this slice.

---

Reviewer note

Status: exploratory proposal, no code changes.
Suggested labels: enhancement, tooling, TransformerBridge, complexity:high

cc @jlarson4 — could you take a look and let me know if this scope and API shape look right?
If this gets a green light, I'd like to pick it up myself. Could you assign this issue to me once it's approved?
Thank you

---

[jlarson4]

Reposting this here – I have too many windows open looking at too many parts of the repo, apologies!

Thanks for putting this together @janmenjayap. Green light on the concept and API shape, with the following conditions.

1. NormalizationBridge integration must be strictly additive. No restructuring of the existing three-way dispatch. On the native path, wrap the original_component(x) call itself in the custom Function that makes your torch.equal validation hold by construction. Include a regression test that rule-context-inactive behavior is byte-identical to today. tests/integration/test_backward_lens.py (one autograd.grad, no fallback warnings) and the #1526 hook-semantics suite must stay green untouched.

2. Define rule-context x fallback precedence, fail closed. Backward hooks on hook_scale/hook_normalized, or a forward-hook edit, while a rule is active must raise. The existing fallbacks would break your bit-identical-forward invariant and silently compose LRP with hook edits. Hooks can attach after context entry, so check per-forward, not at entry.

3. Rule targeting must be positional, not class-based. 21 adapters mount RMSNormalizationBridge on q_norm/k_norm. Make sure to select by mounting point (ln1/ln2). gemma3n, gemma4, and stablelm have placeholder q/k components & the coverage reporter must classify them correctly.

4. Gated-MLP recompute: allowlist per HF module class, not per adapter. And note W_gate/W_in return TL orientation. Please cover orientation explicitly in the analytic VJP tests.

5. Keep the shared backend private until RelP is filed – underscore module under model_bridge, no documented API, so the protocol can still change shape. Module naming can be settled in PR2 review.

[janmenjayap]

Hi @jlarson4, hope you're doing well! Following up on the 2026-09-09 sign-off — thank you again for taking the time to review the plan.

I wanted to flag that this issue is currently closed as completed, but it looks like only the first of the three agreed PRs has landed so far: #1764 (refactor(jacobian_lens): extract estimator-independent fit driver). Its own description notes this too — it's "PR1 of the 3-PR plan," a pure refactor that "ships no new public API and no relevance-lens surface." It extracts a BackwardProvider seam so the ordinary J-lens fit and a future relevance-modified fit can share one driver, and pins a regression fixture for the ordinary path. That's valuable groundwork, but since no RelevanceLens, no LN/Identity/Half rule backend, and no provenance/merge guards exist yet, the feature this issue is tracking hasn't fully shipped yet.

For reference, still outstanding per the plan you gave conditional sign-off on:

• PR2 — the shared LN/Identity/Half custom-VJP backend (NormalizationBridge native-forward integration, the raw-GatedMLPBridge recompute path, fail-closed coverage). This is the risky part and folds in all five of your conditions from 2026-09-09 (strictly-additive NormalizationBridge integration, fail-closed rule/hook precedence checked per forward call, positional rather than class-based rule targeting, per-HF-module-class gated-MLP allowlisting, keeping the backend module private until the RelP issue is filed).
• PR3 — RelevanceLens.fit itself, estimator/rule provenance, merge/registry guards, and docs.

I'd like to reopen this issue so PR2 and PR3 stay tracked against it instead of the remaining scope going untracked. I have detailed, commit-by-commit implementation plans for both ready to go, already incorporating your five conditions from 2026-09-09 — happy to share them here or fold them into the PR descriptions directly when I open PR2.

Let me know if you'd prefer a fresh tracking issue instead of reopening this one — happy to do either

[jlarson4]

@janmenjayap reopened, clearly I was moving to swiftly when closing out all the post-PR issues last week. Everything should be properly opened now

[jlarson4]

@janmenjayap main description of the issue has been updated to add checkboxes for the PRs to help me keep track of where we are in each of your proposals. Thanks again for all your hard work on this, I look forward to seeing PR3 whenever it is ready

[janmenjayap]

Hi @jlarson4 — thanks for reopening and for adding the PR checkboxes, that makes the remaining scope much easier to track.

Confirming PR2 landed as #1785. All four of your review points are folded into the merged version:

the gated-MLP recompute is gone — the Half-rule is now a gradient scale at the product and the Identity-rule rides the HF act_fn module, so Falcon-H1 multipliers / Gemma3n gate sparsity no longer get a silently-wrong gradient, and the MLP's own hooks stay in the backward graph;
ln1_post / ln2_post joined the normalization mounts, so Gemma 2/3/4 and GLM-4 post-norms get the LN-rule and show up in coverage;
requesting a kind with no canonical mount (e.g. attention=True) now raises instead of silently no-opping;
the weight-less OlmoLayerNorm path no longer AttributeErrors on the first rule-active forward.
PR3 is next and is the thin consumer: RelevanceLens.fit driving PR1's _fit_transport_matrices under PR2's use_relevance_rules, plus estimator/rule provenance, merge + registry guards, and a docs section. Branch feat/r-lens off dev, TDD commit-by-commit, gated on make format + uv run mypy . + make test-pr. I'll open it as a draft early so you can see the shape before it's review-ready.

One question so I don't guess: since the shared backend has now landed, do you want the RelP issue filed first (pointing at the landed _relevance_rules module) before PR3 documents the R-lens side, or is it fine to file RelP in parallel? Happy either way — PR3 only consumes the backend internally and won't re-export it publicly.

[jlarson4]

@janmenjayap Totally fine to do them in parallel