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

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.

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. - 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. - 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. - Related to the last point, the proposed associative memory formation lacks modeling of semantics, where the task boils down to tokens and their similarity. This may be a good formulation for models that are not powerful enough and mostly rely on surface cues to generate the next token, but not for models that have strong language understanding and logic. This also relates to the shallow network (one-layer transformer) which is the main target for theoretical analysis. The results may not transfer/be enlightening for improving the current LLMs.

Questions

None

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

It would be better to include a dedicated "limitations" section beyond those in the checklist.

Summary

This paper studies the mechanics of factual recall in transformer-based models, framing the problem as next token prediction. In particular, the authors focus on the brittleness of language models, which can be elicited to provide different answers to factual queries by adding distracting information in the prompt (a procedure which the authors term context hijacking). This phenomenon is shown for various language models, with sizes up to 7B parameters. The authors then formulate a hypothesis for the hijacking phenomenon, according to which the model predicts the next token based on a similarity measure in a latent concept space. They provide a series of theoretical results that explain how a single-layer transformer can solve the latent concept association problem. Finally, the paper includes empirical validation of the theoretical results.

Strengths

- The problem studied is important to improve our understanding of language models' internal mechanisms. - 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? - In the motivation of the problem, the authors show how LLaMA 7B can be “hijacked” to output a wrong answer to the prompt about the location of the Eiffel Tower. However, a text consisting of the sentence “The Eiffel Tower is not in Chicago” repeated eight times represents an input arguably out of distribution with respect to the model’s pre-training data. In this setting, it is possible that the model continues the prompt ”the Eiffel Tower is in the city of” using a different mechanism than it would have without the prepended hijacking context. A test for the authors’ hypothesis would be to prepend a non-templated paragraph about, for example, the city of Chicago (possibly mentioning that the Eiffel Tower is not located there). If the authors’ hypothesis is correct, this should still steer the model towards completing “the Eiffel Tower is in the city of” with “Chicago.” Is this the case? - Minor point: the figures can be improved (e.g., increasing the font size)

Questions

- 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)?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

4

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."

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.

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. - 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. - 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? - I am not fully convinced that the one-layer transformer is a meaningful model the context hijacking phenomenon. For example, in the main example ("the Eiffel tower is not in Chicago"), it's seems like the most plausible mechanisms for either resisting context hijacking, or falling for context hijacking, would involve multi-layer Transformers--for example, the model might predict "the Eiffel tower is in Chicago" due to a kind of induction head/ICL mechanism. I think it would be helpful to expand in more detail on the connection between context hijacking and the toy model (section 5.5).

Questions

- For Efficacy Score (section 3), it seems it should also require that Pr[o_] < Pr[o*] prior to modifying the context. In Fig 2a, the efficacy score seems to be the opposite of what is described in the text--the intervention makes the score go down (i.e., more often that Pr[o_] > Pr[o*]). - In Fig 2a, are these results averaged over all relation types? - In Section 5.4: "This implies that the self-attention layer can mitigate noise and concentrate on the informative conditional distribution π". In this setting, the final token is sampled without noise. Would this also be true if the final token had noise?

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

I think the authors adequately addressed the limitations of their work.

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.

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. 2. The proposed method effectively avoids the complicated effect of multi-layer attention by experimenting with a one-layer attention structure. 3. 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? 2. Given the prevalence of low-bit quantization, I wonder how quantization will change the associative memory. For example, after quantization, will LLM be less distracted? Or perhaps quantization will enhance the effects of the misleading context? 3. For the constructed value matrix (formula 4.1 and 5.1), even if results in section 6.1 show that using the constructed value matrix retains the accuracy, two concerns remain. First, why should a value matrix constructed from an embedding matrix be expected to act as an associative memory (intuitively)? Secondly, how and why is the constructed value matrix different from self-attention’s value matrix (gather information), if any? 4. In the results section, it is unclear how the method performs differently across different LLMs, different datasets, and other SOTA methods. It is unclear if the results can be generalized to other settings.

Questions

1. In Section 5.3, it is unclear how the embeddings are trained/updated. 2. In Section 5.5, the support evidence is not sufficient. Instead of showing across different mixture rates, I would like to see how changing context (for example, adding “The Eiffel Tower is not in Chicago” to the beginning of the prompt) can potentially impact associative memory and self-attention. Another concern regarding Fig C.13 is, how/if the impact will be different should one concatenate additional context (“The Eiffel Tower is not in Chicago”) to the beginning, middle, and end of prompt. 3. In Section 6.1, Figure C.2, the author concludes that the constructed value matrix can be used to replace the self-attention value matrix without performance sacrifices. However, the author fails to provide enough justification on the consistent drop after certain dimensions, across different m. For example, when m=5, the accuracy of using the constructed matrix drops significantly after dim = 128, which hints that the constructed matrix and self-attention’s value matrix are not equivalent.

Rating

6

Confidence

4

Soundness

2

Presentation

2

Contribution

2

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.