Attention Heads of Large Language Models: A Survey
Zifan Zheng Yezhaohui Wang Yuxin Huang Shichao Song Bo Tang Feiyu Xiong Zhiyu Li

Foundation Model or Finetune? Evaluation of few-shot semantic segmentation for river pollution
    Marga Don10009-0001-5435-8935 
Stijn Pinson 2 Blanca Guillen Cebrian 2
Yuki M. Asano 1,3 


    September 9, 2024
===============================================================================================




§ ABSTRACT 
Since the advent of ChatGPT, Large Language Models (LLMs) have excelled in various tasks but remain largely as black-box systems. Consequently, their development relies heavily on data-driven approaches, limiting performance enhancement through changes in internal architecture and reasoning pathways. As a result, many researchers have begun exploring the potential internal mechanisms of LLMs, aiming to identify the essence of their reasoning bottlenecks, with most studies focusing on attention heads.
Our survey aims to shed light on the internal reasoning processes of LLMs by concentrating on the interpretability and underlying mechanisms of attention heads. We first distill the human thought process into a four-stage framework: Knowledge Recalling, In-Context Identification, Latent Reasoning, and Expression Preparation. Using this framework, we systematically review existing research to identify and categorize the functions of specific attention heads. Furthermore, we summarize the experimental methodologies used to discover these special heads, dividing them into two categories: Modeling-Free methods and Modeling-Required methods. Also, we outline relevant evaluation methods and benchmarks. Finally, we discuss the limitations of current research and propose several potential future directions. Our reference list is open-sourced at <https://github.com/IAAR-Shanghai/Awesome-Attention-Heads>.








§ INTRODUCTION

The Transformer architecture <cit.> has demonstrated outstanding performance across various tasks, such as Natural Language Inference and Natural Language Generation. However, it still retains the black-box nature inherent to Deep Neural Networks (DNNs) <cit.>. As a result, many researchers have dedicated efforts to understanding the internal reasoning processes within these models, aiming to uncover the underlying mechanisms <cit.>. This line of research provides a theoretical foundation for models like BERT <cit.> and GPT <cit.> to perform well in downstream applications. Additionally, in the current era where Large Language Models (LLMs) are widely applied, interpretability mechanisms can guide researchers in intervening in specific stages of LLM inference, thereby enhancing their problem-solving capabilities <cit.>.

Among the components of LLMs, attention heads play a crucial role in the reasoning process. Particularly in the recent years, attention heads within LLMs have garnered significant attention, as illustrated in Figure <ref>. Numerous studies have explored attention heads with specific functions. This paper consolidates these research efforts, organizing and analyzing the potential mechanisms of different types of attention heads. Additionally, we summarize the methodologies employed in these investigations.






 §.§ Structure of Our Work

The logical structure and classification method of this paper are illustrated in Figure <ref>. We begin with the background of the problem in Section <ref>, where we present a simplified representation of the LLMs' structures (Section <ref>) and explain the related key terms (Section <ref>). In Section <ref>, we first summarize the four stages of human thought processes from a cognitive neuroscience perspective and apply this framework to analyze the reasoning mechanisms of LLMs. Using this as our classification criterion, we categorize existing work on attention heads, identifying commonalities among heads that contribute to similar reasoning processes (Sections <ref>–<ref>) and exploring the collaborative mechanisms of heads functioning at different stages (Section <ref>).

Investigating the internal mechanisms of models often requires extensive experiments to validate hypotheses. To provide a comprehensive understanding of these methods, we summarize the current experimental methodologies used to explore special attention heads in Section <ref>. We divide these methodologies into two main categories based on whether they require additional modeling: Modeling-Free (Section <ref>) and Modeling-Required (Section <ref>).

In addition to the core sections shown in Figure <ref>, we summarize the evaluation tasks and benchmarks used in relevant studies in Section <ref>.
Furthermore, in Section <ref>, we compile research on the mechanisms of Feed-Forward Networks (FFNs) and Mechanical Interpretability to help deepen our understanding of LLM structures from multiple perspectives.
Finally, in Section <ref>, we offer our insights on the current state of research in this field and outline several potential directions for future research.




 §.§ Comparison with Related Surveys

To the best of our knowledge, there is no survey focused on the mechanisms of LLMs' attention heads. Specifically, <cit.> mainly discusses non-Transformer architectures, with little focus on attention heads. The surveys by <cit.> cover older content, primarily focusing on the various attention computation methods that emerged during the early development of the Transformer. However, current LLMs still use the original scaled-dot product attention, indicating that many of the derived attention forms have become outdated. Although <cit.> focuses on the internal structure of LLMs, it only summarizes experimental methodologies and overlooks research findings related to operational mechanisms.

Compared to the aforementioned surveys, the strengths of our work are:

    
  * Focus on the latest research. Although earlier researchers explored the mechanisms of attention heads in models like BERT, many of these conclusions are now outdated. This paper primarily focuses on highly popular LLMs, such as LLaMA and GPT, consolidating the latest research findings.
    
  * An innovative four-stage framework for LLM reasoning. We have distilled key stages of human thought processes by integrating knowledge from cognitive neuroscience, psychology, and related fields. And we have applied these stages as an analogy for LLM reasoning.
    
  * Detailed categorization of attention heads. Based on the proposed four-stage framework, we classify different attention heads according to their functions within these stages, and we explain how heads operating at different stages collaborate to achieve alignment between humans and LLMs.
    
  * Clear summarization of experimental methods. We provide a detailed categorization of the current methods used to explore attention head functions from the perspective of model dependency, laying a foundation for the improvement and innovation of experimental methods in future research.





 §.§ Out-of-scope Topics


    
  * This paper primarily targets the attention heads within current mainstream LLM architectures, specifically those with a decoder-only structure. As such, we do not discuss early studies related to the Transformer, such as those focusing on attention heads in BERT-based models.
    
  * Some studies on mechanistic interpretability propose holistic operational principles that encompass embeddings, attention heads, and MLPs. However, this paper focuses exclusively on attention heads. Consequently, Sections <ref> through <ref> do not cover the roles of other components within the Transformer architecture; these are only briefly summarized in Section <ref>.





§ BACKGROUND
 



 §.§ Mathematical Representation of LLMs
 
To facilitate the discussion in subsequent sections, we first define the relevant notations[Currently, there are two main layer normalization methods in LLMs: Pre-Norm and Post-Norm <cit.>. However, since these are not the focus of this paper, we will omit Layer Normalization in our discussion.].

As shown in Figure <ref>, a model ℳ consists of an embedding layer, L identical decoder blocks, and an unembedding layer. The input to ℳ are one-hot sentence tokens, with a shape of {0,1}^N × |𝒱|, where N is the length of the token sequence and |𝒱| represents the vocabulary size.

After passing through the embedding layer, which applies semantic embedding 𝐖_𝐄∈ℝ^|𝒱| × d and positional encoding 𝐏_𝐄 (e.g., RoPE <cit.>), the one-hot matrix is transformed into the input 𝐗_0,0∈ℝ^N × d for the first decoder, where d represents the dimension of the token embedding (latent vector).



In the ℓ-th ℓ(1≤ℓ≤ L) decoder block, there are two residual blocks. Each decoder blocks contain H attention heads. The first residual block combines the input matrix 𝐗_ℓ,0∈ℝ^N × d with the output 𝐗_ℓ^attn obtained from the multi-head attention operation, producing 𝐗_ℓ,1 (as shown in Equation <ref>). Subsequently, 𝐗_ℓ,1 serves as the input for the second residual block. Here, Attn_ℓ^h(·) (1≤ℓ≤ L, 1≤ h ≤ H) represents the computation function of the h-th attention head in the ℓ-th layer, where 1 ≤ h ≤ H.

    𝐗_ℓ^attn   = ∑_h=1^HAttn_ℓ^h(𝐗_ℓ,0) 
    𝐗_ℓ,1   = 𝐗_ℓ,0 + 𝐗_ℓ^attn


Similarly, as shown in Equation <ref>, the second residual block combines 𝐗_ℓ,1 with the output 𝐗_ℓ^ffn obtained after passing through the FFN, yielding the final output 𝐗_ℓ+1,0 of the ℓ-th decoder block. This output also serves as the input for the ℓ+1-th decoder block. Here, FFN_ℓ(·) consists of linear layers (and activation functions) such as GLU (Gated Linear Units), SwiGLU <cit.>, or MoE <cit.>.

    𝐗_ℓ^ffn   = FFN_ℓ(𝐗_ℓ,1) 
    𝐗_ℓ+1,0   = 𝐗_ℓ,1 + 𝐗_ℓ^ffn


Here, we will concentrate on the details of Attn_ℓ^h(·). This function can be expressed using matrix operations.
Specifically, each layer's Attn_ℓ^h(·) function corresponds to four low-rank matrices: 𝐖_𝐐_ℓ^h, 𝐖_𝐊_ℓ^h, 𝐖_𝐕_ℓ^h∈ℝ^d ×d/H, 𝐎_ℓ^h∈ℝ^d/H× d. By multiplying 𝐗_ℓ,0 with 𝐖_𝐐_ℓ^h, the query matrix 𝐐_ℓ^h∈ℝ^N ×d/H is obtained. Similarly, the key matrix 𝐊_ℓ^h and the value matrix 𝐕_ℓ^h can be derived.
The function Attn_ℓ^h(·) can then be expressed as Equation <ref> <cit.>.

    Attn_ℓ^h(𝐗_ℓ,0) = softmax(𝐐_ℓ^h⊤·𝐊_ℓ^h) ·𝐕_ℓ^h·𝐎_ℓ^h







 §.§ Glossary of Key Terms
 

This paper introduces several specialized terms related to reasoning mechanism modeling and experimental exploration methods. Below, we provide explanations for the key terms [For more definitions of specialized terms, please refer to the work of <cit.>.].


    
  * Circuits. Circuits are abstractions of the reasoning logic in deep models. The model ℳ is viewed as a computational graph. There are two main approaches to modeling circuits. One approach treats the features in the latent space of ℳ as nodes and the transitions between features as edges <cit.>. The other approach views different components of ℳ, such as attention heads and neurons, as nodes; and the interactions between these components, such as residual connections, as edges <cit.>. A circuit is a subgraph of ℳ. Current research on attention heads primarily uses the second definition.
    
    
  * Residual stream. The residual stream after layer ℓ is the sum of the embedding and the outputs of all layers up to layer ℓ, and is the input to layer ℓ+1. <cit.> offers a perspective of the residual stream as a shared bandwidth. Different layers can transmit information through this shared bandwidth (with lower layers writing information and higher layers reading it), as illustrated in Figure <ref>.
    
    
  * Knowledge circuit. As defined in <cit.>, a knowledge circuit is a critical subgraph in ℳ to view the knowledge mechanism of Transformers. Knowledge circuits focus on how different components collaborate to express internal knowledge.
    
    
  * Activation patching & Ablation study. Activation patching involves replacing the activation values in certain layers of a model to analyze the contribution of these layers to the model's decision-making. Ablation study, on the other hand, involves removing a component from the LLM and observing the changes in the output <cit.>.

    Both methods require calculating the effect towards output after the operation. Specifically, there are three types of effects: direct effect, indirect effect, and total effect, as shown in Figure <ref>.

    The key difference between them is that Activation Patching does not remove components, whereas Ablation Study logically removes a component.
    
    
  * Logit lens. When calculating effects like those shown in Figure <ref>, logit lens can quantify this effect. Specifically, it uses uembedding layer to map an intermediate representation vector to the logits values of the vocabulary, allowing for the comparison of logits differences or other metrics <cit.>.







§ OVERVIEW OF SPECIAL ATTENTION HEADS
 
Previous research has shown that the decoder-only architecture described in Section <ref> follows the Scaling Law, and it exhibits emergent abilities once the number of parameters reaches a certain threshold <cit.>. Many LLMs that have emerged subsequently demonstrate outstanding performance in numerous tasks, even close to human. However, researchers still do not fully understand why these models are able to achieve such remarkable results. To address this question, recent studies have begun to delve into the internal mechanisms of LLMs, focusing on their fundamental structure—a neural network composed of multi-attention heads and FFNs.

We have observed that many studies concentrate on the functions of attention heads, attempting to explain their reasoning processes. Additionally, several researchers have drawn parallels of reasoning methods  between LLMs and human <cit.>. Therefore, in this section, we will use the framework of human cognitive paradigms as a guiding method to classify the functions of different attention heads.






 §.§ How does Brain / Attention Head think?

By summarizing and analyzing existing research, we find that the role of an attention head, as its name suggests, is quite analogous to the functions of the human brain. From a behavioral neuroscience perspective, the process by which the human brain thinks about specific problems can be abstracted into a four-step process: Knowledge Recalling (KR), In-Context Identification (ICI), Latent Reasoning (LR), and Expression Preparation (EP). These four steps can interact and transition between each other, as illustrated in Figure <ref>.

When solving a problem, humans first need to recall the knowledge they have learned that is relevant to the issue at hand. This process is known as Knowledge Recalling (KR). During this stage, the hippocampus integrates memories into the brain's network <cit.> and activates different types of memories as needed <cit.>.
Confronted with the specific text of the problem, humans need to perform In-Context Identification (ICI). This means that the brain not only focuses on the overall structural content of the text <cit.> but also parses the syntactic <cit.> and semantic <cit.> information embedded within it.

Once the brain has acquired the aforementioned textual and memory information, it attempts to integrate this information to derive conclusions, a process known as Latent Reasoning (LR). This stage primarily includes arithmetic operations <cit.> and logical inference <cit.>.
Finally, the brain needs to translate the reasoning results into natural language, forming an answer that can be expressed verbally. This is the Expression Preparation (EP) stage. At this point, the brain bridge the gap between “knowing” and “saying” <cit.>.

As indicated by the arrows in Figure <ref>, these four stages are not executed in a strictly one-direction fashion when humans solve problems; rather, they can jump and switch between each other. For example, the brain may “cycle” through the identification of contextual content (the ICI stage) and then retrieve relevant knowledge based on the current context (the KR stage). Similarly, if latent reasoning cannot proceed further due to missing information, the brain may return to the Knowledge Recalling and In-Context Identification stages to gather more information.

We will now draw an analogy between these four steps and the mechanisms of attention heads, as depicted in Figure <ref>. Previous research has shown that LLMs possess strong contextual learning abilities and have many practical applications <cit.>. As a result, much of the work on interpretability has focused on the ability of LLMs to capture and reason about contextual information. Consequently, the functions of currently known special attention heads are primarily concentrated in the ICI and LR stages, while there are fewer attention heads that operate in the KR and EP stages.



my-box=[
    rectangle,
    draw=hidden-draw,
    rounded corners,
    align=left,
    text opacity=1,
    minimum height=1.5em,
    minimum width=5em,
    inner sep=2pt,
    fill opacity=.8,
    line width=0.8pt,
]

leaf-head=[my-box, minimum height=1.5em,
    draw=gray!80, 
    fill=gray!35,  
    text=black, font=,
    inner xsep=2pt,
    inner ysep=4pt,
    line width=0.8pt,
]

leaf-task=[my-box, minimum height=2.5em,
    draw=red!80, 
    fill=red!20,  
    text=black, font=,
    inner xsep=2pt,
    inner ysep=4pt,
    line width=0.8pt,
]

leaf-paradigms=[my-box, minimum height=2.5em,
    draw=orange!70, 
    fill=orange!15,  
    text=black, font=,
    inner xsep=2pt,
    inner ysep=4pt,
    line width=0.8pt,
]
leaf-others=[my-box, minimum height=2.5em,
    
    draw=yellow!80, 
    fill=yellow!15,  
    text=black, font=,
    inner xsep=2pt,
    inner ysep=4pt,
    line width=0.8pt,
]
leaf-other=[my-box, minimum height=2.5em,
    
    draw=green!80, 
    fill=green!15,  
    text=black, font=,
    inner xsep=2pt,
    inner ysep=4pt,
    line width=0.8pt,
]
leaf-application=[my-box, minimum height=2.5em,
    
    draw=blue!80, 
    fill=blue!15,  
    text=black, font=,
    inner xsep=2pt,
    inner ysep=4pt,
    line width=0.8pt,
]

modelnode-task=[my-box, minimum height=1.5em,
    draw=red!80, 
    fill=red!20,  
    text=black, font=,
    inner xsep=2pt,
    inner ysep=4pt,
    line width=0.8pt,
]

modelnode-paradigms=[my-box, minimum height=1.5em,
    draw=orange!70, 
    fill=orange!15,  
    text=black, font=,
    inner xsep=2pt,
    inner ysep=4pt,
    line width=0.8pt,
]
modelnode-others=[my-box, minimum height=1.5em,
    draw=yellow!80, 
    fill=yellow!15,  
    text=black, font=,
    inner xsep=2pt,
    inner ysep=4pt,
    line width=0.8pt,
]
modelnode-other=[my-box, minimum height=1.5em,
    draw=green!80, 
    fill=green!15,  
    text=black, font=,
    inner xsep=2pt,
    inner ysep=4pt,
    line width=0.8pt,
]
modelnode-application=[my-box, minimum height=1.5em,
    draw=blue!80, 
    fill=blue!15,  
    text=black, font=,
    inner xsep=2pt,
    inner ysep=4pt,
    line width=0.8pt,
]









 §.§ Knowledge Recalling (KR)
 
For LLMs, most knowledge is learned during the training or fine-tuning phases, which is embedded in the model's parameters. This form of knowledge is often referred to as LLMs' “parametric knowledge”. Similar to humans, certain attention heads in LLMs recall this internally stored knowledge—such as common sense or domain-specific expertise—to be used in subsequent reasoning. These heads typically retrieve knowledge by making initial guesses or by focusing on specific content within the context, injecting the memory information into the residual stream as initial data.

In general tasks, <cit.> identified that some attention heads function as associative memories, gradually storing and retrieving knowledge during the model's training phase. The so-called Memory Head <cit.> can retrieve content related to the current problem from the parametric knowledge. This content could be knowledge learned during pre-training or experience accumulated during previous reasoning processes.

In specific task scenarios, such as when LLMs tackle Multiple Choice Question Answering (MCQA) problems, they may initially use Constant Head to evenly distribute attention scores across all options, or they might use Single Letter Head to assign a higher attention score to one option while giving lower scores to others, thereby capturing all potential answers <cit.>.

Additionally, in the context of Binary Decision Tasks (BDT)[A Binary Decision Task is a problem where the solution space is discrete and contains only two options, such as yes-no questions or answer verification.], <cit.> found that LLMs often exhibit a negative bias when handling such tasks. This could be because the model has learned a significant amount of negative expressions related to similar tasks from prior knowledge during training. Consequently, when the model identifies a given text as a binary task, a Negative Head may “preemptively” choose the negative answer due to this prior bias.





 §.§ In-Context Identification (ICI)
 
Understanding the in-context nature of a problem is one of the most critical process to effectively addressing it. Just as humans read a problem statement and quickly pick up on various key pieces of information, some attention heads in LLMs also focus on these elements. Specifically, attention heads that operate during the ICI stage use their QK matrices to focus on and identify overall structural, syntactic, and semantic information within the in-context. This information is then written into the residual stream via OV matrices.



  §.§.§ Overall Structural Information Identification
 
Identifying the overall structural information within a context mainly involves LLMs attending to content in special positions or with unique occurrences in the text.
Previous Head <cit.> (also referred to as Positional Head in <cit.>) attend to the positional relationships within the token sequence. They capture the embedding information of the current token and the previous token.
Rare Words Head <cit.> focus on tokens that appear with the lowest frequency, emphasizing rare or unique tokens.
Duplicate Token Head excel at capturing repeated content within the context, giving more attention to tokens that appear multiple times <cit.>.

Besides, as LLMs can gradually handle long texts, this is also related to the “Needle-in-a-Haystack” capability of attention heads. (Global) Retrieval Head can accurately locate specific tokens in long texts <cit.>. These heads enable LLMs to achieve excellent reading and in-context retrieval capabilities.



  §.§.§ Syntactic Information Identification
 
For syntactic information identification, sentences primarily consist of subjects, predicates, objects, and clauses. Syntactic Head can distinctly identify and label nominal subjects, direct objects, adjectival modifiers, and adverbial modifiers.
Some words in the original sentence may get split into multiple subwords because of the tokenizer (e.g., “happiness” might be split into “happi” and “ness”). The Subword Merge Head focus on these subwords and merge them into one complete word <cit.>.

Additionally, <cit.> proposed the Mover Head cluster, which can be considered as “argument parsers”. These heads often copy or transfer sentence's important information (such as the subject's position) to the [END] position[The [END] position refers to the last token's position in the sentence being decoded by the LLM. Many studies indicate that summarizing contextual information at this position facilitates subsequent reasoning and next-token prediction.].
Name Mover Head and Backup Name Mover Head can move the names in the text to the [END] position.
Letter Mover Head can extract the first letters of certain words in the context and aggregate them at the [END] position <cit.>.
Conversely, Negative Name Mover Head prevent name information from being transferred to the [END] position <cit.>.



  §.§.§ Semantic Information Identification
 
As for semantic information identification,
Context Head <cit.> extract information from the context that is related to the current task.
Further, Content Gatherer Head <cit.> “move” tokens related to the correct answer to the [END] position, preparing to convert them into the corresponding option letter for output.
The Sentiment Summarizer proposed by <cit.> can summarize adjectives and verbs that express sentiment in the context near the [SUM] position[[SUM] position is next to [END] position.], making it easier for subsequent heads to read and reason.

Capturing the message about relationship is also important for future reasoing. Semantic Induction Head <cit.> capture semantic relationships within sentences, such as part-whole, usage, and category-instance relationships.
Subject Head and Relation Head <cit.> focus on subject attributes and relation attributes, respectively, and then inject these attributes into the residual stream.




 §.§ Latent Reasoning (LR)
 
The KR and ICI stages focus on gathering information, while Latent Reasoning (LR) is where all the collected information is synthesized and logical reasoning occurs. Whether in humans or LLMs, the LR stage is the core of problem-solving. Specifically, QK matrices of a head performs implicit reasoning based on information read from the residual stream, and then the reasoning results or signals are written back into the residual stream through OV matrices.



  §.§.§ In-context Learning
 
In-context Learning is one of the most widely discussed areas. It primarily includes two types: Task Recognition (TR) and Task Learning (TL) <cit.>. Both involve learning from context to infer the problem's solution, but they differ in their reliance on semantic information. TR has labels with clear semantics, such as “positive” and “negative” in sentiment classification. In contrast, TL depends on learning the specific mapping function between example-label pairs, where the example and label do not have a semantic connection.

For Task Recognition: Summary Reader <cit.> can read the information summarized at the [SUM] position during the ICI stage and use this information to infer the corresponding sentiment label.
<cit.> proposed that the output of certain mid-layer attention heads can combine into a Function Vector. These heads abstract the core features and logical relationships of a task, based on the semantic information identified during ICI, and thereby trigger task execution.

For Task Learning, the essence of solving these problems is enabling LLMs to inductively find patterns.
Induction Head is among the most widely studied attention heads <cit.>. They capture patterns such as “… [A][B] … [A]” where token [B] follows token [A], and predict that the next token should be [B]. Specifically, Induction Head can get information about all tokens in the context and the previous token from Previous Head. It then match this with information at the [END] position to perform further reasoning.

Induction Head tends to strictly follow a pattern once identified and complete fill-in-the-blank reasoning. However, in most cases, the real problem will not be identical to the examples—just as a student's exam paper will not be exactly the same as their homework. To address this, <cit.> introduced the In-context Head, whose QK matrix calculates the similarity between information at the [END] position and each label. The OV matrix then extracts label features and weights them according to the similarity scores to determine the final answer (take all labels into consideration rather than only one label).



  §.§.§ Effective Reasoning
 
Some studies have identified heads related to reasoning effectiveness. Truthfulness Head <cit.> and Accuracy Head <cit.> are heads highly correlated with the truthfulness and accuracy of answers. They help the model infer truthful and correct results in QA tasks, and modifying the model along their activation directions can enhance LLMs' reasoning abilities.

However, not all heads positively impact reasoning. For example, Vulnerable Head <cit.> are overly sensitive to certain specific input forms, making them susceptible to irrelevant information and leading to incorrect results. During reasoning, it is advisable to minimize the influence of such heads.



  §.§.§ Task Specific Reasoning
 
Finally, some heads are specialized for specific tasks.
In MCQA tasks, Correct Letter Head <cit.> can complete the matching between the answer text and option letters by comparing the index of each option, determining the final answer choice.
When dealing with tasks related to sequential data, Iteration Head <cit.> can iteratively infer the next intermediate state based on the current state and input.
For arithmetic problems, Successor Head <cit.> can perform increment operations on ordinal numbers.

These examples illustrate how various attention heads specialize in different aspects of reasoning, contributing to the overall problem-solving capabilities of LLMs.





 §.§ Expression Preparation (EP)
 
During the Expression Preparation (EP) stage, LLMs need to align their reasoning results with the content that needs to be expressed verbally. Specifically, EP heads may first aggregate information from various stages.
<cit.> proposed the Mixed Head, which can linearly combine and aggregate information written into the residual stream by heads from the ICI and LR stages (such as Subject Heads , Relation Heads, Induction Heads, etc.). The aggregated results are then mapped onto the vocabulary logits value via the OV matrix.

Some EP heads have a signal amplification function. Specifically, they read information about the context or reasoning results from the residual stream, then enhance the information that needs to be expressed as output, and write it back into the stream.
Amplification Head <cit.> and Correct Head <cit.> amplify the signal of the correct choice letter in MCQA problems near the [END] position. This amplification ensures that after passing through the Unembedding layer and softmax calculation, the correct choice letter has the highest probability.

In addition to information aggregation and signal amplification, some EP heads are used to align the model's reasoning results with user's instructions.
In multilingual tasks, the model may sometimes fail to respond in the target language that user wanted. Coherence Head <cit.> ensure linguistic consistency in the generated content. They help LLMs maintain consistency between the output language and the language of user's query when dealing with multilingual inputs.
Faithfulness Head <cit.> are strongly associated with the faithfulness of CoT[CoT stands for Chain-of-Thought <cit.>. Faithfulness refers to whether the model's generated response accurately reflects its internal reasoning process and behavior, i.e., the consistency between output and internal reasoning.]. Enhancing the activation of these heads allows LLMs to better align their internal reasoning with the output, making the CoT results more robust and consistent.

However, for some simple tasks, LLMs might not require special EP heads to refine language expression. At this situation, the information written back into the residual stream during the ICI and LR stages may be directly suitable for output, i.e., skip the EP stage and select token with highest probability. 






 §.§ How Attention Heads Working Together?
 
If we divide the layers of a LLM (e.g., GPT-2 Small) into three segments based on their order—shallow (e.g., layers 1-4), middle (e.g., layers 5-8), and deep (e.g., layers 9-12)—we can map the relationship between the stages where heads act and the layers they are in, according to Section <ref>-<ref>. The figure is illustrated in Figure <ref>.

Some researchers have explored the potential semantic meanings embedded in the query vector 𝐪_ℓ, j^h=𝐐_ℓ^h[:, j] and key vector 𝐤_ℓ, j^h=𝐊_ℓ^h[:, j] when attention heads collaborate. For example, in MCQA problem, during the ICI stage, a Content Gatherer Head moves the tokens of the correct answer text to the [END] position. Then, in the LR stage, the Correct Letter Head uses the information passed by the Content Gatherer Head to identify the correct option. The query vector in this context effectively asks, “Are you the correct label?” while recalling the gathered correct answer text. The key vector represents, “I'm choice [A/B/C/D], with corresponding text [...]”. After matching the right key vector to the query vector, we can get the correct answer choice.

Consider the Parity Problem[The Parity Problem involves determining the parity (odd or even) of the sum of a given sequence. The sequence only consists of 0/1 values. For instance, the sum of the sequence “001011” is 3, so it is an odd sequence. Let s_i indicates the parity of the sum of first i digits. So the corresponding s_[0:t] to “001011” is eeeooeo, where e represents even and o represents odd. When querying an LLM, the prompt format is “[0/1 seq.] [EOI] [s_[0:t]] [END]”, with [EOI] as the End-Of-Input token. The expected answer is the final state s_t.].
During the ICI stage, a Mover Head transmits the position of the [EOI] token, which separates the input sequence and the intermediate state sequence, to the [END] position.
In the LR stage, an Iteration Head first reads the [EOI]'s position index from [END] and uses its query vector to ask, “Are you position t?” The key vector for each token responds, “I'm position t^'.” This querying process identifies the last digit in the input sequence, which, combined with s_t-1, allows the model to calculate s_t.

Further research has explored integrating multiple special attention heads into a cohesive working mechanism. Take the IOI (Indirect Object Identification) task, which tests the model's ability to deduce the indirect object in a sentence, as an example. Figure <ref> outlines the process.

    
  * In the KR stage, the Subject Head and Relation Head focus on “Mary” and “bought flowers for”, respectively, triggering the model to recall that the answer should be a person's name <cit.>.

    
  * Then in the ICI stage, the Duplicate Head identifies “John”, while the Name Mover Head focuses on both “John” and “Mary”.

    
  * During the iterative stages of ICI and LR, the Previous Head and Induction Head work together to attend to “John”. All this information is written to the residual stream. Then the Inhibition Head detects that “John” appears multiple times and is the subject, thereby suppressing the logits value of “John”.
    
  * Finally in the stage of EP, the Amplification Head boosts the logits value for “Mary”.







§ UNVEILING THE DISCOVERY OF ATTENTION HEADS
 
How can we uncover the specific functions of those special heads mentioned in Section <ref>? In this section, we will unveiling the discovery methods. Current research primarily employs experimental methods to validate the working mechanisms of those heads. We categorize the mainstream experimental approaches into two types based on whether they require the construction of new models: Modeling-Free and Modeling-Required. The classification scheme and method examples are shown in Figure <ref>.






 §.§ Modeling-Free
 



Modeling-Free methods do not require setting up new models, making them widely applicable in interpretability research. These methods typically involve altering a latent state computed during the LLMs' reasoning process and then using Logit Lens to map the intermediate results to token logits or probabilities. By calculating the logit (or probability) difference, researchers can infer the impact of the change. Modeling-Free methods primarily include Activation Patching and Ablation Study. However, due to the frequent interchange of these terms in the literature, a new perspective is required to distinguish them. This paper further divides these methods into Modification-Based and Replacement-Based Methods based on how the latent state representation is altered, as summarized in Table <ref>.

Modification-Based methods involve altering the values of a specific latent state while retaining some of the original information.
Directional Addition retains part of the information in the original state and then directionally adds some additional information. For instance, <cit.> input texts containing positive and negative sentiments into LLMs, obtaining positive and negative representations from the latent state. The difference between these two representations can be seen as a sentiment direction in the latent space. By adding this sentiment direction vector to the activation of the attention head under investigation, the effect on the output can be analyzed to determine whether the head has the ability to summarize sentiment.

Conversely, Directional Subtraction retains part of the original state information while directionally removing some of it <cit.>. This method can be used to investigate whether removing specific information from a latent state affects the model's output in a significant way, thereby revealing whether certain attention heads can backup or fix the deleted information.

In contrast to Modification-Based methods, Replacement-Based methods discard all information in a specific latent state and replace it with other values.
Zero Ablation and Mean Ablation replace the original latent state with zero values or the mean value of latent states across all samples from a dataset, respectively. This can logically “eliminate” the head or cause it to lose its special function, allowing researchers to assess its importance.

Naive Activation Patching is the traditional patching method. It involves using a latent state obtained from a corrupted prompt to replace the original latent state at the corresponding position. For example, consider the original prompt “John and Mary went to the store.” Replacing “Mary” with “Alice” results in a corrupted prompt. By systematically replacing the latent state obtained under the original prompt with the one obtained under the corrupted prompt across each head, researchers can preliminarily determine which head has the ability to focus on names based on the magnitude of the impact <cit.>.




 §.§ Modeling-Required
 





Modeling-Required methods involve explicitly constructing models to delve deeper into the functions of specific heads. Based on whether the newly constructed models require training, we further categorize Modeling-Required methods into Training-Required and Training-Free methods, as summarized in Table <ref>.

Training-Required methods necessitate training the newly established models to explore mechanisms.
Probing is a common training-based method. This approach extracts activation values from different heads as features and categorizes heads into different classes as labels. A classifier is then trained on this data to learn the relationship between the activation patterns and the head's function. Subsequently, the trained classifier can serve as a probe to detect which heads within the LLMs possess which functions <cit.>.

Another approach involves training a simplified transformer model on a clean dataset for a specific task. Researchers investigate whether the heads in this simplified model exhibit certain functionalities, which can then be extrapolated whether similar heads in the original model possess the same capabilities. This method reduces computational costs during training and analysis, while the constructed model remains simple and highly controllable <cit.>.

Training-Free methods primarily involve designing scores that reflect specific phenomena. These scores can be viewed as mathematical models that construct an intrinsic relationship between the attributes of components and certain model characteristics or behaviors.
For instance, when investigating Retrieval Heads, <cit.> defined a Retrieval Score. This score represents the frequency with which a head assigns the highest attention score to the token it aims to retrieve across a sample set, as shown in Equation <ref>. A high Retrieval Score indicates that the head possesses a strong “Needle in a Haystack” ability.

Similarly, when exploring Negative Heads, <cit.> introduced the Negative Attention Score (NAS), as shown in Equation <ref>. Here, i denotes the i-th token in the input prompt, and t_Yes and t_No represent the positions of “Yes” and “No” in the prompt, respectively. A high NAS suggests that the head focuses more on negative tokens during decision-making, making it prone to generating negative signals.

    RetrievalScore_ℓ^h = |𝒟_right∩𝒟_all|/|𝒟_all|


    NAS_ℓ^h = ∑_i(Attn_ℓ^h[i, t_Yes] + Attn_ℓ^h[i, t_No]) ·log(Attn_ℓ^h[i, t_No]/Attn_ℓ^h[i, t_Yes])


In addition to scoring, researchers have proposed other novel training-free modeling methods.
<cit.> introduced the concept of an Information Flow Graph, where nodes represent tokens and edges represent information transfer between tokens via attention heads or FFNs. By calculating and filtering the importance of each edge to the node it points to, key edges can be selected to form a subgraph. This subgraph can then be viewed as the primary internal mechanism through which LLMs perform reasoning.





§ EVALUATION
 
This section summarizes the benchmarks and datasets used in the interpretability research of attention heads. Based on the different evaluation goals during the mechanism exploration process, we categorize them into two types: Mechanism Exploration Evaluation and Common Evaluation. The former is designed to evaluate the working mechanisms of specific attention heads, while the latter assesses whether enhancing or suppressing the functions of certain special heads can improve the overall performance of LLMs.











































 §.§ Mechanism Exploration Evaluation

To delve deeper into the internal reasoning paths of LLMs, many researchers have synthesized new datasets based on existing benchmarks. The primary feature of these datasets is the simplification of problem difficulty, with elements unrelated to interpretability, such as problem length and query format, being standardized. As shown in Table <ref>, these datasets essentially evaluate the model's knowledge reasoning and knowledge recalling capabilities, but they simplify the answers from a paragraph-level to a token-level.

Take exploring sentiment-related heads as an example, <cit.> created the ToyMovieReview and ToyMoodStory datasets, with specific prompt templates illustrated in Figure <ref>. Using these datasets, researchers employed sampling methods to calculate the activation differences of each head for positive and negative sentiments. This allowed them to identify heads with significant differences as potential candidates for the role of Sentiment Summarizers.




 §.§ Common Evaluation

The exploration of attention head mechanisms is ultimately aimed at improving the comprehensive capabilities of LLMs. Many researchers, upon identifying a head with a specific function, have attempted to modify that type of head—such as by enhancing or diminishing its activation—to observe whether the LLMs' responses become more accurate and useful. We classify these Common Evaluation Benchmarks based on their evaluation focus, as shown in Table <ref>. The special attention heads discussed in this paper are closely related to improving LLMs' abilities in four key areas: knowledge reasoning, sentiment analysis, long context retrieval, and text comprehension.




§ ADDITIONAL TOPICS
 
In this section, we summarize various works related to the LLMs interpretability. Although these works may not introduce new special heads as discussed in Section <ref>, they delve into the underlying mechanisms of LLMs from other perspectives. We will elaborate on these studies under two categories: FFN Interpretability and Machine Psychology.



 §.§ FFN Interpretability
 
As discussed in Section <ref>, apart from attention heads, FFNs also plays a significant role in the LLMs reasoning process. This section primarily summarizes research focused on the mechanisms of FFNs and the collaborative interactions between attention heads and FFNs.

One of the primary functions of FFNs is to store knowledge acquired during the pre-training phase. 
<cit.> proposed that factual knowledge stored within the model is often concentrated in a few neurons of the MLP.
<cit.> observed that the neurons in the FFN of GPT models can be likened to key-value pairs, where specific keys can retrieve corresponding values, i.e., knowledge.
<cit.> discovered a hierarchical storage of knowledge within the model's FFN, with lower layers storing syntactic and semantic information, and higher layers storing more concrete factual content.

FFNs effectively complement the capabilities of attention heads across the four stages described in Section <ref>. The collaboration between FFNs and attention heads enhances the overall capabilities of LLMs.
<cit.> proposed that attention heads and FFNs can work together to enrich the representation of a subject and then extract its related attributes, thus facilitating factual information retrieval during the Knowledge Recall (KR) stage.
<cit.> found that, unlike attention heads, which focus on global information and perform aggregation, FFNs focus only on a single representation and perform local updates. This complementary functionality allows them to explore textual information both in breadth (attention heads) and depth (FFNs).

In summary, each component of LLMs plays a crucial role in the reasoning process. The individual contributions of these components, combined with their interactions, accomplish the entire process from Knowledge Recalling to Expression.




 §.§ Machine Psychology
 
Current research on the LLMs interpretability often draws parallels between the reasoning processes of these models and human thinking. This suggests the need for a more unified framework that connects LLMs with human cognition. The concept of Machine Psychology has emerged to fill this gap <cit.>, exploring the cognitive activities of AI through psychological paradigms.

Recently, <cit.> and <cit.> have proposed different approaches to studying machine psychology.
Hagendorff's work focuses on using psychological methods to identify new abilities in LLMs, such as heuristics and biases, social interactions, language understanding, and learning. His research suggests that LLMs display human-like cognitive patterns, which can be analyzed to improve AI interpretability and performance.

Johansson's framework integrates principles of operant conditioning <cit.> with AI systems, emphasizing adaptability and learning from environmental interactions. This approach aims to bridge gaps in AGI research by combining insights from psychology, cognitive science, and neuroscience.

Overall, Machine Psychology provides a new perspective for analyzing LLMs. Psychological experiments and behavioral analyses may lead to new discoveries about these models. As LLMs are increasingly applied across various domains of society, understanding their behavior through a psychological lens becomes increasingly important, which offers valuable insights for developing more intelligent AI systems.




§ CONCLUSION
 


 §.§ Limitations in Current Research

Firstly, we observe that the application scenarios explored in current research are relatively simple and limited to specific types of tasks, lacking generalizability. For instance, studies like <cit.> and <cit.> have discovered reasoning circuits in LLMs through tasks such as the IOI task and the Color Object Task. However, these circuits have not been validated across other tasks, making it difficult to prove whether these mechanisms are universally applicable.

Secondly, most research focuses on the mechanisms of individual heads, with only a few researchers delving into the collaborative relationships among multiple heads. As a result, the existing work lacks a comprehensive framework for understanding the coordinated functioning of all the attention heads in LLMs.

Finally, the conclusions of most studies lack mathematical proofs. Many studies start by proposing a hypothesis about a circuit or mechanism based on an observed phenomenon, followed by experiments designed to validate the hypothesis. The downside of this research approach is that experiments cannot establish the theoretical soundness of the mechanism, nor can they determine whether the mechanism is merely coincidental.



 §.§ Future Directions and Challenges

Building on the limitations discussed above and the content presented earlier, this paper outlines several potential research directions for the future:

    
  * Exploring mechanisms in more complex tasks. Investigate whether certain attention heads possess special functions in more complex tasks, such as open-ended question answering, math problems.
    
    
  * Mechanism's robustness against prompts. Research has shown that current LLMs are highly sensitive to prompts, with slight changes potentially leading to opposite responses <cit.>. Future work could analyze this phenomenon through the lens of attention head mechanisms and propose solutions to mitigate this issue.
    
    
  * Developing new experimental methods. Explore new experimental approaches, such as designing experiments to verify whether a particular mechanism is indivisible or whether it has universal applicability.
    
    
  * Building a Comprehensive Interpretability Framework. This framework should encompass both the independent and collaborative functioning mechanisms of most attention heads and other components.
    
    
  * Integrating Machine Psychology. Incorporate insights from Machine Psychology to construct an internal mechanism framework for LLMs from an anthropomorphic perspective, understanding the gaps between current LLMs and human cognition and guiding targeted improvements.





§ LIMITATION
 
Current research on the interpretability of LLMs’ attention heads is relatively scattered, primarily focusing on the functions of individual heads, with a lack of overarching frameworks. As a result, the categorization of attention head functions from the perspective of human cognitive behavior in this paper may not be perfectly orthogonal, potentially leading to some overlap between different stages.



unsrtnat
 







