Last Layer Logits to Logic: Empowering LLMs with Logic-Consistent Structured Knowledge Reasoning

AI summary

Overview

Research area: Natural Language Processing, specifically knowledge graph question answering (KGQA) and the decoding behaviour of large language models (LLMs) when reasoning over structured knowledge.

Technical level: Intermediate. The paper assumes familiarity with LLM decoding, token logits, and knowledge graph triplets, and introduces a finite-automaton formulation of the reasoning process.

Scope: The paper diagnoses a failure mode it calls "Logic Drift" in LLM structured knowledge reasoning and proposes a decoding-time framework, Logits-to-Logic, that edits last-layer logits to align LLM outputs with the logic of the question and the knowledge graph.

What This Paper Is About

LLMs are pretrained on unstructured text and therefore struggle to stay logically consistent when reasoning over structured knowledge such as knowledge graphs, producing reasoning paths that either do not exist in the graph or are semantically unrelated to the question. The authors call this "Logic Drift" and argue that existing agentic and prompt-based methods only guide the model at the input level, leaving the drift unresolved in the output. Their goal is to correct it at the source, by modifying the logits of the last layer during autoregressive generation.

Key Contributions

1. The paper identifies the core challenge of LLMs in structured knowledge reasoning as an inconsistency between model outputs and the logical distributions of the question and the knowledge graph, and proposes addressing this from the output perspective rather than the input perspective.
2. It introduces the Logits-to-Logic framework, whose two core modules, logits strengthening and logits filtering, are designed to correct logical defects during decoding.
3. It unifies the LLM's autoregressive generation and the KG's structure within a state-transition Non-deterministic Finite Automaton (NFA), using the LLM vocabulary as the automaton's input alphabet.
4. It reports state-of-the-art results on multiple KGQA benchmarks and shows the method transfers directly to different knowledge graphs and tasks without modification.

Main Findings

• Logic Drift is measurable and prevalent: Figure 1 reports Logic Drift ratio statistics for ToG, DoG, GCR, and KG-CoT on 100 sampled instances from the CWQ dataset.
• Main KGQA results (Hit@1): Logits-to-Logic with LLaMA3.1-8b reaches 95.4 on WebQSP, 80.8 on CWQ, 82.0 on GrailQA, and 80.3 on SQ. The paper states this surpasses RoG and GoG by 9.7% and 11% on WebQSP, KG-Agent, SymAgent and GCR by 8.6%, 22% and 5% on CWQ, DoG, PoG and DARA by 2%, 5.5% and 5% on GrailQA, and KG-CoT and ToG by 2.5% and 13.6% on SQ.
• Advantage at smaller model scale: The paper notes that these baselines use large models such as ChatGPT and GPT4, while its method uses the smaller LLaMA-3.1-8b model.
• Ablation of the two modules (Hit@1 / F1): Removing logits strengthening gives 93.9 / 63.1 on WebQSP and 74.0 / 44.7 on CWQ; removing logits filtering gives 86.2 / 60.0 and 79.5 / 44.1; removing both gives 57.6 / 56.8 and 59.7 / 38.0. The full method gives 95.4 / 63.3 and 80.8 / 45.0. The paper reports the removal of logits strengthening decreases performance by 1.5% and 6.9% on WebQSP and CWQ, and removing logits filtering causes drops of 9.2% and 1.4%.
• Strengthening coefficient: Performance is optimal at ω = 2.0 across datasets and metrics. A negative value (ω = -1.0) causes outputs to deviate from question logic, and ω > 3.0 degrades performance because excessive amplification disrupts language logic.
• Transfer results (Hit@1): Without modification, the method reaches 59.2 on QALD10-en, 87.5 on T-REx, and 91.3 on Zero-shot RE, covering multi-hop QA, single-hop QA, and slot filling over Wikidata.
• Backbone sensitivity: Across five model families, LLaMA-2-13b achieves the highest F1 (75.2 on WebQSP, 64.1 on CWQ), while Qwen2.5-14b achieves the highest WebQSP Hit@1 at 96.6. Smaller backbones also perform well: Qwen-2-0.5b reaches 81.6 Hit@1 on WebQSP and 75.2 on CWQ.
• Error analysis: Logits-to-Logic achieves the lowest total error count of 14, of which Question-Inconsistent Logic Drift accounts for only 10, while the compared methods (ToG, DoG, KG-CoT, GCR) exceed 27.
• Visualisation: Logits distribution plots show that without strengthening correct tokens rank lower, without filtering KG-inconsistent tokens hold high logits, and removing both produces obvious logic drift.

Methodology in Plain English

The authors treat generation as a state machine. First, in a stage called Logic Compiling, they take the topic entity of a question and compile all legal reasoning paths in the knowledge graph into an NFA, where each intermediate token sequence of a legal path is a state and the transition function specifies which next tokens are acceptable. A sentence-transformer is used to score paths for similarity to the question, and high-scoring paths are retained, so the automaton reflects both graph structure and question semantics.

Second, in Logits Strengthening, the model is run twice: once with the normal prompt and once with the answer paths replaced by MASK tokens. The difference between the two output distributions is computed and scaled by a coefficient ω, then added back to the masked distribution. This boosts the logits of tokens that belong to question-aligned reasoning paths, following the intuition of contrastive decoding.

Third, in Logits Filtering, the NFA transition function is used to suppress tokens that are not part of any legal path, effectively setting their contribution to negative infinity so they cannot be sampled. Together the two modules push the model's output distribution towards the logical distributions of the question and the graph. Experiments use LLaMA-3.1-8b with beam search (beam size 20), ω = 2.0, and a 22M-parameter sentence-transformer-all-MiniLM-L6-v2 for path scoring, with supervised fine-tuning on 1/10 of the randomly sampled CWQ and WebQSP training data to teach the output path format.

Why This Matters

Impact on research: The paper reframes Logic Drift as a mismatch between output distributions and the logical distributions of structured knowledge, moving the intervention point from prompt design to decoding. It offers an alternative to increasingly complex multi-agent pipelines and provides an NFA-based formalism that connects LLM token generation to graph structure.

Real-world applications:

• Question answering over enterprise or public knowledge graphs, such as Freebase and Wikidata, where answers must follow verifiable paths.
• Fact verification and compliance checking, where an answer must be traceable to a specific chain of relations rather than a generated claim.
• Slot filling and structured data extraction, which the paper demonstrates on T-REx and Zero-shot RE.
• Multi-hop query interfaces that need precise, explainable reasoning chains rather than fluent but ungrounded text.

Industry relevance: Because the method works on top of open-source models from 0.5b to 14b parameters and the paper reports strong results with a 7-8b model, it is relevant to deployments where large closed models are impractical. The authors are affiliated with Zhejiang University, the ZJU-Ant Group Joint Lab of Knowledge Graph, and Ant Group, indicating industrial interest in structured knowledge systems.

Future Directions

• Reducing the number of incorrect candidate paths: the stated limitation is that the search space for correct reasoning paths is excessively large, and beam search still introduces a certain number of incorrect paths.
• Improving path scoring, since the current pipeline depends on a 22M-parameter sentence-transformer to rank NFA paths against the question.
• Extending the automaton formulation beyond the 2-hop subgraphs explored by breadth-first search from the topic entity in this work.
• Testing whether the decoding-time correction can be combined with, rather than replace, existing agentic frameworks, given that several agentic baselines were evaluated with larger models than the proposed method.

Target Audience

Researchers and practitioners working on knowledge graph question answering, LLM reasoning, or constrained decoding. It is also relevant to engineers building grounded question-answering systems over enterprise knowledge graphs, and to readers interested in alternatives to prompt-based and multi-agent approaches for improving LLM logical consistency.

Authors’ abstract

Large Language Models (LLMs) achieve excellent performance in natural language reasoning tasks through pre-training on vast unstructured text, enabling them to understand the logic in natural language and generate logic-consistent responses. However, the representational differences between unstructured and structured knowledge make LLMs inherently struggle to maintain logic consistency, leading to \textit{Logic Drift} challenges in structured knowledge reasoning tasks such as Knowledge Graph Question Answering (KGQA). Existing methods address this limitation by designing complex workflows embedded in prompts to guide LLM reasoning. Nevertheless, these approaches only provide input-level guidance and fail to fundamentally address the \textit{Logic Drift} in LLM outputs. Additionally, their inflexible reasoning workflows cannot adapt to different tasks and knowledge graphs. To enhance LLMs' logic consistency in structured knowledge reasoning, we specifically target the logits output from the autoregressive generation process. We propose the \textit{Logits-to-Logic} framework, which incorporates logits strengthening and logits filtering as core modules to correct logical defects in LLM outputs. Extensive experiments show that our approach significantly improves LLMs' logic consistency in structured knowledge reasoning and achieves state-of-the-art performance on multiple KGQA benchmarks.