Set-based training improves classifiers on imbalanced data sets
Using Set Learning to Improve Set-Valued Classifiers on Long-tailed Data
Machine LearningComputer Vision and Pattern Recognition
Summary
When computer programs try to recognize things, they are sometimes more confident about some categories than others, especially if some categories are rare. The paper studies ways to improve how these programs handle rare categories by training them with new methods that consider groups of examples together rather than one at a time. The authors adapted and improved a special training approach called odd-K-out loss to make it more efficient. Their changes helped the programs be better at identifying rare categories without making overly large guesses.
What this means in practice
- •For machine learning engineers: Improve classification models trained on imbalanced datasets by using set-based loss functions to enhance accuracy on minority classes without larger output sets.
- •For medical data teams: Enhance diagnostic prediction models that face rare condition classes by training with set-based learning methods to boost sensitivity on less common diseases.
Authors
Kevin Kasa, Scott C. Lowe, Graham W. Taylor
Abstract
Set-valued classification (SVC) offers a practical approach to communicating uncertainty while maintaining accuracy by predicting variable-sized sets per input rather than fixed-sized sets (e.g. top-1 or top-5). However, set-valued classifiers (SVCs) exhibit significantly degraded performance on long-tailed and imbalanced datasets, a setting for which they are ostensibly well-suited. This becomes apparent under macro-average evaluation, which weighs all classes equally. SVCs are typically formed during a post-hoc calibration stage, so little is known about how training dynamics affect their behaviour. In this work, we investigate how different neural network training strategies and loss functions impact downstream SVC performance on highly imbalanced data. In addition to methods such as label smoothing and class-balancing, we investigate whether loss functions that operate over sets of examples in a batch, rather than treating all examples as independent, lead to better downstream SVCs. We adapt the recently proposed odd-K-out (OKO) loss function, which learns from sets of training samples, improving its sample efficiency whilst introducing algorithmic modifications that reduce its computational overhead from 3.2x to just 1.36x that of standard cross-entropy training. Experiments on highly long-tailed datasets demonstrate our modified OKO approach significantly improves macro-accuracy across multiple SVC methods - including top-k, average-k, and conformal prediction - while maintaining competitive micro-accuracy. This improved performance on minority classes occurs without substantially increasing prediction set sizes, suggesting that carefully constructed set-based training can produce networks better suited to imbalanced data.