Localized Calibrated Uncertainty in Code Language Models

TL;DR

This paper introduces techniques to localize and calibrate uncertainty in code language models, enabling better identification of parts of generated code likely to need editing, with implications for AI oversight.

Contribution

It presents a dataset of minimal intent aligning patches and compares probing techniques for well-calibrated uncertainty estimation in code generation.

Findings

Probes with small models achieve low calibration error.

Calibration techniques generalize somewhat to natural language errors.

Achieved Brier Skill Score of approx 0.2 on code edits.

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.