Paper: Function Vectors in Large Language Models (ICLR-2024)

Reviewer comment (Novelty):
As a related issue, I think it is crucial for the authors to explain why the results are important. Currently, the authors say "Our study of function vectors differs qualitatively from these previous works: rather than training a model to create function representations, we ask whether a pretrained transformer already contains compact function representations that it uses to invoke task execution". In a way, I feel that this is underselling their work. It is clear without doing any research that pre-trained LLMs must have some form of task representations (otherwise they won't be able to do zero/few-shot learning). I think that developing a way to extract those representations is valuable, but a more thorough/careful discussion is crucial to highlight why this is important and how it advances existing work.

Author response:
__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.

Reviewer comment (Evaluation):
I might have missed it, but I wonder if the authors could provide a metric for the "proportion of task representation encoded by FVs". One potential way to visualize that would be to add lines to figure 4 representing best ICL performance, similarly for table 2, a few-shot baseline column would be very helpful. In other words, I do not fully understand how much task information we lose when we move from ICL to FVs.

Author response:
__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)

Reviewer comment (Presentation):
A minor suggestion on presentation: I feel that figures 1 and 2 are introduced much too early. They are confusing without proper context, which is given much later.

Author response:
__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.

Reviewer comment (Writing):
There is some back and forth on using all heads versus only manipulating the last head (am I understanding that correctly?).

Author response:
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.

Reviewer comment (Reproducibility):
FVs appear to be a vector over tasks; can you describe again how these are applied to the LLM to invoke task behavior? Do you mean that we are only looking at the attention heads to the last token (last column)? Then you sum over these heads for each task and add it to the last hidden state?

Author response:
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.

Reviewer comment (Experiments):
Why compute the top attention heads over all tasks and not just have a different distinct set for each task?

Author response:
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.

Reviewer comment (Writing):
In section 2.2, J is mentioned but not defined — does it mean the number of heads?

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

Reviewer comment (Theory):
Though the paper mainly investigates attention heads, I’m still wondering why attention heads contain such information rather than the MLP layers. I find some papers [1], and [2] seem to find that MLP is more critical from the information or memory aspect. Can the paper discuss more on this part, which would better explain the proposed findings?

Author response:
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.

Reviewer comment (Experiments):
I think the step from identifying useful attention heads to the proposed FV vector sounds empirical. For example, can the author explain how to determine the number of selected heads? I find that the paper uses |A| = 10 attention heads for GPT-J and scales the number for larger models. However, this point including the scaling and the initial 10 heads for GPT-J needs more explanation. Also, can the paper give more description about directly summing the head outputs together as defined in Eq. (5)? Has the paper compared with other ways like summing them up with different weights?

Author response:
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.

Reviewer comment (Experiments):
For experiments, I find though the paper tested 40 tasks, they are very simple. Thus, can the paper think of some more difficult ones? Meanwhile, I think if the paper can add some analyses about the FV ability under model scaling laws, it can be better.

Author response:
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.

------

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.