Localized Calibrated Uncertainty in Code Language Models

TL;DR

The paper tackles the challenge of localized, calibrated uncertainty in LLM-generated code by introducing a dataset of Minimal Intent Aligning Patches and evaluating multiple approaches to produce reliable line- and token-level confidences. It compares latent-probe based calibration, multisampling, and reflective prompting, demonstrating that small supervisor probes can achieve meaningful Brier Skill Scores (around $0.2$) on large MUTs, with multisampling delivering strong line- and token-level performance particularly on GitHub-scale tasks. Platt scaling emerges as a practical tool to improve calibration across methods, though some approaches (notably reflective prompting) require data scaling to reach parity with probing-based techniques. The work also explores generalization to non-code text, showing partial transfer to natural language hallucinations under probability scaling, and discusses practical considerations for deployment, cost, latency, and the potential for scalable AI oversight in code generation workflows.

Abstract

Large Language models (LLMs) can generate complicated source code from natural language prompts. However, LLMs can generate output that deviates from what the user wants, requiring supervision and editing. To support this process, we offer techniques to localize where generations might be misaligned from user intent. We first create a dataset of "Minimal Intent Aligning Patches" of repaired LLM generated programs. Each program uses test cases to verify correctness. After creating a dataset of programs, we measure how well various techniques can assign a well-calibrated probability to indicate which parts of code will be edited in a minimal patch (i.e., give a probability that corresponds with empirical odds it is edited). We compare white-box probing (where we propose a technique for efficient arbitrary-span querying), against black-box reflective and self-consistency based approaches. We find probes with a small supervisor model can achieve low calibration error and Brier Skill Score of approx 0.2 estimating edited lines on code generated by models many orders of magnitude larger. We discuss the generalizability of the techniques, and the connections to AI oversight and control, finding a probe trained only on code shows some signs of generalizing to natural language errors if new probability scaling is allowed.

Figures (3)

• Figure 1: Illustrating the basic flow. GPT-4o generates a flawed solution with one line incorrect. Our pipeline estimates the minimal patch. Then we take the diff which gives us labels at the token, line, and problem level. We illustrate sample confidences. From these we evaluate the quality of confidences via the Brier Score (MSE), which taking into account the underlying dataset base rate, becomes the Brier Skill Score (BSS), or ECE.
• Figure 2: Calibration reliability curves of the transfer of a probe trained on code data to HaluBench. Plots show a comparison of the actual probability compared to the predicted probability. The lines plot show quintiles of the data, as predictions are not evenly distributed between buckets.
• Figure 3: Factor breakdown exploring some of the design space when probing for local uncertainty. We gridsearch over options. Box plots show metrics accross different configurations, and help understand both the best achieved results (top of the range), but also how robust the selection of the parameter is to other choices (via the median and spread).