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

## Abstract

Abstract A standard way to evaluate the abilities of LLM involves presenting a multiple-choice question and selecting the option with the highest logit as the model’s predicted answer. However, such a format for evaluating LLMs has limitations, since even if the model knows the correct answer, it may struggle to select the corresponding letter simply due to difficulties in following this rigid format. To address this, we introduce new scores that better capture and reveal model’s underlying knowledge: the Query-Key Score (QK-score), derived from the interaction between query and key representations in attention heads, and the Attention Score, based on attention weights. These scores are extracted from specific select-and-copy heads, which show consistent performance across popular Multi-Choice Question Answering (MCQA) datasets. Based on these scores, our method improves knowledge extraction, yielding up to 16% gain for LLaMA2-7B and up to 10% for larger models on popular MCQA benchmarks. At the same time, the accuracy on a simple synthetic dataset, where the model explicitly knows the right answer, increases by almost 60%, achieving nearly perfect accuracy, therefore demonstrating the method’s efficiency in mitigating MCQA format limitations. To support our claims, we conduct experiments on models ranging from 7 billion to 70 billion parameters in both zero- and few-shot setups.

## 1 Introduction

Questions with multiple answer options are a common form
of benchmarks evaluating question answering (Hendrycks et al., ), common sense (Zellers et al., ), reading comprehension (Huang et al., ), and other abilities of large language models. In multiple choice question answering tasks (MCQA), the model is provided with the question and multiple answer options, e.g. ”Question: How many natural satellites does the Earth have? Options: A. 0. B. 1. C. 2. D. 3. E. None of the above. F. I don’t know.” Sometimes the context that might be helpful to give the answer is added before the question, such as a paragraph or a dialogue for reading comprehension or some common sense reasoning. The model is asked to output the letter denoting the correct answer option. This format is similar to certain real-life students’ exams and shares some benefits with them: it is straightforward to evaluate, using automated tools.

On the other hand, for LLMs, especially smaller ones, understanding and adhering to a multiple-option format is not always trivial. The model’s performance on a given multiple-option dataset depends not only on the ability to solve the task itself but also on its in-context learning or instruction-following capabilities. The model may produce correct answers with formatting issues, which hinders the automatic evaluation of the MCQA task. Consequently, some works delegate answer evaluation to another LLM instead of relying on exact string comparison (Wang et al., ).
When assessing the logits of the model for options, LLMs can follow shallow patterns such as options distribution. Some LLMs are inclined to prefer the answer option “A”, while others tend to choose “D” (Zheng et al., ). All of these issues demonstrate pitfalls in the current MCQA evaluation process, especially for smaller LLMs.

However, the model’s inability to follow the task format does not imply a lack of actual ’knowledge’ regarding the correct answer.
In this work, we show that while small LLMs generally perform poorly on MCQA benchmarks, their intermediate attention states can often provide better insights. Specifically, we introduce the method that uses the queries and keys within individual attention heads to select an answer. We identify certain select-and-copy heads that can choose the option with semantically relevant information and transfer further its representation. Our findings suggest that LLMs process MCQA tasks more effectively in the middle layers but tend to revise this information in the later layers, leading to reduced performance. Our method belongs to the class of ”white-box” techniques, meaning that we use the internal representations of LLMs to extract solutions for given tasks.

Figure: Figure 1: Our method calculates the Query-Key score between the end-of-line token of an answer option and the last token of the prompt for the designated head, from which we derive the answer.
Refer to caption: extracted/5897460/images/QK_new_version_small_shadow.png

We identify the principal elementary algorithmic operation performed by pretrained Transformer models when answering multiple-choice questions. This task is complex, requiring the model to
first compute a representation of semantic information contained in both the question and the options inside such specific heads.
After completing this task, the model selects the most appropriate option using the query-key alignment mechanism, see section 3, and then copy and outputs the option.
Based on this intuition, we propose to identify the heads in the model that perform this select-and-copy operation on the aggregated embeddings of the possible answers. Our results reveal the presence of such heads in all models we examined, ranging from 7 billion to 70 billion parameters. Remarkably, the best few performing heads are the same for different datasets. Moreover, the answers produced by these heads are significantly more accurate than the final output of the model, particularly in zero-shot scenarios.

Our contributions are as follows:
(1) We demonstrate the presence of select-and-copy heads in LLMs of the wide range of 7-70B parameters, performing option selection operation for MCQA task; (2) We introduce QK-score, along with attention score, the option scoring methods based on key and query representations derived from such heads. This scoring leads to 9-16% improvement of the accuracy with task-specific heads;
(3) We demonstrate that our method is more stable than the baselines to option permutations, renaming and also when supplementary options, like “I don’t know”, are added;
(4) Our results provide further support for the hypothesis observed in other papers, e.g. Li et al. (); Stolfo et al. (), that the representations of the semantic meaning of a phrase are encoded in certain heads in the query, key, and value vectors of the phrase’s last tokens—namely, the end-of-sentence punctuation token or the end-of-line token;
(5) We study the attention patterns of select-and-copy heads and their behavior under various conditions. This is a step towards better understanding how the LLMs work in general.

## 2 Related work

Question answering datasets are the standard way to measure capabilities of the Large Language Models to retain the knowledge, understand given texts, and perform reasoning. One can see the results of such testing in a number of technical reports on new LLMs, such as LLaMA2, LLaMA3, gpt-4-o, or Claude 3 Opus (Touvron et al., ; Dubey et al., ; OpenAI, ; Anthropic, ). Many commonly used benchmarks incorporate MCQA tasks, as they are easy to evaluate  (Ye et al., ; Pal et al., ).

Multiple choice prompting (MCP), when we present the question and multiple answer options to model, has many advantages over cloze prompting (CP), when a model is asked to complete partial inputs with a single probable word or phrase (Robinson & Wingate, ). While in CP one uses the normalized answer probabilities
for evaluation, in MCP we can use the probabilities of the options’ single tokens as a proxy. However, recent works highlight some issues
in evaluating models on MCQA tasks. Gupta et al. () and Pezeshkpour & Hruschka () show that permutation of the contents for options can significantly affect the accuracy using MCP. Similarly, Zheng et al. () describe selection bias for different LLMs and propose a debiasing method PriDe to boost the accuracy.

In our work, we investigate the inner mechanisms of LLMs, especially the role of attention heads in MCQA tasks. Functional roles of attention heads were analysed for transformer-based models from the very beginning of encoder-only models (Jo & Myaeng, ; Pande et al., ), and nowadays even more detailed approaches were developed for decoder-only models in the common track of mechanistic interpretability (Elhage et al., ; Olsson et al., ; Bricken et al., ). For example, induction heads identified by Elhage et al. () play an important role in in-context learning (Olsson et al., ; Von Oswald et al., ), indirect object identification (Wang et al., ) and overthinking (Halawi et al., ). Additionally, there is a number of research connecting theoretically constructed networks with real pretrained language models, revealing elements such as constant heads (Lieberum et al., ), negative heads (Yu et al., ), and content gatherer heads (Merullo et al., ), among others. For more information on mechanistic interpretability and attention heads, we refer to Rai et al. () and Zheng et al. ().

In our work, we focus on select-and-copy heads that are used to select the right option for the MCQA task. The special case of such heads looking on option labels was mentioned in (Lieberum et al., ). However, we show that not only are other tokens more representative for MCQA, but they can be used to significantly increase accuracy compared to baseline.

Moreover, our experiments conclude that those heads that outperform the baseline on MCQA are located on the middle layers of LLM. It correlates with previous findings that many information is present in earlier layers, but is somehow lost or revised in later layers (Kadavath et al., ; Azaria & Mitchell, ; Liu et al., ; Zou et al., ; CH-Wang et al., ). These studies mainly focus on linear probes of hidden representations (Ettinger et al., ; Conneau et al., ; Burns et al., ), but we show that the disagreement between model output and inner structures can be captured in the level of query-key interactions and attention maps.

## 3 Attention as select-and-copy algorithm

In this section, we describe how attention mechanism can work as select-and-copy operation.
Suppose we have a sequence of $N$ token embeddings $\{\mathbf{x}_{i}\}_{i=1}^{N}$, which serve as an input to the corresponding attention head of transformer, each $\mathbf{x}_{i}\in\mathbb{R}^{d\times 1}$. In classical transformer architecture Vaswani et al. (), each attention head performs the transform of input embeddings:

$$ ${\bm{o}}_{m}=\sum_{n=1}^{N}a_{m,n}{\bm{v}}_{n},\quad a_{m,n}=\frac{\exp\left( \frac{{\bm{q}}_{m}^{\top}{\bm{k}}_{n}}{\sqrt{d}}\right)}{\sum_{j=1}^{N}\exp \left(\frac{{\bm{q}}_{m}^{\top}{\bm{k}}_{j}}{\sqrt{d}}\right)},$ (1) $$

where ${\bm{q}}_{i}={\bm{W}}_{q}{\bm{x}}_{i}$, ${\bm{k}}_{i}={\bm{W}}_{k}{\bm{x}}_{i}$, ${\bm{v}}_{i}={\bm{W}}_{v}{\bm{x}}_{i}$ and ${\bm{W}}_{q},{\bm{W}}_{k},{\bm{W}}_{v}\in\mathbb{R}^{d_{model}\times d}$ are learned weight matrices. The resulting matrix ${\bm{A}}=\{a_{n,m}\}_{n,m=1}^{N}$ is stochastic, meaning that all its rows sum up to one.
For decoder transformers, causal mask is applied to ${\bm{A}}$ before softmax: $a_{i,j}=0,j>i$.
Thus, from equation 1, for each token position $k$ in decoder transformers we can write

$$ ${\bm{o}}_{m}=\sum_{n\leq m}a_{m,n}{\bm{v}}_{n}$ (2) $$

meaning that the $m$-th token of output embedding is the linear combination of values of the preceding tokens weighted by $m$-th row of the attention matrix ${\bm{A}}$. If all but one component here are close to zero, this transform can be considered as a conditional copy mechanism. Indeed, if $a_{m,j}$ is the only non-zero weight in the $m$-th row, then $a_{m,j}\approx 1$, and ${\bm{o}}_{m}\approx{\bm{v}}_{j}$ (by 2).
Each token position from 0 0 to $m$ can be considered as a cell storing the corresponding value vector; and attention weights $a_{m,0},\dots,a_{m,m}$ are responsible for the choice which cell to copy to the $m$-th output.

Based on this, we came up with the idea
of select-and-copy heads, which implement such copying mechanism. Namely, in this work, we are interested in finding heads in the model, which select the proper option and copy the information from it to the answer. In such heads, the attention of $m$-th row should be concentrated on a few selected tokens, where $m$ is the output answer position.

In modern models, positional encoding information can be represented as the additional transform of queries and keys
in Eq. 1. For example,
in Rotary Position Embedding (RoPE) (Su et al., ) the rotation function $R_{f}(\cdot)$ is applied to them before taking dot product. Pre-softmax logit of the standard attention becomes $R_{f}({\bm{q}}_{m})^{T}R_{f}({\bm{k}}_{n})=R_{g}({\bm{q}}_{m},{\bm{k}}_{n},m-n)$, introducing the dependency on the position shift $m-n$.

In this paper, we evaluate the efficiency of the answer options scoring derived from select-and-copy heads. We aim to choose heads which rely on options semantics rather than the position; to mitigate the effect of the relative position shift, we consider the QK-score, which does not use the positional shift when comparing the queries and keys (see details in the next section).

## 4 Approach

Consider some MCQA task
with the corresponding dataset $\mathcal{D}=\mathcal{D}_{val}\cup\mathcal{D}_{test}$, where each instance represents the request to the model, consisting of prompt, question, and labelled answer options (Fig. 1). Given the request, the model should generate the label of the best option from the request.

To find the heads in the model that implement the described above option selection mechanism, we pick best-performing heads using $\mathcal{D}_{val}$ based on the accuracy there, and then evaluate their performance on much larger $\mathcal{D}_{test}$. If such heads are in fact fully responsible for option selection, the performance just on them should be at least comparable to the performance of the whole model. We prove this claim by experiments in Section 5.4. Another way to select such heads is proposed in Section 6;
we demonstrate by attention maps analysis, that the best-performing heads indeed implement the option selection algorithm described above.

QK-score and Attention-Score.
Given a data sample of MCQA task, we denote by $q$ the question supported with context if applicable, by $o=\{o_{1},o_{2},...,o_{n}\}$ the semantic content of the provided answer options, and the corresponding labels by $d=\{d_{1},d_{2},...,d_{n}\}$ (e.g A/B/C/D); we believe that the labels are default-ordered. All these parts are concatenated to a string
$q\ast d_{1}\ast o_{1}\ast\dots\ast d_{n}\ast o_{n}\ast$,
where $\ast$ stands for any kind of delimiters, usually punctuation marks or newline characters (Fig. 1). The model should estimate $P(d_{i}\mid q,d,o)$ – the probability of option $d_{i}$ given the question $q$ and contents of the answer $o$, concatenated with the answer options $d$.

Let $t_{i},i\in\{1,2,...,n\}$ be the indices of tokens that incorporate knowledge about the corresponding answer options. We call them option-representative tokens. Choosing such tokens properly is important for the success of our algorithm. In most experiments, we use the end-of-line token after the $i$-th option content as $t_{i}$; other possible variants are presented in Fig. 2(a). We study them in Sec. 6.

Let $N$ be the length of the whole text sequence, and consider the head with index $h$ from the layer $l$. Then, given $(q,d,o)$, we can compute QK-score $S^{(l,h)}_{QK}(d_{i})$ for option $d_{i}$ (from query and key vectors), and Attention-score $S^{(l,h)}_{Att}(d_{i})$ (from attention weights):

$$ $S^{(l,h)}_{QK}(d_{i})={\bm{q}}_{N}^{(l,h)\top}{\bm{k}}^{(l,h)}_{t_{i}},\quad S ^{(l,h)}_{Att}(d_{i})=a_{N,t_{i}}^{(l,h)},\quad i\in\{1,2,...,n\}$ (3) $$

Our QK-score of $i$-th option is calculated as a dot product of the $t_{i}$-th key and the last query vector $q_{N}$ (see Fig. 1).
In QK-score we do not apply positional transformation, therefore it is not equal to the attention scores before softmax. The best token by QK-score does not necessarily correspond to the token with maximum attention, see Figure 8 for an example.

For each method, the prediction is straightforward: we take the option, for which the score gives maximum. By applying softmax function to scores we could also estimate the head and score specific probabilities for options.

Choosing the predicting heads. We do not aggregate heads predictions. Instead, we use the scores from the single best head, which is selected by the accuracy on the validation set $\mathcal{D}_{val}$. In Section 5.4, we report the results obtained from the best heads chosen separately for each dataset and each number of shots (i.e. number of examples provided in the prompt). In Section 6, we show that for each model, there exist universal heads working well on the most tasks and number of shots. Furthermore, we demonstrate that such universal heads can be found without access to labelled validation data.

Figure: (a)
Refer to caption: extracted/5897460/images/token_scheme_new.png

## 5 Experiments

### 5.1 Datasets

We experiment on four challenging real-world MCQA datasets from LLM benchmarks: MMLU (Hendrycks et al., ), CosmosQA (Huang et al., ), HellaSwag (Zellers et al., ) and HaluDialogue, which is a ”dialogue” part of HaluEval (Li et al., ). All of them consist of questions with four possible answer options and some additionally have a context to be used to give the answer. More details about each dataset can be found in Appendix A.1. Additionally, we introduce Simple Synthetic Dataset (SSD) created as a synthetic task in MCQA setting that will allow to estimate the ability of the model deal with the bare task format. Tasks from SSD do not require any factual knowledge from the model. The main version of this dataset contains questions of the form “Which of the following options corresponds to “$<$word$>$” ?” and contains 2.500 examples. Options include a word from the question and 3 random words, all mixed in a random order and marked by letters ‘A’-‘D’. Other variations of this dataset have another number of options, sampled and named by the same principle. These other versions are described in more details in Appendix H.

Finally, following Ye et al. () in all five datasets we specially modified questions by adding two extra options “E. None of the above.” and “F. I don’t know.” that are intended to aggregate the uncertainty of LLM. Despite adding these two options, there are NO questions for which ‘E’ or ‘F’ are correct answers.
Examples from all datasets are listed in the Appendix A.2, as well as prompt formatting we used.

Following the previous approach by Zheng et al. (), with fixed $N$-shot setup, we select $\mathcal{D}_{val}$ as $5\%$ of $\mathcal{D}$ for each dataset that is dedicated to assessing each head’s performance. Based on this evaluation, the best head is chosen and applied to other questions in the dataset.

### 5.2 Baselines

The standard approach for MCQA is to use output probabilities from LLM for all options $d_{i}$ to choose the predicted option $\hat{d}$:

$$ $\hat{d}=\operatorname*{arg\,max}_{d_{i}}P(d_{i}\mid q,d,o),$ (4) $$

where $q$ is the question, $o=\{o_{1},o_{2},...,o_{n}\}$ are option contents, and $d=\{d_{1},d_{2},...,d_{n}\}$ are the options labels (e.g A/B/C/D).
In our experiments we refer to this method as Baseline.

In recent work (Zheng et al., ),it was proposed to mitigate the option selection bias, averaging the results over options permutation. The idea is to use the set of all cyclic permutations $\mathcal{I}=\{(i,i+1,...,n,1,...,i-1)\}_{i=1}^{n}$ to calculate the debiased probability:

$$ $\tilde{P}(d_{i}\mid q,d,o)=\frac{1}{|\mathcal{I}|}\sum_{I\in\mathcal{I}}\log P (\pi_{I}(d_{i})\mid q,d,\pi_{I}(o))$ (5) $$

Since computing probabilities for all permutations for each question is expensive, authors propose to estimate the prior distribution for option IDs on test set which is $5\%$ of all samples, and use it to debias new samples. In our experiments we refer to this method as PriDe. The test set is the same as we use for the best heads selection.

### 5.3 Experimental setup

Our main experiments were carried out according to the following pipeline: first, we took a frozen pre-tranied Transformer LLM (its weights were not modified in any of the experiments); then, we passed through it questions from the validation subset and for each head of the model and each question obtained best in terms of QK-score answer. After that, we chose a single head on which the highest accuracy was achieved (if several heads appeared to have equal accuracy scores we chose one from the lower level of the model; although, in our experiments this happened extremely rare).
Then we obtained answer predictions via baseline method and via QK-score on the chosen head. Finally, we perform random shuffle of options in all questions and repeat the abovementioned procedure: it is done to correctly compute the Permutation Accuracy metric.
Note that it may be two different heads that achieve best QK-scores on validation set before and after option permutation.

We report two quality metrics on the test subset: accuracy of predicted answers (from the first run) and Permutation Accuracy (PA) metric. The latter was introduced in Gupta et al. () and is, in a sense, accuracy stable for choice permutation. PA metric is computed as the percentage of questions for which model choose correct choices before and after random permutation of options.
$\text{PA}~{}=~{}\frac{1}{N}\sum_{i=1}^{N}\mathrm{I}_{i}\mathrm{I}^{p}_{i},$
where $N$ is the dataset size, $\mathrm{I}_{i}$ is the indicator value equals to 1 iff model’s answer on question $i$ is correct, while $\mathrm{I}^{p}_{i}$ equals to 1 iff model gives correct answer on question $i$ after its options (their texts not letters) were permuted.
Answer options “E. None of the above.” and “F. I don’t know.” are special and therefore are exempt from shuffling.

The prompt templates we use in our experiments are provided in Appendix A.3.
In few-shot regimes before asking the question we provide model with demonstrations in the same format except that the true answers (single capital letter for the correct option) are given after each example separated by single whitespace. Examples are separated from each other and from the actual question by single line breaks.
The demonstrations are the same for every question in the given dataset. The set of examples for ($k+1$)-shot prompts contains the set of examples for $k$-shot prompts and one new example. The demonstrations were chosen from the first fifteen entries of the validation set, and their choice was mostly arbitrary, but we tried to filter out questions that we considered suboptimal from the perspective of an English-speaking human expert.

### 5.4 Results

Figure 3 demonstrates the results of our method for LLaMA2-7B model. We observe an impressive improvement by 7-16% on all the datasets in zero-shot regime. Although QK-scores is not completely robust to option permutations, it is more stable than the baseline: the relative performance drop by PA metric is less than the baseline on all the datasets. In the few-shot regime, our approach is on par or outperforms other methods, with the most visible improvement on Halu Dialogue dataset by 5-9% depending on the number of shots.

PriDe results are added on Figure 3 for the comparison. PriDe in the most cases performs better than the baseline, but sometimes fails in zero-shot regime. Our analysis reveals that this method is not robust for additional uncertain options ”E” and ”F”. We additionally provide experiments without such options in Appendix B, where PriDe performs better in few-shot regimes, but still loses in 0-shot setup. But overall, in all cases and for any options set, QK score outperforms PriDe.

We also applied our method to larger models of LLaMA family: LLaMA2 (-13B, -70B) and LLaMA3 (-8B, -70B) as well as to their chat/instruction-tuned versions. Table 1 presents the results of our method for large models in zero-shot regime; full version including few-shot regimes is provided in Appendix G. Overall, the results are in line with those obtained for LLaMA2.
For all the smaller models of 8B and 13B size in zero-shot settings, our approach outperforms the baseline on all the datasets, both on accuracy and permutation accuracy, with the improvement up to huge 27%, achieved on HellaSwag dataset with LLaMA3-8B model. With larger models, MMLU is the most difficult benchmark for our method, likely because questions from it are oriented on general knowledge while our method by design focuses more on the semantic relations between the question and the possible answers.

Regarding performance of models on synthetic dataset SSD, in Figure 5(b) we can see that in baseline zero-shot setting LLaMA2-7B struggles to point onto the correct option, meanwhile, our method allows to extract the needed information from the model and gain much better quality. The figure shows accuracy for QK-score from five best heads (we denote them by their (Layer, Head) indices). Three of these heads can also be seen in the Table 4; the other two ((8, 8) and (12, 15)) are unique for this particular dataset.

Figure: Figure 3: Comparison of different methods for LLaMA2-7B (base) on various Q&A datasets. Reported metrics are Accuracy (Acc) and Permutation Accuracy (PA).
Refer to caption: extracted/5897460/images/shot_prompting.png

**Table 1: Comparison of different base models in zero-shot setup on various Q&A datasets. Reported metrics are Accuracy (Acc) and Permutation Accuracy (PA). Best results are highlighted in bold.**
|  |  | LLaMA… | LLaMA… (chat, instruct) |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Method |  | -30B | -65B | 2-13B | 2-70B | 3-8B | 3-70B | 2-13B | 2-70B | 3-8B | 3-70B |
|  |  | MMLU |  |  |  |  |  |  |  |  |  |
| Baseline | Acc | 50.4 | 48.3 | 34.6 | 59.7 | 60.3 | 75.3 | 47.4 | 57.7 | 60.5 | 78.2 |
| PA | 37.9 | 35.7 | 22.4 | 48.5 | 50.4 | 68.8 | 34.6 | 45.9 | 47.7 | 70.1 |  |
| QK-score | Acc | 45.2 | 46.2 | 42.2 | 56.7 | 61.0 | 74.5 | 49.7 | 58.9 | 63.0 | 77.9 |
| PA | 30.7 | 32.1 | 25.9 | 39.2 | 51.5 | 66.0 | 38.3 | 47.1 | 49.3 | 67.9 |  |
|  |  | Cosmos QA |  |  |  |  |  |  |  |  |  |
| Baseline | Acc | 59.9 | 65.7 | 29.6 | 65.5 | 54.9 | 82.0 | 48.1 | 68.5 | 85.4 | 91.6 |
| PA | 47.5 | 53.1 | 19.4 | 56.3 | 39.3 | 75.7 | 36.8 | 58.3 | 71.0 | 82.5 |  |
| QK-score | Acc | 60.1 | 63.5 | 58.2 | 69.5 | 70.6 | 87.6 | 67.7 | 84.8 | 88.6 | 94.1 |
| PA | 44.4 | 50.8 | 44.3 | 56.2 | 60.9 | 81.7 | 51.6 | 75.9 | 75.1 | 88.1 |  |
|  |  | Hellaswag QA |  |  |  |  |  |  |  |  |  |
| Baseline | Acc | 35.2 | 33.4 | 36.8 | 71.6 | 33.5 | 82.5 | 41.6 | 61.4 | 67.4 | 86.8 |
| PA | 16.5 | 13.7 | 17.1 | 62.9 | 15.8 | 76.1 | 25.8 | 49.0 | 27.8 | 71.2 |  |
| QK-score | Acc | 43.9 | 53.8 | 52.9 | 74.9 | 60.9 | 82.1 | 50.8 | 73.0 | 72.5 | 86.3 |
| PA | 21.5 | 35.0 | 38.8 | 63.3 | 50.8 | 75.2 | 37.3 | 64.9 | 36.3 | 72.8 |  |
|  |  | Halu Dialogue |  |  |  |  |  |  |  |  |  |
| Baseline | Acc | 36.3 | 46.7 | 41.0 | 39.4 | 46.6 | 44.3 | 49.4 | 39.4 | 62.1 | 68.8 |
| PA | 21.1 | 29.8 | 22.2 | 25.4 | 29.1 | 33.5 | 32.6 | 26.6 | 42.6 | 63.8 |  |
| QK-score | Acc | 44.8 | 42.4 | 47.2 | 58.4 | 52.3 | 67.8 | 56.2 | 58.1 | 64.7 | 76.7 |
| PA | 27.6 | 22.5 | 30.2 | 42.6 | 36.7 | 57.9 | 42.5 | 42.8 | 46.6 | 65.6 |  |

## 6 Analysis

Choosing option-representative tokens.
To compare our scores, we need to select option-representative tokens $\{t_{i}\}$, where the semantic information about each option semantics is concentrated. Due to the causal nature of the attention in LLMs, the logical choice is the last token after the content of the option, which is the end-of-line token. We use it in most of our experiments, although there are other tokens worth analysing: label itself, period after label and period after option content (see Fig. 2(a)). We also experimented with the mean aggregated score through all tokens in the content of the option, but it gave poor results. The detailed analysis of such variations for attention scores is presented in Fig. 2(b). We observe that the period after content and the end-of-line tokens are the most representative of our scores. There is an interesting finding concerning label token: despite it being almost useless in 0-shot setup as was shown in (Lieberum et al., ) also, we can see the better performance for 5-shot setup, in different heads. We hypothesise that there exist several types of ”select-and-copy” heads, which influence the logits differently.

Select-and-copy heads ablation.
To investigate select-and-copy heads and their relation to the performance of the model,
we use zero-ablation of heads (Olsson et al., ) to analyze the causal relationship between select-and-copy heads and model output. In this method, we replace the output of a selected set of heads with a zero vector, effectively removing their contribution to the residual stream. We experiment with a set of the 10 best heads based on the Attention-score with EOL and label option-representative tokens and report the results in Fig. 4. Additionally, we perform random ablations by selecting 10 random heads and aggregating the results from 5 runs. To ensure a fair comparison, we select random heads only from the middle layers (12 to 20), where the top heads are also located. We observe a significant drop in accuracy, sometimes below random performance, for the most of dataset, when ablating heads by EOL token, indicating a causal relationship with the model’s output. Some additional experiments where logit lens (Nostalgebraist, ) is utilized to further evaluate the layer-wise dynamics of the answer and select-and-copy heads in the model can be found in Appendix F.

Figure: Figure 4: Zero-ablation of heads for LLaMA2-7B (upper) and LLaMA3-8B (lower)
Refer to caption: extracted/5897460/images/heads_ablation.png

Best heads.

Figure: (a)
Refer to caption: extracted/5897460/images/heatmaps/Heatmap_Top_1_Heads_On_Whole_Set_small.png

As choosing the best head on validation set requires a sufficient amount of training data, we would like to determine whether there are universal heads that will be applicable to MCQA across multiple datasets and perform on par with the heads that are chosen for each task separately. Moreover, finding such heads would help mitigate the effects of a poorly chosen validation set, when discrepancies exist between the questions in the validation and test sets.

To illustrate how best-performing heads change in different setups, we select the best heads on the mixes across datasets and across shots, and select the 5% ot the best heads for each mix. The result is shown on Figure 5(a). This heatmap highlights the most stable heads, which appears among the best in several mixed tasks: when “shots” are mixed (“framed” cells), or when datasets are mixed (dark cells).
Most notably, the majority of robust heads in this sense lay within 12 and 21 layers. Then, we repeated the processing, but in this case we combined best-performing heads w.r.t. to the datasets. The most robust heads are (14,24) and (14,20). They appeared in the top 5% pairs in mixed-“shot” setup for all the datasets. They also demonstrate high performance on the synthetic data when the number of options is increased up to 24, as shown on Figure 5(b), while the performance of the baseline method drops below random. These results provide an additional evidence that the selected heads indeed able to perform the option selection task based on option content.
For more detailed analysis for 0-shot performance see analysis using percentiles in Appendix E.

Attention patterns analysis.
Figure 6 reflects the typical attention pattern together with QK scores for our most stable head (14, 24); attention patterns of the other best heads across our tasks - (14, 20), (14, 26), and (14, 13) (right top corner of the Figure 7) are reflected in Appendix, Figure 8. We can see that the attention weights are concentrated on option-representative tokens, namely $\backslash n$ symbols after options, with the highest weight on the correct option, and exactly that is expected from select-and-copy heads. Interestingly, that QK score gives the clearer picture.

Finding best heads without validation labels. Based on this observation, we can propose an algorithm to find such stable heads without a labeled validation set. Namely, such heads should have heavy attention weights on option-representative tokens and high variability in the options they attend to.
Thus, we can score each head using a product of two values: 1) sum of average attention weights to all $\backslash n$ symbols after options on this head; 2) a frequency of “choosing” any option aside of the most popular one (see formal definitions at Appendix I).
If some head doesn’t attend on the options, then the first value is close to zero, meanwhile when the head is “looking” on options, but “chooses” the same option most of the time, the second value is close to zero. Multiplying these two values yields low scores for heads that consistently “look” at the same option or ignore options entirely. Conversely, heads that “look” at diverse options receive high scores.
Sorting all the heads of LLaMA2-7B model by this score, we see that the named four heads have very high score. Namely, they get into top-20 heads if scored across real datasets we use, and they get into top-10 heads if scored on our synthetic dataset, see Figure 25.

Figure: Figure 6: QK-scores after softmax (upper part of the diagram) and attentions (lower part of the diagram) for the last token on the 0-shot MMLU example on (14, 24) head. The task is “Question: What singer appeared in the 1992 baseball film ’A League of Their Own’$\backslash$nOptions:$\backslash$nA. Brandy.$\backslash$nB. Madonna.$\backslash$nC. Garth Brooks.$\backslash$nD. Whitney Houston.$\backslash$nE. I don’t know.$\backslash$nF. None of the above.$\backslash$nAnswer:”. Full version is on Figure 8.
Refer to caption: extracted/5897460/images/attention_maps/attentions_14_24_with_softmax_qk.png

Selection bias.
Following previous studies on selection bias Pezeshkpour & Hruschka (); Zheng et al. (), we investigate our methods towards the tendency to choose specific option rather then choosing a correct answer. We observe that among best heads we also have uneven distribution in predictions, which are corrected as well when increasing the number of shots. However, there is an interesting pattern that two best heads distributions are complementing each other, i.e $S^{(14,20)}_{QK}$ is biased to options ”A” and ”D” an $S^{(14,24)}_{QK}$ - to options ”B” and ”C”. More detailed information can be found in Fig. 26 in Appendix J.

Figure: Figure 7: Accuracy of the best performing heads and of the most robust heads - (14, 24), (14, 20)
Refer to caption: extracted/5897460/images/accuracy_with_best_heads_V3.png

## 7 Conclusion

In this work, we introduced two novel scoring mechanisms: QK-score and Attention-score, derived from internal mechanism of LLM that can help to improve the performance on multiple-choice question answering tasks. Our experiments demonstrated significant improvements (up to 16%) across popular benchmarks, and even more striking results (up to 60%) on a synthetic dataset designed to test the model’s understanding of task format.

We identified a subset of attention heads, which we termed select-and-copy heads that play a critical role in these performance gains. These heads are relatively stable across different datasets and exist universally across model scales, and we explored their causal effect on task performance. Our findings suggest that these specialized heads have the potential to deepen our understanding of LLMs’ capabilities not only for MCQA but for other reasoning tasks as well.

This work opens up new avenues for further research into the internal dynamics of LLMs, including a deeper exploration of attention mechanisms and their role in complex task-solving that requires selection and copying information from the text.

## 8 Limitations

Our method cannot be applied to models without an access to attention matrices.
Also, our method is not applicable on scarce-resource tasks, even though one can utilize the heads we marked as robust enough.
Besides,
MCQA task itself was
criticized for oversimplification (Balepur et al., ).

## Appendix A Datasets

### A.1 Datasets Details

Massive Multitask Language Understanding (MMLU) (Hendrycks et al., ) contains 4-way questions on the variety of topics related to STEM, the humanities, the social sciences, and other fields of knowledge. We sample 10,000 instances from the test set to utilize them in our experiments.

CosmosQA(^1^11[https://wilburone.github.io/cosmos/](https://wilburone.github.io/cosmos/)) (Huang et al., ) together with question and answer options additionally contains text paragraph that is supposed to be used by a model to give the final answer. The purpose is to evaluate reading comprehension and commonsense reasoning capabilities of the model. Similar to MMLU, we sampled 10,000 instances from the test set.

HellaSwag (Zellers et al., ) evaluates the commonsense reasoning capabilities of the model through selecting the best sentence completion for a given sentence prompt, given a short text as a context. We, once again, extracted 10,000 entities from this dataset.

Halu Dialogue is a ”dialogue” part of HaluEval (Li et al., ) dataset with about 10,000 examples. Here a model is asked to choose an appropriate continuation of a dialogue from four possible options.

### A.2 Examples of questions from datasets

Figure: Listing 1: MMLU example

Figure: Listing 2: CosmosQA example

Figure: Listing 3: HellaSwag example

Figure: Listing 4: Halu Dialogue example

Figure: Listing 5: Simple Synthetic Dataset example

### A.3 Prompt Templates and Examples

Variable parts are highlighted in bold; whitespace placing is marked by underscores; position of line-breaks is explicitly shown by symbols ‘$\backslash$n’ (note that the last line always ends without whitespace or line break). In our datasets we ensured that each question ends with question mark, and each choice ends with point (single whitespace before it does not affect the logic of tokenization by LLaMA tokenizer).

Figure: Listing 6: MMLU prompt template

Figure: Listing 7: CosmosQA/HellaSwag/Halu Dialogue prompt template

Following are an example of $1$-shot prompts from MMLU. 2-3-4-5-shot prompts were built in the same way and prompts for dataset with context are built the same way, except each question is preceded by its context. Note that in demonstrations we add a single whitespace between “Answer:” and the correct choice letter; for example, “Answer: A”, but NEVER “Answer:A”. This is done because sequences like “: A” and “:A” are differently split into tokens by LLaMA tokenizer, and the former produces the same tokens corresponding to letter “A” as in the choice option line, while later yields a different version of “A”. From LLaMA’s point of view, these two versions of letters are separate entities and are NOT interchangeable. Removing those symbols of whitespace in many cases leads to noticeable drop in performance.

Figure: Listing 8: An example of 1-shot prompt for a question from MMLU dataset

## Appendix B Some more intuition on options ‘E’ and ‘F’

Figure: Figure 9: Correlation between heads QK-scoring accuracy on questions with 4 (‘A’-‘D’) and 6 (‘A’-‘F’) answer options. Solid red lines mark the accuracy level of $0.25$, dashed red line – $0.167$ (6 options random choice accuracy).
Refer to caption: extracted/5897460/images/appendix_B/cosmos-0shot-4-6-.png

As we mentioned in the main text, inclusion of fictional, though always incorrect, choices “E. None of the above” and “F. I don’t know” in every question was aimed at creating the “uncertainty sinks”. However, they are also beneficial for the analysis of attention head roles, but that is somewhat beyond the scope of this article. Here we would like to provide some intuition to it.

We performed experiments on a modified version of our datasets, where questions include only 4 “meaningful” choices, i.e. options ‘A’-‘D’ only. Scatterplots on Figure 9 show the correlation between accuracy of heads using QK-scores on options without ‘E’-‘F’ (by y-axis) and their accuracy on questions with all 6 options (by x-axis). Here, only validation subsets were used. We present plots for few of the possible setups, but other follow similar pattern. From these charts we can see that if a head reaches good accuracy answering 4-choice questions, it usually will reach nearly the same accuracy on questions with 6 choices and vice versa, see points around the diagonal $y=x$ in the upper-right quadrant.

We can also observe another major trend: horizontal stripe near y-level $0.25$. It can be explained in the following manner: in the data used, ground-truth answers are perfectly balanced – that is, for every choice ‘A’-‘D’ $25\%$ of the questions have it as the correct answer. And if a head reaches 4-choice accuracy of $\approx 25\%$, it falls into one of the three categories:

- 1.
This head chooses only one option in all questions. Usually it is the last one of the list.
- 2.
This head “guesses” answers, choosing options nearly randomly and “independent” from their meanings.
- 3.
This head “understands” questions, but is genuinely bad at answering them.

Addition of choices ‘E’ and ‘F’ drops the performance of the first type heads down to nearly $0\%$, second type – to around $16.7\%$; QK-scoring accuracy of the third type heads, however, usually remains the same.

Thus, we can conclude that choices ‘E’ and ‘F’ cause little effect on performance of good heads, but, at the same time, their inclusion creates separation between heads that are bad at Multiple Choice Question Answering and heads which do not have MCQA in their functionality at all (they may perform other roles for LM).

## Appendix C Numerical results for comparison of QK-score with other methods

Table 2 provides numerical results for our main experiments with QK-scores from heads of LLaMA2-7B model that are presented on Figure 3 in the main text.

**Table 2: Comparison of different methods for LLaMA2-7B (base) on various Q&A datasets. Reported metrics are Accuracy (Acc) and PErmutation Accuracy (PA). Best results are highlighted in bold.**
|  | …-shot prompting |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Method |  | 0 | 1 | 2 | 3 | 4 | 5 |
|  | MMLU |  |  |  |  |  |  |
| Baseline | Acc | 26.7 | 39.1 | 43.1 | 43.7 | 44.1 | 43.8 |
| PA | 8.9 | 21.3 | 26.2 | 27.4 | 28.5 | 28.4 |  |
| PRIDE | Acc | 15.5 | 36.9 | 39.8 | 40.8 | 41.5 | 42.7 |
| PA | 5.7 | 20.8 | 24.2 | 24.6 | 25.6 | 28.9 |  |
| Attention | Acc | 34.8 | 39.9 | 39.8 | 40.5 | 41.0 | 42.1 |
| score | PA | 17.2 | 19.4 | 21.9 | 23.4 | 24.1 | 24.3 |
| QK-score | Acc | 33.6 | 40.7 | 42.1 | 40.5 | 42.0 | 42.7 |
| PA | 17.2 | 21.7 | 23.7 | 22.0 | 23.4 | 24.0 |  |
|  | Cosmos QA |  |  |  |  |  |  |
| Baseline | Acc | 31.1 | 39.3 | 59.1 | 56.9 | 57.9 | 54.7 |
|  | PA | 11.1 | 21.9 | 44.1 | 39.2 | 40.7 | 35.7 |
| PRIDE | Acc | 15.2 | 44.6 | 58.6 | 59.2 | 60.7 | 61.3 |
|  | PA | 6.8 | 25.7 | 42.7 | 43.7 | 45.7 | 46.3 |
| Attention | Acc | 40.6 | 48.5 | 60.9 | 61.8 | 62.3 | 62.3 |
| score | PA | 23.8 | 28.3 | 46.8 | 47.7 | 48.4 | 48.2 |
| QK-score | Acc | 41.4 | 50.0 | 61.5 | 59.3 | 61.5 | 61.0 |
| PA | 25.6 | 33.6 | 47.3 | 43.6 | 47.1 | 46.4 |  |
|  | Hellaswag QA |  |  |  |  |  |  |
| Baseline | Acc | 26.5 | 28.8 | 30.6 | 33.3 | 36.1 | 34.6 |
|  | PA | 7.5 | 9.4 | 11.8 | 13.9 | 17.3 | 15.0 |
| PRIDE | Acc | 17.8 | 32.7 | 35.6 | 38.6 | 39.6 | 41.4 |
|  | PA | 4.9 | 12.9 | 16.0 | 20.5 | 21.8 | 23.2 |
| Attention | Acc | 34.8 | 40.0 | 41.7 | 42.6 | 43.2 | 43.5 |
| score | PA | 18.3 | 22.9 | 24.9 | 27.2 | 26.6 | 26.5 |
| QK-score | Acc | 33.0 | 37.1 | 40.4 | 42.7 | 45.7 | 42.3 |
| PA | 15.9 | 14.3 | 23.0 | 22.2 | 28.5 | 24.6 |  |
|  | Halu Dialogue |  |  |  |  |  |  |
| Baseline | Acc | 21.1 | 30.9 | 34.2 | 36.1 | 34.5 | 35.6 |
|  | PA | 5.4 | 10.2 | 14.3 | 18.9 | 16.8 | 20.7 |
| PRIDE | Acc | 3.0 | 32.0 | 35.5 | 36.3 | 36.5 | 36.1 |
|  | PA | 0.5 | 12.8 | 18.2 | 17.7 | 18.9 | 20.8 |
| Attention | Acc | 31.4 | 39.9 | 39.3 | 41.1 | 42.1 | 39.9 |
| score | PA | 10.9 | 19.4 | 17.5 | 19.0 | 22.3 | 21.3 |
| QK-score | Acc | 37.1 | 36.6 | 40.6 | 42.3 | 45.3 | 42.8 |
| PA | 17.7 | 14.6 | 19.6 | 22.0 | 25.9 | 22.2 |  |

## Appendix D Best Heads

**Table 3: Top 1% heads based on accuracy, intersected for 4 datasets on each setup separately**
| Setup | Best (Layer, Head) |
| --- | --- |
| 0-shot | (14, 24) |
| 1-shot | (15, 5), (15, 23), (14, 20) |
| 2-shot | (14, 24), (15, 5), (15, 4) |
|  | (18, 10), (15, 23), (16, 17) |
| 3-shot | (14, 24), (15, 5), (15, 4) |
|  | (18, 10), (15, 23), (14, 26), (17, 18) |
| 4-shot | (14, 24), (15, 5), (14, 4) |
|  | (15, 4), (18, 10), (15, 23) |
|  | (14, 20), (14, 26), (17, 18), (16, 17) |

## Appendix E Stability of Best Heads

Then, we utilized the minimum of accuracy percentiles as a way to determine stable heads, that can be seen on Figure 10(a). Again, the heads from 14th layer are showing the highest accuracy on almost all percentiles. We also listed the top 1% pairs for all setups based on accuracy in Table 4 and Table 4. There is a noticeable overlap between heads for various setups, and, once again, all of them are middle layers of the model.

If we compare the performance of the “stable” heads with results obtained with preceding calibration on Figure 7, (14, 24) and (14, 20) are frequently chosen from validation set, but even when they do not, their performance is comparable to their validation-chosen counterparts, except for HaluDialogue. Besides, we tested the heads (14, 24), (14, 20), (14, 26), and (14, 13) for a stability against increasing the amount of options in SSD dataset (see Figure 5(b)) and against changing the symbols that denote an options, following Alzahrani et al. () (see Appendix H). We also added other heads that are performing well on SSD dataset to these plots for comparison.

Figure: (a)
Refer to caption: extracted/5897460/images/headvis_llama2.png

## Appendix F Logit lens experiment

We follow Halawi et al. () and track the accuracy in the intermediate layers using logit lens (Nostalgebraist, ). Denoting ${\bm{h}}^{(l)}\in\mathbb{R}^{d}$ as a hidden state corresponding to last token in layer $l$, we extract intermediate probabilities for options $d_{i}$ using:

$$ $P_{l}(d_{i}\mid q,d,o)=\text{Logits}^{(l)}_{t_{i}},\qquad\text{Logits}^{(l)}= \text{Softmax}({\bm{W}}_{U}\cdot\text{LayerNorm}({\bm{h}}^{(l)}))$ (6) $$

Fig. 11 demonstrates the results for LLaMA2-7B base model, which shows some interesting patterns. In most cases, we see the improvements after the 12 layer for all setups excluding 0-shot. As we compare it with the maximal accuracy over Attention-score for two different types of option-representative tokens, we see the similar trend. However, the peak accuracy is seen in the middle layers, after which it degrades. Interestingly, that the logit lens performance demonstrate a sudden performance drop around 20. This indicates that some alternative “thoughts” about the answer emerges at this point, being overlapped further by the correct answer.

Figure: Figure 11: Logit Lens results on LLaMA2-7B base model for 0-shot and few-shot setups (upper) and a comparison to maximal accuracy per layer via Attention-score (lower)
Refer to caption: extracted/5897460/images/logit_lens_new.png

## Appendix G Comprehensive results for experiments on larger models

Here we provide complete results of our experiments with QK-scores on four main datasets (MMLU, CosmosQA, HellaSwag and Halu Dialogue) for larger models. As before, reported metrics are Accuracy and Permutation Accuracy.

- •
Figure 12 contains results for LLaMA2-13B, and Figure 18 for its chat-tuned version
- •
Figure 13 contains results for LLaMA2-70B, and Figure 19 for its chat-tuned version
- •
Figure 14 contains results for LLaMA3-8B, and Figure 20 for its instruct-tuned version
- •
Figure 15 contains results for LLaMA3-70B, and Figure 21 for its instruct-tuned version
- •
Figure 16 contains results for LLaMA-30B
- •
Figure 17 contains results for LLaMA-65B.

Note that in our experiments the accuracy scores for these baseline models are somewhat lower than ones you can see in the original technical reports  (Touvron et al., ; Dubey et al., ). The main reason for this is that we added additional “E” and “F” options that were not used in those reports; some differences in prompts and particular examples for few-shot learning also could play a role.
Also note that in many experiments we focus on zero-shot scenario without chain-of-thoughts prompting, that has received less attention in the original technical reports.

Figure: Figure 12: Comparison of different methods for LLaMA2-13B (base) on various Q&A datasets.
Refer to caption: x1.png

Figure: Figure 13: Comparison of different methods for LLaMA2-70B (base) on various Q&A datasets.
Refer to caption: x2.png

Figure: Figure 14: Comparison of different methods for LLaMA3-8B (base) on various Q&A datasets.
Refer to caption: x3.png

Figure: Figure 15: Comparison of different methods for LLaMA3-70B (base) on various Q&A datasets.
Refer to caption: x4.png

Figure: Figure 16: Comparison of different methods for LLaMA-30B (base) on various Q&A datasets.
Refer to caption: x5.png

Figure: Figure 17: Comparison of different methods for LLaMA-65B (base) on various Q&A datasets.
Refer to caption: x6.png

Figure: Figure 18: Comparison of different methods for LLaMA2-13B-chat on various Q&A datasets.
Refer to caption: x7.png

Figure: Figure 19: Comparison of different methods for LLaMA2-70B-chat on various Q&A datasets.
Refer to caption: x8.png

Figure: Figure 20: Comparison of different methods for LLaMA3-8B-instruct on various Q&A datasets.
Refer to caption: x9.png

Figure: Figure 21: Comparison of different methods for LLaMA3-70B-instruct on various Q&A datasets.
Refer to caption: x10.png

## Appendix H Behaviour of the best heads under the change of options symbols and options amount

Recall that aside of the standard version of Simple Synthetic Dataset (with four essential options and two additional options “E” and “F”) we consider alternative versions of SSD containing various numbers of possible options. For example, the variation of the dataset that corresponds to the number “10” on the x-axis of the Figure 5(b) contains ten essential options - A, B, C, D, E, F, G, H, I, J, - and two special options - “K. I don’t know” and “L. None of the above” (see the Example 9). Also note that in these experiments we used 200 examples from each version of the dataset to get the attention scores.

Figure  22 is an extended version of the Figure  5(b), that includes more heads for LLAMA2-7B (left) and the similar experiment for several heads of LLAMA3-8B (right), four of which are taken from the upper right part of the Figure  10(b) as the most stable across real datasets.

Figure: Listing 9: Modification of SSD with ten options - example

On Figure  23 we return to standard 4-options SSD dataset but change the symbols for option labels. We include (renamed) special options “E” and “F” for the upper plot and omit them for the lower plot for LLaMA2-7B model, and Figure  24 shows the same but for LLaMA3-8B model.

Figure: Figure 22: The results for a various numbers of options in the Simple Synthetic Dataset in zero-shot for LLaMA2-7B (left) and LLAMA3-8B (right). Different colors of the lines correspond to the results of QK dot products from different heads. “Square” markers correspond to the heads, working well across real datasets, and “round” markers correspond to the heads that work well on the synthetic dataset.
Refer to caption: extracted/5897460/images/instability/LLAMA-2-options-num-big.png

Figure: Figure 23: Performance of the QK-score from the best heads of LLaMA2-7B model for different options symbols when ”uncertainty” options (i.e. “I don’t know” and “None of the above”) are presented (upper figure) and not presented (lower figure). The accuracy of the best four heads from the Figure 10(a) declines in these new setups, but the head (12, 0) keeps being stable across all of setups. Another interesting head is the head (5, 11): it’s accuracy is high for all setups with “uncertainty” and for “$&# setup, but drops abruptly for “ABCD” and “ptgU”. Studying such “anomalies” is a matter for future research.
Refer to caption: extracted/5897460/images/instability/LLAMA2-symbols-6.png

Figure: Figure 24: Performance of the QK-score from the best heads of LLaMA3-8B model for different options symbols with and without “uncertainty” options. Interestingly, the best heads of the LLAMA3-8B model (see Figure 10(b)) are significantly more stable across considered setup.
Refer to caption: extracted/5897460/images/instability/LLAMA3-symbols-6.png

Figure: Figure 25: Left: average top heads scores across real datasets (first twenty). Dark blue marks four heads from the right top corner of Figure 10(a). Medium blue marks other heads with minimum of accuracy percentiles on all tasks more then 0.9. As we can see, the first two heads that get the best scores across real datasets, belong to the group of the best heads from the right top corner of Figure 10(a). Right: top heads scores on Simple Synthetic Dataset (first ten). Here, the top-scored head (14, 4) doesn’t appear at the right top corner of Figure 10(a), but it appears at the Figure 7 as one of the best heads for MMLU dataset. Note that for calculating this score we didn’t use the dataset labels.
Refer to caption: extracted/5897460/images/Heads_scoring/Average_heads_scores.png

## Appendix I Head scoring without validation set

Let $\hat{\mathcal{D}}$ be some unlabelled MCQA dataset. Then, for each head we may calculate a score

$$ $HeadScore=\left(\frac{1}{|\hat{\mathcal{D}}|}\sum\limits_{\hat{\mathcal{D}}} \sum\limits_{i=1}^{n}a_{Nt_{i}}\right)\left(\frac{1}{|\hat{\mathcal{D}}|} \mathbb{I}\{\operatorname*{arg\,max}_{i}(a_{Nt_{i}})\neq\hat{i}\}\right),$ $$

where $\hat{i}$ denotes the most frequent option for the given head; head indices $(l,h)$ are omitted. The left component here denotes the average amount of attention, concentrated on the option-representative tokens $t_{i},i=1,\dots,n$. The right component reflects the frequency of the situation, when the largest attention among the options falls on the option other than $\hat{i}$, i.e. any not the most frequent option.

The results of ranking heads according to this scores are presented on Figure 25.

## Appendix J Selection Bias

We investigate our methods towards the tendency to choose specific option rather then choosing a
correct answer. Fig. 26 presents a selection bias for baseline and 3 heads for QK-score in 0-shot and 5-shot regimes.

Table 5 shows the selection bias in terms of recall. We can see that most methods (especially in 0-shot setup) concentrates on single or several options.

Figure: Figure 26: Distribution of predictions across options for different methods on MMLU 0-shot (upper) and 5-shot (lower) setup. $(l,h)$ depicts the distribution for $S^{(l,h)}_{QK}$
Refer to caption: extracted/5897460/images/QKDistr0shot.png

**Table 5: Selection bias for different methods on MMLU 0-shot and 5-shot using LLaMA2-7B. The table compares original accuracy (for task to predict A/B/C/D/E/F) and recall only on subset with single ground truth option (i.e. only questions with answer A).**
| 0-shot | 5-shot |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Method | Orig. | A | B | C | D | Orig. | A | B | C | D |
| Baseline | 26.6 | 73.4 | 7.9 | 0.6 | 28.3 | 43.9 | 41.7 | 53.4 | 38.2 | 42.5 |
|  | (+46.8) | (-18.7) | (-26.0) | (+1.7) |  | (-2.2) | (+9.5) | (-5.7) | (-1.4) |  |
| $S^{(14,20)}_{QK}$ | 31.4 | 43.9 | 9.0 | 10.8 | 60.2 | 37.6 | 78.2 | 40.5 | 25.2 | 12.4 |
|  | (+12.5) | (-22.4) | (-20.6) | (+28.8) |  | (+40.6) | (+2.9) | (-12.4) | (-25.2) |  |
| $S^{(14,24)}_{QK}$ | 33.6 | 30.7 | 58.5 | 33.2 | 14.4 | 43.0 | 14.3 | 49.7 | 53.0 | 51.9 |
|  | (-2.9) | (+24.9) | (-0.4) | (-19.2) |  | (-28.7) | (+6.7) | (+10.0) | (+8.9) |  |
| $S^{(15,23)}_{QK}$ | 26.2 | 30.9 | 2.1 | 5.0 | 63.6 | 36.3 | 71.2 | 36.3 | 14.7 | 27.0 |
|  | (+4.7) | (-24.1) | (-21.2) | (+37.4) |  | (+34.9) | (+0.0) | (-21.6) | (-9.3) |  |

## Peer Reviews

confidence: 5
contribution: 3
initial_rating: 3
presentation: 4
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?
soundness: 2
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.
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).
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).

---

confidence: 5
contribution: 2
initial_rating: 5
presentation: 4
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.
soundness: 4
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
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.
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

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
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?
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
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?
soundness: 3
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.
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.
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.

Meta-review: 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.

Decision: Reject