💬 NLP

Attention-Guided Layer Selection for Contrastive Decoding in Large Language Models

This paper proposes three attention-guided layer selection strategies (Attention-JSD, Attention-Entropy-Max, and Attention-Entropy-Min) that leverage internal self-attention mechanisms to improve the factuality of Large Language Models, demonstrating consistent performance gains over the original DoLa method on TruthfulQA benchmarks.

Original authors: Yusuke Sakai, Natthawut Kertkeidkachorn, Kiyoaki Shirai

Published 2026-07-28

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yusuke Sakai, Natthawut Kertkeidkachorn, Kiyoaki Shirai

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are talking to a super-smart robot that has read almost every book on the internet. It can write poems, solve math problems, and tell jokes. But sometimes, this robot gets a little too confident and starts making things up. It might tell you that the moon is made of green cheese or that a famous historical figure invented the smartphone. In the world of artificial intelligence, this is called "hallucination." It's a big problem because if we want to use these robots for real-world tasks like answering medical questions or writing news, we need them to be truthful, not just creative.

To fix this, scientists have been trying to teach these robots how to check their own work while they are thinking. One popular method is called "Contrastive Decoding." Think of it like a robot having an internal debate. The robot has many layers of "thinking" inside it, kind of like a stack of pancakes. Some layers are "immature" (they think fast and shallowly), and some are "mature" (they think deeply and carefully). The old method of checking their work was to compare what the shallow layers said against what the deep layers said. If the shallow layer was too confident about a wrong answer, the robot would ignore it. But this paper asks a new question: Is looking at the words the robot is considering the only way to know if it's telling the truth? Or can we look at how the robot is paying attention to its own thoughts to find the truth?

This paper, titled "Attention-Guided Layer Selection for Contrastive Decoding in Large Language Models," proposes a clever new way to help these robots tell the truth. The authors, from the Japan Advanced Institute of Science and Technology, suggest that instead of just comparing the final words the robot wants to say, we should look at the robot's "attention." Imagine attention as a spotlight the robot shines on different parts of its memory. The researchers found that by watching where this spotlight moves, they could pick the best "thinking layer" to use for checking facts.

They tested three new strategies:

Attention-JSD: This looks for the layer where the robot's "spotlight" is most different from its final, mature thought. It's like finding the layer that is looking at the problem from a totally different angle.
Attention-Entropy-Max: This picks the layer where the spotlight is spread out very widely, looking at many things at once.
Attention-Entropy-Min: This picks the layer where the spotlight is super focused, zooming in tightly on just a few specific pieces of evidence.

When they tested these ideas on a famous quiz called TruthfulQA, which is designed to trick robots with common myths and misconceptions, the results were exciting. The new methods, especially the ones that looked at how focused or spread out the attention was, did a better job than the old method. They were particularly good at finding multiple correct answers, not just one. For example, on the LLaMA-7B model, the "Attention-Entropy-Min" strategy (the super-focused spotlight) got a score of 62.8% on the multiple-answer test, beating the previous best of 56.5%.

The authors also peered inside the robot's brain to see why this worked. They discovered that not all parts of the robot's attention are created equal. Some specific "attention heads" (tiny sub-processors inside the robot) seem to carry a stronger signal for truthfulness than others. In fact, they found that using just one specific attention head (Head 6) could sometimes make the robot even smarter than using all of them together. This suggests that the robot's internal "spotlight" patterns are a very sensitive signal for fact-checking, perhaps even more so than just looking at the list of words it might say next.

However, the paper is careful to note that this isn't a magic wand that solves everything. The researchers suggest that while these methods work well for multiple-choice questions, we still need to see how they handle free-form writing, where the robot has to generate long stories or essays from scratch. They also point out that they still need to do some extra testing to decide exactly which layers to use for different robot sizes, and they don't fully understand why certain attention heads are so good at spotting facts yet. But for now, this work suggests that if we want our AI to be more honest, we should stop just listening to what it says and start watching how it pays attention.

Technical Summary: Attention-Guided Layer Selection for Contrastive Decoding in Large Language Models

Problem Statement

Large Language Models (LLMs) are prone to hallucination, generating plausible but factually incorrect content. While contrastive decoding methods like DoLa (Decoding by Contrasting Layers) have improved factuality by contrasting the output distributions of "mature" (final) and "premature" (early) layers within a single model, they rely exclusively on divergences in output vocabulary distributions (logits). This approach ignores the structural information embedded in the model's internal self-attention mechanisms, which may provide more sensitive signals for identifying layers where factual knowledge emerges.

Methodology

The authors propose three attention-guided strategies for selecting the premature layer in the DoLa framework. Instead of comparing output logits, these methods utilize the self-attention distributions ( $A_l$ ) of intermediate layers to determine which layer to contrast with the final layer ( $L$ ).

For a given token position $t$ , the methods compute a head-averaged attention row $a^l_t$ representing the attention distribution of layer $l$ over past tokens. The three strategies are:

Attention-JSD: Selects the layer $l^*$ that maximizes the Jensen-Shannon Divergence (JSD) between its attention distribution and the final layer's distribution.
$l^* = \arg \max_{l \in C} \text{JSD}(a^L_t, a^l_t)$
Attention-Entropy-Max: Selects the layer with the maximum attention entropy, hypothesizing that broadly distributed attention serves as a cue for layer selection.
$l^* = \arg \max_{l \in C} H(a^l_t)$
Attention-Entropy-Min: Selects the layer with the minimum attention entropy, hypothesizing that concentrated attention reflects strong reliance on specific evidence tokens.
$l^* = \arg \min_{l \in C} H(a^l_t)$

These strategies preserve the DoLa framework, including the plausibility constraint (restricting candidates to high-probability tokens under the final layer) and the dynamic selection process at each token.

Key Contributions

Novel Selection Signals: The paper introduces the first systematic use of structural information from self-attention mechanisms (divergence and entropy) as dynamic signals for token-level layer selection in contrastive decoding, moving beyond reliance on output vocabulary distributions.
Empirical Validation: Extensive experiments on TruthfulQA and FACTOR benchmarks across six models (LLaMA-7B/13B/33B/65B, Gemma-7B, Mistral-7B) demonstrate that attention-guided strategies consistently outperform the original DoLa, particularly on metrics evaluating multiple correct answers (MC2 and MC3).
Mechanistic Insights: Through head-wise analysis, the authors identify that specific attention heads carry distinct signals conducive to factuality discrimination. They show that factuality improvements are not uniform across all heads; for instance, Head 6 in LLaMA-7B yielded significantly higher MC3 scores than the average of all heads.

Experimental Results

TruthfulQA: The proposed methods consistently outperformed the baseline and DoLa. Attention-Entropy-Min achieved the highest MC2 (62.8%) and MC3 (36.5%) scores for LLaMA-7B, significantly surpassing DoLa. Attention-JSD yielded the best results for LLaMA-13B. While DoLa maintained the highest MC1 (single correct answer) for larger models (33B/65B), the attention-based methods outperformed DoLa in MC2 and MC3, suggesting attention distributions are more sensitive for resolving factual knowledge in multi-answer scenarios.
FACTOR: On long-form text factuality, the methods showed performance comparable to or better than DoLa. For example, Attention-Entropy-Max achieved 62.6% accuracy on the LLaMA-7B Wiki subset (vs. DoLa's 62.1%), and Attention-JSD reached 64.1% on the Gemma-7B Wiki subset (vs. DoLa's 63.3%).
Layer Selection Patterns: Analysis revealed that while DoLa and Attention-JSD tend to select the shallowest available layers, Attention-Entropy-Min adaptively selects layers across a wider, deeper range (e.g., layers 22–30 in LLaMA-7B), suggesting it responds more dynamically to input context.
Score Distribution: DoLa tends to produce "sharp" distributions with high confidence in a single correct answer (beneficial for MC1), whereas attention-based methods produce "flatter" distributions that spread scores across multiple correct options (beneficial for MC2/MC3).

Significance and Claims

The paper claims that internal attention distributions provide a more sensitive signal for resolving factual knowledge than output vocabulary distributions, particularly for tasks requiring the identification of multiple valid answers. By leveraging structural information from self-attention, the proposed methods enhance factuality without requiring additional training or external models.

The authors emphasize that their approach is a practical, inference-time plug-in that maintains computational overhead within a reasonable range. They conclude that while the current methods require prior validation to select candidate layer buckets, the discovery that specific attention heads (rather than just the average) drive factuality suggests a pathway for further optimization through context-dependent head selection.

Limitations

The authors acknowledge that their evaluation is currently limited to multiple-choice benchmarks, and the applicability to free-form generation remains to be verified. Additionally, the selection of candidate premature-layer buckets still requires validation, and the theoretical basis for why specific attention heads excel at factual recall is not yet fully elucidated.