Watching Language Models Think: Real-Time Jacobian Lens Visualization with Subtext

Hook

What if you could watch a language model form its verdict during the reading phase, seconds before it generates a single token of refusal?

Context

Interpretability research has a reproducibility crisis. Papers show beautiful attribution maps on cherry-picked prompts, but practitioners debugging production models can't easily reproduce those insights on their own conversations. The gap between "here's what attention pattern 4.7 does on this one sentence" and "why did my chatbot refuse this reasonable request" remains enormous.

Anthroponic's Jacobian lens method—introduced in their research on workspace hypothesis—offered a breakthrough: instead of explaining which input tokens mattered (attribution), it reveals which output concepts are active in the model's residual stream right now, even when the model hasn't started writing yet. The insight is elegant: treat each layer's hidden state as a compressed representation of future text, then use a learned linear projection (the "lens") to decode it into vocabulary space. Suddenly you can see the model planning "refusal" or "apology" or "technical explanation" while it's still reading your prompt. But Anthropic's implementation lived in research code, worked only on their proprietary models, and provided no way to observe this during actual conversations. Subtext changes that by implementing continuous lens readout during live chat with open models, with an architecture designed around the constraint that real inference uses KV caching.

Technical Insight

The core challenge is that production inference has two completely different execution modes. During prefill (processing user input), the model reads all tokens in parallel—a single forward pass computes hidden states at every position. During generation, it's sequential: one new token per step, reusing cached key-value pairs. Most interpretability tools ignore this split and just record activations, but Subtext's architecture embraces it.

The implementation uses PyTorch forward hooks installed at nine fixed transformer layers. During prefill, hooks fire once but capture a tensor of shape [batch, seq_len, hidden_dim]—every input position. The lens projection runs on all positions at once:

# Simplified from the actual hook logic
def lens_hook(module, input, output):
    residual = output[0]  # [batch, seq_len, hidden_dim]
    
    if self.phase == "prefill":
        # Project every position through the lens
        logits = self.lens_projection(residual)  # [batch, seq_len, vocab_size]
        probs = F.softmax(logits, dim=-1)
        
        # Filter to word-initial tokens only for display
        display_probs = probs * self.word_initial_mask
        
        # Stream top-k per position to client
        for pos in range(residual.size(1)):
            top_k = torch.topk(display_probs[0, pos], k=10)
            self.ws_send({"layer": layer_idx, "pos": pos, 
                         "tokens": top_k.indices, "probs": top_k.values})
    
    elif self.phase == "generate":
        # KV cache is active—only read the newest position
        new_residual = residual[:, -1:, :]  # [batch, 1, hidden_dim]
        logits = self.lens_projection(new_residual)
        # ... same filtering and streaming for single position

This phase-aware design is why Subtext maintains native inference speed. During generation, each hook processes a single 4096-dimensional vector rather than a full sequence matrix, and the lens projection (4096 × 151936 for Qwen models) completes in under 2ms on a 3090. The WebSocket client receives 9 lens frames per generated token—one per hooked layer—fast enough that the visualization updates feel synchronous with text appearance.

The word-initial filtering deserves scrutiny because it affects what you see. BPE tokenization splits "understanding" into ["_understand", "ing"], but only "_understand" (with leading space) is word-initial. The implementation computes probabilities over the full 151,936-token vocabulary, then zeros out non-word-initial entries before extracting top-k:

# Mask construction (done once at startup)
word_initial_mask = torch.zeros(vocab_size)
for token_id, token_str in tokenizer.vocab.items():
    if token_str.startswith(' ') or token_str.startswith('Ġ'):  # GPT-2 style
        word_initial_mask[token_id] = 1.0

This makes the display readable—you see "refusal" not "_ref" "use" "al"—but it's lossy. The model's actual workspace might be activating mid-word fragments that carry meaning (the "use" in "refusal" versus "useful"), and filtering hides them. The reference Anthropic implementation does the same masking for display while acknowledging this limitation.

The client-server split enables the killer feature: session export. Every lens frame gets recorded as JSON with millisecond timestamps:

{"frame": 47, "layer": 6, "position": 12, "phase": "generate",
 "top_tokens": [[" reject", 0.23], [" decline", 0.19], [" refuse", 0.15], ...]}

Replay mode loads this JSON into the browser and re-animates the entire session with no GPU, no Python, just JavaScript manipulating DOM elements. This solves reproducibility: you can export a conversation where the model exhibited surprising behavior, commit the JSON to GitHub, and anyone can replay the exact workspace evolution in their browser. The live demo on the repo is actually just a static replay—no server running.

Validation was critical because the KV-cached inference path differs from the reference implementation's full-pass apply() function. The author compared lens outputs token-by-token against Neuronpedia's reference on identical prompts and verified cosine similarity above 0.99998 between logit vectors, with exact matches on top-5 predictions. This proves the method works under production inference constraints, not just research evaluation.

Gotcha

The pre-fitted lens creates hard model lock-in. Subtext ships with lens weights for Qwen 4B and 27B models only, fitted by Neuronpedia using the jlens library. If you want to use a different architecture—Llama, Mistral, Gemma—you need to fit your own lens by running hundreds of prompts through the model, recording residual streams, and solving a least-squares regression to learn the projection matrices. For a 4B model this takes an hour on a 3090; for larger models it's worse. The repo provides no guidance on prompt selection for fitting (the quality of the lens depends heavily on prompt diversity), and there's no way to assess lens quality except visual inspection of whether the results "look reasonable."

The nine-layer sampling is fixed in code with no configuration surface. You can't choose which layers to observe, adjust granularity based on model depth, or dynamically enable/disable layers during a session. For a 27-layer model, observing layers [4, 7, 10, 13, 16, 19, 22, 25, 27] might be optimal; for a 32-layer model, the same absolute indices could miss important stratification. The implementation assumes you'll use the defaults and doesn't expose the layer selection logic.

Streaming has no buffer or seek controls. During live inference, if the visualization lags and you miss watching a critical frame, it's gone—the ledger records the text output and timestamps, but not the spatial canvas state. You can't pause, rewind, or step through frames during generation. This is fine for post-hoc analysis via replay export, but limits real-time exploratory debugging. If you're investigating why the model generated a specific token and want to inspect the lens state from two tokens prior, you need to re-run the entire prompt.

Verdict

Use if: You're doing interpretability research on Qwen models and need to understand what concepts are active during inference (not just which input tokens have high attribution), you're debugging refusal behavior or safety mechanisms and want to see when verdicts form relative to generation, you need reproducible interpretability artifacts that work in a browser without GPU access, or you're teaching mechanistic interpretability and want live demos that feel magical. Skip if: You work with model families other than Qwen and aren't prepared to spend hours fitting custom lenses, you need to track multi-token concepts or abstract reasoning that doesn't correspond to single vocabulary items, you require fine-grained control over which layers to observe or need to visualize non-greedy generation strategies like beam search, or you're looking for input attribution rather than workspace observation. This is specialized infrastructure for a specific interpretability paradigm—essential if you're in that paradigm, irrelevant if you're not.