Attribution-Guided Decoding

TL;DR

Attribution-Guided Decoding (AGD) treats decoding as a search for the next token that is most justified by a Region of Interest (ROI) defined over inputs or internal components, using post-hoc attributions (notably Layer-wise Relevance Propagation) to score and select among a small candidate set. The method is decoding-time, fine-tuning-free, and can target instruction adherence or factual grounding via task-specific ROI definitions, including guidance toward internal knowledge heads or in-context evidence. Empirical results across instruction following and factuality tasks show that AGD improves adherence and factual recall, with an entropy-gating variant mitigating quality loss and reducing computational cost. The approach also provides interpretability by illustrating why certain tokens receive higher attribution signals. Limitations include reliance on potential tokens proposed by the model and the added backward passes, with future work aimed at more efficient proxies and circuit-level ROI definitions.

Abstract

The capacity of Large Language Models (LLMs) to follow complex instructions and generate factually accurate text is critical for their real-world application. However, standard decoding methods often fail to robustly satisfy these requirements, while existing control techniques frequently degrade general output quality. In this work, we introduce Attribution-Guided Decoding (AGD), an interpretability-based decoding strategy. Instead of directly manipulating model activations, AGD considers a set of high-probability output token candidates and selects the one that exhibits the highest attribution to a user-defined Region of Interest (ROI). This ROI can be flexibly defined over different parts of the model's input or internal components, allowing AGD to steer generation towards various desirable behaviors. We demonstrate AGD's efficacy across three challenging domains. For instruction following, we show that AGD significantly boosts adherence (e.g., improving the overall success rate on Llama 3.1 from 66.0% to 79.1%). For knowledge-intensive tasks, we show that guiding generation towards usage of internal knowledge components or contextual sources can reduce hallucinations and improve factual accuracy in both closed-book and open-book settings. Furthermore, we propose an adaptive, entropy-based variant of AGD that mitigates quality degradation and reduces computational overhead by applying guidance only when the model is uncertain. Our work presents a versatile, more interpretable, and effective method for enhancing the reliability of modern LLMs.

Figures (10)

• Figure 1: Illustration of the Attribution-Guided Decoding (AGD) framework.. (a) A Region of Interest (ROI) is defined over a relevant area, such as the user's instruction. (b) Next, a standard forward pass generates a candidate set of high-probability tokens, both "festival" (most probable) and "royal" are included. (c) An attribution method computes a relevance score for each candidate on the input tokens, quantifying a candidate's dependence on the ROI. (d) AGD selects the token ("royal") with the highest aggregated attribution score on the ROI, thereby satisfying the constraint.
• Figure 2: Analysis of attribution signal for word inclusion and exclusion tasks. (a) AGD improves performance on both task types. (b, c) Layer-wise attribution in the residual stream for Llama 3.1 (8B) at decision points where AGD's choice diverges from the greedy path to satisfy a constraint.
• Figure 3: Heatmap visualizations of attribution scores on input token embeddings. For each task, we compare a token that adheres to the instruction (adherent) with one that does not (non-adherent). Spaces are omitted for clarity. Full examples are in Appendix \ref{['appendix:heatmaps']}.
• Figure 4: Ablation analysis for Llama 3.1 (8B)
• Figure 5: Distribution of token entropy on IHEval dataset for Llama 3.1 (8B)