# Do Uncertainty Signals Help? A Systematic Study of Uncertainty-Aware Decoding with Rollback Mechanisms

arXiv: https://arxiv.org/abs/2608.14653 · HTML: https://mlsanity.com/paper/2608.14653/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 31 of its 41 references that are in the map.
- First posted on arXiv 31 Jul 2026.

## Builds on (most cited first, 30 of 31)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [Self-Refine: Iterative Refinement with Self-Feedback](https://mlsanity.com/paper/2303.17651/) (2023)
- [Competition-Level Code Generation with AlphaCode](https://mlsanity.com/paper/2203.07814/) (2022)
- [Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code Generation](https://mlsanity.com/paper/2305.01210/) (2023)
- [Teaching Large Language Models to Self-Debug](https://mlsanity.com/paper/2304.05128/) (2023)
- [Measuring Coding Challenge Competence With APPS](https://mlsanity.com/paper/2105.09938/) (2021)
- [Semantic Uncertainty: Linguistic Invariances for Uncertainty Estimation in Natural Language Generation](https://mlsanity.com/paper/2302.09664/) (2023)
- [BigCodeBench: Benchmarking Code Generation with Diverse Function Calls and Complex Instructions](https://mlsanity.com/paper/2406.15877/) (2024)
- [CodeT: Code Generation with Generated Tests](https://mlsanity.com/paper/2207.10397/) (2022)
- [CodeRL: Mastering Code Generation through Pretrained Models and Deep Reinforcement Learning](https://mlsanity.com/paper/2207.01780/) (2022)
- [CodeGeeX: A Pre-Trained Model for Code Generation with Multilingual Benchmarking on HumanEval-X](https://mlsanity.com/paper/2303.17568/) (2023)
- [OpenCodeInterpreter: Integrating Code Generation with Execution and Refinement](https://mlsanity.com/paper/2402.14658/) (2024)
- [LEVER: Learning to Verify Language-to-Code Generation with Execution](https://mlsanity.com/paper/2302.08468/) (2023)
- [Is Self-Repair a Silver Bullet for Code Generation?](https://mlsanity.com/paper/2306.09896/) (2023)
- [Planning with Large Language Models for Code Generation](https://mlsanity.com/paper/2303.05510/) (2023)
- [Kernel Language Entropy: Fine-grained Uncertainty Quantification for LLMs from Semantic Similarities](https://mlsanity.com/paper/2405.20003/) (2024)
- [Self-planning Code Generation with Large Language Models](https://mlsanity.com/paper/2303.06689/) (2023)
- [Generalization or Memorization: Data Contamination and Trustworthy Evaluation for Large Language Models](https://mlsanity.com/paper/2402.15938/) (2024)
- [Instruction Tuning for Secure Code Generation](https://mlsanity.com/paper/2402.09497/) (2024)
- [A Survey of Large Language Models for Code: Evolution, Benchmarking, and Future Trends](https://mlsanity.com/paper/2311.10372/) (2023)
- [CodeChain: Towards Modular Code Generation Through Chain of Self-revisions with Representative Sub-modules](https://mlsanity.com/paper/2310.08992/) (2023)
- [Iterative Refinement of Project-Level Code Context for Precise Code Generation with Compiler Feedback](https://mlsanity.com/paper/2403.16792/) (2024)
- [Towards Understanding the Characteristics of Code Generation Errors Made by Large Language Models](https://mlsanity.com/paper/2406.08731/) (2024)
- [ROCODE: Integrating Backtracking Mechanism and Program Analysis in Large Language Models for Code Generation](https://mlsanity.com/paper/2411.07112/) (2024)
- [JavaBench: A Benchmark of Object-Oriented Code Generation for Evaluating Large Language Models](https://mlsanity.com/paper/2406.12902/) (2024)
- [Active Layer-Contrastive Decoding Reduces Hallucination in Large Language Model Generation](https://mlsanity.com/paper/2505.23657/) (2025)
- [The First Prompt Counts the Most! An Evaluation of Large Language Models on Iterative Example-Based Code Generation](https://mlsanity.com/paper/2411.06774/) (2024)

## Abstract (arXiv)

Prediction uncertainty is a widely adopted metric for quantifying model confidence, with downstream applications spanning model explanation, data selection, and prediction rollback. Despite its demonstrated utility, the potential of uncertainty quantification to enhance code generation in large language models (LLMs) remains largely underexplored, raising a critical question: to what extent can uncertainty serve as an effective signal for improving LLM-based code generation? To answer this question, we study uncertainty-aware rollback decoding, an inference-time strategy that uses uncertainty signals to identify unreliable generation regions and roll back to earlier valid prefixes without retraining the model. We evaluate this framework on seven code LLMs, five code generation benchmarks, and eight token-level uncertainty signals under a unified decoding setup. Our results show that the complete rollback framework improves over equal-budget restart across the evaluated benchmarks and model settings, with gains of up to 0.26 in pass@1 and 0.35 in AvgTestPassRate on functional code generation benchmarks, and an absolute improvement of up to 6.4% in Patch-Aligned Safe Rate on Dsec-Python. Among the evaluated signals, information-theoretic measures such as token entropy and negative log-likelihood show the most favorable overall trend, frequently achieving the best or near-best results on standard benchmarks. A component-controlled ablation further shows that feedback-guided rollback provides the main improvement, while uncertainty localization provides an additional gain when checking, budget, rollback, and branch decay are held fixed.