The primary goal of this project is to develop interactive animations that will enable us to gain a deep understanding of the LLMs. At a superficial level, this simply means tracing out the architectures of different models.
What I'm imagining for the superficial version, which should be one tab of the demo, is to use something like this ELIZA demo, from my LLM course. Ideally, like that demo, it could run entirely in a browser. I had started to put together a transformer analog to that demo (as part of this demo), but is not nearly as comprehensive as the ELIZA version. Another source of inspiration could be the animated transformer. But instead of (solely) animating the equations, it would be useful to have animated components that can be clicked on and manipulated. E.g., representing a matrix as a rectangle, where (zoomed out) it could represent a component in a larger system, but clicking on it could enable the user to see actual values. Ideally the user could also select or provide additional documents/data sources (e.g., from huggingface or things they upload) to show how training works and affects everything. And similarly, ideally the user could also provide their own text, system prompts, etc., and the demo should update accordingly (showing how the text is turned into a sequence of tokens, and then how those tokens are processed).
But the deeper point is to understand the effects of going through each step, e.g., from a physical perspective (considering each token as a particle, and matrix multiplications as forces acting on those particles) and a geometric perspective (visualizing embedding spaces as manifolds, rendered as 3D animated surfaces, and matrices acting on those embeddings as vector fields over the manifolds-- visualized as quiver style plots with arrows rendered on the manifolds' surfaces). This latter idea comes in large part from this paper.
Approach
The superficial and "physics/geometry/visualization" demos have fundamentally different goals, and likely we will want to have different (real) models underlying each. For the superficial demo, we should have a real LLM (something like SmolLM2-- small enough to run in a browser, but powerful enough to lightly interact with in a conversation; that model came out in 2024 so there are likely to be better options now!). It needs to be 100% open (architecture AND weights) so that every component can be directly explored and manipulated.
For the deeper demo, a small (probably built from scratch) model will, I think, be more practical. Every component should be chosen specifically to facilitate visualization. So we probably want something like 3-D embeddings, just 3 or 4 transformer layers, relatively short maximum context windows (50 tokens max?), possibly a relatively small vocabulary and/or a simple tokenizer, and so on.
I'm imagining:
- Start with a sphere; embeddings evenly tile the surface. Limited vocabulary (1000 words?).
- Input a prompt. Each new token induces a "vector field" over the sphere-- it tells us: if the next token was this (i.e., the token represented at that position), which direction is the token after that?
- This vector field idea can be examined for each transformer layer individually, OR for the full model (across all layers).
- We should also be able to separate out the effects of the key/query/value matrices. The key/query matrices give us "attention": to what extent is this token's impact influenced by previous tokens in the sequence? I want to be able to visualize the impact-- and also be able to try out (or edit directly) different interesting/illustrative matrices (e.g., identity, toeplitz matrices with a "fuzzy" diagonal, pure random, random + autocorrelated, and so on) for each step, and see what the effect is. The "value" matrix is (according to the Latifi Jabelli paper) what imposes the "vector field" (forces) over the manifold's surface, so we should be able to visualize and edit that too.
- Every step of the visualization needs to map onto a specific transformer step, and ALL steps must be represented. In other words, there needs to be a 1-to-1 correspondence between parts of the visualization and parts of the model. We could use a torch GPT2 model, or something else small, open, and customizable for our purposes.
Other considerations
- A few of the example visualizations already in this repository start to get at some of these ideas. But they're not going in quite the right direction. They're also not smooth enough to interact with and they don't have the visual cleanliness and "pop" that would make them fun to play around with. We need a better rendering engine, and more optimizations of performance.
- Many aspects of these ideas are ambiguous or potentially even pushing in different and even incompatible directions. But in reconciling those ambiguities and inconsistencies, it's critical that we keep the broader goal in mind: feature complete, accurate, and interactive ways of exploring and understanding these models.
- For now, let's focus on standard decoder-only transformer models. But we should design the infrastructure to enable us to add support for other architectures later-- BERT (decoder + encoder), text diffusion models, multi-modal models, mixture-of-experts models, and potentially even some sort of GUI with a drag and drop interface for constructing or manipulating totally new models, and then watching the impact on how sequences are processed and predicted.
The primary goal of this project is to develop interactive animations that will enable us to gain a deep understanding of the LLMs. At a superficial level, this simply means tracing out the architectures of different models.
What I'm imagining for the superficial version, which should be one tab of the demo, is to use something like this ELIZA demo, from my LLM course. Ideally, like that demo, it could run entirely in a browser. I had started to put together a transformer analog to that demo (as part of this demo), but is not nearly as comprehensive as the ELIZA version. Another source of inspiration could be the animated transformer. But instead of (solely) animating the equations, it would be useful to have animated components that can be clicked on and manipulated. E.g., representing a matrix as a rectangle, where (zoomed out) it could represent a component in a larger system, but clicking on it could enable the user to see actual values. Ideally the user could also select or provide additional documents/data sources (e.g., from huggingface or things they upload) to show how training works and affects everything. And similarly, ideally the user could also provide their own text, system prompts, etc., and the demo should update accordingly (showing how the text is turned into a sequence of tokens, and then how those tokens are processed).
But the deeper point is to understand the effects of going through each step, e.g., from a physical perspective (considering each token as a particle, and matrix multiplications as forces acting on those particles) and a geometric perspective (visualizing embedding spaces as manifolds, rendered as 3D animated surfaces, and matrices acting on those embeddings as vector fields over the manifolds-- visualized as quiver style plots with arrows rendered on the manifolds' surfaces). This latter idea comes in large part from this paper.
Approach
The superficial and "physics/geometry/visualization" demos have fundamentally different goals, and likely we will want to have different (real) models underlying each. For the superficial demo, we should have a real LLM (something like SmolLM2-- small enough to run in a browser, but powerful enough to lightly interact with in a conversation; that model came out in 2024 so there are likely to be better options now!). It needs to be 100% open (architecture AND weights) so that every component can be directly explored and manipulated.
For the deeper demo, a small (probably built from scratch) model will, I think, be more practical. Every component should be chosen specifically to facilitate visualization. So we probably want something like 3-D embeddings, just 3 or 4 transformer layers, relatively short maximum context windows (50 tokens max?), possibly a relatively small vocabulary and/or a simple tokenizer, and so on.
I'm imagining:
Other considerations