Get started
Work

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

Year: 2026
Type: preprint
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.
DOI: 10.48550/arxiv.2601.19208
Open Access status: green
Funders: National Science Foundation, University of Wisconsin-Madison, Open Philanthropy Project, Australian Government, Australian Research Council, Office of Naval Research, Air Force Office of Scientific Research
Awards: G1523888516, SLES: Foundations of Safety-Aware Learning in the Wild, Graduate Research Fellowship Program (GRFP), G6098521345, G6108913550, G6602531639, Transfer Learning for Reliable Data Detection in Open-set Environments, CAREER: Foundations of Human-Centered Machine Learning in the Wild, G8876996369
OpenAlex ID: W7125979386
Cites:
Cited by:
Related to: