Listening to the Wise Few: Select-and-Copy Attention Heads for Multiple-Choice QA

Authors: Eduard Tulchinskii, Kristian Kuznetsov, Laida Kushnareva, Anastasia Voznyuk, Andrei Andriiainen, Evgeny Burnaev, Irina Piontkovskaya, Serguei Barannikov
Primary area: foundation or frontier models, including LLMs
Keywords: large language models (LLMs), attention mechanisms, model interpretability, zero-shot learning

Abstract:
Multiple-choice question answering (MCQA) is one of the most widely adopted methods for evaluating large language models (LLMs). In this approach, the model is presented with a question and a set of possible answers, and the answer with the highest logit is selected as the model's prediction. However, this evaluation format has limitations, as even if the model knows the correct answer, it may struggle to select the corresponding option simply due to difficulties in following this rigid format. Methods such as instruction tuning or in-context learning help alleviate this issue but introduce their own biases, such as dependence on the order and semantics of training examples. In this paper, we address this issue by conducting an intrinsic investigation of the LLM’s decision-making process when answering multiple-choice questions. Specifically, we identify and study specific select-and-copy heads responsible for choosing the correct answer. We develop new scores to reveal the underlying knowledge from these heads: the Query-Key Score, which measures the interaction between query and key representations in the selected head, and the Attention Score, which is based on the attention weights. By studying these scores, we found that the most pronounced select-and-copy heads are consistent across four popular Multi-Choice Question Answering (MCQA) datasets. Moreover, our scores enable better knowledge extraction, achieving up to a 16% gain for LLaMA2-7B and up to 10% for larger models on these benchmarks. On a synthetic dataset, where the correct answer is known explicitly, accuracy increases by nearly 60%, confirming the method's effectiveness in overcoming MCQA format limitations. To support our claims, we conduct experiments on models ranging from 1.5 billion to 70 billion parameters, in both zero-shot and few-shot settings.

Decision: reject
Reviewer scores: 3, 5, 6, 5

Reviews and discussion:
<thread>
creation date:
1729237738311

thread_type:
official_review

<thread_object>
creator: reviewer_u9do

<content>
summary:
The authors use a small amount of labeled data to identify "select-and-copy" heads. These are heads where information is being copied (via the attention mechanism) from a token the authors associate with a particular answer option (e.g., the newline after the option text) to the final token that will be used for prediction. This selection is primarily done by max "QK-score" (dot product between the query of the last token and key of the token associated with the answer option). The authors show "select-and-copy" heads are present in a variety of Llama models. They argue that using these heads for prediction leads to better accuracy and is less dependent than baseline on the order of answer options. The authors also run some design and head ablations, and explain an approach to finding "select-and-copy" heads in a label-free way.

strengths:
* ORIGINALITY: I'm not very familiar with mechanistic interpretability work, but in my view this work seems quite novel. The authors find attention heads that seem to have a very well-defined function for MCQA, and show these are present in many models.
* QUALITY: The authors' experiments seem well-designed and their argument (at least regarding the presence of "select-and-copy" heads) is convincing. The authors also have a sizable appendix, suggesting they've tried a lot of things.
* CLARITY: The paper is mostly clear and easy to read.
* SIGNIFICANCE: I think this is significant in that it provides more insight into the mechanisms behind MCQA in LLMs.

weaknesses:
* Primary weaknesses
  * This paper heavily emphasizes the zero-shot case, but I don't think it should be highlighted (I think it should be moved to the appendix if included at all). This is because in the zero-shot case the right way of answering (for the baseline models) is ambiguous. A human wouldn't know whether to respond with a letter vs the answer option text. I think e.g., Table 1, for example, should not show zero-shot results. I don't think the zero-shot setting is a fair setting for comparison.
  * I don't think the "E" and "F" options should be included (or at most this should be moved to appendix). As far as I know, adding the "E" and "F" options is not consistent with the majority of prior work in MCQA, and seems like an added variable that's not justified.
  * The authors pitch QK-score as being better than PriDe, and also having much improved accuracy across answer orders. However, I am not convinced of either of these. In the 1+ shot, no "E"/"F" setting PriDE seems as good or better. Also in e.g., Figure 3, the drop for PA is substantial for QK-score (just as substantial as for the alternatives). It seems like, from the appendix, the baseline is actually better than QK-score in many cases. Just to be clear, I don't think the authors' method needs to be more accurate than alternatives for the paper to be useful or accepted. I'm just saying the authors could maybe reassess their claims a little bit.
* The authors only consider Llama models, so it's unclear if these results apply to other LLMs.
* I don't fully follow the "Best Heads" part and Figure 5a in Section 6 despite having read it a few times. I definitely get the point being made, but I couldn't reproduce the result based on the description. To improve understanding and reproducibility, it might be nice to include a step-by-step description or pseudocode.
* I didn't take this into account in my rating, but I think the paper could benefit from another solid pass just for grammar. There are just enough errors that at times it was a bit distracting.
* See questions for more things I think could use clarification/improvement.

questions:
* Is RoPE applied when using attention score?
* Why is "stochastic" used to imply "sums to one" on line 147 (I may be missing something)?
* For Llama base vs chat models was the same prompt used? Would this lead to worse performance for the baseline?
* Why the big difference in accuracy for e.g., HaluDialogue vs small difference in accuracy for MMLU? I don't find the argument on 322-323 convincing.
* Why is attention score included? It seems like QK-score is used as the default, and attention score is barely mentioned. My inclination would be to move the attention score parts to the appendix to prevent confusion over when which score is being used. At the very least, there could be more clarification on exactly when each is being used.
* In the unsupervised head finding part, what accuracies do the top heads achieve? I'm curious if they're like 90% as good as the best ones, or if they're much worse because their function matches but they're doing something entirely different.
* Why not ensemble heads?
* I'm curious why accuracy remains quite high (despite the drop) in the head removal ablation (especially in higher shot setting). Is it just that there are more than 10 "select-and-copy" heads? Or do "select-and-copy" heads only explain part of what's going on?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730654224704

thread_type:
official_review

<thread_object>
creator: reviewer_nfpp

<content>
summary:
This work presents a new method for improving the evaluation of LLMs in MCQA by recognizing and using select-and-copy heads, which are particular attention heads. These attention heads consistently extract relevant information and improve response selection using the Query-Key Score (QK-score) and Attention Score. The strategy significantly improves MCQA benchmarks and a synthetic dataset for understanding. The study emphasizes the importance of intermediate attention states for disclosing underlying knowledge, particularly in smaller LLMs where typical output-based evaluation may understate the model's capabilities.

strengths:
The paper introduces the concept of attention heads that are adept at copying information relevant to MCQA tasks, advancing the interpretability of LLMs.
QK-score and Attention Score are presented as innovative metrics that provide deeper insights into model decision-making processes.
Strong experimental setup with results across different models and settings enhances the credibility of the findings.

weaknesses:
The approach focuses heavily on MCQA and may not generalize to open-ended or complex QA tasks.
Evaluating individual attention heads may be resource-intensive, especially for larger models.
While improving robustness, the paper does not fully address biases inherent to specific head selections.
Performance can differ based on head choice, potentially introducing instability in applications without careful selection.

questions:
1. Investigate the relevance of the methodology to a broader range of QA formats and practical open-domain tasks.
2. Suggest ways or instruments that facilitate the selection and utilization of appropriate heads for enhanced adoption.
What precautions were implemented to prevent the identified select-and-copy heads from introducing unintentional biases in model outputs?
4. How is the effectiveness of these attention heads different for different model types, such as encoder-only vs. decoder-only?
5. Is it possible to scale cross-lingual or multilingual multiple-choice question answering evaluation?
6. How well do the QK-score and Attention Score work when used in models that have been fine-tuned for specific topic tasks?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730719794779

thread_type:
official_review

<thread_object>
creator: reviewer_4bv1

<content>
summary:
This work introduces two new metrics—the Query-Key Score (QK-score) and the Attention Score—that utilize select-and-copy attention heads within the models to better capture their underlying knowledge. The authors argue that relying solely on logit scores to select answers can be misleading, especially for smaller models struggling with rigid formats. By using intermediate attention representations, this method reveals deeper insights into the model’s understanding, yielding accuracy gains of up to 16% on MCQA benchmarks such as MMLU and HellaSwag. The study finds that middle-layer attention heads are particularly effective, whereas later layers tend to revise and diminish performance. Overall, this work contributes an approach that not only improves MCQA accuracy but also enhances interpretability of LLMs.

strengths:
1. The authors introduce QK-score and Attention Score for deeper evaluation of LLMs.
2. This method yields significant gains in MCQA tasks, with up to 16% improvement on some benchmarks which is quite significant.
3. The work leverages internal attention heads, offering transparent answer selection.
4. The authors demonstrate the effectiveness of middle-layer attention heads over final layers.
5. This method is tested across models ranging from 7B to 70B parameters

weaknesses:
1. While the experiments have been performed across different generations of llama models, showing generalization across model families could be important
2. Although the method is effective, there is complexity in terms of implementation. The applicability of the method to various practical scenarios remains questionable

questions:
1. Some experiments / results on other model families
2. Comment on usability. (refer to comment #2 in weakness)
3. Given the complexity of the method, will be interesting to see latency analysis when compared to baseline.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730727890533

thread_type:
official_review

<thread_object>
creator: reviewer_jnc8

<content>
summary:
The widely used evaluation for large language models, multiple-choice question answering (MCQA), is very brittle, especially for small models -- existing works show that even if models know the answer, it often cannot output the correct A/B/C/D due to all sorts of bias. This work proposes to tackle the problem by looking at a novel QK-score: they first select certain "select-and-copy" attention heads based on a validation set, and then calculate the query-key dot product between the option and the question (there are many possible ways, and the authors conducted thorough ablations). 

The authors conducted comprehensive experiments on commonly used datasets, with zero-shot/many-shot experiments across model scales. The proposed method significantly improved over the standard MCQA baseline and some previously proposed methods. The analysis revealed interesting aspects, such as the meaning of a phrase is often encoded in the last token of the phrase.

My main concern is:

(1) Cloze completion has been widely used and has shown to be much more stable than MCQA in most standard evaluations. There is almost no discussion on it and also no empirical comparison. Since the work's main goal is to make evaluation more reliable, I found the lack of comparison significantly undermines this work's contribution.

(2) Improving the score doesn't make one evaluation better -- the authors should show that it reflects a better comparison that is more consistent with human evaluation or some intuition (for example, previous evaluations show much higher variance or reversed trends like an 80B model is worse than 7B; this new method fixed it).

strengths:
(1) The brittleness of MCQA is well known and is a problem in evaluation. The proposed method is intuitive, simple, and effective. 

(2) The authors conducted a comprehensive evaluation and interesting analysis that demonstrated the effectiveness of the method.

(3) The proposed method can be used beyond standard evaluation, especially in interpretability applications.

weaknesses:
My main concern is as shown below

(1) Cloze completion has been widely used and has shown to be much more stable than MCQA in most standard evaluations. There is almost no discussion on it and also no empirical comparison. Since the work's main goal is to make evaluation more reliable, I found the lack of comparison significantly undermines this work's contribution.

(2) Improving the score doesn't make one evaluation better -- the authors should show that it reflects a better comparison that is more consistent with human evaluation or some intuition (for example, previous evaluations show much higher variance or reversed trends like an 80B model is worse than 7B; this new method fixed it).

questions:
Please see the "weaknesses" section + the question below

(3) How does the variance of each method look like for the main table/figure, especially when sampling different in-context examples + different orders?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732304576173

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Reply to the common concern about the generalisation capabilities of our method to models other than LLaMA 1-3, 7-70B

comment:
We appreciate reviewers for comprehensive feedback! The main question that was expressed by the most reviewers is the generalisation capabilities of our method to model families other than LLaMA. Here we would like to address this concern.

We applied our method to smaller models: Qwen 2.5-1.5B (-Instruct and -Base) and Phi-3.5-mini-Instruct (3.8B parameters). We identified which attention heads within these models utilize the select-and-copy mechanism to answer the given questions and assessed their capability in doing so. Additionally, we found that the QK-scores for these heads are much closer to the baseline in terms of accuracy and permutation score, compared to the results observed for LLaMA 7B-70B. The results are shown in the tables below.

In particular, we confirmed that Qwen 2.5-1.5B-base has several heads that are consistently good across real datasets and synthetic dataset, such as heads (20, 4) and (21, 11). The consistency of these good heads is similar to that observed in LLaMA-2-7B and other models in the LLaMA family. However, the layers with the best heads in Qwen 2.5 are closer to the final layer than those in the LLaMA-family models. Specifically, the best heads in this model are concentrated around layers 16-22, while the model has a total of 28 layers (see Appendix M “Best heads on synthetic dataset for Qwen 2.5-1.5B”). Nevertheless, the general pattern still holds: very early and very late layers do not contain very strongly pronounced select-and-copy heads, and there are some consistently good heads across several datasets, as we indeed stated in our paper.

Metrics of our method applied to Qwen2.5-1.5B-Instruct:

```
|                | MMLU            | Cosmos          | Hellaswag       | HaluDialogue    |
|----------------|-----------------|-----------------|-----------------|-----------------|
|                | 0-shot | 5-shot | 0-shot | 5-shot | 0-shot | 5-shot | 0-shot | 5-shot |
| Baseline (acc) | 58.3   | 59.4   | 74.1   | 76.7   | 54.8   | 61.4   | 29.7   | 45.3   |
| QK-score (acc) | 57.5   | 58.3   | 76.0   | 75.8   | 59.4   | 59.0   | 36.0   | 44.3   |
| Baseline (PA)  | 49.2   | 50.2   | 67.8   | 70.8   | 46.2   | 52.4   | 19.3   | 34.2   |
| QK-score (PA)  | 41.1   | 46.0   | 69.1   | 68.7   | 51.0   | 48.5   | 12.8   | 25.6   |
```

Metrics of our method applied to Qwen2.5-1.5B-base:

```
|                | MMLU            | Cosmos          | Hellaswag       | HaluDialogue    |
|----------------|-----------------|-----------------|-----------------|-----------------|
|                | 0-shot | 5-shot | 0-shot | 5-shot | 0-shot | 5-shot | 0-shot | 5-shot |
| Baseline (acc) | 58.9   | 58.5   | 77.0   | 77.6   | 60.6   | 56.1   | 41.4   | 41.0   |
| QK-score (acc) | 57.3   | 56.6   | 74.8   | 77.2   | 56.9   | 54.0   | 42.8   | 43.6   |
| Baseline (PA)  | 49.1   | 47.7   | 70.1   | 71.7   | 49.8   | 43.3   | 29.0   | 29.7   |
| QK-score (PA)  | 42.3   | 47.6   | 65.5   | 70.7   | 43.4   | 44.0   | 30.1   | 30.5   |
```

Metrics of our method applied to Phi3.5-mini-instruct:

```
|                | MMLU            | Cosmos          | Hellaswag       | HaluDialogue    |
|----------------|-----------------|-----------------|-----------------|-----------------|
|                | 0-shot | 5-shot | 0-shot | 5-shot | 0-shot | 5-shot | 0-shot | 5-shot |
| Baseline (acc) | 65.9   | 68.1   | 77.5   | 81.1   | 78.0   | 75.7   | 59.7   | 58.2   |
| QK-score (acc) | 66.1   | 65.8   | 81.8   | 81.6   | 78.4   | 75.1   | 63.2   | 65.5   |
| Baseline (PA)  | 58.7   | 60.6   | 71.7   | 75.9   | 73.0   | 70.0   | 50.2   | 47.6   |
| QK-score (PA)  | 56.2   | 57.5   | 75.5   | 74.9   | 71.7   | 69.0   | 52.2   | 52.7   |
```

Regarding the proximity of the QK-score to the baseline score in these setups, we hypothesise that this may be because these models were more specifically fine-tuned for multiple-choice question answering (MCQA). This hypothesis is supported by the observation that these models (even relatively small ones with 1.5B parameters) achieve much better performance than LLaMA 2 with 7B parameters in baseline setups.

Additionally, we observed that the performance of the baseline is very similar in both 0-shot and 5-shot setups for these three models, particularly for Qwen Base and Instruct. It seems that these models are so well-accustomed to the MCQA format that they do not even require example prompts to understand how to answer MCQA questions to the best of their ability. We hypothesise that this is connected with very effective propagation of the signal from the select-of-copy heads, which provides accuracy similar to baseline scoring from the last layer.

We will include the complete results in all setups (including 1-, … 4- shot prompting) in the Appendix L of our paper before the rebuttal period ends. We also plan to add the results from bigger models.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This study introduces a novel method aimed at enhancing the evaluation of large language models (LLMs) in multiple-choice question answering (MCQA) by leveraging select-and-copy heads, which are specific attention heads. These heads consistently extract pertinent information, thereby improving response selection through the use of the Query-Key Score (QK-score) and Attention Score. The proposed strategy leads to significant advancements in MCQA benchmarks as well as on a synthetic dataset designed for comprehension.

However, the reviewers have raised significant concerns regarding the experimental setup, the generalizability of the proposed method, and its comparison with previous work. The authors should address these issues to enhance the paper's credibility and persuasiveness.