Set Learning for Accurate and Calibrated Models

TL;DR

This work targets model miscalibration and overconfidence by proposing odd-$k$-out learning (OKO), a set-based training framework that minimizes cross-entropy over sets of size $k+2$ instead of individual examples. By aggregating logits across sets and using either a soft or hard loss, OKO captures inter-example correlations, yielding smoother, more reliable probability estimates without adding calibration hyperparameters or changing the architecture. The authors provide theoretical calibration insights, inclucing a per-datapoint excess-confidence score and an entropy-based relative cross-entropy measure, and demonstrate substantial empirical gains in both accuracy and calibration, especially in low-data and heavy-tailed settings across MNIST, Fashion-MNIST, and CIFAR benchmarks. OKO remains a general, plug-in training paradigm with linear training cost and preserves single-sample inference at test time, suggesting practical impact for safety-critical and data-scarce domains. The findings indicate that set-based objectives can align predictive accuracy with probabilistic calibration, reducing the need for post-hoc calibration pipelines and improving robustness to class imbalance.

Abstract

Model overconfidence and poor calibration are common in machine learning and difficult to account for when applying standard empirical risk minimization. In this work, we propose a novel method to alleviate these problems that we call odd-$k$-out learning (OKO), which minimizes the cross-entropy error for sets rather than for single examples. This naturally allows the model to capture correlations across data examples and achieves both better accuracy and calibration, especially in limited training data and class-imbalanced regimes. Perhaps surprisingly, OKO often yields better calibration even when training with hard labels and dropping any additional calibration parameter tuning, such as temperature scaling. We demonstrate this in extensive experimental analyses and provide a mathematical theory to interpret our findings. We emphasize that OKO is a general framework that can be easily adapted to many settings and a trained model can be applied to single examples at inference time, without significant run-time overhead or architecture changes.

Figures (16)

• Figure 1: A: OKO minimizes cross-entropy on sets of examples rather than on single examples and naturally yields smoothed logits after training. At inference time it can be applied to single examples without additional computational overhead. B: Expected calibration error as a function of the classification error. Each point in the graph represents the performance of a single seed; there are five for every number of training data points. For each dataset, every model was evaluated on the same test set. Dashed diagonal lines indicate a linear regression fit. Top: Uniform class distribution during training. Bottom: Heavy-tailed class distribution during training.
• Figure 2: Test set accuracy in % as a function of different numbers of data points used during training. Error bands depict 95% CIs and are computed over five random seeds for all training settings and methods. Top: Uniform class distribution during training. Bottom: Heavy-tailed class distribution.
• Figure 3: Reliability diagrams for balanced datasets. Confidence and accuracy scores were averaged over random seeds and the number of training data points. Dashed diagonal lines indicate perfect calibration.
• Figure 4: Here, we show the distribution of entropies of the predicted probability distributions for individual test data points across all heavy-tailed training settings partitioned into correct and incorrect predictions respectively.
• Figure 5: For different numbers of training data points, OKO achieves a substantially lower MAE for the average cross-entropy error between true and predicted class distributions and the average entropy of the predictions --- across both uniform and heavy-tailed class distributions during training.

Theorems & Definitions (17)

• Theorem 1
• Definition 1
• Lemma 1
• Lemma 2
• Proposition 1
• proof
• Proposition 2
• Proposition 3
• Lemma 3
• proof : Proof of Proposition \ref{['prop:risk-minima']}