WeZZard/jlens-qwen36

TL;DR: jlens-qwen36 is a local visual debugger for Qwen3.6-27B (4-bit) on Apple Silicon that fits a Jacobian lens — Anthropic's recent upgrade to the logit lens — and renders a live position × layer grid of what the model is "thinking about" at every layer, even when those concepts never reach the output. It goes beyond read-only interpretability: you can edit the latent workspace by hand or run a backward search that finds the edit producing a reply you name. The engineering (a custom Metal VJP kernel for Gated DeltaNet, full-depth lens fitting in ~2.75h on an M4 Pro) makes it the most practical way to actually play with this line of interpretability work right now.

What it is

The project ports Anthropic's Jacobian-lens reference implementation (from Verbalizable Representations Form a Global Workspace in Language Models) to Apple MLX, wraps it in a web UI, and points it at a quantized Qwen3.6-27B running locally. The headline demo: in an email-blackmail setup where the model writes a calm, compliant reply, the workspace band lights up with blackmail / suicide / murder / threatening / fictional across the middle layers. The concept is present in the residual stream but suppressed from the output — you can literally watch what the model doesn't say out loud. There's also a read-only hosted version at jlens.wezzard.com with shareable conversation URLs.

How the lens works

A logit lens naively projects intermediate residuals through the unembedding, which is known to be unfaithful at earlier layers because the basis drifts. The Jacobian lens fixes this: for each layer ℓ, fit a matrix J_ℓ ∈ R^{d×d} — the network's average input→output Jacobian over a prompt corpus — which maps h_ℓ into the final-layer basis before unembedding: softmax(W_U · norm(J_ℓ h_ℓ)). Fitting is done by chaining per-layer Jacobians, J_ℓ = J_{ℓ+1} · M_ℓ.

The hard implementation problem: Qwen3.6-27B has 48 Gated DeltaNet linear-attention layers, and MLX's fused GDN kernel has no VJP; the ops fallback is ~22× slower. The author wrote a custom Metal backward kernel (custom_gdn_vjp.py) plus analytic branch-Jacobian assembly, getting full-depth (63-layer) lens fitting down to ~2.75h on an M4 Pro. This is real engineering depth, not a thin paper port — and it makes fitting your own lens on 100+ prompts affordable.

The interesting part: writing to the workspace

The lens is bidirectional. Click a cell and Replace/Add/Remove/Erase a concept, scoped to one cell, a layer band, or the whole reply; a Baseline/Intervened toggle diffs the two runs. Replacing France with China across the band rewrites the answer.

More novel is the backward search: click a word in the reply, type what it should say instead, and the app searches for a latent edit that produces that reply, collecting results as verified "recipes." When no direct edit works, it looks for the premise behind the reply — to move ⟨Paris⟩→⟨Beijing⟩, it finds the ⟨France⟩→⟨China⟩ swap upstream. That's a small automated causal-mediation tool disguised as a UI feature, and a genuinely nice framing.

What's shipped and caveats

Shipped: code (Apache-2.0), a pre-fitted full-depth lens (3.3 GB, GitHub release), the Metal kernels, docs (docs/lenses.md, docs/interventions.md, docs/perf/), and the hosted demo. The model itself auto-downloads from HuggingFace (~15 GB).

Caveats, stated candidly in the README:

• The bundled lens is demo-grade — fit on only 20 prompts. Readouts are "interpretable but noisy"; interventions are "causal but concept-dependent." For research use, load Neuronpedia's n=1000 lens (the paper's fitting scale) via JLENS_PATH, or fit your own.
• Apple Silicon / MLX only, and the GDN kernel is qwen3_5-architecture-specific.
• Single-token concepts only; multi-token needs the paper's extension.
• Chat runs with enable_thinking=False, deliberately forcing computation into the latent stream — a sensible choice for the tool's purpose, but note it changes model behavior relative to default Qwen chat.
• No quantitative faithfulness numbers in the README — before drawing scientific conclusions, verify that lens readouts actually track behavior (the intervention diff toggle is your friend here).

Try it

Needs an Apple Silicon Mac with ~24 GB free RAM, or just use the hosted read-only demo at jlens.wezzard.com.

git clone https://github.com/WeZZard/jlens-qwen36.git
cd jlens-qwen36 && uv sync

gh release download v0.2-fulldepth --repo WeZZard/jlens-qwen36 \
  --pattern '*.npz.part-*' --dir data/lens/
cat data/lens/*.npz.part-* > data/lens/lens.npz && rm data/lens/*.part-*

uv run python -m uvicorn jlens_qwen.serve:app --host 127.0.0.1 --port 8765
# open http://127.0.0.1:8765/