Differentiable Top-k Classification Learning

Felix Petersen, Hilde Kuehne, Christian Borgelt, Oliver Deussen

Keywords: top-k, top-5, imagenet

Abstract: The top-k classification accuracy is one of the core metrics in machine learning. Here, k is conventionally a positive integer, such as 1 or 5. In this work, we relax this assumption and propose to draw k from a probability distribution for training. Combining this with recent advances in differentiable sorting and ranking, we propose a new family of differentiable top-k cross-entropy classification losses. We find that relaxing k does not only produce better top-5 accuracies, but also makes models more robust, which leads to top-1 accuracy improvements. When fine-tuning publicly available ImageNet models, we achieve a new state-of-the-art on ImageNet for publicly available models with an 88.36% top-1 and a 98.71% top-5 accuracy.

Community Implementations: [ 1 code implementation](https://www.catalyzex.com/paper/arxiv:2206.07290/code)