# Differentiable Top-k Classification Learning

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

## Key facts

- Cited by 21 papers in the map (data through 14 Oct 2026).
- Builds on 33 of its 43 references that are in the map.
- First posted on arXiv 15 Jun 2022.

## Builds on (most cited first, 30 of 33)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Scaling Up Visual and Vision-Language Representation Learning With Noisy Text Supervision](https://mlsanity.com/paper/2102.05918/) (2021)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Self-training with Noisy Student improves ImageNet classification](https://mlsanity.com/paper/1911.04252/) (2019)
- [Scaling Vision Transformers](https://mlsanity.com/paper/2106.04560/) (2021)
- [Revisiting Unreasonable Effectiveness of Data in Deep Learning Era](https://mlsanity.com/paper/1707.02968/) (2017)
- [Big Transfer (BiT): General Visual Representation Learning](https://mlsanity.com/paper/1912.11370/) (2019)
- [ImageNet-21K Pretraining for the Masses](https://mlsanity.com/paper/2104.10972/) (2021)
- [Exploring the Limits of Weakly Supervised Pretraining](https://mlsanity.com/paper/1805.00932/) (2018)
- [CoAtNet: Marrying Convolution and Attention for All Data Sizes](https://mlsanity.com/paper/2106.04803/) (2021)
- [Pervasive Label Errors in Test Sets Destabilize Machine Learning Benchmarks](https://mlsanity.com/paper/2103.14749/) (2021)
- [Meta Pseudo Labels](https://mlsanity.com/paper/2003.10580/) (2020)
- [Fast Differentiable Sorting and Ranking](https://mlsanity.com/paper/2002.08871/) (2020)
- [Differentiable Ranks and Sorting using Optimal Transport](https://mlsanity.com/paper/1905.11885/) (2019)
- [Stochastic Optimization of Sorting Networks via Continuous Relaxations](https://mlsanity.com/paper/1903.08850/) (2019)
- [Learning with Differentiable Perturbed Optimizers](https://mlsanity.com/paper/2002.08676/) (2020)
- [Differentiable Patch Selection for Image Recognition](https://mlsanity.com/paper/2104.03059/) (2021)
- [SoftSort: A Continuous Relaxation for the argsort Operator](https://mlsanity.com/paper/2006.16038/) (2020)
- [Learning with Average Top-k Loss](https://mlsanity.com/paper/1705.08826/) (2017)
- [Smooth Loss Functions for Deep Top-k Classification](https://mlsanity.com/paper/1802.07595/) (2018)
- [Loss Functions for Top-k Error: Analysis and Insights](https://mlsanity.com/paper/1512.00486/) (2015)
- [On the Consistency of Top-k Surrogate Losses](https://mlsanity.com/paper/1901.11141/) (2019)
- [Differentiable Sorting Networks for Scalable Sorting and Ranking Supervision](https://mlsanity.com/paper/2105.04019/) (2021)
- [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)
- [Monotonic Differentiable Sorting Networks](https://mlsanity.com/paper/2203.09630/) (2022)
- [Learning with Algorithmic Supervision via Continuous Relaxations](https://mlsanity.com/paper/2110.05651/) (2021)

## Built on by (most cited first, 21 of 21)

- [Grounding Everything: Emerging Localization Properties in Vision-Language Transformers](https://mlsanity.com/paper/2312.00878/) (2023)
- [Fast, Differentiable and Sparse Top-k: a Convex Analysis Perspective](https://mlsanity.com/paper/2302.01425/) (2023)
- [Learning by Sorting: Self-supervised Learning with Group Ordering Constraints](https://mlsanity.com/paper/2301.02009/) (2023)
- [LapSum -- One Method to Differentiate Them All: Ranking, Sorting and Top-k Selection](https://mlsanity.com/paper/2503.06242/) (2025)
- [Dual-Encoders for Extreme Multi-Label Classification](https://mlsanity.com/paper/2310.10636/) (2023)
- [Holistic chemical evaluation reveals pitfalls in reaction prediction models](https://mlsanity.com/paper/2312.09004/) (2023)
- [Context Guided Transformer Entropy Modeling for Video Compression](https://mlsanity.com/paper/2508.01852/) (2025)
- [KDk: A Defense Mechanism Against Label Inference Attacks in Vertical Federated Learning](https://mlsanity.com/paper/2404.12369/) (2024)
- [Welfare as a Guiding Principle for Machine Learning -- From Compass, to Lens, to Roadmap](https://mlsanity.com/paper/2502.11981/) (2025)
- [Provable Robust Saliency-based Explanations](https://mlsanity.com/paper/2212.14106/) (2022)
- [Learning to Summarize Videos by Contrasting Clips](https://mlsanity.com/paper/2301.05213/) (2023)
- [Optimizing Predictive AI in Physical Design Flows with Mini Pixel Batch Gradient Descent](https://mlsanity.com/paper/2402.06034/) (2024)
- [Relative-Based Scaling Law for Neural Language Models](https://mlsanity.com/paper/2510.20387/) (2025)
- [MOCLIP: A Foundation Model for Large-Scale Nanophotonic Inverse Design](https://mlsanity.com/paper/2511.18980/) (2025)
- [Differentiable Knapsack and Top-k Operators via Dynamic Programming](https://mlsanity.com/paper/2601.21775/) (2026)
- [Evaluation Metrics for DNNs Compression](https://mlsanity.com/paper/2305.10616/) (2023)
- [Self-supervised Learning for Anomaly Detection in Computational Workflows](https://mlsanity.com/paper/2310.01247/) (2023)
- [Analyzing Transformer Models and Knowledge Distillation Approaches for Image Captioning on Edge AI](https://mlsanity.com/paper/2506.03607/) (2025)
- [Scalable and Precise Patch Robustness Certification for Deep Learning Models with Top-k Predictions](https://mlsanity.com/paper/2507.23335/) (2025)
- [TT-Sparse: Learning Sparse Rule Models with Differentiable Truth Tables](https://mlsanity.com/paper/2603.07606/) (2026)
- [Learning Permutation from Structure Without Supervision](https://mlsanity.com/paper/2605.25551/) (2026)

## Abstract (arXiv)

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, leading to top-1 or top-5 training objectives. In this work, we relax this assumption and optimize the model for multiple k simultaneously instead of using a single k. Leveraging recent advances in differentiable sorting and ranking, we propose a differentiable top-k cross-entropy classification loss. This allows training the network while not only considering the top-1 prediction, but also, e.g., the top-2 and top-5 predictions. We evaluate the proposed loss function for fine-tuning on state-of-the-art architectures, as well as for training from scratch. We find that relaxing k does not only produce better top-5 accuracies, but also leads to top-1 accuracy improvements. When fine-tuning publicly available ImageNet models, we achieve a new state-of-the-art for these models.
