Interpreting the output

At every generation step the model produces a [batch, vocab] logit tensor. The probe converts those logits into probabilities and records, for both the raw and the sampling-processed distributions:

Metric Meaning
top_ids / top_probs / top_logprobs Top-k candidate tokens and their probabilities
k_used Adaptive top-k: smallest k whose cumulative mass ≥ mass_threshold, clamped to [min_k, max_k]
top_mass_used Cumulative probability mass covered by k_used
entropy Shannon entropy of the full next-token distribution (nats)
selected_id / selected_prob / selected_logprob The token actually sampled and its likelihood under the analyzed distribution
selected_rank Rank of the selected token across the full vocabulary

The adaptive top-k matters because a fixed top-20 wastes detail on confident steps and clips it on uncertain ones:

Step 1: top-8 tokens cover 95 % of mass
Step 2: top-45 tokens needed for 95 %
Step 3: top-12 tokens cover 95 %

The probe runs twice per step — once on the raw temperature-scaled logits, once on the post-sampling logits (after top_p / top_k filtering) — so you can see how sampling reshapes the distribution.

Reading the plots

Signal Meaning
Low k_used (≈ min_k) Confident step — distribution is concentrated
High k_used (≈ max_k) Uncertain / diffuse — many plausible continuations
Low selected_prob Sampling drew a low-probability token (high temp / wide top_p / open-ended step)
High selected_rank The chosen token wasn't among the top candidates — useful for debugging sampling settings
Big gap between raw and processed k_used Sampling filters are aggressively pruning the natural distribution

Direct Logit Attribution (the Attribution lens)

With --capture-full-activations, each generated token's logit is decomposed into additive contributions through the unembedding, with the final norm folded as a fixed scale (standard "direct logit attribution"):

logit(target) ≈ embed + Σ_layers ( attn_L + mlp_L ) + error

where attn_L is layer L's attention block (o_proj output) and mlp_L its MLP block (mlp_out). The Attribution lens shows these as diverging bars — orange promotes the token, blue suppresses it — sorted by impact.

  • error / unexplained bar — for RMSNorm models (Qwen / Llama / Mistral / …) the decomposition is exact, so error ≈ 0; a large error means the fold-norm linearization or a norm variant isn't fully captured. Always read it — it is how much of the logit the bars do not explain.
  • Per-head — expand an attention bar to split it into per-head contributions, W_O[:, head] · z_head folded through the norm. These sum exactly to the layer's attn bar, so you can see which head wrote the token (induction / name-mover heads, etc.).

This is the direct (OV) path only: it explains how information already at the final position maps to the logit, not how attention patterns formed (QK circuits). It is correlational — to confirm a contribution is causal, ablate it.

Interventions / ablation

Direct logit attribution is correlational; ablation is the causal test. Each attribution bar is a hypothesis you can check from the Python library / CLI (llm_token_heatmap.intervention, examples/dla_causal_validation.py): re-run the forward pass with a component or head's last-position output zeroed (or scaled), then read how the next-token distribution moved:

Readout Meaning
KL (nats) Divergence between the baseline and patched next-token distributions — how much the ablation moved the output overall
P(target) before → after The realized token's probability change; a top contributor's ablation should drop it noticeably
Top-token flips Where the ranked candidates reordered (e.g. #1 " Paris" → " London")

A faithful attribution predicts the intervention: ablating a high-attn head should drop the target probability more than a random head. Ablation holds the attention patterns fixed (it removes a block's write to the residual at the analyzed position). Interactive ablation will return to the viewer once the CLI precomputes ablations into the trace.

Worked example — validating one head

configs/recall-probe.yaml is designed to make this concrete: a factual prompt ("The capital of France is" → Paris) at low temperature with full capture. examples/dla_causal_validation.py runs the argument on the small model in a few forward passes (no GPU). On Qwen/Qwen2.5-0.5B-Instruct:

  • The decomposition is exact — error ≈ 0, and the per-head contributions sum to the layer's attention bar to floating point (Δ ≈ 1e-7).
  • DLA isolates one dominant promoter of Paris: layer 21, head 6 (attn ≈ +3.58), far above the rest.
  • Ablating it drops P(" Paris") 0.302 → 0.232 (Δ −0.070, KL 0.051); a near-zero-DLA control head (L1 h3) moves it only −0.002 — ~36× less. Ablating the whole L21 attention block drops it to 0.097, so head 6 is the largest piece of a layer-21 fact-writing computation.

The attribution predicted the intervention — the difference between a suggestive chart and a causal claim.

This single-head crispness is scale-dependent — and that's the honest part. Running the same argument (MODEL=… examples/dla_causal_validation.py) up the Qwen2.5 family:

Model top-head DLA single-head ablation ΔP block ablation ΔP
0.5B L21 h6 +3.58 (dominant) −0.070 (≈36× a control) −0.206
7B L23 h6 +0.70 −0.019 (≈4× a control) −0.160
14B L40 h23 +0.83 (tied with h22) −0.015 (≈ control, in the noise) −0.039

Two things hold at every scale: the decomposition stays faithful (per-head sums to the layer bar to floating point; error small), and block-level attribution stays causally crisp (ablating the top attention block always moves the answer most). What washes out with size is the single-head story: bigger models spread fact-writing across many redundant heads, so no one head is load-bearing. Two compounding reasons: (1) redundancy/superposition grows with scale, and (2) a single-head ablation mixes the head's direct contribution (≈ its DLA) with a second-order final-norm rescaling (removing a vector shifts the RMSNorm denominator, rescaling all logits) — which dominates the comparison when direct contributions are small. So: trust per-head DLA as a decomposition; reach for block-level ablation (or ablating the top-k heads together) when you want a crisp causal handle on a larger model.

Manifold metrics

token-heatmap manifold (CLI) treats the captured activations for one (layer, submodule) as a cloud of points — one per token position — and measures its geometry. The motivation comes from “When Models Manipulate Manifolds”: models often encode a scalar (there, characters-until-line-break) on a smooth, low-dimensional, curved manifold, so the interesting structure is geometric, not per-neuron.

Metric Meaning What to look for
Participation ratio (Σλ)² / Σλ² over the PCA eigenvalues — a smooth "effective number of dimensions". A small value (e.g. 2–3) despite a wide hidden dim ⇒ the cloud really lives on a low-dimensional manifold.
Intrinsic dim (TwoNN) Geometry-based dimension estimate that, unlike PCA, sees curvature: a 1-D curve coiled in 3-D reads as ≈1. Lower than the participation ratio ⇒ the cloud is a curved manifold, not a flat subspace. Unreliable for very short / regularly-spaced traces — prefer the participation ratio there.
Trajectory curvature Mean turning of the position-ordered path through PCA space. ≈0 ⇒ a straight sweep; large ⇒ the representation bends sharply as generation proceeds.
Periodicity (period · power) Dominant period of the leading projection component (FFT) and its normalized power. High power at a clean period is the signature of a circular / helical coordinate — the line-break "counting" manifold.
Variance spectrum (scree) Explained-variance fraction per principal component, with the cumulative curve. Variance collapsing into the first few bars ⇒ low-dimensional structure.

The Manifold tab in the web app shows the 2-D PCA projection (coloured by step, with the trajectory drawn through it) alongside these metrics, so you can see the manifold and read its summary numbers together.

For local CPU / small GPU:

  • Qwen/Qwen2.5-0.5B-Instruct
  • Qwen/Qwen2.5-1.5B-Instruct
  • TinyLlama/TinyLlama-1.1B-Chat-v1.0
  • microsoft/phi-2

For better output quality (GPU recommended):

  • Qwen/Qwen2.5-3B-Instruct
  • Qwen/Qwen2.5-7B-Instruct
  • meta-llama/Llama-3.1-8B-Instruct (gated — set HUGGINGFACE_HUB_TOKEN)

Tip: pass use_chat_template=True (or --use-chat-template on the CLI) when using instruct models, otherwise the model sees a malformed prompt and tends to ramble.