💬 NLP

How Do Transformers Learn to Associate Tokens: Gradient Leading Terms Bring Mechanistic Interpretability

This paper provides a mechanistic interpretation of how transformers learn semantic associations by deriving closed-form expressions for early-stage weights as compositions of bigram, token-interchangeability, and context mappings based on training gradient dynamics.

Original authors: Shawn Im, Changdae Oh, Zhen Fang, Sharon Li

Published 2026-05-14

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Shawn Im, Changdae Oh, Zhen Fang, Sharon Li

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine a Transformer (the brain behind modern AI chatbots) as a giant, blank notebook. When we start training it, the pages are empty. The paper you're asking about asks a simple but profound question: How does this notebook learn to connect words like "bird" and "flew" just by reading millions of sentences?

Instead of looking at the finished, complex AI, the authors decided to zoom in on the very first few moments of learning. They used a mathematical "magnifying glass" to look at the first few steps the AI takes when it starts updating its memory.

Here is the breakdown of their discovery using simple analogies:

1. The "First Step" Shortcut

Training an AI is like trying to predict the next word in a sentence. Usually, the math behind how the AI updates its memory is incredibly messy and complex.

The authors realized that at the very beginning, the AI doesn't need to do complex calculus. It just needs to follow the strongest, most obvious signal in the data. They call this the "leading term." Think of it like a hiker starting a journey up a mountain. At the very start, the path is obvious and straight. It's only later, as they get higher, that the path gets twisty and complex. The authors proved that for a surprisingly long time, the AI's "memory weights" stay very close to that simple, straight initial path.

2. The Three Building Blocks

The paper reveals that the AI's memory isn't a random mess. Instead, it is built by stacking three specific types of " Lego blocks" (which the authors call basis functions). These blocks are derived directly from the statistics of the text the AI reads:

The "Next-Word" Block (Bigram Mapping):
- Analogy: Imagine a child learning that "The" is almost always followed by a noun like "cat" or "dog."
- What it does: This block simply records: "If you see word A, what is the most likely word B to come next?" It captures simple, immediate connections.
The "Swappable" Block (Interchangeability Mapping):
- Analogy: Think of a thesaurus. If you can swap "car" with "truck" in a sentence without breaking the grammar, they are "swappable."
- What it does: This block learns that certain words play the same role. If "red" usually describes a "car," and "fast" also describes a "car," this block helps the AI realize that "red" and "fast" might be related because they both hang out with the same words. It groups words by their function, not just their immediate neighbors.
The "Context" Block (Context Mapping):
- Analogy: Imagine reading a story about a "fish." Even if the word "pond" isn't right next to "fish," the story might mention "water," "lake," or "boat" earlier.
- What it does: This block looks further back in the sentence. It learns that if a word appears in a specific context (like a story about nature), it is likely to be associated with other words from that same context, even if they are far apart.

3. How They Work Together

The paper shows that the AI's internal "weights" (the numbers that determine how it thinks) are just combinations of these three blocks.

The Output Layer (the part that guesses the next word) is mostly just the Next-Word Block. It's the AI saying, "Based on what I just saw, what comes next?"
The Attention Layer (the part that decides what to focus on) is a mix of the Swappable and Context blocks. It's the AI saying, "I see the word 'fish.' I should look back at the word 'pond' because they often appear together in similar stories, even if they aren't touching."

4. The Proof: It Actually Works

The authors didn't just write equations; they tested this on real AI models.

They trained a small AI on a dataset of children's stories.
They compared the AI's actual learned memory against their mathematical formula.
The Result: The match was incredibly close (over 99% similarity in some cases). Even when they looked at a massive, real-world AI (Pythia-1.4B), the same patterns held true. The AI was indeed organizing its memory using these three simple statistical building blocks.

The Big Takeaway

The paper argues that we don't need to view AI as a mysterious "black box" that magically understands language. Instead, at its core, it is simply statistically organizing the world by:

Noting what usually follows what.
Grouping words that do the same job.
Linking words that share a background story.

By understanding these three simple rules, we can finally see how the machine learns to associate "bird" with "flew" and "country" with "capital." It's not magic; it's a very structured way of summarizing the text it has read.

Technical Summary: How Do Transformers Learn to Associate Tokens: Gradient Leading Terms Bring Mechanistic Interpretability

Problem Statement

Large Language Models (LLMs) based on self-attention have demonstrated remarkable capabilities in capturing factual knowledge and generating coherent text, largely driven by their ability to learn semantic associations (e.g., the link between "bird" and "flew" or the interchangeability of "car" and "truck"). While recent studies have identified specific mechanisms like induction heads or linear semantic relations, a principled, theory-grounded account of how these associations emerge during the training of attention-based transformers on real-world natural language data remains lacking.

Prior theoretical works often rely on unrealistic assumptions, such as synthetic structured language, simplified architectures (lacking positional encodings or residual connections), or non-standard training regimes (e.g., sequential component-wise training). These departures from practice limit the generalizability of their insights to modern LLMs. This paper addresses the gap by asking: How do semantic associations emerge during the training of attention-based language models on natural language data?

Methodology

The authors propose a theoretical framework grounded in training dynamics, specifically analyzing the leading-term approximation of gradients during the early stages of training.

Model Setup: The study focuses on a standard $L$ -layer attention-based transformer with positional encodings (relative), causal masking, and residual streams, trained via standard full-batch gradient descent on natural language corpora.
Leading-Term Approximation: Recognizing that transformers acquire core behaviors (including semantic relations) early in training, the authors derive closed-form expressions for the weight matrices (Output $W_O$ , Value $V^{(l)}$ , Query-Key $W^{(l)}$ , and Positional Encoding $P^{(l)}$ ) by approximating the gradient updates. They show that for a sufficient number of steps $s$ (where $s \leq \eta^{-1} \min(\dots)$ ), the weights remain close to their gradient leading terms.
Basis Function Decomposition: The theoretical analysis reveals that the learned weights can be expressed as simple compositions of three fundamental basis functions derived from corpus statistics:
- Bigram Mapping ( $\bar{B}$ ): Captures next-token dependencies ( $P(e_j|e_i)$ ).
- Interchangeability Mapping ( $\Sigma \bar{B} = \bar{B}^\top \bar{B}$ ): Captures functional similarity between tokens based on the similarity of their preceding token distributions (e.g., synonyms or shared grammatical roles).
- Context Mapping ( $\bar{\Phi}$ ): Encodes longer-range prefix-suffix co-occurrence, representing the average context in which a token appears.

Key Contributions

The paper makes three primary contributions:

Explicit Weight Characterization: It presents the first explicit, closed-form characterization of weights in attention-based transformers trained on real-world text corpora under the next-token prediction loss. The weights are shown to be compositions of the three basis functions mentioned above.
Mechanistic Interpretation of Semantic Associations: The authors interpret the features learned in the weights as specific compositions of bigram, interchangeability, and context mappings. They demonstrate how these compositions allow the model to capture semantic associations (e.g., how "fish" relates to "pond" via context mapping, or how "red" relates to "truck" via bigram mapping).
Empirical Validation: The theoretical predictions are validated on both a toy 3-layer attention model trained on the TinyStories dataset and a practical, large-scale LLM (Pythia-1.4B) trained on OpenWebText.

Results

Theoretical Alignment: Experiments on a 3-layer transformer trained on TinyStories show that the cosine similarity between the theoretically predicted leading terms and the actually learned weights remains extremely high (e.g., >0.99 for Attention and Value matrices) even after 100 epochs of training.
Semantic Structure: Qualitative analysis confirms that the basis functions capture meaningful linguistic structures. For instance, the bigram mapping $\bar{B}$ correlates adjectives with the nouns they modify, while the context mapping $\bar{\Phi}$ correlates entities with their typical environments (e.g., "fish" with "pond").
Generalization to Real LLMs: When applied to Pythia-1.4B, the theoretical leading-term features show strong agreement with the learned embeddings and attention correlations, particularly in the early stages of training. The analysis reveals that while weights drift toward richer knowledge representations as training progresses, they retain the foundational associative features predicted by the theory for a significant duration.
Component Cooperation: The analysis of the full model computation (Eq. 12) illustrates how the residual stream provides an average next-token prediction, while the self-attention block (driven by $\bar{Q}$ and $\bar{\Phi}^\top \Sigma \bar{B}$ ) refines this prediction by selectively focusing on tokens most indicative of the next token based on corpus statistics.

Significance

The paper claims to provide a mechanistic foundation for representation learning in large language models by connecting deep learning training dynamics with linguistic theory. By grounding the analysis in naturalistic text distributions and standard training procedures, the work significantly reduces the gap between formal theoretical analysis and practical LLM usage.

The authors argue that their findings offer a rigorous way to contextualize the function of each transformer component. Rather than viewing the model as a black box, the theory suggests that the model's behavior develops by first establishing these fundamental associative features (bigrams, interchangeability, context), which serve as anchors for learning more complex semantic relationships. This work opens pathways for interpretability research by suggesting that weight matrices across different components can be decomposed into simple functions of shared corpus statistics, allowing for broader hypotheses about how concepts arise in models beyond specific mechanisms.