ST$_k$: A Scalable Module for Solving Top-k Problems

The cost of ranking becomes significant in the new stage of deep learning. We pro-pose ST k , a fully differentiable module with a single trainable parameter, designed to solve the Top-k problem without requiring additional time or GPU memory. Due to its fully differentiable nature, ST k can be embedded end-to-end into neural networks and optimize the Top-k problems within a unified computational graph. We apply ST k to the Average Top-k Loss (AT k ), which inherently faces a Top-k problem. The proposed ST k Loss outperforms AT k Loss and achieves the best average performance on multiple benchmarks, with the lowest standard deviation. With the assistance of ST k Loss, we surpass the state-of-the-art (SOTA) on both CIFAR-100-LT and Places-LT leaderboards.