head-level-jlens

J^z — head-level Jacobian lenses for Llama-3.1-8B-Instruct, Qwen3.5-4B and Gemma-3-4B-IT

A Jacobian lens fitted at o_proj's input rather than at the block output, so that a single attention head's additive write can be transported into the final-layer basis and read. Produced by head-level-jlens with the estimator of Verbalizable Representations Form a Global Workspace in Language Models (Gurnee, Sofroniew, Lindsey et al., 2026). One file per checkpoint; each is the artifact that repository fits for that checkpoint, not a copy of anything else.

Published so the GPU-hours of the fits are optional (~15 for Llama, ~22 for Qwen, ~6 for Gemma).

What the anchor changes

Two consequences, both load-bearing:

• The columns are indexed by o_proj's input, so head h's slice is [head_dim*h : head_dim*(h+1)] — n_heads × head_dim columns, which equals d_model on Llama and does not on the other two (4096 ≠ 2560 on Qwen, 2048 ≠ 2560 on Gemma). Slicing a block-anchored lens the same way is meaningless; the block output has already been through the layer's MLP.
• It covers the top layer too, because a head writes before its own layer's residual add and MLP. On Qwen3.5 only every fourth block is softmax attention (the rest are Gated DeltaNet, with no o_proj), so its lens holds the 8 attention layers 3, 7, …, 31 and nothing else. The list of source layers is stored in the file and the consuming code checks it against the checkpoint.

Files

Each .pt holds {"J": {layer: tensor}, "n_prompts", "source_layers", "d_model"}, fp16 on disk and upcast to fp32 by JacobianLens; every tensor is [d_model, n_heads × head_dim].

meta-llama/Llama-3.1-8B-Instruct

Qwen/Qwen3.5-4B

google/gemma-3-4b-it

Use

Through the project, which resolves a local fit first and digest-checks whatever it loads:

from src.config import LENSES, jz_path
from src.model import load_jz

spec = LENSES["llama3.1-8b-it"]          # or "qwen3.5-4b", "gemma-3-4b-it"
jz = load_jz(spec, jz_path(spec, n_prompts=1000))
head_15_3 = jz.jacobians[15][:, spec.head_slice(3)]   # [d_model, head_dim]

Or standalone, with only jlens:

from jlens import JacobianLens

jz = JacobianLens.from_pretrained(
    "fmerlo/head-level-jlens", filename="llama3.1-8b-it_Jzfit_n1000.pt"
)

from_pretrained skips the digest check, since it returns a lens rather than a path. Verify the sha256 above if the file's identity matters to your result.

A lens is only valid for the checkpoint it was fitted on. JacobianLens carries no model identity, so a file loads happily against another model of the same width and produces plausible-looking nonsense. Match the file to its base_model.

How they were fitted

wikitext-103-raw-v1 (Salesforce/wikitext), the corpus the published lenses were fitted on — the same 1000 prompts for every checkpoint (seed 0, truncated to 2000 characters and 128 tokens, the first 16 positions skipped as attention sinks). Each run as four disjoint 250-prompt shards merged by JacobianLens.merge, an n_prompts-weighted mean, so the merged estimate is exact rather than approximate. bf16 forward pass, matching the published lenses' fit.

1000 prompts is a fixed budget: the estimator averages with no stopping rule, as jlens.fitting.fit does, and the per-prompt trace beside each shard in the repository's results/fit/ is monitoring only. On Qwen3.5 the fit differentiated 32 output coordinates per backward pass instead of 64 (its DeltaNet blocks run a pure-PyTorch fallback whose graph does not fit an 80 GB card at 64); that changes memory and wall-clock, not the estimate. Everything the Llama pipeline assumed and the other two checkpoints break — Gemma's RMSNorm between o_proj and the residual add, Qwen's gated attention output and hybrid stack, the 1 + w gain convention, the multimodal wrappers — is listed in the repository's paper.md, Supplementary Section S7.

torch 2.13.0+cu130 and transformers 5.15.0 throughout.

Citation

@misc{merlo2026headleveljlens,
  author = {Merlo, Filippo},
  title  = {head-level-jlens: a head-anchored Jacobian lens for Llama-3.1-8B-Instruct, Qwen3.5-4B and Gemma-3-4B-IT},
  year   = {2026},
  url    = {https://github.com/filippo-merlo/head_level_jlens}
}