When to Answer and When to Defer: A Decision Framework for Reliable Code Predictions

TL;DR

This paper presents a unified framework for improving the reliability of code language models by enabling better uncertainty estimation, abstention, and external validation, supporting safer deployment.

Contribution

It introduces a novel integrated approach combining uncertainty, calibration, and tool-based abstention for code models, addressing limitations of existing methods.

Findings

Framework enables models to estimate correctness probabilities reliably.

Supports abstention and external validation to improve safety.

Applicable to both classification and generation tasks.

Abstract

Code language models are increasingly adopted for both understanding and generative tasks. Despite their success, these models frequently produce overconfident incorrect predictions and underconfident correct predictions, undermining their reliability in deployment. Practical deployment demands three capabilities: accurately estimating the likelihood of correctness, abstaining on uncertain predictions, and invoking external mechanisms to validate or repair abstained outputs. Existing calibration and uncertainty estimation methods, primarily developed for natural language tasks, do not readily transfer to code. Notably, post-hoc calibration techniques often reduce probability misalignment but fail to improve the ranking of predictions by correctness likelihood-a requirement for selective prediction under partial coverage. Furthermore, most approaches treat uncertainty as a passive indicator rather than an actionable signal. This work introduces a unified framework that integrates uncertainty estimation, model calibration, and tool-based abstention handling for code models. The proposed design enables models to assign reliable correctness probabilities, abstain under uncertainty, and invoke lightweight program analysis procedures to process abstained cases. By combining these components within a single deployment-oriented workflow, this framework supports risk-aware, coverage-controlled use of code models across both classification and generation settings.