Rethinking Associative Memory Mechanism in Induction Head
It addresses the problem of understanding coordination between in-context learning and pretrained knowledge in LLMs for researchers, but appears incremental as it builds on existing induction head mechanisms.
The paper investigates how a two-layer transformer captures in-context information and balances it with pretrained bigram knowledge in next token prediction, using associative memory theory and experimental prompts to evaluate alignment with theoretical results.
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.