Answer, Assemble, Ace: Understanding How LMs Answer Multiple Choice Questions

Thank you for your positive assessment of our paper!

"...the OLMo 0724 7B Instruct (Figure 9) study used 0-shot accuracy…It would be interesting to see 3-shot results for the study or to hear a justification for use of 0-shot accuracy.”

Our primary motivation for showing the 0-shot learning curve is that it’s a lower bound on 3-shot performance, and we see in Fig. 9 that the model clearly learns the task (without needing in-context examples) early in training. But you raise a good point, and we’ll update Fig. 9 (and additionally, Figs. 21-22) to be 3-shot for consistency with the rest of the paper.

"The findings in this paper rely on some assumptions inherent in activation patching and vocabulary projection…it is something the authors could address in the paper if desired.”

Thanks for pointing this out. This is exactly why we include both methods, because they are complementary (lines 238-240). We initially had more discussion of their strengths & weaknesses that got cut for space– we will add this back in the Appendix of the updated PDF (coming shortly).

Thank you for acknowledging many positive aspects of our work; we can alleviate many of the weaknesses you mention with clarifications presented here and writing updates to our PDF. We've updated text in diagrams to be as large as possible given the page constraint, and made another pass to fix typos and put citations in brackets; thank you for these catches. Updated PDF coming soon!

"The claim in the lines 518-519…cannot be ​​deduced from the analysis of the checkpoints of only one particular Olmo 0724 7B base model (however, at lines 79-81 you make much more humble version of this claim, with which I agree).”

Thank you for pointing this out; we will update this sentence in our new PDF (coming shortly) to instead reflect the more humble version of the claim.

"In the lines 76-78…you only showed this effect for one non-standard letter set, and, again, I see only Olmo 7B experiments here (please, correct me if I'm wrong).”

When we say “the model’s hidden state”, we are referring to the Olmo 7B Instruct model mentioned earlier in the sentence. We are explicit in the paper that this result is for the Olmo model (lines 409 and 416), and thus represents one means by which models produce OOD letters (line 41). We'll further clarify this in the updated PDF though, and you're right that it will be valuable for us to add experiments for Llama and Qwen.

As for other non-standard answer symbols, we did observe this effect consistently in our preliminary experiments, but did not include these results in the paper. This is valuable feedback and we will add results on another randomly-selected letter set to show that our findings are robust.

Your questions:

"I don't understand the caption for the Figure 1, especially the sentence "Finally, when we switch to more unusual answer choice symbols, one behavior by which models adjust is to initially operate in the space of more standard tokens, such as A/B/C/D, before aligning probability to the correct symbols at a late layer".”

This is referring to our finding in lines 76-78 and in Section 5 lines 407-419 under the heading “Some answer symbols are produced in a two-stage process”: that for the Olmo 7B Instruct model, hidden states initially assign high logit values to expected answer symbols (A/B/C/D) before switching to the symbols given in the prompt. We'll try to adjust the phrasing to make this clearer.

"Fig.9: Do you have such curve for any other model, aside of Olmo 0724 7B base?”

No, and unfortunately we are unable to produce one. The Llama and Qwen developers did not release any intermediate model checkpoints, and the Olmo 1B model doesn’t do well enough on the synthetic task at the end of training to warrant constructing one (see Fig. 2).

"Fig.4: I don't understand the purpose of making this figure. What should it show to us?”

The purpose of the figure is to show that trends in promoting answer choices in vocabulary space are largely similar across tasks, despite the fact that the difficulty and subject matter of the tasks vary substantially and model performance differs as a result (i.e., from left to right: 55%, 51%, and 100% accuracy). We elaborate on this in lines 392-394, but if additional clarification or discussion would be valuable, we are happy to add. We’ve also added the accuracies to the figure caption.

"Fig.8: Can you, please, elaborate the captions a), b), c)?”

Plots b) and c): this is the difference in logits on the answer choice tokens produced by a hidden state projected to the vocabulary space (described in Section 4.2, specifically lines 254-255). This is equivalent to the “logit difference” lines in Figures 5a and 5b, except now at the more fine-grained level of individual attention heads instead of hidden state outputs. We plot this to show the extent to which attention heads specialize by letter (by comparing how the two plots differ). Additionally, while the model is ultimately producing positive values (blue) for the instances in both graphs that result in correct predictions, the plots show that not all attention heads contribute positively towards the model predicting the correct answers.

Plot a): this is the sum of logit values for all answer choices (logit of A + logit of B + logit of C + logit of D) using the same vocabulary projection method. This is equivalent to summing the A, B, C, and D lines in Figure 5a, except now at the more fine-grained level of individual attention heads instead of hidden state outputs. We plot this to show the overlap between plots a), b) and c): many attention heads appear to play a role in both generating any valid answer choice symbol and outputting the correct symbol (lines 478-479).

We will include additional context for these plots in the paper.

We have updated the PDF with the requested changes (see general response for the details).

Thank you for your positive comments about our work and valuable questions!

"The motivation for selecting the patching format is not fully substantiated. By "patching format," I refer to the reordering of the correct answer to a new position. Although this approach is intuitively clear, I believe the motivation could be strengthened by exploring different patching formats (for example, removing correct answer from sample).”

As we mention in lines 231-233, preliminary conditions for obtaining meaningful interpretations from patching is that 1) the model predicts the correct answer for both patching formats, and 2) the answer changes after a patch. Removing the correct answer from the sample does not result in a valid new correct answer, making it difficult to interpret the role of various model components since the task has fundamentally changed. The two methods we use (reordering the correct answer to a new position and changing the answer choice symbols) ensure the correct answer has changed but is still valid, allowing us to isolate mechanisms for predicting a specific correct answer.

"The circuits are shown to be dependent on the nodes (parts of model) for patching [1]. Patching the layers is interesting, but for full picture it would be good to examine the attention maps by patching.”

Good suggestion; we will include a plot for attention head patching in the updated PDF (coming soon). (We assume you mean "attention head" by "attention map" here; but please let us know if this is not what you mean).

Your questions:

"Could you please specify what values are presented by Fig. 8? Is taken from logit lens applied to MSHA components by heads?”

Yes, this is exactly what is plotted, following Yu et al (2023)’s methodology. We describe the decomposition of the MHSA component into attention heads in Appendix A.2 (lines 835-846) and briefly in lines 456-457, but we’ll add more details in the Figure 8 caption so it’s clearer.

"What conclusion do you make from the fact that in Fig. 7c (MHSA) change in answers are visible only in 24 layer, while in Fig.3 it maintains further. It will be interesting to describe in detail that mechanism.”

This provides supporting evidence for the two-stage prediction process for the Q/Z/R/X prompt. When the correct answer symbol (A to D) and location (index 0 to index 3) change in Figure 7c, but not the vocabulary space of answer choices (A,B,C,D), layer 24 is strongly causally implicated in the model switching its prediction from A to D– this may be encoded in the residual stream at layer 24 as either positional information saying “the correct answer is at index 3” or as the actual letter, “D”, both of which flip the model’s prediction to “D”. But when only the answer symbol (A to Q) and the vocabulary space (Q,Z,R,X) change and not location (index 0 stays correct), we observe layers 27-29 playing a key role. This means that layer 24 is neither encoding that 1) Q represents position 0 and position 0 is correct nor that 2) Q represents the selected answer. Aligned with Figure 6a, Fig. 3 depicts that there is a delayed layerwise effect in assigning non-negligible probability to Q.

"It is relatively small issue, but do you consider using another variant of logit lens? The reason is that logit lens by itself could be potentially worse for earlier layers [2]”

This is the primary reason we also employ causal tracing, as it can discover key mechanisms in early layers that logit lens cannot. We do not use the tuned lens approach [2] because it trains a linear probe instead of using the vocabulary space defined by the model’s unembedding matrix, and thus cannot answer our research question of how predictions form in the model’s vocabulary space (lines 239-241).

Thank you for your patience. We completed the attention patching results on each attention head for the experimental setup plotted in Figure 7(c) from layer 22 onward (i.e., where we first observe an effect from the MHSA function). The results demonstrate that the promotion of specific answer choices (the spikes in Fig. 7c) are attributed to 1-4 attention heads per layer. To give an example, you can see the result for patching each of the 32 attention heads at layer 24 here, where heads 1 and 19 give a cumulative effect to produce Fig 7c's pattern: https://imgur.com/a/9EfbAUz

We have added the graphs (depicting layers 22 to 31) as an additional figure to the Appendix and referenced this figure in our section 6 paragraph titled "Answer symbol production is driven by a sparse portion of the network" (unfortunately we are no longer able to update the PDF here for your viewing, but hope the above image temporarily suffices as evidence that we have done this). These results are valuable additional evidence to complement our vocabulary projection plots in Fig. 8 and demonstrate causal evidence of the unique roles of individual attention heads; thank you for the suggestion.

Our current PDF version that we uploaded earlier this week includes an updated caption for Figure 8 with more details about the experiments as you requested-- please see our general response above for more information. We are happy to answer any more questions or suggestions you may have.

Thanks for your detailed engagement with our paper & consideration of our rebuttal.

For each patch, we perform 2 inference runs on GPU in parallel (i.e., on prompts $x_A$ and $x_B$) and then patch a hidden state from one run over to the other (i.e., $x_B\rightarrow x_A$) before completing the inference run on input $x_A$ to see how the score changes as a result of the patch. For layerwise analysis, this is, 32 paired inference runs over the dataset for a 32-layer model like Olmo 7B. In the naive implementation, patching each attention head adds a number-of-attention-heads multiplicative factor, so for Olmo 7B, you would now run the inference procedure over the dataset 32 heads * 32 layers number of times. This can be sped up slightly by caching the hidden states up to the layer where you are performing the intervention. There is also some nascent work on speeding patching up by using gradient-based approximation techniques (https://www.neelnanda.io/mechanistic-interpretability/attribution-patching), but we stuck to the purist approach here to avoid approximating anything.

Thanks for all of your positive feedback!

"The range of considered model sizes is limited (0.5B-7B)”

The only other sizes of Llama 3.1 available are 70B and 405B, and Olmo does not have any model sizes apart from those we tested (1B and 7B). We included Qwen 2.5 0.5B and 1.5B sizes in the paper to provide a size contrast to the 7B and 8B models; we feel these 3 model families and 5 sizes are fairly representative.

While we unfortunately did not have the resources to scale to super large (70B+) sizes, we note this is the case in virtually all current academic mechanistic interpretability research. However there are some promising up-and-coming initiatives (such as https://arxiv.org/abs/2407.14561) to allow for analysis at larger scales soon, so we hope we can scale in future work.

Your questions:

"why the threshold on 20%+ random (i.e. 45%) is chosen?”

Good question. 20% represented a reasonable trade-off to us between selecting high-performing models and having a sufficiently representative sample of models, but the exact threshold doesn't matter. We’ll reword this in the updated PDF (coming soon) to be that we selected the best-performing model from each model family for further analysis, and they all do all 3 tasks sufficiently above random (in this case, >20% above random).

"Evalusting 3-shot accuracy with alternative symbols (e.g. QXYZ), which symbols are used for in-context examples?”

We use the same alternative symbols for the in-context examples. See footnote 4 on line 214.

"Do you have any ideas which of the observed phenomena can be extended to other tasks, beyond MCQA?”

This is an open question! We don’t want to extrapolate too much, but there is some initial evidence that models share circuit components across tasks (such as https://arxiv.org/abs/2310.08744). This is an important direction for future work.

New results for Qwen 2.5 and LLaMA 3.1 are very interesting. I wonder, why LLaMA behaves differently, than OLMO and Qwen in the experiments with Q/Z/R/X symbols for answer choices. Maybe Qwen and OLMO were over-tuned on some MCQA tasks with A/B/C/D letters, and this is why they "want" to assign some scores to these "familiar" letters before going to Q/Z/R/X, and LLaMA is less "familiar" with such tasks and thus less inclined to A/B/C/D format? I don't know for sure, it's just a hypothesis.

Anyway, I increased my score from 5 to 6 due to new results that make paper more valuable.