Function Vectors in Large Language Models

Thank you for your thoughtful review. We are glad you found our paper highly relevant and clearly written. We have posted an updated version of the paper including the improvements that you have suggested. We also discuss several of your suggestions and questions here:

---

On the previous work on Lampinen on task representations, as well as Clark on attention heads:

Our work builds on that previous work and is consistent with the findings from the existing body of work, and it adds some new insights. We have added this discussion to our updated related work section.

In detail: our work builds upon the attention-weight-based analyses of attention as done by Voita 2018, Clark 2019, Voita 2019, Kovaleva 2019, Reif 2019, Lin 2019, Htut 2019 who found that attention patterns follow many linguistic dependency relationships across a variety of tasks; we also find that attention heads in ICL systematically follow structural dependencies, as shown in Figure 3b. However, that is not the focus of our inquiry: inspired by the observations that attention patterns alone do not fully explain the mechanisms (Jain 2019, Wiegreffe 2019, Kobayashi 2020, Bibal 2022), we focus on another way of understanding the role of attention heads.

The focus of our paper is to ask another question about the specific role of attention heads in ICL contexts: what is the content of the information transported by the attention heads? We pose this question in a way that differs from the probing classifiers proposed in Clark et al. which examine attention maps to characterize attention head behavior for specific linguistic relationships. In contrast we use causal mediation analysis, which does not rely directly on the attention pattern, but on the content of the attention output, and is thus not directly comparable to their approach. Our finding on the transported function-identification information in ICL offers a new window into the nontrivial and interpretable role that attention can play.

---

Clarifying the importance of our work.

Thank you for the suggestions. We agree that the importance of the finding should be clarified, so based on your suggestion we have made several changes to the paper to clarify the framing and importance of the result.

• We have changed the title of our work to “Function Vectors in Large Language Models”.
• We have updated the introduction to include a clearer contrast of our work and simple word vector arithmetic, including a detailed discussion in Appendix A.
• We have also updated our discussion and related works section.

---

How much ICL performance do function vectors recover?

We mention briefly at the beginning of section 3.1 that “we only include those test queries for which the LLM answers correctly given a 10-shot ICL prompt; all accuracy results are reported on this filtered subset.“ Thus, the results presented in figure 4 and table 2 can be understood to be the proportion of the task representation encoded by the FVs, relative to the original model’s ICL performance - which corresponds to 100% for the reported results. However, we agree that we should make this more clear and have added some clarification of the implications of this choice in the revised version of the paper.

Additionally, we have included in Appendix D an analysis where we compare the originally-reported filtered result to the result we get without filtering the test set and include the ICL performance as an oracle baseline for comparison. We find that the performance of our FVs remains very similar under both settings (filtering the test set vs. no filtering)

---

Regarding Figures 1 & 2

Based on your suggestions, we have updated the content of figure 1, and the location of figure 2 to be more in-line with the text describing its setup and results in the updated version of our paper.

I have read other reviews and your response. On the one hand, other reviews bring up some valid points/concerns. I especially resonate with R 9YEq's comment on how a deeper look into the FV mechanisms is desirable. On the other hand, I believe that the authors address many of the reviewers' concerns completely or partially.

Overall, therefore, I am leaving my score unchanged.

I would like to encourage the authors to further refine the revised discussion that has to do with previous work (Analyzing the Attention Mechanism section especially). For example, the new sentence "our work draws motivation from the observation that attention patterns alone do not fully explain mechanisms" is a little grammatically confusing (mechanisms of what?). The continuation is also a little confusing "for example Kobayashi et al. (2020) notes the magnitude of transported data is important for insight" - it's not immediately clear how this advances the argument and what is meant by "insight". It would benefit the paper to further rework this paragraph to include a more clear/concrete example of how previous attention analyses failed. I, however, appreciate the ending of that paragraph since it tries to succinctly formulate how their proposed approach allows to gain deeper insight into LLM attention/representations, as compared to previous works.

Thank you for your review, we’re glad you found our paper well-motivated, intuitive, and creative. We have updated our paper submission to answer your questions and incorporate your feedback. We also address your questions here, including pointers to corresponding updates in the paper:

---

Clarifying how function vectors are computed and applied to invoke task behavior, and attention heads used

We have added a section to the appendix (Appendix B) to clarify how attention head outputs are added together to form a function vector, as well as how we add a function vector to the hidden state of a particular layer to invoke task behavior. To summarize here - because a function vector is a sum over a set of attention head outputs, adding it to the hidden state is done in the same way that attention outputs are added to the hidden state vector. If the function vector is $v_t$, and the hidden state is $h_l$ we compute the resulting hidden state we get when adding the function vector as h_l^\{'\} = h_\{l-1\} + m_l + \sum_\{j \leq J\}\{a_\{lj\}\} + v_t. This can also be written as h_l^\{'\} = h_\{l\} + v_t.

You are correct: we focus on the computations being done at the last token of the prompt, so all of these attention heads are located in this token’s stream of predictions.

For our causal analysis we perform a preliminary search over all attention heads (at the final token) to determine their causal effect, but only select a subset of heads to use when creating a function vector. We use multiple tasks to determine this small set of important heads, but we use the activations from only a single task to compute function vectors. In other words, a function vector is only a vector over a single task.

---

Could we have not just added the task value (without summing in equation 5) to each attention head leading to computing the next hidden state?

Yes - This is a good question, of a similar flavor to that of Reviewer 9YEq. We have run this experiment and report the results in Appendix M. We find that if we add the attention outputs to its corresponding layer (or at its corresponding ($l, j$) position), we find that this generally matches the performance of our function vector. So the causal effects and results we see are comparable between this alternate formulations.

---

Also, why compute the top attention heads over all tasks and not just have a different distinct set for each task?

We found that because there was very substantial overlap in the attention heads that were important for a majority of the individual tasks, having a distinct set for each task performs about the same as using a universal set of heads over all tasks. In Appendix G we include a figure showing the indirect effect of each head, and the overlap of these heads across many tasks.

---

In 2.2 J is mentioned but not defined: Number of heads?

Yes, thanks for pointing this out. We clarify this in the update to the paper.

Thank you for your review. We’re glad you found our paper clear and easy to follow. We have updated our paper submission to answer your questions and incorporate your feedback. We also address your questions here, including pointers to corresponding updates in the paper:

---

Though the paper mainly investigates attention heads, I’m still wondering why attention heads contain such information rather than the MLP layers. … Can the paper discuss more on this part, which would better explain the proposed findings?

Our work is consistent with previous findings (Meng 2022, Geva 2023) that MLP layers do play a role in enriching hidden states with information stored in the model. Yet both Meng and Geva also found direct evidence that attention heads at the final token transport some important information; in our current work, we choose to investigate the transport of information through the model. The attention mechanism we study is interesting because unlike information storage that utilizes different MLPs for different pieces of information (See Meng 2022 appendix B Figs 10, 14), we find a tiny set of attention heads that serve as common mediators for many very different ICL tasks (Appendix G).

In addition, our research is asking an open question that is not answered by Meng et al: what kind of information is being transported by attention heads? We find that within ICL contexts, attention heads carry a kind of information that was not previously noticed by Meng: we find an encoding of the task itself, rather than a specific answer to a task.

We have added an extended discussion in Appendix A comparing and contrasting our results to the previous "vocabulary space" findings of Geva 2022.

---

For example, can the author explain how to determine the number of selected heads?

For GPT-J we tested using 1-30 attention heads to construct function vectors and found that the increase in performance for many tasks begins to plateau when using about |A| = 10 attention heads.

For the larger models, the number of attention heads is much larger, but we maintain about the same ratio of ~2% of the attention heads. The total number of attention heads for each model is: GPT-J (6B) = 448, LlaMa 2 (7B) = 1024, LlaMa 2 (13B) = 1600, GPT-NeoX (20B) = 2816, LlaMa 2 (70B) = 5120. For these models we use 10, 20, 50, 50, and 100 heads respectively, corresponding to roughly 2.23%, 1.95%, 3.12%, 1.78%, and 1.95% of the total heads.

We have added information to Appendix C with a figure showing these experiments and additional discussion of this choice of the number of selected heads.

---

Clarifying the creation of function vectors in equation 5

We have added an appendix section (Appendix B) to clarify how attention head outputs are added together to form a function vector and how we add a function vector to the hidden state of a particular layer.

Because the function vector is a sum of average contributions of attention heads, these attention head outputs are already naturally weighted by the transformer. We did perform some preliminary investigations into weighting the function vector when adding it to a particular layer, finding that a weight of 1 was close to optimal. Thus, weighting seemed to not lend any advantage. Additionally, not weighting the FV allows us to keep the function vector as close to the original hidden state distribution as possible.

---

For experiments, I find though the paper tested 40 tasks, they are very simple. Thus, can the paper think of some more difficult ones?

While we do not investigate these tasks extensively in the main paper due to space constraints, we do include a few standard challenging NLP benchmark tasks among the set of tasks we test. For example, in the appendix we show results for a sentiment analysis task (SST2), an entity recognition task (CoNLL2003), a question-answering task (commonsenseQA), and a text classification task (AG News).

These can be found in Appendix E.

In addition, reviewer 9YEq has suggested analyzing tasks with cyclic structure that cannot be done with a simple vector offset, and we have created two new datasets of this flavor. We present results on these datasets in Appendix L.

---

Meanwhile, I think if the paper can add some analyses about the FV ability under model scaling laws, it can be better.

We agree, thank you for this suggestion.

In Appendix J in our updated submission we show how FVs perform for different model sizes of the same model family (Llama 2). In general, the ability of FVs to trigger task execution seems to improve with scale. In addition, we have added to the appendix (Appendix E.3) additional results for many individual tasks for Llama 2 (13B) and Llama 2 (70B) for a deeper analysis of these scaling effects.

Thank you for your review and helpful suggestions. We have updated our paper submission to answer your questions and incorporate your feedback. We also address your questions here, with pointers to the corresponding updates in the paper:

---

The paper is relatively phenomenological rather than mechanistic; we see that there effects, but there is relatively little analysis of how they actually occur.

In our paper we focus on how LLMs trigger task execution via a functional mechanism. We test the causal effects of states under interventions as we add a function vector, so we’re not just looking at correlations. Furthermore, the tests are done under dramatic domain shifts, which is evidence of very strong generalizable mechanistic effects. We’ve rewritten the discussion in the paper to clarify this contribution.

---

Clarifying how function vectors are created and added to a hidden state.

We have added an appendix section (appendix B) to clarify how attention head outputs are added together to form a function vector and how we add a function vector to the hidden state of a particular layer.

In short, the attention head outputs can be viewed as having the same dimension as the hidden states. Because a function vector is a sum over a set of attention head outputs, adding it to the hidden state is done in the same way that attention outputs are added to the hidden state vector.

---

How was the ‘direct decoding’ of FV’s performed?

To decode a function vector, we use the transformer’s decoder (denoted as $D$ in the paper). Formally, this means that we pass the function vector through the model’s final layer norm and subsequent unembedding matrix to turn it into a distribution over the model’s vocabulary space, similar to how the output of the final transformer layer is decoded to vocabulary space via the same decoder $D$. We would write this as $D(v_t) =$ Unembedding(LayerNorm($v_t$)), where $v_t$ is a function vector.

---

What is the ‘optimization’ used to reconstruct a $v_\{tk\}$ that matches the distribution $Q_\{tk\}$ when decoded?

The optimization process we use to reconstruct a $v_\{tk\}$ starts with a randomly initialized vector $v \in \mathbb{R}^d$. We then minimize the Cross-Entropy Loss between the decoded vocabulary distribution of the random vector, $D(v)$, and our resampled top-k distribution $Q_\{tk\}$ using gradient descent w/ the Adam optimizer. $Q_\{tk\}$ is a distribution containing the top $k$ tokens with the highest probability from $Q_t=D(v_t)$; the rest of the vocabulary is zeroed out. In equation 6 we use the term “CE” but realized that we never pointed out this was Cross-Entropy Loss. We have clarified this in the updated version of the paper.

---

What happens if the FV is the SET of averaged $a_\{lj\}$, with each inserted into the appropriate $l,j$ position? If this isn’t better than ‘adding the FV to the hidden state at layer $l$’, why not?

Great question. We have added this analysis to appendix M in the updated submission of the paper. A function vector can be thought of as adding several attention head outputs to a single layer’s hidden state. But, if we instead add the attention outputs to its corresponding layer (or at its corresponding ($l,j$) position), we find that this generally matches the performance of our function vector.

---

The paper would be strengthened if the authors could provide some results from carefully constructed test sets in which the set of outputs is the same set as the set of inputs (e.g matched antonym and synonym pairs, months, days of week, letters of the alphabet, progression of planets, chemicals in the periodic table) afford possibilities.

This is an excellent suggestion. We had similar thoughts with regards to tasks that have a cyclic nature, where the input and output spaces are similar—for example, antonyms and some of these others you have mentioned which cannot be done via a simple vector offset like perhaps some other tasks can. We have added some discussion about this idea in an updated version of the paper, in addition to an appendix section (appendix A) discussing this in more detail. We also discuss how this contributes to the idea that FVs are not just simple word vector offsets.

We also created two new tasks - (1) “next item” which includes pairs such as (monday, tuesday), (three, four), (g, h), (February, March) etc. and (2) “previous item” which has similar pairs, but in the other direction: (tuesday, monday), (four, three), (h,g), (January, December). We find that function vectors are able to trigger this type of cyclic behavior when the input and output spaces are similar. We have added our analysis of function vectors in this setting to Appendix L in the updated version of the paper.

---

Examining function vectors in vocabulary space

We analyze this in Appendix N, and we find that function vectors have causal effects beyond what is captured by top-1 accuracy.

Thank you for your interest in the paper, and your helpful feedback! We address your additional comments below:

---

please be explicit in the text with a phrase that a_lj is not the head output as in Vaswani et al but its projection through a subsequent matrix that scales it up to the model dimension, then refer to the appendix for the details.

We have updated the formulation in section 2.2 to make this explicit.

---

Additional Cyclic Tasks

The cyclic tasks are two additional tasks. Next Item and Previous Item. Both tasks are collected over a heterogeneous set of sequential data that includes cyclic types such as days of the week, months of the year, and letters of the alphabet, as well as non-cyclic types such as numbers and roman numerals. For example, a single ICL example for the Previous-Item task might look like “Q: six\nA: five\n\nQ: a\nA:z\n\nQ: VII\nA:VI\n\nQ: September\nA:”. The model would ideally be able to answer “August” given this ICL prompt.

We collect function vectors from ICL examples such as the one above to get a “Previous-Item” function vector. We then test it on zero-shot examples such as “Q: Monday\nA:”, which when we add the function vector, the model outputs “Sunday”. And for similar queries we observe the following behavior: Monday+FV = Sunday, Sunday+FV=Saturday, Saturday+FV=Friday, Friday+FV=Thursday, Thursday+FV=Wednesday, Wednesday+FV=Tuesday, Tuesday+FV=Monday.

We call the input and output spaces similar because the task defines a function that maps numbers to numbers, days of weeks to days of weeks, and months to months. In addition, we say the function has cyclic structure because if the function is applied repeatedly it ends up arriving at the same data point for certain subsets. For example, Monday will map back to Monday again. This cyclic structure supports the arguments in Appendix A (point 1), and at the bottom of page 48.

---

Causality and Mechanisms

We appreciate the reviewer’s clarity around the distinction between causality and full explication of mechanisms. Our paper has taken one step in elucidating a mediating mechanism of ICL: function vectors. But we agree that we have not characterized the inner workings of these function vectors themselves, which we leave for future work. We have updated the Discussion section to reflect this.

As described in our updated discussion, prior to this work the mechanisms underlying ICL were largely unknown. Our identification of function vectors can be thought of as analogously moving from knowing “going to a party causes headaches” to “red wine causes headaches”. Although we haven’t elucidated every detail of how language models perform ICL, the identification of function vectors as a mediator is a key step forward and opens the path for follow-up work.

---

Tasks that include choosing a word from a list

Apologies for overlooking this comment. We have updated section 3.3 to include some discussion about the difference between tasks that require "word-selection" and "word-transformation" and potential implications for future research.