AdaDec: A Uncertainty-Guided Lookahead Decoding Framework for LLM-Based Code Generation

1. Problem Statement

Large Language Models (LLMs) have revolutionized code generation, but their performance is highly sensitive to the decoding strategy used to select tokens.

• The Core Issue: Standard strategies like Greedy Decoding treat all tokens uniformly, often committing to suboptimal choices at high-uncertainty decision points. In code generation, a single incorrect token (e.g., choosing return instead of for at a critical logic junction) can cause an irreversible "logic drift," leading to syntactically correct but semantically broken code.
• Limitations of Existing Methods:
  • Beam Search: Improves accuracy by exploring multiple paths but incurs high computational overhead by searching globally at every step.
  • AdapT: Adjusts sampling temperature based on line-start positions but relies on coarse heuristics and randomness, leading to unstable gains.
  • Uncert-CoT: Uses uncertainty to trigger Chain-of-Thought reasoning but is limited to line boundaries and lacks dynamic adaptation to model-specific uncertainty patterns.
• Gap: There is a lack of a decoding framework that dynamically identifies specific high-uncertainty tokens (both at line starts and within lines) and applies targeted, efficient corrections without the cost of global search.

2. Methodology: AdaDec Framework

The authors propose AdaDec, an uncertainty-guided adaptive decoding framework that integrates a "Pause-then-Rerank" mechanism. The system operates in three main stages:

A. Empirical Foundation

Before designing the framework, the authors conducted an empirical study on models (DeepSeek-Coder, Qwen, etc.) and the HumanEval dataset.

• Finding 1: Many generation errors stem from local ranking mistakes where the correct token exists in the candidate set but is not ranked #1. These "logic drift points" occur not just at line starts but also within lines.
• Finding 2: Shannon Entropy is a strong signal for detecting these drift points. High entropy correlates strongly with the ground-truth token having a low rank. However, a fixed entropy threshold is suboptimal across different models.

B. Core Components of AdaDec

1. Model-Specific Entropy Threshold Learning ($\tau_{LM}$):

   • Instead of using a fixed threshold, AdaDec learns a model-specific threshold using Logistic Regression.
   • Training: A binary classifier is trained to predict whether the ground-truth token is the top-1 prediction based solely on the entropy value.
   • Data: Uses ground-truth guided next-token prediction on a code corpus (BigCodeBench) to generate positive (correct top-1) and negative (incorrect top-1) samples.
   • Outcome: A dynamic threshold $\tau_{LM}$ that accurately identifies high-uncertainty steps for a specific model.

2. Pause-then-Rerank Mechanism:

   • Pause Check: During inference, at each step $t$, the model computes the Shannon entropy $H(p_t)$.
     • If $H(p_t) \le \tau_{LM}$: Proceed with the base decoding strategy (e.g., Greedy).
     • If $H(p_t) > \tau_{LM}$: Pause normal decoding and trigger the reranking routine.
   • Lookahead Reranking:
     • Select the top-$B$ candidate tokens (e.g., $B=3$).
     • For each candidate, generate a lookahead trajectory of length $L$ (e.g., $L=5$) using greedy decoding.
     • Scoring: Calculate the average log-probability of the entire trajectory (current token + $L$ future tokens) to estimate long-term coherence.
     • Selection: Choose the candidate with the highest trajectory score as the next token.

3. Key Contributions

1. Empirical Discovery: Demonstrated that code generation errors are frequently caused by local ranking errors at high-uncertainty steps (detected via Shannon entropy), occurring both at line starts and within lines.
2. AdaDec Framework: Introduced a novel uncertainty-guided adaptive decoding framework that learns model-specific thresholds and employs a lookahead-based reranking strategy only when necessary.
3. Efficiency-Accuracy Trade-off: Proved that selective intervention (pausing only ~6.75% of the time) yields significant accuracy gains without the massive latency overhead of global Beam Search.

4. Experimental Results

Experiments were conducted on HumanEval+, MBPP+, and DevEval benchmarks using models ranging from 0.6B to 8B parameters.

• Accuracy (Pass@1):
  • AdaDec improves Pass@1 by up to 20.9% (absolute) over Greedy Decoding (e.g., on MBPP+ with Qwen3-0.6B).
  • It consistently outperforms Beam Search (with $B=3$) and the state-of-the-art AdapT method across all benchmarks.
  • On the repository-level DevEval benchmark, AdaDec showed robust improvements (+4.53% average) where other methods struggled.
• Efficiency:
  • Pause Rate: The framework triggers the expensive reranking mechanism only ~6.75% of the time on average.
  • Latency: While slightly slower than Greedy decoding, AdaDec is significantly faster than Beam Search (e.g., on DevEval, Beam Search adds ~24s latency, while AdaDec adds only ~7.5s).
  • Counter-intuitive Finding: For some small models, AdaDec was faster than Greedy because it prevented the model from generating long, erroneous loops or repetitive text, resulting in shorter, correct solutions.
• Ablation Studies:
  • Learned vs. Fixed Threshold: Using the learned threshold consistently outperformed a fixed threshold (e.g., 1.2), proving the necessity of model-specific calibration.
  • Hyperparameters: The system is robust to variations in lookahead length ($L$) and threshold shifts, with optimal performance at $L=5$.

5. Significance

• Reliability: AdaDec addresses the critical issue of "logic drift" in code generation by intervening precisely where the model is uncertain, significantly reducing semantic errors.
• Scalability: By learning thresholds and using selective pausing, the method scales efficiently across different model sizes (0.6B to 8B) without requiring massive computational resources like global beam search.
• Generalizability: The framework is model-agnostic and can be integrated with various base decoding strategies (Greedy, Sampling), making it a versatile plug-in for improving LLM-based code generation reliability.
• Future Direction: The paper highlights that uncertainty-aware adaptive decoding is a promising direction for balancing the trade-off between generation quality and inference cost in software engineering tasks.

In conclusion, AdaDec represents a shift from static or globally expensive decoding strategies to a dynamic, data-driven approach that leverages model uncertainty to make targeted corrections, achieving state-of-the-art performance in code generation with minimal overhead.