Reflection-Window Decoding: Text Generation with Selective Refinement

2025-02-05

AI Summary

Autoregressive decoding in large language models (LLMs) lacks backtracking capability, leading to suboptimal sequences that deviate from the globally optimal joint probability distribution. To address this, we propose an uncertainty-aware selective refinement framework. Our method introduces token-level uncertainty estimation, integrated with a sliding reflective window and a dynamic pausing mechanism, enabling localized resampling and re-decoding during generation. The framework is plug-and-play—requiring no architectural modifications—and preserves high inference efficiency while approaching joint-probability optimality. Evaluated across diverse open-ended generation and reasoning benchmarks, it significantly improves factual consistency and fluency, yielding BLEU and ROUGE gains of 2.1–4.3 points, with inference latency overhead under 12%.

Technology Category

• Reasoning under Uncertainty: Sequential Decision Making
• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

The autoregressive decoding for text generation in large language models (LLMs), while widely used, is inherently suboptimal due to the lack of a built-in mechanism to perform refinement and/or correction of the generated content. In this paper, we consider optimality in terms of the joint probability over the generated response, when jointly considering all tokens at the same time. We theoretically characterize the potential deviation of the autoregressively generated response from its globally optimal counterpart that is of the same length. Our analysis suggests that we need to be cautious when noticeable uncertainty arises during text generation, which may signal the sub-optimality of the generation history. To address the pitfall of autoregressive decoding for text generation, we propose an approach that incorporates a sliding reflection window and a pausing criterion, such that refinement and generation can be carried out interchangeably as the decoding proceeds. Our selective refinement framework strikes a balance between efficiency and optimality, and our extensive experimental results demonstrate the effectiveness of our approach.

Problem

• Improves text generation in LLMs
• Addresses suboptimal autoregressive decoding
• Introduces selective refinement mechanism

Innovation

• sliding reflection window
• pausing criterion
• selective refinement framework