GitHub - hiskuDN/open-r-lens: Independent open-source implementation of R-lens: LRP rules for the Jacobian lens, validated against the authors' released lenses at cosine 0.996

An independent open-source implementation of R-lens — the Jacobian lens with layer-wise-relevance-propagation rules in the backward pass, from R-lens: Making J-lens More Faithful on Early Layers (Blank, Bhatia & Nanda, August 2026).

Not affiliated with or endorsed by the authors of that post. The authors released fitted lens tensors (camilablank/workspace-lenses) but not the fitting code. This repo is that missing piece.

At a glance

What you get: the three rules as a verified reversible model patch, a sharded Modal fitting pipeline, a pass@k evaluator that keeps per-layer ranks, and tooling that checks the result against the authors' released lenses.

flowchart LR
    C["pile-10k<br/>25 prompts"] --> F["jlens.fit"]
    M["HF model"] --> F
    M --> P["lrp_rules()<br/>3 stop-gradients"]
    P --> F
    F -->|unpatched| JL["J-lens"]
    F -->|patched| RL["R-lens"]
    JL --> E["pass@k harness<br/>per-layer ranks"]
    RL --> E
    JL --> X["released_compare<br/>2×2 vs HF lenses"]
    RL --> X

Contents

• The three rules · Verification
• Usage · Layout
• Validation against the released lenses
• Scale curve · Known negatives
• Methodology notes · Cost
• The other released pairs · Ablation and extensions
• Engineering notes · Caveats

The three rules

R-lens is J-lens with three stop-gradients inserted into the backward pass used to estimate the Jacobian. The forward pass is untouched, so the only thing that changes is J_l:

All three are algebraically value-preserving, which gives a strong invariant: patching must not change the forward pass at all. lrp_rules verifies this per module against the dtype's rounding budget and raises if a rewrite drifts, so an unrecognised norm convention fails loudly rather than silently producing a wrong lens.

The identity rule is a custom autograd.Function that passes the real fused kernel's output through and overrides only the backward. Writing it the obvious way — z * sigmoid(z).detach() — recomputes SiLU from unfused primitives and rounds differently, which measures as 2.9% relative drift in the final hidden state of a bf16 Qwen3-0.6B accumulated over 28 layers. With the autograd.Function the drift is exactly zero, so a J-lens/R-lens comparison differs in the backward pass and nothing else.

Verification

verify patches a real model and asserts the forward pass is unchanged while the layer-0 gradient moves. Measured on Qwen/Qwen3-0.6B:

Note the last row. The rules leave the forward pass bitwise identical but rotate the backward pass a long way — cosine 0.38 against the true Jacobian. So R-lens is not an approximation of the first-order causal effect that J-lens is defined as; it is a different readout, which has to earn its keep empirically.

Usage

Because the rules are a model patch, fitting is upstream's jlens.fit called unchanged:

import jlens
from rlens.rules import RuleConfig, lrp_rules

j_lens = jlens.fit(model, prompts)              # baseline
with lrp_rules(hf_model, RuleConfig()):
    r_lens = jlens.fit(model, prompts)          # same call, patched backward

Requires Python 3.12+ and a Modal account. pytest needs no GPU.

uv sync
uv run pytest                                    # 25 tests, no GPU

# gate: cheap, run before paying for a fit
uv run modal run modal_app.py::verify   --model-id Qwen/Qwen3-0.6B

# fit a matched J/R pair, sharded
uv run modal run modal_app.py::fit      --model-id Qwen/Qwen3-0.6B --variant both

# score it (writes per-layer ranks to the volume)
uv run modal run modal_app.py::evaluate --model-id Qwen/Qwen3-0.6B

# fit under the released lenses' exact provenance, then compare
uv run modal run modal_app.py::replicate       --model-id Qwen/Qwen3.5-4B
uv run modal run modal_app.py::released_compare --model-id Qwen/Qwen3.5-4B \
                                                --released-slug qwen3.5-4b

Where analysis runs matters. Anything touching a lens.pt (400 MB – 3.5 GB) stays on Modal, where the file already is. Only the small per-layer rank JSONs come down:

for v in j-lens r-lens; do
  modal volume get rlens-outputs \
    "Qwen__Qwen3-0.6B/$v/eval-ranks.json" "results/Qwen__Qwen3-0.6B/$v/"
done
uv run python scripts/compare.py results/Qwen__Qwen3-0.6B   # k-sweep, per-layer curves

Layout

rlens/rules.py                 the three rules, as a reversible verified model patch
rlens/evaluate.py              pass@k with per-layer ranks retained
modal_app.py                   verify / fit / replicate / evaluate / inspect / diagnostics
scripts/compare.py             J-vs-R tables, k-sweep, per-layer curves, first-hit shift
scripts/compare_to_released.py offline 2×2 (prefer modal_app::released_compare)
tests/test_rules.py            analytic checks that each rule is the specified map
data/evaluations/              the six upstream eval sets (551 items), vendored, Apache-2.0

Validation against the released lenses

Everything below is our measurement, with our harness, of the authors' public artifacts. None of it is a figure quoted from the post, and absolute values are not directly comparable to the post's (see caveats).

Provenance is recorded, so replication can be exact

Each released lens carries a provenance dict. For qwen3.5-4b:

{'model_id': 'Qwen/Qwen3.5-4B', 'dataset_id': 'NeelNanda/pile-10k',
 'target_layer': 30, 't_max': 128, 'n_prompts': 25, 'docs_consumed': 25,
 'skip_first': 4, 'weighting': 'uniform', 'corpus_mode': 'pretrain',
 'config_json': '{"estimator": "relp", "rules": {"ln_rule": true,
   "identity_rule": true, "half_rule": true, "half_rule_beta": 0.5,
   "include_qk_norms": false, "gated_norms": false}}'}

Three settings differ from the upstream J-lens defaults and would silently produce a non-comparable lens: the corpus is pile-10k (not WikiText), skip_first is 4 (not 16), and n_prompts is 25. target_layer: 30 on a 32-block model is the penultimate block, and their source_layers run 0–30 with J_30 exactly the identity.

RuleConfig.as_provenance() emits our settings in the same shape. include_qk_norms: false corresponds to this repo's ln_scope="residual" default (Qwen3-style per-head q_norm/k_norm excluded), and half_rule_beta: 0.5 to the even split. gated_norms has no counterpart here; it is off in every released lens.

Weights: 2×2 cross-similarity

Fitting Qwen3.5-4B under that exact provenance, over 30 comparable layers (the identity target layer excluded, since including it inflates every score):

Clean diagonal dominance: our R matches their R far better than their J, and vice versa. The sharpest single number is the rule-delta test — cosine between our R − J and their R − J, which cancels the shared J baseline and isolates the backward-pass rules from the corpus: median 0.9949, mean 0.9804.

Per-layer relative error falls monotonically with depth — J from 0.43 at layer 0 to 0.012 at layer 29, R from 0.204 to 0.009 — with no layer flagged as an outlier (R relerr never exceeds 4× the median) and Frobenius-norm ratios inside 0.7% from layer 3 on.

Is the early-layer residual just corpus sampling noise? Testable for free: the five shards were each fitted on five disjoint prompts, so sampling noise predicts shard-vs-shard (n=5) exceeds ours-vs-theirs (n=25) by exactly √5 = 2.24×. Measured: 6.5×. Our lens agrees with theirs about 3× better than independent corpus draws would allow, so the same documents are substantially being drawn; the remaining gap points at document selection or chunking rather than the estimator.

Behaviour: our 4B lenses vs theirs, same harness

Matrix agreement shows the weights match. This shows the whole chain matches:

Every absolute value agrees within 0.01–0.03 and the early-layer means agree to 0.004 — two separately fitted lens pairs converging on the same behaviour, not merely the same matrices.

Scale curve

evaluate_released needs only inference, so the post's central quantitative claim — that gains grow with model scale — is testable without fitting anything. Mean pass@10 delta (R minus J) over the six categories, from the released lenses:

Early-layer gains grow monotonically with scale, roughly doubling from 4B to 27B. That is the post's claim, reproduced independently. The all-layer metric does not show it — it dips at 9B — so the choice of window is what makes the effect visible. Reporting only the aggregate would understate R-lens at scale.

Lenses fitted in this repo, for comparison (--rules-preset all, paper recipe for the 0.6B):

At 27B the single largest effect is typo on early layers, +0.333 (0.083 → 0.417), followed by multihop +0.104 (0.014 → 0.118). No category is negative at 27B in either window.

Known negatives

Both survive a k-sweep, so neither is a metric artefact:

• multilingual at 0.6B — a head/tail trade-off. Negative for k ≤ 100 (−0.021 at k=10), positive for k ≥ 500 (+0.070 at k=1000). R-lens degrades the top of the ranking while lifting the tail. Scale-dependent: +0.014 at 4B with our fit, +0.021 at 4B with theirs, +0.002 at 27B.
• association at 9B — a uniform regression on the released lenses, so a property of R-lens rather than of this implementation. R below J at every cutoff (Δ = −0.059, −0.010, −0.029, −0.039, −0.029, −0.020 for k = 10, 50, 100, 500, 1000, 5000). Does not persist at 27B (+0.020).

Methodology notes

pass@10 saturates, and it hides large effects. Where the model is far from the answer, a top-10 cutoff registers a big improvement as +0.000. On Qwen3-0.6B, poetry's target sits at median rank 1778 under J-lens and 319 under R-lens — R-lens better on 95 of 98 items — yet pass@10 reports 0.000 vs 0.010. Sweeping k:

scripts/compare.py prints a k-sweep by default and warns when the headline k is saturated. Do not read a single k.

The logit-lens floor. The floor any fitted lens must clear is doing no transport at all (use_jacobian=False). On Qwen3-0.6B under the paper recipe, R-lens clears it on every category at every k, with the single exception of poetry at k=10 — which is the saturation artefact above:

Fitting corpus matters here: with WikiText instead of pile-10k, the J-lens sat below the floor on order-ops (−0.064), which inflated the apparent R-lens gain. Under the paper's corpus it clears (+0.018).

R-lens is a lower-variance estimator at early layers. Per-shard ||J_l|| over five disjoint 5-prompt shards (Qwen3.5-4B), coefficient of variation:

R-lens more than halves shard-to-shard variance at early layers and converges to identical variance by layer 10. This is a candidate mechanism for why the rules help exactly where they do — a lower-variance estimator, not only less error accumulation — and it is measurable from artifacts already on disk.

Cost

Fitting cost is 2 · N_active · d_model · seq_len FLOPs per prompt — one backward pass per residual dimension, so d_model is a linear multiplier and dominates the choice of target model.

Budget 2× for a matched pair. The accumulator is len(source_layers) · d_model² · 4 bytes and stays resident; trimming source_layers saves memory and disk but not compute, since the backward pass reaches layer 0 either way. Fitting is sharded over disjoint prompt slices and merged with JacobianLens.merge, so wall-clock scales with shard count.

dim_batch replicates the prompt along the batch axis and does not change total FLOPs, but it does drive peak memory: Qwen3.5-4B at dim_batch=32 OOMs a 48 GB L40S, because the retained graph spans 31 blocks and intermediate_size is 9216. dim_batch=8 on an H100 is comfortable.

The other released pairs

All eight models ship matched j-lens/r-lens pairs. A lens file is n_layers · d_model² · 2 bytes, so its size pins L·d²:

qwen3.5-27b and qwen3.6-27b are dimensionally identical — two model generations at matched size, both with released pairs.

Ablation and extensions

Each rule toggles independently, tests/test_rules.py asserts each one moves the backward pass on its own, and --rules-preset fits any subset:

uv run modal run modal_app.py::fit --variant r-lens --rules-preset ln
uv run modal run modal_app.py::fit --variant r-lens --rules-preset identity
uv run modal run modal_app.py::fit --variant r-lens --rules-preset half
uv run modal run modal_app.py::evaluate --variant j-lens,r-lens-ln

--ln-scope switches the LN-rule between residual (the default: only norms at d_model, so Qwen3's per-head q_norm/k_norm are left alone) and all (every norm, as AttnLRP does). On Qwen3-0.6B that is 57 norms versus 113.

RuleConfig also declares attn_half_rule and router_identity_rule, which AttnLRP would apply to the attention softmax and the MoE router. These are not implemented; doing so needs an eager attention path to patch, since SDPA/flash fuse the tensors away. lrp_rules raises NotImplementedError rather than ignoring them, because a silently-inert flag would make an ablation return the baseline's numbers and read as a null result.

Engineering notes

Compare lenses where they already live. modal_app.py::released_compare runs the 2×2 against the released lenses inside Modal, reading ours from the volume and theirs from an HF cache hop. scripts/compare_to_released.py does the same locally and is kept for offline analysis, but the Modal path is the default for a reason:

torch.load accepts a partially-written .pt as long as the zip central directory is intact, and the tensors it returns look entirely plausible. modal volume get returns before all bytes have landed, and a byte-size check is necessary but not sufficient. A partially-downloaded 786 MB lens produced a confident, reproducible, and completely spurious 20% norm deficit at a single layer — reproducible because the corruption was baked into the local file, so re-reading it agreed with itself. Moving the comparison onto Modal removes the failure mode instead of trying to detect it. verify_merge provides the independent check: a merged lens must equal a fresh mean of its own shards at every layer.

Modal volumes are not coherent across container reuse. fit_shard commits each shard, but a merge_shards container reused from an earlier merge holds the volume state it saw at mount time and fails with FileNotFoundError on shards that modal volume ls lists. merge_shards calls outputs.reload() first; merge_existing merges shards already on the volume so a merge failure never costs a refit.

Caveats

• Lenses fitted here use a single seed with no error bars, on 55–107 items per category. Treat directions as informative and magnitudes as indicative.
• The order-ops synonym table (ORDER_OPS_SYNONYMS) is reconstructed from the eval README's description, since the upstream table is not published. It is applied identically to every lens under comparison, so it cannot manufacture a J-vs-R delta, but absolute order-ops numbers are not comparable to the post's.
• Our pass@k implementation reproduces no published number directly. Its external check is that the released lenses show the paper's early-layer claim under it.
• Untested: MoE models, a gated_norms counterpart, and anything above 27B.

Credits

• J-lens and the eval sets: anthropics/jacobian-lens (Apache-2.0), from Verbalizable Representations Form a Global Workspace in Language Models.
• The LRP rules: RelP (Rezaei Jafari et al., arXiv:2508.21258) and AttnLRP / LXT.
• R-lens: Blank, Bhatia & Nanda (2026).

License

MIT — see LICENSE. The vendored evaluation sets in data/evaluations/ are Apache-2.0 and remain under that license; see NOTICE and data/evaluations/LICENSE.