Do LLMs dream of elephants (when told not to)? Latent concept association and associative memory in transformers

Authors: yibo  jiang
goutham  rajendran
pradeep  ravikumar
bryon  aragam
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 01:58:00

Abstract: Large Language Models (LLMs) have the capacity to store and recall facts. Through experimentation with open-source models, we observe that this ability to retrieve facts can be easily manipulated by changing contexts, even without altering their factual meanings. These findings highlight that LLMs might behave like an associative memory model where certain tokens in the contexts serve as clues to retrieving facts. We mathematically explore this property by studying how transformers, the building blocks of LLMs, can complete such memory tasks. We study a simple latent concept association problem with a one-layer transformer and we show theoretically and empirically that the transformer gathers information using self-attention and uses the value matrix for associative memory.

Main paper: 
 introduction: 
 What is the first thing that would come to mind if you were asked not to think of an elephant? Chances are, you would be thinking about elephants. What if we ask the same thing to Large Language Models (LLMs)? Obviously, one would expect the outputs of LLMs to be heavily influenced by tokens in the context [Bro+20] . Could such influence potentially prime LLMs into changing outputs in a nontrivial way? To gain a deeper understanding, we focus on one specific task called fact retrieval [Men+22; Men+23] where expected output answers are given. LLMs, which are trained on vast amounts of data, are known to have the capability to store and recall facts [Men+22; Men+23; DCAT21; Mit+21; Mit+22; Dai+21]. This ability raises natural questions: How robust is fact retrieval, and to what extent does it depend on semantic meanings within contexts? What does it reveal about memory in LLMs?In this paper, we first demonstrate that fact retrieval is not robust and LLMs can be easily fooled by varying contexts. For example, when asked to complete "The Eiffel Tower is in the city of", GPT-2 [Rad+19] answers with "Paris". However, when prompted with "The Eiffel Tower is not in Chicago. The Eiffel Tower is in the city of", GPT-2 responds with "Chicago". See Figure 1 for more examples, including Gemma and LLaMA. On the other hand, humans do not find the two sentences factually confusing and would answer "Paris" in both cases. We call this phenomenon context hijacking. Importantly, these findings suggest that LLMs might behave like an associative memory model. Specifically, we refer to an associative memory model in which LLMs rely on certain tokens in contexts to guide the retrieval of memories, even if such associations formed are not inherently semantically meaningful. This contrasts with the ideal behavior, where LLMs would generalize by understanding new contexts, reasoning through them, and integrating prior knowledge. This associative memory perspective raises further interpretability questions about how LLMs form such associations. Answering these questions can facilitate the development of more robust LLMs. Unlike classical models of associative memory in which distance between memory patterns are measured directly and the associations between inputs and outputs are well-specified, fact retrieval relies on a more nuanced notion of similarity measured by latent (unobserved) semantic concepts.To model this, we propose a synthetic task called latent concept association where the output token is closely related to sampled tokens in the context but wherein similarity is measured via a latent space of semantic concepts. We then investigate how a one-layer transformer [Vas+17] , a fundamental component of LLMs, can tackle this memory retrieval task in which various context distributions correspond to distinct memory patterns. We demonstrate that the transformer accomplishes the task in two stages: The self-attention layer gathers information, while the value matrix functions as associative memory. Moreover, low-rank structure also emerges in the embedding space of trained transformers. These findings provide additional theoretical validation for numerous existing low-rank editing and fine-tuning techniques [Men+22; Hu+21] .Contributions Specifically, we make the following contributions:1. We systematically demonstrate context hijacking for various open source LLM models including GPT-2 [Rad+19] , LLaMA-2 [Tou+23] and Gemma [Tea+24] , which show that fact retrieval can be misled by contexts (Section 3), reaffirming that LLMs lack robustness to context changes [Shi+23; Pet+20; CSH22; Yor+23; PE21].2. We propose a synthetic memory retrieval task termed latent concept association, allowing us to analyze how transformers can accomplish memory recall (Section 4). Unlike classical models of associative memory, our task creates associations in a latent, semantic concept space as opposed to directly between observed tokens. This perspective is crucial to understanding how transformers can solve fact retrieval problems by implementing associative memory based on similarity in the latent space.3. We theoretically (Section 5) and empirically (Section 6) study trained transformers on this latent concept association problem, showing that self-attention is used to aggregate information while the value matrix serves as associative memory. And moreover, we discover that the embedding space can exhibit a low-rank structure, offering additional support for existing editing and fine-tuning methods [Men+22; Hu+21] . 
literature review: 
 Associative memory Associative memory has been explored within the field of neuroscience [Hop82; Seu96; BYBOS95; Ska+94; SS22]. The most popular models among them is the Hopfield network [Hop82] and its modern successors [Ram+20; Mil+22; Zha23; Hu+24d; Wu+23; Hu+24b; Hu+24c; Wu+24a; Hu+24a] are closely related to the attention layer used in transformers [Vas+17] .In addition, the attention mechanism has also been shown to approximate another associative memory model known as sparse distributed memory [BP21] . Beyond attention, Radhakrishnan et al. [RBU20] and Jiang and Pehlevan [JP20] show that overparameterzed autoencoders can implement associative memory as well. This paper studies fact retrieval as a form of associative memory. Another closely related area of research focuses on memorization in deep neural networks. Henighan et al. [Hen+23] shows that a simple neural network trained on toy model will store data points in the overfitting regime while storing features in the underfitting regime. Feldman [Fel20] and Feldman and Zhang [FZ20] study the interplay between memorization and long tail distributions while [Shi+23] , question answering [Pet+20; CSH22; Yor+23] and factual reasoning [PE21] . Although phenomena akin to context hijacking have been reported in different instances, the goals of this paper are to give a systematic robustness study for fact retrieval, offer a framework for interpreting it in the context of associative memory, and deepen our understanding of LLMs. 
context hijacking in llms: 
 In this section, we run experiments on LLMs including GPT-2 [Rad+19] , Gemma [Tea+24] (both base and instruct models) and LLaMA-2-7B [Tou+23] to explore the effects of context hijacking on manipulating LLM outputs. As an example, consider Figure 1 . When we prompt the LLMs with the context "The Eiffel Tower is in the city of", all 4 LLMs output the correct answer ("Paris"). However, as we see in the example, we can actually manipulate the output of the LLMs simply by modifying the context with additional factual information that would not confuse a human. We call this context-hijacking. Due to the different capacities and capabilties of each model, the examples in Figure 1 use different hijacking techniques. This is most notable on LLaMA-2-7B, which is a much larger model than the others. Of course, as expected, the more sophisticated attack on LLaMA also works on GPT-2 and Gemma. Additionally, the instruction-tuned version of Gemma can understand special words like "not" to some extent. Nevertheless, it is still possible to systematically hijack these LLMs, as demonstrated below.We explore this phenomenon at scale with the COUNTERFACT dataset introduced in [Men+22], a dataset of difficult counterfactual assertions containing a diverse set of subjects, relations, and linguistic variations. COUNTERFACT has 21, 919 samples, each of which are given by a tuple (p, o * , o _ , s, r).From each sample, we have a context prompt p with a true target answer o * (target_true) and a false target answer o _ (target_false), e.g. the prompt p = "Eiffel Tower can be found in" has true target o * = "Paris" and false target o _ = "Guam". Additionally, the main entity in p is the subject s (s = "Eiffel Tower") and the prompt is categorized into relations r (for instance, other samples with the same relation ID as the example above could be of the form "The location of {subject} is", "{subject} can be found in", "Where is {subject}? It is in"). For additional details on how the dataset was collected, see [Men+22] .For a hijacking scheme, we report the Efficacy Score (ES) [Men+22] , which is the proportion of samples for which the token probabilities satisfy P r[o _ ] > P r[o * ] after modifying the context, that is, the proportion of the dataset that has been successfully manipulated. We experiment with two hijacking schemes for this dataset. We first hijack by prepending the text "Do not think of {target_false}" to each context. For instance, the prompt "The Eiffel Tower is in" gets changed to "Do not think of Guam. The Eiffel Tower is in". In Figure 2a , we see that the efficacy score rises significantly after hijacking. Here, we prepend the hijacking sentence k times for k = 0, . . . , 5 where k = 0 yields the original prompt. We see that additional prepends increase the score further.In the second scheme, we make use of the relation ID r to prepend factually correct sentences. For instance, one can hijack the example above to "The Eiffel Tower is not located in Guam. The Eiffel Tower is in". We test this hijacking philosophy on different relation IDs. In particular, Figure 2b reports hijacking based on relation ID P 190 ("twin city"). And we see similar patterns that with more prepends, the ES score gets higher. It is also worth noting that one can even hijack by only including words that are semantically close to the false target (e.g., "France" for false target "French"). This suggests that context hijacking is more than simply the LLM copying tokens from contexts. Context hijacking indicates that fact retrieval in LLMs is not robust and that accurate fact recall does not necessarily depend on the semantics of the context. As a result, one hypothesis is to view LLMs as an associative memory model where special tokens in contexts, associated with the fact, provide partial information or clues to facilitate memory retrieval [Zha23] . To better understand this perspective, we design a synthetic memory retrieval task to evaluate how the building blocks of LLMs, transformers, can solve it. 
problem setup: 
 In the context of LLMs, fact or memory retrieval, can be modeled as a next token prediction problem. Given a context (e.g., "The capital of France is"), the objective is to accurately predict the next token (e.g., "Paris") based on the factual relation between context and the following token.Previous papers [Ram+20; Mil+22; BP21; Zha23] have studied the connection between attention and autoassociative and heteroassociative memory. For autoassociative memory, contexts are modeled as a set of existing memories and the goal of self-attention is to select the closest one or approximations to it. On top of this, heteroassociative memory [Mil+22; BP21] has an additional projection to remap each output to a different one, whether within the same space or otherwise. In both scenarios, the goal is to locate the closest pattern within the context when provided with a query (up to a remapping if it's heteroassociative).Fact retrieval, on the other hand, does not strictly follow this framework. The crux of the issue is that the output token is not necessarily close to any particular token in the context but rather a combination of them and the "closeness" is intuitively measured by latent semantic concepts. For example, consider context sentence "The capital of France is" with the output "Paris". Here, none of the tokens in the context directly corresponds to the word "Paris". Yet some tokens contain partial information about "Paris". Intuitively, "capital" aligns with the "isCapital" concept of "Paris", while "France" corresponds to the "isFrench" concept linked to "Paris" where all the concepts are latent. To model such phenomenon, we propose a synthetic task called latent concept association where the output token is closely related to tokens in the context and similarity is measured via the latent space.We propose a synthetic prediction task where for each output token y, tokens in the context (denoted by x) are sampled from a conditional distribution given y. Tokens that are similar to y will be favored to appear more in the context, except for y itself. The task of latent concept association is to successfully retrieve the token y given samples from p(x|y). The synthetic setup simplifies by not accounting for the sequential nature of language, a choice supported by previous experiments on context hijacking (Section 3). We formalize this task below.To measure similarity, we define a latent space. Here, the latent space is a collection of m binary latent variables Z i . These could be viewed as semantic concept variables. Let Z = (Z 1 , ..., Z m ) be the corresponding random vector, z be its realization, and Z be the collection of all latent binary vectors. For each latent vector z, there's one associated token t ∈ [V ] = {0, ..., V -1} where V is the total number of tokens. Here we represent the tokenizer as ι where ι(z) = t. In this paper, we assume that ι is the standard tokenizer where each binary vector is mapped to its decimal number. In other words, there's a one to one map between latent vectors and tokens. Because the map is one to one, we sometimes use latent vectors and tokens interchangeably. We also assume that every latent binary vector has a unique corresponding token, therefore V = 2 m .Under the latent concept association model, the goal is to retrieve specific output tokens given partial information in the contexts. This is modeled by the latent conditional distribution:p(z|z * ) = ωπ(z|z * ) + (1 -ω)Unif(Z)whereπ(z|z * ) ∝ exp(-D H (z, z * )/β) z ∈ N (z * ), 0 z / ∈ N (z * ).Here D H is the Hamming distance, N (z * ) is a subset of Z\{z * } and β > 0 is the temperature parameter. The use of Hamming distance draws a parallel with the notion of distributional semantics in natural language: "a word is characterized by the company it keeps" [Fir57] . In words, p(z|z * ) says that with probability 1 -ω, the conditional distribution uniformly generate random latent vectors and with probability ω, the latent vector is generated from the informative conditional distribution π(z|z * ) where the support of the conditional distribution is N (z * ). Here, π represents the informative conditional distribution that depends on z * whereas the uniform distribution is uninformative and can be considered as noise. The mixture model parameter ω determines the signal to noise ratio of the contexts.Therefore, for any latent vector z * and its associated token, one can generate L context token words with the aforementioned latent conditional distribution:• Uniformly sample a latent vector z *• For l = 1, ..., L -1, sample z l ∼ p(z|z * ) and t l = ι(z l ).• For l = L, sample z ∼ π(z|z * ) and t L = ι(z).Consequently, we have x = (t 1 , .., t L ) and y = ι(z * ). The last token in the context is generated specifically to make sure that it is not from the uniform distribution. This ensures that the last token can use attention to look for clues, relevant to the output, in the context. Let D L be the sampling distribution to generate (x, y) pairs. The conditional probability of y given x is given by p(y|x).With slight abuse of notation, given a token t ∈ [V ], we defineN (t) = N (ι -1 (t)). we also define D H (t, t ′ ) = D H (ι -1 (t), ι -1 (t ′ ))for any pair of tokens t and t ′ .For any function f that maps the context to estimated logits of output labels, the training objective is to minimize this loss of the last position:E (x,y)∈D L [ℓ(f (x), y)]where ℓ is the cross entropy loss with softmax. The error rate of latent concept association is defined by the following:R D L (f ) = P (x,y)∼D L [argmax f (x) ̸ = y]And the accuracy is 1 -R D L (f ).Given a context x = (t 1 , .., t L ) which consists of L tokens, we define X ∈ {0, 1} V ×L to be its one-hot encoding where V is the vocabulary size. Here we use χ to represent the one-hot encoding function (i.e., χ(x) = X). Similar to [LLR23; Tar+23a; Li+24], we also consider a simplified one-layer transformer model without residual connections and normalization:f L (x) = W E T W V attn(W E χ(x)) :L (4.1)whereattn(U ) = U σ (W K U ) T (W Q U ) √ d a ,W K ∈ R da×d is the key matrix, and W Q ∈ R da×d is the query matrix and d a is the attention head size. σ : R L×L → (0, 1) L×L is the column-wise softmax operation. W V ∈ R d×d is the value matrix and W E ∈ R d×V is the embedding matrix. Here, we adopt the weight tie-in implementation which is used for Gemma [Tea+24] . We focus solely on the prediction of the last position, as it is the only one relevant for latent concept association. For convenience, we also use h(x) to mean attn(W E χ(x)) :L , which is the hidden representation after attention for the last position, and f L t (x) to represent the logit for output token t. 
theoretical analysis: 
 In this section, we theoretically investigate how a single-layer transformer can solve the latent concept association problem. We first introduce a hypothetical associative memory model that utilizes self-attention for information aggregation and employs the value matrix for memory retrieval. This hypothetical model turns out to mirror trained transformers in experiments. We also examine the role of each individual component of the network: the value matrix, embeddings, and the attention mechanism. We validate our theoretical claims in Section 6.In this section, we show that a simple single-layer transformer network can solve the latent concept association problem. The formal result is presented below in Theorem 1; first we require a few more definitions. Let W E (t) be the t-th column of the embedding matrix W E . In other words, this is the embedding for token t. Given a token t, define N 1 (t) to be the subset of tokens whose latent vectors are only 1 Hamming distance away from t's latent vector:N 1 (t) = {t ′ : D H (t ′ , t)) = 1} ∩ N (t).For any output token t, N 1 (t) contains tokens with the highest probabilities to appear in the context.The following theorem formalizes the intuition that a one-layer transformer that uses self-attention to summarize statistics about the context distributions and whose value matrix uses aggregated representations to retrieve output tokens can solve the latent concept association problem defined in Section 4.1.Theorem 1 (informal). Suppose the data generating process follows Section 4.1 where m ≥ 3, ω = 1, and N (t) = V \ {t}. Then for any ε > 0, there exists a transformer model given by (4.1) that achieves error ε, i.e. R D L (f L ) < ε given sufficiently large context length L.More precisely, for the transformer in Theorem 1, we will have W K = 0 and W Q = 0. Each row of W E is orthogonal to each other and normalized. And W V is given byW V = t∈[V ] W E (t)( t ′ ∈N1(t) W E (t ′ ) T ) (5.1)A more formal statement of the theorem and its proof is given in Appendix B (Theorem 7).Intuitively, Theorem 1 suggests having more samples from p(x|y) can lead to a better recall rate. On the other hand, if contexts are modified to contain more samples from p(x|ỹ) where ỹ ̸ = y, then it is likely for transformer to output the wrong token. This is similar to context hijacking (see Section 5.5).The construction of the value matrix is similar to the associative memory model used in [Bie+24; CSB24] , but in our case, there is no explicit one-to-one input and output pairs stored as memories.Rather, a combination of inputs are mapped to a single output.While the construction in Theorem 1 is just one way that a single-layer transformer can tackle this task, it turns out empirically this construction of W V is close to the trained W V , even in the noisy case (ω ̸ = 1). In Section 6.1, we will demonstrate that substituting trained value matrices with constructed ones can retain accuracy, and the constructed and trained value matrices even share close low-rank approximations. Moreover, in this hypothetical model, a simple uniform attention mechanism is deployed to allow self-attention to count occurrences of each individual tokens. Since the embeddings are orthonormal vectors, there is no interference. Hence, the self-attention layer can be viewed as aggregating information of contexts. It is worth noting that, in different settings, more sophisticated embedding structures and attention patterns are needed. This is discussed in the following sections.The construction in Theorem 1 relies on the value matrix acting as associative memory. But is it necessary? Could we integrate the functionality of the value matrix into the self-attention module to solve the latent concept association problem? Empirically, the answer seems to be negative as will be shown in Section 6.1. In particular, when the context length is small, setting the value matrix to be the identity would lead to subpar memory recall accuracy. This is because if the value matrix is the identity, the transformer would be more susceptible to the noise in the context. To see this, notice that given any pair of context and output token (x, y), the latent representation after self-attention h(x) must live in the polyhedron S y to be classified correctly where S y is defined as:S y = {v : (W E (y) -W E (t)) T v > 0 where t ̸ ∈ [V ]\ {y}} Note that, by definition, for any two tokens y and ỹ, S y ∩ S ỹ = ∅. On the other hand, because of the self-attention mechanism, h(x) must also live in the convex hull of all the embedding vectors:CV = Conv(W E (0), ..., W E (|V | -1))In other words, for any pair (x, y) to be classified correctly, h(x) must live in the intersection of S y and CV . Due to the stochastic nature of x, it is likely for h(x) to be outside of this intersection. The remapping effect of the value matrix can help with this problem. The following lemma explains this intuition.Lemma 2. Suppose the data generating process follows Section 4.1 where m ≥ 3, ω = 1 andN (t) = {t ′ : D H (t, t ′ )) = 1}.For any single layer transformer given by (4.1) where each row of W E is orthogonal to each other and normalized, if W V is constructed as in (5.1), then the error rate is 0. If W V is the identity matrix, then the error rate is strictly larger than 0.Another intriguing phenomenon occurs when the value matrix is the identity matrix. In this case, the inner product between embeddings and their corresponding Hamming distance varies linearly. This relationship can be formalized by the following theorem.Theorem 3. Suppose the data generating process follows Section 4.1 where m ≥ 3, ω = 1 and N (t) = V \ {t}. For any single layer transformer given by (4.1) with W V being the identity matrix, if the cross entropy loss is minimized so that for any sampled pair (x, y), p(y|x) = p(y|x) = softmax(f L y (x)) there exists a > 0 and b such that for two tokens t̸ = t ′ , ⟨W E (t), W E (t ′ )⟩ = -aD H (t, t ′ ) + bThe hypothetical model in Section 5.1 requires embeddings to form an orthonormal basis. In the overparameterization regime where the embedding dimension d is larger than the number of tokens V , this can be approximately achieved by Gaussian initialization. However, in practice, the embedding dimension is typically smaller than the vocabulary size, in which case it is impossible for the embeddings to constitute such a basis. Empirically, in Section 6.2, we observe that with overparameterization (d > V ), embeddings can be frozen at their Gaussian initialization, whereas in the underparameterized regime, embedding training is required to achieve better recall accuracy.This raises the question: What kind of embedding geometry is learned in the underparameterized regime? Experiments reveal a close relationship between the inner product of embeddings for two tokens and the Hamming distance of these tokens (see Figure 3b Approximately, we have the following relationship:⟨W E (t), W E (t ′ )⟩ = b 0 t = t ′ -aD H (t, t ′ ) + b t ̸ = t ′ (5.2)for any two tokens t and t ′ where b 0 > b and a > 0. One can view this as a combination of the embedding geometry under Gaussian initialization and the geometry when W V is the identity matrix (Theorem 3). Importantly, this structure demonstrates that trained embeddings inherently capture similarity within the latent space. Theoretically, this embedding structure (5.2) can also lead to low error rate under specific conditions on b 0 , b and a, which is articulated by the following theorem.Theorem 4 (Informal). Following the same setup as in Theorem 1, but embeddings obey (5.2), then under certain conditions on a, b and if b 0 and context length L are sufficiently large, the error rate can be arbitrarily small, i.e. R D L (f L ) < ε for any 0 < ε < 1.The formal statement of the theorem and its proof is given in Appendix B (Theorem 8).Notably, this embedding geometry also implies a low-rank structure. Let's first consider the special case when b 0 = b. In other words, the inner product between embeddings and their corresponding Hamming distance varies linearly.Lemma 5. If embeddings follow (5.2) and b = b 0 andN (t) = V \ {t}, then rank(W E ) ≤ m + 2.When b 0 > b, the embedding matrix will not be strictly low rank. However, it can still exhibit approximate low-rank behavior, characterized by an eigengap between the top and bottom singular values. This is verified empirically (see Figure D .9-D.12 in Appendix D.4).As of now, attention does not play a significant role in the analysis. But perhaps unsurprisingly, the attention mechanism is useful in selecting relevant information. To see this, let's consider a specific setting where for any latent vector z * , N (z * ) = {z : z * 1 = z 1 } \ {z * }. Essentially, latent vectors are partitioned into two clusters based on the value of the first latent variable, and the informative conditional distribution π only samples latent vectors that are in the same cluster as the output latent vector. Empirically, when trained under this setting, the attention mechanism will pay more attention to tokens within the same cluster (Section 6.3). This implies that the self-attention layer can mitigate noise and concentrate on the informative conditional distribution π.To understand this more intuitively, we will study the gradient of unnormalized attention scores. In particular, the unnormalized attention score is defined as:u t,t ′ = (W K W E (t)) T (W Q W E (t ′ ))/ d a .Lemma 6. Suppose the data generating process follows Section 4.1 and N (z * ) = {z : z * 1 = z 1 } \ {z * }. Given the last token in the sequence t L , then∇ ut,t L ℓ(f L ) = ∇ℓ(f L ) T (W E ) T W V (α t pt W E (t) -pt L l=1 pt l W E (t l ))where for token t, α t = L l=1 1[t l = t] and pt is the normalized attention score for token t. Typically, α t is larger when token t and t L belong to the same cluster because tokens within the same cluster tend to co-occur frequently. As a result, the gradient contribution to the unnormalized attention score is usually larger for tokens within the same cluster.In light of the theoretical results on latent concept association, a natural question arises: How do these results connect to context hijacking in LLMs? In essence, for the latent concept association problem, the differentiation of output tokens is achieved by distinguishing between the various conditional distributions p(x|y). Thus, adding or changing tokens in the context x so that it resembles a different conditional distribution can result in misclassification. In Appendix D.5, we present experiments showing that mixing different contexts can cause transformers to misclassify. This partially explains context hijacking in LLMs (Section 3). On the other hand, it is well-known that the error rate is related to the KL divergence between conditional distributions of contexts [Cov99] . The closer the distributions are, the easier it is for the model to misclassify. Here, longer contexts, primarily composed of i.i.d samples, suggest larger divergences, thus higher memory recall rate. This is theoretically implied by Theorem 1 and Theorem 4 and empirically verified in Appendix D.6. Such result is also related to reverse context hijacking (Appendix C) where prepending sentences including true target words can improve fact recall rate. 
experiments: 
 The main implications of the theoretical results in the previous section are:1. The value matrix is important and has associative memory structure as in (5.1).Training embeddings is crucial in the underparameterized regime, where embeddings exhibit certain geometric structures.3. Attention mechanism is used to select the most relevant tokens.To evaluate these claims, we conduct several experiments on synthetic datasets. Additional experimental details and results can be found in Appendix D.In this section, we study the necessity of the value matrix W V and its structure. First, we conduct experiments to compare the effects of training versus freezing W V as the identity matrix, with the context lengths L set to 64 and 128. Figure 3a and Figure D .1 show that when the context length is small, freezing W V can lead to a significant decline in accuracy. This is inline with Lemma 2 and validates it in a general setting, implying the significance of the value matrix in maintaining a high memory recall rate.Next, we investigate the degree of alignment between the trained value matrix W V and the construction in (5.1). The first set of experiments examines the similarity in functionality between the two matrices. We replace value matrices in trained transformers with the constructed ones like in (5.1) and then report accuracy with the new value matrix. As a baseline, we also consider randomly constructed value matrix, where the outer product pairs are chosen randomly (detailed construction can be found in Appendix D.1). Figure D.2 indicates that the accuracy does not significantly decrease when the value matrix is replaced with the constructed ones. Furthermore, not only are the constructed value matrix and the trained value matrix functionally alike, but they also share similar low-rank approximations. We use singular value decomposition to get the best low rank approximations of various value matrices where the rank is set to be the same as the number of latent variables (m). We then compute smallest principal angles between low-rank approximations of trained value matrices and those of constructed, randomly constructed, and Gaussian-initialized value matrices. Figure D .3 shows that the constructed ones have, on average, smallest principal angles with the trained ones.In this section, we explore the significance of embedding training in the underparamerized regime and embedding structures. We conduct experiments to compare the effects of training versus freezing embeddings with different embedding dimensions. The learning rate is selected as the best option from {0.01, 0.001} depending on the dimensions. As suggested in Section 5.3, such embedding geometry (5.2) can lead to low rank structures. We verify this claim by studying the spectrum of the embedding matrix W E . As illustrated in Appendix D.4, Figure D .9-D.12 demonstrate that there are eigengaps between top and bottom singular values, suggesting low-rank structures.In this section, we examine the role of attention pattern by considering a special class of latent concept association model as defined in Section 5.4. Figure 3c Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results? Answer: [Yes] Justification: Experimental details are included in the paper. Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
conclusions: 
 In this work, we first presented the phenomenon of context hijacking in LLMs, which suggested that fact retrieval is not robust against variations of contexts. This indicates that LLMs might function like associative memory where tokens in contexts are clues to guide memory retrieval. To investigate this perspective further, we devised a synthetic task called latent concept association and examined theoretically and empirically how single-layer transformers are trained to solve this task. These results provide further insights into the inner workings of transformers and LLMs, and can hopefully stimulate further work into interpreting and understanding the mechanisms by which LLMs predict tokens and recall facts.Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments? Answer: [Yes] Justification: Standard errors are provided for the experiments. Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [NA]Justification: The paper does not use existing assets.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: The context hijacking experiments were only conducted on open-source models and not on commercial models like GPT-4. Nevertheless, even in the official GPT-4 technical report [Ach+23], there is an example similar to context hijacking (the Elvis Perkins example). In that example, the prompt is "Son of an actor, this American guitarist and rock singer released many songs and albums and toured with his band. His name is "Elvis" what?". GPT-4 answers with Presley, even though the answer is Perkins (Elvis Presley is not the son of an actor). GPT-4 can be viewed as distracted by all the information related to music and answers Presley. In fact, it is known that LLMs can be easily distracted by contexts in use cases other than fact retrieval such as problem-solving [Shi+23] . So we reasonably suspect that similar behavior still exists in larger models but is harder to exploit. On the other hand, the theoretical section only focuses on single-layer transformer network. While single-layer networks already demonstrate some interesting phenomena including low-rank structures, the functionality of multi-layer transformers is much different compared to single-layer transformers with the notable emergence of induction head [Elh+21] .B.1 Proofs for Section 5.1Theorem 1 can be stated more formally as follows:Theorem 7. Suppose the data generating process follows Section 4.1 where m ≥ 3, ω = 1, and N (t) = V \ {t}. Assume there exists a single layer transformer given by (4.1) such that a) W K = 0 and W Q = 0, b) Each row of W E is orthogonal to each other and normalized, and c) W V is given byW V = i∈[V ] W E (i)( j∈N1(i) W E (j) T ). Then if L > max{ 100m 2 log(3/ε) (exp(-1 β )-exp(-2 β )) 2 , 80m 2 |N (y)| (exp(-1 β )-exp(-2 β )) 2 } for any y, then R D L (f L ) ≤ ε,where 0 < ε < 1.Proof. First of all, the error is defined to be:R D L (f L ) = P (x,y)∼D L [argmax f L (x) ̸ = y] = P y P x|y [argmax f L (x) ̸ = y]Let's focus on the conditional probabilityP x|y [argmax f L (x) ̸ = y].By construction, the single layer transformer model has uniform attention. Therefore,h(x) = i∈N (y) α i W E (i) where α i = 1 L L k=1 1{t k = i}which is the number of occurrence of token i in the sequence. By the latent concept association model, we know that p(i|y) = exp(-D H (i, y)/β) Z where Z = i∈N (y) exp(-D H (i, y)/β).Thus, the logit for token y isf L y (x) = i∈N1(y) α iAnd the logit for any other token ỹ isf L ỹ (x) = i∈N1(ỹ) α iFor the prediction to be correct, we needmax ỹ f L y (x) -f L ỹ (x) > 0By Lemma 3 of [Dev83] , we know that for all ∆ ∈ (0, 1), if |N (y)| L ≤ ∆ 2 20 , we haveP max i∈N (y) |α i -p(i|y)| > ∆ ≤ P i∈N (y) |α i -p(i|y)| > ∆ ≤ 3 exp(-L∆ 2 /25) Therefore, if L ≥ max{ 25 log(3/ε) ∆ 2 , 20|N (y)| ∆ 2}, then with probability at least 1 -ε, we have,max i∈N (y) |α i -p(i|y)| ≤ ∆ f L y (x) -f L ỹ (x) = i∈N1(y) α i - j∈N1(ỹ) α j = i∈N1(y) α i - i∈N1(y) p(i|y) + i∈N1(y) p(i|y) - j∈N1(ỹ) p(j|y) + j∈N1(ỹ) p(j|y) - j∈N1(ỹ) α j ≥ i∈N1(y) p(i|y) - j∈N1(ỹ) p(j|y) -2m∆ ≥ exp(- 1 β ) -exp(- 2 β ) -2m∆Note that because of Lemma 10, there's no neighboring set that is the superset of another.Therefore as long as∆ < exp(-1 β )-exp(-2 β ) 2m , f L y (x) -f L ỹ (x) > 0 for any ỹ. Finally, if L > max{ 100m 2 log(3/ε) (exp(-1 β )-exp(-2 β )) 2 , 80m 2 |N (y)| (exp(-1 β )-exp(-2 β )) 2 } for any y, then P x|y [argmax f L (x) ̸ = y] ≤ ε And R D L (f L ) = P (x,y)∼D L [argmax f L (x) ̸ = y] = P y P x|y [argmax f L (x) ̸ = y] ≤ ε B.2 Proofs for Section 5.2Lemma 2. Suppose the data generating process follows Section 4.1 where m ≥ 3, ω = 1 and N (t) = {t ′ : D H (t, t ′ )) = 1}. For any single layer transformer given by (4.1) where each row of W E is orthogonal to each other and normalized, if W V is constructed as in (5.1), then the error rate is 0. If W V is the identity matrix, then the error rate is strictly larger than 0.Proof. Following the proof for Theorem 7, let's focus on the conditional probability:P x|y [argmax f L (x) ̸ = y]By construction, we have h(x) = i∈N1(y)α i W E (i)where α i = 1 L L k=1 1{t k = i} which is the number of occurrence of token i in the sequence. Let's consider the first case where W V is constructed as in (5.1). Then we know that for some other tokenỹ ̸ = y, f L y (x) -f L ỹ (x) = i∈N1(y) α i - i∈N1(ỹ) α i = 1 - i∈N1(ỹ) α iBy Lemma 10, we have that for any token ỹ ̸ = y, f L y (x) -f L ỹ (x) > 0 Therefore, the error rate is always 0. Now let's consider the second case where W V is the identity matrix. Let j be a token in the set N 1 (y). Then there is a non-zero probability that context x contains only j. In that case,h(x) = W E (j)However, we know that by the assumption on the embedding matrix,f L y (x) -f L j (x) = (W E (y) -W E (j)) T h(x) = -∥W E (j)∥ 2< 0 This implies that there's non zero probability that y is misclassified. Therefore, when W V is the identity matrix, the error rate is strictly larger than 0.Theorem 3. Suppose the data generating process follows Section 4.1 where m ≥ 3, ω = 1 and N (t) = V \ {t}. For any single layer transformer given by (4.1) with W V being the identity matrix, if the cross entropy loss is minimized so that for any sampled pair (x, y), p(y|x) = p(y|x) = softmax(f L y (x)) there exists a > 0 and b such that for two tokens t̸ = t ′ , ⟨W E (t), W E (t ′ )⟩ = -aD H (t, t ′ ) + bProof. Because for any pair of (x, y), the estimated conditional probability matches the true conditional probability. In particular, let's consider two target tokens y 1 , y 2 and context x = (t i , ..., t i ) for some token t i such that p(x|y 1 ) > 0 and p(x|y 2 ) > 0, thenp(y 1 |x) p(y 2 |x) = p(x|y 1 )p(y 1 ) p(x|y 2 )p(y 2 ) = p(x|y 1 ) p(x|y 2 ) = p(x|y 1 ) p(x|y 2 ) = exp((W E (y 1 ) -W E (y 2 )) T h(x))The second equality is because p(y) is the uniform distribution. By our construction,p(x|y 1 ) p(x|y 2 ) = p(t i |y 1 ) L p(t i |y 2 ) L = exp((W E (y 2 ) -W E (y 1 )) T h(x)) = exp((W E (y 1 ) -W E (y 2 )) T W E (t i ))By the data generating process, we have thatL β (D H (t i , y 2 ) -D H (t i , y 1 )) = (W E (y 1 ) -W E (y 2 )) T W E (t i ) Let t i = y 3 such that y 3 ̸ = y 1 , y 3 ̸ = y 2 , then L β D H (y 3 , y 1 ) -W E (y 1 ) T W E (y 3 ) = L β D H (y 3 , y 2 ) -W E (y 2 ) T W E (y 3 )For simplicity, let's defineΨ(y 1 , y 2 ) = L β D H (y 1 , y 2 ) -W E (y 1 ) T W E (y 2 )Therefore, Ψ(y 3 , y 1 ) = Ψ(y 3 , y 2 )Now consider five distinct labels: y 1 , y 2 , y 3 , y 4 , y 5 . We have, Ψ(y 3 , y 1 ) = Ψ(y 3 , y 2 ) = Ψ(y 4 , y 2 ) = Ψ(y 4 , y 5 )In other words, Ψ(y 3 , y 1 ) = Ψ(y 4 , y 5 ) for arbitrarily chosen distinct labels y 1 , y 3 , y 4 , y 5 . Therefore,Ψ(t, t ′ ) is a constant for t ̸ = t ′ .For any two tokens t̸ = t ′ , L β D H (t, t ′ ) -W E (t) T W E (t ′ ) = C Thus, W E (t) T W E (t ′ ) = - L β D H (t, t ′ ) + C B.3 Proofs for Section 5.3Theorem 4 can be formalized as the following theorem.Theorem 8. Following the same setup as in Theorem 7, but embeddings follow (5.2) then if b > 0,∆ 1 > 0, 0 < ∆ < exp(-1 β )-exp(-2 β ) 2m , L ≥ max{ 25 log(3/ε) ∆ 2 , 20|N (y)| ∆ 2} for any y, and0 < a < 2 exp( 1 β ) (|V | -2)m 2 and b 0 > max{ a(m -2)m + ∆ 1 exp(-1 β ) -exp(-2 β ) -2m∆ +b, (b -a)∆ 1 -|V |-2 2 abm 2 exp(-1 β ) + |V |-2 2 a 2 (m -2)m 2 1 -|V |-2 2 am 2 exp(-1 β ) } we have R D L (f L ) ≤ ε where 0 < ε < 1.Proof. Following the proof of Theorem 7, let's also focus on the conditional probabilityP x|y [argmax f L (x) ̸ = y]By construction, the single layer transformer model has uniform attention. Therefore,h(x) = i∈N (y) α i W E (i) where α i = 1 L L k=1 1{t k = i}which is the number of occurrence of token i in the sequence. For simplicity, let's define α y = 0 such thath(x) = i∈[V ] α i W E (i) Similarly, we also have that if L ≥ max{ 25 log(3/ε) ∆ 2 , 20|N (y)| ∆ 2}, then with probability at least 1 -ε, we have, maxi∈[V ] |α i -p(i|y)| ≤ ∆Also define the following:ϕ k (x) = j∈N1(k) W E (j) T i∈[V ] α i W E (i) v k (y) = W E (y) T W E (k)Thus, the logit for token y isf L y (x) = |V |-1 k=0 v k (y)ϕ k (x)Let's investigate ϕ k (x). By Lemma 9,ϕ k (x) = i∈[V ] α i ( j∈N1(k) W E (j) T W E (i)) = (b 0 -b) j∈N1(k) α j + i∈[V ] α i (-a(m -2)D H (k, i) + (b -a)m) Thus, for any k 1 , k 2 ∈ [V ], ϕ k1 (x) -ϕ k2 (x) = (b 0 -b)( j1∈N1(k1) α j1 - j2∈N1(k2) α j2 ) + i∈[V ] α i a(m -2)(D H (k 2 , i) -D H (k 1 , i)) Because -m ≤ D H (k 2 , i) -D H (k 1 , i) ≤ m, we have (b 0 -b)( j1∈N1(k1) α j1 - j2∈N1(k2) α j2 ) -a(m -2)m ≤ ϕ k1 (x) -ϕ k2 (x) ≤ (b 0 -b)( j1∈N1(k1) α j1 - j2∈N1(k2) α j2 ) + a(m -2)mFor prediction to be correct, we needmax ỹ f L y (x) -f L ỹ (x) > 0This also means thatmax ỹ |V |-1 k=0 v k (y) -v k (ỹ) ϕ k (x) > 0 One can show that for any k, if ι -1 ( k) = ι -1 (y) ⊗ ι -1 (ỹ) ⊗ ι -1 (k) where ⊗ means bitwise XOR, then v k (y) -v k (ỹ) = v k(ỹ) -v k(y) (B.1)First of all, if k = y, then k = ỹ, which meansv k (y) -v k (ỹ) = v k(ỹ) -v k(y) = b 0 + aD H (y, ỹ) -b If k ̸ = y, ỹ, then (B.1) implies that D H (k, y) -D H (k, ỹ) = D H ( k, ỹ) -D H ( k, y)We know that D H (k, y) is the number of 1s in ι -1 (k) ⊗ ι -1 (y) and,ι -1 ( k) ⊗ ι -1 (y) = ι -1 (y) ⊗ ι -1 (ỹ) ⊗ ι -1 (k) ⊗ ι -1 (y) = ι -1 (ỹ) ⊗ ι -1 (k) Similarly, ι -1 ( k) ⊗ ι -1 (ỹ) = ι -1 (y) ⊗ ι -1 (k)Therefore, (B.1) holds and we can rewritef L y (x) -f L ỹ (x) as f L y (x) -f L ỹ (x) = |V |-1 k=0 v k (y) -v k (ỹ) ϕ k (x) = (b 0 -b + aD H (y, ỹ))(ϕ y (x) -ϕ ỹ (x)) + k̸ =y,ỹ,D H (k,y)≥D H (k,ỹ) a(D H (k, y) -D H (k, ỹ))(ϕ k (x) -ϕ k(x))We already know that b 0 > b > 0 and a > 0, thus, b 0 -b + aD H (y, ỹ) > 0 for any pair y, ỹ.We also want ϕ y (x) -ϕ ỹ (x) to be positive. Note thatϕ y (x) -ϕ ỹ (x) ≥ (b 0 -b)(exp(- 1 β ) -exp(- 2 β ) -2m∆) -a(m -2)m We need ∆ < exp(-1 β )-exp(-2 β ) 2mand for some positive ∆ 1 > 0, b 0 needs to be large enough such thatϕ y (x) -ϕ ỹ (x) > ∆ 1 which implies that b 0 > a(m -2)m + ∆ 1 exp(-1 β ) -exp(-2 β ) -2m∆ + b (B.2)On the other hand, for k ̸ = y, ỹ, we haveϕ k (x) -ϕ k(x) ≥ (b 0 -b)( j1∈N1(k) α j1 - j2∈N1( k) α j2 ) -a(m -2)m ≥ (b 0 -b)(-(m -1) exp(- 1 β ) -exp(- 2 β ) -2m∆) -a(m -2)m ≥ (b 0 -b)(-(m -1) exp(- 1 β ) -exp(- 2 β ) + exp(- 2 β ) -exp(- 1 β )) -a(m -2)m ≥ -(b 0 -b)m exp(- 1 β ) -a(m -2)mThen, we havef L y (x) -f L ỹ (x) ≥ (b 0 -b + a)∆ 1 - |V | -2 2 (b 0 -b)am 2 exp(- 1 β ) + a 2 (m -2)m 2 ≥ 1 - |V | -2 2 am 2 exp(- 1 β ) b 0 -(b -a)∆ 1 + |V | -2 2 abm 2 exp(- 1 β ) - |V | -2 2 a 2 (m -2)m 2The lower bound is independent of ỹ, therefore, we need it to be positive to ensure the prediction is correct. To achieve this, we want1 - |V | -2 2 am 2 exp(- 1 β ) > 0 which implies that a < 2 exp( 1 β ) (|V | -2)m 2 (B.3)And finally we needb 0 > (b -a)∆ 1 -|V |-2 2 abm 2 exp(-1 β ) + |V |-2 2 a 2 (m -2)m 2 1 -|V |-2 2 am 2 exp(-1 β ) (B.4) To summarize, if b > 0, ∆ 1 > 0, 0 < ∆ < exp(-1 β )-exp(-2 β ) 2m , L ≥ max{ 25 log(3/ε) ∆ 2 , 20|N (y)| ∆ 2} for any y, and0 < a < 2 exp( 1 β ) (|V | -2)m 2 and b 0 > max{ a(m -2)m + ∆ 1 exp(-1 β ) -exp(-2 β ) -2m∆ +b, (b -a)∆ 1 -|V |-2 2 abm 2 exp(-1 β ) + |V |-2 2 a 2 (m -2)m 2 1 -|V |-2 2 am 2 exp(-1 β ) } we have R D L (f L ) ≤ εwhere 0 < ε < 1.Lemma 5. If embeddings follow (5.2) and b = b 0 and N (t) = V \ {t}, then rank(W E ) ≤ m + 2.Proof. By (5.2), we have that⟨W E (i), W E (j)⟩ = -aD H (i, j) + b Therefore, (W E ) T W E = -aD H + b11 TLet's first look at D H which has rank at most m + 1. To see this, let's consider a set of m + 1 tokens: {e 0 , e 1 , ..., e m } ⊆ V where e k = 2 k . Here e 0 is associated with the latent vector of all zeroes and the latent vector associated with e k has only the k-th latent variable being 1.On the other hand, for any token i, we have that,i = k:ι -1 (i) k =1 e kIn fact,D H (i) = k:ι -1 (i) k =1 D H (e k ) -D H (e 0 ) + D H (e 0 )where D H (i) is the i-th row of D H , and for each entry j of D H (i), we have thatD H (i, j) = k:ι -1 (i) k =1 D H (e k , j) -D H (e 0 , j) + D H (e 0 , j)This is becauseD H (e k , j) -D H (e 0 , j) = +1 if ι -1 (j) k = 0 -1 if ι -1 (j) k = 1Thus, we can rewrite D H (i, j) asD H (i, j) = k:ι -1 (i) k =1 1[ι -1 (i) k = 1, ι -1 (j) k = 0] -1[ι -1 (i) k = 1, ι -1 (j) k = 1)] + D H (e 0 , j) = m k=1 1[ι -1 (i) k = 1, ι -1 (j) k = 0] -1[ι -1 (i) k = 1, ι -1 (j) k = 1)] + m k=1 1[ι -1 (i) k = 0, ι -1 (j) k = 1] + 1[ι -1 (i) k = 1, ι -1 (j) k = 1)] = m k=1 1[ι -1 (i) k = 1, ι -1 (j) k = 0] + 1[ι -1 (i) k = 0, ι -1 (j) k = 1] = D H (i, j)Therefore, every row of D H can be written as a linear combination of {D H (e 0 ), D H (e 1 ), ..., D H (e m )}. In other words, D H has rank at most m + 1.Therefore, rank((W E ) T W E ) = rank(W E ) ≤ m + 2.Lemma 9. Let z (0) and z (1) be two binary vectors of size m where m ≥ 2. Then,z:D H (z (0) ,z)=1 D H (z, z (1) ) = (m -2)D H (z (0) , z (1) ) + mProof. For z such that D H (z, z (0) ) = 1, we know that there are two cases. Either z differs with z (0) on a entry but agrees with z (1) on that entry or z differs with both z (0) and z (1) .For the first case, we know that there are D H (z (0) , z (1) ) such entries. In this case, D H (z, z (1) ) = D H (z (0) , z (1) ) -1. For the second case, D H (z, z (1) ) = D H (z (0) , z (1) ) + 1.Therefore,z:D H (z,z (0) )=1 D H (z, z (1) ) = D H (z (0) , z (1) )(D H (z (0) , z (1) ) -1) + (m -D H (z (0) , z (1) ))(D H (z (0) , z (1) ) + 1) = (m -2)D H (z (0) , z (1) ) + m Lemma 10. If m ≥ 3 and N (t) = V \ {t}, then N 1 (t) ̸ ⊆ N 1 (t ′ ) for any t, t ′ ∈ [V ].Proof. For any token t, N 1 (t) contains any token t ′ such that D H (t, t ′ ) = 1 by the conditions. Then given a set N 1 (t), one can uniquely determine token t. This is because for the set of latent vectors associated with N 1 (t), at each index, there could only be one possible change.B.4 Proofs for Section 5.4Lemma 6. Suppose the data generating process follows Section 4.1 and N (z * ) = {z : z * 1 = z 1 } \ {z * }. Given the last token in the sequence t L , then∇ ut,t L ℓ(f L ) = ∇ℓ(f L ) T (W E ) T W V (α t pt W E (t) -pt L l=1 pt l W E (t l ))where for token t, α t = L l=1 1[t l = t] and pt is the normalized attention score for token t.Proof. Recall that,f L (x) = W E T W V attn(W E χ(x)) :L = W E T W V L l=1 exp(u t l ,t L ) Z W E (t l )where Z is a normalizing constant.Define pt l = exp(ut l ,t L ) Z. Then we havef L (x) = W E T W V L l=1 pt l W E (t l ) Note that if t l = t then, ∂ pt l ∂u t,t L = pt l (1 -pt l ) Otherwise, ∂ pt l ∂u t,t L = -p t l ptBy the chain rule, we know that∇ ut,t L ℓ(f L ) = ∇ℓ(f L ) T (W E ) T W V ( L l=1 1[t l = t]p t l W E (t) - L l=1 pt l pt W E (t l ))Therefore,∇ ut,t L ℓ(f L ) = ∇ℓ(f L ) T (W E ) T W V (α t pt W E (t) -pt L l=1 pt l W E (t l ))whereα t = L l=1 1[t l = t].In this section, we show the results of additional context hijacking experiments on the COUNTERFACT dataset [Men+22] .Reverse context hijacking In Figure 2a , we saw the effects of hijacking by adding in "Do not think of {target_false}." to each context. Now, we measure the effect of the reverse: What if we prepend "Do not think of {target_true}." ?Based on the study in this paper on how associative memory works in LLMs, we should expect the efficacy score to decrease. Indeed, this is what happens, as we see in Hijacking based on relation IDs We first give an example of each of the 4 relation IDs we hijack in Table 1 . Hijacking without exact target words So far, the experiments use prompts that either contain true or false target words. It turns out, the inclusion of exact target words are not necessary. To see this, we experiment a variant of the generic hijacking and reverse hijacking experiments. But instead of saying "Do not think of {target_false}" or "Do not think of {target_true}". We replace target words with words that are semantically close. Specifically, for relation P1412, we replace words representing language (e.g., "French") with their associated country name (e.g., "France"). As shown in In this appendix section, we present additional experimental details and results from the synthetic experiments on latent concept association.Experimental setup Synthetic data are generated following the model in Section 4.1. Unless otherwise stated, the default setup has ω = 0.5, β = 1 and N (i) = V \ {i} and L = 256. The default hidden dimension of the one-layer transformer is also set to be 256. The model is optimized using AdamW [LH17] where the learning rate is chosen from {0.01, 0.001}. The evaluation dataset is drawn from the same distribution as the training dataset and consists of 1024 (x, y) pairs. Although theoretical results in Section 5 may freeze certain parts of the network for simplicity, in this section, unless otherwise specified, all layers of the transformers are trained jointly. Also, in this section, we typically report accuracy which is 1error.In this section, we provide additional figures of Section 6.1. Specifically, Figure D .1 shows that fixing the value matrix to be the identity will negatively impact accuracy. Figure D.2 indicates that replacing trained value matrices with constructed ones can preserve accuracy to some extent. Figure D.3 suggests that trained value matrices and constructed ones share similar low-rank approximations. For the last two sets of experiments, we consider randomly constructed value matrix, where the outer product pairs are chosen randomly, defined formally as follows:W V = i∈[V ] W E (i)( {j}∼Unif([V ]) |N 1 (i)| W E (j) T )This section provides additional figures from Section 6.2. This section provides additional figures from Section 6.3. In this section, we want to simulate context hijacking in the latent concept association model. To achieve that, we first sample two output tokens y 1 (true target) and y 2 (false target) and then generate contexts x 1 = (t 1 1 , ..., t 1 L ) and x 2 = (t 2 1 , ..., t 2 L ) from p(x 1 |y 1 ) and p(x 2 |y 2 ). Then we mix the two contexts with rate p m . In other words, for the final mixed context x = (t 1 , ..., t L ), t l has probability 1 -p m to be t 1 l and p m probability to be t 2 l . Figure D .13 shows that, as the mixing rate increases from 0.0 to 1.0, the trained transformer tends to favor predicting false targets. This mirrors the phenomenon of context hijacking in LLMs.As alluded in Section 5.5, the memory recall rate is closely related to the KL divergences between context conditional distributions. Because contexts contain mostly i.i.d samples, longer contexts imply larger divergences. This is empirically verified in Figure D.14 which demonstrates that longer context lengths can lead to higher accuracy. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 


Decision: Summary: 
The paper studies the brittleness of factual recall among a variety of LLMs, where the authors show that by simply manipulating the words within the context of a prompt (without changing their semantics), one can force the model to generate wrong answers (a.k.a., context hijacking). The authors then propose a theoretical formulation of a hypothesis that could explain such a phenomenon and show its empirical applicability on single-layer transformers (the building blocks of LLMs). Specifically, they show that single-layer transformers gather information using self-attention and use the value matrix as an associative memory.  

Strengths: 
Understanding the workings of LLMs is an important problem, and the ML community at large would greatly benefit from it. The paper takes an important step in that direction by showing a specific failure case and unraveling its potential cause. The authors propose a simple theory and convincingly prove its validity empirically, albeit on single-layer transformer models. The paper is well-organized, well-written, and easy to understand. The theory presented is clear, and the empirical evidence is sufficient. 

Weaknesses: 
One weakness of the paper, which should be addressed in the final version (if accepted), is that the authors do not discuss the limitations of their work. Specifically, under what circumstances will their proposed theory fail? On a related note, for instance, the experiments are done only on single-layer transformer models. While they are indeed the building blocks of LLMs, there’s likely a lot more going on in full LLM models. The paper would be strengthened if the authors could comment on how their ideas could be generalized to explain more complex models. 

Overall Recommendation: 
Given the importance of the problem the paper is addressing, the soundness of the solution proposed, and the overall quality of the presentation, the reviewers unanimously recommend that this paper be accepted. The wider ML community will benefit from the dissemination of this research.

Reviews: review: paper_summary: paper_summary: The paper presents a way to study associative memory in Transformer blocks. Specifically, the author presents a method to construct a value matrix representing associative memory and suggests its equivalence to self-attention’s value matrix. Through experiments based on synthetic data, the author proposes that the transformer gathers information through self-attention, while the value matrix stores associative memory.
strength_weakness: strength_weakness: 1. Besides the findings based on pre-trained LLMs, I would like the author to dig deeper into other methods to solve context hijacking. For example, will or to what extent can supervised finetuning correct the distracted focus back to the correct context?

2. Given the prevalence of low-bit quantization, I wonder how quantization will change the associative memory. For example, after quantization, will LLM be less distracted? Or perhaps quantization will enhance the effects of the misleading context?

3. For the constructed value matrix (formula 4.1 and 5.1), even if results in section 6.1 show that using the constructed value matrix retains the accuracy, two concerns remain. First, why should a value matrix constructed from an embedding matrix be expected to act as an associative memory (intuitively)? Secondly, how and why is the constructed value matrix different from self-attention’s value matrix (gather information), if any?

4. In the results section, it is unclear how the method performs differently across different LLMs, different datasets, and other SOTA methods. It is unclear if the results can be generalized to other settings.
questions: questions: 1. In Section 5.3, it is unclear how the embeddings are trained/updated.

2. In Section 5.5, the support evidence is not sufficient. Instead of showing across different mixture rates, I would like to see how changing context (for example, adding “The Eiffel Tower is not in Chicago” to the beginning of the prompt) can potentially impact associative memory and self-attention. Another concern regarding Fig C.13 is, how/if the impact will be different should one concatenate additional context (“The Eiffel Tower is not in Chicago”) to the beginning, middle, and end of prompt.

3. In Section 6.1, Figure C.2, the author concludes that the constructed value matrix can be used to replace the self-attention value matrix without performance sacrifices. However, the author fails to provide enough justification on the consistent drop after certain dimensions, across different m. For example, when m=5, the accuracy of using the constructed matrix drops significantly after dim = 128, which hints that the constructed matrix and self-attention’s value matrix are not equivalent.
limitations: limitations: The author didn't specify any overall limitations or potential negative impact of their method. Some discussions of insufficient/abnormal behavior of the results (see the section strength/weakness and the section questions) would be helpful.
paper_summary: paper_summary: This paper investigates the mechanisms underlying factual recall in transformer language models. First, the paper demonstrates a "context hijacking" phenomenon, where distractor sentences lead language models to output the wrong answer to factual questions. The paper conducts a theoretical analysis of a one-layer transformer on a noisy associative memory tasks, showing how the context hijacking phenomenon could arise. These findings are supported with experiments on synthetic data.
strength_weakness: strength_weakness: - A key argument of the paper is that an LLM can be seen as an "associative memory model", but I feel that this term lacks a precise definition. For example, one definition is that "tokens in contexts guide the retrieval of memories, even if such associations formed are not inherently semantically meaningful". It seems that the first part of this sentence would apply to any question answering model, and the notion of "not inherently semantically meaningful" needs to be defined. I think it would be especially helpful to give some examples of what an alternative model would be--associative memory, as opposed to what?
- I am not fully convinced that the one-layer transformer is a meaningful model the context hijacking phenomenon. For example, in the main example ("the Eiffel tower is not in Chicago"), it's seems like the most plausible mechanisms for either resisting context hijacking, or falling for context hijacking, would involve multi-layer Transformers--for example, the model might predict "the Eiffel tower is in Chicago" due to a kind of induction head/ICL mechanism. I think it would be helpful to expand in more detail on the connection between context hijacking and the toy model (section 5.5).
questions: questions: - For Efficacy Score (section 3), it seems it should also require that Pr[o_] < Pr[o*] prior to modifying the context. In Fig 2a, the efficacy score seems to be the opposite of what is described in the text--the intervention makes the score go down (i.e., more often that Pr[o_] > Pr[o*]).
- In Fig 2a, are these results averaged over all relation types?
- In Section 5.4: "This implies that the self-attention layer can mitigate noise and concentrate on the informative conditional distribution π". In this setting, the final token is sampled without noise. Would this also be true if the final token had noise?
limitations: limitations: I think the authors adequately addressed the limitations of their work.
paper_summary: paper_summary: This paper studies the mechanics of factual recall in transformer-based models, framing the problem as next token prediction. In particular, the authors focus on the brittleness of language models, which can be elicited to provide different answers to factual queries by adding distracting information in the prompt (a procedure which the authors term context hijacking). This phenomenon is shown for various language models, with sizes up to 7B parameters.
The authors then formulate a hypothesis for the hijacking phenomenon, according to which the model predicts the next token based on a similarity measure in a latent concept space. They provide a series of theoretical results that explain how a single-layer transformer can solve the latent concept association problem.
Finally, the paper includes empirical validation of the theoretical results.
strength_weakness: strength_weakness: - The authors do not discuss any limitations of their work. For instance, the authors assume that each latent concept is associated to one token only, how would the theoretical results look without this assumption? Moreover, the study is motivated by the behavior of multi-billion-parameter models but focuses on a single-layer transformer, how do the authors expect their results to generalize to larger models?
- In the motivation of the problem, the authors show how LLaMA 7B can be “hijacked” to output a wrong answer to the prompt about the location of the Eiffel Tower. However, a text consisting of the sentence “The Eiffel Tower is not in Chicago” repeated eight times represents an input arguably out of distribution with respect to the model’s pre-training data. In this setting, it is possible that the model continues the prompt ”the Eiffel Tower is in the city of” using a different mechanism than it would have without the prepended hijacking context. A test for the authors’ hypothesis would be to prepend a non-templated paragraph about, for example, the city of Chicago (possibly mentioning that the Eiffel Tower is not located there). If the authors’ hypothesis is correct, this should still steer the model towards completing “the Eiffel Tower is in the city of” with “Chicago.” Is this the case?
- Minor point: the figures can be improved (e.g., increasing the font size)
questions: questions: - Maybe I am missing something, but isn’t the efficacy score expected to increase as the hijacking context gets longer (i.e., the more times the sentence “Do not think of {target_false}.” gets repeated, the more likely the model is to assign higher probability mass to target_false)?
limitations: limitations: The authors do not address any limitation of their work. Their answer to the checklist question "Does the paper discuss the limitations of the work performed by the authors?" the authors motivate their "yes" answer with a single sentence: "Studying single-layer transformers is limited."
paper_summary: paper_summary: This paper demonstrates the phenomenon of "context hijacking" in LLMs, where repeated mentions of sentences in the context could negatively influence a model's factual recall. Motivated by this, the authors then formulate an associative memory recall task and prove theoretically certain properties of one-layer transformer models on this task, which shows that transformers are capable of performing the task and the roles of their individual components in the task. Further experiments verify the results.
strength_weakness: strength_weakness: - It is unclear to what degree the context hijacking phenomenon still exists in more powerful models (e.g., GPT-4) with better semantic understanding. For example, I tried the proposed attack methods on GPT-4o and it never influenced the model negatively.

- Related to the last point, the proposed associative memory formation lacks modeling of semantics, where the task boils down to tokens and their similarity. This may be a good formulation for models that are not powerful enough and mostly rely on surface cues to generate the next token, but not for models that have strong language understanding and logic. This also relates to the shallow network (one-layer transformer) which is the main target for theoretical analysis. The results may not transfer/be enlightening for improving the current LLMs.
questions: questions: None
limitations: limitations: It would be better to include a dedicated "limitations" section beyond those in the checklist.

Comments: comment: Thank you for taking the time to review our work and for adjusting your score. We greatly appreciate your effort and consideration. Regarding value matrices, we think that the constructed value matrices are only approximations. Achieving a more precise approximation could likely enhance accuracy further and is an interesting future direction.

For Section 5.5, thanks for the suggestion! We haven't yet examined the attention differences before and after context hijacking, and we agree that this is a very interesting direction to explore in future work. 

For the second half, we have only tried putting misleading prompts at the beginning. This is because we want to keep the original query prompt intact. Inserting misleading prompts in the middle could cause the original prompt to lose its meaning. Placing them at the end would alter the original next token prediction task.
Thank you for your responses. I will raise my evaluation.

Re. Value matrix 
Thank you for the clarifications. Given that those accuracies after replacement do not match exactly in Fig C.2, it is worthwhile to address this, as supposed to L337 “Figure C.2 indicates that the accuracy does not significantly decrease when the value matrix is replaced with the constructed ones”

Re. section 5.5 
Regarding the first half of the original comment, the question was, besides Efficacy Scores or Accuracy, have the author explored/proved a shift in attention after context hijacking? Regarding the second half, I wonder if the author has explored context hijacking to the middle, or end of the prompt?
Thanks for your reply! Although the term “associative memory” is used loosely in the literature, we focus on a particular type of associative memory (L123-L125). To analyze this rigorously, we precisely define the latent concept association task in Sec 4.1 which concretely formalizes the idea that tokens with shared latent concepts would co-occur more (L148-184, in particular L182-184 where the final objective is defined). This is complemented by a rigorous theoretical analysis in Sec 5 (Theorems 1, 4, see also Theorem 7, 8 in App A) and detailed experiments in Sec 6. This is all in precise language that can be falsified.

Furthermore, to motivate the precise task we defined, we first conducted systematic experiments on context hijacking to show that prepending the same misleading prompts more can cause LLMs to perform worse (Figure 2). This motivated us to hypothesize that LLMs might pay attention to the frequency of certain tokens in the context as supposed to understand the factual meaning of the context, which leads to the precise task we analyze in this paper.

On the other hand, in this paper, we are interested in studying a **failure mode** of LLMs – context hijacking. Since we show via experiments that the failure occurs in large models like LLaMA, one can reasonably expect it to persist in smaller models as well. Because smaller models allow us to do carefully controlled experiments, it is _more meaningful_ to study how this problem persists in smaller models as a starting point, which is what we provide evidence for with latent concept association in single-layer transformers. This is different from in-context learning and induction head which are shown to work mostly for models with at least two layers. 


We hope this clarifies your concerns, and appreciate the opportunity to discuss these details with you.
Thank you for responding to my questions. I will keep my score as it is (5). While I appreciate the clarifications, I still have concerns that the concept of "associative memory model" is imprecise. For me to give a higher score, I think this hypothesis needs to be defined more precisely and contrasted with some other possible explanation for the observed phenomena. Similarly, I am still not convinced that a one-layer transformer is a very useful model for thinking about context hijacking, given that most mechanisms I can think of that might be relevant in this setting would require at least two transformer layers.
Thanks for your support! We will certainly incorporate these discussions into the final version of the paper. They have been very helpful in clarifying our work.
We are pleased that your concerns have been addressed. We truly appreciate your time and effort to engage with our work, and for updating your score accordingly.
Thank you for your responses.

Re. Limitations: I recommend including these comments in the final version of the paper.

Re. Hijacking prompt: Thank you for conducting this quick test. I believe the paper would benefit from including a discussion on the limitations and potential side effects of hijacking using a repeated template, particularly in relation to the function of induction heads, as raised by Reviewer bm96.
Thank you for the response, which addresses my concerns to some degree. I will raise my evaluation.