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

TL;DR

The work presents a gradient-dynamics based theory for how semantic associations emerge in attention-based transformers trained on natural language. By focusing on the leading term of gradient updates, the authors derive closed-form weight expressions that decompose into three basis functions: the bigram mapping $\mathbf{\bar{B}}$, the interchangeability mapping $\mathbf{\Sigma}_{\mathbf{\bar{B}}}$, and the context mapping $\mathbf{\bar{\Phi}}$, with higher-layer weights forming compositions involving $\mathbf{\bar{Q}}$ and $\Delta$ and the positional term $\Delta$. They prove that, under realistic training conditions (small initialization, standard cross-entropy loss, and early training steps), transformer weights approximate these leading-term forms, and they validate these predictions on TinyStories and OpenWebText-based models like Pythia-1.4B. The results demonstrate both a principled mechanistic interpretation of associative features and practical correlations between theoretical leading terms and actual learned weights, offering a path toward transparent and interpretable representation learning in large language models.

Abstract

Semantic associations such as the link between "bird" and "flew" are foundational for language modeling as they enable models to go beyond memorization and instead generalize and generate coherent text. Understanding how these associations are learned and represented in language models is essential for connecting deep learning with linguistic theory and developing a mechanistic foundation for large language models. In this work, we analyze how these associations emerge from natural language data in attention-based language models through the lens of training dynamics. By leveraging a leading-term approximation of the gradients, we develop closed-form expressions for the weights at early stages of training that explain how semantic associations first take shape. Through our analysis, we reveal that each set of weights of the transformer has closed-form expressions as simple compositions of three basis functions (bigram, token-interchangeability, and context mappings), reflecting the statistics of the text corpus and uncovering how each component of the transformer captures semantic associations based on these compositions. Experiments on real-world LLMs demonstrate that our theoretical weight characterizations closely match the learned weights, and qualitative analyses further show how our theorem shines light on interpreting the learned associations in transformers.

Figures (8)

• Figure 1: To understand the emergence of associative features, we analyze the training dynamics of Transformers by focusing on the gradient leading terms for weights, which allows us to identify interpretable basis functions that characterize each weight by their compositions. Empirical validation confirms that our weight characterizations match the actual ones learned in practical transformers.
• Figure 2: Illustration of theoretical results. We characterize weight matrices of the attention-only transformer as compositions of three basis functions: bigram mapping, interchangeability mapping, and context mappings. We illustrate how these mappings are composed across weight matrices to learn semantic associations between a given query token and its surrounding text.
• Figure 3: An example of $\mathbf{\bar{\Phi}}$ with arrows pointing to prefix tokens for "fish" with context summary scores on edges. Larger values indicate the token appears more frequently in the context of "fish".
• Figure 4: Cosine similarity between theoretical and learned weights. Results from a 3-layer transformer model trained on TinyStories.
• Figure 5: Selected tokens from the top 30 correlated tokens under different basis features from TinyStories. The characterized features actually capture both grammatical and semantic structures.

Theorems & Definitions (26)

• Definition 3.1: Attention-Based Transformer
• Theorem 4.1
• Lemma D.1: General Gradient Form
• proof
• Lemma D.2: Softmax Jacobian Norm
• proof
• Lemma D.3: First Gradient Step
• proof
• Lemma D.4: Second Gradient Step
• proof