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

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/