🤖 AI

Uncertainty Quantification for LLM-based Code Generation

This paper introduces RisCoSet, a novel framework that leverages multiple hypothesis testing to construct risk-controlling prediction sets for LLM-based code generation, effectively overcoming the limitations of existing methods by accommodating multiple valid outputs and significantly reducing unnecessary code removal while maintaining high confidence in correctness.

Original authors: Senrong Xu, Yuhao Tan, Yanke Zhou, Guangyuan Wu, Zenan Li, Yuan Yao, Taolue Chen, Feng Xu, Xiaoxing Ma

Published 2026-05-13

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Senrong Xu, Yuhao Tan, Yanke Zhou, Guangyuan Wu, Zenan Li, Yuan Yao, Taolue Chen, Feng Xu, Xiaoxing Ma

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you ask a very talented but occasionally overconfident AI to write a computer program for you. Sometimes, the AI writes perfect code. Other times, it "hallucinates," mixing in lines that look right but actually break the program.

The problem is: How do you know which parts of the AI's code are safe to keep, and which parts are dangerous?

This paper introduces a new method called RISCOSET to solve this. Think of it as a "Safety Net" for AI-generated code.

The Problem with Old Methods

Previously, researchers tried to use a method called PAC (which stands for "Probably Approximately Correct").

The Analogy: Imagine you are trying to find a specific key in a giant, messy drawer full of keys. The old method (PAC) had a strict rule: "You can only look at keys that are smaller than the last one you picked."
The Flaw: This rule was too rigid. It forced the researchers to throw away too many potentially good keys (or in this case, lines of code) just to follow the rule. It also assumed there was only one correct answer, but in coding, there are often many different ways to write a program that does the exact same thing.

The New Solution: RISCOSET

The authors propose RISCOSET, which uses a smarter approach called LTT (Learn Then Test).

1. The "Partial Program" (The Skeleton)
Instead of trying to guess the entire perfect program, RISCOSET builds a "skeleton" or a "partial program."

The Analogy: Imagine the AI writes a story, but you aren't sure about the ending. RISCOSET doesn't delete the whole story. Instead, it highlights the uncertain sentences and removes only those specific parts, leaving you with a solid framework (a partial story) that you know is safe.
The Goal: This skeleton is guaranteed (with high statistical confidence) to be a part of a correct solution. You can then fill in the missing pieces later.

2. The "Multi-Label" Advantage
The old method assumed there was only one "correct" answer. RISCOSET understands that in coding, there are many valid solutions.

The Analogy: If you ask for a recipe for a cake, there isn't just one "right" way to do it. You can use butter or oil; you can bake it in a round pan or a square one. RISCOSET accepts that there are many "correct" versions, so it doesn't throw away good code just because it looks slightly different from the first guess.

3. The "Selective Execution" (Saving Time and Money)
To make sure their "skeletons" are safe, the system needs to test the code. But testing every single piece of code is slow and expensive (like running a full car engine test for every single bolt).

The Analogy: RISCOSET uses a "Selective Execution" strategy. It's like a quality inspector who looks at a product. If the product looks very clean and high-quality (low uncertainty), the inspector skips the full test. If it looks messy (high uncertainty), the inspector runs the full test.
The Result: This saves a huge amount of time and computing power while still keeping the error rate very low.

What Did They Find?

The researchers tested this on three different AI models and three different coding datasets.

Less Waste: Compared to the best existing methods, RISCOSET removed 24.5% fewer lines of code. This means you get to keep more of the AI's useful work.
Safety First: Even though they kept more code, they still maintained the same high level of safety guarantees. The "skeleton" they built was just as likely to contain a correct solution as the old methods, but with much less unnecessary deletion.

In Summary

RISCOSET is a new tool that helps us trust AI-generated code better. Instead of throwing away large chunks of code because we are afraid of mistakes, it carefully carves out a safe "skeleton" that is guaranteed to be part of a working solution. It's smarter, less wasteful, and respects the fact that there are many ways to write good code.

Technical Summary: RISCOSET for Uncertainty Quantification in LLM-based Code Generation

Problem Statement

Large Language Models (LLMs) are increasingly used for code generation but suffer from hallucinations and the generation of toxic or biased outputs, limiting their deployment in high-stakes applications. While Conformal Prediction (CP) offers a framework for uncertainty quantification with distribution-free guarantees, adapting it to structured generation tasks like code is challenging.

Existing approaches, such as PAC (Probably Approximately Correct) prediction sets for code (Khakhar et al., 2023), face two critical limitations:

Monotonicity Constraints: They rely on a strong monotonicity assumption regarding risk and prediction set size. This restricts the space of admissible partial programs, often leading to suboptimal solutions that remove too many code nodes.
Single-Label Assumption: They are designed for single-label classification. However, code generation is inherently a multi-label problem, where multiple syntactically different programs can possess identical, correct semantics. Existing methods fail to accommodate this multiplicity, severely limiting the candidate space.

Methodology: RISCOSET

The authors propose RISCOSET, a method that constructs risk-controlling prediction sets for LLM-based code generation by reframing the problem as a multi-label classification task and leveraging the Learn Then Test (LTT) framework (Angelopoulos et al., 2025).

Core Components

Multi-Label Reframing & Structured Risk:
Instead of enforcing monotonicity, RISCOSET defines a structured risk function based on semantic correctness. For a given input, the goal is to produce a prediction set (represented as a partial program or AST subtree) that contains at least one correct program from an enhanced label set $\tilde{Y}$ . The empirical risk is defined as the fraction of calibration instances where the prediction set fails to intersect with the set of correct programs.
Optimization for Partial Programs:
To generate the prediction set, the method solves an optimization problem to prune nodes from the Abstract Syntax Tree (AST) of the generated code.
- Objective: Retain as many nodes as possible while ensuring the total uncertainty (negative log probability) of the remaining partial program does not exceed a threshold $\lambda$ .
- Constraints: The solution must satisfy tree connectivity (if a parent is removed, children must be removed) and a limit on the number of permissible subtree removals ( $t_{max}$ ). Crucially, this optimization removes the strict monotonicity constraint found in PAC methods, allowing for a larger feasible label space.
Multiple Hypothesis Testing (LTT):
RISCOSET employs the LTT framework to select the optimal parameter $\hat{\lambda}$ that controls the risk.
- It treats the selection of $\lambda$ as a multiple-hypothesis testing problem.
- For a set of candidate $\lambda$ values, it computes $p$ -values based on the empirical structured risk.
- Using a Family-Wise Error Rate (FWER) controlling algorithm, it identifies a valid set of parameters $\Lambda_{valid}$ that guarantees the risk remains below a user-specified tolerance $\alpha$ with high confidence ( $1-\delta$ ).
Selective Execution Strategy:
Since code generation is multi-label, the calibration process requires sampling multiple candidate programs and verifying them via test-case execution, which is computationally expensive. To mitigate this, RISCOSET introduces a selective execution strategy:
- It uses uncertainty scores (e.g., perplexity) to decide which sampled programs require test execution.
- Programs with low uncertainty are accepted as correct without execution, while high-uncertainty programs are verified.
- This strategy uses importance sampling and Hoeffding's inequality to estimate an error upper bound, ensuring that the label error rate remains within a controlled bound $\epsilon$ while significantly reducing computational costs.

Key Contributions

RISCOSET Framework: The first application of the LTT framework to structured code generation, reframing prediction set construction as a multi-label problem with weaker constraints than PAC-based approaches.
Structured Risk Optimization: A novel optimization formulation for generating partial programs (AST subtrees) that maximizes retained information while satisfying risk control guarantees, free from restrictive monotonicity assumptions.
Efficient Calibration: A selective execution strategy that balances the trade-off between risk control and the computational cost of label verification, supported by non-asymptotic error bounds.
Empirical Validation: Comprehensive experiments demonstrating the method's effectiveness across three LLMs (Deepseek-Coder, Qwen2.5-Coder, Llama3.1) and three datasets (HumanEval, MBPP, APPS).

Experimental Results

The authors evaluated RISCOSET against state-of-the-art baselines, including the PAC prediction set approach and a greedy variant.

Node Removal Reduction: RISCOSET significantly reduces the number of AST nodes removed compared to baselines while maintaining the same level of risk control. For instance, on the MBPP dataset with $\alpha=0.15$ , the method reduced code removal by 24.5% compared to the best competitor.
Coverage Guarantees: The constructed prediction sets consistently achieved coverage rates exceeding the target bound ( $1-\alpha$ ), validating the theoretical risk-controlling guarantees.
Model Scalability: The method performed well across both code-specialized models (Qwen2.5-Coder) and general-purpose models (Llama3.1), showing that performance improves with more powerful generative models.
Selective Execution Efficiency: The selective execution strategy successfully reduced the proportion of programs requiring test-case execution (e.g., avoiding execution for ~30% of programs at a label error rate of $\epsilon=0.1$ ) while maintaining coverage above the target threshold.

Significance and Claims

The paper claims to be the first work to introduce the LTT framework into structured code generation tasks. Its primary significance lies in addressing the limitations of existing PAC-based methods by:

Eliminating the need for strong monotonicity assumptions, thereby expanding the space of candidate programs and reducing unnecessary code removal.
Naturally accommodating the multi-label nature of code generation (multiple valid solutions for a single task).
Providing rigorous statistical guarantees for the correctness of the generated code sets without relying on the independence assumptions that fail in structured domains like code.

The authors conclude that while their focus is on code generation, the proposed method is adaptable to other structured prediction problems where uncertainty quantification is critical. They acknowledge the computational overhead of the calibration phase but emphasize that the selective execution strategy effectively mitigates this cost.