Fully Interpretable Minimal Transformers: From Geometry to Algorithm
cs.LG, cs.AI
Submitted: 2026-10-07
Updated: 2026-10-07
Comments: 27 pages, 15 figures, 2 tables. Code and training-dynamics animations: https://github.com/Raneem-mahajne/creating_transformer
Code: https://github.com/Raneem-mahajne/creating_transformer
License: http://creativecommons.org/licenses/by/4.0/
The gist: We present a framework for building and interpreting minimal transformer models.
Terminology
Abstract
We present a framework for building and interpreting minimal transformer models. By constraining a transformer's embedding dimension and head size to 2, we enable full two-dimensional visualization of its internal representations. Embeddings, query/key/value transforms, attention outputs, residual streams, and decision boundaries can all be seen directly. Our central claim is that the learned geometry implies an algorithm; the arrangement of points and boundaries in R squared can be read as a step-by-step procedure. We train a transformer on a simple task where it must produce the most recently observed even number whenever the '+' operator appears in a sequence of digits. Once trained, we visually walk through every step of the transformer's computation. We show how the model embeds the tokens and their respective positions in the sequence, transforms them via the Q, K, and V matrices, uses the dot product between the Q and K representations to form the attention matrix, and uses the attention matrix to select values that move the representation of each input token to the region of the domain of the output layer that will correctly predict the next token. We introduce a suite of interpretability visualizations that make the algorithmic interpretation of this procedure explicit. Our framework offers a pedagogical and experimental testbed to explore how transformers use informational geometry to implement next-token prediction.
Sources
- Eliciting Latent Predictions from Transformers with the Tuned Lens
- What Does BERT Look At? An Analysis of BERT's Attention
- Learning Transformer Programs
- Grokking modular arithmetic
- How does GPT-2 compute greater-than?: Interpreting mathematical abilities in a pre-trained language model
- Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task
- Tracr: Compiled Transformers as a Laboratory for Interpretability
- Towards Understanding Grokking: An Effective Theory of Representation Learning
- UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction
- Clustering and Alignment: Understanding the Training Dynamics in Modular Addition
- Progress measures for grokking via mechanistic interpretability
- The Linear Representation Hypothesis and the Geometry of Large Language Models
- Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets
- Understanding Addition in Transformers
- Attention Is All You Need
- Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small
- The Clock and the Pizza: Two Stories in Mechanistic Explanation of Neural Networks
- What Algorithms can Transformers Learn? A Study in Length Generalization
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks