fmerlo/head-level-jlens

J^z — a head-level Jacobian lens for Llama-3.1-8B-Instruct

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).

Published so the ~15 GPU-hours of the fit are optional. It is the artifact that repository fits, not a copy of anything else.

What the anchor changes

	Anchor	Covers	Reads
J_l — neuronpedia/jacobian-lens	block output x_{l+1}	layers 0..30	a whole layer
J^z_l — this repo	o_proj input z_l	layers 0..31	one head

Two consequences, both load-bearing:

• The columns are indexed by o_proj's input, so head h's slice is [128*h : 128*(h+1)] — 32 heads × head_dim 128. Slicing a block-anchored lens the same way is meaningless; the block output has already been through the layer's MLP.
• It covers all 32 layers, one more than the published lens, because a head writes before its own layer's residual add and MLP. That layer count is the only thing distinguishing the two families on disk, and the consuming code checks it.

File

llama3.1-8b-it_Jzfit_n1000.pt	1,073,752,149 bytes
sha256	07a9a25aead37574c16829c3569ca1b2be567e7246bd71b14040d102a2caa323
Contents	32 tensors of [4096, 4096], fp16 on disk, upcast to fp32 by JacobianLens

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"]
jz = load_jz(spec, jz_path(spec, n_prompts=1000))
head_15_3 = jz.jacobians[15][:, spec.head_slice(3)]   # [4096, 128]

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 this file loads happily against another model and produces plausible-looking nonsense. It is for meta-llama/Llama-3.1-8B-Instruct only.

How it was fitted

wikitext-103-raw-v1 (Salesforce/wikitext), the corpus the published lenses were fitted on — 1000 prompts, seed 0, truncated to 2000 characters and 128 tokens, with the first 16 positions skipped as attention sinks. 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 is the budget, not a converged count: the stop_at_delta = 0.002 early stop that the published lenses used never fired here, because it is a materially stricter target at this anchor than at the block output it was calibrated on. Convergence is instead answered by an independent fit on the disjoint corpus[1000:1400]. The derivation is in the repository's METHODS.md; the per-prompt convergence traces are in its results/fit/.

Fitted 2026-08-13 with torch 2.13.0+cu130 and transformers 5.15.0.

Citation

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