LensLapse: Watching Language Models Learn with an In-Browser Logit Lens over Training Checkpoints

Tatsuki Okada · Yusuke Sakai · Taro Watanabe

Nara Institute of Science and Technology (NAIST), Japan

Scrub across Pythia's public training checkpoints and watch next-token predictions crystallize from noise into knowledge — layer by layer, with zero backend.

🚀 Live Demo 💻 Code 🤗 Checkpoints

Abstract

The logit lens is a standard instrument for inspecting how a transformer's next-token prediction forms across layers — yet every hosted lens tool inspects a single finished model. LensLapse turns training time into a first-class axis of token-level interpretability: scrub a slider across 38 public checkpoints of Pythia-70M and watch per-layer, per-position lens predictions crystallize from early frequency bias into factual knowledge. The same architecture-generic recipe also converts three multilingual checkpoint suites — MAP-Neo-250M, BAAI Aquila-135M, and BLOOM-560M — with no changes to the lens itself. Curated prompts render instantly from precomputed data, while any free-text prompt is probed live in the browser via WebGPU/WASM ONNX inference — no server, no API key, and your prompt never leaves your device. The in-browser lens matches fp32 PyTorch with 100% top-1 agreement on our evaluation suite (where int8 quantization degrades to as low as 52%), and warm probes complete in 63–316 ms on Chromium/WebKit across three model sizes (14M–160M).

Highlights

A time axis for the logit lens

Existing lens tools show one finished model. LensLapse scrubs 38 checkpoints of Pythia-70M (step 0 → 143k) with exact P(token) trajectories — revealing fact acquisition that is visibly non-monotone: learned, lost, and sometimes relearned.

Fully in-browser, zero backend

Free-text prompts run via onnxruntime-web (WebGPU, with WASM fallback) or an optional local probe server — one click now probes every checkpoint automatically, so ▶ playback works for anything you type, just like the curated prompts. A static site plus versioned weights on the HF Hub means unlimited concurrent users at zero hosting cost — the demo cannot rot.

Your checkpoints, one pip install

pip install then lenslapse server serves the app and the probe API on one local port. Point it at any HF model or a local Trainer run — no ONNX conversion, no config files. The recipe is architecture-generic: GPT-NeoX and GPT-2-style LayerNorm models, plus Llama-style, Mistral-style, and Gemma-style (plus-one-weight) RMSNorm models, all pass the parity check — shipped for MAP-Neo, Aquila, BLOOM, SmolLM2, Qwen3, OPT, and Gemma 3 alongside Pythia and GPT-2. A layout the generic heuristic can't reach registers an explicit override in one line, with no change to the resolver itself.

Fidelity you can verify

By construction, lens(h_L) equals the model's logits — asserted per checkpoint at export. fp16 weight storage keeps 100% top-1 agreement with fp32 PyTorch on our evaluation suite; dynamic int8 drops to 52%, so we do not ship it.

Watchable by design

One click plays the whole time-lapse, cells flashing as their top-1 changes. A top-10 bar race and the classic p-vs-layer profile animate at the pinned cell; an acquisition map shows when each cell first learned its final answer, a diff view shows what changed since any reference checkpoint, and a compare mode renders a second model in lockstep — falling back to a live probe of its own when the compared model has no precomputed match for the prompt.

Scriptable end to end

The core actions have CLI twins — lenslapse probe/trace/models — and the app can copy the exact CLI or cURL command that reproduces the current view. Track any token across training, share permalinks, or embed a live grid anywhere with ?embed.

Screenshots

The full LensLapse interface: logit-lens grid, training-step slider, trajectory panel, and emergence stories — The full interface at step 8,000 (Pythia-70M): the logit-lens grid with its view toggles (left), the log-scale training-step slider with the play button (top), and — for the pinned cell — exact probability trajectories, the p-vs-layer profile, and the top-10 race (right), above curated emergence stories.

Logit-lens grid at training step 8,000 where the token Tokyo first reaches top-1 at the last layer — The lens grid for “The capital of Japan is the city of” at step 8,000 — “ Tokyo” first reaches top-1 at the last layer, while lower layers still predict echoed or generic continuations.

How it works

Each checkpoint is exported once into a pair of ONNX graphs: a backbone that returns the pre-layer-norm residual stream after the embedding and every block, and a small lens head (final layer norm + unembedding) applicable to any layer's hidden state. Curated prompts are precomputed with fp32 PyTorch into static JSON shards; a live prompt is tokenized once and then probed against every checkpoint in turn — one forward pass each, pushing every (layer, position) hidden state through the lens head in a single batched matmul — so the same time-lapse playback works for it too. Weights are stored fp16 and cast to fp32 at session load.

Pythia checkpoint (HF Hub, revision step{N})
   ├─ backbone.f16.onnx   input_ids → hidden states [L+1, T, H]
   ├─ lens.f16.onnx       hidden [N, H] → logits [N, V]
   └─ precompute          → static JSON shards (top-10 per cell + trajectories)

browser (static site, no backend)
   ├─ precomputed mode: JSON shard → canvas grid + SVG trajectories
   └─ live mode: onnxruntime-web (WebGPU → WASM) + transformers.js tokenizer

BibTeX

@misc{okada2026lenslapse,
  title  = {LensLapse: Watching Language Models Learn with an In-Browser Logit Lens over Training Checkpoints},
  author = {Tatsuki Okada and Yusuke Sakai and Taro Watanabe},
  year   = {2026},
  url    = {https://github.com/iamtatsuki05/lenslapse}
}