💻 computer science

Task Abstention for Large Language Models in Code Generation

This paper proposes a theoretically grounded, distribution-free method for Large Language Models to abstain from code generation tasks likely to produce hallucinations by assessing generation consistency through code execution outcomes, thereby enabling safer and more robust automated coding without relying on external test cases.

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

Published 2026-05-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Yanke Zhou, Yuhao Tan, Senrong Xu, Zenan Li, Yuan Yao, Taolue Chen, 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

The Core Problem: The "Overconfident Chef"

Imagine you hire a brilliant but overconfident chef (the Large Language Model or LLM) to cook a complex dish based on a vague recipe description. Sometimes, the chef is amazing and creates a perfect meal. But often, the chef gets confident and creates a dish that looks delicious but tastes terrible or is actually poisonous (this is called "hallucination").

Currently, most safety checks happen after the chef has already cooked the dish. They taste the food to see if it's good. If it's bad, they throw it away. But by then, the chef has already wasted time and ingredients, and you've been served a bad meal.

This paper asks a different question: Can we teach the chef to say, "I don't know how to cook this," before they even start chopping vegetables?

This is called Task Abstention. It's about knowing when to refuse a job to avoid making a mistake.

The Solution: CODEREFUSER

The authors built a system called CODEREFUSER. Think of it as a "Safety Manager" for the chef. This manager doesn't just look at the recipe; it runs a simulation to see if the chef is likely to succeed before giving the green light.

Here is how the Safety Manager works, broken down into three simple steps:

1. The "Practice Run" (Calibration Phase)

Before the Safety Manager can make decisions, it needs to learn the chef's limits.

The Analogy: Imagine the manager gives the chef 100 practice recipes. For each one, the chef tries to cook it 64 times.
The Twist: The manager doesn't just look at the final dish. It also asks the chef to invent a "taste test" (test cases) for each recipe.
The Problem: Sometimes the chef is so bad at inventing taste tests that the test itself is broken (e.g., asking for a negative number of eggs).
The Fix (Sample-Test Dual Filtering): The manager uses a clever trick called Sample-Test Dual Filtering. It looks at the chef's 64 attempts. If the chef's "taste test" causes 64 different, chaotic results (some explode, some crash), the manager realizes, "Hey, this test is broken, not the chef!" It throws out the bad test and keeps the good ones. This ensures the manager isn't fooled by the chef's own confusion.

2. Setting the Rules (The "Risk Tolerance")

Once the manager has practiced, it sets the rules for the real world.

The Goal: The manager wants to ensure that if it says "Go ahead and cook," there is a very high chance the chef will succeed.
The Math (Simplified): The manager uses a statistical method called "Multiple Hypothesis Testing." Imagine the manager has a list of 1,000 different "stop signs" (thresholds). It tests them all against the practice data to find the specific signs that guarantee the chef won't fail more than, say, 20% of the time.
The Result: The manager creates a strict rulebook: "If the chef's confidence is below X, or if the test results are too messy, STOP."

3. The Real Job (Testing Phase)

Now, a real customer orders a dish.

The chef tries to cook it a few times (generates code samples).
The chef invents a new taste test for this specific dish.
The Safety Manager runs the Sample-Test Dual Filtering again to clean up any broken tests.
The Manager checks the rulebook.
- Scenario A: The chef's attempts are all consistent and pass the tests. The Manager says, "Admit: Go ahead, you can cook this."
- Scenario B: The chef's attempts are all over the place, or the tests are too confusing. The Manager says, "Abstain: I'm sorry, I cannot let you cook this. It's too risky."

Why This is Better Than Other Methods

The paper compares their method to other ways of checking for mistakes:

Static Methods (The "Grammar Police"): These methods just read the code to see if it looks right, like checking spelling. The paper shows this is useless for code because a sentence can be spelled perfectly but still make no sense (e.g., "The blue color ate the number").
Old Execution Methods: These try to run the code but don't clean up the bad tests. They often get confused by the chef's own bad test questions and refuse to let the chef cook even when they could have succeeded.

The Result: CODEREFUSER is much better at knowing when to say "I don't know." In tests, it improved the ability to spot dangerous tasks by about 26% compared to the best existing methods.

The Bottom Line

This paper introduces a way to make AI coders safer. Instead of waiting for the AI to make a mistake and then fixing it, this system acts like a smart supervisor that knows when the AI is in over its head. It forces the AI to say, "I don't know," before it generates any code, preventing the creation of broken or dangerous software.

Key Takeaway: It's better for an AI to admit it can't do a task than to confidently do it wrong. CODEREFUSER teaches the AI how to make that admission safely and reliably.

Technical Summary: Task Abstention for Large Language Models in Code Generation

Problem Definition

The paper addresses the critical issue of hallucination in Large Language Models (LLMs) applied to code generation. While existing research largely focuses on sample hallucination (detecting individual code snippets that fail to execute or meet requirements), this work identifies a more fundamental source of error: task-level hallucination. This occurs when an LLM is inherently unlikely to solve a specific coding task due to ambiguous prompts or the model's inherent limitations, regardless of how many samples it generates.

The authors formalize the Task Abstention problem: determining whether an LLM should refuse to answer a specific code generation prompt to avoid producing functionally incorrect code. The goal is to define a refusal function $r(x)$ that identifies tasks where the probability of generating a correct solution is low, specifically when the metric $H@k$ (the probability that $k$ randomly generated samples are all incorrect) exceeds a user-defined risk tolerance $\alpha$ .

Methodology: CODEREFUSER

The proposed approach, CODEREFUSER, is built upon the Learn Then Test (LTT) framework, utilizing multiple hypothesis testing to provide rigorous, distribution-free theoretical guarantees. The method operates in two phases:

1. Calibration Phase

Calibration Set Construction: Unlike standard NLP tasks, the calibration set consists of task prompts paired with oracle test cases (ground truth).
Admission Risk Definition: The risk is defined as the probability that a task is admitted (not refused) but the LLM fails to generate a correct solution within $k$ attempts ( $H@k$ ).
Score Functions: To handle the syntactic diversity of semantically equivalent code, the authors propose execution-based score functions rather than static linguistic metrics. These rely on clustering code samples based on their execution outputs against test cases:
- Cluster Ratio (CR): Measures the proportion of samples belonging to the dominant semantic cluster (model consensus).
- Semantic Entropy (SE): Quantifies the uncertainty of the task by measuring the diversity of the resulting cluster distribution.
Threshold Selection: Using multiple hypothesis testing, the method identifies a set of valid thresholds ( $\Lambda_{valid}$ ) that guarantee the admission risk remains below the tolerance $\alpha$ with high confidence ( $1-\delta$ ).

2. Testing Phase

Dual Generation: For a new test prompt, the LLM generates both code solutions and corresponding test cases.
Sample-Test Dual Filtering (STDF): A novel mechanism to address the issue that LLM-generated test cases may be invalid or flawed.
- Error Rate Filtering: Removes test cases that cause widespread execution failures across the sample population.
- Output Diversity Filtering: Removes test cases that induce high semantic entropy (indicating undefined behavior or invalid inputs) while limiting the pruning ratio to avoid discarding valid tests.
Decision Making: The refined test cases are used to cluster the code samples. The abstention decision is made by applying the calibrated thresholds to the Cluster Ratio or Semantic Entropy scores. If the criteria are not met, the model abstains.

Key Contributions

Formalization of Task Abstention: The paper introduces and formalizes the concept of task-level abstention in code generation, distinguishing it from sample-level hallucination detection.
Theoretically Guaranteed Framework: It proposes a method grounded in multiple hypothesis testing that provides a rigorous, distribution-free guarantee that the risk of accepting a "bad" task is controlled below a specified threshold.
Execution-Based Scoring with STDF: The authors introduce a scoring scheme based on runtime execution outcomes and a Sample-Test Dual Filtering (STDF) mechanism. This allows the system to operate without reliance on oracle test cases during the testing phase while mitigating the noise introduced by model-generated tests.
Minimal Reliance on Oracles: The approach requires oracle test cases only during the calibration phase, making it practical for deployment where ground-truth tests are unavailable for every new task.

Experimental Results

The authors evaluated CODEREFUSER on the HumanEval and MBPP benchmarks using four open-source code LLMs (DeepSeek-Coder, Qwen2.5-Coder, CodeLlama, and WizardCoder).

Performance: CODEREFUSER significantly outperformed existing baselines, including static methods (PPL, NLI, CLM) and execution-based detectors (CodeHalu, Execution). On average, it achieved a 26.5% absolute improvement in abstention precision compared to the best existing competitor.
Static vs. Execution: Experiments demonstrated that static methods (even with quadrupled sample sizes) failed to match the performance of execution-based methods, confirming that running code is necessary for accurate risk estimation in this domain.
Transferability: The method showed robust transferability, maintaining high performance when calibration was performed on MBPP and testing on HumanEval, despite the distribution shift.
Risk Control: Empirical results confirmed the theoretical guarantees, showing that the admission risk was successfully controlled below the target tolerance $\alpha$ in over 90% of cases. Furthermore, CODEREFUSER achieved this risk control with a significantly lower abstention rate (refusing fewer valid tasks) compared to baseline methods like PPL.

Significance and Claims

The paper claims to represent progress toward provably correct LLM code generation. By enabling models to confidently say "I don't know" for tasks they are unlikely to solve, the approach provides a reliable mechanism for safer and more robust code generation.

The authors modestly acknowledge limitations:

Distribution Shift: The theoretical guarantee relies on the assumption that calibration and test data are independent and identically distributed (i.i.d.). While experiments show robustness to moderate shifts, extreme real-world deviations could weaken guarantees.
Determinism: The method assumes tasks have deterministic solutions. It may not apply directly to non-deterministic or interactive software systems.
Language Scope: Current evaluation is limited to Python; effectiveness in other languages (e.g., C, Java) with different memory or pointer characteristics requires further evaluation.

The work does not propose new applications but rather a foundational safety layer that can be integrated into existing code generation workflows to mitigate the risk of hallucination.