Research map / Architectures and efficiency / RNNs and normalization
ST$_k$: A Scalable Module for Solving Top-k Problems
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes STk, a fully differentiable smoothed Top-k module with one trainable threshold λ built on a uniformly convergent smoothed ReLU, and uses it to turn the Average Top-k loss into a sort-free STk loss that is trained end-to-end.
Problem
Top-k problems (finding the k-th largest element and sum of top-k values) are costly in deep learning because sorting is O(n log n) or worse and is non-differentiable, which hampers end-to-end optimization of Average Top-k losses for imbalanced, long-tailed, and regression tasks.
Approach
Rewrite the Top-k sum as a min over a threshold λ of a hinge-based objective (Ogryczak and Tamir), replace the non-differentiable ReLU with a uniformly convergent Smoothed ReLU (SReLU) controlled by δ, and optimize λ jointly with network weights via SGD, so no sorting is needed.
Key results
- On CIFAR-100-LT STk reaches 89.8 accuracy versus 89.1 for the prior SOTA, and on Places-LT 53.7 versus 52.2.
- On the Gaussian synthetic task the Top-k sum via STk takes about 0.0127 s for 10,000 samples, versus 0.1243 s for HeapSort and 20.42 s for BubbleSort.
- On binary classification STk attains the lowest misclassification rate on most of the 7 KEEL/UCI datasets, with the lowest standard deviation across them.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| replace ATk aggregation with STk aggregation | long-tailed visual classification · CIFAR-100-LT | accuracy: 89.3 -> 89.8 | Helped |
| replace ATk aggregation with STk aggregation | long-tailed visual classification · ImageNet-LT | accuracy: 78.1 -> 79.1 | Helped |
| replace ATk aggregation with STk aggregation | long-tailed visual classification · Places-LT | accuracy: 52.3 -> 53.7 | Helped |
| replace ATk aggregation with STk aggregation | translation models · IWSLT2014 test | BLEU: 35.45 -> 35.80 | Helped |
| replace ATk aggregation with STk aggregation | translation models · WMT2017 Newstest2016 test | BLEU: 35.05 -> 35.42 | Helped |
| use STk instead of ATk for the Top-5 sum | Top-k computation · 10,000 standard normal samples | time: 0.2167 ± 0.1528 -> 0.0127 ± 0.0020 | Helped |
| add STk to cost-sensitive learning (ablation) | MAE backbone long-tailed learning · ImageNet-LT | accuracy: 69.884 -> 70.140 | Helped |
| add STk to cost-sensitive learning (ablation) | MAE backbone long-tailed learning · CIFAR-100-LT | accuracy: 82.135 -> 83.048 | Helped |
| add STk to PEL (ablation) | CLIP backbone long-tailed learning · ImageNet-LT | accuracy: 78.296 -> 79.148 | Helped |
| add STk to PEL (ablation) | CLIP backbone long-tailed learning · CIFAR-100-LT | accuracy: 89.103 -> 89.833 | Helped |
| set smoothing coefficient δ to 0.1 rather than 0.01 (ablation) | STk · ImageNet-LT | accuracy: 79.09 -> 76.11 | Hurt |
| set smoothing coefficient δ to 0.05 rather than 0.01 (ablation) | STk · WMT2017 Newstest2016 test | BLEU: 35.42 -> 35.28 | Hurt |
Surprising findings
- Adding λ to the computational graph and optimizing it jointly with SGD performs as well as separately solving for λ, so the threshold need not be solved exactly.
- Smoothing the ReLU is reported to be more effective than other ReLU variants (ELU, SoftPlus, Leaky-ReLU) on the synthetic task, with SReLU reaching 76.1% accuracy.
Limitations
- The authors state that STk has only been evaluated for smoothing the ATk loss and not in other application scenarios.
- The approximation error is bounded by δ/2 and performance is sensitive to δ (e.g., accuracy drops at δ=0.1 on CIFAR-100-LT).
- Convergence guarantees for SGD-STk are cited from prior work rather than proved for the neural setting.
Experimental setup
Datasets: CIFAR-100-LT ImageNet-LT Places-LT IWSLT2014 WMT2017 KEEL binary classification datasets UCI datasets Sinc Housing Abalone Cpusmall Gaussian synthetic dataset
Benchmarks: CIFAR-100-LT ImageNet-LT Places-LT IWSLT2014 WMT2017 (Newstest2016) KEEL/UCI binary classification Sinc/Housing/Abalone/Cpusmall regression
Metrics: accuracy BLEU misclassification rate ParaF1 RMSE wall-clock time
Baselines: Average loss Maximum loss ATk loss MATk ELU SoftPlus Leaky-ReLU PEL
Builds on: Average Top-k loss (ATk) Ogryczak and Tamir Top-k formulation Parameter-Efficient Long-Tailed Recognition (PEL) Logit Adjust loss
Closest papers by meaning
- Fast, Differentiable and Sparse Top-k: a Convex Analysis Perspective (ICML 2023 · similarity 0.87)
- LapSum - One Method to Differentiate Them All: Ranking, Sorting and Top-k Selection (ICML 2025 · similarity 0.84)
- RTop-K: Ultra-Fast Row-Wise Top-K Selection for Neural Network Acceleration on GPUs (ICLR 2025 · similarity 0.83)
- Matching the Statistical Query Lower Bound for $k$-Sparse Parity Problems with Sign Stochastic Gradient Descent (NeurIPS 2024 · similarity 0.83)
- Scalable Neural Network Kernels (ICLR 2024 · similarity 0.83)
- Enhancing Logits Distillation with Plug&Play Kendall's $\tau$ Ranking Loss (ICML 2025 · similarity 0.82)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.