attention visualization

Explore how your phone visualizes AI attention with each word it chooses

Watching your phone visualize attention as it generates each word feels like peeking inside a working brain.

3 min readMachine Learning

The most honest thing about Talha Tahir's attention visualizer is that it looked broken at first. Running SmolLM2-135M entirely in the browser, generating one token at a time on a phone, he expected to see the model's focus spread intelligently across the words it had already written. Instead, nearly everything pointed back to the first token, a phenomenon he initially assumed was a bug in his own code. It wasn't. It was the attention sink, a quirk of modern transformer training where models learn to dump excess attention onto a single, often meaningless, position. That discovery, tucked inside a neat interactive demo, is worth pausing on.

What Tahir built is genuinely useful for anyone who has ever stared at a spreadsheet of perplexity scores or a training loss curve and wondered what the model is actually doing. His Attention Galaxy lets you type a prompt, watch the model propose its next token with probabilities, and then click any generated word to see exactly which earlier tokens it attended to, ranked and weighted. No server, no GPU, no setup. Just a browser tab and a curious mind. This fits a broader pattern we have been tracking: tools that make model internals legible to non-specialists. Consider how hierarchical routing cuts memory while preserving long-context accuracy or how an AI learns Clash Royale defense in a live browser demo. Each of these projects shares the same instinct: if you can see the mechanism, you can start to trust it.

The attention sink is not just a footnote. It changes how you interpret any attention visualization you come across, including Tahir's own. When you see a strong link from a later token back to position one, you should not assume semantic importance. It might just be the model's way of saying "I need to store some extra energy here," a structural artifact rather than a linguistic one. That is a practical lesson for anyone building on transformer outputs, whether you are debugging a retrieval pipeline or explaining a model's reasoning to a stakeholder. The demo does not merely show you attention; it teaches you to read it skeptically.

Our take is simple: this is the kind of tool that should be standard in every ML onboarding deck. The interactive training visualizer we covered recently showed backpropagation step by step; Tahir's project does the same for inference-time behavior. The next step would be extending this beyond a 135M model to something larger, or adding a toggle to compare attention sinks across different architectures. For now, the concrete takeaway is this: when you see a model attending heavily to its first token, do not assume it is confused. It is probably just doing housekeeping. That single insight, gained by watching a phone generate a sentence one word at a time, is worth more than another hour of reading papers.

From Machine Learning

I made a little thing that visualizes attention. It runs SmolLM2-135M fully in the browser (onnxruntime-web, int8, no server, works on my phone).

You give it a prompt, it generates one token at a time, and for each step you get the next-token candidates with probabilities, plus the attention from the current token back to the earlier ones drawn as links with the weights labeled. You can click any word to see where its attention actually went, ranked.

Read the original at Machine Learning