Llama See, Llama Do: A Mechanistic Perspective on Contextual Entrainment and Distraction in LLMs
Authors: Jingcheng Niu, Xingdi Yuan, Tong Wang, Hamidreza Saghir, Amir H. Abdi
Accepted at: acl25

Llama See, Llama Do: A Mechanistic Perspective on Contextual Entrainment and Distraction in LLMs

Jingcheng Niu*, Xingdi Yuan*, Tong Wang*, Hamidreza Saghir*, Amir H. Abdi*
*University of Toronto, UKP Lab, Technical University of Darmstadt, Microsoft Research
niu@cs.toronto.edu, {eric.yuan,tong.wang,hsaghir,amirabdi}@microsoft.com

Abstract
We observe a novel phenomenon, contextual entrainment, across a wide range of language models (LMs) and prompt settings, providing a new mechanistic perspective on how LMs become distracted by “irrelevant” contextual information in the input prompt. Specifically, LMs assign significantly higher logits (or probabilities) to any tokens that have previously appeared in the context prompt, even for random tokens. This suggests that contextual entrainment is a mechanistic phenomenon, occurring independently of the relevance or semantic relation of the tokens to the question or the rest of the sentence. We find statistically significant evidence that the magnitude of contextual entrainment is influenced by semantic factors. Counterfactual prompts have a greater effect compared to factual ones, suggesting that while contextual entrainment is a mechanistic phenomenon, it is modulated by semantic factors.

We hypothesise that a circuit of attention heads — the entrainment heads — that corresponds to the contextual entrainment phenomenon. Using a novel entrainment head discovery method based on differentiable masking, we identify these heads across various settings. When we “turn off” these heads, i.e., set their output to zero, the effect of contextual entrainment is significantly attenuated, causing the model to generate output that capitulates to what it would produce if no distracting context were provided. Our discovery of contextual entrainment, along with our investigation into LM distraction via the entrainment heads, marks a key step towards mechanistic analysis and mitigation of the distraction problem.1

1 Introduction
Language models (LMs), especially large language models (LLMs), can sophisticatedly utilise contextual information provided in prompts to a surprising degree. Brown et al. (2020) was among the first to identify this capability and coin the term in-context learning (ICL) to describe this capability. Subsequent work has demonstrated that LMs can process and process and utilise contextual information provided in prompts across various settings.

Nonetheless, LMs can also misuse contextual information in prompts (Figure 1). Shi et al. (2023) experimented with inserting distracting, irrelevant information into grade-school maths problems and found that it successfully diverted the model from reaching the correct answer. Their work shed light on a fundamental issue within LMs, which they termed distraction. Since then, distraction has been recognised as one of the most challenging and widespread issues for RAG (Yoran et al., 2023; Cuconasu et al., 2024; Wu et al., 2024), prompting the development of distraction mitigation strategies. Notably, Yoran et al. (2023) proposed leveraging an NLI model to remove irrelevant context from the prompt as a solution to this problem.

Distraction, however, is a phenomenon in LMs that is easy to grasp but difficult to define precisely. Most prior work defines distraction using the term “(ir)relevant,” framed in RAG and information retrieval terms; i.e., whether the context prompt contains the information needed to answer the question correctly. While this provides an adequate general description of the problem, we identify challenges when examining it in greater detail. First, relevance is too broad a concept. Consider the following context prompts: (1) Messi is a football player, (2) Japan is in Asia, (3) Greece is in Asia, and (4) Colorless green ideas sleep furiously. According to the earlier definition, all of these are ‘irrelevant’ since they do not contain the information needed to correctly answer the question Greece is located on the continent of ___. However, it is evident that they differ drastically in how they might influence an LM’s response, and therefore, a more precise definition and more fine-grained taxonomy of dis-
(a) When no context is provided to the model, it can confidently generate the correct response.

(b) However, when context from a related topic is provided, the model may become distracted and produce incorrect responses.

(c) The model continues to exhibit distractions across different formats and paraphrasing of contextual information.

(d) Just a single token, Asia, can distract the model. This indicates that the phenomenon of LLM distraction is more superficial and fundamental than previously suggested.

Figure 1: LLMs can be distracted by various types of context. Each sub-figure illustrates Llama-3.1-8B’s output probability of the top ten tokens given the input prompt. The model inputs are displayed in green.

traction is needed. Moreover, we find evidence that these “irrelevant” context prompts can benefit the LMs’ performance. Although they do not contain the exact answer, they may provide useful implicit information about the question.

Regardless of the debate on the precise definition of distraction, we observe a phenomenon in LMs related to how they use and misuse information from the context prompt. In particular, we identify contextual entrainment — LMs consistently assign significantly higher probabilities (or logits) to any tokens that have appeared earlier in the context prompt, regardless of their relevance or semantic relation to the question or the prompt. Simply put: llama see, llama do. If a token appears in the context prompt, even a randomly token, the model assigns it a higher probability or logit. Our experiments show that various LMs exhibit contextual entrainment across a wide range of configurations.

At first glance, this phenomenon shares some similarities with the inductive pattern-repeating phenomenon identified by (Elhage et al., 2021; Olsson et al., 2022), the differences are far more substantial. First, contextual entrainment does not require the reappearance of a prefix as a trigger, unlike inductive pattern-repeating, which depends on first encountering a pattern [A][B] and then seeing the prefix [A] again to generate [B]. Instead, contextual entrainment occurs when a token has previously appeared in the context. Second, unlike the inductive pattern-repeating phenomenon, which is largely independent of semantic factors and token statistics (Olsson et al., 2022), the magnitude of the effect of contextual entrainment is influenced by semantic factors. In particular, we find that counterfactual prompts have a significantly greater effect. Thus, we identify contextual entrainment as a novel phenomenon that plays a crucial role in LMs’ use and misuse of contextual information in prompts. It may be a major factor contributing to the distraction problem.

Lastly, in contrast to the previous discussion on how this phenomenon differs from inductive pattern-repeating, here we identify a similarity. Similar to the induction heads identified by Olsson et al. (2022), we find that 3–10% of the attention heads are associated with contextual entrainment, which we refer to as entrainment heads. When these entrainment heads are disabled, i.e., their output is set to zero, the effect of contextual entrainment is drastically suppressed and the LM generates outputs similar to that produced when no context prompts are given. This provides an interesting insight into the mechanism governing the contextual entrainment phenomenon, and we hope our work can serve as a starting point for investigating the problem of distraction mechanistically.

Contributions In this paper, we identify contextual entrainment, a novel phenomenon that plays a crucial role in LM distraction (§3). We provide evidence that the phenomenon can be considered mechanistic and occurs commonly across various LMs and settings (§3.2). However, it is also influenced by semantic factors: counterfactual context more effectively induce contextual entrainment (§3.2). Finally, we identify entrainment heads (§4) using a novel method based on differentiable masking which, when “turned off,” the effect of contextual entrainment is drastically suppressed.

2 Related Work

Distraction LMs are known to be susceptible to distractions caused by contextual information
<table>
  <tr>
    <th>Context Setting</th>
    <th>Context Prompt</th>
    <th>Query Prompt</th>
    <th></th>
    <th></th>
  </tr>
  <tr>
    <td>Related</td>
    <td><b>On the inside, bananas are white.</b></td>
    <td>What color are mangoes on the inside? They are</td>
    <td>white</td>
    <td>orange</td>
  </tr>
  <tr>
    <td>Irrelevant</td>
    <td>The capital of Canada is Ottawa.</td>
    <td>What color are mangoes on the inside? They are</td>
    <td>Ottawa</td>
    <td>orange</td>
  </tr>
  <tr>
    <td>Random</td>
    <td>Promotion</td>
    <td>What color are mangoes on the inside? They are</td>
    <td>Promotion</td>
    <td>orange</td>
  </tr>
  <tr>
    <th colspan="5">Distraction over Counterfactual Context</th>
  </tr>
  <tr>
    <td>Counterfactual</td>
    <td><b>On the inside, bananas are green.</b></td>
    <td>What color are mangoes on the inside? They are</td>
    <td>green</td>
    <td>white</td>
    <td>orange</td>
  </tr>
</table>

Table 1: Prompt Setup. The emojis represent the target tokens: ☺: counterfactual; ☹: distracting; 😊: correct.

in prompts. For instance, Shi et al. (2023) found that while LMs can accurately solve grade-school maths problems, they may fail when provided with additional information in the prompt. They, however, did not explore the mechanisms underlying these distractions. More research has since confirmed that LMs can be easily distracted (Yoran et al., 2023; Wu et al., 2024; Cuconasu et al., 2024, inter alia). This problem has received remarkable attention in the RAG community due to its apparent connection to retrieval robustness. Since retrievers cannot always retrieve perfectly relevant documents, the LLM within a RAG system should be as resistant to distraction as possible.

Mechanistic Interpretability & Induction Heads
Our research setup shares similarities with efforts to understand ICL, but, as discussed earlier, the differences are more significant. Elhage et al. (2021); Olsson et al. (2022) successfully identified an inductive pattern-repeating phenomenon, a crucial step toward understanding how LMs perform ICL. They observed that if an LM has encountered a sequence of tokens — even for random tokens — it will repeat the sequence if it appears again in the prompt. For instance, given the input Category 40 ids node struction ... Category 40 ids node, the model predicts struction as the most probable next token. Therefore, they concluded that LMs possess some inductive capability and “are not memorising a fixed table of n-gram statistics.” They also identified certain attention heads as induction heads in two-layer toy transformer models, which they claim perform pattern completion. More recently, Crosbie and Shutova (2024) identified induction heads in real-world LMs such as Llama-3-8B.

3 Context Entrainment

We present experiments that confirm the existence of the contextual entrainment phenomenon in this section. We find that contextual entrainment is both a mechanistic phenomenon but influenced by
• We present the models with four context prompt settings: Related, Irrelevant, Random and Counterfactual.
• Based on the setting, we observe LM’s output on two (or three) specific tokens: the counterfactual token (☺), the distracting token (☹), and the correct token (😊).
• In Section 4, we experiment with the original LM (🞄) and compare it to the modified LM (🞅), in which the identified entrainment heads are removed.

Figure 2: Glossary and Emoji-Based Notation.

semantic factors. The phenomenon is observed in all prompt settings, including those with completely randomly sampled tokens, suggesting that its existence is independent of semantic factors (i.e., contextual entrainment is mechanistic); however, the exact magnitude of its impact depends on semantic factors, as demonstrated by the significantly greater effect of counterfactual context.

3.1 Experimental Setup

Prompts & Data Table 1 shows how we constructed our prompts using facts from the LRE dataset (Hernandez et al., 2024), and Figure 2 presents a summary of the glossary and shorthands used throughout this paper. The LRE dataset contains facts in the triplet format: ⟨source, target, relation⟩ or ⟨s, t, r⟩. For example, ⟨Canada, Ottawa, capital⟩ corresponds to the fact that Canada’s capital is Ottawa. Each fact in the dataset contains several prompt templates that we leverage to construct the context and query portion of the prompts.2

We present the model with a context and a query in the prompt, separated by a single space character (e.g., <context> <query>). Given a query generated from a fact ⟨s, t, r⟩, there are four context prompt settings: related, where facts ⟨s′, t′, r⟩ are sampled from the same relation type r but differ in source s′ and target t′; irrelevant, where facts ⟨s′, t′, r′⟩ are sampled from a completely different relation type r′ without domain or range overlap; random, where the context consists of a single randomly chosen token; and counterfactual, where the tar-

2Appendix A presents more details of the LRE dataset.
The bar height indicates the average logits of:
<table>
  <tr>
    <th></th>
    <th>with context</th>
    <th>without context</th>
  </tr>
  <tr>
    <td>😊</td>
    <td>blue</td>
    <td>orange</td>
  </tr>
  <tr>
    <td>😉</td>
    <td>black</td>
    <td>grey</td>
  </tr>
</table>

(a) Average logits for different tokens across prompt settings. The colour indicates the token type (blue: 😊, orange: 😉), and the shade indicates the context setting (light: with context, dark: w/o context). Observation: The logits of the 😊 tokens increase in the relevant† setting and decrease in the irrelevant, and random↓ settings, while the logits of the 😉 tokens increase across all settings†.

The bar height indicates the difference in logits with or without the distracting context across tokens:
<table>
  <tr>
    <th></th>
    <th>with context</th>
    <th>without context</th>
  </tr>
  <tr>
    <td>😊</td>
    <td>blue</td>
    <td>orange</td>
  </tr>
  <tr>
    <td>😉</td>
    <td>black</td>
    <td>grey</td>
  </tr>
</table>

(b) Now, let us zoom in on the magnitude of the increase or decrease: the difference in logits with or without the context across different tokens. Again, colour indicates token type (blue: 😊, orange: 😉). Observation: The direction of change confirms the observation in (a). Notably, the magnitude of change for the 😊 tokens is greater than that of the 😉 tokens.

The bar height indicates the relative probability differences with or without the context across tokens:

\[
\Delta_p(\text{😊}) = \frac{p(\text{😊 w/ctx.}) - p(\text{😊 w/o ctx.})}{p(\text{😊 w/o ctx.})}
\]
\[
\Delta_p(\text{😉}) = \frac{p(\text{😉 w/ctx.}) - p(\text{😉 w/o ctx.})}{p(\text{😉 w/o ctx.})}
\]

(c) The results observed in the logits carry over to the probability space. This subfigure shows the relative difference in probability. After applying softmax, we observe the same trend as in the logits.

Figure 3: Experimental Results for the llama-3.1-8B Model. Logit and probability values change consistently after the context prompt is provided to the model. The LM assigns significantly higher logits and probabilities to tokens that appear in the context prompts. All shifts in probabilities and logits are statistically significant, with \( p < 0.0001 \) according to paired t-tests.

get in the fact \( \langle s', t'_{cf}, r \rangle \) is replaced by another target sampled from the same relation. The random tokens are sampled from the Brown corpus (Francis and Kucera, 1979). For larger relations that yield more than 100,000 combinations, we cap the size at 100,000 through random sampling.

Then, we feed in the input prompt (e.g., On the inside, bananas are white. What color are mangoes on the inside? They are), and the model outputs a logit score (or a probability after applying softmax) for every token in the vocabulary. We are interested in the logit or probability assigned to two (or three) tokens based on the input prompt: the distracting (😊) token (white) that appears in the context prompt with a similar sentence format, and the correct (😉) token (orange) that answers the actual query prompt. In the counterfactual context prompt setting (e.g., On the inside, bananas are green. What color are mangoes on the inside? They are), we further distinguish between the counterfactual (😉) token (green) (which introduces incorrect information in the context prompt), the distracting (😊) token (white) (which is the actual correct token in context), and the correct (😉) token (orange).

Language Models We experiment with GPT2 xl (Radford et al., 2019) and 4 LLaMA models (Touvron et al., 2023): llama-3.1-8B, Llama-3.1-8B-Instruct, Llama-2-7b-hf, and Llama-2-13b-hf.

3.2 Experiment Results

Figure 3 presents the experimental results for three types of distracting contexts: distracting context from a related topic (Related), irrelevant topic (Irrelevant), and random token (Random), for the llama-3.1-8B model. The results of other LMs are shown in Figure 4, which support the same observations and findings. The figures show the averaged results across all LRE relations. The full list of experimental results, which supports the same conclusion, can be found in Appendix B.

Finding 1: Contextual Entrainment: LMs assign higher logits and probabilities to tokens that appear in the context. When a model is given distracting context, there is a significant increase in the logits and probabilities of the corresponding distracting tokens. For example, when asked the question Greece is located in ___, distracting context prompts such as Japan is in Asia, Bananas are yellow, or even a single randomly sampled token, Promotion, can cause the model to assign higher logits to the distracting tokens: Asia, yellow, and Promotion, respectively. When normalised with
softmax, logit increases translate into higher probabilities. Notably, the model typically assigns very low probabilities (\(10^{-5}\) to \(10^{-3}\)) to these tokens, but with distracting context, their probabilities can increase by a factor of 10 to 100. Paired Student’s (1908) t-tests confirm that these increases are statistically significant across all LMs, regardless of their size, family, or instruction-tuning status.

Finding 2: “Distracting” context prompts can be beneficial when relevant. While the probabilities and logits of the distracting tokens (😊) consistently increase, the direction of change for the correct token (☺) varies based on topic relevance. Except for the instruction-tuned Llama-3.1-8B-Instruct model. There is a small but statistically significant decrease in the correct answer token’s logit when context information from irrelevant or random context prompts are provided.

While prior work typically groups “irrelevant context” into a single category and considers it detrimental, our findings suggest the need for a more nuanced classification. Although distracting context may not contain the exact correct answer, the implicit hints it provides can be beneficial, increasing the likelihood of the model generating the correct response. Figure 5 illustrates an example where “distracting” contextual information proves helpful, particularly in cases of question ambiguity. For instance, in the question In Argentina, people speak the language of from the country language relation in the LRE dataset, the term language can be interpreted metaphorically, leading to responses like love or football. However, the distracting context In Russia, the primary language is Russian can guide the model towards the correct answer.

Figure 4: Experimental Results for Other LMs. The figures follow the same plotting scheme as Figure 3. The results confirm the same observations shown in Figure 3 and the findings discussed in Section 3.2 across a wider range of models. Again, all shifts in probabilities and logits are statistically significant, with \(p < 0.0001\) according to paired t-tests.

Figure 5: Providing “distracting” yet relevant context can be beneficial. For instance, when a question is ambiguous, such context may help clarify its intended meaning or guide interpretation. For example, “distracting” context provided, Llama3.1-8B’s top responses to the prompt In Argentina, people speak the language of shifted from love and football to the language-related tokens: Spanish, Spain and Cast (the first word-piece of “Castilian Spanish”).

Discussion: Contextual Entrainment — A Novel Mechanistic Phenomenon Thus, a new perspective emerges for analysing the phenomenon of distraction from a mechanistic angle, which we term Contextual Entrainment. Specifically, the model assigns a higher probability to tokens that appear within the context prompt. The fact that the model assigns higher probabilities to completely random tokens underscores the mechanistic nature of this phenomenon, as no linguistic or factual factors can plausibly account for the increase in logits and probabilities of the random tokens.

Our coinage of the term Contextual Entrainment does not imply any connection to human cognitive or psycholinguistic phenomena, such as brain entrainment (Poeppel and Assaneo, 2020; Pérez et al., 2022) or lexical entrainment (Garrod and Anderson, 1987; Brennan and Clark, 1996), nor does it suggest that LMs in any way replicate human brains or cognition. Rather, we use this term
Figure 6: Counterfactual context prompts consistently cause greater distraction than factual context prompts.

because entrainment most accurately describes the phenomenon we have observed. It refers solely to the output patterns exhibited by LMs in response to contextual input, without making any claims about the underlying cognitive mechanisms or their resemblance to human cognitive processes.

3.3 Counterfactual Experiment Results

Figure 6 shows the results with counterfactual context prompts. Using these counterfactual context prompts results in a significantly greater impact compared to previously identified factual context prompts. This suggests that, while we previously established that contextual entrainment is a “mechanistic” phenomenon, it is still subject to semantic factors in determining its magnitude of impact.

Finding 3: Counterfactual context prompts consistently cause greater distraction than factual context prompts. We present the model with two types of distracting context prompts: a factual (:) one (Japan is in Asia) and a counterfactual (:) one (Japan is in Africa). After the context prompt, we query the LM with a question (Greece is located in __) and observe how the context prompt changes the logits and probability of the (counterfactual token - Africa), the (: distracting token - Asia), and the (correct token - Europe). The : setting, serving as a control group, is identical to the distraction prompt setting in Section 3.2.

The amount of distraction — in other words, the magnitude of contextual entrainment — is greater for counterfactual context prompts (:) than factual prompts (:). This is because the absolute logits of the : token when : prompts are provided are significantly lower than those of the : token when : prompts are provided (height of blue bars in Figure 6). Moreover, the extent of change is greater for counterfactual prompts. With a : prompt, the : token’s logits increases smaller compared to no context, while a : prompt causes a much larger increase, showing that counterfactual prompts create stronger distractions and greater shifts in the model’s output (orange bar height in Figure 6).

Discussion: Contextual Entrainment — A mechanistic phenomenon affected by semantic factors. We have established that the presence of contextual entrainment is independent of semantic factors, given that it occurs even with random tokens. However, in this subsection, we also find that this “mechanistic” phenomenon is nevertheless modulated by semantic factors. In particular, counterfactual prompts induce a greater effect on contextual entrainment than factual context prompts.

The mechanism through which LMs utilise information from prompts is not yet fully understood. There is an ongoing debate regarding whether this capability arises from mere memorisation (Golchin et al., 2024) or from the implementation of an algorithm within the LMs weights and parameters during pre-training (Olsson et al., 2022; Lindner et al., 2023). Our findings suggest that this may not be a strict dichotomy; rather, it could be a compositional phenomenon in which both processes operate concurrently.

Furthermore, the fact that counterfactual prompts can cause greater effects in contextual entrainment suggests that current models are more prone to distraction from counterfactual context prompts. This highlights the potential threat of dis- and misinformation.

4 Entrainment Heads

Recent research presents the argument that attention heads play the crucial role in controlling the LMs’ utilisation of context (Wang et al., 2022; Meng et al., 2022; Jin et al., 2024; Crosbie and Shutova, 2024; Yu et al., 2024a, inter alia). Notably, Jin et al. (2024) studied a similar phenomenon, termed knowledge conflicts, which examines how models react when information from the context prompt contradicts the information acquired during pre-training. They use the terms internal memory and external context to refer to these two types of information. Furthermore, they identify two types of attention heads — memory heads and context heads — that correspond to the LM’s utilisation of these distinct sources of information.

Knowledge conflict appears very similar to our counterfactual experiment; however, our research differs in several places. First and foremost, our prompt setting does not present a conflict. While we might both include a piece of counterfactual
information in the context (e.g., The capital city of Germany is Moscow), we will ask the model to answer a question unrelated to either Germany or Moscow (e.g., we would query the model with The capital of Nigeria is the city of __); whereas knowledge conflict would query the model with a question directly related to Germany or Moscow (e.g., The capital of Germany/Russia is the city of __). Second, counterfactual experiments are part of our investigation into the contextual entrainment phenomenon, whereas researchers who study information conflict focus solely on scenarios where such conflicts arise. Third, while identifying this novel phenomenon of contextual entrainment, we seek to understand the process by which LMs utilise information from the context prompt. In contrast, studies on information conflict focus more on the applicational aspect, where the desideratum is to find a way to ensure that LLMs can effectively resolve conflicting information and generate outputs that align with the intended factuality or coherence of the given context. Nevertheless, there are several aspects in which our research can mutually inform and benefit from one another.

In particular, Jin et al. (2024) found that attention heads play a key role in utilizing contextual information. As we will show later in this section, our experimental results further confirm this finding. However, their method is limited to investigating each attention head in isolation and do not consider the interaction between attention heads. Moreover, an increasing number of studies (Niu et al., 2024; Yu et al., 2024b; Bhaskar et al., 2024) have identified issues with this individual approach, as it disregards the intricate structures of transformer LMs, and have advocated for a more holistic analysis of the entire computational “circuit.” Inspired by this line of research, we adapt the differentiable masking based approach proposed by Yu et al. (2024b); Bhaskar et al. (2024) to identify the set of attention heads responsible for contextual entrainment. Our approach yields better results than the method proposed by Jin et al. (2024), suggesting that a circuit of attention heads — the “entrainment heads” — may have been formed in the model to process contextual information.

4.1 Entrainment Heads Discovery

Inspired by Yu et al. (2024b), we propose an automatic method to identify attention heads corresponding to contextual entrainment for each task configuration. We “turn off” specific heads by setting their contribution to the residual stream (Elhage et al., 2021)3 to zero. To achieve this, we introduce a binary mask \( m_j \) for each head \( h_j \), which selectively activates or deactivates heads (\( \sum_{h_j \in H_i} m_j h_j(x_i) \)). This mask is made differentiable by converting the sampled variable \( s_i \) from a Gumbel-sigmoid distribution using the straight-through estimator (Bengio et al., 2013):

\[
s_i = \sigma(\frac{l_i - \log \frac{\log U_1}{\log U_2}}{\tau}); m_i = [\mathbf{1}_{s_i > \frac{1}{2}} - s_i]_{\text{detach}} + s_i,
\]

where \( \tau \in (0, \infty) \) is a temperature hyperparameter, \( l_i \) is a learnable logit of the sigmoid distribution \( \sigma(\cdot) \), and \( U_1, U_2 \sim \text{Uniform}(0, 1) \) are random variables drawn from a uniform distribution.

We can then apply gradient descent to a dataset to identify the optimal combinations of attention heads to disable in order to suppress contextual entrainment. Our objective is to determine the set of attention heads that contribute the most to contextual entrainment while minimising the number of heads used, using the following loss function:

\[
\mathcal{L} = \underbrace{\ell(\mathbb{I}) - \ell(\mathbb{J})}_{\text{Logits } \Delta} + \lambda \cdot \underbrace{\frac{1}{|H|} \sum_{i=1}^{|H|} \sigma(l_i)}_{\text{Sparsity Loss}} .
\]

Experimental Setup  We conduct our entrainment head experiments using llama-3.1-8B, which has 1,024 attention heads (32 layers × 32 heads). Each LRE relation is split into training (80%), development (10%), and test (10%) sets. Entrainment heads are identified using the training set over 500 epochs,4 selecting the epoch with the best effect and fewest heads.5 All results are reported on the test set, which the model has neither seen nor used for hyperparameter search or checkpoint selection.

4.2 Experiment Results & Analysis

We first present our findings through the case study using the country–capital city relation. The remaining relations support the same findings, we will present them collectively at the end of this section.

“Turning off” the entrainment heads drastically reduce contextual entrainment.  Our algorithm identified 36 entrainment heads for the country–capital city relation.6 When “turning these entrain-

3We briefly review residual stream in Appendix C.
4We use the AdamW optimiser (Loshchilov and Hutter, 2019) with \( \lambda = 1.0, \tau = 1.0 \), and a learning rate of 1.0.
5Specifically, we use the epoch with the maximum logit difference + number of heads × 0.1.
6The entrainment heads identified is generally stable between runs, see Appendix D.
Figure 7: Effects of “Removing” the Entrainment Heads. The figures show the top 10 token probabilities for their respective settings. (a,b) In the original model, when a piece of counterfactual information is presented, the model assigns higher probabilities to the distractions: Berlin and Moscow. (c) After setting the output of the identified entrainment heads to zero, however, the effect of contextual entrainment is drastically attenuated. (d) This operation of removing the entrainment heads has only a small impact on other capabilities; the model can still correctly answer questions in other domains.

<table>
  <tr>
    <th>Measure</th>
    <th>No 😊</th>
    <th>With 😊</th>
    <th>No ☺</th>
    <th>With ☺</th>
  </tr>
  <tr>
    <td>\( \ell(\text{😊}) \)</td>
    <td>19.51</td>
    <td>20.68</td>
    <td>19.49</td>
    <td>21.21</td>
  </tr>
  <tr>
    <td>\( \ell(\text{☺}) \)</td>
    <td>8.75</td>
    <td>12.99</td>
    <td>7.87</td>
    <td>8.01</td>
  </tr>
  <tr>
    <td>\( \Delta = \ell(\text{😊}) - \ell(\text{☺}) \)</td>
    <td>10.76</td>
    <td>7.69</td>
    <td>11.62</td>
    <td>13.20</td>
  </tr>
  <tr>
    <td>Avg. 😊 Token Rank</td>
    <td>1.00</td>
    <td>1.00</td>
    <td>1.00</td>
    <td>1.00</td>
  </tr>
  <tr>
    <td>Avg. ☺ Token Rank*</td>
    <td>1756.7</td>
    <td>37.5</td>
    <td>1707.3</td>
    <td>1289.6</td>
  </tr>
</table>

Table 2: Effects of “Removing” the Entrainment Heads across the Entire Country–Capital City Relation Test Set. Removing the entrainment heads caused a significant effect across logits delta and the ranks of the 😊 tokens, making them capitulate the situation when no distracting context is provided. *: \( p < 6.9 \times 10^{-54} \) according to paired t-tests conditions between 😊 and ☺.

<table>
  <tr>
    <th>Relation</th>
    <th># Heads (Density)</th>
    <th>\( \ell(\text{😊}) - \ell(\text{☺}) \)</th>
  </tr>
  <tr>
    <td>company hq</td>
    <td>90 (8.8%)</td>
    <td>3.94 \( \Rightarrow \) 14.68</td>
  </tr>
  <tr>
    <td>country capital city</td>
    <td>36 (3.5%)</td>
    <td>7.69 \( \Rightarrow \) 13.20</td>
  </tr>
  <tr>
    <td>country currency</td>
    <td>42 (4.1%)</td>
    <td>4.73 \( \Rightarrow \) 11.67</td>
  </tr>
  <tr>
    <td>country language</td>
    <td>30 (2.9%)</td>
    <td>6.20 \( \Rightarrow \) 8.95</td>
  </tr>
  <tr>
    <td>country largest city</td>
    <td>33 (3.2%)</td>
    <td>8.68 \( \Rightarrow \) 13.35</td>
  </tr>
  <tr>
    <td>food from country</td>
    <td>38 (3.7%)</td>
    <td>3.98 \( \Rightarrow \) 9.95</td>
  </tr>
  <tr>
    <td>fruit inside color</td>
    <td>56 (5.5%)</td>
    <td>0.97 \( \Rightarrow \) 11.16</td>
  </tr>
  <tr>
    <td>fruit outside color</td>
    <td>80 (7.8%)</td>
    <td>2.14 \( \Rightarrow \) 13.82</td>
  </tr>
  <tr>
    <td>landmark in country</td>
    <td>59 (5.8%)</td>
    <td>3.93 \( \Rightarrow \) 9.68</td>
  </tr>
  <tr>
    <td>landmark on continent</td>
    <td>52 (5.1%)</td>
    <td>2.51 \( \Rightarrow \) 9.14</td>
  </tr>
  <tr>
    <td>product by company</td>
    <td>110 (10.7%)</td>
    <td>3.62 \( \Rightarrow \) 16.47</td>
  </tr>
  <tr>
    <td>star constellation name</td>
    <td>72 (7.0%)</td>
    <td>1.07 \( \Rightarrow \) 8.87</td>
  </tr>
  <tr>
    <td>task done by tool</td>
    <td>66 (6.4%)</td>
    <td>4.70 \( \Rightarrow \) 12.31</td>
  </tr>
  <tr>
    <td>task person type</td>
    <td>41 (4.0%)</td>
    <td>6.51 \( \Rightarrow \) 12.47</td>
  </tr>
  <tr>
    <td>work location</td>
    <td>68 (6.6%)</td>
    <td>3.17 \( \Rightarrow \) 12.68</td>
  </tr>
</table>

Table 3: Entrainment Head Identified across All LRE Relations. A small set of attention heads (3.2% to 10.7%) can substantially increase the gap between the logits of 😊 and ☺ tokens, i.e., attenuate contextual entrainment. This may suggest that these heads play a crucial role in the presence of contextual entrainment and can have broader implication in understanding how LMs utilise context information from prompts.

ment heads off,” i.e., setting their output to zero, the model attenuates the effect of contextual entrainment, as illustrated in Figure 7. Normally, Llama-3.1-8B is distracted by the counterfactual context The capital city of Germany is Moscow, confirming our previous findings in §3.3. However, after “turning off” the entrainment heads, the contextual entrainment effect is substantially attenuated. The rankings of the tokens Berlin and Moscow dropped from 2nd and 4th to 53rd and 68th, respectively. Additionally, the difference in logits between the correct 😊 token (Abu)7 and the distracting ☺ tokens (Berlin and Moscow) increased from 0.86 and 2.00 to 7.54 and 7.79, respectively. This attenuation is not a fluke. Table 2 shows that removing the entrainment heads significantly shifts the logits difference, probability difference, and the ranks of the distracting ☺ tokens towards the values observed when no distracting context is provided, across the entire country–capital city relation test set.

Table 3 shows that our differentiable-masking-

7The first wordpiece of Abuja, the capital city of Nigeria.

based entrainment head discovery method applies to other relations in the LRE dataset, highlighting the generalisability of our approach. We observe an increase in logit differences on the same scale as our country–capital city case study, further supporting its findings. This result also suggests that entrainment head are relation-specific

Entrainment heads are task-specific (or relation-specific), not model-specific. Again, as we have shown in Figure 3, our method has identified different and different amount of entrainment heads for each LRE relation. This suggests that the entrainment heads are task-specific (or relation-specific), rather than model-specific. We do observe some
level of overlap in the function of entrainment heads across different relations (Appendix F).

This finding unveils a limitation of prior work that relies on individual approaches. Specifically, Jin et al. (2024) classified attention heads as either memory heads or context heads, implying a strong and discrete assignment of functional roles to specific heads in a model-specific manner. This is an assumption that cannot account for our results.

Removing the entrainment heads has only a small effect on other LM capabilities. While removing the entrainment heads significantly impacts contextual entrainment, it has only a negligible to small effect on other LM capabilities. First, we use other relations from the LRE dataset to evaluate whether the LM can still interpret the query and recall factual information, as well as perform ICL. Table 4 shows the performance of the original model (\( \mathbb{E} \)) and the modified model (\( \mathbb{E} \)) on all other relations without distracting context, demonstrating that the model can still perform factual recall. We report both strict (the correct answer appears within the top-3 predicted tokens) and credulous (top-10) accuracy, as multiple correct answers may exist. Relying solely on whether the gold-standard token is the most probable leads to unstable results (Appendix E). Moreover, we experiment with the three ICL tasks identified by Brown et al. (2020): arithmetic, spelling correction, and translation (Appendix G). After removing the entrainment heads (\( \mathbb{E} \)), the model exhibits only a small performance decrease (0.2–3%) and continues demonstrate strong ICL capabilities with high accuracy in the same ballpark with \( \mathbb{E} \). Appendix H shows removing the entrainment heads also have a small effect on the LMs’ general capabilities.

This finding supports our hypothesis that this “circuit” of entrainment heads collectively corresponds to contextual entrainment rather than other capabilities and phenomena, such as factual recall, and is not strongly related to how LMs process and utilise contextual information or perform ICL more broadly. Thus, contextual entrainment and its connection to entrainment heads provides a novel perspective for understanding distraction. While current mitigation strategies (Yoran et al., 2023; Cuconasu et al., 2024; Wu et al., 2024) focus on methods external to the model — either modifying the context prompt or prompting the model to self-correct through reasoning — our findings suggest there could be a way to mitigate distraction by directly modifying or monitoring the internal mechanisms of LMs when performing RAG.

<table>
  <tr>
    <th>Relation</th>
    <th colspan="2">Strict Acc.</th>
    <th colspan="2">Credulous Acc.</th>
  </tr>
  <tr>
    <th></th>
    <th>\( \mathbb{E} \)</th>
    <th>\( \mathbb{E} \)</th>
    <th>\( \mathbb{E} \)</th>
    <th>\( \mathbb{E} \)</th>
  </tr>
  <tr><td>company hq</td><td>83.5%</td><td>90.0%</td><td>88.0%</td><td>90.0%</td></tr>
  <tr><td>country capital city</td><td>100.0%</td><td>100.0%</td><td>100.0%</td><td>100.0%</td></tr>
  <tr><td>country currency</td><td>83.7%</td><td>100.0%</td><td>100.0%</td><td>100.0%</td></tr>
  <tr><td>country language</td><td>85.7%</td><td>100.0%</td><td>100.0%</td><td>100.0%</td></tr>
  <tr><td>country largest city</td><td>100.0%</td><td>100.0%</td><td>100.0%</td><td>100.0%</td></tr>
  <tr><td>food from country</td><td>92.0%</td><td>98.5%</td><td>100.0%</td><td>98.5%</td></tr>
  <tr><td>fruit inside color</td><td>77.0%</td><td>100.0%</td><td>98.0%</td><td>100.0%</td></tr>
  <tr><td>fruit outside color</td><td>38.0%</td><td>84.0%</td><td>82.0%</td><td>84.0%</td></tr>
  <tr><td>landmark in country</td><td>89.5%</td><td>91.0%</td><td>95.0%</td><td>91.0%</td></tr>
  <tr><td>landmark on continent</td><td>88.5%</td><td>91.0%</td><td>97.0%</td><td>93.0%</td></tr>
  <tr><td>product by company</td><td>95.0%</td><td>96.0%</td><td>98.0%</td><td>96.0%</td></tr>
  <tr><td>star constellation name</td><td>84.7%</td><td>89.3%</td><td>92.3%</td><td>89.3%</td></tr>
  <tr><td>task done by tool</td><td>78.0%</td><td>91.0%</td><td>93.5%</td><td>91.0%</td></tr>
  <tr><td>task person type</td><td>78.5%</td><td>80.0%</td><td>80.0%</td><td>80.0%</td></tr>
  <tr><td>work location</td><td>60.5%</td><td>75.0%</td><td>75.0%</td><td>75.0%</td></tr>
  <tr><td>arithmetic 0-shot</td><td>100.0%</td><td>100.0%</td><td>100.0%</td><td>100.0%</td></tr>
  <tr><td>spelling correction 1-shot</td><td>73.6%</td><td>72.0%</td><td>78.6%</td><td>76.8%</td></tr>
  <tr><td>spelling correction 2-shot</td><td>94.6%</td><td>91.6%</td><td>97.0%</td><td>94.8%</td></tr>
  <tr><td>spelling correction 5-shot</td><td>99.0%</td><td>98.4%</td><td>100.0%</td><td>100.0%</td></tr>
  <tr><td>translation 1-shot</td><td>74.4%</td><td>73.0%</td><td>78.4%</td><td>76.8%</td></tr>
  <tr><td>translation 2-shot</td><td>94.0%</td><td>93.0%</td><td>97.0%</td><td>96.2%</td></tr>
  <tr><td>translation 3-shot</td><td>98.6%</td><td>97.2%</td><td>99.6%</td><td>99.4%</td></tr>
</table>

Table 4: Removing the entrainment heads of the country–capital city relation has a small to negligible effect on other LM capabilities. This table compares the strict (answer in top-3) and credulous (answer in top-10) accuracy of the original model (\( \mathbb{E} \)) and the model with country–capital city entrainment heads removed (\( \mathbb{E} \)). Removing these heads has a negligible effect on the LM’s performance across other relations, with no obvious differences between \( \mathbb{E} \) and \( \mathbb{E} \).

5 Conclusion

Llama see, llama do. We observe and confirm contextual entrainment, a novel phenomenon. If a token has appeared previously in the prompt, the model assigns a higher logit to that token, even for random tokens. Thus, a novel mechanistic effect may be at play in governing how LMs process and utilise information from the prompt—an effect that is analogous to but distinct from previously identified phenomena, such as the inductive pattern repetition effect observed by Olsson et al.’s (2022). However, we also discover that contextual entrainment is influenced by semantic factors. This finding highlights the potential threat of dis- and misinformation, which may be more severe than mere mistakes generated by the model. It also suggests that there may not be a strict dichotomy between mechanistic and statistical interpretations of LMs. Our identification of the entrainment heads suggests that interpretability techniques could provide crucial insights for real-world applications, such as the study of distraction.
Limitations

We did not conduct our experiments using larger LMs such as Llama-3.1-70B and Llama-3.1-405B due to resource limitations. However, we have used Llama-3.1-8B, Llama-3.1-8B-Instruct, and Llama-2-13B-hf, as these models are sufficiently large and powerful for our experiments. Moreover, we observed no differences in findings between larger and smaller models, such as GPT-2 XL and Llama-2-13B-hf, suggesting that our results are not significantly affected by model scale within this range. We encourage others to reproduce our work using larger models to further validate our findings.

Our experimental setup is rigorous. However, since RAG is most relevant to the problem of distraction, we conducted experiments using only the LRE dataset in this setting. We did not use standard RAG datasets (e.g., Simple-QA (Wei et al., 2024)), as they are difficult to control and compare fairly. Nonetheless, our experiment with random token inputs provides strong evidence — if such a setup yields successful results, then more structured approaches are unlikely to fail. Once again, we encourage others to reproduce our results using these datasets to further validate our findings.

Finally, while demonstrating two novel and insightful findings — the contextual entrainment phenomenon and the identification of entrainment heads — we do not propose an application to mitigate the distraction problem. We believe our findings serve as foundational steps toward addressing this issue. Given the depth of contributions presented in this work, we leave such applications for future research.

Acknowledgements

We thank Alessandro Sordoni and Lucas Caccia (Microsoft Research) as well as Gerald Penn and Lei Yu (University of Toronto) for their valuable discussions and feedback. We are also grateful to Zhanao Fu (University of Toronto) for introducing us to relevant background on brain entrainment. Finally, we thank the reviewers and area chairs for their insightful suggestions and engaging discussions that helped improve this work.

References

Yoshua Bengio, Nicholas Léonard, and Aaron Courville. 2013. Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation. Preprint, arXiv:1308.3432.

Lukas Berglund, Meg Tong, Maximilian Kaufmann, Mikita Balesni, Asa Cooper Stickland, Tomasz Korbak, and Owain Evans. 2024. The Reversal Curse: LLMs trained on “A is B” fail to learn “B is A”. In The Twelfth International Conference on Learning Representations.

Adithya Bhaskar, Alexander Wettig, Dan Friedman, and Danqi Chen. 2024. Finding transformer circuits with edge pruning. In The Thirty-Eighth Annual Conference on Neural Information Processing Systems.

Susan E. Brennan and Herbert H. Clark. 1996. Conceptual pacts and lexical choice in conversation. Journal of Experimental Psychology: Learning, Memory, and Cognition, 22(6):1482–1493.

Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. 2020. Language Models are Few-Shot Learners. arXiv:2005.14165 [cs].

Joy Crosbie and Ekaterina Shutova. 2024. Induction Heads as an Essential Mechanism for Pattern Matching in In-context Learning. Preprint, arXiv:2407.07011.

Florin Cuconasu, Giovanni Trappolini, Federico Siciliano, Simone Filice, Cesare Campagnano, Yoelle Maarek, Nicola Tonellotto, and Fabrizio Silvestri. 2024. The Power of Noise: Redefining Retrieval for RAG Systems. In Proceedings of the 47th International ACM SIGIR Conference on Research and Development in Information Retrieval, pages 719–729.

Yanai Elazar, Nora Kassner, Shauli Ravfogel, Abhilasha Ravichander, Eduard Hovy, Hinrich Schütze, and Yoav Goldberg. 2021. Measuring and Improving Consistency in Pretrained Language Models. Transactions of the Association for Computational Linguistics, 9:1012–1031.

Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. 2021. A mathematical framework for transformer circuits. Transformer Circuits Thread.

W. N. Francis and H. Kucera. 1979. Brown corpus manual. Technical report, Department of Linguistics, Brown University, Providence, Rhode Island, US.
Simon Garrod and Anthony Anderson. 1987. Saying what you mean in dialogue: A study in conceptual and semantic co-ordination. Cognition, 27(2):181–218.

Shahriar Golchin, Mihai Surdeanu, Steven Bethard, Eduardo Blanco, and Ellen Riloff. 2024. Memorization in In-Context Learning. Preprint, arXiv:2408.11546.

Evan Hernandez, Arnab Sen Sharma, Tal Haklay, Kevin Meng, Martin Wattenberg, Jacob Andreas, Yonatan Belinkov, and David Bau. 2024. Linearity of relation decoding in transformer language models. In The Twelfth International Conference on Learning Representations.

Zhuoran Jin, Pengfei Cao, Hongbang Yuan, Yubo Chen, Jiexin Xu, Huaijun Li, Xiaojian Jiang, Kang Liu, and Jun Zhao. 2024. Cutting Off the Head Ends the Conflict: A Mechanism for Interpreting and Mitigating Knowledge Conflicts in Language Models. In Findings of the Association for Computational Linguistics: ACL 2024, pages 1193–1215, Bangkok, Thailand. Association for Computational Linguistics.

Zhoubo Li, Ningyu Zhang, Yunzhi Yao, Mengru Wang, Xi Chen, and Huaun Chen. 2024. Unveiling the pitfalls of knowledge editing for large language models. In The Twelfth International Conference on Learning Representations.

David Lindner, Janos Kramar, Sebastian Farquhar, Matthew Rahtz, Thomas McGrath, and Vladimir Mikulik. 2023. Tracr: Compiled transformers as a laboratory for interpretability. In Thirty-Seventh Conference on Neural Information Processing Systems.

Ilya Loshchilov and Frank Hutter. 2019. Decoupled Weight Decay Regularization. In ICLR 2019.

Thomas McGrath, Matthew Rahtz, Janos Kramar, Vladimir Mikulik, and Shane Legg. 2023. The Hydra Effect: Emergent Self-repair in Language Model Computations. Preprint, arXiv:2307.15771.

Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. 2022. Locating and editing factual associations in GPT. Advances in Neural Information Processing Systems, 36.

Jingcheng Niu, Andrew Liu, Zining Zhu, and Gerald Penn. 2024. What does the knowledge neuron thesis have to do with knowledge? In The Twelfth International Conference on Learning Representations.

Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. 2022. In-context learning and induction heads. Transformer Circuits Thread.

Alejandro Pérez, Matthew H. Davis, Robin A.A. Ince, Hanna Zhang, Zhanao Fu, Melanie Lamarca, Matthew A. Lambon Ralph, and Philip J. Monahan. 2022. Timing of brain entrainment to the speech envelope during speaking, listening and self-listening. Cognition, 224:105051.

David Poeppel and M. Florencia Assaneo. 2020. Speech rhythms and their neural foundations. Nature Reviews Neuroscience, 21(6):322–334.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. 2019. Language Models are Unsupervised Multitask Learners. OpenAI Blog, page 24.

Freda Shi, Xinyun Chen, Kanishka Misra, Nathan Scales, David Dohan, Ed H. Chi, Nathanael Schärli, and Denny Zhou. 2023. Large Language Models Can Be Easily Distracted by Irrelevant Context. In Proceedings of the 40th International Conference on Machine Learning, pages 31210–31227. PMLR.

Student. 1908. The Probable Error of a Mean. Biometrika, 6(1):1–25.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. 2023. LLaMA: Open and Efficient Foundation Language Models. Preprint, arXiv:2302.13971.

Kevin Ro Wang, Alexandre Variengien, Arthur Conny, Buck Shlegeris, and Jacob Steinhardt. 2022. Interpretability in the Wild: A Circuit for Indirect Object Identification in GPT-2 Small. In The Eleventh International Conference on Learning Representations.

Jason Wei, Nguyen Karina, Hyung Won Chung, Yunxin Joy Jiao, Spencer Papay, Amelia Glaese, John Schulman, and William Fedus. 2024. Measuring short-form factuality in large language models. Preprint, arXiv:2411.04368.

Siye Wu, Jian Xie, Jiangjie Chen, Tinghui Zhu, Kai Zhang, and Yanghua Xiao. 2024. How Easily do Irrelevant Inputs Skew the Responses of Large Language Models? In First Conference on Language Modeling.

Ori Yoran, Tomer Wolfson, Ori Ram, and Jonathan Berant. 2023. Making Retrieval-Augmented Language Models Robust to Irrelevant Context. In The Twelfth International Conference on Learning Representations.

Lei Yu, Meng Cao, Jackie CK Cheung, and Yue Dong. 2024a. Mechanistic Understanding and Mitigation of Language Model Non-Factual Hallucinations. In Findings of the Association for Computational Linguistics: EMNLP 2024, pages 7943–7956, Miami, Florida, USA. Association for Computational Linguistics.
Lei Yu, Jingcheng Niu, Zining Zhu, and Gerald Penn. 2024b. Functional Faithfulness in the Wild: Circuit Discovery with Differentiable Computation Graph Pruning. Preprint, arXiv:2407.03779.
<table>
  <tr>
    <th>Relation</th>
    <th># Samples</th>
    <th>Context Templates</th>
    <th>Query Templates</th>
  </tr>
  <tr>
    <td>company hq</td>
    <td>674</td>
    <td>The headquarters of {} is in the city of<br>Where are the headquarters of {}? It is in the city of</td>
    <td>{} is headquartered in the city of<br>The headquarters of {} are in the city of</td>
  </tr>
  <tr>
    <td>country capital city</td>
    <td>24</td>
    <td>The capital of {} is the city of<br>What is the capital of {}? It is the city of</td>
    <td>The capital city of {} is<br>The capital of {} is</td>
  </tr>
  <tr>
    <td>country currency</td>
    <td>30</td>
    <td>What is the official currency of {}? It is called the {}’s official currency is called the<br>The name of {}’s currency is the</td>
    <td>The official currency of {} is the<br>{}’s official currency is the</td>
  </tr>
  <tr>
    <td>country language</td>
    <td>24</td>
    <td>{}, where most people speak<br>In {}, people speak the language of<br>People in {} speak the language of</td>
    <td>People in {} speak<br>The language used in {} is<br>In {}, the primary language is</td>
  </tr>
  <tr>
    <td>country largest city</td>
    <td>24</td>
    <td>What is the largest city in {}?<br>It is the city of<br>The largest city in {} is the city of</td>
    <td>The largest city in {} is<br>The biggest city in {} is</td>
  </tr>
  <tr>
    <td>food from country</td>
    <td>30</td>
    <td>What is the country of origin for {}? It originates from<br>{} originates from the country of</td>
    <td>{} originates from<br>{} is from the country of</td>
  </tr>
  <tr>
    <td>fruit inside color</td>
    <td>36</td>
    <td>What color are {} on the inside? They are</td>
    <td>On the inside, {} are</td>
  </tr>
  <tr>
    <td>fruit outside color</td>
    <td>30</td>
    <td>What color are {} on the outside? They are the color of</td>
    <td>On the outside, {} are</td>
  </tr>
  <tr>
    <td>landmark in country</td>
    <td>836</td>
    <td>What country is {} in? It is in {} is in the country of</td>
    <td>{} is in the country of</td>
  </tr>
  <tr>
    <td>landmark on continent</td>
    <td>947</td>
    <td>What continent is {} on? It is on<br>{} is on the continent of</td>
    <td>{} is on the continent of</td>
  </tr>
  <tr>
    <td>product by company</td>
    <td>522</td>
    <td>Which company developed {}? It was developed by</td>
    <td>{} was created by<br>{} is a product of</td>
  </tr>
  <tr>
    <td>star constellation name</td>
    <td>362</td>
    <td>What is the name of the constellation that {} is part of?<br>It is part of<br>{} is part of the constellation named<br>What is the name of the constellation that {} belongs to?<br>It belongs to</td>
    <td>{} is part of the constellation named</td>
  </tr>
  <tr>
    <td>task person type</td>
    <td>32</td>
    <td>The task of {} would be best performed by someone with the role of a<br>The professional role most suited to handle {} is a</td>
    <td>{} is best suited for someone with the role of a</td>
  </tr>
  <tr>
    <td>task done by tool</td>
    <td>52</td>
    <td>What tool is used for {}? Usually, you need a<br>To accomplish {}, you need a tool called a</td>
    <td>The tool used for {} is called a</td>
  </tr>
  <tr>
    <td>work location</td>
    <td>38</td>
    <td>A {} typically works at a<br>You can usually find a {} working in a</td>
    <td>A {} typically works at a</td>
  </tr>
</table>

Table 5: Selected LRE Relations.
<table>
  <tr>
    <th rowspan="2">Setting</th>
    <th colspan="3">\( \ell(\smile) \)</th>
    <th colspan="3">\( \ell(\frown) \)</th>
    <th colspan="3">\( P(\smile) \)</th>
    <th colspan="3">\( P(\frown) \)</th>
  </tr>
  <tr>
    <th>No CTX</th><th>With CTX</th><th>\( \Delta \)</th>
    <th>No CTX</th><th>With CTX</th><th>\( \Delta \)</th>
    <th>No CTX</th><th>With CTX</th><th>\( \Delta \)</th>
    <th>No CTX</th><th>With CTX</th><th>\( \Delta \)</th>
  </tr>
  <tr>
    <th colspan="12">Llama-3.1-8B</th>
  </tr>
  <tr>
    <td>Distraction</td><td>16.72</td><td>17.47</td><td>0.75</td><td>8.94</td><td>13.02</td><td>4.08</td><td>0.39</td><td>0.47</td><td>0.08</td><td>4.83e-03</td><td>0.03</td><td>0.02</td>
  </tr>
  <tr>
    <td>Irrelevant</td><td>16.68</td><td>15.45</td><td>-1.23</td><td>3.05</td><td>7.96</td><td>4.92</td><td>0.38</td><td>0.35</td><td>-0.03</td><td>8.44e-05</td><td>2.51e-03</td><td>2.42e-03</td>
  </tr>
  <tr>
    <td>Random</td><td>16.69</td><td>15.52</td><td>-1.17</td><td>3.19</td><td>7.19</td><td>4.01</td><td>0.38</td><td>0.34</td><td>-0.05</td><td>3.41e-04</td><td>3.55e-03</td><td>3.19e-03</td>
  </tr>
  <tr>
    <th colspan="12">Llama-3.1-8B-Instruct</th>
  </tr>
  <tr>
    <td>Distraction</td><td>15.75</td><td>17.63</td><td>1.88</td><td>8.64</td><td>12.07</td><td>3.43</td><td>0.27</td><td>0.43</td><td>0.16</td><td>5.13e-03</td><td>0.02</td><td>0.01</td>
  </tr>
  <tr>
    <td>Irrelevant</td><td>15.74</td><td>15.84</td><td>0.11</td><td>2.19</td><td>6.78</td><td>4.58</td><td>0.27</td><td>0.34</td><td>0.07</td><td>1.31e-05</td><td>7.45e-04</td><td>7.32e-04</td>
  </tr>
  <tr>
    <td>Random</td><td>15.74</td><td>14.78</td><td>-0.95</td><td>2.66</td><td>7.00</td><td>4.33</td><td>0.27</td><td>0.24</td><td>-0.03</td><td>3.54e-05</td><td>3.55e-03</td><td>3.52e-03</td>
  </tr>
  <tr>
    <th colspan="12">Llama-2-13b-hf</th>
  </tr>
  <tr>
    <td>Distraction</td><td>14.33</td><td>15.17</td><td>0.84</td><td>7.76</td><td>11.39</td><td>3.63</td><td>0.32</td><td>0.44</td><td>0.12</td><td>6.35e-03</td><td>0.04</td><td>0.04</td>
  </tr>
  <tr>
    <td>Irrelevant</td><td>14.28</td><td>13.56</td><td>-0.72</td><td>3.14</td><td>6.98</td><td>3.84</td><td>0.32</td><td>0.31</td><td>-0.20e-03</td><td>2.53e-04</td><td>3.35e-03</td><td>3.10e-03</td>
  </tr>
  <tr>
    <td>Random</td><td>14.29</td><td>13.19</td><td>-1.10</td><td>3.41</td><td>7.32</td><td>3.91</td><td>0.32</td><td>0.26</td><td>-0.06</td><td>3.58e-04</td><td>9.73e-03</td><td>9.37e-03</td>
  </tr>
  <tr>
    <th colspan="12">Llama-2-7b-hf</th>
  </tr>
  <tr>
    <td>Distraction</td><td>16.67</td><td>17.87</td><td>1.19</td><td>9.44</td><td>13.54</td><td>4.10</td><td>0.40</td><td>0.47</td><td>0.07</td><td>5.67e-03</td><td>0.04</td><td>0.03</td>
  </tr>
  <tr>
    <td>Irrelevant</td><td>16.60</td><td>15.54</td><td>-1.06</td><td>5.06</td><td>8.11</td><td>3.05</td><td>0.40</td><td>0.36</td><td>-0.03</td><td>3.57e-04</td><td>3.74e-03</td><td>3.39e-03</td>
  </tr>
  <tr>
    <td>Random</td><td>16.61</td><td>14.43</td><td>-2.18</td><td>5.50</td><td>7.39</td><td>1.89</td><td>0.40</td><td>0.32</td><td>-0.08</td><td>4.27e-04</td><td>5.53e-03</td><td>5.10e-03</td>
  </tr>
  <tr>
    <th colspan="12">GPT2 XL</th>
  </tr>
  <tr>
    <td>Distraction</td><td>9.17</td><td>10.13</td><td>0.96</td><td>4.18</td><td>7.41</td><td>3.23</td><td>0.18</td><td>0.26</td><td>0.08</td><td>7.76e-03</td><td>0.05</td><td>0.05</td>
  </tr>
  <tr>
    <td>Irrelevant</td><td>9.16</td><td>9.06</td><td>-0.11</td><td>-1.55</td><td>3.04</td><td>4.59</td><td>0.18</td><td>0.17</td><td>-0.01</td><td>6.86e-05</td><td>5.28e-03</td><td>5.21e-03</td>
  </tr>
  <tr>
    <td>Random</td><td>9.16</td><td>8.79</td><td>-0.37</td><td>-1.72</td><td>-0.28</td><td>1.44</td><td>0.18</td><td>0.16</td><td>-0.02</td><td>1.84e-04</td><td>2.31e-03</td><td>2.13e-03</td>
  </tr>
</table>

Table 6: Average Logits and Probabilities across All LRE Relations and Prompt Settings.
<table>
  <tr>
    <th rowspan="2">Relation</th>
    <th rowspan="2">Setting</th>
    <th colspan="3">\( \ell(\mathbb{1}) \)</th>
    <th colspan="3">\( \ell(\mathbb{0}) \)</th>
    <th colspan="3">\( P(\mathbb{1}) \)</th>
    <th colspan="3">\( P(\mathbb{0}) \)</th>
  </tr>
  <tr>
    <th>No CTX</th>
    <th>With CTX</th>
    <th>\( \Delta \)</th>
    <th>No CTX</th>
    <th>With CTX</th>
    <th>\( \Delta \)</th>
    <th>No CTX</th>
    <th>With CTX</th>
    <th>\( \Delta \)</th>
    <th>No CTX</th>
    <th>With CTX</th>
    <th>\( \Delta \)</th>
  </tr>
  <tr>
    <td rowspan="3">city in country</td>
    <td>dstr</td>
    <td>19.56</td><td>20.73</td><td>1.17</td><td>9.84</td><td>14.37</td><td>4.53</td><td>0.75</td><td>0.81</td><td>0.06</td><td>6.38e-03</td><td>0.03</td><td>0.02</td>
  </tr>
  <tr>
    <td>irr</td>
    <td>19.49</td><td>19.53</td><td>0.04</td><td>5.54</td><td>9.80</td><td>4.27</td><td>0.75</td><td>0.77</td><td>0.03</td><td>5.49e-05</td><td>2.38e-03</td><td>2.32e-03</td>
  </tr>
  <tr>
    <td>random</td>
    <td>19.52</td><td>16.96</td><td>-2.56</td><td>3.68</td><td>7.47</td><td>3.79</td><td>0.75</td><td>0.63</td><td>-0.12</td><td>2.33e-05</td><td>1.95e-03</td><td>1.93e-03</td>
  </tr>
  <tr>
    <td rowspan="3">company hq</td>
    <td>dstr</td>
    <td>16.59</td><td>16.93</td><td>0.34</td><td>7.63</td><td>13.31</td><td>5.68</td><td>0.45</td><td>0.46</td><td>6.30e-03</td><td>1.50e-03</td><td>0.06</td><td>0.06</td>
  </tr>
  <tr>
    <td>irr</td>
    <td>16.66</td><td>15.43</td><td>-1.23</td><td>4.34</td><td>9.64</td><td>5.30</td><td>0.46</td><td>0.36</td><td>-0.09</td><td>1.22e-04</td><td>6.17e-03</td><td>6.05e-03</td>
  </tr>
  <tr>
    <td>random</td>
    <td>16.63</td><td>15.86</td><td>-0.77</td><td>2.74</td><td>5.59</td><td>2.85</td><td>0.45</td><td>0.41</td><td>-0.04</td><td>1.37e-04</td><td>3.26e-03</td><td>3.12e-03</td>
  </tr>
  <tr>
    <td rowspan="3">country capital city</td>
    <td>dstr</td>
    <td>18.42</td><td>20.03</td><td>1.60</td><td>7.48</td><td>11.95</td><td>4.47</td><td>0.83</td><td>0.91</td><td>0.08</td><td>1.47e-04</td><td>1.12e-03</td><td>9.73e-04</td>
  </tr>
  <tr>
    <td>irr</td>
    <td>18.43</td><td>17.33</td><td>-1.10</td><td>3.73</td><td>9.07</td><td>5.35</td><td>0.82</td><td>0.74</td><td>-0.08</td><td>2.18e-05</td><td>1.42e-03</td><td>1.40e-03</td>
  </tr>
  <tr>
    <td>random</td>
    <td>18.44</td><td>17.29</td><td>-1.16</td><td>3.59</td><td>8.68</td><td>5.09</td><td>0.82</td><td>0.74</td><td>-0.09</td><td>3.28e-05</td><td>1.98e-03</td><td>1.95e-03</td>
  </tr>
  <tr>
    <td rowspan="3">country currency</td>
    <td>dstr</td>
    <td>17.76</td><td>18.29</td><td>0.53</td><td>8.08</td><td>12.59</td><td>4.51</td><td>0.19</td><td>0.51</td><td>0.32</td><td>4.07e-04</td><td>7.58e-03</td><td>7.17e-03</td>
  </tr>
  <tr>
    <td>irr</td>
    <td>17.72</td><td>15.25</td><td>-2.48</td><td>2.47</td><td>6.46</td><td>3.99</td><td>0.19</td><td>0.18</td><td>-0.01</td><td>3.72e-04</td><td>2.00e-03</td><td>1.63e-03</td>
  </tr>
  <tr>
    <td>random</td>
    <td>17.73</td><td>16.41</td><td>-1.33</td><td>3.30</td><td>6.53</td><td>3.24</td><td>0.19</td><td>0.18</td><td>-0.01</td><td>8.41e-03</td><td>1.20e-03</td><td>1.13e-03</td>
  </tr>
  <tr>
    <td rowspan="3">country language</td>
    <td>dstr</td>
    <td>14.79</td><td>18.32</td><td>3.53</td><td>8.29</td><td>12.14</td><td>3.85</td><td>0.23</td><td>0.56</td><td>0.34</td><td>8.41e-03</td><td>0.01</td><td>2.97e-03</td>
  </tr>
  <tr>
    <td>irr</td>
    <td>14.79</td><td>17.50</td><td>2.71</td><td>3.04</td><td>8.10</td><td>5.05</td><td>0.23</td><td>0.56</td><td>0.34</td><td>8.68e-05</td><td>4.16e-03</td><td>4.07e-03</td>
  </tr>
  <tr>
    <td>random</td>
    <td>14.72</td><td>13.66</td><td>-1.07</td><td>3.26</td><td>8.64</td><td>5.38</td><td>0.22</td><td>0.18</td><td>-0.04</td><td>8.84e-04</td><td>0.01</td><td>0.01</td>
  </tr>
  <tr>
    <td rowspan="3">country largest city</td>
    <td>dstr</td>
    <td>18.11</td><td>19.68</td><td>1.58</td><td>7.38</td><td>11.16</td><td>3.78</td><td>0.67</td><td>0.79</td><td>0.12</td><td>1.22e-04</td><td>8.22e-04</td><td>7.00e-04</td>
  </tr>
  <tr>
    <td>irr</td>
    <td>18.13</td><td>17.15</td><td>-0.98</td><td>3.50</td><td>8.18</td><td>4.69</td><td>0.67</td><td>0.60</td><td>-0.08</td><td>1.58e-05</td><td>8.34e-04</td><td>8.19e-04</td>
  </tr>
  <tr>
    <td>random</td>
    <td>18.13</td><td>17.53</td><td>-0.60</td><td>3.73</td><td>9.40</td><td>5.67</td><td>0.67</td><td>0.64</td><td>-0.04</td><td>1.29e-04</td><td>3.13e-03</td><td>3.00e-03</td>
  </tr>
  <tr>
    <td rowspan="3">food from country</td>
    <td>dstr</td>
    <td>18.14</td><td>18.06</td><td>-0.08</td><td>8.99</td><td>12.73</td><td>3.73</td><td>0.54</td><td>0.55</td><td>0.08</td><td>1.70e-03</td><td>0.01</td><td>0.01</td>
  </tr>
  <tr>
    <td>irr</td>
    <td>18.12</td><td>16.25</td><td>-1.87</td><td>4.03</td><td>9.89</td><td>5.86</td><td>0.54</td><td>0.41</td><td>-0.13</td><td>3.28e-04</td><td>6.26e-03</td><td>5.93e-03</td>
  </tr>
  <tr>
    <td>random</td>
    <td>18.10</td><td>16.56</td><td>-1.54</td><td>3.40</td><td>6.25</td><td>2.85</td><td>0.53</td><td>0.43</td><td>-0.10</td><td>1.09e-04</td><td>6.43e-04</td><td>5.34e-04</td>
  </tr>
  <tr>
    <td rowspan="3">fruit inside color</td>
    <td>dstr</td>
    <td>15.78</td><td>15.57</td><td>-0.21</td><td>13.35</td><td>14.58</td><td>1.23</td><td>0.12</td><td>0.17</td><td>0.05</td><td>0.02</td><td>0.07</td><td>0.05</td>
  </tr>
  <tr>
    <td>irr</td>
    <td>15.77</td><td>12.82</td><td>-2.95</td><td>1.54</td><td>6.28</td><td>4.74</td><td>0.12</td><td>0.10</td><td>-0.02</td><td>7.85e-07</td><td>1.63e-04</td><td>1.62e-04</td>
  </tr>
  <tr>
    <td>random</td>
    <td>15.77</td><td>14.78</td><td>-0.99</td><td>3.02</td><td>6.86</td><td>3.84</td><td>0.12</td><td>0.14</td><td>0.01</td><td>1.06e-03</td><td>1.31e-03</td><td>2.47e-04</td>
  </tr>
  <tr>
    <td rowspan="3">fruit outside color</td>
    <td>dstr</td>
    <td>14.57</td><td>14.53</td><td>-0.03</td><td>11.47</td><td>12.05</td><td>0.58</td><td>0.06</td><td>0.08</td><td>0.02</td><td>9.13e-03</td><td>0.01</td><td>2.67e-03</td>
  </tr>
  <tr>
    <td>irr</td>
    <td>14.51</td><td>12.97</td><td>-1.55</td><td>3.40</td><td>8.52</td><td>5.12</td><td>0.06</td><td>0.05</td><td>-5.39e-03</td><td>4.31e-06</td><td>2.00e-03</td><td>1.99e-03</td>
  </tr>
  <tr>
    <td>random</td>
    <td>14.52</td><td>12.95</td><td>-1.57</td><td>3.58</td><td>7.92</td><td>4.35</td><td>0.06</td><td>0.04</td><td>-0.02</td><td>1.77e-03</td><td>6.01e-03</td><td>4.24e-03</td>
  </tr>
  <tr>
    <td rowspan="3">landmark in country</td>
    <td>dstr</td>
    <td>17.40</td><td>17.60</td><td>0.19</td><td>8.45</td><td>13.83</td><td>5.37</td><td>0.51</td><td>0.49</td><td>-0.02</td><td>1.23e-03</td><td>0.04</td><td>0.04</td>
  </tr>
  <tr>
    <td>irr</td>
    <td>17.44</td><td>15.82</td><td>-1.62</td><td>3.97</td><td>9.36</td><td>5.40</td><td>0.52</td><td>0.45</td><td>-0.07</td><td>1.93e-04</td><td>6.79e-03</td><td>6.59e-03</td>
  </tr>
  <tr>
    <td>random</td>
    <td>17.42</td><td>15.79</td><td>-1.62</td><td>3.30</td><td>6.55</td><td>3.25</td><td>0.51</td><td>0.43</td><td>-0.08</td><td>1.81e-04</td><td>8.17e-04</td><td>7.66e-04</td>
  </tr>
  <tr>
    <td rowspan="3">landmark on continent</td>
    <td>dstr</td>
    <td>17.59</td><td>17.46</td><td>-0.13</td><td>11.90</td><td>15.64</td><td>3.74</td><td>0.36</td><td>0.28</td><td>-0.08</td><td>8.79e-03</td><td>0.07</td><td>0.06</td>
  </tr>
  <tr>
    <td>irr</td>
    <td>17.06</td><td>15.59</td><td>-1.47</td><td>4.15</td><td>8.50</td><td>4.36</td><td>0.32</td><td>0.24</td><td>-0.09</td><td>2.39e-05</td><td>8.56e-04</td><td>8.32e-04</td>
  </tr>
  <tr>
    <td>random</td>
    <td>17.07</td><td>16.16</td><td>-0.91</td><td>3.43</td><td>6.48</td><td>3.04</td><td>0.33</td><td>0.29</td><td>-0.04</td><td>1.68e-04</td><td>5.98e-04</td><td>4.30e-04</td>
  </tr>
  <tr>
    <td rowspan="3">product by company</td>
    <td>dstr</td>
    <td>15.99</td><td>15.03</td><td>-0.96</td><td>6.83</td><td>11.13</td><td>4.31</td><td>0.49</td><td>0.47</td><td>-0.02</td><td>1.95e-03</td><td>0.03</td><td>0.03</td>
  </tr>
  <tr>
    <td>irr</td>
    <td>16.03</td><td>14.29</td><td>-1.74</td><td>2.67</td><td>5.43</td><td>2.77</td><td>0.50</td><td>0.41</td><td>-0.08</td><td>7.22e-05</td><td>5.18e-04</td><td>4.45e-04</td>
  </tr>
  <tr>
    <td>random</td>
    <td>16.03</td><td>15.29</td><td>-0.74</td><td>2.28</td><td>4.74</td><td>2.46</td><td>0.50</td><td>0.46</td><td>-0.05</td><td>5.94e-04</td><td>1.30e-03</td><td>7.09e-04</td>
  </tr>
  <tr>
    <td rowspan="3">star constellation name</td>
    <td>dstr</td>
    <td>17.57</td><td>16.33</td><td>-1.24</td><td>11.83</td><td>15.43</td><td>3.60</td><td>0.28</td><td>0.26</td><td>-0.02</td><td>4.51e-03</td><td>0.04</td><td>0.03</td>
  </tr>
  <tr>
    <td>irr</td>
    <td>17.56</td><td>14.94</td><td>-2.62</td><td>2.88</td><td>7.87</td><td>4.98</td><td>0.28</td><td>0.20</td><td>-0.07</td><td>2.74e-05</td><td>1.34e-03</td><td>1.31e-03</td>
  </tr>
  <tr>
    <td>random</td>
    <td>17.58</td><td>16.48</td><td>-1.10</td><td>3.40</td><td>7.27</td><td>3.87</td><td>0.28</td><td>0.21</td><td>-0.07</td><td>1.91e-04</td><td>2.13e-03</td><td>1.93e-03</td>
  </tr>
  <tr>
    <td rowspan="3">task done by tool</td>
    <td>dstr</td>
    <td>15.82</td><td>16.56</td><td>0.74</td><td>6.65</td><td>12.38</td><td>5.73</td><td>0.28</td><td>0.34</td><td>0.06</td><td>1.60e-03</td><td>0.03</td><td>0.03</td>
  </tr>
  <tr>
    <td>irr</td>
    <td>15.82</td><td>12.67</td><td>-2.15</td><td>0.83</td><td>5.32</td><td>4.50</td><td>0.28</td><td>0.19</td><td>0.09</td><td>5.52e-06</td><td>5.57e-04</td><td>5.51e-04</td>
  </tr>
  <tr>
    <td>random</td>
    <td>15.85</td><td>14.72</td><td>-1.13</td><td>2.90</td><td>7.51</td><td>4.60</td><td>0.28</td><td>0.24</td><td>-0.04</td><td>4.21e-05</td><td>6.53e-03</td><td>6.49e-03</td>
  </tr>
  <tr>
    <td rowspan="3">task person type</td>
    <td>dstr</td>
    <td>14.71</td><td>16.92</td><td>2.21</td><td>6.70</td><td>11.91</td><td>5.21</td><td>0.21</td><td>0.39</td><td>0.18</td><td>9.63e-04</td><td>0.01</td><td>0.01</td>
  </tr>
  <tr>
    <td>irr</td>
    <td>14.70</td><td>14.32</td><td>-0.38</td><td>0.99</td><td>6.74</td><td>5.75</td><td>0.21</td><td>0.20</td><td>-0.09</td><td>8.39e-06</td><td>1.70e-03</td><td>1.69e-03</td>
  </tr>
  <tr>
    <td>random</td>
    <td>14.73</td><td>13.90</td><td>-0.84</td><td>2.87</td><td>8.11</td><td>5.24</td><td>0.21</td><td>0.17</td><td>-0.04</td><td>3.44e-05</td><td>7.84e-03</td><td>7.80e-03</td>
  </tr>
  <tr>
    <td rowspan="3">work location</td>
    <td>dstr</td>
    <td>14.77</td><td>17.51</td><td>2.74</td><td>8.18</td><td>13.08</td><td>4.90</td><td>0.22</td><td>0.39</td><td>0.18</td><td>2.85e-03</td><td>0.04</td><td>0.03</td>
  </tr>
  <tr>
    <td>irr</td>
    <td>14.70</td><td>14.34</td><td>-0.35</td><td>1.68</td><td>8.23</td><td>6.55</td><td>0.22</td><td>0.20</td><td>-0.02</td><td>1.31e-05</td><td>2.96e-03</td><td>2.94e-03</td>
  </tr>
  <tr>
    <td>random</td>
    <td>14.74</td><td>14.01</td><td>-0.73</td><td>2.53</td><td>7.09</td><td>4.56</td><td>0.22</td><td>0.18</td><td>-0.04</td><td>4.47e-05</td><td>2.89e-03</td><td>2.85e-03</td>
  </tr>
</table>

Table 7: Average Logits and Probabilities for each LRE Relation. LM: Llama-3.1-8B.
A  LRE Dataset

We construct our experimental prompts using commonsense and factual data from the LRE dataset (Hernandez et al., 2024). This dataset comprises 47 relations with over 10,000 instances, spanning four categories: factual associations, commonsense knowledge, implicit biases, and linguistic knowledge. The dataset was created by filtering out ambiguous or noisy triples from ParaRel (Elazar et al., 2021), ensuring that each relation is well-defined and distinct. Additionally, the dataset was expanded by incorporating examples from structured knowledge sources, increasing coverage across various relation types.

However, during our experimentation, we identify several issues with certain relations. Some relations are too obscure for all of our chosen models to complete, resulting in highly noisy logits and probabilities for the correct tokens in some models. For example, the superhero archnemesis relation includes instances such as (Martian Manhunter; Despero, superhero archnemesis), which may not be well-represented across models. Additionally, some relations in the linguistics and bias domains are not particularly relevant to the LM distraction setting. As a result, we select 15 relations, with their statistics listed in Table 5.

B  Context Entrainment Experiment Supplementary Results

Table 6 shows the results across all relations to generate Figure 3. There is a small amount of variance for the logits of the correct token (\( \ell(\Theta) \)) No CTX) because of we cap the amount of samples to 100,000 by random sampling, as described in Section 3.1. Table 7 shows the breakdown for each LRE relation for the Llama-3.1-8B model.

C  Background: The Residual Stream

Here, we review the concept of the residual stream (Elhage et al., 2021), which provides the relevant background for our entrainment head discovery method. The basic idea is to view the computation of transformer LMs as maintaining a communication channel between model components via a shared vector space: the residual stream. Instead of operating independently, attention heads and MLP modules in each layer continuously pass and modify information through residual connections. Starting with the word embedding \( x_0 \), each layer (residual block) modifies it to become \( x_i \), and the final result, \( x_{-1} \), is converted into the output probability distribution by the unembed module. For each residual block at layer \( i \), the output of the previous layer is \( x_{i-1} \). Let \( H_i \) denote the set of all attention heads in the layer. The block’s output, \( x_i \), is then computed as described in Equation (3):

\[
x_i^{\text{mid}} = x_{i-1} + \sum_{h \in H_i} h(x_{i-1});
\]
\[
x_i = x_i^{\text{mid}} + \text{MLP}(x_i^{\text{mid}}).
\]

D  Stochasticity between Runs

The entrainment head discovery method involve several stochastic processes that can introduce noise to the results. Therefore, for the country–capital city relation, we re-run the entrainment head discovery process for 5 runs and report their Jaccard overlap. I.e., how many heads are identified by both runs divided by how many heads are identified by at least one run (\( \frac{A \cap B}{A \cup B} \)). Table 8 shows the Jaccard overlap result.

<table>
  <tr>
    <th>Overlaps</th>
    <th>Run 1</th>
    <th>Run 2</th>
    <th>Run 3</th>
    <th>Run 4</th>
    <th>Run 5</th>
    <th># Heads</th>
  </tr>
  <tr>
    <td>Run 1</td>
    <td>1.000</td>
    <td>0.381</td>
    <td>0.542</td>
    <td>0.372</td>
    <td>0.516</td>
    <td>96 (9.38%)</td>
  </tr>
  <tr>
    <td>Run 2</td>
    <td>0.381</td>
    <td>1.000</td>
    <td>0.419</td>
    <td>0.548</td>
    <td>0.322</td>
    <td>60 (5.86%)</td>
  </tr>
  <tr>
    <td>Run 3</td>
    <td>0.542</td>
    <td>0.419</td>
    <td>1.000</td>
    <td>0.500</td>
    <td>0.351</td>
    <td>89 (6.84%)</td>
  </tr>
  <tr>
    <td>Run 4</td>
    <td>0.372</td>
    <td>0.548</td>
    <td>0.500</td>
    <td>1.000</td>
    <td>0.317</td>
    <td>70 (6.84%)</td>
  </tr>
  <tr>
    <td>Run 5</td>
    <td>0.516</td>
    <td>0.322</td>
    <td>0.351</td>
    <td>0.317</td>
    <td>1.000</td>
    <td>92 (8.98%)</td>
  </tr>
</table>

Table 8: Jaccard overlap between the five runs.
<table>
  <tr>
    <th>Relation</th>
    <th>Exact Acc. (Top-1)</th>
    <th>Strict Acc. (Top-3)</th>
    <th>Credulous Acc. (Top-10)</th>
  </tr>
  <tr>
    <td>company hq</td>
    <td>71.0% 63.5%</td>
    <td>83.5% 90.0%</td>
    <td>88.0% 90.0%</td>
  </tr>
  <tr>
    <td>country capital city</td>
    <td>94.0% 94.0%</td>
    <td>100.0% 100.0%</td>
    <td>100.0% 100.0%</td>
  </tr>
  <tr>
    <td>country currency</td>
    <td>18.0% 19.7%</td>
    <td>83.7% 100.0%</td>
    <td>100.0% 100.0%</td>
  </tr>
  <tr>
    <td>country language</td>
    <td>37.0% 60.3%</td>
    <td>85.7% 100.0%</td>
    <td>100.0% 100.0%</td>
  </tr>
  <tr>
    <td>country largest city</td>
    <td>97.0% 97.0%</td>
    <td>100.0% 100.0%</td>
    <td>100.0% 100.0%</td>
  </tr>
  <tr>
    <td>food from country</td>
    <td>78.0% 79.0%</td>
    <td>92.0% 98.5%</td>
    <td>100.0% 98.5%</td>
  </tr>
  <tr>
    <td>fruit inside color</td>
    <td>49.0% 50.0%</td>
    <td>77.0% 100.0%</td>
    <td>98.0% 100.0%</td>
  </tr>
  <tr>
    <td>fruit outside color</td>
    <td>0.0% 0.0%</td>
    <td>38.0% 84.0%</td>
    <td>82.0% 84.0%</td>
  </tr>
  <tr>
    <td>landmark in country</td>
    <td>72.0% 54.5%</td>
    <td>89.5% 91.0%</td>
    <td>95.0% 91.0%</td>
  </tr>
  <tr>
    <td>landmark on continent</td>
    <td>41.5% 38.5%</td>
    <td>88.5% 83.0%</td>
    <td>97.0% 83.0%</td>
  </tr>
  <tr>
    <td>product by company</td>
    <td>86.0% 81.0%</td>
    <td>95.0% 96.0%</td>
    <td>98.0% 96.0%</td>
  </tr>
  <tr>
    <td>star constellation name</td>
    <td>22.3% 26.0%</td>
    <td>84.7% 89.3%</td>
    <td>92.3% 89.3%</td>
  </tr>
  <tr>
    <td>task done by tool</td>
    <td>58.0% 58.5%</td>
    <td>78.0% 91.0%</td>
    <td>93.5% 91.0%</td>
  </tr>
  <tr>
    <td>task person type</td>
    <td>67.0% 53.5%</td>
    <td>78.5% 80.0%</td>
    <td>80.0% 80.0%</td>
  </tr>
  <tr>
    <td>work location</td>
    <td>49.0% 49.0%</td>
    <td>60.5% 75.0%</td>
    <td>75.0% 75.0%</td>
  </tr>
  <tr>
    <td>arithmetic 0-shot</td>
    <td>95.8% 97.9%</td>
    <td>100.0% 100.0%</td>
    <td>100.0% 100.0%</td>
  </tr>
  <tr>
    <td>spelling correction 1-shot</td>
    <td>58.2% 54.8%</td>
    <td>73.6% 72.0%</td>
    <td>78.6% 76.8%</td>
  </tr>
  <tr>
    <td>spelling correction 2-shot</td>
    <td>77.6% 74.0%</td>
    <td>94.6% 91.6%</td>
    <td>97.0% 94.8%</td>
  </tr>
  <tr>
    <td>spelling correction 5-shot</td>
    <td>87.4% 86.2%</td>
    <td>99.0% 98.4%</td>
    <td>100.0% 100.0%</td>
  </tr>
  <tr>
    <td>translation 1-shot</td>
    <td>58.4% 57.0%</td>
    <td>74.4% 73.0%</td>
    <td>78.4% 76.8%</td>
  </tr>
  <tr>
    <td>translation 2-shot</td>
    <td>74.8% 73.8%</td>
    <td>94.0% 93.0%</td>
    <td>97.0% 96.2%</td>
  </tr>
  <tr>
    <td>translation 5-shot</td>
    <td>85.4% 84.8%</td>
    <td>98.6% 97.2%</td>
    <td>99.6% 99.4%</td>
  </tr>
</table>

Table 9: Removing the entrainment heads of the country–capital city relation has a small to negligible effect on other LM capabilities with exact (must be top-1 response) included. The exact accuracy metric is highly unstable and therefore lacks reference value.

E Evaluation Method

We report both strict (the correct answer appears within the top-3 predicted tokens) and credulous (top-10) accuracy, as multiple correct answers may exist. Relying solely on whether the gold-standard token is the most probable leads to unstable results. For example, the LRE dataset includes the example On the outside, apples are ___, where red is the only correct answer. In some cases, however, the LM selects green as the most probable next token, which is also a valid response.

Table 9 shows Table 4’s results with the exact (must be top-1 response) metrics included. The exact accuracy metric is highly unstable due to the aforementioned reasons and therefore lacks reference value.

F Functional Overlap of Entrainment Heads

Table 10 shows the effect of removing the country-capital entrainment heads compared to removing the task’s heads. We can observe that removing the entrainment heads for country–capital relation causes a similar effect in suppressing contextual entrainment. This suppressing effect is, unsurprisingly, less consistent and smaller compared to removing the actual relation’s entrainment heads.

Therefore, while there is some degree of cross-relation overlap in entrainment heads, it is still a relation-, task-, or domain-specific effect. We hypothesize that entrainment heads are part of the general task-solving mechanism or “circuit.” Recent work has highlighted a high level of idiosyncrasy in how LMs solve tasks. For example, Li et al. (2024), Niu et al. (2024), and Berglund et al. (2024) found that LMs process “A is B” differently from “B is A.” Our observation aligns with these findings.

G ICL Tasks

We employ the three ICL tasks identified by Brown et al. (2020) to evaluate the effect of removing the entrainment heads on the LM’s overall capability. In particular, Brown et al. (2020) proposed three tasks: arithmetic, spelling correction, and translation. Figure 8 shows Brown et al. (2020)’s illustration of these tasks. However, Brown et al. (2020) did not release the full dataset used for evaluation, so we recreate it.

Arithmetic  We randomly sampled 1000 prompts from all possible two digit summations. E.g., \( 23 + 18 = 41 \).
<table>
  <tr>
    <th rowspan="2">Relation</th>
    <th colspan="2">Removing Task Entrainment Heads</th>
    <th colspan="2">Removing Country–Capital Entrainment Heads</th>
  </tr>
  <tr>
    <th>⇒</th>
    <th>⇒</th>
    <th>⇒</th>
    <th>⇒</th>
  </tr>
  <tr>
    <td>company hq</td>
    <td>3.94 ⇒ 14.68</td>
    <td></td>
    <td>3.94 ⇒ 6.81</td>
    <td></td>
  </tr>
  <tr>
    <td>country capital city</td>
    <td>7.69 ⇒ 13.20</td>
    <td></td>
    <td></td>
    <td></td>
  </tr>
  <tr>
    <td>country currency</td>
    <td>4.73 ⇒ 11.67</td>
    <td></td>
    <td>4.73 ⇒ 9.04</td>
    <td></td>
  </tr>
  <tr>
    <td>country language</td>
    <td>6.20 ⇒ 8.95</td>
    <td></td>
    <td>6.20 ⇒ 8.28</td>
    <td></td>
  </tr>
  <tr>
    <td>country largest city</td>
    <td>8.68 ⇒ 13.35</td>
    <td></td>
    <td>8.68 ⇒ 13.27</td>
    <td></td>
  </tr>
  <tr>
    <td>food from country</td>
    <td>3.98 ⇒ 9.95</td>
    <td></td>
    <td>3.98 ⇒ 8.41</td>
    <td></td>
  </tr>
  <tr>
    <td>fruit inside color</td>
    <td>0.97 ⇒ 11.16</td>
    <td></td>
    <td>0.97 ⇒ 4.59</td>
    <td></td>
  </tr>
  <tr>
    <td>fruit outside color</td>
    <td>2.14 ⇒ 13.82</td>
    <td></td>
    <td>2.14 ⇒ 7.45</td>
    <td></td>
  </tr>
  <tr>
    <td>landmark in country</td>
    <td>3.93 ⇒ 9.68</td>
    <td></td>
    <td>3.93 ⇒ 6.14</td>
    <td></td>
  </tr>
  <tr>
    <td>landmark on continent</td>
    <td>2.51 ⇒ 9.14</td>
    <td></td>
    <td>2.51 ⇒ 8.14</td>
    <td></td>
  </tr>
  <tr>
    <td>product by company</td>
    <td>3.62 ⇒ 16.47</td>
    <td></td>
    <td>3.62 ⇒ 7.42</td>
    <td></td>
  </tr>
  <tr>
    <td>star constellation name</td>
    <td>1.07 ⇒ 8.87</td>
    <td></td>
    <td>1.07 ⇒ 5.95</td>
    <td></td>
  </tr>
  <tr>
    <td>task done by tool</td>
    <td>4.70 ⇒ 12.31</td>
    <td></td>
    <td>4.70 ⇒ 6.53</td>
    <td></td>
  </tr>
  <tr>
    <td>task person type</td>
    <td>6.51 ⇒ 12.47</td>
    <td></td>
    <td>6.51 ⇒ 8.21</td>
    <td></td>
  </tr>
  <tr>
    <td>work location</td>
    <td>3.17 ⇒ 12.68</td>
    <td></td>
    <td>3.17 ⇒ 8.28</td>
    <td></td>
  </tr>
</table>

Table 10: Removing the entrainment heads for country-capital relation, for example, causes a similar effect in suppressing contextual entrainment.

![Illustration of three ICL tasks: arithmetic, spelling correction and translation.](page_324_1042_1096_312.png)

Figure 8: Brown et al.’s (2020) illustration of the three ICL tasks: arithmetic, spelling correction and translation.
<table>
  <tr>
    <th colspan="2">(a) Prompt to Obtain the Spelling Correction Data</th>
  </tr>
  <tr>
    <td colspan="2">
      Give me 200 simple, random English words with 1 letter scrambled. For example:<br>
      gaot => goat<br>
      sakne => snake<br>
      brid => bird<br>
      fsih => fish<br>
      dcek => duck<br>
      cmihp => chimp<br>
      organize the results in a Python list:<br>
      [('gaot', 'goat'), ('sakne', 'snake'), ...]
    </td>
  </tr>
  <tr>
    <th colspan="2">(b) Prompt to Obtain the Translation Data</th>
  </tr>
  <tr>
    <td colspan="2">
      Give me 200 simple, random English words with their French translations. For example:<br>
      thanks => merci<br>
      hello => bonjour<br>
      mint => menthe<br>
      wall => mur<br>
      otter => loutre<br>
      bread => pain<br>
      Avoid using French accent marks or letters that does exist in English<br>
      organize the results in a Python list: [('thanks', 'merci'), ('hello', 'bonjour'), ...]
    </td>
  </tr>
</table>

Table 11: Prompts Used to Obtain ICL Evaluation Data.
gaot ⇒ goat, sakne ⇒ snake, brid ⇒ bird, fsih ⇒ fish, dck ⇒ duck, cmihp ⇒ chimp, hosre ⇒ horse, tiegr ⇒ tiger, zbera ⇒ zebra, muose ⇒ mouse, lino ⇒ lion, bera ⇒ bear, wlof ⇒ wolf, fxo ⇒ fox, dree ⇒ deer, frgo ⇒ frog, girafef ⇒ giraffe, donkey ⇒ donkey, bunyn ⇒ bunny, sehep ⇒ sheep, whela ⇒ whale, shrak ⇒ shark, eagel ⇒ eagle, corw ⇒ crow, sawn ⇒ swan, gosoe ⇒ goose, pengiun ⇒ penguin, ostirch ⇒ ostrich, moneky ⇒ monkey, koala ⇒ koala, haed ⇒ head, hnaad ⇒ hand, foto ⇒ foot, lge ⇒ leg, amr ⇒ arm, era ⇒ ear, yee ⇒ eye, lpi ⇒ lip, teo ⇒ toe, hiar ⇒ hair, aplpe ⇒ apple, baanna ⇒ banana, ornage ⇒ orange, maong ⇒ mango, garpe ⇒ grape, pecah ⇒ peach, pera ⇒ pear, plmu ⇒ plum, kwii ⇒ kiwi, leomn ⇒ lemon, carrto ⇒ carrot, pottao ⇒ potato, onoin ⇒ onion, tomtao ⇒ tomato, ltetuce ⇒ lettuce, rdaish ⇒ radish, spianch ⇒ spinach, cucumberr ⇒ cucumber, ppeper ⇒ pepper, celeyr ⇒ celery, rde ⇒ red, bleu ⇒ blue, geren ⇒ green, yellwo ⇒ yellow, purple ⇒ purple, balck ⇒ black, wihte ⇒ white, pnik ⇒ pink, borwn ⇒ brown, gary ⇒ gray, franecc ⇒ france, sapin ⇒ spain, chnia ⇒ china, inida ⇒ india, itlay ⇒ italy, jaapn ⇒ japan, caanda ⇒ canada, barzil ⇒ brazil, geramny ⇒ germany, russai ⇒ russia, teaxs ⇒ texas, manie ⇒ maine, oihio ⇒ ohio, iowa ⇒ iowa, uath ⇒ utah, nevaad ⇒ nevada, aalaska ⇒ alaska, haawii ⇒ hawaii, flordia ⇒ florida, gerogia ⇒ georgia, tabel ⇒ table, cahir ⇒ chair, phoen ⇒ phone, clcok ⇒ clock, wacth ⇒ watch, lihgt ⇒ light, doro ⇒ door, windwo ⇒ window, sopon ⇒ spoon, frok ⇒ fork, rnu ⇒ run, jmup ⇒ jump, wlak ⇒ walk, tlak ⇒ talk, raed ⇒ read, wrtie ⇒ write, eta ⇒ eat, slepe ⇒ sleep, drvie ⇒ drive, siwm ⇒ swim, bgi ⇒ big, smal ⇒ small, fsat ⇒ fast, solw ⇒ slow, hto ⇒ hot, clod ⇒ cold, nwe ⇒ new, odl ⇒ old, hpapy ⇒ happy, sda ⇒ sad, cta ⇒ cat, dgo ⇒ dog, cpu ⇒ cup, pne ⇒ pen, snu ⇒ sun, mono ⇒ moon, satr ⇒ star, teer ⇒ tree, rokc ⇒ rock, blal ⇒ ball, oepn ⇒ open, clsoe ⇒ close, puhs ⇒ push, plul ⇒ pull, lifit ⇒ lift, dorp ⇒ drop, crary ⇒ carry, hodl ⇒ hold, thorw ⇒ throw, cacth ⇒ catch, doctro ⇒ doctor, laweyr ⇒ lawyer, teacehr ⇒ teacher, nusre ⇒ nurse, drierv ⇒ driver, artsit ⇒ artist, sinegr ⇒ singer, writre ⇒ writer, chfe ⇒ chef, pliot ⇒ pilot, freind ⇒ friend, emmey ⇒ enemy, hosue ⇒ house, hoem ⇒ home, family ⇒ family, moeny ⇒ money, waetr ⇒ water, frie ⇒ fire, earht ⇒ earth, widn ⇒ wind, yse ⇒ yes, on ⇒ no, pu ⇒ up, dwon ⇒ down, lft ⇒ left, rigth ⇒ right, ni ⇒ in, otu ⇒ out, dya ⇒ day, nighit ⇒ night, ciyt ⇒ city, tonw ⇒ town, roda ⇒ road, steret ⇒ street, sohp ⇒ shop, sotre ⇒ store, bnak ⇒ bank, csah ⇒ cash, hosiptal ⇒ hospital, clinci ⇒ clinic, questoin ⇒ question, anwser ⇒ answer, probelm ⇒ problem, solutoin ⇒ solution, loev ⇒ love, haet ⇒ hate, peac ⇒ peace, wra ⇒ war, truhit ⇒ truth, lei ⇒ lie, music ⇒ music, moive ⇒ movie, boko ⇒ book, paeg ⇒ page, paep ⇒ paper, pecnil ⇒ pencil, deks ⇒ desk, soaf ⇒ sofa, pillwo ⇒ pillow, blnaket ⇒ blanket.

Table 12: 200 Spelling Correction Pairs Used for ICL Capability Evaluation.
thanks ⇒ merci, hello ⇒ bonjour, mint ⇒ menthe, wall ⇒ mur, otter ⇒ loutre, bread ⇒ pain, water ⇒ eau, friend ⇒ ami, love ⇒ amour, cat ⇒ chat, dog ⇒ chien, house ⇒ maison, horse ⇒ cheval, cow ⇒ vache, cheese ⇒ fromage, family ⇒ famille, black ⇒ noir, white ⇒ blanc, red ⇒ rouge, green ⇒ vert, blue ⇒ bleu, boy ⇒ garcon, girl ⇒ fille, night ⇒ nuit, day ⇒ jour, morning ⇒ matin, evening ⇒ soir, sun ⇒ soleil, moon ⇒ lune, star ⇒ etoile, sky ⇒ ciel, flower ⇒ fleur, car ⇒ voiture, city ⇒ ville, country ⇒ pays, beach ⇒ plage, forest ⇒ foret, river ⇒ riviere, mountain ⇒ montagne, desert ⇒ desert, island ⇒ ile, table ⇒ table, chair ⇒ chaise, window ⇒ fenetre, door ⇒ porte, book ⇒ livre, pen ⇒ stylo, pencil ⇒ crayon, letter ⇒ lettre, store ⇒ magasin, restaurant ⇒ restaurant, coffee ⇒ cafe, tea ⇒ the, juice ⇒ jus, milk ⇒ lait, egg ⇒ oeuf, butter ⇒ beurre, sugar ⇒ sucre, salt ⇒ sel, pepper ⇒ poivre, chicken ⇒ poulet, beef ⇒ boeuf, fish ⇒ poisson, bird ⇒ oiseau, snake ⇒ serpent, frog ⇒ grenouille, turtle ⇒ tortue, rabbit ⇒ lapin, pig ⇒ cochon, sheep ⇒ mouton, goat ⇒ chevre, fox ⇒ renard, wolf ⇒ loup, lion ⇒ lion, tiger ⇒ tigre, bear ⇒ ours, phone ⇒ telephone, computer ⇒ ordinateur, keyboard ⇒ clavier, screen ⇒ ecran, mouse ⇒ souris, camera ⇒ camera, photo ⇒ photo, movie ⇒ film, music ⇒ musique, song ⇒ chanson, dance ⇒ danse, poem ⇒ poeme, library ⇒ bibliotheque, museum ⇒ musee, school ⇒ ecole, university ⇒ universite, teacher ⇒ professeur, student ⇒ etudiant, office ⇒ bureau, job ⇒ travail, money ⇒ argent, bank ⇒ banque, street ⇒ rue, road ⇒ route, building ⇒ batiment, tall ⇒ grand, small ⇒ petit, short ⇒ court, big ⇒ gros, new ⇒ nouveau, old ⇒ vieux, happy ⇒ heureux, sad ⇒ triste, angry ⇒ fache, tired ⇒ fatigue, busy ⇒ occupe, free ⇒ libre, open ⇒ ouvert, closed ⇒ ferme, expensive ⇒ couteux, cheap ⇒ , yes ⇒ oui, no ⇒ non, maybe ⇒ , never ⇒ jamais, always ⇒ toujours, often ⇒ souvent, sometimes ⇒ parfois, rarely ⇒ rarement, early ⇒ tot, late ⇒ tard, now ⇒ maintenant, soon ⇒ bientot, yesterday ⇒ hier, today ⇒ aujourd), tomorrow ⇒ demain, hour ⇒ heure, minute ⇒ minute, second ⇒ seconde, time ⇒ temps, moment ⇒ moment, week ⇒ semaine, month ⇒ mois, year ⇒ annee, monday ⇒ lundi, tuesday ⇒ mardi, wednesday ⇒ mercredi, thursday ⇒ jeudi, friday ⇒ vendredi, saturday ⇒ samedi, sunday ⇒ dimanche, spring ⇒ printemps, summer ⇒ ete, autumn ⇒ automne, winter ⇒ hiver, police ⇒ police, fire ⇒ feu, help ⇒ aide, problem ⇒ probleme, question ⇒ question, answer ⇒ reponse, truth ⇒ verite, lie ⇒ mensonge, idea ⇒ idee, important ⇒ important, interesting ⇒ interessant, possible ⇒ possible, impossible ⇒ impossible, difficult ⇒ difficile, easy ⇒ facile, strong ⇒ fort, weak ⇒ faible, light ⇒ lumiere, dark ⇒ sombre, direction ⇒ direction, left ⇒ gauche, right ⇒ droite, straight ⇒ tout), back ⇒ arriere, up ⇒ haut, down ⇒ bas, in ⇒ dans, out ⇒ dehors, on ⇒ sur, under ⇒ sous, behind ⇒ derriere, next ⇒ prochain, near ⇒ pres, far ⇒ loin, between ⇒ entre, each ⇒ chacun, all ⇒ tous, some ⇒ quelques, none ⇒ aucun, every ⇒ chaque, anyway ⇒ de toute facon example ⇒ exemple, reason ⇒ raison, mistake ⇒ erreur, gift ⇒ cadeau, party ⇒ fete, plan ⇒ plan, goal ⇒ objectif, success ⇒ succes.

Table 13: 200 Translation Pairs (English ⇒ French) Used for ICL Capability Evaluation.
<table>
  <tr>
    <th>MMLU Task</th>
    <th>Before Removing</th>
    <th>After Removing</th>
  </tr>
  <tr>
    <td>high school european history</td>
    <td>70.9</td>
    <td>71.5</td>
  </tr>
  <tr>
    <td>business ethics</td>
    <td>55.0</td>
    <td>51.0</td>
  </tr>
  <tr>
    <td>clinical knowledge</td>
    <td>58.4</td>
    <td>58.8</td>
  </tr>
  <tr>
    <td>medical genetics</td>
    <td>62.0</td>
    <td>63.0</td>
  </tr>
  <tr>
    <td>high school us history</td>
    <td>65.6</td>
    <td>68.1</td>
  </tr>
  <tr>
    <td>high school physics</td>
    <td>35.7</td>
    <td>35.7</td>
  </tr>
  <tr>
    <td>high school world history</td>
    <td>68.7</td>
    <td>67.0</td>
  </tr>
  <tr>
    <td>virology</td>
    <td>51.2</td>
    <td>46.3</td>
  </tr>
  <tr>
    <td>high school microeconomics</td>
    <td>58.4</td>
    <td>57.1</td>
  </tr>
  <tr>
    <td>econometrics</td>
    <td>36.8</td>
    <td>36.8</td>
  </tr>
  <tr>
    <th>Average</th>
    <th>56.3</th>
    <th>55.5</th>
  </tr>
</table>

Table 14: Effect on removing the entrainment heads on the MMLU tasks.

of entrainment heads collectively corresponds more to contextual entrainment in specific settings, and less to the model’s overall performance.

Spelling Correction  We obtain the spelling correction data by prompting ChatGPT. Specifically, we use ChatGPT-o18 with the prompt shown in Table 11(a). Table 12 lists the generated spelling correction pairs. We then randomly sample 1,000 instances under three different settings (1-shot, 2-shot, and 5-shot) from these 200 pairs.

Translation  Similarly, for translation, we also sampled English to French word pairs by prompting ChatGPT-o1. The prompt used is shown in Table 11(a). Table 13 lists the generated spelling correction pairs. We then randomly sample 1,000 instances under three different settings (1-shot, 2-shot, and 5-shot) from these 200 pairs.

H  General Capabilities Affected by Removal of Heads

We run the first 10 tasks from the MMLU benchmark using the llama-3.1-8B model with the deepeval9 package under the default setting (no chain-of-thought). The entrainment heads are the same ones identified using the country–capital relation, consistent with the rest of the evaluation in Section 3.2. Results depicted in Table 14.

Overall, we see the same trend as in our evaluation on other LRE relations and the ICL tasks: a small decrease in overall performance (56.3% to 55.5%), but greater variance in the individual breakdown. There is no significant change in the model’s performance, and the LMs’ general ability remains in the same ballpark as the original model after removing the entrainment heads. These results further support our hypothesis that this circuit

8https://openai.com/index/openai-o1-system-card/
9https://www.deepeval.com/

Reviewer paper summaries: The paper investigates "contextual entrainment"  in large language models, focusing on how irrelevant or distracting context influences model outputs. The authors identify that LLMs exhibit a bias toward previously seen tokens, even when those tokens are semantically unrelated to the task at hand. And more importantly even with fine-tuned models (one model). 

 The paper continues with two contributions, a- Constructing prompts with different types of context to test the entrainment. and b- identifies entrainment heads (subset of induction heads?, see below), that are able to suppress the contextual entrainment strength.
Summary of strengths: * The paper provides valuable mechanistic insights into the behavior of LLMs, particularly in understanding contextual entrainment. The comparison to phenomena like in-context learning (ICL) and repetition further enriches its contribution, making it a noteworthy study.

* By employing a wide range of context settings (such as irrelevant, random, and counterfactual prompts) the authors enable a thorough investigation into how different types of context influence the model's outputs. 

* The paper demonstrates successful interventions, such as disabling specific attention heads, which effectively suppress distracting tokens from the model's outputs while preserving overall performance.

* The paper is well-written, methodologically sound, and highly relevant to ongoing efforts to improve the reliability and robustness of LLMs (except the emojis, see below).
Summary of weaknesses: * The study focuses primarily on GPT-2 XL and LLaMA models. It remains unclear whether the findings generalize to other architectures or (good) fine-tuned models.

* The use of Gumbel-sigmoid sampling introduces randomness into the process of identifying entrainment heads, however having 500 epoch is equibalent to run the process this number of times, but I would not only guarantees the best run is guaranteed to be the scholastically robust results, but another metric (like std) of these runs would be important to have stable results.  

* The loss function used to identify entrainment heads bears a conceptual resemblance to the approach used for induction heads. I would be interested in seeing whether the suppression of contextual entrainment arises specifically from the identified entrainment heads or if it also involves induction heads, i.e., whether the same heads are responsible for both phenomena, (small discussion is enough and no further experiments needed). 

* Lastly, the sparsity consider |m| terms, is it a mistake or there is a subset of head that are included in the term?
Comments, suggestions and typos: This paper makes a significant contribution to the field of mechanistic interpretability by shedding light on the phenomenon of contextual entrainment  in LLMs. However, I find it hard to follow (memorize) the meaning of the variables, so I strongly suggest the equations are rewritten in a proper way (except maybe illustrative equations).
Overall assessment scores: 4