Source-linked AI summary

Uncertainty-Guided Chain-of-Thought for Code Generation with LLMs

Yuqi Zhu, Ge Li, Xue Jiang, Jia Li, Hong Mei, Zhi Jin, Yihong Dong

arXiv:2503.15341v1cs.SE

TL;DR

Existing CoT methods can overthink code-generation tasks, adding reasoning that may introduce errors and consume computation unnecessarily. UnCert-CoT uses uncertainty measures to invoke CoT selectively, and reports improved PassRate on HumanEval and MHPP.

  • Problem

    Existing CoT can add unnecessary reasoning that distorts correct answers and spends computation on tasks solvable directly.

    Source: §1 · p. 1
  • Method

    UnCert-CoT uses entropy and the probability gap between the top two predictions to trigger multi-path, confidence-selected CoT reasoning when uncertainty is high.

    Source: §3.1 · p. 3, §1 · p. 3, §3.3 · p. 5
  • Results

    UnCert-CoT improves PassRate over state-of-the-art baselines by 3.5% on HumanEval and 6.1% on MHPP.

    Source: §5 · p. 6
  • Takeaways & Limitations

    The method outperforms state-of-the-art CoT methods across various LLMs and shows some robustness to hyperparameter settings.

    Source: §9 · p. 10

    Evaluation uses two benchmarks, HumanEval and MHPP, so external validity depends in part on dataset quality and scope.

    Source: §8 · p. 9

Abstract

from arXiv · show

Chain-of-Thought (CoT) reasoning has been demonstrated as an effective technique for improving the problem-solving capabilities of large language models (LLMs) in the context of code generation. However, existing CoT methods often exhibit a tendency toward "overthinking", where the LLM consistently applies reasoning strategies without adequately considering the task's underlying complexity. This results in the LLMs allocating excessive computational resources, in terms of tokens, to relatively simple tasks or problems where the correct answer is already evident. Additionally, this overthinking may lead LLMs down incorrect reasoning paths, resulting in incorrect code generation. In this paper, we introduce UnCertainty-Aware Chain-of-Thought (UnCert-CoT), an LLM-based approach designed to enhance code generation by incorporating an uncertainty-aware CoT reasoning mechanism, which focuses computational resources on targeting points where LLMs are more prone to error. We propose two confidence-based uncertainty measures: Entropy-based and Probability Differential-based methods. When uncertainty is high, UnCert-CoT activates CoT-decoding to generate multiple reasoning paths and selects the final code that exhibits the highest likelihood of correctness. In contrast, LLM directly generates the code when uncertainty is low. This uncertainty judgment mechanism allows LLMs to prioritize complex tasks and avoid unnecessary steps in simpler cases, thereby improving overall efficiency and accuracy in code generation. Our experimental results demonstrate that UnCert-CoT significantly enhances code generation accuracy on challenging benchmark MHPP(Mostly Hard Python Problems), it achieves improvements up to 6.1% on PassRate accuracy, particularly in situations where traditional LLMs are prone to errors.

1 Introduction

Existing CoT code-generation methods can overthink tasks, adding unnecessary reasoning that may distort correct answers and waste computation. UnCert-CoT uses uncertainty to decide when to invoke reasoning, and reports gains on HumanEval and MHPP.

1 Introduction

  • Overthinking occurs when unnecessary reasoning distorts a model’s answer and consumes computation even when direct code generation could succeed.The paper argues resources should instead go to areas requiring deeper reasoning.Source: §1 · p. 1
  • UnCert-CoT checks uncertainty at the first non-indentation token of a new code line, invoking CoT-decoding above a threshold and otherwise generating code directly.CoT-decoding explores multiple reasoning paths and selects a high-confidence result.Source: §1 · pp. 1–2
  • Entropy of the first non-indentation token and the probability gap between the top two predictions are the proposed uncertainty measures.Both measures are used to detect high uncertainty and trigger reasoning.Source: §1 · p. 2
  • The approach targets line-level reasoning, preserving direct generation at low uncertainty while selecting confident outcomes to mitigate error accumulation in complex tasks.The paper presents this as a way to preserve efficiency and reduce error propagation.Source: §1 · p. 2
  • UnCert-CoT outperforms baselines by 3.5% on HumanEval and 6.1% on MHPP, with consistent gains across model sizes and families.The paper also reports robustness to hyperparameter configurations.Source: §1 · p. 2

2 Motivating Example

The motivating example contrasts correct direct code generation with CoT reasoning that introduces a greedy-algorithm error in a dynamic-programming task. It illustrates why reasoning should be applied selectively rather than assumed to help every task.

2 Motivating Example

  • The example illustrates overthinking: intermediate reasoning can introduce errors that would not occur with direct code generation.The authors describe this as a counterintuitive possibility in code generation.Source: §2 · p. 3
  • Direct code generation correctly solves the product-pricing task with dynamic programming, while CoT reasoning switches to a flawed greedy approach.The greedy approach fails to consider alternative product combinations and incorrectly treats a negative remainder as proof of impossibility.Source: §2 · p. 3
  • The authors argue for selectively applying reasoning based on problem characteristics rather than assuming CoT universally improves code generation.Source: §2 · p. 3

3 Methodology

UnCert-CoT measures uncertainty at the start of each new code line and selectively invokes reasoning when uncertainty exceeds a threshold. It uses entropy or the gap between the top two token probabilities, then selects the highest-confidence reasoning result.

3.1 Overview

  • UnCert-CoT detects uncertainty at new-line boundaries, where prediction difficulty is higher because the line shapes the subsequent logical structure.When uncertainty exceeds the threshold, it invokes CoT-decoding to guide the next code line.Source: §3.1 · p. 3

3.2 Uncertainty Computation

  • The uncertainty component uses entropy or the difference between the most- and second-most-likely token probabilities to decide when to invoke CoT.Entropy measures output-distribution randomness; the probability-differential measure compares the top two predicted probabilities.Source: §3.1 · p. 3

3.2.1 Entropy-based uncertainty measurement.

  • Entropy-based uncertainty uses the LLM’s prediction distribution over possible tokens to calculate uncertainty for the next generated token.The passage describes the distribution as conditional on the input and previously generated context.Source: §3.2.1 · p. 4

3.2.2 Probability Differential-based uncertainty measurement.

  • The probability-differential method uses the difference between the highest and second-highest predicted probabilities as an uncertainty indicator.A larger U_d(p) value indicates greater certainty at the current step.Source: §3.2.2 · p. 4

3.3 CoT-decoding

  • For high uncertainty, CoT-decoding generates multiple reasoning paths and selects the result with the highest confidence score.For code generation, few-shot examples guide sampling reasoning paths with temperature t.Source: §3.3 · p. 5

3.4 UnCert-CoT

  • At each new code line, UnCert-CoT invokes CoT-decoding above a predefined uncertainty threshold and otherwise generates the line greedily.The low-uncertainty branch selects the maximum-probability token at each token step.Source: §3.1 · p. 3, §3.4 · p. 5

4 Study Design

The study evaluates UnCert-CoT through comparisons with existing CoT methods, different LLMs, an ablation, and uncertainty-threshold settings. It uses HumanEval and MHPP and describes the code LLMs used in the experiments.

4 Study Design

  • The study is presented as a large-scale evaluation of UnCert-CoT, with details of datasets, baselines, models, and evaluation metrics.Table 1 compares the performance of CoT methods on code-generation tasks.Source: §4 · p. 5, Table 1 · p. 6

4.1 Research Questions

  • The study evaluates UnCert-CoT against state-of-the-art CoT methods, across LLMs, through a CoT-decoding ablation, and under varying uncertainty thresholds.These are the four research questions posed for the study.Source: §4.1 · p. 5

4.2 Datasets

  • HumanEval contains 164 Python programming problems; MHPP contains 210 curated problems focused on specification comprehension, restrictions, multistep reasoning, and coding knowledge.MHPP’s evaluation pipeline was developed to mitigate data leakage, and the passage describes it as a challenging benchmark.Source: §4.2 · p. 5

4.3 Models

  • The models section describes DeepSeekCoder and CodeLlama and specifies DeepSeekCoder-6.7B-Base as the model used.The provided passage also describes DeepSeekCoder as a series ranging from 1B to 33B parameters.Source: §4.3 · p. 5

UnCert-CoT

This section describes the Qwen2.5-Coder-7B-Base model, defines the comparison baselines, and specifies PassRate as the evaluation metric.

UnCert-CoT

  • The study uses the Qwen2.5-Coder-7B-Base model from the QwenCoder series.The passage describes QwenCoder as a series designed for coding tasks, based on Qwen2.5 models and further fine-tuned on coding instruction data.Source: UnCert-CoT · p. 6

4.4 Baselines

  • The baselines are greedy Base Model, Zero-shot CoT, few-shot Self-planning, and CoT-Decoding with multiple reasoning samples.CoT-Decoding selects the answer with the highest confidence value.Source: §4.4 · p. 6

4.5 Evaluation Metric

  • PassRate measures functional correctness by executing test cases, dividing the number of problems answered correctly by the total number of problems.The passage denotes the correct-answer count as C and the total problem count as N.Source: §4.5 · p. 6

5 RESULTS AND ANALYSES

UnCert-CoT outperforms the reported baselines on code-generation benchmarks, with larger gains on MHPP than HumanEval. Its results also vary with the uncertainty threshold, and removing CoT decoding particularly reduces performance on difficult problems.

RQ1

  • UnCert-CoT improves PassRate over state-of-the-art baselines by 3.5% on HumanEval and 6.1% on MHPP.The authors report the larger improvement on MHPP, which requires more complex code reasoning.Source: §5 · p. 6

RQ2

  • Across DeepSeekCoder, Qwen-Coder, and CodeLlama, UnCert-CoT improves performance over the base models on HumanEval and MBPP.Reported average gains on HumanEval are 16.0%, 1.2%, and 10.2%, respectively; on MBPP they are 5.1%, 2.0%, and 14.5%.Source: §5 · p. 7

RQ3

  • Removing CoT decoding lowers average performance by 5.5% on HumanEval and 9% on MHPP, with a particularly pronounced decline on harder problems.The authors attribute the loss to the reduced ability to refine reasoning through multiple steps and paths.Source: §5 · pp. 7–8

RQ4

  • When τ is below 0.1, UnCert-CoT performs below the baseline; performance rises to a peak as τ increases, then declines.The authors attribute the low-threshold result to unnecessary reasoning when the solution is already clear.Source: Figure 3 · p. 8, Figure 4 · p. 8, §5 · p. 8
  • The Entropy-based method improves over thresholds τ∈[0.2, 0.3], while the Probability Differential-based method improves over τ∈[0.2, 0.7].The authors report favorable results across multiple threshold settings and different optimal ranges for the two methods.Source: §5 · p. 8

6 Discussion

The discussion distinguishes UnCert-CoT from Structured Chain-of-Thought (SCoT) in prompting and decoding, and presents a case where UnCert-CoT passes all tests while two alternatives fail.

6 Discussion

  • Unlike SCoT’s curated structured prompts, UnCert-CoT uses a few-shot code-and-plan example without specialized prompt engineering.SCoT also uses reflection in a two-stage generation process; the paper therefore does not directly compare it with UnCert-CoT.Source: §6 · pp. 8–9
  • In the case study, UnCert-CoT passes all test cases, while the base model and self-planning method fail.The case is presented in Figure 5.Source: §6 · p. 9

7 Related Work

Related work covers code-generation techniques, the development of Chain-of-Thought prompting, and uncertainty estimation in language models, including its relation to code-token correctness.

7.1 Code Generation with LLMs

  • Code-generation work includes prompting, domain-specific fine-tuning, and decoding methods that use tests, runtime feedback, or sampling.The review describes Code LLMs as trained on large codebases and lists applications such as completion, debugging, and refactoring.Source: §7.1 · p. 9

7.2 Chain-of-Thought Prompting

  • CoT prompting generates intermediate reasoning steps; later work includes self-consistency voting across paths and code-generation methods such as self-planning and SCoT.The review also notes zero-shot demonstrations, the “Let’s think step by step” instruction, and CoT capabilities in OpenAI o1 and Deepseek-R1 [15].Source: §7.2 · p. 9

7.3 Uncertainty of LLMs

  • Uncertainty describes variability across model outputs, while confidence concerns the likelihood that a specific prediction is correct.Prior code-completion work reports that higher entropy-based uncertainty is associated with a greater likelihood of incorrect generated tokens [34].Source: §7.3 · p. 9

8 Threats to Validity

The authors identify dataset quality and hyperparameter settings as validity threats. Their evaluation uses two human-curated benchmarks and examines the effect of changing τ, while noting that other settings could alter performance.

8 Threats to Validity

  • External validity is assessed on HumanEval and MHPP, whose human-curated problems span varied programming tasks; PassRate checks generated programs against test cases.The experiments also use base models from different LLM families and sizes.Source: §8 · pp. 9–10
  • Different hyperparameter settings could produce performance variation, although the experiments keep k, m, and t the same across approaches.The authors characterize hyperparameter choice as a minor threat.Source: §8 · p. 10
  • Varying τ is used to examine the upper bound of the method and how the portion of code using CoT affects final results.Source: §8 · p. 10

9 Conclusion

The conclusion presents UnCert-CoT as an uncertainty-aware decoding method that applies CoT at points requiring reasoning. It reports improved code generation across varied LLMs and some robustness to hyperparameter settings.

9 Conclusion

  • UnCert-CoT dynamically adapts CoT to generate intermediate reasoning steps based on model uncertainty during code generation.Source: §9 · p. 10
  • The method uses Entropy-based and Probability Differential-based confidence measures to guide uncertainty-aware code generation.Source: §9 · p. 10
  • The experiments report that UnCert-CoT substantially outperforms state-of-the-art CoT methods across various LLMs and shows some robustness to hyperparameter settings.Source: §9 · p. 10
Loading 2503.15341v1…