# Rethinking Associative Memory Mechanism in Induction Head

arXiv: https://arxiv.org/abs/2412.11459 · HTML: https://mlsanity.com/paper/2412.11459/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Reasoning benchmarks and evaluation, one of 17 directions in LLM reasoning.
- Builds on 36 of its 63 references that are in the map.
- First posted on arXiv 16 Dec 2024.

## Placed in

- LLM reasoning / [Reasoning benchmarks and evaluation](https://mlsanity.com/lr/reasoning-benchmarks-and-evaluation/)

## Builds on (most cited first, 30 of 36)

- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Efficient Estimation of Word Representations in Vector Space](https://mlsanity.com/paper/1301.3781/) (2013)
- [DeBERTa: Decoding-enhanced BERT with Disentangled Attention](https://mlsanity.com/paper/2006.03654/) (2020)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [Self-Attention with Relative Position Representations](https://mlsanity.com/paper/1803.02155/) (2018)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)
- [Progress measures for grokking via mechanistic interpretability](https://mlsanity.com/paper/2301.05217/) (2023)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Dissecting Recall of Factual Associations in Auto-Regressive Language Models](https://mlsanity.com/paper/2304.14767/) (2023)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Dense Associative Memory for Pattern Recognition](https://mlsanity.com/paper/1606.01164/) (2016)
- [Rethinking Positional Encoding in Language Pre-training](https://mlsanity.com/paper/2006.15595/) (2020)
- [One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention](https://mlsanity.com/paper/2307.03576/) (2023)
- [Analyzing Transformers in Embedding Space](https://mlsanity.com/paper/2209.02535/) (2022)
- [What and How does In-Context Learning Learn? Bayesian Model Averaging, Parameterization, and Generalization](https://mlsanity.com/paper/2305.19420/) (2023)
- [The Evolution of Statistical Induction Heads: In-Context Learning Markov Chains](https://mlsanity.com/paper/2402.11004/) (2024)
- [How Transformers Learn Causal Structure with Gradient Descent](https://mlsanity.com/paper/2402.14735/) (2024)
- [Randomized Positional Encodings Boost Length Generalization of Transformers](https://mlsanity.com/paper/2305.16843/) (2023)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (2023)
- [Improve Transformer Models with Better Relative Position Embeddings](https://mlsanity.com/paper/2009.13658/) (2020)
- [Length Generalization in Arithmetic Transformers](https://mlsanity.com/paper/2306.15400/) (2023)
- [Encoding word order in complex embeddings](https://mlsanity.com/paper/1912.12333/) (2019)
- [Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers](https://mlsanity.com/paper/2409.10559/) (2024)
- [Learning Transformer Programs](https://mlsanity.com/paper/2306.01128/) (2023)
- [SHAPE: Shifted Absolute Position Embedding for Transformers](https://mlsanity.com/paper/2109.05644/) (2021)
- [Scaling Laws for Associative Memories](https://mlsanity.com/paper/2310.02984/) (2023)

## Abstract (arXiv)

Induction head mechanism is a part of the computational circuits for in-context learning (ICL) that enable large language models (LLMs) to adapt to new tasks without fine-tuning. Most existing work explains the training dynamics behind acquiring such a powerful mechanism. However, the model's ability to coordinate in-context information over long contexts and global knowledge acquired during pretraining remains poorly understood. This paper investigates how a two-layer transformer thoroughly captures in-context information and balances it with pretrained bigram knowledge in next token prediction, from the viewpoint of associative memory. We theoretically analyze the representation of weight matrices in attention layers and the resulting logits when a transformer is given prompts generated by a bigram model. In the experiments, we design specific prompts to evaluate whether the outputs of the trained transformer align with the theoretical results.