fmerlo/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

	Anchor	Covers	Reads
J_l — neuronpedia/jacobian-lens	block output x_{l+1}	every layer but the top one	a whole layer
J^z_l — this repo	o_proj input z_l	every layer with a head anchor, top one included	one head

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

llama3.1-8b-it_Jzfit_n1000.pt	1,073,752,149 bytes
sha256	07a9a25aead37574c16829c3569ca1b2be567e7246bd71b14040d102a2caa323
Contents	32 tensors of [4096, 4096], layers 0..31; 32 heads × head_dim 128
Fitted	2026-08-13

Qwen/Qwen3.5-4B

qwen3.5-4b_Jzfit_n1000.pt	167,775,903 bytes
sha256	249929f24139c4f2722645cc0308ff5fd509455240db5890391c369094a74fa6
Contents	8 tensors of [2560, 4096], layers 3, 7, 11, 15, 19, 23, 27, 31 (the softmax-attention blocks); 16 heads × head_dim 256
Fitted	2026-09-15

google/gemma-3-4b-it

gemma-3-4b-it_Jzfit_n1000.pt	356,526,669 bytes
sha256	5c9f6f1263f496a4569f7c46d0c3469fcefb61f1633e9ce549338544269dc6e9
Contents	34 tensors of [2560, 2048], layers 0..33; 8 heads × head_dim 256
Fitted	2026-09-16

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}
}