Skip to yearly menu bar Skip to main content


Poster

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

Hanchen Xia · Weidong Liu · Xiaojun Mao


Abstract: The cost of ranking becomes significant in the new stage of deep learning. We propose 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 embeded 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 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.

Live content is unavailable. Log in and register to view live content