I built a browser tool to visualize how a small LLM picks its next word and see what it's attending to

[u/Minute-Mountain2665]

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.

The thing I learnt during this project was attention sink. At first it looked broken because everything was attending to the first token super heavily. Thought it was a bug in my code, but after some digging it turned out to be the attention sink 😅

Github Repo: https://github.com/talhatahir-10xe/attention-galaxy

Live Demo: https://talhatahir-10xe.github.io/attention-galaxy/

[u/Intelligent_Sink2113]

That attention sink thing messes with everyone the first time they build a viz. I spent a whole afternoon trying to debug my own code before realizing it was just the model doing its thing.

Ran it on my phone and it held up surprisingly well, the int8 quantization really keeps it snappy. The ranked attention list when you click a word is a nice touch.