The Jacobian lens. Carries the eager-attention requirement (fused SDPA has no forward-AD), rotary position_embeddings threading, the module-tree search that survived Gemma-3/4 and found Qwen3.5's linear_attn, and model-only trust (tokenizer trust broke Yi).
raw: jlens_jacobian.py · annotate via ../marginalia/code-jlens_jacobian-py.json
#!/usr/bin/env python3 class="s">"""jlens_jacobian.py — the lens nemo actually specified: lens(h) = softmax(W_U · norm(J · h)). The first readout used a LOGIT lens, softmax(W_U · norm(h)) — the tractable baseline. This is the Jacobian refinement from the lab brief. THE TRACTABILITY TRICK. The full Jacobian J of the remaining network w.r.t. a hidden state is [hidden × hidden] per layer per position — hopeless to materialise (4096² per position). But the lens never needs J itself, only the product **J · h**, and a Jacobian-vector product is exactly what forward-mode autodiff computes in ONE extra forward pass: _, Jh = torch.func.jvp(f_rest, (h,), (h,)) where f_rest runs layers ℓ+1..N on h and returns the final normed hidden. The tangent is h itself, so the output tangent is J·h — "how the network's final state moves as this hidden state moves along itself". Then lens = softmax(W_U · norm(J·h)). Interpretation, and why it should beat the logit lens: the logit lens asks *what does this hidden state look like if you decode it right now*, which at mid layers is dominated by high-norm outlier and attention-sink directions (we measured exactly that: the raw top-k was CJK fragments and <|im_end|>). The Jacobian lens asks *what does this hidden state PUSH the output toward* — it reads the causal direction rather than the residual's current contents, so sink directions that the network does not act on should fall away. python3 jlens_jacobian.py --model 01-ai/Yi-1.5-34B-Chat --run-label yi-1.5-34b \ --questions /tmp/jlens_lab_harness/data/lab_questions_20.csv --limit 20 \ --out-dir /mnt/cache/lens_yi_jac class="s">""" from __future__ import annotations import argparse, csv, json, os, sys os.environ.setdefault(class="s">"HF_HOME", class="s">"/mnt/assets/models/huggingface") os.environ.setdefault(class="s">"HF_HUB_OFFLINE", class="s">"1") os.environ.setdefault(class="s">"TOKENIZERS_PARALLELISM", class="s">"false") sys.path.insert(0, os.environ.get(class="s">"JLENS_HARNESS", class="s">"/tmp/jlens_lab_harness")) import torch from torch.func import jvp from transformers import AutoModelForCausalLM, AutoTokenizer from geo.common import prompts NULL_PROMPTS = [ class="s">"What are the public transport options in Melbourne?", class="s">"Explain how to renew an Australian passport.", class="s">"What is the weather like in Sydney in winter?", class="s">"Describe the process of enrolling a child in primary school in NSW.", class="s">"How do I recycle electronic waste in Australia?", ] def parts(model): class="s">"""(decoder layers, final norm, unembedding), found by SEARCHING the module tree. Do not read the config for this. Multimodal configs (Gemma3Config, Gemma4) have no top-level num_hidden_layers — it hides under config.text_config — and a lens that trusts the config dies on exactly the newest models we most want to measure. The module tree always knows: find the deepest ModuleList of decoder blocks and the norm that sits beside it. Ask the model, not its description of itself.class="s">""" import torch.nn as nn best = None for name, mod in model.named_modules(): if isinstance(mod, nn.ModuleList) and len(mod) >= 4: blk = mod[0] # "linear_attn": Qwen3.5's Gated-DeltaNet hybrid names its mixer linear_attn, and the # whole five-rung ladder failed "could not locate a decoder stack" on exactly the # newest lineage this function's docstring promises to handle. Match any child whose # name contains attn/attention instead of enumerating spellings. if any(hasattr(blk, at) for at in (class="s">"self_attn", class="s">"attention", class="s">"attn", class="s">"linear_attn")) \ or any((class="s">"attn" in n) or (class="s">"attention" in n) for n, _ in blk.named_children()): if best is None or len(mod) > len(best[1]): best = (name, mod) if best is None: raise RuntimeError(class="s">"could not locate a decoder stack in this architecture") stack_name, layers = best parent = model for part in stack_name.split(class="s">".")[:-1]: parent = getattr(parent, part) norm = (getattr(parent, class="s">"norm", None) or getattr(parent, class="s">"final_layernorm", None) or getattr(parent, class="s">"final_layer_norm", None)) if norm is None: # last resort: any norm module directly under the parent import torch.nn as _nn for n, m in parent.named_children(): if class="s">"norm" in n.lower(): norm = m; break head = model.get_output_embeddings() if norm is None or head is None: raise RuntimeError(fclass="s">"found {len(layers)} blocks at {stack_name} but no final norm / lm head") return layers, norm, head def build_messages(query, locale=class="s">"en-AU"): return [{class="s">"role": class="s">"system", class="s">"content": prompts.BASE_SYSTEM_PROMPT}, {class="s">"role": class="s">"user", class="s">"content": prompts.BASE_USER_TEMPLATE.format(locale=locale, query=query)}] def merge_system_into_user(messages): class="s">"""Fold a system turn into the first user turn, VERBATIM. Some chat templates (Gemma-2's among them) raise `TemplateError: System role not supported`. The probe runner already handles this; the lens did not, so every Gemma-2 model died at template time after paying the full model load. Only the delivery structure changes — the prompt text is byte-identical, which is what keeps the GEO prompts verbatim as specified. class="s">""" sys_txt = class="s">"\n\n".join(m[class="s">"content"] for m in messages if m[class="s">"role"] == class="s">"system") rest = [m for m in messages if m[class="s">"role"] != class="s">"system"] if sys_txt and rest and rest[0][class="s">"role"] == class="s">"user": rest = [{class="s">"role": class="s">"user", class="s">"content": sys_txt + class="s">"\n\n" + rest[0][class="s">"content"]}] + rest[1:] return rest def apply_template(tok, messages, **kw): class="s">""class="s">"apply_chat_template, retrying without a system role for templates that reject one."class="s">"" try: return tok.apply_chat_template(messages, **kw) except Exception as e: if class="s">"system" not in str(e).lower(): raise print(fclass="s">"[template] system role rejected ({e}); folding it into the first user turn", file=sys.stderr) return tok.apply_chat_template(merge_system_into_user(messages), **kw) @torch.no_grad() def hidden_states(model, ids): return model(ids, output_hidden_states=True).hidden_states def jacobian_lens_mass(model, tok, ids, a0, layers_sel, topk=64, max_pos=64, track_ids=None): class="s">"""J·h at each selected layer over the answer span, decoded through the unembedding. track_ids: optional {name: set(token_ids)} — accumulate FULL-softmax probability over those ids regardless of topk. Exists because the identity experiment's anchor families sat below the top-64 cutoff and read as zero: an anchor absent from a topk dict is not an anchor with no mass, it is an anchor the readout could not see. Tracked mass is returned under class="s">"__tracked__:<name>" keys so existing consumers are unaffected. class="s">""" layers, norm, head = parts(model) hs = hidden_states(model, ids) n_layer = len(layers) acc = torch.zeros(head.weight.shape[0], device=model.device, dtype=torch.float32) count = 0 for L in [l for l in layers_sel if 0 < l < n_layer]: h_full = hs[L] # [1, seq, hidden] entering block L span = h_full[:, a0:, :] if span.shape[1] == 0: continue # subsample positions: the JVP is a forward pass per call, so cap the cost honestly step = max(1, span.shape[1] // max_pos) sel = span[:, ::step, :].contiguous() # transformers 5.x decoder blocks REQUIRE position_embeddings (rotary cos/sin); calling # blk(hidden) alone returns None. Build them for the true positions of the subsampled tokens # so the rotary phase is right rather than silently zeroed. pos_ids = torch.arange(a0, h_full.shape[1], device=h_full.device)[::step][None, :] rot = getattr(getattr(model, class="s">"model", model), class="s">"rotary_emb", None) pe = rot(sel, pos_ids) if rot is not None else None def f_rest(x): # run the REMAINING blocks on x, then the final norm. Positions are independent here # (no attention across the subsample) — an approximation we state rather than hide: # it reads each position's own causal push, not its interaction with its neighbours. y = x for blk in layers[L:]: out = blk(y, position_embeddings=pe) if pe is not None else blk(y) y = out[0] if isinstance(out, tuple) else out return norm(y) try: _, Jh = jvp(f_rest, (sel,), (sel,)) # forward-mode: one extra pass, no J built except Exception as e: print(fclass="s">" [layer {L}] jvp failed ({type(e).__name__}: {e}) — skipping", file=sys.stderr) continue logits = head(Jh).float() # softmax(W_U · norm(J·h)) p = torch.softmax(logits, dim=-1).sum(dim=1).squeeze(0) acc += p; count += 1 if not count: return {} acc /= count vals, idx = torch.topk(acc, min(topk, acc.numel())) mass = {} for v, i in zip(vals.tolist(), idx.tolist()): piece = tok.decode([i]).strip().lower() if piece and v > 0: mass[piece] = mass.get(piece, 0.0) + float(v) if track_ids: for name, ids_ in track_ids.items(): tot = float(sum(acc[i].item() for i in ids_ if 0 <= i < acc.numel())) mass[fclass="s">"__tracked__:{name}"] = tot return mass def main(): ap = argparse.ArgumentParser() ap.add_argument(class="s">"--model", required=True); ap.add_argument(class="s">"--run-label", required=True) ap.add_argument(class="s">"--questions", required=True); ap.add_argument(class="s">"--limit", type=int, default=20) ap.add_argument(class="s">"--out-dir", required=True); ap.add_argument(class="s">"--layers", default=class="s">"") ap.add_argument(class="s">"--device", default=class="s">"cuda:0"); ap.add_argument(class="s">"--max-new", type=int, default=512) ap.add_argument(class="s">"--trust", action=class="s">"store_true") # internlm/hunyuan/baichuan need custom code a = ap.parse_args() os.makedirs(a.out_dir, exist_ok=True) rows = list(csv.DictReader(open(a.questions)))[: a.limit] print(fclass="s">"loading {a.model} …", file=sys.stderr, flush=True) tok = AutoTokenizer.from_pretrained(a.model) # NOT trust here — it broke Yi earlier # EAGER attention is REQUIRED here, not a preference: torch's fused SDPA/flash kernel has no # forward-mode AD derivative ("Trying to use forward AD with _scaled_dot_product_flash_attention # that does not support it"), so the JVP cannot differentiate through it. Eager attention is plain # PyTorch ops all the way down, which forward AD handles. Slower; it is the price of the Jacobian. model = AutoModelForCausalLM.from_pretrained(a.model, dtype=torch.bfloat16, device_map={class="s">"": a.device}, attn_implementation=class="s">"eager", trust_remote_code=a.trust) # belt and braces: some transformers builds ignore the experts_implementation kwarg # silently (the container did — the retry still hit _grouped_mm). The method call is the # version that provably takes effect; warn loudly where the API is absent entirely. if hasattr(model, class="s">"set_experts_implementation"): try: model.set_experts_implementation(class="s">"eager") except Exception as e: print(fclass="s">"[experts] set failed: {e}", file=sys.stderr) elif class="s">"moe" in type(model).__name__.lower(): print(class="s">"[experts] WARNING: no experts_implementation API — MoE jvp will fail on this build", file=sys.stderr) model.eval() _layers, _, _ = parts(model) n_layer = len(_layers) sel = [int(x) for x in a.layers.split(class="s">",") if x] or \ sorted({int(n_layer * f) for f in (0.4, 0.55, 0.7, 0.85)}) print(fclass="s">"loaded; {n_layer} layers, JACOBIAN lens at {sel}", file=sys.stderr, flush=True) def readout(query): ids = apply_template(tok, build_messages(query), add_generation_prompt=True, return_tensors=class="s">"pt", return_dict=True)[class="s">"input_ids"].to(model.device) with torch.no_grad(): gen = model.generate(ids, max_new_tokens=a.max_new, do_sample=False, pad_token_id=(tok.pad_token_id or tok.eos_token_id)) full = gen return jacobian_lens_mass(model, tok, full, ids.shape[1], sel) print(class="s">"null baseline …", file=sys.stderr, flush=True) base = {} for q in NULL_PROMPTS: for k, v in readout(q).items(): base[k] = base.get(k, 0.0) + v base = {k: v / len(NULL_PROMPTS) for k, v in base.items()} for q in rows: qid = q[class="s">"query_id"] mass = readout(q[class="s">"query"]) rec = {class="s">"query_id": qid, class="s">"condition": class="s">"baseline", class="s">"run_label": a.run_label, class="s">"model": a.model, class="s">"layers": sel, class="s">"lens": class="s">"jacobian_jvp", class="s">"token_mass": mass, class="s">"baseline_token_mass": base} with open(os.path.join(a.out_dir, fclass="s">"{qid}.json"), class="s">"w") as f: json.dump(rec, f) top = sorted(mass.items(), key=lambda kv: -kv[1])[:8] print(fclass="s">"[{qid}] jacobian top: {[t for t, _ in top]}", file=sys.stderr, flush=True) print(fclass="s">"wrote -> {a.out_dir}", file=sys.stderr, flush=True) if __name__ == class="s">"__main__": main()