Think, Align, Select: Query–Key Scores for LLM Reasoning

Mark Obozov, Eduard Tulchinskii, Kristian Kuznetsov, Michael Diskin, Serguei Barannikov

Keywords: Query–Key alignment, attention heads, white-box selection, white-box verification, chain-of-thought (CoT), self-consistency, permutation robustness

Abstract: We demonstrate that a "think-first" phase via chain-of-thought (CoT) prompting systematically strengthens internal query–key (QK) alignment, improving ability to select and verify answers directly from model activations rather than decoded tokens. Building on multiple-choice evaluation with MMLU-Pro and extending to free-form reasoning on MATH-500, GSM8K, and our variant of Humanity's Last Exam, we evaluate three settings: (i) MCQA vs MCQA+CoT with QK-based selection; (ii) candidate generation with/without CoT followed by QK-based selection among self-proposed answers; and (iii) QK-based verification of LLM solutions. We analyze QK-score accuracy, permutation robustness, and diagnostics relating alignment strength to correctness. This yields a white-box, computation-efficient decision rule that turns CoT from a purely generative aid into a deliberation-then-selection mechanism grounded in the model's own representations. By leveraging this internal signal, we surpass preference-optimized LLMs on fundamental reasoning tasks, achieving performance gains up to 22\% across various benchmarks and models.

Submission Number: 241