Calibration-Aware Policy Optimization for Reasoning LLMs

2026-04-14

AI Summary

This work addresses the calibration degradation commonly observed in large language models during policy optimization, where overconfidence leads to lower perplexity on incorrect answers than on correct ones, undermining reliability. The authors propose Calibration-Aware Policy Optimization (CAPO), a framework that reveals, for the first time, how GRPO-style algorithms suffer from poor calibration due to neglecting uncertainty. CAPO integrates a logistic AUC surrogate loss, an uncertainty-aware advantage estimator, and a noise masking mechanism to jointly enhance both accuracy and calibration. Experiments demonstrate that CAPO-1.5B achieves up to a 15% improvement in calibration on multiple mathematical reasoning benchmarks while matching or surpassing GRPO in accuracy, with further gains of up to 5% on reasoning extrapolation tasks. In settings allowing abstention, CAPO attains a Pareto-optimal trade-off between accuracy and coverage.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Stochastic Optimization
• Natural Language Processing: Safety and Robustness

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Group Relative Policy Optimization (GRPO) enhances LLM reasoning but often induces overconfidence, where incorrect responses yield lower perplexity than correct ones, degrading relative calibration as described by the Area Under the Curve (AUC). Existing approaches either yield limited improvements in calibration or sacrifice gains in reasoning accuracy. We first prove that this degradation in GRPO-style algorithms stems from their uncertainty-agnostic advantage estimation, which inevitably misaligns optimization gradients with calibration. This leads to improved accuracy at the expense of degraded calibration. We then propose Calibration-Aware Policy Optimization (CAPO). It adopts a logistic AUC surrogate loss that is theoretically consistent and admits regret bound, enabling uncertainty-aware advantage estimation. By further incorporating a noise masking mechanism, CAPO achieves stable learning dynamics that jointly optimize calibration and accuracy. Experiments on multiple mathematical reasoning benchmarks show that CAPO-1.5B significantly improves calibration by up to 15% while achieving accuracy comparable to or better than GRPO, and further boosts accuracy on downstream inference-time scaling tasks by up to 5%. Moreover, when allowed to abstain under low-confidence conditions, CAPO achieves a Pareto-optimal precision-coverage trade-off, highlighting its practical value for hallucination mitigation.

Problem

• calibration
• overconfidence
• reasoning LLMs
• policy optimization
• AUC

Innovation

• Calibration-Aware Policy Optimization
• relative calibration
• AUC surrogate loss
• uncertainty-aware advantage estimation
• noise masking