Rethinking Associative Memory Mechanism in Induction Head

Thank you very much for your helpful feedback and support for the paper!

Q.1 Oversight can mean either "unintentional failure" which is your meaning but also can mean "action of supervising something". Is there some other term available for “oversight of patterns”?

A.1 Thank you for the helpful comment. As you correctly understand our intention to the term “oversight”, we used it to describe that the model fails to look at important information in a long context, but readers can be confused if they interpret the word as supervision of pattern. We will replace oversight of pattern to the following:

(before) oversight of pattern

(after) neglect of in-context pattern

Q.2 Definition 1, Line 120: Is there a typo? “A matrix is called…” Do you mean W?

A.2 Thank you for pointing this out. Yes, we define a matrix W as an associative memory if W satisfies the condition written in Definition 1. We will revise the sentence to make the subject explicit and avoid confusion:

(before) A matrix is called associative memory when W is expressed as the sum of the products of ui and vj

(after) A matrix W is called an associative memory if it can be expressed as a weighted sum of outer products of u_i​ and v_j​

Q.3 Proposition 1: I don't understand the purpose of this for this paper?

A.3 In the context of length generalization in transformers, models without positional encodings (NoPE) are often discussed [1, 2]. For this reason, we considered the construction of an induction head under a NoPE setting.

Prior to Proposition 1, we observe that positional encodings have a significant effect on the induction head’s ability to generalize to longer sequences, as shown in the comparison between APE and RPE. However, it remained unclear whether an induction head mechanism could be implemented at all in a NoPE transformer. To address this, Proposition 1 provides a constructive proof that shows an induction head can, in principle, be realized even in a NoPE transformer. However, we note that this is purely a constructive result; we did not analyze whether such a mechanism can emerge through training and whether it generalizes to a long sequence. Thus, the proposition may have been confusing.

(references)

[1] Kazemnejad, Amirhossein, et al. "The impact of positional encoding on length generalization in transformers." Advances in Neural Information Processing Systems 36 (2023): 24892-24928.

[2] Wang, Jie, et al. "Length generalization of causal transformers without position encoding." arXiv preprint arXiv:2404.12224 (2024).

Thank you for the helpful feedback on our work. We write [FIX] for the revised version of our paper in the thread.

Q.1 Can the authors also include Nichani et al. 2024 and Friedman et al. 2023 they discussed in Section 2 in the table?

A.1 Thank you for pointing this out. We will revise the table to include additional rows for Nichani et al. (2024) and Friedman et al. (2023): [FIX1].

Q.2 The sentence in lines 59-60 is incomplete.

A.2 We will fix the grammatical errors and replace the sentence. Thank you!

Modern Hopfield networks (Krotov & Hopfield, 2016) are associative memory models that store patterns in an energy function and retrieve them via an updating rule.

Q.3 There’s a typo in the equation below line 97.

A.2 Thank you for pointing that out! The second term in Equation (1) is incorrect. This term corresponds to the residual connection, and it should be x_t^{(0,v)}​. We will fix this in the revised version.

R.1 The paper only considers a simple data model and one real dataset; it can be strengthened by including some (more) experiments on real data with larger models to support the key findings, and investigate how general they are.

C.1 We appreciate the suggestion to expand our evaluation to larger models and additional real-world datasets. Our goal in this work is to strike a careful balance between theory, experiment, and interpretability. In practice, pursuing generality often comes at the cost of transparency and analytical tractability. We believe that overly complex settings can obscure the underlying mechanisms we aim to understand. Our current setup with a simplified data model and a two-layer transforme is deliberately chosen to isolate and study the induction head mechanism in depth. This follows a similar spirit to prior work that employs even simpler architectures, such as linear transformers or other minimal designs, to reveal core principles.

We view our work as a first step toward understanding how positional encodings affect the formation and function of induction heads. We hope that the theoretical insights and targeted experiments provided here will guide future investigations into more complex settings.

R.2 Currently, the title and abstract don’t clearly reflect the main message of the paper. It seems to have two different focus points - length generalization and global vs in-context knowledge, but they seem to be treated as two separate things and are not connected well together.

C.2 We believe the current title does reflect the main message of the paper. Both length generalization and the balance between global and in-context knowledge are analyzed through the unified lens of associative memory and the induction head mechanism. The use of "Rethinking" in the title is meant to highlight our novel findings that go beyond earlier work—particularly regarding the impact of relative positional encoding (RPE), which was not previously examined, and the interplay between global and contextual knowledge in token prediction.

That said, we acknowledge that the current abstract may not clearly convey this framing. In particular, the role of associative memory, a central conceptual component of our analysis, is not sufficiently emphasized. Additionally, terms used in the abstract such as "in-context information over long context" and "how a transformer thoroughly captures in-context information" may obscure the relevance to length generalization, especially for readers who are scanning quickly. To address this, we propose the following revised version of the abstract, which makes the connection to associative memory and length generalization more explicit: [FIX2]

Thank you for the response.

Thanks for the clarifications and for addressing the concerns about the writing. I think the updated abstract reads better (except a typo in "how a transformer extract information from long contexts and then use it to", please fix) and clearly reflects the interesting contribution of the work. However, the response to the concern about limited empirical validation is not satisfactory, so I will maintain my score.

We thank the reviewer for the thoughtful review. We write [FIX] for the revised version of our paper in the thread.

R.1 Unclear whether / how the results generalize beyond bigrams, since both theoretical and empirical analyses depend on either text generated by a bigram model or word-pair tasks. More discussion of how the simplifying assumptions in the theoretical results related to actual transformer implementations and natural language data would be helpful.

C.1 Even in this simplified setup, we are able to demonstrate that absolute positional encodings can negatively affect the functionality of the induction head. This finding may help explain why many real-world transformer architectures favor relative positional encodings. It is common in theoretical work to use simplified architectures or datasets, or to impose strong assumptions in order to enable formal proofs. We believe this is a necessary and accepted trade-off. Importantly, if a phenomenon cannot be understood in a small or simplified model, it is unlikely to be interpretable in larger or more complex systems. In this sense, our paper aims to offer a first step toward understanding more general transformer behaviors by carefully analyzing a tractable setting.

We would also like to echo the reasoning articulated in Bietti et al. [1], who write:

“Our goal was to simplify the model as much as possible to ease the understanding of what is happening, while ensuring that the model is still rich enough to capture the relevant phenomena to solve the task. It would definitely be much more cumbersome to illustrate the memory viewpoint, and theoretically study training dynamics on a model where many more components are trained. We hope the simplicity of our architecture can help provide better intuition for what we believe to be a key internal mechanism in all transformer models.”

We share this philosophy. The simplicity of our setting is not a limitation, but a deliberate design choice to make induction heads and associative memory more interpretable, while still preserving the essential dynamics of transformer behavior.

[1] https://openreview.net/forum?id=3X2EbBLNsk#:~:text=Our%20goal%20was,all%20transformer%20models.

R.1 While it does seem independent and a new contribution, a lot of the paper is hard to follow if the reader is not intimately familiar with that earlier paper.

We admit that the earlier paper can enhance the understanding of our work, even though we do provide contexts and explanations of the notations and assumptions used in our paper. That said, we agree that including a brief summary of the earlier work—particularly Bietti et al.—can offer helpful context and background knowledge to readers, and make the motivation and contributions of our paper more transparent.

To address this, we plan to include the following short summary paragraph about Bietti et al. in the appendix: [FIX2]

Q.1 I encourage the authors to check all their references for canonical published versions.

A.1 Thank you for the helpful suggestion. We will revise the reference section to ensure that cited papers include their canonical published versions, with the corresponding conference or journal names where applicable. The following papers are included: [FIX1].

Thank you very much for your helpful feedback and support for the paper!

Q.1 What form do you consider for the absolute position encodings p_t? Are these learned d-dimensional embeddings with Gaussian entries? If not, do they depend on the value of t?

A.1 No, they are not learned. The position encodings p_t​ are fixed d-dimensional embeddings with Gaussian entries, so they do not depend on the value of t. This setting is what Bietti et al., (2024) used for their analysis. In our experiments, we use fixed embeddings to ensure near-orthogonality, which is necessary for the memory recall metric to function properly.

Q.2 in this line “if the input sequence contains out-of-distribution tokens, the induction head mechanism does not activate”, how are out-of-distribution tokens defined?

A.2 Out-of-distribution tokens are defined as tokens in the vocabulary that never appear in any of the training sequences. Since these tokens are never observed during training, the model does not learn their relationship with the relative position embedding r_{-1}. In other words, the associative memory in the first attention layer does not store any information about the pair consisting of such a token and the previous position embedding.

Q.3 In lemma 1, Q in W_Q and Q in \pi_Q are not the same right?

A.3 Yes, they are different. W_Q refers to the query matrix in the transformer, which is used to compute query vectors in the attention mechanism. On the other hand, \pi_Q​ refers to the bigram conditional distribution used to generate trigger tokens in our data generation process. While we used the notation \pi_Q​ to be consistent with Bietti et al. (2024), we acknowledge that this may cause confusion due to the overload of the symbol "Q." We plan to revise the notation in the revised version to avoid ambiguity.

Q.4 In equation 8, what does the prime symbol over the logit symbols denote?

A.4 The prime symbol denotes the logits before applying the softmax function, as stated in line 249. More specifically, Equation (4) defines the logits \xi = \sigma(W_U x_t), where \sigma is the softmax function. We define \xi’ as W_U x_t​, i.e., the raw logits before softmax is applied. By comparing the entries of this vector, we can analyze how the logit difference between two vocabulary items (e.g., v_1​ and v_2 in Proposition 2​) arises.

Q.5 In equation 8, shouldn’t the difference be between the logarithm of the logits to be equal to the value in the right-hand side?

A.5 Yes, the value on the right-hand side should indeed correspond to the difference between the logarithms of the softmax probabilities, which in turn equals the difference between the unnormalized logits before softmax. This is what is explained in Q.4 and A.4.

Q.6 suggest the authors define induction heads at least informally before talking about them in the introduction.

A.6 Thank you for the suggestion. We agree that introducing the concept of induction heads earlier in the paper would improve clarity for readers unfamiliar with the term. We will revise the introduction to include a brief explanation of induction heads:

(before) It is known that the induction head mechanism can emerge in two-layer transformers 26 (Elhage et al., 2021), and many theoretical studies have focused on analyzing two-layer 27 architectures to understand how induction head develops (Edelman et al., 2024; Nichani 28 et al., 2024; Chen et al., 2024b)

(after) It is known that the induction head mechanism, which is a pattern of attention in a transformer that enables the model to copy and reuse information from earlier in the context, can emerge in two-layer transformers (Elhage et al., 2021), and many theoretical studies have focused on analyzing two-layer architectures to understand how induction head develops (Edelman et al., 2024; Nichani et al., 2024; Chen et al., 2024b)

We appreciate the reviewer’s thorough review. The [FIX]s are for the revised version of our paper and displayed in another comment.

Q.1 shouldn’t the trigger token-output token pairs mostly be specified by the context?

A.1 Yes, we sample the trigger token and the output token for each sequence individually. While lines 142–144 mention how these tokens are sampled, what we mean is that the trigger token is sampled from the conditional distribution \pi_q​ for each sequence. Thus, as you correctly point out, the trigger token–output token pairs should indeed be specified by the context. We agree that lines 142–144 may be unclear, so we propose revising them: [FIX1]. Also, we add the following section in the Appendix: [FIX2].

Q.2 quite confused about the mention of the latter half near line 165 and the comparison between sequence position <128 and >=128. Is there a sharp transition at half the sequence length?

A.2 The transition does not occur exactly at position 128. In Fig. 3 of Appendix H, we visualize the attention pattern in the first layer of a trained transformer. The figure illustrates that the model loses its ability to attend to the previous token around the midpoint of the sequence. This observation motivated our use of the term "latter half" to describe the general region where this degradation becomes apparent. Thus, in the experiment shown in Fig. 1(a), we divide the sequence into two halves at position 128 and measure the memory recall associated with the previous token in each half.

Q.3 Comparison with NOPE: I think a comparison with no positional encoding (2) might be useful.

A.3 We conducted an experiment to see whether 3-layer transformer with no positional encoding learns previous token head. We examined the 2nd layer attention and the result is in the figure [1]. It attends to previous token for the first appearance of trigger-output tokens but does not learn to always attend to previous token.

[1] https://www.anonfile.la/ad77c6

R.1 Heavily relying on theory of simple architectures: Theorems proven on a two-layer transformer may not generalize to larger or more realistic models.

C.1 In our opinion, using a two-layer transformer allowed us to analyze the induction head mechanism with greater precision and clarity. We believe that analyzing such small, interpretable models is a necessary step toward understanding the behavior of larger, more complex ones. In fact, we were able to provide a tractable yet meaningful explanation of how relative positional encoding facilitates the implementation of induction heads. This, in turn, offers a theoretical justification for the widespread use of RPE in practice.

R.2 Lack of surprisingness: It isn’t particularly surprising that RPE yields more stable induction heads

C.2 We believe the surprising aspect lies in the fact that these behaviors can be theoretically explained in terms of associative memory, while we agree that some of the observed phenomena, such as the superiority of RPE over APE, may be expected or well-known empirically. The core contribution of our work is to formally prove that RPE is better suited for implementing induction heads. We believe this shift from empirical intuition to theoretical understanding is highly meaningful for the community, as it not only justifies a common design choice but also deepens our insight into how transformers operate.

R.3 Setting description could be improved: The paper could more clearly specify exactly how the synthetic training data are constructed and sampled.

C.3 we addressed this concern in A.1. We fully agree that improving readability is important to help readers follow the logic more easily.

R.4 Presentation of results: Figure 1a in its current form is difficult to interpret, what the main findings are.

C.4 The main finding illustrated in Fig.1a is that our theoretical insights about the behavior of RPE and APE are empirically validated. This figure plots how the memory recall metric, defined in lines 297–298, evolves throughout training. We will add the paragraph as a supplementary explanation: [FIX3].

R.5 Practical relevance: It is unclear how relevant the study of positional encoding’s effect on ICL is when in practice, APE is almost never used compared to RPE.

C.5 Our study helps explain why APE is rarely used in practice and RPE is the default choice in many real-world applications. Our work addresses this gap by providing a theoretical explanation for why APE performs poorly in long-context settings. We believe this insight is practically relevant because it offers interpretability for a widely observed empirical trend. Shedding light on why certain design choices (like RPE) lead to better in-context learning not only justifies current best practices. In that sense, our work contributes to the broader goal of increasing transparency and understanding of in-context learning, a core capability of modern language models.