# A two-head loss function for deep Average-K classification

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 14 of its 34 references that are in the map.
- First posted on arXiv 31 Mar 2023.

## Builds on (most cited first, 14 of 14)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [ImageNet Large Scale Visual Recognition Challenge](https://mlsanity.com/paper/1409.0575/) (2014)
- [Densely Connected Convolutional Networks](https://mlsanity.com/paper/1608.06993/) (2016)
- [An overview of gradient descent optimization algorithms](https://mlsanity.com/paper/1609.04747/) (2016)
- [Conformal Prediction: a Unified Review of Theory and New Challenges](https://mlsanity.com/paper/2005.07972/) (2020)
- [The Emerging Trends of Multi-Label Learning](https://mlsanity.com/paper/2011.11197/) (2020)
- [Multi-Label Learning from Single Positive Labels](https://mlsanity.com/paper/2106.09708/) (2021)
- [Smooth Loss Functions for Deep Top-k Classification](https://mlsanity.com/paper/1802.07595/) (2018)
- [On the Consistency of Top-k Surrogate Losses](https://mlsanity.com/paper/1901.11141/) (2019)
- [Top-k Multiclass SVM](https://mlsanity.com/paper/1511.06683/) (2015)
- [Set-valued classification -- overview via a unified framework](https://mlsanity.com/paper/2102.12318/) (2021)
- [Stochastic smoothing of the top-K calibrated hinge loss for deep imbalanced classification](https://mlsanity.com/paper/2202.02193/) (2022)
- [eProduct: A Million-Scale Visual Search Benchmark to Address Product Recognition Challenges](https://mlsanity.com/paper/2107.05856/) (2021)

## Abstract (arXiv)

Average-K classification is an alternative to top-K classification in which the number of labels returned varies with the ambiguity of the input image but must average to K over all the samples. A simple method to solve this task is to threshold the softmax output of a model trained with the cross-entropy loss. This approach is theoretically proven to be asymptotically consistent, but it is not guaranteed to be optimal for a finite set of samples. In this paper, we propose a new loss function based on a multi-label classification head in addition to the classical softmax. This second head is trained using pseudo-labels generated by thresholding the softmax head while guaranteeing that K classes are returned on average. We show that this approach allows the model to better capture ambiguities between classes and, as a result, to return more consistent sets of possible classes. Experiments on two datasets from the literature demonstrate that our approach outperforms the softmax baseline, as well as several other loss functions more generally designed for weakly supervised multi-label classification. The gains are larger the higher the uncertainty, especially for classes with few samples.