Uncertainty Weighted Gradients for Model Calibration

2025-03-26

AI Summary

Deep neural networks often suffer from poor calibration due to miscalibrated predictive probabilities; existing loss functions are hampered by optimization inconsistency and biased uncertainty estimation. To address this, we propose a gradient-aware uncertainty-weighted optimization framework. Our method is the first to incorporate the Brier Score into a gradient weighting mechanism, mitigating both uncertainty estimation bias and gradient scaling misalignment inherent in focal loss. Furthermore, we introduce a multi-logit uncertainty estimator and a unified loss formulation to enable fine-grained, sample-level uncertainty modeling. Extensive experiments across diverse architectures and benchmark datasets demonstrate substantial improvements in calibration performance: our approach achieves state-of-the-art results on key metrics including Expected Calibration Error (ECE), while preserving or improving classification accuracy.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Stochastic Optimization
• Computer Vision: Learning & Optimization for CV

Application Category

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

Abstract

Model calibration is essential for ensuring that the predictions of deep neural networks accurately reflect true probabilities in real-world classification tasks. However, deep networks often produce over-confident or under-confident predictions, leading to miscalibration. Various methods have been proposed to address this issue by designing effective loss functions for calibration, such as focal loss. In this paper, we analyze its effectiveness and provide a unified loss framework of focal loss and its variants, where we mainly attribute their superiority in model calibration to the loss weighting factor that estimates sample-wise uncertainty. Based on our analysis, existing loss functions fail to achieve optimal calibration performance due to two main issues: including misalignment during optimization and insufficient precision in uncertainty estimation. Specifically, focal loss cannot align sample uncertainty with gradient scaling and the single logit cannot indicate the uncertainty. To address these issues, we reformulate the optimization from the perspective of gradients, which focuses on uncertain samples. Meanwhile, we propose using the Brier Score as the loss weight factor, which provides a more accurate uncertainty estimation via all the logits. Extensive experiments on various models and datasets demonstrate that our method achieves state-of-the-art (SOTA) performance.

Problem

• Addressing model calibration in deep neural networks
• Improving uncertainty estimation in classification tasks
• Optimizing loss functions for better calibration performance

Innovation

• Uncertainty weighted gradients for calibration
• Brier Score as loss weight factor
• Reformulated optimization focusing uncertain samples