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

---

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

---