A new visualization tool demonstrates the attention mechanism in transformer-based large language models (LLMs). During the generation phase, LLMs utilize information from previous tokens selectively, which is crucial for effective text generation. The visualization allows users to interact with generated tokens to see which past tokens influenced the current generation. This is achieved by calculating attention weights, scaled by the magnitude of value vectors, across all attention heads and layers. The resulting visualization highlights the most significant past tokens, with the largest values having full opacity.
For instance, in the 'Office Move Summary' prompt, users can observe how specific tokens, such as addresses and dates, are copied verbatim. This raises questions about the probabilistic nature of LLMs and their ability to accurately reproduce information without errors. The visualization clarifies that LLMs do not predict sequences solely from limited internal states; they can draw from all past tokens, minimizing the likelihood of mistakes.
In another example, the model demonstrates its capability to reproduce a JavaScript function with minor modifications, although it requires hints to identify issues. The visualization also shows how the model combines information from different phrases, enhancing its output.
The tool is built as a basic React application using Transformers.js for text generation. Due to the need for additional data extraction from the model, the standard generation loop could not be used, necessitating a custom implementation. To improve user experience, pre-generated prompts are loaded instantly, avoiding delays associated with downloading large models.
The visualization also faces challenges in accessing certain internal values, as the model uses .onnx files that encapsulate the entire computation graph. Modifications to the .onnx file were necessary to expose these internal values, leading to the creation of a separate instrumented model hosted on Hugging Face. The code for the visualization is available in a GitHub repository.