# Beyond Probability Partitions: Calibrating Neural Networks with Semantic Aware Grouping

arXiv: https://arxiv.org/abs/2306.04985 · HTML: https://mlsanity.com/paper/2306.04985/

## Key facts

- Cited by 10 papers in the map (data through 14 Oct 2026).
- Builds on 44 of its 60 references that are in the map.
- First posted on arXiv 8 Jun 2023.

## Builds on (most cited first, 30 of 44)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Very Deep Convolutional Networks for Large-Scale Image Recognition](https://mlsanity.com/paper/1409.1556/) (2014)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [Densely Connected Convolutional Networks](https://mlsanity.com/paper/1608.06993/) (2016)
- [On Calibration of Modern Neural Networks](https://mlsanity.com/paper/1706.04599/) (2017)
- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [A Survey on Bias and Fairness in Machine Learning](https://mlsanity.com/paper/1908.09635/) (2019)
- [Concrete Problems in AI Safety](https://mlsanity.com/paper/1606.06565/) (2016)
- [ShuffleNet: An Extremely Efficient Convolutional Neural Network for Mobile Devices](https://mlsanity.com/paper/1707.01083/) (2017)
- [End to End Learning for Self-Driving Cars](https://mlsanity.com/paper/1604.07316/) (2016)
- [Aleatoric and Epistemic Uncertainty in Machine Learning: An Introduction to Concepts and Methods](https://mlsanity.com/paper/1910.09457/) (2019)
- [A Review of Uncertainty Quantification in Deep Learning: Techniques, Applications and Challenges](https://mlsanity.com/paper/2011.06225/) (2020)
- [A Survey of Uncertainty in Deep Neural Networks](https://mlsanity.com/paper/2107.03342/) (2021)
- [A Simple Baseline for Bayesian Uncertainty in Deep Learning](https://mlsanity.com/paper/1902.02476/) (2019)
- [Bayesian Deep Learning and a Probabilistic Perspective of Generalization](https://mlsanity.com/paper/2002.08791/) (2020)
- [Accurate Uncertainties for Deep Learning Using Calibrated Regression](https://mlsanity.com/paper/1807.00263/) (2018)
- [On Fairness and Calibration](https://mlsanity.com/paper/1709.02012/) (2017)
- [Revisiting the Calibration of Modern Neural Networks](https://mlsanity.com/paper/2106.07998/) (2021)
- [Beyond temperature scaling: Obtaining well-calibrated multiclass probabilities with Dirichlet calibration](https://mlsanity.com/paper/1910.12656/) (2019)
- [Calibrating Deep Neural Networks using Focal Loss](https://mlsanity.com/paper/2002.09437/) (2020)
- [Verified Uncertainty Calibration](https://mlsanity.com/paper/1909.10155/) (2019)
- [On Mixup Training: Improved Calibration and Predictive Uncertainty for Deep Neural Networks](https://mlsanity.com/paper/1905.11001/) (2019)
- [Mix-n-Match: Ensemble and Compositional Methods for Uncertainty Calibration in Deep Learning](https://mlsanity.com/paper/2003.07329/) (2020)
- [Evaluating model calibration in classification](https://mlsanity.com/paper/1902.06977/) (2019)
- [Improving model calibration with accuracy versus uncertainty optimization](https://mlsanity.com/paper/2012.07923/) (2020)
- [Calibration of Neural Networks using Splines](https://mlsanity.com/paper/2006.12800/) (2020)
- [Calibrated Learning to Defer with One-vs-All Classifiers](https://mlsanity.com/paper/2202.03673/) (2022)
- [Intra Order-preserving Functions for Calibration of Multi-Class Neural Networks](https://mlsanity.com/paper/2003.06820/) (2020)
- [Better Uncertainty Calibration via Proper Scores for Classification and Beyond](https://mlsanity.com/paper/2203.07835/) (2022)

## Built on by (most cited first, 10 of 10)

- [Feature Clipping for Uncertainty Calibration](https://mlsanity.com/paper/2410.19796/) (2024)
- [Consistency Calibration: Improving Uncertainty Calibration via Consistency among Perturbed Neighbors](https://mlsanity.com/paper/2410.12295/) (2024)
- [Sample Margin-Aware Recalibration of Temperature Scaling](https://mlsanity.com/paper/2506.23492/) (2025)
- [SoC: Semantic Orthogonal Calibration for Test-Time Prompt Tuning](https://mlsanity.com/paper/2601.08617/) (2026)
- [Quantile Adaptive Temperature Scaling for Confidence Calibration](https://mlsanity.com/paper/2606.21749/) (2026)
- [Towards Precise Prediction Uncertainty in GNNs: Refining GNNs with Topology-grouping Strategy](https://mlsanity.com/paper/2412.14223/) (2024)
- [FedLAS: Feature-Modulated Bidirectional Label Smoothing for Neural Network Calibration](https://mlsanity.com/paper/2606.28654/) (2026)
- [Deep Probability Segmentation: Are segmentation models probability estimators?](https://mlsanity.com/paper/2409.12535/) (2024)
- [MARCO: Click-Intent Decomposition for Calibrated Ads Conversion Prediction](https://mlsanity.com/paper/2608.10562/) (2026)
- [Signal-Routed Temperature Scaling: Low-Capacity Risk-Conditioned Calibration for Small Validation Budgets](https://mlsanity.com/paper/2609.38936/) (2026)

## Abstract (arXiv)

Research has shown that deep networks tend to be overly optimistic about their predictions, leading to an underestimation of prediction errors. Due to the limited nature of data, existing studies have proposed various methods based on model prediction probabilities to bin the data and evaluate calibration error. We propose a more generalized definition of calibration error called Partitioned Calibration Error (PCE), revealing that the key difference among these calibration error metrics lies in how the data space is partitioned. We put forth an intuitive proposition that an accurate model should be calibrated across any partition, suggesting that the input space partitioning can extend beyond just the partitioning of prediction probabilities, and include partitions directly related to the input. Through semantic-related partitioning functions, we demonstrate that the relationship between model accuracy and calibration lies in the granularity of the partitioning function. This highlights the importance of partitioning criteria for training a calibrated and accurate model. To validate the aforementioned analysis, we propose a method that involves jointly learning a semantic aware grouping function based on deep model features and logits to partition the data space into subsets. Subsequently, a separate calibration function is learned for each subset. Experimental results demonstrate that our approach achieves significant performance improvements across multiple datasets and network architectures, thus highlighting the importance of the partitioning function for calibration.
