PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Xianzong Wu, Xiaohong Li, Yuejun Guo, Xinyang Liu, Tianlin Li, Junjie Wang, Qiang Hu

arXiv:2608.14653v1cs.LGcs.AI

TL;DR

Code generation remains vulnerable to early errors that standard decoding cannot revise, so this paper studies uncertainty-aware rollback decoding as an inference-time remedy. Across evaluated models and benchmarks, the complete framework outperforms equal-budget restart, while feedback-guided rollback provides the main improvement and uncertainty localization adds a further gain.

  • Problem

    Early code-generation errors can propagate irreversibly under left-to-right decoding, motivating evidence on whether uncertainty can identify unreliable regions for rollback.

  • Method

    The framework incrementally generates statement-like chunks, uses execution feedback to detect failures, and applies uncertainty signals to select rollback locations before resuming decoding.

  • Results

    The complete rollback framework improves over equal-budget restart across evaluated models and benchmarks, with gains up to 0.26 in pass@1 on functional benchmarks.

  • Takeaways & Limitations

    Information-theoretic signals show the most stable overall trend, while feedback-guided rollback supplies the main gain and uncertainty localization adds further improvement under matched conditions.

  • Takeaways & Limitations

    Because experiments use single runs without confidence intervals or significance tests, reported improvements are empirical trends rather than statistically conclusive effect estimates.

Abstract

from arXiv · show

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.

I. INTRODUCTION

The introduction frames code generation as fragile because early errors can invalidate programs, motivating uncertainty-aware rollback as a decoding-time control strategy. The study evaluates rollback, uncertainty, and branch-decay components across models, benchmarks, signals, and error types.

  • Early syntactic, semantic, or execution mistakes can invalidate an entire generated program, making code generation fundamentally fragile.This fragility distinguishes code generation from natural language generation, where constraints are less strict.
  • Uncertainty has mainly measured final-answer reliability, so the study investigates whether it can localize unreliable regions for decoding-time rollback and regeneration.The framework incrementally generates and checks code, rolls back after detected failures, and compares information-theoretic with sampling-based signals.
  • The complete rollback framework is compared with equal-budget restart across datasets and models, while controlled ablations isolate rollback, uncertainty, and branch-decay contributions.The study also examines uncertainty across rollback-triggering error types through distributional analysis.
  • Feedback-guided rollback provides the main improvement, uncertainty localization adds further gains, and information-theoretic signals offer more stable benefits than sampling-based alternatives.These findings summarize the matched component ablation and the cross-setting comparison of uncertainty signal families.
  • The paper formulates uncertainty-aware rollback decoding and contributes unified comparisons, multi-benchmark multi-model evaluation, controlled ablations, and diagnostics of localization, token cost, and same-budget alternatives.These contributions position uncertainty-aware rollback as a decoding-time strategy for mitigating error accumulation in LLM code generation.

II. PROBLEM STATEMENT · III. UNCERTAINTY SIGNALS · A. Single-Pass Token-Level Signals

The paper frames code generation as autoregressive decoding where irreversible prefix errors can cause later failures, motivating rollback guided by execution feedback and uncertainty localization. It distinguishes online single-pass token-level signals and evaluates six such measures computed from the model’s conditional token distribution.

  • II. PROBLEM STATEMENT: Autoregressive decoding generates tokens sequentially from a prompt, with each token conditioned on the previously generated prefix.The output sequence is modeled as y1:T, with y<t denoting the prefix before step t.
  • II. PROBLEM STATEMENT: Irreversible prefix errors can propagate into syntactic, semantic, or execution failures in generated code.Standard left-to-right decoding cannot revise an erroneous prefix.
  • II. PROBLEM STATEMENT: Rollback-based decoding revises a generated prefix after failure, using execution feedback to trigger reconsideration and uncertainty to localize the rollback position.A rollback position r satisfies r < t when the partial program fails an execution check or test case.
  • II. PROBLEM STATEMENT: The study keeps the rollback mechanism fixed while varying uncertainty signals to compare standard decoding, uncertainty-agnostic rollback, and uncertainty-aware rollback.This design isolates whether uncertainty helps determine where rollback should occur.
  • III. UNCERTAINTY SIGNALS: Because rollback operates online, uncertainty estimators are separated into single-pass token-level signals and multi-pass signals requiring repeated forward passes.Single-pass measures use one predictive distribution, whereas multi-pass estimates use stochastic inference or committee-style disagreement.
  • A. Single-Pass Token-Level Signals: Single-pass token-level signals are computed directly from p(· | x, y<t) without additional stochastic sampling, and six such signals are evaluated.These signals are designed for efficient computation at each decoding step.
  • A. Single-Pass Token-Level Signals: Avg NLL measures average token surprisal, while Max Token Entropy uses the largest token-level entropy within the current statement or span.Both are single-pass token-level uncertainty measures.
  • A. Single-Pass Token-Level Signals: Confidence-based signals convert token confidence into uncertainty, with Max Probability taking a statement maximum, Least Confidence taking a statement average, and Margin of Confidence averaging 1 −(p(1) −p(2)).Avg NLL, Least Confidence, Margin of Confidence, and Gini Impurity use statement means; Max Token Entropy and converted Max Probability use statement maxima, and larger values indicate higher uncertainty.

B. Multi-Pass Signals · C. Usage in Rollback Decoding

The study compares single-pass and multi-pass uncertainty signals within one unchanged rollback framework. Multi-pass estimators use repeated stochastic inference and cost more, while uncertainty complements execution feedback by localizing rollback points.

  • B. Multi-Pass Signals: Multi-pass signals require repeated stochastic forward passes, and the study evaluates Predictive Entropy and Variation Ratio.Predictive Entropy uses the entropy of the mean predictive distribution; Variation Ratio is 1 − f/N, where f is modal-token frequency over N samples.
  • B. Multi-Pass Signals: 1 − f/N defines Variation Ratio, with f denoting modal-token frequency across N stochastic samples.These signals require repeated stochastic inference and therefore cost more than single-pass signals.
  • C. Usage in Rollback Decoding: Uncertainty complements execution feedback: checks identify problematic prefixes, while uncertainty helps locate where to roll back.This division of roles keeps execution feedback and uncertainty signals functionally distinct within rollback decoding.
  • C. Usage in Rollback Decoding: For single-pass signals, token scores are aggregated over the current statement span before statement-level ranking.Avg NLL uses mean token surprisal, Max Token Entropy uses maximum token entropy, and confidence-style signals are oriented so larger values indicate lower confidence.
  • C. Usage in Rollback Decoding: For multi-pass signals, token scores come from multiple stochastic forward passes and are averaged within each statement before rollback triggers.Table I summarizes the evaluated uncertainty signals and groups them by computation type.
  • C. Usage in Rollback Decoding: All signals use identical decoding, checking, and rollback conditions, enabling comparison of which estimators best support rollback-based correction.The framework is held constant because the goal is estimator comparison rather than proposing a new estimator.
  • C. Usage in Rollback Decoding: Single-pass signals are emphasized for online use, whereas multi-pass estimators serve as highercost references.The distinction reflects the additional repeated stochastic inference required by multi-pass signals.

IV. UNCERTAINTY-AWARE DECODING FRAMEWORK

The framework incrementally generates newline-bounded chunks, checks executable partial programs, and revises failures through rollback while preserving verified prefixes. Feedback localizes errors when possible, while uncertainty guides fallback rollback and branch decay discourages repeating failed continuations.

  • Incremental statement generation: The decoder extends an accepted prefix with newline-bounded chunks, enabling validation after each local extension.Each chunk is generated token by token and treated as a statement-like unit in practice.
  • Program assembly and feedback checking: Partial candidates are normalized for execution and checked for syntax errors, public-test failures, and repetition before being accepted or revised.Normalization can complete indentation blocks and insert a temporary break inside an open while block.
  • Generation trace and branch decay: The trie records generated tokens, uncertainty values, and explored branches, while decaying failed-suffix edges to reduce exact repetition without blocking future alternatives.The penalty preserves flexibility when the search revisits an earlier prefix.
  • Rollback strategy: Rollback preserves the verified prefix and uses syntax offsets, repetition locations, or runtime error lines to determine where revision begins.If feedback localization is invalid or repeatedly fails, the decoder selects the highest-uncertainty statement on the current trace.
  • Framework roles: Intermediate checking, prefix-preserving rollback, and uncertainty-guided fallback provide distinct complementary roles during decoding.Uncertainty is consulted only when explicit feedback cannot reliably localize the underlying error.

V. EXPERIMENTAL SETUP

The experiments evaluate models without retraining while keeping the rollback mechanism, checking procedure, and decoding loop fixed across conditions.

  • Models are evaluated without retraining, with the rollback mechanism, checking procedure, and decoding loop fixed across conditions.

A. Benchmarks · B. Decoding Setup · C. Intermediate Checking and Rollback

The study evaluates uncertainty-aware rollback across diverse code-generation and security-repair benchmarks using online statement-level checking. Decoding combines fixed sampling and budget settings with hybrid feedback- and uncertainty-guided rollback, while separating online checks from hidden final evaluation.

  • A. Benchmarks: Five code-generation benchmarks and one security-oriented benchmark cover function synthesis, multilingual execution, competition problems, and secure code repair.The suite includes HumanEval, MBPP, HumanEval-CPP, CodeForces2305, APPS, and Dsec-Python.
  • A. Benchmarks: HumanEval contains 164 Python function-level tasks with hidden unit tests, using executable public inputs for intermediate checking when available.Final evaluation follows the benchmark-standard pipeline.
  • A. Benchmarks: The benchmark suite extends beyond function synthesis through structured Python tasks, stricter C++ compilation, long-range competition problems, difficult hidden-test tasks, and patch-aligned security repair.These roles are represented by MBPP, HumanEval-CPP, CodeForces2305, APPS, and Dsec-Python, respectively.
  • B. Decoding Setup: Generation proceeds statement by statement up to a newline or local structural boundary, with immediate checking enabling online rather than draft-level rollback.After each statement, benchmark-specific checks either accept the prefix or trigger rollback and regeneration from an earlier prefix.
  • B. Decoding Setup: Default decoding uses temperature = 0.0, top-k = 50, top-p = 1.0, one sample per task, a 768-token maximum, and a total budget of twice that length.Error-line rollback is limited to two consecutive failures at one location before uncertainty-guided localization activates, and regeneration applies path decay 0.9.
  • C. Intermediate Checking and Rollback: Online lightweight and public-test checks are strictly separated from hidden tests, which are accessed only after generation terminates for final evaluation.Lightweight checks include syntax validation, executable-prefix checks, and repetition detection; public examples guide online rollback decisions.
  • C. Intermediate Checking and Rollback: Hybrid rollback uses explicit failure locations when valid, then selects the statement containing the highest-uncertainty token after invalid or repeated failures.Path decay downweights previously failed continuations rather than blocking them, while execution or analysis feedback and uncertainty localization serve distinct roles.

D. Uncertainty Signals · E. Baselines and Ablation Settings · F. Evaluation Metrics

The study evaluates token-level uncertainty signals under controlled multi-pass estimation, compares rollback against equal-budget restart and component ablations, and reports task-specific functional and patch-safety metrics. Diagnostics separately quantify uncertainty-selected rollback behavior while treating line distance and AssertionError as operational proxies rather than manually labeled categories.

  • D. Uncertainty Signals: Multi-pass uncertainty signals use N = 5 stochastic forward passes with MC dropout and sampling temperature 1.0.Uncertainty estimation temporarily enables training mode with gradients disabled, then restores evaluation mode for generation.
  • E. Baselines and Ablation Settings: The main RQ1 baseline is equal-budget restart, which allows two 768-token attempts within a 1,536-token budget without feedback, rollback, path decay, or uncertainty localization.
  • E. Baselines and Ablation Settings: Component ablations separate rollback, uncertainty localization, and branch decay by testing rollback without uncertainty guidance and uncertainty scoring without rollback.
  • E. Baselines and Ablation Settings: Same-budget comparisons on HumanEval and CodeForces2305 use CodeLlama-7B, temperature = 0.2, a 1536-token budget, and shared task-indexed seeds.Restart from scratch discards failed attempts and regenerates complete solutions from the original prompt without error feedback.
  • F. Evaluation Metrics: For HumanEval, MBPP, HumanEval-CPP, CodeForces2305, and APPS, evaluation reports pass@1 and AvgTestPassRate from one generated program per task.pass@1 requires passing all hidden tests, while AvgTestPassRate macro-averages the fraction of hidden tests passed per task.
  • F. Evaluation Metrics: For Dsec-Python, PASR requires line-level IoU ≥0.5 and normalized similarity ≥0.6 in the modified region, with success if any of at most three candidates meets both criteria.The equal-budget restart and full framework share the matched total token budget, and PASR is computed only after generation.
  • F. Evaluation Metrics: Rollback diagnostics retain uncertainty_fallback and uncertainty_refinement events and report selected-event counts, comparable-error-line subsets, within-3-line rates, and next-check success fractions.Line distance remains an operational proxy rather than a manually annotated measure.
  • F. Evaluation Metrics: Functional-error analysis uses AssertionError as a proxy for programs that execute but fail test assertions, not as a manually labeled semantic-error category.

G. Implementation Details · VI. RESULTS · equal-budget restart across datasets and models?

The study evaluates rollback decoding across diverse models, benchmarks, and uncertainty signals under a documented implementation setup. Compared with equal-budget restart, the complete framework improves code-generation results across datasets and model scales, while information-theoretic signals show the strongest overall trend without being universally optimal.

  • G. Implementation Details: The main experiments used PyTorch and HuggingFace Transformers on four NVIDIA Tesla V100-SXM2 GPUs, 28 CPU cores, and 754GB RAM.The environment used Python 3.10, PyTorch 2.8.0, Transformers 4.40.0, and CUDA 12.7.
  • VI. RESULTS: The study evaluates overall effectiveness, uncertainty signals, model-scale effects, rollback diagnostics, component ablations, and error types using Tables II–VI and Figures 2–3.These analyses cover the main results and complementary diagnostic evaluations.
  • equal-budget restart across datasets and models?: 0.04/0.14 to 0.08/0.26 and 0.11/0.30 to 0.15/0.40 are the CodeForces2305 gains for StarCoder2-7B and CodeLlama-34B, respectively.On Dsec-Python, the framework improves from 0.50 to 0.56 for DeepSeek-7B and from 0.52 to 0.57 for Qwen2.5-7B, with smaller margins than on execution-based benchmarks.
  • equal-budget restart across datasets and models?: Max Token Entropy and Avg NLL show the strongest overall trend among uncertainty signals on standard code-generation benchmarks.The signal comparison is not equal-compute because multi-pass signals require five forward passes per token position.
  • equal-budget restart across datasets and models?: No single uncertainty signal is universally optimal, with competitive alternatives such as Gini Impurity and model-dependent best signals on Dsec-Python.Differences among signals become smaller on harder settings including CodeForces2305, APPS-Competition, and Dsec-Python.
  • equal-budget restart across datasets and models?: Improvements appear across DeepSeek-7B, StarCoder2-7B, Qwen2.5-7B, Mistral-7B, CodeLlama-7B, and SafeCoder-7B, indicating broad backbone coverage.The framework also remains effective for the stronger CodeLlama-34B model, including gains on HumanEval, MBPP, and HumanEval-CPP.
  • equal-budget restart across datasets and models?: The overall rollback-decoding trend persists across different parameter scales and both functional and security-oriented benchmarks.The paper also reports that the fine-tuned variant improves over its base model under the same rollback configuration.

B. RQ2: What are the contributions of rollback, uncertainty, and branch decay? · 1) Rollback Diagnostics and Controlled Comparison:

The controlled diagnostics show that entropy-selected rollback often localizes failures closely and enables successful subsequent checks, while the full framework outperforms restart and direct error-line rollback under the same token budget. These findings support rollback’s contribution while highlighting that diagnostic localization is not a counterfactual comparison and that uncertainty-score overhead is not end-to-end decoding overhead.

  • 1) Rollback Diagnostics and Controlled Comparison:: Entropy-selected rollback targets the highest-entropy token’s line after invalid reports or repeated failures, without using a fixed entropy threshold.The diagnostic uses CodeLlama-7B, Max Token Entropy, temperature 0.0, and seed 0 on HumanEval and CodeForces2305.
  • 1) Rollback Diagnostics and Controlled Comparison:: Uncertainty-selected points are often near observed failures and frequently make the next check pass, but this diagnostic is not a counterfactual comparison with alternative localization rules.The measurements isolate entropy localization from feedback-only rollback.
  • 1) Rollback Diagnostics and Controlled Comparison:: 11.52 and 20.24 average rollbacks per task occur on HumanEval and CodeForces2305, respectively.
  • 1) Rollback Diagnostics and Controlled Comparison:: 3.99% and 1.86% of runtime are spent computing uncertainty scores on HumanEval and CodeForces2305, respectively.These figures do not measure end-to-end overhead relative to standard decoding.
  • 1) Rollback Diagnostics and Controlled Comparison:: 41.46% pass rate from the full framework exceeds 23.17% for restart and 32.32% for direct error-line rollback on HumanEval under a 1536-token budget.The full framework also uses fewer average generated tokens than both alternatives.
  • 1) Rollback Diagnostics and Controlled Comparison:: 2.22% pass rate is achieved by both restart and error-line rollback on CodeForces2305, while the full framework follows the same improvement trend.The comparison uses the same 1536-token budget.
  • 1) Rollback Diagnostics and Controlled Comparison:: 110 to 76 and 73 to 67 final AssertionError proxy cases decrease on HumanEval and CodeForces2305 when the full method replaces error-line rollback.This temperature-0.2 same-budget result is a test-level proxy rather than manually annotated semantic-error ground truth; Table III compares uncertainty signals using pass@1, AvgTestPassRate, and PASR.

2) Component Ablation: · C. RQ3: How does uncertainty behave across rollback-triggering error types?

Component ablation shows feedback-guided rollback provides the main improvement, with uncertainty localization adding further gain under controlled decoding conditions. Across rollback-triggering errors, uncertainty is higher for structurally unstable failures, supporting its use to localize rollback regions rather than predict correctness alone.

  • 2) Component Ablation:: Feedback-guided rollback reaches 0.33 pass@1 and 0.53 AvgTestPassRate on HumanEval, compared with 0.25 and 0.23 pass@1 for feedback-only detection and vanilla decoding.The comparison is reported for CodeLlama-7B.
  • 2) Component Ablation:: Uncertainty localization adds gain when feedback, rollback, branch decay, checking, and budget are held fixed.Feedback-Guided Rollback and the Full Framework differ only in uncertainty localization.
  • C. RQ3: How does uncertainty behave across rollback-triggering error types?: Rollback handles both localized repairs and broader structural revisions across rollback-triggering error types.NameError and SyntaxError use small fragment replacements, whereas IndexError and RepeatPatternError require larger statement or control-flow rewrites.
  • C. RQ3: How does uncertainty behave across rollback-triggering error types?: NameError (n = 283) and TypeError (n = 202) are among the most frequent rollback-triggering error categories.They are followed by RepeatPatternError, IndexError, and ValueError.
  • C. RQ3: How does uncertainty behave across rollback-triggering error types?: 1.85 is the highest mean uncertainty for SyntaxError, while IndentationError and TypeError each have 1.51.Common local semantic errors show lower means: NameError 0.94, IndexError 0.83, and ValueError 0.70.
  • C. RQ3: How does uncertainty behave across rollback-triggering error types?: Uncertainty is more useful as a cue for rollback localization than as a standalone predictor of final correctness.The observed distribution supports prioritizing regions where the decoding trajectory becomes unstable instead of simply revisiting the most recently generated tokens.
  • C. RQ3: How does uncertainty behave across rollback-triggering error types?: Categories with fewer than five events are reported descriptively and should be interpreted cautiously.Error bars summarize dispersion and are not confidence intervals.
  • C. RQ3: How does uncertainty behave across rollback-triggering error types?: Rollback-triggering uncertainty tends to be higher for structurally unstable errors than for common local semantic mistakes.This pattern associates uncertainty with structural instability rather than merely error frequency.

VII. THREATS TO VALIDITY · VIII. CONCLUSION

The study’s findings are empirical and context-dependent because of statistical, construct, external, and methodological limitations. Nevertheless, the complete rollback framework improves over equal-budget restart across evaluated settings, with reported gains on functional benchmarks and Dsec-Python.

  • VII. THREATS TO VALIDITY: Single-run evaluation omitted repeated trials, confidence intervals, and formal significance tests, so improvements are empirical trends rather than statistically conclusive effect estimates.This caution is especially important when differences between uncertainty signals are small.
  • VII. THREATS TO VALIDITY: Functional correctness metrics and Patch-Aligned Safe Rate do not fully capture generation quality or subtle security improvements.Dsec-Python uses patch alignment and lightweight static validation rather than end-to-end exploit-level verification.
  • VII. THREATS TO VALIDITY: The limited model coverage and Python-centered benchmarks constrain external validity, so findings apply to evaluated settings rather than universally.The study includes several 7B-scale open-source models and one 34B model.
  • VII. THREATS TO VALIDITY: Observed uncertainty-signal effectiveness is specific to the fixed rollback framework, intermediate-checking design, and branch-decay policy.Different rollback strategies, checking mechanisms, or search policies may change relative estimator behavior.
  • VII. THREATS TO VALIDITY: Diagnostic and same-budget runtime results came from a separate machine and should not be directly compared with the main experimental server.Lightweight intermediate checking may also miss some failure modes unequally.
  • VII. THREATS TO VALIDITY: Future work targets adaptive rollback, alternative intermediate verification, larger proprietary code models, and stronger execution feedback for realistic programming environments.The passage specifically anticipates stronger execution feedback improving uncertainty-guided localization.
  • VIII. CONCLUSION: 0.26 in pass@1, 0.35 in AvgTestPassRate, and 6.4% in Patch-Aligned Safe Rate were the maximum reported gains over equal-budget restart.These gains came from the complete framework across evaluated models and benchmarks, with the first two on functional benchmarks and the latter on Dsec-Python.
Loading 2608.14653v1…