Command-line interface¶

After pip install -e . (or conda env create -f environment.yml) the token-heatmap command is on your PATH.

YAML config files¶

All trace flags can be set in a YAML file so you don't have to repeat long command lines. CLI flags always override YAML values.

pip install pyyaml          # one-time; only needed for --config
token-heatmap trace --config configs/example.yaml
token-heatmap trace --config configs/example.yaml --max-new-tokens 128   # override one field

configs/example.yaml (included in the repo):

model: Qwen/Qwen2.5-0.5B-Instruct
prompt: "Explain what a large language model is in one sentence."
max_new_tokens: 64
temperature: 0.8
top_p: 0.95
out: outputs/example-run
capture_logit_lens: true

All keys are optional — any missing key falls back to the CLI default.

Basic trace¶

token-heatmap trace \
  --model Qwen/Qwen2.5-0.5B-Instruct \
  --prompt "Explain diffusion models." \
  --max-new-tokens 80 \
  --temperature 0.8 \
  --top-p 0.95 \
  --min-k 8 --max-k 64 --mass-threshold 0.95 \
  --out outputs/

Output directory:

outputs/
├── generated.txt
├── adaptive_token_trace.csv
├── adaptive_token_trace.json
├── adaptive_heatmap.png
├── entropy.png
└── selected_probability.png

Run token-heatmap trace --help for the full flag list.

Viewing the result¶

trace only writes the bundle to disk. To explore it, drag adaptive_token_trace.json onto the web viewer (cd app && npm run dev), or open it in the desktop app. See web-app.md.

Inspecting attention and the logit lens¶

AttentionProbe and LogitLens are wired through token-heatmap trace as opt-in flags. They are off by default because forcing eager attention is significantly slower than the SDPA / FlashAttention kernel.

token-heatmap trace \
  --model Qwen/Qwen2.5-0.5B-Instruct \
  --prompt "Explain attention in one sentence." \
  --max-new-tokens 12 \
  --capture-attention --attention-layers all --attention-top-k 8 \
  --capture-full-attention \
  --capture-logit-lens --lens-layers 0,3,7,11 --lens-top-k 5 \
  --out outputs/inspect/

Flag	Meaning
--capture-attention	Attach an AttentionProbe. Forces eager attention (slow).
--attention-layers	all (default) or comma-separated layer indices (e.g. 0,3,7,11).
--attention-top-k	Top-k attended positions kept inline per head (default 8).
--capture-full-attention	Also write a Tier-2 attention.<step>.npz per step under <out>/attention/. Implies --capture-attention.
--capture-logit-lens	Attach a LogitLens.
--lens-layers	Same syntax as --attention-layers.
--lens-top-k	Top-k tokens retained per layer (default 8).

Extra output files when these flags are active:

• adaptive_token_trace.json — includes inline attention aggregates, per-layer logit-lens projections, and attention_sidecar_ref pointers.
• attention_layer_head_grid.png — per-step layer × head entropy grid.
• logit_lens.png — per-layer top-k table (first step).
• selected_rank_heatmap.png — selected-token rank by layer × step.

The Logit Lens tab in the web app shows this data interactively, synced to the heatmap cursor.

Capturing activations¶

ActivationProbe captures per-layer / per-submodule summary statistics and, optionally, full activation tensors as .npz sidecars.

token-heatmap trace \
  --model Qwen/Qwen2.5-0.5B-Instruct \
  --prompt "Explain attention in one sentence." \
  --max-new-tokens 12 \
  --capture-activations \
  --activation-layers all \
  --activation-submodules residual_post,mlp_out,o_proj \
  --activation-top-k 8 \
  --out outputs/activations/

Add --capture-full-activations to also write activations/activation.<step>.npz (full hidden-state tensors) and embed activation_sidecar_ref pointers in the JSON trace. It additionally embeds, inline in the trace, the TWERA neuron ranking and Direct Logit Attribution — each generated token's logit decomposed into per-layer attention (o_proj) and MLP (mlp_out) contributions, expandable to per attention head — which the web app's Activations and Attribution lenses read. Implies --capture-activations.

Flag	Meaning
--capture-activations	Attach an ActivationProbe. Off by default.
--activation-layers	all (default) or comma-separated decoder layer indices.
--activation-submodules	Comma-separated submodule keys (default residual_post,mlp_out,o_proj). Supported: resid_pre/residual_pre, resid_post/residual_post, mlp_out/mlp.down_proj, o_proj.
--activation-top-k	Top-k highest-magnitude neurons retained per (layer, submodule) (default 8).
--capture-full-activations	Write full tensors as .npz sidecars under <out>/activations/. Implies --capture-activations.

Comparing two activation traces¶

token-heatmap diff outA/adaptive_token_trace.json outB/adaptive_token_trace.json \
  --out diff/ \
  --metric l2

The subcommand projects each input's activation subset, calls compare_activations with align="auto", and writes:

• activation_diff.json — schema-shaped diff payload.
• activation_delta.png — stacked layer × step heatmap, one subplot per captured submodule.

The CLI refuses to diff (non-zero exit) when the two parent traces have different metadata.prompt values or when zero steps align between them.

Manifold analysis¶

Analyze the geometry of the captured activation clouds — inspired by “When Models Manipulate Manifolds”, which finds a model encoding a scalar on a low-dimensional, curved (helical) manifold. For each (layer, submodule) the analysis stacks the full per-token activation vectors into a matrix and computes PCA spectrum + participation ratio, TwoNN intrinsic dimension, a 2-D/3-D projection, trajectory curvature, and FFT periodicity.

It reads the full activation vectors from the sidecars, so generate the trace with --capture-full-activations first:

# 1. generate with full activation sidecars
token-heatmap trace --config configs/example.yaml \
  --capture-activations --capture-full-activations

# 2. add the manifold analysis to the trace
token-heatmap manifold --trace outputs/adaptive_token_trace.json

This writes a top-level manifold field back into the trace (in place by default), which the web app's Manifold tab then renders. The analysis is pure-numpy — no torch needed to run it.

Flag	Default	Meaning
--trace	required	Path to a trace JSON that has activation_sidecar_refs.
--out	overwrite --trace	Write the augmented trace elsewhere instead of in place.
--layers	all captured	Subset of layer indices to analyze.
--submodules	all captured	Subset of submodule names to analyze.
--components	3	Number of PCA projection components to keep.

Exits non-zero when the trace has no activation_metadata, carries no activation_sidecar_ref (i.e. was generated without --capture-full-activations), or when no (layer, submodule) cloud has at least two positions to analyze.

See interpreting.md for what the metrics mean.

After augmenting a trace with manifold, drag its adaptive_token_trace.json onto the web viewer (cd app && npm run dev), or open it in the desktop app — see web-app.md.

Building the frontend (web build)¶

Run npm install + npm run build in app to produce a static dist/. The viewer is backend-free, so you can serve dist/ from any static file server on a host with no Node.js.

token-heatmap web build                 # output: app/dist/
# then serve it anywhere, e.g.:
python -m http.server -d app/dist 8080
# open http://localhost:8080/ and drag a trace JSON onto the page

HPC: build the GPU venv (hpc setup)¶

Idempotently build the dedicated cu124 torch venv on the HPC so token-heatmap runs on the GPU instead of silently falling back to CPU.

token-heatmap hpc setup            # build/update the GPU venv
token-heatmap hpc setup --verify   # also run a real GPU matmul check (queues a short srun)

Flag	Default	Meaning
--verify	off	Also run a real GPU matmul check.
--ssh-host	j7zang-gpu	SSH host alias.
--remote-repo	/work/j7zang/Token-Heatmap	Repo checkout on the HPC.
--remote-venv	/work/j7zang/th-gpu	GPU venv path.
--anaconda-python	base interpreter	Base interpreter used to create the venv.

HPC round-trip (hpc run)¶

One command from your laptop. It uploads the config to the HPC, submits a Slurm GPU job (the only remote step) that runs trace + manifold, polls it to completion, then rsyncs the whole outputs/<name>/ folder back — so you view it locally with no GPU and no tunnel (drag outputs/<name>/adaptive_token_trace.json onto the web viewer, or open it in the desktop app). A pre-flight check refuses runs that won't fit the GPU's VRAM before submitting.

token-heatmap hpc run configs/wrap-text.yaml --model Qwen/Qwen2.5-14B-Instruct \
  --capture activations --probe line_position --extra "--max-new-tokens 320"
# 32B on one GPU: add  --4bit

Flag	Default	Meaning
config	required	Trace config YAML (its basename is the default run name).
--name	config basename	Run name → outputs/NAME locally + on the HPC.
--model	from config	Override the model id.
--gpu	rtx6000	GPU type (rtx6000 or l40s); both 48 GB.
--qos	qos_rtx6000_max / normal	Slurm qos (defaults by GPU type).
--mem	64G / 28G	Host memory (default depends on qos).
--time	01:00:00	Walltime HH:MM:SS.
--capture	full	full = +attention (slower); activations = manifold-only.
--probe	none	Add a supervised manifold probe scalar (e.g. line_position).
--extra	none	Extra trace flags (e.g. --max-new-tokens 320).
--4bit	off	Load in 4-bit NF4 (for 32B+).
--no-manifold	off	Skip the manifold pass.
--no-sync	off	Don't git pull the HPC repo first.
--no-pull	off	Leave outputs on the HPC (no rsync back).
--setup	off	Build/verify the GPU venv on the HPC first (one-time).
--force	off	Skip the pre-flight "won't fit in VRAM" size check.

Connection / path overrides also exist (--ssh-host, --remote-repo, --remote-venv, --anaconda-python, --remote-bin-gpu, --poll-seconds); run token-heatmap hpc run --help for the full list.