Fine-grained Approaches for Confidence Calibration of LLMs in Automated Code Revision

1. Problem Statement

Large Language Models (LLMs) are increasingly used for Automated Code Revision (ACR) tasks, such as program repair, vulnerability repair, and code refinement. However, these models are prone to generating incorrect outputs, which can disrupt developer workflows. A critical issue is that post-trained LLMs do not inherently produce well-calibrated confidence scores; their predicted probabilities of correctness often do not match their empirical accuracy.

Existing calibration methods, such as global Platt-scaling applied to sequence-level confidence scores (e.g., average token probability), have proven effective for general code generation but are unreliable for ACR tasks. The authors hypothesize that these conventional methods fail because:

1. Coarse-grained nature: ACR correctness is often determined by specific local edit decisions rather than the entire sequence. Sequence-level scores dilute these critical signals with high-probability tokens from unchanged code.
2. Sample-dependent miscalibration: The relationship between uncalibrated scores and true correctness varies significantly across different samples (heterogeneity), which a single global calibrator cannot capture.

2. Methodology

The study proposes a two-pronged fine-grained approach to address these issues:

A. Fine-Grained Confidence Scores

Instead of aggregating probabilities over the entire output sequence, the authors propose three intrinsic, single-inference metrics that focus on local edit decisions:

1. Minimum Token Probability ($P_{min}$): The lowest probability assigned to any token in the sequence. This captures the "weakest link" in the revision, as a single critical error can invalidate the code.
2. Lowest-K Token Probability ($P_{low-K}$): The arithmetic mean of the $K$ tokens with the lowest probabilities. $K$ is dynamically selected using the Kneedle algorithm to identify the "elbow" where probabilities transition from high to low, filtering out outliers.
3. Attention-Weighted Uncertainty ($P_{attn-\omega}$): An extension of $P_{low-K}$ that weights tokens based on their attention mass (influence on subsequent tokens). Tokens with low probability but high downstream influence are weighted more heavily, as errors there are more catastrophic.

B. Local Platt-Scaling

The authors introduce Local Platt-scaling to replace the conventional global approach:

• Global Platt-scaling: Trains a single logistic regression model (calibrator) for all samples in a task.
• Local Platt-scaling:
  1. Clustering: Uses HDBSCAN to cluster input/output embeddings (concatenated with uncalibrated confidence scores) into distinct groups based on similarity.
  2. Local Calibration: Trains a separate logistic regression calibrator for each cluster.
  3. Inference: Assigns new samples to their nearest cluster and applies the corresponding local calibrator. Outliers use a backoff strategy (either the global calibrator or raw score).

C. Experimental Setup

• Tasks: Automated Program Repair (DCF-Bug), Vulnerability Repair (DCF-Vul), and Automated Code Refinement (CR-Trans).
• Models: 14 open-source LLMs (Llama-3.1, CodeLlama, Qwen2.5, DeepSeek) ranging from 7B to 72B parameters.
• Metrics:
  • Correctness: Exact Match (EM), Edit Progress (EP), Checks Passed (CP).
  • Calibration: Expected Calibration Error (ECE), Brier Score (B), and Bin Coverage (BC).

3. Key Contributions

1. First Fine-Grained Calibration for ACR: The study is the first to propose and validate fine-grained confidence scores for ACR tasks, demonstrating that local edit decisions contain the most salient uncertainty signals.
2. Introduction of Local Platt-Scaling: It proposes a sample-aware calibration method that corrects for non-uniform, sample-dependent miscalibration patterns that global methods miss.
3. Comprehensive Empirical Analysis: The paper provides a rigorous evaluation across 3 tasks, 3 correctness metrics, and 14 models, establishing new baselines for confidence calibration in software engineering.

4. Key Results

RQ1: Separation of Correct vs. Incorrect Revisions

• Token Skewness: Analysis revealed strong negative skewness in token probabilities for ACR tasks (most tokens are high probability, few are low).
• Separation Power: Fine-grained scores (especially Minimum Token Probability) showed significantly higher Wasserstein Distance ($W_1$) and Kendall's $\tau$ between correct and incorrect revisions compared to sequence-level scores. This indicates they better separate correct from incorrect outputs.

RQ2: Effectiveness of Global Platt-Scaling with Fine-Grained Scores

• Performance: Fine-grained scores consistently outperformed sequence-level scores across all tasks.
• Bin Collapse: Sequence-level scores frequently suffered from "single bin collapse" (BC=1), where the calibrator failed to learn a meaningful mapping, resulting in uniform confidence scores. Fine-grained scores maintained high bin coverage.
• Best Metric: Minimum Token Probability was the top performer across all tasks and models, achieving the lowest ECE and highest bin coverage.

RQ3: Impact of Local Platt-Scaling

• Program & Vulnerability Repair (DCF-Bug/Vul): Global Platt-scaling with fine-grained scores was generally sufficient (low ECE). Local Platt-scaling offered only marginal improvements (ECE reduction of ~0.01–0.07) but significantly improved bin coverage in vulnerability repair.
• Code Refinement (CR-Trans): This task exhibited high error heterogeneity. Global Platt-scaling failed to achieve acceptable calibration (ECE > 0.14). Local Platt-scaling was essential, reducing ECE by up to 0.16 and enabling reliable decision-making.
• Optimal Strategy: The combination of Minimum Token Probability + Local Platt-Scaling yielded the best results, particularly for code refinement.

5. Significance and Practical Recommendations

The study provides actionable guidelines for deploying LLMs in software engineering:

• For Program/Vulnerability Repair: If inference latency is a priority, use Minimum Token Probability with Global Platt-scaling. It offers sufficiently low calibration error with minimal overhead. If the absolute lowest error is required, Local Platt-scaling can be added at the cost of latency.
• For Automated Code Refinement: Local Platt-scaling with Minimum Token Probability is mandatory. Without it, confidence scores are severely miscalibrated, leading to unreliable developer decisions.
• Latency Trade-off: While Local Platt-scaling requires embedding generation and clustering, it remains significantly faster than alternative uncertainty quantification methods like stochastic sampling or P(True).

Conclusion

This paper demonstrates that the "coarse-grained" nature of traditional confidence calibration is ill-suited for automated code revision. By shifting focus to local edit decisions (via fine-grained scores) and accounting for sample heterogeneity (via local Platt-scaling), the authors achieve significantly more trustworthy confidence estimates. This work bridges the gap between LLM capabilities and developer trust, enabling safer and more efficient AI-assisted software development.