Set Learning for Accurate and Calibrated Models

TL;DR

This paper introduces odd-$k$-out learning (OKO), a novel set-based training method that improves model accuracy and calibration by capturing data correlations, especially in limited or imbalanced datasets, without extra calibration tuning.

Contribution

The paper proposes OKO, a new set-based learning framework that enhances calibration and accuracy, with theoretical analysis and broad applicability.

Findings

OKO improves calibration even with hard labels.

OKO outperforms standard methods in limited data regimes.

No additional calibration tuning needed with OKO.

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.