stanleytheli/qwen3.6-35B-A3B-jlens

Jacobian lens for Qwen3.6-35B-A3B

A fitted Jacobian lens (J-lens) for Qwen/Qwen3.6-35B-A3B, following "Verbalizable Representations Form a Global Workspace in Language Models" (Anthropic, 2026) and its reference implementation, jacobian-lens.

The lens reads out what an internal activation is disposed to make the model say: it linearly transports a residual-stream vector at any layer/position into the final-layer basis with the average input–output Jacobian, then decodes it with the model's own unembedding:

lens_l(h) = unembed( J_l @ h ),   J_l = E[∂h_final / ∂h_l]

Files

file	fit corpus	note
lens.pt	1000 WikiText-103 prompts	paper setting; use this one
lens-100prompt.pt	100 WikiText-103 prompts	matches lens.pt within noise (quality saturates fast, §9.3)

Both cover source layers 0–38 (target layer 39) of the 40-layer decoder, d_model=2048, stored fp16 (~330 MB).

Usage

import transformers, jlens

hf = transformers.AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3.6-35B-A3B", dtype="bfloat16", device_map="cuda")
tok = transformers.AutoTokenizer.from_pretrained("Qwen/Qwen3.6-35B-A3B")
model = jlens.from_hf(hf, tok)

lens = jlens.JacobianLens.from_pretrained("stanleytheli/qwen3.6-35B-A3B-jlens")
lens_logits, model_logits, _ = lens.apply(
    model, "Fact: The currency used in the country shaped like a boot is",
    positions=[-1])
for layer, logits in sorted(lens_logits.items()):
    print(layer, [tok.decode([t]) for t in logits[0].topk(5).indices])

Fitting details

• Estimator: cotangents summed over all valid target positions, averaged over source positions (positions ≥ 16 of 128-token prompts), one forward + d_model/dim_batch backward passes per prompt with all parameters frozen.
• Fitted on 8 H200s (dim_batch=16, ~101 GiB peak each, ~55–80 s/prompt), shards merged by prompt-count-weighted mean.
• Gradients flow through the sparse 256-expert MoE dispatch and the hybrid gated linear-attention blocks via the pure-PyTorch (autograd-compatible) transformers path.

Evaluation vs the classic logit lens

Rank of the model's own final output token under each lens (Qwen3.6-35B-A3B, median over samples):

layer	J-lens (WikiText)	logit lens (WikiText)	J-lens (multihop)	logit lens (multihop)
8	182	11,793	15,691	73,541
16	106	42,247	162	166,944
24	61	28,634	127	225,956
32	47	672	443	143,817

On two-hop prompts ("the currency of the country shaped like a boot"), the lens reads the latent middle hop (Italy) at rank ≲300 through the mid-network band where the logit lens ranks it ~40,000–60,000; the middle hop's top-1 probability peaks near layer 36 and recedes at the final layers as the answer takes over.

License

Apache-2.0, matching the reference implementation. The lens contains no Qwen weights — only averaged Jacobians of the base model.