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

2026-07-25

AI Summary

This work addresses the limitation of existing contrastive decoding methods, such as DoLa, which rely solely on output-layer token distribution discrepancies and struggle to effectively capture factual knowledge. The study introduces, for the first time, structural information from the self-attention mechanisms of large language models and proposes three attention-distribution-based dynamic layer selection strategies: Attention-JSD, Attention-Entropy-Max, and Attention-Entropy-Min. These strategies leverage Jensen-Shannon divergence and entropy extrema to refine layer selection during decoding. Experimental results demonstrate that the proposed approach significantly outperforms DoLa on the TruthfulQA benchmark, with particularly notable improvements on the MC2 and MC3 multi-choice metrics, thereby enhancing the factual consistency of generated text.

Abstract

Contrastive decoding methods such as DoLa improve the factuality of Large Language Models (LLMs) by contrasting the output distributions of mature and premature layers. However, DoLa's dynamic layer selection relies solely on divergences in output vocabulary distributions. In this work, we propose three attention-guided strategies: Attention-JSD, Attention-Entropy-Max, and Attention-Entropy-Min, which leverage structural information carried by internal self-attention mechanisms as a signal for layer selection. Experimental results on TruthfulQA demonstrate that our strategies, particularly Attention-JSD and Attention-Entropy-Min, consistently outperform the original DoLa. We observe significant gains on multi-answer metrics (MC2 and MC3), suggesting that attention distributions can provide a more sensitive signal for resolving factual knowledge than output vocabulary distributions.

Problem

• contrastive decoding
• layer selection
• attention mechanism
• factuality
• large language models

Innovation

• attention-guided layer selection
• contrastive decoding
• DoLa
• self-attention mechanism
• factuality in LLMs