大型语言模型中注意力机制的可视化。
基于 Transformer 的大型语言模型有一个有趣的特点：在生成阶段，它能够从之前的任意一个 token 中提取信息。但它需要有选择性；如果每个 token 对生成都有相等影响，效果就不会很好。这个过程需要一个机制来决定某个 token 对下一个 token 的影响程度。
事实证明，我们可以把这个机制可视化！
你可以点击或悬停在任何生成的 token 上，查看影响*该生成的过去 token。
*“影响”可能不完全准确，因为这个可视化被高度简化。它计算注意力权重，再乘以值向量的大小，聚合所有注意力头，并在所有层中求和。然后用它来控制前面 token 的透明度。最大值的透明度始终为 1，其余的则进行插值处理。
为了将可视化限制为每个过去 token 仅一个数值，许多信息必须被舍弃。正因为如此，当我开始实现这个功能时，我实际上认为它可能不易理解。但它实际上可以产生一些有趣的模式！
例如，在默认的“办公室搬迁总结”提示中，你可以悬停在像地址和日期这样的逐字复制的文本上。你会发现原始数据非常突出，因为生成的 token 从源数据中获取了很多信息。
这解释了我之前觉得大型语言模型不直观的一点。如果它们通过概率预测下一个 token 来工作，为什么它们在复制粘贴内容方面表现得如此出色？难道它们最终不会仅凭随机机会而犯错吗？
但通过这个机制，你可以看到它并不是从一些有限的内部状态预测整个序列。由于它可以访问所有过去的标记，它可以决定从过去的哪些标记中复制，因此错误的概率可以非常低。在“调试平均函数”示例中，你可以看到这个相当小的模型（6亿参数）可以轻松复现整个 JS 函数，除了预期的修改部分。（虽然它实际上无法自己找到问题，所以需要一些提示。）
另一个有趣的部分是，当你将鼠标悬停在“办公室搬迁总结”提示中的“remain”的时候，你会看到它从“Existing employee access cards will work”中的“work”和“company phone numbers will stay the same”中的“stay the same”中提取信息。所以它有点像是将两个短语中的信息结合起来，我觉得这非常酷。
可视化本身是一个相当基础的 React 应用，使用 Transformers.js 来生成文本。但是，由于我们需要从模型中提取更多数据以进行可视化，它无法使用常规的生成循环。我不得不在应用中自定义编写生成循环，这样我们才能真正跟踪数值以进行可视化。
尽管使用了较小的模型，它仍然有数百兆字节，等待它下载完再显示任何内容是不行的。因此，我预生成了一些提示，可以立即加载和查看。
另一个棘手的问题是，可视化中的一些内容实际上并不打算被读取，因此它们未被定义为输出。我想如果你使用 Python ML 库实现它，仍然很容易访问这些内容。但 Transformers.js 使用包含完整计算图的 .onnx 文件。模型加载和计算逻辑是在 wasm 中实现的，所以据我所知，除了预定义的输出之外，没有简单的方法可以访问其他内容。
最后，我使用了一个小脚本来修改 onnx 文件，仅足以暴露那些内部值。但这意味着我不能仅使用常规的 .onnx 模型。由于我想要一个基于浏览器的生成功能，我必须将一个单独的插装模型上传到我自己的 Hugging Face 仓库：https://huggingface.co/ishamf/Qwen3-0.6B-ONNX-Instrumented，并在应用中指向那里。
您可以在 GitHub 存储库中找到代码：https://github.com/ishamf/llm-visualizer。
A visualization of the attention mechanism in LLMs.
One interesting thing about transformer-based large language models are that, during the generation phase, it is able to draw information from any of its previous tokens. But it needs to be selective; if every token affects the generation equally, it won't be very effective. This process needs a mechanism to decide how much a token affects the next token.
Turns out, we can visualize this mechanism!
You can tap or hover over any of the generated tokens to see the past tokens that affected* the generation.
* "Affected" might not be fully accurate, as this visualization is highly simplified. It's calculating the attention weight, scaled by the magnitude of the value vector, aggregated across all attention heads, and summed across all layers. This is then used to control the opacity of the previous tokens. The largest values always have an opacity of 1 and the rest are interpolated.
A lot of information had to be thrown away to limit the visualization to just one numeric value per past token. Because of that, when I started implementing this, I actually thought it might not be comprehensible. But it actually can produce some interesting patterns!
For example, in the default "Office Move Summary" prompt, you can hover over the text that are copied verbatim like the address and dates. You can then see the original data stand out quite a bit, because the generated token takes up a lot of the information from the source data.
This addresses one thing that I've previously found unintuitive about LLMs. If they work by predicting the next tokens probabilistically, why are they somehow so good at copy-pasting stuff? Won't they eventually make a mistake just by random chance?
But with this mechanism, you can see that it doesn't predict the entire sequence from some limited internal states. Since it has access to all past tokens, it can just decide which past tokens to draw from when copying, and so the probability of errors can be very low. In the "Debugging an Average Function" example, you can see that this quite small model (600 million parameters) can easily reproduce an entire JS function except for the intended modification. (Although it's not actually capable of finding the issue by itself, so it needed some hints.)
Another interesting part is when you hover over the "remain" in "Existing access cards and phone numbers remain " in the "Office Move Summary" prompt. You can see that it draws from "work" in "Existing employee access cards will work " and "stay the same" in "company phone numbers will stay the same ". So it's kind of combining the information from the words in both phrases, which I find quite cool.
The visualization itself is a pretty basic React app using Transformers.js to generate the text. But, since we need to pull more data out of the model to visualize it, it can't use the regular generation loop. I had to vibe-code the generation loop in the app so we can actually keep track of the values to visualize.
Despite using a smaller model for this, it's still hundreds of megabytes, and waiting for it to download before showing anything just won't work. So I pre-generated a bunch of prompts that can be loaded and viewed instantly.
Another tricky thing is that some of the things in the visualization are not actually meant to be read, so they're not defined as outputs. I suppose if you're implementing this using Python ML libraries, it would still be easy to access them. But Transformers.js uses .onnx files that contains the entire computation graph. The model loading and computation logic is implemented in wasm, so there's no easy way to access anything other than the predefined outputs, as far as I can tell.
In the end, I used a small script to modify the onnx file just enough to expose those internal values. But that means I can't just use the regular .onnx model. Since I want to have a browser-based generation feature, I have to upload a separate instrumented model to my own Hugging Face repo：https://huggingface.co/ishamf/Qwen3-0.6B-ONNX-Instrumented and point the app there.
You can find the code in the GitHub repo：https://github.com/ishamf/llm-visualizer.
扫码直达文章详情