Rethinking Associative Memory Mechanism in Induction Head

• URL: http://arxiv.org/abs/2412.11459v2
• Date: Tue, 08 Jul 2025 13:14:01 GMT
• Title: Rethinking Associative Memory Mechanism in Induction Head
• Authors: Shuo Wang, Issei Sato,
• Abstract summary: This paper investigates how a two-layer transformer thoroughly captures in-context information and balances it with pretrained bigram knowledge in next token prediction.<n>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.
• Score: 37.93644115914534
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Provable In-Context Learning of Nonlinear Regression with Transformers [58.018629320233174]
  In-context learning (ICL) is the ability to perform unseen tasks using task-specific prompts without updating parameters.<n>Recent research has actively explored the training dynamics behind ICL.<n>This paper investigates more complex nonlinear regression tasks, aiming to uncover how transformers acquire in-context learning capabilities.
  arXiv  Detail & Related papers  (2025-07-28T00:09:28Z)
• Beyond Induction Heads: In-Context Meta Learning Induces Multi-Phase Circuit Emergence [28.260455480198047]
  Transformer-based language models exhibit In-Context Learning (ICL), where predictions are made adaptively based on context.<n>We experimentally clarify how such meta-learning ability is acquired by analyzing the dynamics of the model's circuit during training.
  arXiv  Detail & Related papers  (2025-05-22T13:59:30Z)
• CoT-ICL Lab: A Petri Dish for Studying Chain-of-Thought Learning from In-Context Demonstrations [11.907286102852957]
  CoT-ICL Lab is a framework and methodology to generate synthetic tokenized datasets. We systematically study chain-of-thought (CoT) in-context learning (ICL) in language models.
  arXiv  Detail & Related papers  (2025-02-21T01:24:54Z)
• Emergence of Abstractions: Concept Encoding and Decoding Mechanism for In-Context Learning in Transformers [18.077009146950473]
  Autoregressive transformers exhibit adaptive learning through in-context learning (ICL) We propose concept encoding-decoding mechanism to explain ICL by studying how transformers form and use internal abstractions in their representations. Our empirical insights shed light into better understanding the success and failure modes of large language models via their representations.
  arXiv  Detail & Related papers  (2024-12-16T19:00:18Z)
• Interpreting Affine Recurrence Learning in GPT-style Transformers [54.01174470722201]
  In-context learning allows GPT-style transformers to generalize during inference without modifying their weights. This paper focuses specifically on their ability to learn and predict affine recurrences as an ICL task. We analyze the model's internal operations using both empirical and theoretical approaches.
  arXiv  Detail & Related papers  (2024-10-22T21:30:01Z)
• How Transformers Get Rich: Approximation and Dynamics Analysis [11.789846138681359]
  We provide both approximation and dynamics analyses of how transformers implement induction heads.<n>In the em approximation analysis, we formalize both standard and generalized induction head mechanisms.<n>For the em dynamics analysis, we study the training dynamics on a synthetic mixed target, composed of a 4-gram and an in-context 2-gram component.
  arXiv  Detail & Related papers  (2024-10-15T10:22:27Z)
• In-Context Learning with Representations: Contextual Generalization of Trained Transformers [66.78052387054593]
  In-context learning (ICL) refers to a capability of pretrained large language models, which can learn a new task given a few examples during inference. This paper investigates the training dynamics of transformers by gradient descent through the lens of non-linear regression tasks.
  arXiv  Detail & Related papers  (2024-08-19T16:47:46Z)
• Parameter-Efficient and Memory-Efficient Tuning for Vision Transformer: A Disentangled Approach [87.8330887605381]
  We show how to adapt a pre-trained Vision Transformer to downstream recognition tasks with only a few learnable parameters. We synthesize a task-specific query with a learnable and lightweight module, which is independent of the pre-trained model. Our method achieves state-of-the-art performance under memory constraints, showcasing its applicability in real-world situations.
  arXiv  Detail & Related papers  (2024-07-09T15:45:04Z)
• How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? [82.51626700527837]
  Transformer-based large language models displayed impressive in-context learning capabilities, where a pre-trained model can handle new tasks without fine-tuning. We analyze how the mechanics of how Transformer to achieve ICL contribute to the technical challenges of the training problems in Transformers.
  arXiv  Detail & Related papers  (2024-02-23T21:07:20Z)
• Towards Understanding How Transformers Learn In-context Through a Representation Learning Lens [9.590540796223715]
  In this paper, we attempt to explore the in-context learning process in Transformers through a lens of representation learning. The ICL inference process of the attention layer aligns with the training procedure of its dual model, generating token representation predictions. We extend our theoretical conclusions to more complicated scenarios, including one Transformer layer and multiple attention layers.
  arXiv  Detail & Related papers  (2023-10-20T01:55:34Z)
• How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations [98.7450564309923]
  This paper takes initial steps on understanding in-context learning (ICL) in more complex scenarios, by studying learning with representations. We construct synthetic in-context learning problems with a compositional structure, where the label depends on the input through a possibly complex but fixed representation function. We show theoretically the existence of transformers that approximately implement such algorithms with mild depth and size.
  arXiv  Detail & Related papers  (2023-10-16T17:40:49Z)
• Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection [88.23337313766353]
  This work first provides a comprehensive statistical theory for transformers to perform ICL. We show that transformers can implement a broad class of standard machine learning algorithms in context. A emphsingle transformer can adaptively select different base ICL algorithms.
  arXiv  Detail & Related papers  (2023-06-07T17:59:31Z)
• Birth of a Transformer: A Memory Viewpoint [25.294093283819443]
  Large language models based on transformers have achieved great empirical successes. As they are deployed more widely, there is a growing need to better understand their internal mechanisms in order to make them more reliable. We study how transformers balance these two types of distributions of knowledge by considering a synthetic setup where tokens are generated from either global or context-specific bigrams.
  arXiv  Detail & Related papers  (2023-06-01T15:30:33Z)
• Dynamic Inference with Neural Interpreters [72.90231306252007]
  We present Neural Interpreters, an architecture that factorizes inference in a self-attention network as a system of modules. inputs to the model are routed through a sequence of functions in a way that is end-to-end learned. We show that Neural Interpreters perform on par with the vision transformer using fewer parameters, while being transferrable to a new task in a sample efficient manner.
  arXiv  Detail & Related papers  (2021-10-12T23:22:45Z)
• Multilinear Compressive Learning with Prior Knowledge [106.12874293597754]
  Multilinear Compressive Learning (MCL) framework combines Multilinear Compressive Sensing and Machine Learning into an end-to-end system. Key idea behind MCL is the assumption of the existence of a tensor subspace which can capture the essential features from the signal for the downstream learning task. In this paper, we propose a novel solution to address both of the aforementioned requirements, i.e., How to find those tensor subspaces in which the signals of interest are highly separable?
  arXiv  Detail & Related papers  (2020-02-17T19:06:05Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.