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

Hanchen Xia, Weidong Liu, Xiaojun Mao
2/3/2026

No confident code implementation found for this paper. 4 possible matches can be shown below.

Abstract

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.

DOISemantic Scholar

Code Implementations

No confident code match yet

We couldn't find an author-owned or strongly-evidenced community implementation for this paper. 4 weaker matches are hidden by default — verify before relying on them.

No code implementations found yet.

Know of an implementation? Let us know in the comments below!

Cite this paper

@article{xia2026stk,
  title  = {ST$_k$: A Scalable Module for Solving Top-k Problems},
  author = {Hanchen Xia and Weidong Liu and Xiaojun Mao},
  year   = {2026},
  doi    = {10.52202/079017-0534},
  url    = {https://doi.org/10.52202/079017-0534},
  journal = {NEURIPS 2024 2024}
}

Discussion