Quantifying Logical Consistency in Transformers via Query-Key Alignment

Large language models have become remarkably adept at mimicking human conversation, writing poetry, and summarizing complex texts. They achieve this by predicting the next word in a sentence based on vast amounts of training data. However, when asked to solve a multi-step logic puzzle, these systems often stumble. They might generate a convincing-sounding chain of reasoning that leads to a completely wrong conclusion, or they might get distracted by irrelevant details in the prompt. While researchers have found ways to coax better performance out of these models by asking them to "think step-by-step," a fundamental problem remains: the models themselves do not have a built-in way to check if their own reasoning is actually coherent. They can produce a logical-sounding path without truly understanding if the steps connect correctly.

This uncertainty creates a blind spot in artificial intelligence. We can see the final answer a model gives, but we cannot easily see the internal signals that indicate whether the model is actually following the rules of logic or just guessing. To address this, a team of researchers from the Skolkovo Institute of Science and Technology and other institutions has developed a new way to peek inside the model's "mind" during the reasoning process. Instead of waiting for the model to finish its answer and then checking if it is right, they look at the specific internal connections the model makes while it is still thinking. They found that certain parts of the model's architecture act as reliable checkpoints, signaling whether a logical step is valid long before the final answer is spoken.

The researchers focused on how these models process information internally. Inside a large language model, information flows through layers of mathematical operations. At the heart of this process are "attention heads," which act like tiny switches that decide which parts of the input are important to focus on at any given moment. The team discovered that they could measure the alignment between two specific signals within these heads: the signal representing the statement being tested and the signal representing the potential answer. By calculating a simple score based on how well these two signals line up, they could determine if the model recognized the logical connection between a premise and a conclusion.

To test this idea, the researchers ran experiments on a wide variety of logical reasoning tasks. They used synthetic datasets where the rules were clear and the answers were definitively true or false, such as determining if a character named Polly is opaque based on a set of fictional rules about rompus and jompus. They also tested the method on more complex datasets involving real-world language and multiple steps of deduction. In these tests, they compared the performance of their new internal scoring method against the model's standard way of answering, which is simply to pick the option with the highest probability.

The results showed that the internal scoring method was significantly more reliable than the model's own final guess. In scenarios where the model was presented with distracting, irrelevant information, the standard model often got confused and gave the wrong answer. However, the specific internal signals the researchers were tracking remained steady, correctly identifying the valid logical path even when the model's final output was wrong. This suggests that the model actually possesses the correct reasoning capability deep within its layers, but that capability can sometimes be obscured by later stages of processing or by the noise of the distracting context.

The team tested this approach on models ranging from very small ones with 1.5 billion parameters to massive ones with 70 billion parameters. Across the board, they found that specific attention heads consistently acted as "verification anchors." These heads could distinguish between valid and invalid inferences with high accuracy, regardless of how many steps the reasoning required or how many distractions were present. In some cases, the internal score was correct over 90 percent of the time, while the model's final answer was correct less than 60 percent of the time. This indicates that the model is often "knowing" the right answer internally but failing to express it correctly in its final output.

One of the most significant findings was that this method does not require changing the model or retraining it. The researchers could simply run the model once, look at the internal signals, and derive a score that was often more accurate than the model's own prediction. This offers a lightweight and efficient way to evaluate logical consistency without the heavy computational cost of previous methods, which often involved disabling parts of the model to see what happened. The study suggests that by identifying these specific internal components, we can gain a clearer window into how artificial intelligence processes complex relationships, moving us closer to systems that are not just fluent in language, but also reliable in their reasoning.

The researchers noted that their method works best when there is a sufficient amount of data to calibrate which specific internal signals to trust for a given task. They also acknowledged that this approach does not mean these specific parts are the only ones responsible for logic, nor does it solve every problem with reasoning in artificial intelligence. However, the work provides a concrete, measurable way to see where the logic holds up and where it breaks down, offering a new tool for building more transparent and trustworthy language models. By focusing on the alignment of internal signals rather than just the final output, the study reveals that the capacity for sound reasoning often exists within the model, waiting to be properly read.