Identifying Semantic Induction Heads to Understand In-Context Learning (2024)

Abstract: Although large language models (LLMs) have demonstrated remarkable performance, the lack of transparency in their inference logic raises concerns about their trustworthiness. To gain a better understanding of LLMs, we conduct a detailed analysis of the operations of attention heads and aim to better understand the in-context learning of LLMs. Specifically, we investigate whether attention heads encode two types of relationships between tokens present in natural languages: the syntactic dependency parsed from sentences and the relation within knowledge graphs. We find that certain attention heads exhibit a pattern where, when attending to subject tokens, they recall object tokens and increase the output logits of those object tokens. More crucially, the formulation of such semantic induction heads has a close correlation with the emergence of the in-context learning ability of language models. The study of semantic attention heads advances our understanding of the intricate operations of attention heads in transformers, and further provides new insights into the in-context learning of LLMs.

## introduction in recent years, the transformer-based large language models (llms) @xcite @xcite have rapidly emerged as one of the mainstreams in the field of natural language processing (nlp). while these models demonstrate emergent abilities as they scale @xcite , they become less interpretable due to the vast number of parameters and complex architectures, which emphasizes llms' safety and trustworthiness @xcite @xcite . thus, beyond classical gradient-based explanations @xcite , and perturbation-based explanations @xcite @xcite , recent studies in mechanistic interpretability @xcite attempt to reverse engineer the computations in transformers (particularly attention layers). the mechanistic interpretability on transformer language models was first performed by @xcite . they disentangle two circuits from the operation of each attention head in transformers: query-key circuit (determines which token the head prefers to attend to) and output-value circuit (determines how the head affects the output logits of the next token). then, @xcite discover that some attention heads prefer to search for a previous occurrence of the current token in context and copy the next token associated with that occurrence, as shown in figure 1 . the attention heads performing such operations are termed induction heads. taking a step further, @xcite ; @xcite have discovered that the presence of induction heads has a close correlation with the in-context learning (icl) ability of llms. this finding highlights the importance of understanding the behavior of attention heads to the overall learning capabilities of llms. on the other hand, semantic relationships have a vital importance on natural language understanding and processing. however, @xcite only focus on whether the attention heads copy the attended token, without studying semantic relationships between tokens. another major limitation of previous studies is that @xcite does not explain the popular few-shot in-context learning schema. instead, they study the loss decreasing along with the increase of token indices. this setting does not fully capture the complete ability of llms to learn from the context. in this work, beyond simple copying, we delve deeper into high-level relationships encoded in at-induction head for copying semantic induction head for syntactic dependency triplet: (pen, mod, nice) semantic induction head for semantic relationship triplet: (pen, used-for, writing) i have a nice pen for writing . he also wants a nice prefix matching copying i have a nice pen for writing . the pen is nice attend to the head token raise the tail token following the dependency relationship i have a nice pen for writing . i can use it raise the tail token following the semantic relationship to write attend to the head token figure 1 : induction heads and semantic induction heads. for the sequence "... a nice ... a", an induction head finds a place where the current token "a" occurred, attends to its next token "nice" (prefix matching), and then copies "nice" to the output (copying). in contrast, the semantic induction head raises the output logits of tail tokens ("nice" in the mod dependency and "writing" in the used-for relation) when attending to the head token "pen". inspired by the study of induction heads and incontext learning, we further explore the correlation between semantic induction heads and in-context learning. we first categorize the in-context learning ability into three basic levels: loss reduction, format compliance, and pattern discovery. these three levels progressively increase in difficulty, with each subsequent level building upon the achievements of the previous one. the experimental results are consistent with our hypothesis, demonstrating the emergence of three levels of icl in a sequential manner. specifically, we observe the emergence of loss reduction from the beginning of the training, followed by the emergence of format compliance at around 1.6b tokens, and finally, the emergence of pattern discovery after training on approximately 4b tokens. moreover, we find semantic induction heads mainly emerge around the same time as pattern discovery. based on this finding, we infer that the emergence of semantic induction heads plays a crucial role in facilitating the icl of llms. our contributions can be summarized as follows. ## related works in this section, we provide an overview of recent advancements in the interpretability of neural networks, particularly mechanistic interpretability. on the other hand, previous studies @xcite have also explored the topic of semantic relationships in models. however, our study distinguishes itself by focusing on mechanistic interpretability. previous studies in interpretability can be roughly categorized into the following four types: estimating the attribution of input features to the network output @xcite @xcite @xcite , discovering interaction patterns between input features @xcite @xcite , extracting concepts from intermediate-layer features @xcite @xcite , and designing self-explainable architectures @xcite . as transformer-based models become mainstream, recent works focus on understanding the attention mechanism. the most direct approach is to visualize the attention using bipartite graphs @xcite or heatmaps @xcite . another line of research aims to reverse engineer the operation of attention heads, called mechanistic interpretability @xcite . @xcite proposed the circuit analysis (introduced in section 3) to examine the operation of attention heads, and they found induction heads in attention-only models. @xcite further investigated the correlation between the formation of induction heads and icl. @xcite observed an overlap between the set of induction heads and the set of important attention heads for icl. using circuit analysis, wang et al. (2023) also found some attention heads performing the function of identifying/removing names in the indirect object identification task. other studies @xcite @xcite intervened the attention or ffn layers to study their functions. in this paper, we leverage the circuit analysis to investigate semantic relationships in attention heads. on the other hand, the matrix w h @xmath0, termed the output-value (ov) circuit, computes the independent output of each head at the current token regardless of the attention pattern. the output of the ov circuit can be projected back to the vocabulary as xw h ov w u by the unembedding transformation w @xmath1. the projected vector represents the influence of the attention head on the output. importantly, according to @xcite , both the qk circuit and ov circuit are directly performed on input embeddings, facilitating the understanding of operations in attention heads. based on the above decomposition, elhage et al. (2021) identify a specific behavior in attention heads, which they refer to as induction heads. they observe this behavior in attention heads when presented with sequences like "[a] [b] • • • [a]". in these induction heads, the qk circuit causes the attention head to attend to the token @xcite , which appears next to the previous occurrence of the current token @xcite . this behavior is termed prefix matching. then, the ov circuit increases the output logit of the attended token [b], termed copying. this mechanism is shown in figure 1 . ## in-context learning and semantic induction heads in this section, we investigate the correlation between in-context learning and semantic induction heads. we first categorize the icl ability into three levels and observe the gradual emergence of different levels of icl. then, we investigate the formation of semantic induction heads in the training process to understand the emergence of icl. ## more discussions about relationships encoded in attention heads a notable distinction between copying and relationship is that the relationship between tokens depends on the input context, while copying is context-agnostic. therefore, different inputs may utilize different attention heads. thus, we propose to find a common group of attention heads that represent specific relationships in different inputs. for each triplet t , we identify the attention head that has the largest value of e j [a h,j t ] (larger than 0). then, we count the occurrence of each head having the largest value among all triplets t for each dependency relationship. figure 10 shows the number of occurrences of each attention head having the largest relation index for syntactic dependencies. please refer to appendix g for results on semantic relationships. we find that for each type of dependency/relationship, there are around 5∼15 attention heads frequently activated by different inputs. besides, different relationships tend to share some common attention heads (e.g., layer2, head4 for syntactic dependencies). this may indicate that these human-defined relationships are not mutually exclusive from the llms' point of view. in other words, there exists a many-to-many mapping between human-defined relationships and attention heads in llms. ## conclusion previous studies @xcite in mechanistic interpretability only studied the simple functions in very specific tasks @xcite . in this study, we extended the conventional induction heads to analyze high-level relationships between words/entities in natural languages. our experiments revealed that specific attention heads encode syntactic dependencies and semantic relationships in natural languages. furthermore, we identified three levels of the in-context learning ability of llms, and experimental results showed they are progressively learned during the training process. finally, we observed a close correlation between the formation of semantic induction heads and incontext learning ability, strengthening our understanding of in-context learning.