Do LLMs dream of elephants (when told not to)? Latent concept association and associative memory in transformers
Authors: Yibo Jiang
Goutham Rajendran
Pradeep Kumar Ravikumar
Bryon Aragam
Venue: NeurIPS 2024 poster (2024)
Keywords: Transformer
Associative Memory
Large Language Models
Interpretability
Fact retrieval
Primary area: interpretability_and_explainability
TL;DR: We study a 1-layer transformer to understand the mechanisms of associative memory in LLMs

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.

Do LLMs dream of elephants (when told not to)? Latent concept association and associative memory in transformers
Yibo Jiang', Goutham Rajendran', Pradeep Ravikumar", and Bryon Aragam'
Department of Computer Science, University of Chicago Machine Learning Department, Carnegie Mellon University
Booth School of Business, University of Chicago
Abstract
Large Language Models (LLMs) have the capacity to store and recall facts. Through experimentation with open-source models, we observe that this abil- ity 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 simple latent concept association problem with 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.
1 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 nontrivial way? To gain 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, first demonstrate that fact retrieval is not robust and LLMs can be casily 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 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 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, reasoming 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.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Contest Hijacking
MODEL.
CONTENT
NEXT
models
Eiflel
city
GPT-2 /iemma-28
The Eiftel Tower not in Chicago, Therefore, the Eiflel Tower in the city of
Chicagn
The Eiffel Tower is not in Chicago However, the Chicago river is in Chicago Therefore, the Hifflel Tower the city of
Gemma-28-IT 28-IT
Chicago
The Eiffel Tower
Chicago. The Biffel Tower Chicago.
Chicago, Eiffel Hiffel Tower not Chicago. Hiffll not Chicago. The Biffel Tower is not Chicago. The Giflel Tower is not in Chicagoo The Eifiel Tower not in Chicago. Therefore, the Eiflel Tower in
LLaMA-7B
Chicago
Figure 1: Examples of context hijacking for various LLMS,, showcasing that fact retrieval is not robust.
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 more nuanced notion of similarity measured by latent (unobserved) semantic concepts. To model this, we propose synthetic task called latent concept association where the output token is closely related to sampled tokens in the context but wherein similarity measured via a latent space of semantic concepts. We then investigate how one-layer transformer [Vas++1] 17], 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-tur tinettnnig techniques [Men+22; Hu+21].
Contributions Specifically, make the following contributions:
1. We systematically demonstrate context hijacking for various open source LLM models including GPT-2 [Rad+19], LLaMA-2 [Tou++33] and Gemma [Teaa+24], which show that fact retrieval can be misled by contexts (Section 3), reaffirming that LLMs lack robustness context changes [Shi+23; Per+20; CSH22; Yor+23; PE21J.
2. We propose 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 latent, semantic concept space as opposed to directly between observed tokens. This perspective crucial solve fact retrieval problems by implementing
understanding how transformers 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 exhibit low-rank structure, offering additional support for existing editing and fine-tuning methods [Men+22; Hu+21].
2 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 well. This paper studies fact retrieval form of associative memory. Another closely

---

Prepending 'Do think (target false).
Hijacking based P190
10-
0.7
04
Efficacy Score
2-4
Efficacy Score
9.5
0.6
caena communtryppe
openai communiyynnn googhogemmm op:akeppmima-2ee
gagglogenme
01-
Number of prepends
Number of prepends
(a) Hijacking generically
(b) Hijacking based on Relation ID P190
Figure 2: Context hijacking can cause LLMs output false target. The figure shows efficacy versus the number of prepends for various LLMs on the COUNTERFACT dataset under two hijacking schemes.
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 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 Kim et al. [KKM22] and Mahdavi et al. [MLT23] study the memorization capacity of transformers.
Interpreting transformers and LLMs There's a growing body work on understanding how transformers and LLMs work [LLR23; AZL23a; AZL23b; AZL24; EI+24; Tar+23b; Tar+23a; Li+24], including training dynamics [Tia+23a; Tia+23b; She+24] and in-context learning [Xie+21;; Gar+22; Bai+24; Bai+24].. Recent papers have introduced synthetic tasks to better understand the mechanisms of transformers [Cha22; Liu+22; Nan+23; Zha+22; Zho+24], such as those focused on Markov chains [Bie+24; Ede+24; NDL.2; Mak+24]. Most notably, Bietti al. [Bie+24] and subsequent works [CDB23; CSB24] study weights in transformers as associative memory but their focus on understanding induction head [Ols+22b] and one-to-one map between input query and output memory. An increasing amount of research is dedicated to understanding the internals of pre- trained LLMs, broadly categorized under the term "mechanistic interpretability" [Elh+21; Ols+22a; Gev+23; Men+22; Men+23; Jia+24; Raj+24; Has+24; Wan+22; McG+23; Gei+21; Gei+22; Gei+24; Wu+24b].
Knowledge editing and adversarial attacks on LLMs Fact recall and knowledge editing have been extensively studied [Men+22; Men+23; Has+24; Sak+23; DCAT21; Mit+21; Mit422; Dai+21; Zha+23; Tia+24; Jin+23]] including the use of in-context learning to edit facts [Zhe+233]. This paper aims to explore a different aspect by examining the robustness of fact recall to variation in prompts. A closely related line of work focuses on adversarial attacks on LLMs [see Cho+24, for a review]. Specifically, prompt-based adversarial attacks [Xu+23; Zhu+23; Wan+23b) focus on the manipulation of answers within specific classification tasks while other works concentrate on safety issues [Liu+23a; PR22; Zou+23; Apr+22; Wan+23a; Si+22; Rao+23; SMR23; Liu+23b). Yu al. [Yu+24] and Luo et al. [Luo+24] also study jailbreak phenomena within the context of modern Hopfield network. There are also works showing LLMS can distracted by irrelevant contexts in problem solving [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 systematic robustness study for fact retrieval, offer framework for interpreting it in the context of associative memory, and deepen our understanding of LLMs.
3 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 human. We call
3

---

this context- Due to the different capacities and capabilties of each model, the examples in Figure use different hijacking techniques. This is most notable on ALaMA-2-7BB which is a much larger model than the others. Of course, as expected, the more sophisticated attack on LLaMM also works on GPT-2 and Gemma. Additionally, the instruction-tuner 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 linguis- tic variations. FOUNTERFACT has 21, 919 samples, each of which are given by a tuple (p. O Q_, s. r). From each sample, we have a context prompt p with a true target answer o, (target true) and 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 (8 = "Fiffel Tower") and the prompt is categorized into relations r (for instance, other samples with the same relation 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 Pr[o] Pr(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_falee "asssaalle""" 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"" Figure 2a, we see that the efficacy score rises significantly after hijacking. Here, we prepend the hijacking sentence &o times for k = 0....,5 5 where yields the original prompt. We see that additional prepends increase the score further.
In the second scheme, we make use of the relation ID to prepend factually correct sentences. For instance, one can hijack the example above to "The Eiffel Tower is not located in Guam. The Biffel Tower is in". We test this bijacking philosophy on different relation IDs. In particular, Figure 2b reports hijacking based on relation ID P190 ("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 "French"", This suggests that context hijacking is more than simply the LLM copying tokens from contexts. Additional details and experiments for both hijacking schemes and for other relation IDs are in Appendix C.
These experiments show that context hijacking changes the behavior of LLMs, leading them to output incorrect tokens, without altering the factual meaning of the context. It is worth noting that similar fragile behaviors of LLMs have been observed in the literature in different contexts [Shi+23;; Pet+20; CSH22; Yor+23; PE21]. See Section 2 for more details.
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 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 synthetic memory retrieval task evaluate how the building blocks of LLMs, transformers, can solve it.
4 Problem setup
In the context of LLMs, fact or memory retrieval, can be modeled as a next token prediction problem. Given 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 different one, whether within the same space otherwise. In both scenarios, the

---

goal is to locate the closest pattern within the context when provided with a query (up to remapping
it's heteroassociativve.
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 any particular token in the context but rather 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 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.
4.1 Latent concept association
We propose a synthetic prediction task where for each output token y, tokens in the context (denoted by x) are sampled from conditional distribution given y- Tokens that similar to y will be favored to appear more in the context, except for itself. The task of latent concept association is to successfully retrieve the token y given samples from p(xly). 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₁- These could be viewed as semantic concept variables. Let Z = (Zi,... Zm) be the corresponding random vector, 2 be its realization, and Z be the collection of all latent binary vectors. For each latent vector z, there's one associated token E [V] {0, ..., V - 1} where V is the total number of tokens. Here we represent the tokenizer as where u(z) = l. 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 's 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 = 2mm
Under the latent concept association model, the goal is to retrieve specific output tokens given partial information in the contexts. This modeled by the latent conditional distribution:
p(z|2*) = wT (2|====) + (1 (11-w) Un)f((
where
(exp(-Du(z,
2 E N(()2"),
(2|2²)) x 10
z $ (z*)).
Here Dy is the Hamming distance, (=*))) a of 2 and ,3 > 0s the temperature parame- ter. The use of Hamming distance draws a parallel with the notion of distributional semantics in natural language: "a word is characterized by the company keeps" [Fir57]. In words, p(z|z*) says that with probability 1 1- w, the conditional distribution uniformly generate random latent vectors and with prob- ability w. the latent vector is generated from the informative conditional distribution π T22(2*) where the support of the conditional distribution is N(²*). Here, represents the informative conditional dis- tribution that depends on 2 whereas the uniform distribution is uninformative and can be considered as noise. The mixture model parameter w determines the signal to noise ratio of the contexts.
Therefore, for any latent vector 2* and its associated token, one can generate context token words with the aforementioned latent conditional distribution:
Uniformly sample latent vector 2*
For I 1,.., ...,, L 1, sample 2 p(2|2") and f2 ((1).
For ( L, sample = ( 2²*) and tr = ((2).
Consequently, we have x = (41, (,) and y = ((="). The last token in the context 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 D4 be the sampling distribution to generate (z, y) pairs. The conditional probability of y given æ is given by p(y(z).
5

---

With slight abuse of notation, given token t E [V], we define (tt) = we also define Du t') = (
for any pair tokens 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:
(A (x))()))))
where is the cross entropy loss with softmax. The error rate latent concept association is defined by the following:
Ro (f) = =P (ma)))
(argmax f(x) (arag=a//argaaxxxx #
And the accuracy Rp (f).
4.2 Transformer network architecture
Given a context x = (41,--,44) (T) which consists of I. tokens, we define E {0, 1) to be its
one-hot encoding where V is the vocabulary size. Here use X represent the one-hot encoding function (i.e., x() = X). Similar to [LLR23; Tar+23a; Li+24], we also consider simplified one-layer transformer model without residual connections and normalization:
(²(x) = WE WWyttn(Wex(0))
(4.1)
:
where
autn(U) = Uo (WxU)
Vdn
WK E xd is the key matrix, and Wq E Rda xd is the query matrix and dg is the attention head size. a LXL + (0, 1) is the column-wise softmax operation. Wy E [gd is the value matrix and WE E RddV 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 only one relevant for latent concept association. For convenience, we also use h(x)) to mean (attnWW>> xx)))) .1* which is the hidden representation after attention for the last position, and f! (z) to represent the logit for output token t.
5 Theoretical analysis
In this section, theoretically investigate how a single-layer transformer 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.
5.1 Hypothetical associative memory model
In this section, show that simple Angle-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 (t) be the 1-th column of the embedding matrix WE- In other words, this is the embedding for token t. Given a token f, define Nt( to be the subset of tokens whose latent vectors are only 1 Hamming distance away from f's latent vector: NG (t) {t' D (t', t)) 1] (t). For any output token I, Ni (1) contains tokens with the highest probabilities to appear in the context.
The following theorem formalizes the intuition that 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. 2 3. 1, 1, and (()) = V {t). Then for any & >> there exists a transformer model given by (4.1) that achieves error i.e. RD (fb) given sufficiently large context length L.

---

More precisely, for the transformer in Theorem 1, we will have WK = 0 and Wq = 0. Each row of WE orthogonal each other and normalized. And Wv is given by
W = Σ We(()(
Σ WE(()))
(5.1)
rE|V]
MENi()
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(xly) can lead to better recall rate. On the other hand, if contexts are modified to contain more samples from p(xlū) where ÿ * y. then is likely for transformer output the wrong token. This is similar context hijacking (see Section 5.5). The construction of the value matrix is similar to the associative memory model used in [Bic+24; CSB24J, but in our case, there is explicit one-to-one input and output pairs stored as memories. Rather, a combination of inputs are mapped to single output.
While the construction in Theorem is just one way that single-layer transformer can tackle this task, it turns out empirically this construction of Wy is close to the trained Wy, even in the noisy case (w # 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, 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 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 needed. This is discussed in the following sections.
5.2
On the role of the value matrix
The construction in Theorem relies on the value matrix acting as associative memory. But 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 (ar, y), the latent representation after self-attention h(x) must live in the polyhedron Sy to be classified correctly where Sy defined as:
Sy (u: (WE(v) W
where & [V] {y}}
Note that, by definition, for any two tokens y and J. S, n S3 = 0. 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 Conn(WF((0,,... W (IVI-1
In other words, for any pair (z, 3) to be classified correctly, hxx) must live in the intersection of Sy and CV. Due to the stochastic nature of æ, 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 Suppose the data generating process follows Section 4.1 where m 2 3, ((()) (t' (n(tt (((,,t)) 1}. = For any single layer transformer given by (4.1) where each row of We orthogonal each other and normalized, if WW is constructed as in (5.1), then the error rate is O. Wf is the identity matrix, then the error rate is strictly larger than 0.
and
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 2 3, w = 1 and N((1 = V (t). For any single layer transformer given by (4.1) 1) with Wy being the identity matrix, if the cross entropy loss is minimized so that for any sampled pair (I,4),
p(y]z) = +(ylæz) = softmax (f)fmaa(ff((())) (x))
there exists a > and such that for two tokens t + t',
(WE(t), (We()),W((t)) -aDD -a, (t,
7

---

5.3 Embedding training and geometry
The 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 constitute such a basis. Empirically, in Section 6.2, we observe that with overparameetrrization (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 close relationship between the inner product of embeddings for two tokens and the Hamming distance of these tokens (see Figure 3b and Figure D.5 in Appendix D.2). Approximately, we have the following relationship:
(We(t), WE((') = bo
(aD (1a0)))--(mmm)
1=1
(5.2)
t+t
for any two tokens & and l' where bn > and > 0. One can view this combination of the embedding geometry under Gaussian initialization and the geometry when Wv 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 bo, b and which is articulated by the following theorem.
Theorem 4 (Informal). Following the same setup in Theorem but embeddings obey (5.2), then under certain conditions on a, and if bo and context length L are sufficiently large, the error rate can be arbitrarily small, iee. Ry (f²)
for any 0 1.
The formal statement of the theorem and its proof given in Appendix B (Theorem 8).
Notably, this embedding geometry also implies a low-rank structure. Let's first consider the special case when bn 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 = bo and ((()) = V (4), then rank(W S 2.
When bn b, the embedding matrix will not be strietly 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 12 in Appendix D.4).
5.4 The role of attention selection
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 2", (N(2") = {2 2 21} {2*}.
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:
Ue. (Wx k W (T)) (WoWe(t))/Vd.
Lemma Suppose the data generating process follows Section 4.1 and N(2") {2 #1} } Given the last token in the sequence LI, then
Vur,e =(ft)= = Ve(Ft) W (a.p.We(t) P. De. We(tt)
where for token t, 04 E1= 1|tt- t] and Pe is the normalized attention score for token t.
8

---

Quay Takono
frod
24
Concepts
Hamming
Say SaUOTS
training
(b) Emhedding
(c)
Figure 3: Key components the single-layer transformer working together on the latent concept association problem. (a) Fixing the value matrix Wy as the identity matrix results lower accuracy compared to training
Wy. The figure reports average accuracy for both fixed and trained Wy with I. 64. (b) When training in the underparameterized regime, the embedding structure is approximated by (5.2). The graph displays the average inner product between embeddings of two tokens against the corresponding Hamming distance between these tokens when m (c) The self-attention layer can select tokens within the same cluster. The ligure shows average attention score heat map with m 8 and the cluster structure from Section 5.4.
Typically, oa is larger when token t and lr. belong to the same cluster because tokens within the same cluster tend to co-occur frequently. As result, the gradient contribution the unnormalized attention score is usually larger for tokens within the same cluster.
5.5 Context hijacking and the misclassification of memory recall
In light of the theoretical results on latent concept association, a natural question arises: How these results connect 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(ally) Thus, adding or changing tokens in the context 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 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 and Theorem and empirically verified in Appendix D.6. Such result also related to reverse context hijacking (Appendix C) where prepending sentences including true target words can improve fact recall rate.
that it resembles different
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).
2. Training embeddings is crucial in the underparameterizee 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 experi- mental details and results can be found in Appendix D.
6.1 On the value matrix Wy
In this section, we study the necessity of the value matrix Wy and its structure. First, we conduct ex- periments to compare the effects of training versus freezing Wy as the identity matrix, with the context lengths L set to 64 and 128. Figure 3a and Figure D.I show that when the context length is small, freez- ing Wy can lead to significant decline in accuracy. This is inline with Lemma and validates it in general setting, implying the significance of the value matrix in maintaining high memory recall rate.

---

Next, we investigate the degree of alignment between the trained value matrix Wy and the con- struction 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.I). 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 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-initializee value matrices, Figure D.3 shows that the constructed ones have, on average, smallest principal angles with the trained ones.
6.2 On the embeddings
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 the dimensions. Figure D.4 clearly shows that when the dimension is smaller than the vocabulary size (d < V), embedding training is required. is not necessary in the overparameterized regime (d > V), partially confirming Theorem because if embeddings initialized from a lige-diiensional multi-variate Gaussian, they are approximately orthogonal to each other and have the same norms.
The next question is what kind of embedding structures are formed for trained transformers in the underparamerized regime. From Figure 3b and Figure D.5, it is evident that the relationship between the average inner product of embeddings for two tokens and their corresponding Hamming distance roughly aligns with (5.2). Perhaps surprisingy. if we plot the same graph for trained transformers with a fixed identity value matrix, the relationship is mostly linear as shown in Figure D.6, confirming theory (Theorem 3).
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 WE. 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.
6.3 On the attention selection mechanism
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 and Figure D.7 clearly show that the self- attention select tokens in the same clusters. This suggests that attention can filter out noise and focus on the informative conditional distribution и. We extend experiments to consider cluster structures that depend on the first two latent variables (detailed construction can be found in Appendix D.3) and Figure D.8 shows attention pattern as expected.
7 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, devised 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 L.MMs, and can hopefully stimulate further work into interpreting and understanding the mechanisms by which LLMs predict tokens and recall facts.
10

---

Acknowledgments We thank Victor Veitch for insightful discussions that helped shape the initial idea of this work. We acknowledge the support of AFRL and DARPA via FA8750-23-2 0--75-033-2-01-- ONR via N00014-23-1-2388
NSF via IIS-1909816, 6, IIS- IIS-1956330, and NIH R01GM14046) We also acknowledge the support of the Robert H. Topel Faculty Research Fund at the University of Chicago Booth School of Business.
References
[Ach+23]
J. Achiam, S. Adler, S. Agarwal, L. Ahmad, I. Akkaya, F. L. Aleman, D. Almeida, Altenschmidt, S. Altman, S. Anadkat, al. "Gpt-4 technical report". arXiv preprint 0rXiv:2303.08874 (2023) (cit. on 17).
[AZL23a]
Z. Allen-Zhu and Y. Li. Physics of language models: part 3.1, knowledge storage and extraction, 2023, arXiv: 2309. 14316 [cs.CL) (cit. on p. 3).
[AZL23b]
Z. Allen-Zhu and Y. Li. Physics of language models: part 3.2, knowledge manipulation.
2023, arXiv: 2309. 14402 [Cs.CL] (cit. on p. 3).
[A71.24]
7. Allen-Zhu and Y. Li. Physics of language models: part 3.3, knowledge capacity scaling laws. 2024. arXiv: 2404.05405 [cs.Cl, (cit. on p.3).
[Apr+22]
G. Apruzzese, H. S. Anderson, S. Dambra, D. Freeman, F. Pierazzi, and K. A. Roundy. "real attackers don't 't compute gradients". : bridging the gap between adversarial ml research and practice. 2022. arXiv: 2212. 14315 [cs.Ch) (cit. on 3).
[Bai+24]
Y. Bai, Chen, H. Wang, C. Xiong, and Mei. "Transformers as statisticians: provable in-context learning with in-context algorithm selection".Advannee in neural information processing systems (2024) (cit. on p. 3).
[BYBOS95] R. Ben- Yishai, R. L. Bar-Or, and H. Sompolinsky. "Theory of orientation tuning in visual cortex." " Proceedings of the National Academy of Sciences (1995) (cit. on p. 2).
[Bie+24]
A. Bietti, V. Cabannes, D. Bouchacourt, H. Jegou, and L. Bottou. "Birth of a trans- former: a memory viewpoint". Advances in Neural Information Processing Systems (2024) (cit. pp. 3, 7).
[BP21]
T. Bricken and C. Pehlevan. "Attention approximates sparse distributed memory". Advances in Neural Information Processing Systems (2021) (cit. on pp. 2, 4).
[Bro+20]
T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss G. Krueger, T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei. Language models are few-shot learners. 2020. arXiv: 2005. 14165 (cs.CL) (cit. on p. 1).
(CDB23]
V. Cabannes, F. Dohmatob, and A. Bietti. "Scaling laws for associative memories". arXiv preprint arXiv:23 02984 (2023) (cit. on p. 3).
[CSB24]
V. Cabannes, B. Simsek, and A. Bietti. "Learning associative memories with gradient descent". arXiv preprint arXiv: 18724 (2024) (cit. on pp. 3, 7).
[Cha22]
F. Charton. "What is my math transformer doing?-three results on interpretability and generalization". arXiv preprint arXiv:2211.00/ (2022) (cit. on p. 3).
[Cho+24]
Shezan, V. Jain, and A. Chadha. "Breaking down the defenses: comparative survey of attacks on large language models". arXiv preprint arXiv:2403.04786 (2024) (cit. on p. 3).
A. G. Chowdhury, M. M. Islam, V. Kumar,
[Cov99]
T. M. Cover. Elements of information theory. 1999 (cit. on p. 9).
[CSH22]
A. Creswell, M. Shanahan, and I. Higgins. "Selection-iiferencce exploiting large
language models for interpretable logical reasoning". arXiv preprint arXiv:2205.097/2
(2022) (cit. on pp. 2-4).
[Dai+21]
D. Dai, L. Dong, Y. Hao, Z. Sui, B. Chang, and F. Wei. "Knowledge neurons in pretrained transformers" arXiv preprint arXiv:2104 08696 (2021) (cit. on pp. 1, 3).
[DCAT21]
N. De Cao, W. Aziz, and I. Titov. "Editing factual knowledge in language models". arXiv preprint 2rXiv::104 08164 (2021) (cit. on pp. 1,3).
[Dev83]
L. Devroye. "The equivalence of weak, strong and complete convergence in for kernel density estimates".. The Annals of Statistics 3 (1983) (cit. on p. 18).
11

---

[Ede+24]
B. L. Edelman, E. Edelman, S. Goel, E. Malach, and N. Tsilivis. "The evolution of statistical induction heads; in-context learning markov chains". arXiv preprint arXiv:2402. / 1004 (2024) (cit, on p. 3).
[Eth+21]
N. Elhage, N. Nanda, C. Olsson, T. Henighan, N. Joseph, B. Mann, A. Askell, Y. Bai, A. Chen, T. Conerly, et al. "A mathematical framework for transformer circuits". Transformer Circuits Thread (2021) (cit. pp. 3, 17).
[EI+24]
M Emrullah Ildiz, Y. Huang, Y. Li, A. Singh Rawat, and S. Oymak. "From self- attention to markov models: unveiling the dynamics of generative transformers", arXiv e-prints (2024) (cit. on onp.3. 3).
[Fel20]
V. Feldman. "Does learning require memorization? a short tale about a long tail". In: Proceedings of the 52nd Annual ACM SIGACT Symposium on Theory of Computing. 2020 (cit. on p. 3).
[FZ20]
V. Feldman and C. Zhang. "What neural networks memorize and why: discovering the long tail via influence estimation". Advances in Neural Information Processing Systems (2020) (cit. on p. 3).
[Fir57]
J. Firth. "A synopsis of linguistic theory, 1930-1955". -troiis in linguistic analysis
(1957) (cit. on p. 5).
[Gar+22]
S. Garg, D. Tsipras, P. S. Liang, and G. Valiant. "What can transformers learn in- context? case study of simple function classes". Advances in Neural Information Processing Systems (2022) (cit. on p. 3).
[Gei+21]
A. Geiger, H. Lu, T. Icard, and C. Potts. "Causal abstractions of neural networks". Advances in Neural Information Processing Systems (2021) (cit. on p. 3).
[Gei+22]
A. Geiger, Z. Wu, H. Lu, J. Rozner, E. Kreiss, T. Icard, N. Goodman, and C. Potts. "Inducing causal structure for interpretable neural networks". In: International Con- ference on Machine Learning. PMLR. 2022 (cit. on p. 3).
[Gei+24]
A. Geiger, Z. Wu, C. Potts, T. Icard, and N. Goodman. "Finding alignments between interpretable causal variables and distributed neural representations" In: Causal Learning and Reasoning. PMLR. 2024 (cit. on p. 3).
[Gev+23]
M. Geva, J. Bastings, K. Filippova, and A. Globerson. "Dissecting recall of factual associations auto-regressive language models". arXiv preprint urXiv:2304. 14767
(2023) (cit. p. 3).
[Has+24)
P. Hase, M. Bansal, B. Kim, and A. Ghandeharioun. "Does localization inform edit- ing? surprising differences in causality-based localization vs. knowledge editing in language models". Advances in Neural Information Processing Systems (2024) (cit. on p. 3).
[Hen+23]
T. Henighan, S. Carter, T. Hume, N. Elhage, R. Lasenby, S. Fort, N. Schiefer, and C. Olah. "Superposition, memorization, and double descent". Transformer Circuits Thread (2023) (cit. 3).
[Hop82]
J. Hopfield. "Neural networks and physical systems with emergent collective putational abilities." " Proceedings of the national academy of sciences 8 (1982) (cit. on p. 2).
[Hu+21]
E. J. Hu, Shen, Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen. Lora: low-rand adaptation of large language models. 2021. arXiv: 2106. 09685 [cs. CL] (cit, on p. 2).
[Hu+24a]
J. Y.C. Hu, P.-H. Chang, R. Luo, H.-Y. Chen, W. Li, W.-P Wang. and H. Liu. "Outlier-effcieet hopfield layers for large transformer-basee models". arXiv preprint arXiv:2404.0388 03888888888888888 (2024) (cit. on 2).
[Hu+24b]
J. Y.-C Hu, B.-Y. Chen, D. Wu, F. Ruan, and H. Liu. "Nonparametric modern hopfield models". arXiv preprint arXiv:2404. 039900 (2024) (cit. on p. 2).
[Hu+24c]
Y.-. Hu, Lin, Z. Song, and H. Liu. "On computational limits modern hopfield models: fine-grained complexity analysis". arXiv preprint 4rXiv:2402.04520 (2024)
(cit. on p. 2).
[Hu+24d]
J. Y.-C. Hu, D. Yang, D. Wu, C. Xu, B.-Y. Chen, and H. Liu. "On sparse modern hopfield model", Advances in Neural Information Processing Systems (2024) (cit. on
p. 2).
12

---

[JP20]
Y. Jiang and C. Pehlevan. "Associative memory in iterated overparamettrrzee sigmoid autoencoders". In: International conference on machine learning. PMLR. 2020 (cit, on p. 2).
[Jia+24]
Y. Jiang, G. Rajendran, P. Ravikumar, B. Aragam, and V. Veitch. "On the origins of linear representations in large language models". arXiv preprint WxXiv:2403.0380
(2024) (cit. p. 3).
[Jin+23]
T. Jin, N. Clement, X. Dong, V. Nagarajan, M. Carbin, J. Ragan-Kelley, and G. K. Dziugaite. "The cost of down-scaling language models: fact recall deteriorates before in-context learning". arXiv preprint arXiv:2310.044: 1..00660 (2023) (cit. on p. 3).
[KKM22]
J. Kim, M. Kim, and B. Mozafari. "Provable memorization capacity transformers".
In: The Eleventh International Conference on Learning Representations. 2022 (cit. on p. 3).
[Li+24]
Y. Li, Y. Huang, M. E. Ildiz, A. S. Rawat, and S. Oymak. "Mechanics of next token prediction with self-attention". . In: International Conference on Artificial Intelligence and Statistics. PMLR. 2024 (cit. on pp. 3, 6).
[LLR23]
Y. Li, Y. Li, and A. Risteski. "How do transformers learn topic structure: towards mechanistic understanding". In: International Conference on Machine Learning.
PMLR. 2023 (cit. on pp. 3, 6).
[Liu+23a]
Y. Liu, G. Deng, Y. Li, K. Wang, T. Zhang, Y. Liu, H. Wang. Y. Zheng, and Y. Liu. "Prompt injection attack against Imm-iteggrated applications". arXiv preprint wxXiv2236.054490 (2023) (cit. on 3).
[Liu+23b]
Y. Liu, G. Deng, Z. Xu, Y. Li, Y. Zheng, Y. Zhang, L. Zhao, T. Zhang, and Y. Liu. "Jailbreaking chatgpt via prompt engineering: an empirical study". arXiv preprint arXiv:2305. 13860 (2023) (cit. on p. 3).
[Liu+22]
Z. Liu, O. Kitouni, N. S. Nolte, E. Michaud, M. Tegmark, and M. Williams. "Towards understanding grokking: effective theory of representation learning". Advances in Neural Information Processing Systems (2022) (cit. on p. 3).
[LHI7]
I. Loshchilov and F. Hutter. "Decoupled weight decay regularization". arXiv preprint arXiv:1711 1.11.011100011 .00100 (2017) (cit. on p. 27).
[Luo+24]
H. Luo, J. Yu, W. Zhang, J. Li, J. Y.-C. Hu, X. Xin, and H. Liu. "Decoupled alignment for robust plug-and-pn -od--and-pay adaptation", arXiv preprint arXiv:2406.0157 01554 (2024) (cit. on
p. 3).
[MLT23]
S. Mahdavi, R. Liao, and C. Thrampoulidis. "Memorization capacity of multi-head attention in transformers". arXiv preprint arXiv:2306. 02010 (2023) (cit. on p. 3).
[Mak+24]
A. V. Makkuva, M. Bondaschi, A. Girish, A. Nagle, M. Jaggi, H. Kim, and M. Gastpar. Attention with markov: a framework for principled analysis of transformers via markov chains. 2024. arXiv: 2402. 04161 [cs.LG] LG] (cit. on p. 3).
[McG+23]
T. McGrath, M. Rahtz, Kramar, V. Mikulik, and S. Legg. "The hydra effect: emergent self-repair in language model computations". arXiv preprint arXiv:2307.15777 (2023)
(cit. on p. 3).
[Men+22]
K. Meng, D. Bau, A. Andonian, and Y. Belinkov. "Locating and editing factual associations in "apt..Avances. in Neural Information Processing Systems (2022) (cit. on pp. 14,25)..
[Men+23]
K. Meng, A. S. Sharma, A. Andonian, Y. Belinkov, and D. Bau. Mass-editing memory in a transformer. 2023. arXiv: 2210. 07229 [cs CL. (cit. on pp. 1)
[Mil+22]
B. Millidge, T. Salvatori, Y. Song, T. Lukasiewicz, and R. Bogacz. "Universal hop- field networks: general framework for single-shot associative memory models". In: International Conference on Machine Learning. PMLR. 2022 (cit. on pp. 2, 4).
[Mit+21]
E. Mitchell, C. Lin, A. Bosselut, C. Finn, and C. D. Manning. "Fast model editing at scale". arXiv preprint arXiv:2110. 10. 1309 (2021) (cit. on pp. 1, E. Mitchell, C. Lin, A. Bosselut, C. D. Manning, and C. Finn. "Memory-based" model editing at scale". In: International Conference on Machine Learning. PMLR. 2022 (cit. on pp. 3).
[Mit+22]
[Nan+23]
N. Nanda, L. Chan, T. Lieberum, J. Smith, and J. Steinhardt. "Progress measures for grokking via mechanistic interpretability". arXiv preprint arXiv:2301.052 (2023) (cit. on p. 3).
13

---

[NDL24]
E. Nichani, A. Damian, and J. D. Lee. How transformers learn causal structure with gradient descent. 2024. arXiv: 2402. 14735 [cs. LG] (cit. on p. 3).
[OIs+22a]
C. Olsson, N. Elhage, N. Nanda, N. Joseph, N. DasSarma, T. Henighan, B. Mann, A. Askell, Y. Bai, A. Chen, T. Conerly, D. Drain, D. Ganguli, Z. Hatfield-Dodds,
D. Hernandez, S. Johnston, A. Jones, J. Kernion, L. Lovitt, K. Ndousse. D. Amodei,
T. Brown, J. Clark, J. Kaplan, S. McCandlish, and C. Olah. In-context learning and induction heads. 2022. arXiv: 2209 11895 [ca. LG] (cit. on p. 3).
[OIs+22b]
C. Olsson, N. Elhage. N. Nanda, N. Joseph, N. DasSarma, T. Henighan, B. Mann,
A. Askell, Y. Bai. A. Chen, et al. "In-context learning and induction heads". arXiv preprint arXiv: 1885 (2022) (cit. on p. 3).
[PE21]
L. Pandia and A. Ettinger. "Sorting through the noise: testing robustness of information processing in pre-trained language models". arXiv preprint wxXiv:210911233 (2021) (cit. on pp. 2-4).
[PR22]
F. Perez and Ribeiro. "Ignore previous prompt: attack techniques for language models". arXiv preprint wrXiv:2211.09522 (2022) (cit. on p. 3).
[Pet+20]
F. Petroni, P. Lewis, A. Piktus, T. Rocktäschel, Y. Wu, A. H. Miller, and S. Riedel. "How context affects language models' factual predictions". arXiv preprint arXiv:2005. 0441:1005.046111 (2020) (cit. on pp. 2-4).
[Rad+19]
A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, al. "Language models are unsupervised multitask learners". OpenAl blog 8 (2019) (cit. on pp. 1-3).
[RBU20]
A. Radhakrishnan, M. Belkin, and C. Uhler. "Overparameterizeed neural networks implement associative memory". Proceedings of the National Academy of Sciences 44 (2020) (cit. on p. 2).
[Raj+24]
G. Rajendran, S. Buchholz, B. Aragam, B. Schölkopf, and P. Ravikumar. "Learning interpretable concepts: unifying causal representation learning and foundation models". arXiv preprint ar arXiv:2402.092.
(2024) (cit. on 3)
[Ram+20]
H. Ramsauer, Schäfl, J. Lehner, P. Seidl, M. Widrich, T. Adler, L. Gruber, M. Holzleitner, M. Pavlović, G. K. Sandve, et al. "Hopfield networks is all you need". arXiv preprint arXiv:2008. 022177 (2020) (cit. on pp. 4).
[Rao+23]
A. Rao, S. Vashistha, A. Naik, S. Aditya, and M. Choudhury. "Tricking llms into disobedience: understanding, analyzing, and preventing jailbreaks". arXiv preprint arXiv:2305.14965 (2023) (cit. on p. 3).
[Sak+23]
M. Sakarvadia, A. Ajith, A. Khan, D. Grzenda, N. Hudson, A. Bauer, K. Chard, and Foster. "Memory injections: correcting multi-hop reasoning failures during inference in transformer-based language models". arXiv preprint 230900300
(2023) (cit. on p. 3).
[Seu96]
H. S. Seung. "How the brain keeps the eyes still". Proceedings of the National Academy of Sciences 23 (1996) (cit. on p. 2).
[SMR23]
M. Shanahan, K. McDonell, and I.. Reynolds. "Role play with large language models". Nature 7987 (2023) (cit. on p. 3).
[She+24]
H. Sheen, S. Chen, T. Wang, and H. H. Zhou. "Implicit regularization of gradient flow on one-layer softmax attention". arXiv preprint arXiv:2403.08869 (2024) (cit. on p. 3).
[Shi+23]
F. Shi, X. Chen, K. Misra, N. Scales, D. Dohan, E. H. Chi, N. Schärli, and D. Zhou. "Large language models can be easily distracted by irrelevant context". In: Interna- tional Conference on Machine Learning. PMLR. 2023 (cit. on pp. 2-4, 17).
[Si+22]
W. M. Si, M. Backes, J. Blackburn, E. De Cristofaro, G. Stringhini, S. Zannettou, and Y. Zhang. "Why so toxic? measuring and triggering toxic behavior in open-domain chatbots". In: Proceedings of the 2022 ACM SIGSAC Conference on Computer and Communications Security. 2022 (cit. on p. 3).
[Ska+94]
W. Skaggs, J. Knierim, H. Kudrimoti, and B. McNaughton. "A model of the neural basis of the rat's 's sense of direction". Advances in neural information processing systems (1994) (cit. on 2.
[SS22]
J. Steinberg and H. Sompolinsky. "Associative memory of structured knowledge". Scientific Reports 1 (2022) (cit. on p. 2).
14

---

[Tar+23a]
D. A. Tarzanagh, Y. Li, C. Thrampoulidiss and S. Oymak. "Transformers as support vector machines". arXiv preprint arXiv:2308. 16898 (2023) (cit. on pp. 3, 6).
[Tar+33b]
D. A. Tarzanagh, Y. Li, X. Zhang, and S. Oymak. "Margin maximization in attention mechanism". arXiv preprint arXiv:2306. 13596 (2023) (cit. on p. 3).
[Tea+24]
G. Team, T. Mesnard, C. Hardin, R. Dadashi, S. Bhupatiraju, S. Pathak, L., Sifre, M. Rivière, M. S. Kale, J. Love, et al. "Gemma: open models based on gemini research and technology". arXiv preprint arXiv:2403. 08295 (2024) (cit. on pp. 2. 3, 6).
[Tia+24]
B. Tian, S. Cheng, X. Liang, N. Zhang. Hu, K. Xue, Y. Gou, X. Chen, and H. Chen. "Instructedit: instruction-based knowledge editing for large language models". arXiv preprint arXiv:2402. 16123 (2024) (cit. on p. 3)).
[Tia+23a]
Y. Tian, Y. Wang, B. Chen, and S. S. Du. "Scan and snap: understanding training dynamics and token composition in I-layer transformer". Advances in Neural Infor- mation Processing Systems (2023) (cit. on p. p)
[Tia+23b]
Y. Tian, Y. Wang. Z. Zhang, B. Chen, and S. Du. "Joma: demystifying multilayer trans- formers via joint dynamics of mip and attention". arXiv preprint arXiv:23
(2023) (cit. on p. 3).
[Tou+23]
H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov,
S. Batra, P. Bhargava, S. Bhosale, et al. "Ilama 2: open foundation and fine-tuned chat models". arXiv preprint arXiv:2307.09288 (2023) (cit. on pp. 2)..
[Vas+17]
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, N. Gomez, L. Kaiser, and
I. Polosukhin. "Attention is all you need", Advances in neural information processing systems (2017) (cit. on p. 2).
[Wan+23a]
B. Wang. W. Chen, H. Pei, C. Xie, M. Kang, C. Zhang, C. Xu, Z. Xiong, R. Dutta, R. Schaeffer, et al. "Decodingtrust: a comprehensive assessment of trustworthiness in gpt models". arXiv preprint arXiv: /698 (2023) (cit. on p. 3).
[Wan+23b]
J. Wang, X. Hu, W. Hou, H. Chen, R. Zheng, Y. Wang, L. Yang, H. Huang, W. Ye, X. Geng, al. "On the robustness of chatgpt: an adversarial and out-of-distribution perspective". arXiv preprint arXiv: 12095 (2023) (cit. on p. 3).
[Wan+22]
K. Wang, A. Variengien, A. Conmy, B. Shlegeris, and J. Steinhardt. "Interpretability
in the wild: a circuit for indirect object identification in gpt-2 small", arrivv preprint
arXiv:2211.00 12211.0000 2.00059 (2022) (cit. on p. 3).
[Wu+24a]
D. Wu, J. Y.-C. Hu, T.-Y. Hsiao, and H. Liu. "Uniform memory retrieval with larger capacity for modern hopfield models". arXiv preprint arXiv:2404.03827 (2024) (cit.
on p. 2).
[Wu+23]
D. Wu, Y.-C. Hu, W. Li, B.-Y. Chen, and H. Liu. "Stanhop: sparse tandem hopfield time series prediction". arXiv preprint arXiv:2312.17344
model for memory-enhanee (2023) (cit. on p. 2).
[Wu+24b]
Z. Wu, A. Geiger, T. Icard, C. Potts, and N. Goodman. "Interpretabbiiity at scale: identifying causal mechanisms in alpaca". Advances in Neural Information Processing Systems (2024) (cit. on p. 3).
[Xie+21]
S. M. Xie, A. Raghunathan, P. Liang, and T. Ma. "An explanation of in-context learning as implicit bayesian inference". arXiv preprint arXiv:2111.02088 (2021)
(cit. on p. 3).
[Xu+23]
X. Xu, K. Kong, N. Liu, L. Cui, D. Wang, J. Zhang, and M. Kankanhalli. "An Ilm can fool itself: a prompt-basee based adversarial attack".. arXiv preprint arXiv:2310. 13343 (2023) (cit. 3).
[Yor+23]
O. Yoran, T. Wolfson, O. Ram, Jnd Berant. "Making eetrieval-augmenteee language models robust to irrelevant context". arXiv preprint arXiv:2310.0155. (2023) (cit. on
pp. 2-4).
[Yu+24]
Yu, H. Luo, J. Yao-Chieh, W. Guo, H. Liu, and X. Xing. "Enhancing jailbreak attack
against large language models through silent tokens". arXiv preprinl arXiv:2405.20653 (2024) (cit. on p. 3).
[Zha+22]
Y. Zhang, A. Backurs, S. Bubeck, R. Eldan, S. Gunasekar, and T. Wagner. "Unveiling transformers with lego: a synthetic reasoning task". arXiv preprint arXiv:2206 2233002 (2022) (cit. on p. 3).
15

---

[Zha+23]
Z. Zhang, M. Fang, L. Chen, M.-R. Namazi- and Wang. "How do large language models capture the ever-changing world knowledge? a review of recent advances". arXiv preprint arXiv:2310.07343 (2023) (cit. on 3).
[Zha23]
J. Zhao. "In-context exemplars clues to retrieving from large associative memory". arXiv preprint arXiv:2311.0449 1.0449 (2023) (cit. on pp. 2, 4).
[Zhe+23]
C. Zheng, L. Li, Q. Dong, Y. Fan, Z. Wu, J. Xu, and B. Chang. "Can we edit factual knowledge by in-context learning?" arXiv preprint arXiv:2305.12744 (2023) (cit. on
p. 3).
[Zho+24]
Z. Zhong, Z. Liu, M. Tegmark, and J. Andreas. "The clock and the pizza: two stories in mechanistic explanation of neural networks". Advances in Neural Information Processing Systems (2024) (cit. p. 3).
[Zhu+23]
K. Zhu, Wang, Zhou, Z. Wang, H. Chen, Y. Wang, L. Yang, W. Ye, N. Z. Gong,
Y. Zhang, et al. "Promptbench: towards evaluating the robustness of large language models adversarial prompts".. arXiv preprint arXiv:2306.0455 04528 (2023) (cit.
p. 3).
[Zou+23]
A. Zou, Z. Wang, N. Carlini, M. Nasr, J. Z. Kolter, and M. Fredrikson. Universal and transferable adversarial attacks on aligned language models. 2023, arXiv: 2307. 15043 [cs.CL] (cit. on 3).
16

---

A Limitations
The context hijacking experiments only conducted open-source models and not on commer- cial models like GPT-4. Nevertheless, even in the official GPT-4 technical report [Ach+23], there is an example similar 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 "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 music and answers Presley. In fact, easily distracted by contexts in use cases other than fact retrieval such as problem-solving [Shi+233]. So we reasonably suspect that similar behavior still exists in larger models but harder to exploit. On the other hand, the theoretical section only focuses 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].
known that LLMs can be
B Additional Theoretical Results and Proofs
B.1 Proof's for Section 5.1
Theorem be stated more formally as follows:
Theorem 7. Suppose the duta generuting process follows Section 4.1 where i 3, w = 1, and ((t) V {t). Assume there exists single layer transformer given by (4.1) such that a) WK
and Wq = b) Each row of We g is orthogonal to each other and normalized, and c) Wy is given by
Wv =
Σ WE(
Σ We(d)")
iE[V]
JENi()
Then if > max 100m log(3/e)
80m (W(x)
} for any then
(exp(- expl
31)
(exp( 5) --xp(--5p)
RD 19)88.
where 0 <E<1.
Proof. First of all, the error is defined be:
RD (ft) P (( Jargmax (argmaxx f*(x) # yl
Py s/y [argmax (x) & u]
Let's focus on the conditional probability P. Jargmax (x) + v). By construction, the single layer transformer model has uniform attention. Therefore,
h(x)=
Σ
atWa(1)
(EN(9)
where a; - Et= 1fth- =t} which is the number of occurrence of token in the sequence. By the latent concept association model, we know that
p(ily) (i w)/
where Z = EiEN(y) Thus, the logit for token y is
(1) = Σ
SEN((M)
And the logit for any other token y is
(x) = Σ
17

---

For the prediction to be correct, we need
(c) >0
y
By Lemma [Dev83], we know that for all A E (0, 1). if Wiyl
4² 20 we have
P( max lay p(i > A) (e)(
P
Σ
aa -(iy)
A) - 3
(EN(s)
Therefore, if L 2 max
(ng(3/c)
(0)
}. then with probability at least 1 - E, we have,
A²
A
max (EN(w)
04
(x)
Σ
ai Σ aj
(ENily)
JENi(0)
Σ ai
Σ
p(ily)
+
Σ p(ily)
(ENily)
(ENi(
iENi(u)
Σ
p(ily)
+
Σ
p(jly)
Σ Oj
jEN((
jEN((0)
JENi
Σ (ENi(y)
p(ily)
Σ
p(ily)
2m4
jENi(
2
2 exp(
exp(
2mA
Note that because of Lemma 10, there's no neighboring set that is the superset of another. Therefore as long as A
oxp(- -xx(- 3)
2m
(x) 1/ (x)
for any ÿ-
2 (N) (exp( 2) ) expl ">) (expl - 3)1 for any y. then
100m 2 log(3/e)
Finally, if L > max
P Patp f /argmax fz + S
And
RD (ft) = P (u (argmax fr(xx) # vl PaPy /argmax fL((x) + y) Se
B.2 Prool's for Section 5.2
Lemma 2. Suppose the data generating process follows Section 4.1 where m 2 3, w and
(t = (t' (t((,t'))= 13. For any single layer transformer given by (4.1) where each of
WE is orthogonal to each other and normalized, if Wy is constructed as in (5.1), then the error rate
Wy is the identity matrix, then the error rate strictly larger than 0.
Proof. Following the proof for Theorem 7, let's focus on the conditional probability:
IP Palp Jargmax f + v]
By construction, we have
h(x)= = Σ (&))
(ENT(
18

---

where 1 E 1ftx = e} which is the number of occurrence of token in the sequence.
Let's consider the first where Wy is constructed in (5.1). Then know that for other token ÿ #
Σ
Σ
Σ
-
iENi(D)
vEV1(0)
SEN1
By Lemma 10, we have that for any token ÿ a
(x) - (+)
Therefore, the error rate is always O.
Now let's consider the second case where Wy is the identity matrix. Let j be a token in the set (v).. Then there is non-zero probability that context I contains only j. that case,
h(x:) WF(i)
However, we know that by the assumption on the embedding matrix,
( (Wr(w) (=())) = -IW=(0)| <
This implies that there's non zero probability that y is misclassified. Therefore, when Wy the identity matrix, the error rate is strictly larger than O.
D
Theorem 3. Suppose the data generating process follows Section 4.1 where m 2 w and N(t) v {t). For any single layer transformer given by (4.1) with Wy being the identity matrix, if the cross entropy loss is minimized so that for any sampled pair (x, y).
plyla) = -(ul = softmax(s (x)
there exists > 0 and such that for two tokens t + t',
(E())) = -a
Proof. Because for any pair of (ar, y). the estimated conditional probability matches the true condi- tional probability. In particular, let's consider two target tokens y1, y2 and context x = (t;..., tt) for some token t, such that p(rl1)) 0 and p(zly2) O, then
p(u1k)) (alw))p((1)) p(=lv1) p(=[v1)
exp((WE(v)) WE()2))
(( ((22)) p(zly2) p(x|y²)
The second equality is because p(v) is the uniform distribution. By our construction,
p(rlv)) p((l)1)4
W h(()) exp((Wz((1) - ( Wx((z))
p 12)
p(t,)12)
By the data generating process, we have that
L (D (1) DH (1), = (We(11) We
Let 4 = ya such that ys + yı, Us + y2, then
L Du (ys. mn) WE(n)) We(us))We(ss)
L Du (ys. y2) WE(v2) ( (w(vs)
3
3
For simplicity, let's define
(11, 92) = DH( v2) - WE(n)) We(ve)
3
Therefore,
V (ys, v1) = v (ya, y2)
Now consider five distinct labels: y1, U2 U3. V1 Us. We have,
V/ (y3, v1) v (33, v2) v (vs, (2) (ys, y6)
19

---

In other words, (33..8) 91) (34,, Us) for arbitrarily chosen distinct labels y1. U4. V5- Therefore, (t, t') is constant for + t'.
For any two tokens + t',
D Waatt))--a())Wa())> W((() =
Thus,
WE(t)" = - +
B.3 Proofs for Section 5.3
Theorem can be formalized as the following theorem.
Theorem 8. Following the same setup as in Theorem but embeddings follow (5.2) then if expl -xxp 3)
A 0,00
max
25 log(3/e)
20 N(w)]] for any y, and
2m
A
2 2exp(H) (t)
0<a<
(IV) 2)m²
and
(b - (b-a)A1 1V abm? exp( (-)) + 1V 2 -2)m²
a(m 2)m + D
bu max
+b,
exp(-1 +1) exp(- 2m
1 am2 exp(
have
R (ft) SE
where 0 <
Proof. Following the proof of Theorem 7, let's also focus on the conditional probability
Palp Jargmax (x) A y]
By construction, the single layer transformer model has uniform attention. Therefore,
h(x)
Σ (;W
(EN())
where a; = + E 1f4x = i} which is the number of occurrence of token in the sequence. For simplicity, let's define ay 0 such that
h(x) Σ (())
EMM
Similarly, we also have that max we have,
25 log(3/e)
20l A
}, then with probability at least E,
A
max EE(V]
las - p(il)) s A
Also define the following:
dx(x)) Σ WE Σ a.WE(i))
JEN((R)
<E[V]
WE(T) (g())
Thus, the logit for token y is
V1-1
(+) Σ
(x())x((x)
K=U
20

---

Let's investigate (x).). By Lemma 9,
xx(x) E atl Σ
WE()
(E(V] (K)
= (bo - b)
Σ
a + Σ
aj
as(-a(m- - (k, (b (b)
JENi
Thus, for any kı, l22 E [V]
(x)(2) - -x, = (bo - b)(
Σ Oj
Σ ajz)
(EEN((()
Σ
a.a(m (D i) (H1,)
Because- -m < (k2,2 - (k1, i)
S m, we have
(bo -))
Σ Σ
aj
azz) a(m 2)m
(E4)
12 (12)
I 0k, (x) - 0k (2) S
(bo 6)
Σ
a/ E
a3,) + a(m 2)m
(ks))(
(k2))
For prediction to be correct, we need
max 11 (6))
8
This also means that
-
Σ
(w) - un (z)) 0
max
8
k=
One can show that for any k if i then
1aff-1(p6 (b) = - (0))(1111) (()) where @ means bitwise XOR,
vx(y) - un (o) -
(B.1)
First of all, if = y, then k = y. which means
(v( - uk (ū) = Ug (0) - v2(y) = bn + aDn(v. 0)
If l: # v. then (B.1) implies that
Du (k. v) Du (k: 1) DH (k. 0) - (1 v)
We know that DH (Ri, y) is the number of is in (1k) 8 (1)) and,
(()) = 1 (7)
Similarly,
(i) 8 = (0)
Therefore, (B.1) holds and we can rewrite f1((x) f (z) as
M
(x) - (1))))(())-
Σ (vx(v) 2/ (z)
k=
(bo b + aD (v. (b, 0g(x)))
+
Σ
a(D - Ok -))())))))))(())))
- (x(x))
21

---

We already know that by > b > 0 and >> O, thus, bu - b aD > O for any pair y, g.
We also want (5(x) Og(2)) to be positive. Note that
0y(x) 0g xr) (bo (expp( } exp( 2 - --2mA) a(m 2)m
exp( - 2) exp( and for some positive A1 > 0, bo needs to be large enough such that
We need A
2m
0y(x) - Φz(x) > A
which implies that
a(m 2)m + A +
by
(B.2)
exp( }) expl 1/} 2n
On the other hand, for le f U, V.
we have
(x(x) - (x(x) 2 (bo 6)( Σ an
E
0/22) a(m - 2)m
JI (M)
DOEN((()
2
2 (bg b)(-(m (m 1) exp(
-exp(
-2mA) -a(m 22m
8
+ expl - (/] -exp(- }) - a(m 2)m
2 (bo 6)( (1((mm - 1) exp(
3 - exp(
B
(bo b)m exp( 1 ) a(m 2)m
Then, we have
(x) (x) 2 (bu b+ [V (bu b)am² exp(p 1 3) + a² a² -2)m²)
V 2 am?
1
V - abm² exp( -xp(->
V a² 2)m
2 expl - b0 (b a A
2
2
The lower bound is independent of ÿ. therefore, we need it to be positive ensure the prediction is correct. To achieve this, we want
V
am exp(
2
which implies that
2 eep((1)
< (V 22)m²
(B.3)
And finally need
A) (b-a)
bb -
1V|-2 abm? exp(- 1)) + 1V1-2 a 2)m
bo
(B.4)
V1-2 am? exp(-
To summarize, i > 0, A, > 0,0 any y, and
exp( -)-exp( ()) L 2 max
25 lo/(/3
201N(y)]
} for
2m
A
A
2 ecxp(}) (1}
0<a< ([V) -2)m²
and
(b (b- (b-a)Ai 1V|-2 abm? exp 3) + 1V-2 am(m 2)m²
a(m- 2)m= + A
bo > max exp(
+b,
2)) - 2m
V 2
2 am? 2 exp(- ( 1
we have
Rp ( Se
where 0 < < 1.
22

---

Lemma 5. If emmeedding follow (5.2) and b = by and ((t) = V {t), then rank(Wæ) S +
Proof. By (5.2), we have that
= (i)
Therefore,
(WE) -aD + 611
Let's first look at DH which has rank at most m + 1. To see this, let's consider a set of m + 1 tokens: feo, Cl,-,., Cm V where eg 2 Here eo is associated with the latent vector of all zeroes and the latent vector associated with eg has only the k-th latent variable being 1.
On the other hand, for any token i, we have that,
Σ
Ck (+)==1
In fact.
Dz(1)
Σ
(Dnlew) DOnle)-DD(eeo)) (alea)-Dir(eee + (e(eo)
ka 11)=m
where (i) the i-th row of DH, and for each entry j DH (i), we have that
(n(ij)) Σ
(Du(eu.d) Di((eo.)) + (H.eo.
Ri (6))=1
This is because
+1 if if)= =
De(ek.)) -(al:.))DD((e..)))
-1 1f)x
Thus, we can rewrite DH (i, as
Du(i.)) Σ
1,, = 111) + Du((o)))
A: 11)==1
= 000---0,00000000000000000--000-- -1, = =1)1
A=
+ Σ
- -,, = +| 11 = 1,0 = 1))
k=1
m
Σ 1 (1)x = 1,, -0) -0 = 0, L. -1)) =1]
Dn(iij)
Therefore, every row of DH can be written as linear combination of (Dm(eo), (ere),,.... Dm(em)). In other words, DH has rank at most + 1.
Therefore,
rank((We) (W)) rank(We) S m +2
Lemma 9. Let 2(0) and (11) be two binary vectors of size m where m 2 2. Then,
Σ
D2(2,z(") = (m D (z (2))) + m
8: (2 x)==1
23

---

Proof. For 2 such that (2, 20)) ) = 1, we know that there are two cases. Either differs with 2
on a entry but agrees with 21) on that entry or differs with both 20)) and 2(1)
For the first case, we know that there are D=(2((), 20)) such entries. In this case, Dz(2, (=((""), 2(1) 1. For the second case, (2, 2(1) = Du (=(") =0)) +
Therefore.
Σ Dz(2,20)
Dr(z , (Dm( 20))) - 1) + (m Dr(= 2 ()=((( 20 +1)
(m (2 (1))) + m
Lemma 10. If 2 3 and N(1) V (4) then Ni(1) g NG(() for any 1, E [V].
Proof. For any token f, MN (t) contains any token t' such that Dy (t, t') by the conditions. Then given a (t), one can uniquely determine token This is because for the set of latent vectors associated with (t), at each index, there could only be one possible change.
B.4 Proofs for Section 5.4
Lemma 6. Suppose the data generating process follows Section 4.1 and N(=*) {=
21 { Given the last token in the sequence Lt, then
V Vunee e(ft) (WE W (a.p.We(t) P. (e(ttt))))
1=1
where for token I, 0z -f= 1[ = 1] and Dr is the normalized attention score for token t.
Proof. Recall that,
+x(x)) We Wyattt aiyaatnWFxxl (Wex(x))
exp(u (t)
WE Wv
Z
1=1
where Z is a normalizing constant.
Define Pr. oxp("e,t)) Then have
(*(()) -W= W Wv
Σ pr, Wa(k)
Note that if f f then,
Opu De, (1 - -pr)
Jurst
Otherwise,
Opu.
Pr.Pr
Ju,,e.
By the chain rule, we know that
I. 1[th WE(()) - E (e.ttWW((t))
I
V Vunnz e(*)) VE(FL)
1=1
Therefore,
V Vu,,, Uu, = (W) (aspaWW() - P. Σ p. We((t)
1=1
where a; -E 1[t = t].
24

---

C Additional experiments - context hijacking
In this section, show the results of additional context hijacking experiments the COUNTER FACT dataset [Men+22].
Reverse context hijacking In Figure 2a, we saw the effects of hijacking by adding in "Do not think false)." )." to each context. Now, we measure the effect of the reverse: What we prepend "Do not think of [target_true)."
of [targetffalse)""
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, we see in Figure C.I.
Prepending 'Do not think of (target true).'
0.30
openal-commmuityyggtz
google/gemma-26
0.25
google/gemma-2b..
meta-llamma/llmm--22bo 2-7-hh
Efficacy Score
0.20
0.15
0.10
0.05
5
Number of prepends
Figure C.I: Prepending 'Do not think of [target_true]." can increase the chance of LLMs to output correct tokens. This figure shows efficacy score versus the number of prepends for various LLMs on the COUNTERFACT dataset with the reverse context hijacking scheme.
Hijacking based on relation IDs We first give an example of each of the relation IDs we hijack in Table 1.
Table 1: Examples contexts in Relation IDs from COUNTERFACT
RELATION
CONTEXT P
TRUE TARGET O.
FALSE TARGET O_
P190
Kharkiv a twin city The native language of Anatole France Hank Aaron professionally plays the sport Kalamazoo County can be found in
Warsaw French baseball Michigan
Athens English
P103
P641
baskethall
P131
Indiana
Table 2: Examples of hijack and reverse hijack formats based Relation IDs
RELATION
CONTEXT HUACK SENTENCE
REVERSE CONTEXT HIACK SENTENCE
P190
twin city of (subject) not [target false)
The twin city of (subject] [target true] [subject] can speak [target true]
P103
(subject) cannot speak (target false]
P641
[subject) does not play (target false] (subject) is not located in (target_false)
[subject] plays (target true)
P131
(subject) is located in (target_true)
Similar Figure 2b, we repeat the hijacking experiments where we prepend factual sentences generated from the relation ID. We use the format illustrated in Table 2 for the prepended sentences.
25

---

Hljacking based on P103
Hijjcking based on P131
0.6
Efficacy Score
spenal community/ggtt
pesglegemmaad2
Elliaty
0.4
ilama/Lamo 70
0.4
0.3
0.4
0.2
openakcoommr Ryyiap2
progle/gemmee 2e
0.1
googlelgemmo 2t
0.0
Number prepends
Number of prepends
(a) Relation P103
(b) Relation P132
Hijacking based on P190
Hijacking based on P641
0.3
0.8
0.8
Efficacy Score
0.6
upeca--communt wight?
0.7
googleigemim 26
Elficty
googhhigemmm-220
0.6
0.4
comter unilyyigazz
0.4
oosgle/gemm-22
peogleggmmme
0.
Number of prepends
Number of prepends
(c) Relation P190
(d) Relation P641
Figure C.2: Context hijacking based on relation IDs can result in LLMs output incorrect tokens. This figure shows efficacy score the number of prepends for various LLMs the COUNTERFACT dataset with hijacking scheme presented in Table 2.
We experiment with 3 other relation IDs and we see similar trends for all the L.LMs in Figure C.2a, C.2b, and C.2d. That is, the efficacy score rises for the first prepend and as we increase the number of prepends, the trend of ES rising continues. Therefore, this confirms our intuition that LLMs can be hijacked by contexts without changing the factual meaning.
Similar to Figure C.I, we experiment with reverse context hijacking where we give the answers based on relation IDs, as shown in Table 2. We again experiment with the same 4 relation IDs and the results are in Figure C.3a - C.3d. We see that the efficacy score decreases when we prepend the answer sentence, thereby verifying the observations of this study.
Hijacking without exact target words So far, the experiments use prompts that either contain true false target words. turns out, the inclusion exact target words are not necessary. To see this, experiment a variant of the generie hijacking and reverse hijacking experiments. But instead of saying "Do not think of [target_false]] false]" or "Do not think of [targetttrue)". 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"). shown in Figure C.4, context hijacking and reverse hijacing still work in this case.
26

---

Prepending answer based on P103
Prepending answer based on P131
spera conter connii-comuunnllyiyyy unilyigatz gesgleigemme-20 peoglegemmme22 201t
0.20
google progie/genmmar 20
0.12
mattalllamtt lame 2-Tish
goog http///gemmmaa...m ab-it
meta Haw.-lllmmmmmmmmmmmmm 2-w.lllmmm--llttt
Elicacy Scote
Efficacy Score
015
0.00
0.00
0.00
2
Number of prepends
Number of prepends
(a) Relation P103
(b) Relation P132
Prepending answer based on P190
Prepending inn answer based on P641
oponai comer unitygott
ope cpenal-communtn hylgpo2
google elgermmmee
geogleigemme 2b-it
googr eeweeeegemmaaa..m 20-1
mata-llama lama f
mata- mata-lamaa waw.-laaaaaaaaammmmm 2----1
0.3
Efficacy Score
0.29
Eflicaty
0.2
0.15
0.1
005
0.
0.0
Number of prepends
Number of prepends
(c) Relation P190
(d) Relation P641
Figure C.3: Reverse context hijacking based on relation IDs can result in LLMs to be more likely to be correct. This figure shows efficacy score versus the number of prepends for various LLMs the COUNTERFACT dataset with the reverse hijacking scheme presented in Table 2.
Hijacking based on P1412
Reverse Hi lljacking based on P1412
openak-commun
0.7
0.6
Eficaay Score
Efficacy Score
0.3
0.3
0.-
gemeiccommuuittyggt 2eoglelgemne-20
geogleggemmm-2bit
meteellmmnt lame-2-Tbet
Number of prepends
Number of prepends
(a) Hijacking P1412
(b) Reverse bijacking P1412
Figure C.4: Hijacking and reverse hijacking experiments relation P1412 show that context hijacking does not require exact target word to appear the context. This figure shows efficacy score versus the number prepends for various LLMs on the COUNTERFACH FACT dataset.
D Additional experiments and figures latent concept association
In this appendix section, we present additional experimental details and results from the synthetic experiments latent concept association.
Experimental setup Synthetic data are generated following the model Section 4.1. Unless otherwise stated, the default setup has w 0.5, 3 1 and AV(i) V {i} and L = 256. The default hidden dimension of the one-layer transformer is also set be 256. The model optimized using AdamW [LH17] where the learning rate is chosen from {0.01, 0.001). The evaluation dataset
27

---

1.
08-
Accuracy
0.6
0.2
idenbit WV Tratned Vey
Idericiy
Number of Latent Concepts
Number Latent Concepts
(a) L 64
(b) 128
Figure D.I: Fixing the value matrix Wy as the identity matrix results in lower accuracy compared to training Wv. especially for smaller context length L. The ligure reports accuracy for both fixed and trained Wv settings, with standard errors calculated over 10 runs.
drawn from the same distribution as the training dataset and consists of 1024 (xt, y) pairs. Although theoretical results in Section 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, typically report accuracy which is emror.
D.1 On the value matrix Wy
In this section, we provide additional figures Section 6.1. Specifically, Figure D.1 shows that fixing the value matrix 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:
Wv = Σ WE(1)(
E
WE(1)²
EI|
(V))
D.2
On the embeddings
This section provides additional figures from Section 6.2. Figure D.4 shows that the underparame- terized regime, embedding training is required. Figure D.5 indicates that the embedding structure in the underparameterized regime roughly follows (5.2). Finally Figure D.6 shows that, when the value matrix is fixed to the identity, the relationship between inner product of embeddings and their corresponding Hamming distance is mostly linear.
D.3
On the attention selection mechanism
This section provides additional figures from Section 6.3. Figure D.7-D.8 show that attention mechanism selects tokens in the same cluster as the last token. In particular, for Figure D.8, we extend experiments to consider cluster structures that depend on the first latent variables. In other words, for any latent vector z", we have
N(=') (a: and 2 22 {=="}
D.4 Spectrum of embeddings
We display several plots of embedding spectra (Figure D.9, Figure D.10, Figure D.11, Figure D.12) that exhibit eigengaps between the top and bottom eigenvalues, suggesting low-rank structures.
28

---

Accuracy
0.6
Oraginal
Accuracy
0.6
Orignall Constructed Wv
Co Constructed We
Random Vv
Random Wv
04
0.0
0.0
10
512
128
256
512
Dimensions
Dimensions
(a)
(b) m 6
08
0.8
Accurac
0.0
Original
Accuracy
0.6
Originl
Co Constructed We
Constructed Wy
dondomm We
Random Wv
02
0.
0.0
15
32
120
250
512
128
256
532
Dimensions
Dimensions
(c)
(d)
Figure D.2: When the value matrix is replaced with the constructed one in trained transformers, the accuracy does not significantly decrease compared replacing the value matrix with randomly constructed ones. The graph reports accuracy under different embedding dimensions and standard errors runs.
D.5 Context hijacking in latent concept association
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' (true target) and y² (false target) and then generate contexts _ (1), (2) and x.² ( 127) from Mc-ly') and p(>??y'' Then we mix the two contexts with rate Pm- In other words, for the final mixed context
(f1, fz), & has probability
1 Pm to be t/ and Pm probability to t?. 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 LLMs.
D.6 On the context lengths
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 higher accuracy.
29

---

Average Smallest Principa) Ange e
Condructed WV Random We
Average Smallest Principal Angle
Canstracced Us Random Wy
70
Gaussian WV
70
Causs
50-
50-
20
20
26
32
04 Dimensions
128
256
512
16
22
128
255
312
Dimensions
(8)
(b)
6
Average Smallest Principal Angle
Gondmirbed WV Rancom
Average Smallest Principal Angle
80-
Canstrareed Us
70
Gaussian Wy
70-
Gauss
Rondom Wy
40
50
20
20
16
32
04
120
254
512
16
22
64
120
255
532
Dimensions
Dimensions
(c)
(d) 8
Figure D.3: The constructed value matrix Wy has similar low rank approximation with the trained value matrix. figure displays average smallest principal angles between low low-rank approximations trained value matrices and those of constructed, randomly constructed, and Gaussian iiissian-nittiaiized value matrices. Standard errors are over runs.
30

---

0.90
0.8
Accuracy
Accuracy
0.7
0.8
0.80
beddings
0.65
16
32
64
128
256
512
128
236
512
Dimensions
Dimensions
(a)
(b) 6
Accuracy
0.
Accuracy
0.6
Frozon Embedcing
Trained tmbeidzings
0.4
0.4
02
16
32
256
512
64
128
256
512
Dimensions
Dimensions
(c)
(d) 8
Figure D.4: In the underparameteriieed regime (d < V), freezing embeddings to initializations causes significant decrease in performance. The graph reports accuracy with different embedding dimensions and the standard errors are over runs. Red lines indicate when V.
0mm:16
Dm-36
0lm-32
Inner Products
Inner Products
00
40
-70
4
Hamming Distance
Hamming Distance
Figure D.5: The relationship between inner products embeddings and corresponding Hamming distances tokens can be approximated by (5.2). The graph displays the average inner product between embeddings of two tokens against the corresponding Hamming distance hetween these tokens. Standard errors are over 5 runs.
31

---

Dtm 32
0km 32
80
Dim
96
Dir
60-
Inner Products
Inner Products
0-
-20
-20
-40
-00
2
Hamming Distance
Hamming Distance
(8) 5
(b) 6
100
Utm 52
Dim
100
Dim
USm =2
75-
inner Products
Inner Products
50
30-
25
0-
-25
-25
-30
3 4
2
8
Hamming Distance
Hamming Distance
(e)
(d) 8
Figure D.6: The relationship hetween inner products of embeddings and corresponding Hamming distances tokens is mostly linear when the value matrix Wy fixed the identity. The graph displays the average inner product between embeddings of two tokens against the corresponding Hamming distance between these tokens. Standard errors are over 10 runs.
32

---

0010
Query Tokens
Query Tokens
-.003
15
40
0.002
0.011
Key Tokens
Key Tokens
(a) m
(b) m 6
0.008
0.005
Query Tokens
0.004
Query Tokens
0.00150
100
0.00125
0.003
0.00100
80
0.002
0.00775
100
200
0,001
170
Key Tokens
Key
(c) m
(d) m
Figure D.7: The attention patterns show the underlying cluster structure of the data generating pracess. Here, for any latent vector, we have A((²) {2: z =1) (=*). The figure shows attention score heat maps that are averaged over 10
33

---

0.00
0.000
Query Tokens
30
Query Tokens
0.006
15
0.005
30
0.004
40
0.00
0.002
0.002
Key Tokens
Key
(a) m
(b) m 6
0.004
0.0025
QuEry Tokens
0.003
Query Tokens
0.0030
100
0.0023
0.002
80
0.0015
100
200
0.0010
170
Key Tokens
Key
(c) m
(d) m
Figure D.8: The attention patterns show the underlying cluster structure of the data generating pracess. Here, for any latent vector, we have (22") {a 2 and 21 = 22) {=*}. The figure shows attention score heat maps that averaged 10
34

---

Spectrum of Fmbeddings
Spectrum of Embeddings
30
Singular Values
Singular Values
10
i9
(a) Sample
(b) Sample 2
Spectrum of Embeddings
Spectrum of Embeddings
Singular Values
Singular Values
04
14-
(c) Sample 3
(d) Sample 4
Figure D.9: The spectrum of embedding matrix WE has eigengaps between the top and bottom eigenvalues, indicating low rank structures. The figure shows results from 4 experimental runs. Number of latent variable is 7 and the embedding dimension is 32.
35

---

Spectrum of Embeddings
Spectrum of Embeddings
16-
Singular Values
Singular Values
12
10
(a) Sample
(b) Sample 2
Spectrum of Embeddings
Spectrum of Embeddings
Singular Values
Singular Values
16-
14
14
(c) Sample 3
(d) Sample 4
Figure D.10: The spectrum of embedding matrix WE has eigengaps between the top and bottom eigenvalues, indicating low rank structures. The figure shows results from 4 experimental runs. Number of latent variable is 7 and the emhedding dimension is 64.
36

---

Spectrum of Embeddings
Spectrum of Embeddings
Singular Values
23.0
Singular Values
15.0
23.5
22.9
20.0
15.8
165
12.5
(a) Sample
(b) Sample 2
Spectrum of Embeddings
Spectrum of Embeddings
Singular Values
Singular Values
23.0
20.0
00.9
15.0
(c) Sample 3
(d) Sample 4
Figure D.11: The spectrum of embedding matrix WE has eigengaps between the top and bottom eigenvalues, indicating low rank structures. The figure shows results from 4 experimental runs. Number of latent variable 8 and the emhedding dimension is 32.
37

---

Spectrum of Embeddings
Spectrum of Embeddings
375
25.0
Singular Values
Singular Values
12.5
(a) Sample
(b) Sample 2
Spectrum of Embeddings
Spectrum of Embeddings
275
23.0
Singular Values
Singular Values
22.
20.0
20.0
17.5
37.3
15.0
15.0
10.0
15
(c) Sample 3
(d) Sample 4
Figure D.12: The spectrum of embedding matrix WE has eigengaps between the top and bottom eigenvalues, indicating low rank structures. The figure shows results from 4 experimental runs. Number of latent variable 8 and the emhedding dimension is 64.
38

---

Accuracy
Accuracy
65
Mixing Rate
Mixing Rate
6
Accuracy
Accuracy
62
65
03
02
M
Mixing Rate
Mixing Rate
Figure D.13: Mixing contexts can cause misclassification. The figure reports accuracy true target false target under various context mixing rate. Standard errors are over runs.
39

---

18
Dim
Dlm
0.
Accuracy
0.
Accuracy
0.6
0.6
0.5
04
0.3
16
128
Context Lengths
Context Lengths
1.0
Rim
Dim
Accuracy
0.6
Accuracy
0.6
04
0.
16
128
Context Lengths
Context Lengths
Figure D.14: Increasing context lengths can improve accuracy. figure reports accuracy across various context lengths and dimensions. Standard errors are over runs.
40

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Context hijacking is verified empirically. And the study how single-layer transformers can solve latent concept association is done both empirically and theoretically.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
limitation section is included in the appendix.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
41

---

Justification:
The assumptions are clearly stated in the theorem statements and proofs are given in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The experimental details of both LLM and synthetic experiments are provided in the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
42

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
Code is provided in the supplemental material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
wttps.cc/ppblccc//cccccccpcppppp gewic//gugee...ee CoodeSummii agionPolicy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, 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 level of detail that is necessary to appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
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, confi- dence 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.)
43

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative emrr rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Synthetic experiments can just be done on CPUs.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: http///eerr....eee
Answer: [Yes]
Justification: authors have reviewed the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification:
The paper is mostly theoretical and has no societal impacts.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
44

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is 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 being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 best faith effort.
12. Lieenses 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 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 URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.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, paperswi thcode. com com/dataseta has curated licenses for some datasets. Their licensing guide can help determine the license of 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.
45

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
46


Reviews:
[{"review_id": "GG2Cf0Bm4U", "reviewer": "Reviewer_rSoN", "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.", "questions": "1. In Section 5.3, it is unclear how the embeddings are trained/updated.\n\n2. 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.\n\n3. 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": "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.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. Given the heat of LLM, and the importance of prompt engineering, the problem brought by the authors matches the field's concern. Namely, what is an important part of the architecture of LLMs that is susceptible to noise, causing errors in generation.\n\n2. The proposed method effectively avoids the complicated effect of multi-layer attention by experimenting with a one-layer attention structure.\n\n3. The proposed embedding’s low-rank property provides a potential way to reduce computation complexity.", "weaknesses": "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?\n\n2. 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?\n\n3. 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?\n\n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your thoughtful review and questions! We are glad that you find the problem of the paper matches the field's concern.\n\n**_Solving context hijacking_**\n\nWe apologize if the main objective of our paper was not communicated clearly. Our primary interest lies in understanding the inner mechanisms of _pre-trained_ LLMs. To this end, our experiments on context hijacking are designed to stress test LLMs, observe their failure points, and hypothesize how LLMs work from such observations. Through these experiments, we hypothesize that LLMs might use certain tokens in contexts as clues for memory retrieval. While solving the issue of context hijacking is important, we consider it to fall outside the scope and goal of this paper.\n\t\n**_Quantization_**\n\nThanks for the suggestion! Quantization is indeed a fascinating topic to explore. However, since our current paper aims to focus on understanding how LLMs achieve fact retrieval, we believe it falls outside the scope of this work. It would, however, make an excellent direction for future research.\n\n\n**_Value matrix_**\n\nThis is a good question! We study a simplified one-layer transformer network. The output from self-attention is a combination of embeddings by definition. On the other hand, the unembedding (the last linear layer) is tied with the embedding (L192-L193). Therefore, intuitively, it makes sense to construct the value matrix using embeddings, as the value matrix lies between the self-attention mechanism and the unembedding layer, both of which are related to embeddings. Similar construction is also shown in other papers as well [1, 2]. \n\nOn the other hand, the constructed value matrix is different from the trained value matrix as we can see those accuracies after replacement do not match exactly in Fig C.2. This difference arises because the constructed value matrix is only an approximation and a simplified model, whereas the trained value matrix results from more complex training dynamics.\n\n\n[1] Bietti, Alberto, et al. \"Birth of a transformer: A memory viewpoint.\" Advances in Neural Information Processing Systems 36 (2024).\n\n[2] Cabannes, Vivien, Elvis Dohmatob, and Alberto Bietti. \"Scaling laws for associative memories.\" arXiv preprint arXiv:2310.02984 (2023).\n\n**_Comparison to other methods_**\n\nWe apologize if we misunderstand your question, however, we have not proposed any new methods in this paper. The experiments on context hijacking serve only as a robustness test, and we conducted them across different LLMs and various prompts. There is no new method proposed in this paper.\n\n\n**_Embeddings training_**\n\nSorry about the confusion. All the experiments on single-layer transformers, unless otherwise stated, are trained jointly with AdamW. Training details can be found in Appendix C. We will revise the paper to clarify this point earlier in the paper.\n\n\n\n\n**_Question about section 5.5_**\n\n\nOur original experiments on context hijacking do add distracting prompts like “The Eiffel Tower is not in Chicago” at the beginning of the prompt (Figure 1, L103-121). On the other hand, experiments in section 5.5 and Fig C.13 are on synthetic datasets generated from the latent concept association tasks, rather than real sentences. These experiments are designed to simulate context hijacking. We only do experiments on simulated data as opposed to real sentences in this case because it allows us to do controlled experiments like changing the mixture rate.\n\n\n**_Question about Figure C.2 in Section 6.1_**\n\nWe didn’t claim that the constructed value matrices and trained value matrices are equivalent, nor should one use the construction as a method to replace trained value matrices. Rather, the goal is to show that the constructed value matrices are similar to trained value matrices (L224-L225; “it turns out empirically that this construction of Wv is _close_ to the trained Wv, even in the noisy case”) so that one can use the constructed ones to gain insight into how trained value matrices work. Indeed, even with the accuracy drops at certain dimensions, such approximation is still nontrivial compared to the baselines which are based on randomly constructed matrices composed of embeddings (see green line in Figure C.2).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "8jxbB4OuBN", "reviewer": "Reviewer_bm96", "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.", "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*]).\n- In Fig 2a, are these results averaged over all relation types?\n- 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": "I think the authors adequately addressed the limitations of their work.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper documents an interesting failure case for factual retrieval with LLMs (\"context hijacking\"). Similar \"distractor\" effects have been documented in prior work (which the authors cite), but I have not seen this result for the factual retrieval setting.\n- The theoretical analysis presents a simple model that could give rise to the empirical phenomena, and these results are supported by a variety of experiments and analysis.\n- In general, I think it is a useful contribution to provide more theoretical tools for understanding the learning dynamics of attention models, and to try to connect these analyses to real world failure cases (like context hijacking).", "weaknesses": "- 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?\n- 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the support and thorough review! We’re pleased that you think the paper makes a \"useful contribution”.\n\n**_Clarification on associative memory_**\n\t\nThis is a really good question! In the literature, the definition of associative memory is usually quite broad. While any prediction can be thought of as a form of association, the key issue here is how a model can generalize beyond the associations formed by its training sets. An ideal model would understand and reason about new context sentences, compose existing knowledge and produce correct outputs. However, due to the phenomenon of context hijacking, we hypothesize that LLMs might instead rely on the appearances of certain tokens related to the output to guide memory retrieval. In this case, the association is more statistical rather than based on semantic or factual meaning. \"Not semantically meaningful\" in the context of fact retrieval refers to incorrect associations, as demonstrated by context hijacking examples. We will clarify this in the revised version.\n\n**_Connection to induction/ICL_**\n\nThis is an excellent question. We have thought quite a bit about the connection between induction heads and ICL and cite many related works in the reference (L65-69). \n\nFirst, falling for context hijacking is not necessarily a problem of the induction head. In the literature, the induction head, whether it's the direct copy type [1,2] or the statistical type that relies on bigram statistics [3] requires the output token to be present in the context. However, context hijacking is slightly different. As mentioned in Section 3 (L115-116), \"one can even hijack by only including words semantically close to the false target (e.g., 'France' for false target 'French').\" In other words, context hijacking is about latent concept association rather than copying tokens, making it distinct from existing works on induction heads.\n\nSecond, resisting context hijacking is not necessarily a property of larger models either. We have already shown that larger models like LLaMA-7B and instruction-tuned models can exhibit context hijacking in Section 3. While in some examples larger models are harder to hack (Figure 1), it is not always the case (Figure 2(b), Figure B.2(b)).\n\nFinally, in our toy model, we simplified the problem to a hypothesis-testing type problem. If the input context is modified to look like it is coming from a different distribution, then it is harder for models of different sizes to predict, though larger models would theoretically have more capacity to distinguish closer distributions. This is also not specific to the induction head mechanism. \n\nWe will expand on these points more in the updated version.\n\n\n[1] Elhage, Nelson, et al. \"A mathematical framework for transformer circuits.\" Transformer Circuits Thread 1.1 (2021): 12.\n\n[2]  Bietti, Alberto, et al. \"Birth of a transformer: A memory viewpoint.\" Advances in Neural Information Processing Systems 36 (2024).\n\n[3] Edelman, Benjamin L., et al. \"The evolution of statistical induction heads: In-context learning markov chains.\" arXiv preprint arXiv:2402.11004 (2024).\n\n\n**_Questions on efficacy score_**\n\nThanks for catching the typo! It should be  Pr[o_] < Pr[o*]. On the other hand, we didn’t require Pr[o_] < Pr[o*] prior to modifying the context because we are examining data at the population level (i.e., the percentage of prompts affected). By comparing the efficacy score with no prepend (no modification) to multiple prepends, it is evident that adding misleading contexts can cause LLMs to output incorrect tokens  (for example, in Figure 2, the efficacy score drops after just one prepend).\n\n**_Fig 2a_**\n\nThe results in Fig 2a are averaged across all prompts in the CounterFact dataset. We also included the standard error, although it is somewhat difficult to see.\n\n**_What if the last token is noisy_**\n\nThis is a great question! The noise we refer to in section 5.4 is about occurrences of random, somewhat irrelevant tokens in the context. But we didn’t allow the last token to be sampled uniformly. Intuitively, this is because the final token is directly linked to the next token prediction and often in real sentences, not totally random. On the other hand, if the final token is sampled from the noisy distribution, then what is considered noise in the context is not that well defined anymore if the final token is from the uniform mixture. This is because the attention mechanism is to select tokens most relevant for the final token and thus randomness should be defined relative to the final token as well.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4u8FlK87Ng", "reviewer": "Reviewer_vgcR", "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.\nThe 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.\nFinally, the paper includes empirical validation of the theoretical results.", "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": "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.\"", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "- The problem studied is important to improve our understanding of language models' internal mechanisms.\n- The paper presents an interesting theoretical analysis of factual recall and provides solid empirical evidence to support it.", "weaknesses": "- 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?\n- 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?\n- Minor point: the figures can be improved (e.g., increasing the font size)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive feedback and questions. We're glad that you found the paper to be important, interesting, and solid!\n\n\n**_Limitations_**\n\nSorry for the confusion! We don't assume that each latent concept is associated with only one token. Instead, the latent concept space consists of multiple latent variables, with each latent vector associated with one token. Therefore, each token can represent multiple latent concepts.\n\nThe study is motivated by the behavior of multi-billion-parameter models. But this behavior – context hijacking – is a failure case. It is thus reasonable to expect this behavior to translate to smaller models as well. However, studying single-layer transformers allows us to conduct detailed, controlled experiments. Moreover, current theoretical work on multi-layer transformers is limited and remains an active research area. Even with single-layer transformers, we have obtained many interesting results including approximated value matrix structure (L331-L343) and low-rank structure of embeddings (L359-L362). In particular, the existence of the low-rank structure provides the theoretical grounding for many editing and fine-tuning methods including LORA and ROME that exploit low-rank structure. Note that such low-rank structures naturally emerge by just training single-layer transformers.\n\n**_Question on templated hijacking prompt_**\n\nThis is a great question! First of all, in our systematic experiments, we saw that even after one prepend, arguably more in line with the training distributions and less templated, there's still a performance downgrade. \n\nBut **to answer your question,** we conducted a new test by placing a description of Chicago from Wikipedia at the beginning of the prompt: \"Chicago is the most populous city in the U.S. state of Illinois and in the Midwestern United States. With a population of 2,746,388 as of the 2020 census, it is the third-most populous city in the United States. Therefore, The Eiffel Tower is in the city of\". When presented with this prompt, all four models responded with \"Chicago.\" This suggests the hypothesis still makes sense. On the other hand, we are interested in out-of-distribution generalization in this paper because LLM is supposed to encounter sentences it has not seen before. This question is inherently interesting. Out-of-distribution generalization might involve a different mechanism, but that is a separate research question.\n\n**_Question on efficacy score_**\n\nThat’s a good catch. There’s a typo. We will fix it in the final version.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xYS9gyj6TO", "reviewer": "Reviewer_MLxQ", "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.", "questions": "None", "limitations": "It would be better to include a dedicated \"limitations\" section beyond those in the checklist.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Studying how language models' generation is influenced by the context is an important direction toward better understanding and improving LLMs (e.g., hallucinations). The proposed context hijacking is interesting, and could be thought of as another way of stress testing LMs.\n\n- Analyzing the theoretical properties of transformers in associative memory recall tasks could inspire future work and formations toward a better understanding of transformers and their actuality. The roles of different components in the proposed task also help open up the black box of transformers and could potentially inspire future investigations of the role of real LMs' components.\n\n- The paper is generally well-written.", "weaknesses": "- 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.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for your support and useful comments! Your suggestion is really valuable in helping us clarify our paper.\n\n\n**_Questions on larger models like GPT-4 and limitations of single-layer transformer_**\n\n\nThis is a great question! First of all, we didn’t test this on GPT-4 because, as a closed-source model, it’s unclear what is behind its predictions. For example, what kind of instruction tuning is trained and does it rely on external sources (web search + rag) for basic fact prediction. \n\n\nNonetheless, even in the official GPT-4 technical report [1], we see an example similar to context hijacking (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 [2] which we reference many in the paper (L78-82, L120-121). So we reasonably suspect that similar behavior still exists in larger models but is harder to exploit. This is similar to the “Chicago'' example we give in Figure 1, where a larger model like LLaMA-7B requires more prepends but is still hackable and we provide more detailed experiments in section 3 as well. \n\n\nOn the other hand, in the literature, theoretical works on multi-layer transformers are limited and remain an active research area while studying single-layer transformers allows us to easily do carefully controlled experiments. Even with single-layer transformers, we have obtained many interesting results, such as approximated value matrix structure (L331-L343) and low-rank structure of embeddings (L359-L362). In particular, the approximated value matrix structure is closely related to earlier works on multi-layer transformers as well [3]. Moreover, the existence of the low-rank structure provides the theoretical grounding for many existing editing and fine-tuning methods including LORA and ROME that exploit low-rank structure. A lot of these methods work on current LLMs. Note that such low-rank structures naturally emerge by just training single-layer transformers. We hope these results can further lay the foundation for future research on multi-layer transformers on similar topics.\n\n\n\n\n[1] Achiam, Josh, et al. \"Gpt-4 technical report.\" arXiv preprint arXiv:2303.08774 (2023).\n\n\n[2] Shi, Freda, et al. \"Large language models can be easily distracted by irrelevant context.\" International Conference on Machine Learning. PMLR, 2023.\n\n\n[3] Bietti, Alberto, et al. \"Birth of a transformer: A memory viewpoint.\" Advances in Neural Information Processing Systems 36 (2024).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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.
Decision: Accept (poster)