Research map / Large language models / Language, safety and interpretability

Listening to the Wise Few: Query–Key Alignment Unlocks Latent Correct Answers in Large Language Models

NeurIPS 2026 · Large language models › Language, safety and interpretability · approach: LLM pruning and scaling data · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Shows that a Query-Key score from select-and-copy attention heads in middle layers exposes latent correct MCQA answers that LLMs fail to output, and proposes an unsupervised HeadScore for choosing such heads.

Problem

LLMs often fail to output the correct option in multiple-choice QA even though the answer is encoded internally; output-based evaluation is also sensitive to formatting and positional bias.

Approach

Score each option by the inner product between the last-token query and the end-of-line key after option i, computed before RoPE, using select-and-copy heads in middle layers; pick the argmax as the answer. An unsupervised HeadScore (attention mass and variance on option tokens) ranks heads without labeled validation data.

Key results

• On HellaSwag, a single head's QK-score exceeds model zero-shot accuracy by up to +27.4 pp (abstract); on HaluDialogue by up to +49.8 pp.
• LLaMA-3.1-8B zero-shot MMLU accuracy: QK-score 62.5 vs baseline 59.1 (permutation accuracy 50.7 vs 48.6).
• Zero-ablating select-and-copy heads collapses MCQA accuracy to near-random; QK-score on SSD achieves near-perfect accuracy across tested models and languages.

Surprising findings

• On MMLU and CosmosQA, PriDe debiasing slightly outperforms QK-score in the LLaMA-3.1-8B comparison.
• Label tokens are almost useless in zero-shot but perform well in five-shot for certain heads.

Limitations

• Requires access to attention matrices, so not applicable to closed models.
• Not applicable to scarce-resource tasks per authors.
• Head selection on synthetic data transfers poorly to longer-option tasks like HellaSwag and HaluDialogue.

Experimental setup

Datasets: MMLU CosmosQA HellaSwag HaluDialogue (HaluEval) Simple Synthetic Dataset (SSD)

Benchmarks: MMLU CosmosQA HellaSwag HaluDialogue Simple Synthetic Dataset

Metrics: accuracy permutation accuracy

Baselines: Baseline (output probabilities) PriDe Attention-score Synthetic head selection

Builds on: induction heads PriDe RoPE

Closest papers by meaning

• Answer, Assemble, Ace: Understanding How LMs Answer Multiple Choice Questions (ICLR 2025 · similarity 0.89)
• Rethinking Generative Large Language Model Evaluation for Semantic Comprehension (ICML 2024 · similarity 0.89)
• Latent Self-Consistency for Reliable Majority-Set Selection in Short- and Long-Answer Reasoning (AAAI 2026 · similarity 0.89)
• Sample Smart, Not Hard: Correctness-First Decoding for Better Reasoning in LLMs (ICLR 2026 · similarity 0.88)
• A Semantic-Sampling Framework for Evaluating Calibration in Open-Ended Question Answering (NeurIPS 2026 · similarity 0.88)
• Representation Consistency for Accurate and Coherent LLM Answer Aggregation (NeurIPS 2025 · similarity 0.88)