Source-linked AI summary
Quantifying Logical Consistency in Transformers via Query-Key Alignment
Eduard Tulchinskii, Anastasia Voznyuk, Laida Kushnareva, Andrei Andriiainen, Irina Piontkovskaya, Evgeny Burnaev, Serguei Barannikov
TL;DR
LLMs still lack reliable multi-step logical reasoning, and chain-of-thought prompting does not assess transition coherence. This paper evaluates logical consistency through query-key alignment in transformer heads using a single forward pass, finding robust performance across reasoning depths and distractor settings, with transfer gains on one additional benchmark.
Problem
Chain-of-thought prompting improves intermediate-step generation but lacks a mechanism to assess the coherence of logical transitions.
Method
The paper uses QK-score interactions in transformer heads to evaluate logical transitions and identifies relevant heads in a single run.
Results
The selected heads consistently outperform the baseline across reasoning depths and distractor counts, while three of five selected heads exceed baseline accuracy by more than 10% on PARARULE Plus.
Takeaways & Limitations
QK-scores provide a tractable lens for identifying attention heads that act as reasoning checkpoints in complex logical relationships.
Takeaways & Limitations
Head selection was performed only on two deduction-rule sets, and the scope of heads across other logical principles may require further experiments.
Abstract
from arXiv · showhide
Large language models (LLMs) have demonstrated impressive performance in various natural language processing tasks, yet their ability to perform multi-step logical reasoning remains an open challenge. Although Chain-of-Thought prompting has improved logical reasoning by enabling models to generate intermediate steps, it lacks mechanisms to assess the coherence of these logical transitions. In this paper, we propose a novel, lightweight evaluation strategy for logical reasoning that uses query-key alignments inside transformer attention heads. By computing a single forward pass and extracting a "QK-score" from carefully chosen heads, our method reveals latent representations that reliably separate valid from invalid inferences, offering a scalable alternative to traditional ablation-based techniques. We also provide an empirical validation on multiple logical reasoning benchmarks, demonstrating improved robustness of our evaluation method against distractors and increased reasoning depth. The experiments were conducted on a diverse set of models, ranging from 1.5B to 70B parameters.
1 Introduction
The paper introduces QK-score evaluation to assess logical-transition correctness, addressing the lack of coherence checks in chain-of-thought reasoning. It identifies useful attention heads efficiently and evaluates the method across logical benchmarks, especially multi-step and distractor-rich settings.
- Problem: Chain-of-thought prompting generates intermediate reasoning steps but lacks a mechanism to assess their coherence.
- Core motivation and contribution: The method processes all attention heads in a single run, making head identification scalable to large models.
- Core motivation and contribution: QK-score interactions between selected query and key vectors assess the correctness of logical transitions.The method uses internal transformer-head representations as a proxy for logical consistency.
- Evaluation: Experiments on multiple logical benchmarks target improved consistency in multi-step reasoning and contexts containing reasoning distractors.
2 Related Work
Prior work improves logical reasoning with chain-of-thought prompting and studies transformer-head behavior, but component-level analyses can be computationally expensive. The paper’s approach is positioned against these limitations.
- Logical reasoning: Chain-of-thought methods generate intermediate reasoning steps but do not assess the coherence of logical transitions.
- Mechanistic interpretability: Mechanistic interpretability research has examined attention-head roles across knowledge recall, latent reasoning, and expression preparation.
- Component analysis: Some studies disable model components to analyze behavior, but these approaches are often computationally expensive or limited to simpler tasks.
3 Approach
The approach scores alignment between statement and candidate-answer representations in every transformer head, using contexts, statements, and true-or-false candidates. A single forward pass identifies heads that assess logical consistency without head-by-head ablations.
- Inputs: Each input contains a premise-providing context, a candidate conclusion, and a true-or-false candidate answer.
- Scoring: For every layer and head, the QK-score measures alignment between the statement representation and candidate-answer representation.
- Scoring: The candidate-answer token supplies the query vector, while the statement-ending token supplies the key vector.
- Head selection: Evaluating all heads in one forward pass identifies reliable logical-consistency assessors without extensive model modifications or head-by-head ablations.
4 Experiments
The experiments evaluate QK-score heads across three logical-reasoning benchmarks, using zero-shot frozen models and calibration-based head selection. The setup tests performance across reasoning depths, distractors, and cross-dataset generalization.
- Datasets: The evaluation covers ProntoQA-OOD, PARARULE Plus, and Extended-Multi-LogiEval logical-reasoning benchmarks.ProntoQA-OOD includes repeated Modus Ponens and six-rule deduction settings; PARARULE Plus uses fixed-depth true/false questions, while Extended-Multi-LogiEval expands and balances Multi-LogiEval.
- Datasets: ProntoQA-OOD uses 600 calibration examples and more than 1,000 evaluation examples per setup, with distractor counts varied for robustness testing.Calibration data contains 300 true and 300 false examples, partitioned by reasoning depth.
- Experimental setup: All experiments use zero-shot frozen pretrained LLMs ranging from 1.5B to 70B parameters.The models are evaluated without additional training or modification.
- Experimental setup: The Baseline selects the higher-probability option between a0 and a1 from the model’s output probabilities.For each ProntoQA-OOD setup, the best QK-score head is selected on calibration data and evaluated against this baseline on held-out data.
- Experimental setup: Five heads selected from top-performing ProntoQA-OOD setups are evaluated on PARARULE Plus and Extended-Multi-LogiEval to assess generalization.The heads are chosen to cover as many ProntoQA-OOD setups as possible.
5 Results
QK-score-selected heads consistently outperform the baseline across ProntoQA-OOD reasoning depths and distractor levels. Cross-dataset transfer is strongest on PARARULE Plus, while gains are less consistent on Extended-Multi-LogiEval.
- In-Domain Evaluation: The BEST HEAD consistently outperforms the baseline across ProntoQA-OOD reasoning depths and different numbers of distractors.Selected heads also maintain stable performance across varying numbers of reasoning steps.
- Transfer Learning Evaluation: Three of five selected heads exceed the baseline by more than 10% on PARARULE Plus.These heads were selected using ProntoQA-OOD experiments and then evaluated cross-dataset.
- Transfer Learning Evaluation: Only two selected heads outperform the baseline on Extended-Multi-LogiEval.The reported contrast may reflect differences between binary yes/no and true/false question formats.
- Overall findings: Overall, QK-score identifies heads that reliably assess logical transitions within ProntoQA-OOD and across cross-dataset evaluations.The reported robustness covers both in-domain and transfer settings.
6 Analysis
QK-scores provide a semantic-alignment view of logical structure that is distinct from raw attention weights. Selected heads can outperform final model probabilities and remain stable under distractors.
- QK-score analysis: QK-scores are independent of positional encodings and focus on semantic alignment between statements and candidate answers.This distinguishes them from raw attention weights in the paper’s analysis.
- QK-score analysis: Heads selected by QK-scores often outperform the model’s final probabilities, suggesting reasoning signals can be obscured by later processing.The paper characterizes these heads as verification anchors when distractors are present.
- QK-score analysis: High-QK-scoring heads consistently identify valid inferences despite distractors in the input.Their stability motivates viewing these heads as verification anchors for logical consistency.
7 Conclusion
The paper introduces QK-score as a single-pass framework for identifying attention heads that capture logical validity. Experiments indicate that selected heads can outperform final-layer predictions, generalize across datasets, and provide a tractable interpretability signal.
- Conclusion: The QK-score framework identifies attention heads that consistently capture logical validity.The method is presented as a lightweight alternative to ablation-based analysis.
- Conclusion: In multi-step inferences with distractors, certain heads outperform the model’s final-layer predictions.The conclusion links these heads to reasoning checkpoints within the model.
- Conclusion: The single-pass procedure avoids the computational overhead of ablation-based methods and generalizes relatively well across datasets.The paper presents this as a tractable window into complex logical relationships.
8 Limitations
The method requires substantial, balanced, debiased calibration data and has only been tested across limited rule sets and model types. The authors also caution that selected heads are not necessarily the only heads responsible for logical inference.
- Calibration data: At least ≈400 reasoning questions are required for calibration, with balanced logical-rule coverage and debiasing against answer-form shortcuts.These safeguards aim to prevent selecting heads that guess answers from question form.
- Interpretation: The identified heads are not claimed to be the only heads responsible for logical inference.
- Rule coverage: Head selection was tested on only two deduction-rule sets, so attention-head scope across other logical principles remains uncertain.The authors state that broader experiments may be needed.
- Model scope: Experiments considered only Instruct- and Chat-tuned models, leaving changes from base to fine-tuned models for future work.
C Pararule-plus: additional details
The PARARULE-PLUS section illustrates fixed-depth reasoning prompts, while the cross-dataset evaluation reports that QK-score-selected heads often outperform the baseline. The extended Multi-LogiEval data add samples without changing the original reasoning schemes.
- PARARULE-PLUS: Three of five heads selected on ProntoQA-OOD exceed the PARARULE-PLUS baseline, usually by more than 10%.Head (22, 16) from DeepSeek-R1 instead achieves below 0.35 accuracy, indicating reversed QK-score discrimination in that setup.
E Extended results of in-domain evaluation on ProntoQA-OOD
The in-domain evaluation compares BEST HEAD and BASELINE methods across models and reasoning depths on ProntoQA-OOD, including settings with distractors. The listed inference rules define the logical schemes used in the extended benchmark.
- Evaluation results: Tables 5 and 6 report in-domain numerical results on ProntoQA-OOD for BEST HEAD and BASELINE methods across various LLMs.
- Inference rules: The extended benchmark defines inference schemes including Modus Ponens, Modus Tollens, Hypothetical Syllogism, and several dilemma, composition, and quantifier rules.
- Evaluation results: Table 5 compares model performance across different reasoning depths, with the best results highlighted.
- Evaluation results: Table 6 examines how added distractors affect ProntoQA-OOD performance under Modus Ponens inference.