Project Goal
I wanted to understand what a language model outputs before it turns that output into a word. The useful questions were concrete: how text becomes token IDs, how a model scores an entire vocabulary, how much probability belongs to the alternatives, why some decisions are uncertain while others are almost forced, and how the KV cache keeps the loop moving one token at a time.
The product goal was to make those steps visible without inventing data. Every probability and attention weight comes from a real model forward pass, and every point comes from a committed projection of Qwen’s input embeddings. The map is a high-level view of the vocabulary, not a literal diagram of the transformer’s internal computation.
| Model | Qwen2.5-0.5B-Instruct · 494M parameters |
|---|---|
| Atlas | 151,665 decodable vocabulary tokens |
| Live data | Top 200 probabilities + five attention positions per step |
| Projection | 896 dimensions to 3 · 77.73% measured preservation |
| Selection | Greedy argmax at temperature 1 |
How One Token Is Chosen
I wrote a manual autoregressive decode loop instead of callingmodel.generate()because a completed string does not contain the intermediate data the visualization needs. The first pass processes the tokenized prompt. Every later pass receives only the token that was just chosen while the KV cache carries the earlier context.
- The final logit row provides one score for every vocabulary row.
- Softmax at temperature 1 converts those scores into probabilities.
- The 200 highest decodable candidates are kept for the interface.
- Greedy argmax selects the highest-probability token.
- The backend averages the final layer across all 14 attention heads and keeps the five highest-weight context positions.
- A WebSocket event sends the decision to the browser.
The top 200 are display data, not a top-k sampling rule. The shipped demo is deliberately greedy so a run is reproducible. Attention is also presented narrowly: it shows which earlier positions received more weight under one stated aggregation, not why the model chose a token or a hidden chain of thought.
Building the 3D Atlas
Each vocabulary item begins as an 896-number input-embedding vector. I L2-normalized the first 151,665 decodable rows and compared the token neighbourhoods in the original space with the neighbourhoods produced by several 3D projections.
The first implementation used PCA because it was deterministic and easy to explain. Visually it looked structured, but the measurement showed that it preserved only 7.2% of the relevant neighbourhoods. Three-dimensional UMAP reached 77.7%, so the quality gate changed the implementation.
| Projection | Neighbourhood preservation |
|---|---|
| Random placement | 1.0% |
| PCA in 3D | 7.2% |
| UMAP in 3D | 77.7% |
UMAP normally changes when it is refit, so the projection runs offline once with a fixed seed and the coordinates are committed. A token keeps the same position across prompts, sessions, deployments, and timeline scrubbing. The browser receives the field as a 1.8 MB little-endian Float32 buffer instead of parsing 151,665 coordinate objects from JSON.
Rendering 151,665 Tokens
Rendering one React component or Three.js object per token was never practical. The static field uses oneTHREE.Pointsobject backed by typed arrays and buffer geometry. The measured scene needs roughly four to five draw calls, including the live candidate, attention, and path overlays.
A custom point shader turns each sprite into a shaded sphere impostor, writes depth so nearby points occlude distant ones, and keeps point size constant in screen pixels. This lets zoom reveal individual tokens without every node growing at the same time. Cool colour belongs to the static map; amber identifies the chosen token; violet carries attention; and the generated sequence stays warm so it remains distinct from the field.
Playback and Inspection
Model generation and visual playback use separate clocks. The backend can produce a local answer at about 32 tokens per second, but that is too fast to study. The browser stores every step as it arrives, then a playback controller reveals those events at 0.5×, 1×, 2×, or 4×. A visitor can pause, restart, scrub to any token, or inspect one of the ranked alternatives without rerunning the model.
The candidate nucleus is adaptive rather than fixed. Token Atlas lights the smallest ranked set that reaches the selected probability coverage. In the recorded France example, 11 candidates hold 99% of the mass; for several later steps, one token is enough. That makes uncertainty visible instead of making every decision look equally open.
Shipping the Full Stack
The shipped system has three data paths: an offline projection creates fixed vocabulary coordinates, FastAPI runs Qwen one token at a time and streams events over a WebSocket, and React stores those events while Three.js renders the selected playback step. There is no database, account system, or prompt persistence. The only permanent data is the projection, token labels, a real recorded fallback run, and the model weights baked into the backend image.
The frontend is deployed on Vercel and served through the portfolio’s/projects/token-atlas/demopath. A host-neutral Docker image runs FastAPI and CPU-only PyTorch on Modal. Baking the 953 MB model weights into the image reduced model load time after boot to about 1.2 seconds. The measured release used about 690 MB at peak and had a 15.3-second scale-to-zero cold boot.
Public inference also changed the engineering work. The repository includes strict request validation, exact WebSocket origin checks, per-client and global rate limits, queue and connection limits, and a hard container ceiling. The interface names connecting, waking, loading, and ready as separate states, while a real recorded generation keeps the atlas useful before the cloud model is warm.
Decisions and Limitations
Static neighbour links
I removed a 400,139-edge graph because lines already represented attention. Reusing that visual language made the live data harder to read.
Post-processing bloom
The dependency produced a black canvas with the versions in use. A larger, faint copy of only the meaningful live points creates a cheaper and more controlled glow.
Retrieval scope
I cut the planned RAG phase after generation worked end to end. Finishing one honest system was stronger than shipping two shallow ones.
Small-model accuracy
The 494M-parameter model keeps CPU hosting practical, but it can be confidently wrong. The Mudkip preset deliberately leaves that limitation visible.
The points are vocabulary tokens, not neurons. Their positions come from the model’s static input embeddings, not context-dependent hidden states. Three-dimensional UMAP loses some of the original structure, and the default radial spread intentionally changes absolute distance for legibility while preserving direction and radial order. The raw coordinates remain available at a spread setting of zero.
The interface is intentionally desktop-only, but the supported layout includes explicit focus states, reduced-motion behavior, named controls and live regions, a keyboard-navigable candidate list, and non-colour differences for attention weight and the current token.
Result and What I Learned
Token Atlas turned concepts I had mostly understood at the API level into an end-to-end system I could inspect: tokenization, logits, probability distributions, greedy decoding, attention, embeddings, and KV caching. The main data-visualization lesson was that every visual property needs one defensible meaning. Measuring the projection before trusting its appearance was just as important as drawing it.
The final project combines a real open-weight model, a manual streaming decode loop, a measured 3D projection, GPU buffer rendering, independent playback, cold-start-aware product states, and public-service controls. Nineteen automated tests cover the event contract, projection transforms, precision, motion, request validation, and access controls, alongside frontend linting, a production build, and a zero-vulnerability runtime dependency audit at the project’s August 25, 2026 check.