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

Authors: hanchen  xia
weidong  liu
xiaojun  mao
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 15:32:26

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 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.

Main paper: 
 introduction: 
 The ranking problem is quite common in the field of AI. For imbalanced datasets, the Average Top-k (AT k ) Loss is more suitable than the conventional Average Loss [Lyu et al., 2020] . In the context of ambiguous classification tasks, Top-k Learning allows the ground truth to fall within the largest k probabilities, enhancing the model's generalizability [Lapin et al., 2017 , Berrada et al., 2018 , Petersen et al., 2022] . For language models, the Top-k sampling method helps the models select the top k most probable words during text generation, producing more fluent and coherent sentences. Distributed learning systems employ Top-k sparsification with error compensation (Top-k SGD) to reduce communication traffic without noticeably impacting model accuracy [Chen et al., 2018 , Lin et al., 2018 , Shi et al., 2019] . However, as deep learning models continue to grow in size, the cost of the ranking process becomes increasingly significant. For example, on a single NVIDIA H800 GPU, tuning a Llama-8B model using the LoRA [Hu et al., 2022] method with a batch size of 4096 and 1000 iterations takes 990.14 seconds. In particular, performing QuickSort on the individual losses consumes 86.33 seconds.In this work, we propose ST k (Smoothed Top-k), a scalable module for Top-k problems. By adding only a single trainable parameter, ST k is able to solve the Top-k problem in O(n + k) steps. Due to the fully differentiable nature of ST k , it can be embedded end-to-end as a layer. Experiments show that we can even add λ to the computational graph for unified optimization using Stochastic Gradient Descent (SGD), which means that we do not need to consider the time cost of solving the Top-k problems within the computational graph. However, we can still enjoy the performance improvements that Top-k optimization provides. The contributions of this work are summarized as follows:• We propose a uniformly convergent approximation of the ReLU function.• We propose an efficient and robust Smoothed Top-k module, ST k .• We apply ST k Module to smooth AT k Loss, resulting in performance improvement and refreshing the state-of-the-art methods on two long-tailed learning leaderboards.• We design an imbalanced classification dataset with a theoretical decision boundary.For experiments, by applying the ST k Module, we smooth the AT k Loss into ST k Loss. The computation time of ST k Loss is almost identical to that of the average aggregating method, significantly faster than the sorting-required AT k Loss, and it exhibits the best performance. Experiments on synthetic datasets demonstrate that models trained with ST k Loss most closely approximate the theoretical decision boundary. On benchmarks of imbalanced binary classification, models trained by ST k Loss exhibit the lowest average misclassification rate and the lowest standard deviation. On regression datasets, models trained with ST k Loss exhibit the lowest (RMSE). Experiments on large real-world datasets demonstrate that ST k Loss, as an aggregating trick for individual losses, is a scalable technique that improves accuracy on long-tailed benchmarks. With the help of ST k Loss, we surpass the state-of-the-art (SOTA) on both the CIFAR-100-LT and Places-LT leaderboards.In current research, although there are many studies on individual loss, the general characteristics of aggregate loss are often overlooked. In the existing machine learning literature, a related line of work is the data subset selection problem [Wei et al., 2015] , which aims to select a subset from a large training dataset for model training while minimizing average loss. Curriculum learning [Bengio et al., 2009] and self-paced [Kumar et al., 2010] learning are two recent learning schemes. They organized the training process into several iterative stages, gradually including training data from easy to difficult to learn, where the difficulty level is measured by individual loss. Therefore, each training session in these methods corresponds to the average aggregate loss in the selected subset.The difficulties encountered by the Average Loss when dealing with imbalanced data, as discussed in [Shalev-Shwartz and Wexler, 2016, Huang et al., 2020] , prompted the exploration of more robust aggregate losses. Among these, Lyu et al. [2020] introduced the AT k loss, which averages the k largest individual losses, and exhibits advanced performance on imbalanced datasets.Many classification tasks in the real world have inherent label confusion, as mentioned in Berrada et al. [2018] . This confusion may arise from various factors, such as incorrect labels, incomplete annotations, or some fundamental ambiguities that even confuse the true labels for human experts. Therefore, some works proposed the concept of Top-k Learning in the field of image classification to address the issues of multiple semantics and semantic confusion in images [Lapin et al., 2017] . Berrada et al. [2018] proposed a method to partially smooth the Top-k Learning loss function, but did not completely solve the sorting problem in the loss function and introduced a sorting computation of C k n . Petersen et al. [2022] proposed a Split Selection Network (SSN) based on sorting networks, which made the Top-k process differentiable and achieved the state of the art on ImageNet-1K at that time. However, the computation required by this method is cumbersome and multilayered. Sorting networks are similar to algorithms like QuickSort with time complexities of O(n log n).Let N n denotes the set {1, ..., n} and I {a} denotes the indicator function (which is 1 when the proposition a is true and 0 otherwise). Thus, the sign function can be defined as: sign(x) = I {x>0} -I {x<0} . The Hinge function can be defined as: [x] + = max {0, x}. We use ∥x∥ 1 , ∥x∥ 2 , and ∥x∥ ∞ to represent the ℓ 1 , ℓ 2 , and ℓ ∞ norms of x, respectively. For the set L = ℓ 1 , ℓ 2 , ..., ℓ n , ℓ [k] denotes the k-th largest element, so we have ℓ[1] ≥ ℓ [2] ≥ ... ≥ ℓ [n]. In supervised learning problems, our training set typically contains an input set and a target set, the input set coming from the input domain X , and the target set from the target domain Y, and we use their joint domain Z = X × Y to represent the range of the dataset. The training set S = z 1 , z 2 , ..., z n is a subset of Z, where z i = (x i , y i ). Our task is to find a predictor f : X → Y from the function family H that can predict the corresponding target y based on the new input x. To evaluate the effect of the predictor, we need to introduce an individual loss function ℓ : Y × Y → R + , where ℓ = ℓ(f (x), y) usually reflects some distance between the prediction ŷ = f (x) and the true value y. The training process of the predictor can be described as using the gradient descent algorithm to optimize an objective function (usually minimizing the loss function). The objective function can generally be written as L(f, S) + Ω(f ), where L(f, S) is the aggregated individual loss function, and Ω(f ) is a regularization term (ℓ 1 or ℓ 2 regularization term). The loss function L(f, S) is usually Average Loss, that is, L avg (f, S) = 1 n n i=1 ℓ(f (x i ), y i ). However, recent works, have shown some drawbacks of the average loss in adapting to imbalanced data distributions [Shalev-Shwartz and Wexler, 2016] , and explored choices other than the average loss for the aggregate loss formed from individual losses, e.g., the maximum (aggregate) loss, L max (f, S) = max i∈Nn ℓ(f (x i ), y i ).Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: The abstract and/or introduction clearly state the claims made.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
st k architecture: 
 Smoothed Top-k Module This process can certainly be achieved through conventional sorting and summation. However, in the worst-case scenario, the cost of solving the Top-k problem can reach O(n 2 ). Ogryczak and Tamir [2003] proposed an equivalent optimization form to solve the Top-k problem and proved its linear convergence.EQUATIONHowever, this surrogate objective function suffers from a non-differentiable point at e i = λ, which makes it challenging to optimize. The key to solving this problem lies in designing a function approximating [•] + , which can be regarded as a rectified linear unit (ReLU) function. Here, we introduce the Smoothed ReLU (SReLU).EQUATION)where δ is a hyperparameter (usually we set δ = 0.01). It can be observed from Figure 2 that as δ decreases, SReLU increasingly approximates the ReLU function. In fact, SReLU converges uniformly to ReLU as δ → 0 + ; a detailed proof of this uniform convergence is provided in Proposition 1 of Appendix A.1. With the help of SReLU, the objective function (1) can be smoothed as:EQUATIONthe optimal λ * = e [k] which is the k-th largest element. In the following section, we will introduce an application scenario for ST k . 3 From AT k Loss to ST k LossIn the training process of a neural network, we first choose a form of individual loss (e.g., logistic loss, hinge loss, mean square loss, or cross-entropy loss). Then, we aggregate all individual losses to calculate their average, which is the most common aggregate loss function: Average Loss. The Average Loss is widely used in a myriad of deep learning tasks. This widespread application stems from the robust theoretical foundations [Bartlett et al., 2006 , De Vito et al., 2005] . However, Average Loss tends to overfit the training data, especially on imbalanced datasets [Shalev-Shwartz and Wexler, 2016, Huang et al., 2020] . This has inspired the motivation to find other forms of aggregate loss, such as the maximum value among individual losses (referred to as Maximum Loss). The Average Top-k (AT k ) Loss was introduced by Lyu et al. [2020] :EQUATIONwhich represents the average of the largest k individual losses. According to the derivation in [Ogryczak and Tamir, 2003] , this ranking loss can be written in the following equivalent form:EQUATIONWith the help of the ST k Module, L mat-k (f, S) can be reconstructed asEQUATIONwhere ℓ i = ℓ(f wmodel (x i ), y i ) is the individual loss of sample i, and w model represents the set of parameters of the predictor f . It is easy to verify that lim δ→0 + L st-k = L mat-k , given Proposition 1 in Appendix A.1.kλ n + 1 n n i=1 [ℓ i -λ] + - kλ n + 1 n n i=1 1 2 (ℓ i -λ) + δ (ℓ i -λ) 2 δ 2 + 1 -1 = n i=1 1 2n (ℓ i -λ) 2 -(ℓ i -λ) 2 + δ 2 + δ = n i=1 1 2n δ - δ 2 (ℓ i -λ) 2 + (ℓ i -λ) 2 + δ 2 = n i=1 1 2n   δ   1 - 1 (ℓi-λ) 2 δ 2 + (ℓi-λ) 2 δ 2 + 1     < δ 2The approximation error between the smoothed loss function and the original loss function can be uniformly bounded by δ/2, ∀ℓ.When ℓ is convex, Equation ( 6) exhibits joint convexity with respect to the parameters (w model , λ), making the problem a special case of the non-linear multiple choice knapsack problem [Zemel, 1984] , which has at most q = 2 roots. These roots can be found in constant time, allowing the problem to be solved in O(n • ln q) = O(n) time when q is fixed [Megiddo, 1984] . Therefore, it can be iteratively updated using relatively simple algorithms. For example, in the case of batch learning, the block coordinate descent (BCD) method [Nocedal and Wright, 1999 ] can be employed, where w model and λ are updated alternately after initialization.EQUATION)The convergence of the above coordinate descent algorithms can be found in Luo and Tseng [1992] , Saha and Tewari [2013] , Tseng [2001] .Furthermore, empirical evidence suggests that we do not need to spend extra time optimizing λ separately, incorporating λ into the computational graph of w model for unified optimization using Stochastic Gradient Descent (SGD) [Bottou and Bousquet, 2008 , Shamir, 2011 , Srebro and Tewari, 2010] , performance improvements can still be achieved.k : λ (t+1) ← λ (t) -η • ∂ λ L st-k ; w (t+1) ← w (t) -η • ∇ w L st-k .where η t is the size of the update step, and when η t ∼ 1 √ t ′ , the stochastic gradient descent method can ensure convergence to a local minimum of Equation ( 6) [Shamir, 2011, Srebro and Tewari, 2010] . By eliminating points where the gradients are discontinuous, the training process becomes more stable, and converges faster, as experimentally demonstrated by the standard deviations reported in Tables 3, 4 , and 9, and the time cost reported in Table 2 .Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: The limitations can be found in Section 6.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
theory assumptions and proofs: 
 Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes] Answer: [Yes] Justification: Main results can be reproduced by supplementary material.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
time cost: 
 We first perform experiments to compare the time costs of two standard sorting algorithms, AT k , and ST k , in calculating the ranking average. The experimental setup involves finding the Top-k (k=5) sum from 10,000 standard normally distributed samples. For both AT k and ST k , we iterate until the error is less than 10 -2 . For each algorithm, we run 50 experiments and record the average time taken.Average Time(s)BubbleSort O(n 2 ) 20.42196 ± 3.7015 HeapSort O(n log(n)) 0.1243 ± 0.0446 AT k O(n + k) 0.2167 ± 0.1528 ST k (Ours) O(n + k) 0.0127 ± 0.0020Table 1 : Performance comparison of different sorting algorithms and our method.As shown in Table 1 , ST k indeed demonstrates linear time complexity while also exhibiting a stable optimization process.To illustrate the capability of ST k Loss in approximating the ideal decision boundary, we design a Gaussian-distributed dataset. The dataset is generated by a pre-set covariance matrix Σ ∈ R d×d , and mean vectors µ 1 , µ 2 ∈ R d corresponding to two categories. We set d = 200, Σ jk = 0.25 0.5+|j-k| , µ 1 = 0 d and µ 2 = (1, 1, ..., 1, 0, 0, ..., 0) ⊤ , which has 10 ones. For the predictor, we use a LogisticRegression (LR) model f (x) = sigmoid(ω ⊤ x + b).The detailed derivation process of the decision boundary can be found in the Appendix A.2. Here, we present the conclusion directly. For two normal populations with a given covariance matrix Σ, and means µ 1 and µ 2 , the LR model has a theoretical decision boundary:EQUATIONFigure 3 provide an example on 2D plane, we setΣ (j,k) = 0.8 |j-k| , µ 1 = [0, 0] ⊤ , µ 2 = [2, 2] ⊤ , thus according to our derivation above, ω * = [1.111, 1.111] ⊤ , b * = -0.836, where the black dashed line is the theoretical boundary.We use ParaF 1 [Tian and Gu, 2017] to measure the overlap of estimated supports and true supports: whereParaF 1 = 2 • precision • recall precision + recall ,precision = |support( ω)∩support(ω * )| |support( ω)| , recall = |support( ω)∩support(ω * )| |support(ω * )|, where ω is the estimated value of the parameters obtained by the predictor, and ω * is the theoretical value of the parameters mentioned above. The operator | • | is used to find the cardinality of the set, which is the number of elements, and support(•) is the support set of the vector, which refers to the set of indices of its non-zero elements.Other basic settings are as follows. The loss function we use is the binary cross-entropy loss. To obtain a sparse solution, we add an ℓ 1 regularization term. The number of samples is 10,000 for the training set and 2,500 each for the validation set and the test set. We implement early stopping of the iteration based on the accuracy curve in the validation set; that is, we break when the increase in validation accuracy within 200 steps is less than 10 -4 . For each meta-experiment, we repeat it 50 times and take the average. Now we conduct two groups of experiments as follows.Aggregate Losses and ReLU Varients. In these experiments, we set the ratio of positive to negative samples at 8:2 for the training set, while the validation set and the test set remain 1:1. We also compare SReLU with other variants of ReLU. The smoothing coefficient for SReLU is set to 0.01, and the settings for the other ReLU variants remain at their defaults. We use Adam as our optimizer, setting the batch size to 512, while keeping the other hyperparameters as default. According to Equation (8), we could compute the theoretical value of ω * , which is a vector with only the first 10 elements non-zero. Sensitive Analysis. In these experiments, we only adjust the positive-negative ratio and retain the rest of the settings from the previous experiment. as follows:Logistic : ℓ(f (x), y) = log(1 + exp(-yf (x))); Hinge : ℓ(f (x), y) = [1 -yf (x)] + .The prediction model is a two-layer MLP with 10 nodes in the hidden layer, activated by the ReLU function between the two fully connected layers. To increase the stability of the training process, we added an ℓ 2 regularization term to the loss function Ω(w) = 1 2C ∥w∥ 2 2 . We divided the datasets into training, validation, and test sets in a 0.5 : 0.25 : 0.25 ratio. Nowadays, SGD is generally replaced by the mini-batch method instead of stochastic gradient descent with a single sample point, which is faster and more robust. To accommodate datasets of varying sizes in the article, we set the batch size to 16.The hyperparameters in the experiment include k in MAT k and AT k , the coefficient of the regularization term C, the initial learning rate η, and the smoothing coefficient δ. These hyperparameters will be selected based on their convergence accuracy performance on the validation set (for each combination of hyperparameters, we repeat the experiment fifty times and take the average of their prediction accuracy on the validation set as the basis for parameter selection). The search spaces for several hyperparameters are as follows: k ∈ {1} ∪ [0.1 : 0.1 : 1]; C ∈ {10 0 , 10 1 , 10 2 , 10 3 , 10 4 , 10 5 }; η ∈ {0.1, 0.05, 0, 01, 0.005, 0.001}; δ ∈ {0.1, 0.01, 0.001, 0.0001}.Since traditional gradient descent is too sensitive to the choice of step size, which is not conducive to our pure comparison of the convergence speed and accuracy of the loss function, we use the AdaGrad algorithm to iteratively update the learning rate. During the learning process, to avoid overfitting the model, we record the accuracy of the MLP predictor on the validation set after each iteration, and perform an early stop when the accuracy does not increase (the increase in the accuracy of the predictor on the validation set is less than 10 -6 after 50 steps) and roll back the model to the checkpoint with the highest accuracy on the validation set during the entire training process; otherwise, we continue training until convergence. Tables 3 and 4 , respectively, show the average probability of misclassification (%) of the models trained to convergence in 50 experiments in the test set under individual Logistic and Hinge loss (the standard deviation of the 50 experiments is listed in parentheses).In the results shown in Tables 3 and 4 . Almost all the lowest misclassification rates appear in the ST k trained models. Furthermore, models trained with ST k Loss exhibit the lowest standard deviation, which to some extent demonstrates the robustness of the training process. 
long-tailed classification: 
 The interest in long-tailed classification tasks has increased with the advent of large vision-language models such as contrastive language-image pre-training (CLIP) Radford et al. [2021] . Long-tailed versions of recognized datasets were built by the community. Using the Pareto distribution (α = 6), the ImageNet-1K and Places datasets can be sampled to create ImageNet-LT and Places-LT datasets, as described in [Liu et al., 2019] . Consider two types of imbalance [Cui et al., 2019 , Buda et al., 2018] , Cao et al. [2019] built CIFAR-10-LT (ρ = 10) and CIFAR-100-LT (ρ = 100), where ρ = max i {n i }/ min i {n i } represents the imbalance ratio, defined as the ratio between the sample sizes of the most frequent and least frequent class. Due to varying word frequencies, machine translation, or more generally, text generating, can inherently be considered as long-tailed classification tasks. We choose the WMT20174 and IWSLT20145 datasets for the experiments. Settings For visual classification tasks, we use Parameter-Efficient Long-Tailed (PEL) Recognition [Shi et al., 2023] , which primarily leverages information from a text encoder to adjust classification probabilities. The text encoder in PEL is derived from the pre-trained CLIP model, and the backbone is a Vision Transformer (ViT-Base) [Dosovitskiy et al., 2021] pre-trained in ImageNet-21K. We freeze the backbone and train the branch model of Parameter Efficient Fine-Tuning (PEFT) [Jia et al., 2022 , Hu et al., 2022 , Houlsby et al., 2019] . We simply replaced the Average aggregation of individual losses in the Logit Adjust (LA) [Menon et al., 2021] loss with AT k or ST k . For k, we set it at 0.9 × batch size, and δ = 0.01. Then we report the accuracy of the balanced test sets. Visual classification experiments can be implemented on a single L20 GPU, with varying durations ranging from 20 to 200 minutes.For translation tasks, we use the training workflows provided by OpenNMT6 [Klein et al., 2018] and FAIR-Seq7 [Ott et al., 2019] , which are very easy to reproduce. Models are essentially different sizes of Transformers (see the footnotes). In the training process of language models, the actual batch size, calculated as the sequence length multiplied by the number of sequences, is often dynamic. Consequently, we set k/batch size = 0.96. We simply replaced the Average aggregation of individual losses in the Cross Entropy (CE) loss with AT k or ST k . We calculate the bilingual evaluation understudy (BLEU) scores for the models. The BLEU for IWSLT2014-trained models are derived from their test set, while for the WMT2017-trained models, we evaluate their BLEU on the Newstest2016 test set. Training these Transformer-based models would take about 10 hours on a single RTX 4090.Overall Results ST k Loss, as an aggregating trick for individual losses, is a scalable technique that improves accuracy on long-tailed benchmarks. The models trained by ST k Loss surpass SOTA on the CIFAR-100-LT and Places-LT leaderboards. Table 5 : Results on large long-tailed datasets are presented, where the first three columns indicate accuracy, and the last two columns show the BLEU scores. Values marked with an "*" represent the state-of-the-art (SOTA) on the leaderboard.For descriptions of the four regression datasets, see Table 9 . In regression tasks, the individual loss is generally the distance (ℓ 1 or ℓ 2 distance) between the predicted value and the ground truth. Next, we present the specific settings of the experiment. For the Housing, Abalone, and Cpusmall datasets, we normalize the output to between [0, 1]; for the Sinc dataset, we first randomly sample 1000 points (x i , y i ) from the function, where x i ∈ [-10, 10], then we use the Radial Basis Function (RBF) kernel to map x i into the kernel space. We select 10 RBF kernels, resulting in a 10-dimensional input x = [k(x, c 1 ), ..., k(x, c 10 )],where k(x, c j ) = exp -(x -c j ) 2 . Furthermore, we add Gaussian noise N (0, 0.2 2 ) to the target output y.Datasets Square Individual Loss Absolute Individual Loss Average MAT k ST k (Ours) Average MAT k ST k (Ours)Sinc 0.1603±0.0065 0.1591±0.0068 0.1446±0.0070 0.1446±0.0070 0.1446±0.0070 0.1686±0.0064 0.1560±0.0067 0.1457±0.0043 0.1457±0.0043 0.1457±0.0043 Housing 0.0896±0.0083 0.0816±0.0100 0.0810±0.0096 0.0810±0.0096 0.0810±0.0096 0.0781±0.0099 0.0670±0.0083 0.0626±0.0089 0.0626±0.0089 0.0626±0.0089 Abalone 0.0739±0.0018 0.0739±0.0018 0.0523±0.0014 0.0523±0.0014 0.0523±0.0014 0.0740±0.0018 0.0733±0.0018 0.0517±0.0021 0.0517±0.0021 0.0517±0.0021 Cpusmall 0.0613±0.0012 0.0313±0.0005 0.0296±0.0005 0.0296±0.0005 0.0296±0.0005 0.0613±0.0012 0.0313±0.0005 0.0244±0.0006 0.0244±0.0006 0.0244±0.0006 For individual losses, we choose the absolute value loss (ℓ 1 ) and the squared loss (ℓ 2 ). The prediction model is the same as above, a two-layer MLP with 10 nodes in the hidden layer, activated by the ReLU function between the two fully connected layers. As in binary classification experiments, we randomly divide the dataset into training, validation, and test sets, with hyperparameters chosen based on the performance in the validation set, repeated 50 times, and statistics on the test set are reported.Tables 6 show the average RMSE and standard deviation of 50 experiments in regression datasets with different aggregate losses combined with absolute individual loss and squared individual loss (given the experience of the previous experiment, we no longer compare with Maximum Loss here), and ST k shows further performance improvement after smoothing. 
conclusion and limitation: 
 The proposed ST k Module can effectively solve the Top-k problem within neural networks, with no additional GPU memory or ranking time. Due to its fully differentiable nature, the training process are stable. By applying ST k Module to the Average Top-k Loss, we achieve significant improvements across numerous benchmarks.The limitation of this study is that we have not yet evaluated the ST k Module for additional application scenarios. The effectiveness of the ST k Module is only demonstrated in smoothing AT k loss. We hope that future research will further explore the potential utility of the ST k Module.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: The paper has specified all the training and test details.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: We provide the Std of the results, where the proposed method exhibits the lowest Std and surpasses other methods in terms of mean performance.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: The computational resource requirements of experiments on large real-world datasets can be found in 5.2, while the remaining experiments can be produced on CPU.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes] Justification: Sure about that.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification:Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification:Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We ensure that the creators or original owners of assets used in the paper are properly credited, and the license and terms of use are explicitly mentioned and properly respected.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.• For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [NA] Justification: Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: A.1 Uniform Convergence [Proposition 1] [Proposition 1] [Proposition 1]We have that SReLU uniformly converges to ReLU as δ → 0 + , that is,lim δ→0 + sup x |SReLU δ (x) -[x] + | = 0. Proof. Proof. Proof. ReLU(x) = [x] + = 1 2 (x + |x|); SReLU(x) = 1 2 x + δ x 2 δ 2 + 1 -1 . Note that D(x) = 2(ReLU(x) -SReLU(x)) = |x| -δ x 2 δ 2 + 1 -1 , we have D(x) = √ x 2 -x 2 + δ 2 + δ = δ - δ 2 √ x 2 + √ x 2 + δ 2 = δ 1 - 1 x 2 δ 2 + x 2 δ 2 + 1.For any x and δ, the form 0 < 1 -1 x 2 δ 2 + x 2 δ 2 +1< 1, therefore, for any ϵ > 0, if δ < ϵ, then ∀x we shall have D(x) < ϵ. ω * = (µ 1 -µ 0 ) ⊤ Σ -1 ; b * = 1 2 µ ⊤ 0 Σ -1 µ 0 + µ ⊤ 1 Σ -1 µ 1 + ln Pr(Y = 1) Pr(Y = 0) .Proof. Proof. Proof. Usually, we estimate the parameters ω and b by maximizing the likelihood function1 n n i=1 ln[Pr(y i |x i )].This is to estimate parameters in the "most data-fitting manner" when the distribution is unknown. Now we know that two classes are normally distributed, we even know their mean and covariance matrix.In Logistic Regression Here, since we inherently know the data is normally distributed, we havePr(Y = 1|x) = 1 1 + e -(ω ⊤ x+b) ; Pr(Y = 0|x) = e -ω ⊤ x+b 1 + e -(ω ⊤ x+b) .ln Pr(Y = 1|x) Pr(Y = 0|x) = (µ 1 -µ 0 ) ⊤ Σ -1 x + 1 2 µ ⊤ 0 Σ -1 µ 0 + µ ⊤ 1 Σ -1 µ 1 + ln Pr(Y = 1)Pr(Y = 0) .(10) Then, comparing Equation ( 10) with Equation ( 9), we'll have the answer.We conduct an ablation study on "how ST k complements other long-tailed learning algorithms." The backbone here is a Vision Transformer (ViT) pre-trained by MAE / CLIP [He et al., 2022 , Radford et al., 2021] . "CS" represents cost-sensitive learning, and "PEL" refers to the method provided by [Shi et al., 2023] . The batch size was set to 2048 and the models were trained in 30,000 steps. The smoothing coefficient δ, 0.01, is a grid search determined value for all datasets in Section 5.1 (see page 7) and was adopted in all other experiments. We conducted an ablation study on the five real-world datasets with all other settings remaining unchanged, except for δ. In Table 9 , we provide the statistical information for two benchmark sets discussed in Section 5, where c, n, d are the number of classes, samples and features, respectively. The left part provides the information of four regression datasets in Section 5.3. The "Sinc" dataset is a synthetic dataset, sampled from the sinc function y = sin(x)/x. The right part provides the information of seven binary classification datasets in Section 5.1.http://sci2s.ugr.es/keel/datasets.phphttps://archive.ics.uci.edu/ml/datasets.htmlhttps://www.statmt.org/wmt17/https://workshop2014.iwslt.org/https://github.com/OpenNMT/OpenNMT-py/blob/master/docs/source/examples/wmt17/Translation.mdhttps://fairseq.readthedocs.io/en/latest/_modules/fairseq/models/transformer/transformer_legacy.html 


Decision: The paper received the following ratings:  R1 - 7 (accept), R2 - 4 (borderline reject, no reaction to rebuttal), R3 - 4 (borderline reject), R4 - 6 (weak accept).

Overall reviewers highlighted that paper is clearly written and easy to follow, that the proposed top-k loss is an important problem, and that the experiments are clear and show the improvement by the provided loss. 

The main weaknesses listed are that the evaluation might be a bit limited and that more theoretical would be helpful. The authors provided a solid rebuttal which provided additional information. 
While I share the concerns of the reviewers for the original paper, I think that the paper can significantly improve if all comments and additional experiments from the rebuttal would be incorporated in the main paper and/or the supplementary. Given the rebuttal, I'm also confident that the authors are interested in improving the paper and will do this. 

I therefore recommend accepting the paper.

Reviews: review: paper_summary: paper_summary: Authors introduce a novel differentiable module (ST_k) for efficiently solving top-k problems. Their method relies on optimizing a differentiable form of an equivalent optimization problem, by proposing an approximation of ReLU differentiable everywhere. This equivalent optimization problem contains a single parameter \lambda which is supposed to match the k-th largest element in the optimal case, and in practice is optimized over the dataset. Authors  experiment on synthetic data, showing that their loss objective results in a model that more closely matches the true decision boundary in this setting. Next, authors experiment with real-world applications on binary classification and long-tailed classification datasets, showing consistent improvement over other optimization losses.
strength_weakness: strength_weakness: **Weaknesses**

- Although the improvements shown by the authors are consistent, they are also marginal - only improving slightly in for example CIFAR-100-LT classification and ImageNet-LT classification compared to average aggregation. 
- As authors mention, they only apply their module to a single Top-K Loss; ATK, making it unclear how ST_k could be applied in other scenarios and how well it generalizes to different settings.
questions: questions: **Questions**
- How do you tune the smoothing coefficient in practice? Is this a hyperparameter that needs to be sweeped over? That would somewhat deteriorate the efficiency of your module. I would like to see an ablation over the values of delta, i think it is important to see what/if any impact this choice has on model performance.
- Did you experience any training instabilities with the two-step optimization scheme?
limitations: limitations: The authors mention as limitation the fact that they only applied ST_k to the average top-k loss. I think it would be good if they more explicitly mention future research directions that they are interested in/ that would be of value (there is space in the manuscript).
paper_summary: paper_summary: The authors proposed a differentiable layer to approximate the top-k loss in deep learning. The proposed layer is motivated from Eq. (1), replacing the ReLU with a smoothed ReLU (SReLU in Eq. (2)). The authors showed that the proposed layer is point-wise convergent to top-k loss. Numerical experiments validate the superiority of the proposed method.
strength_weakness: strength_weakness: * The biggest concern is that most of the real-world experiments are done on small scale datasets. It would be more convincing if the authors could demonstrate their method on large-scale ranking problems, especially, large deep learning models.

* The proposed method lacks strong theoretical guarantees. Point-wise convergence is somehow weak. As easy to expect, the convergence rate and error bound of the proposed method should depend on the data distribution, which is not discussed in the paper. It would be nice if the authors could show that under what conditions, the proposed method will converge faster than $AT_k$ and / or with smaller approximation error.

* It is important to check how well the proposed method approximates top-k loss in real-world problems. However, most real-world experiments report model accuracy, which is an indirect metric of top-k loss approximation. Please consider to measure the top-k loss approximation error directly.
questions: questions: * In real-world problems, is it possible to design experiments to measure the top-k loss approximation error directly?
* What is the theoretical advantage of $ST_k$ v.s. $AT_k$
limitations: limitations: N/A
paper_summary: paper_summary: This paper proposes a differentiable module for solving the top-k problem. Specifically, the paper proposes to approximate the hinge function with a new differentiable function.   
Experiments on binary classification, long-tailed classification, and regression tasks with $ST_k$ loss show some improvements over baselines, like Average loss, and $AT_k$ loss.
strength_weakness: strength_weakness: (1) For the implementation of $AT_k$, is it implemented by a sort algorithm or just Eq. (5)?
     From my point of view, the performance with $ST_k$ loss should be bounded by the performance with $AT_k$ loss.
     However, Tables 1,4 show that $ST_k$ outperforms $AT_k$.
     Why is $ST_k$ loss better than $AT_k$ loss?

(2) The proposed $ST_k$ loss is flexible and can be directly deployed on the traditional training paradigm.
     It would be better to validate the effectiveness of $ST_k$ on large-scale data, like ImageNet, and iNaturalist.
    
(3) The authors claim that the average loss is easily overfitted on imbalanced data. With $ST_k$ loss and CLIP models, it achieves the state-of-the-art on ImageNet-LT and Places-LT.
     Is the $ST_k$ loss complementary to other state-of-the-art long-tailed learning methods, like GPaCo[2] and BCL[1]?    
     [1] Balanced Contrastive Learning for Long-Tailed Visual Recognition. CVPR 2022.     
     [2] Generalized Parametric Contrastive Learning. TPAMI 2023. 

(4) The models trained with $ST_k$ loss could be optimized alternatively or simultaneously with SGD for model parameters and $\lambda$. The ablation study is encouraged to be included.
questions: questions: (1) Explanation of why $ST_k$ outperforms $AT_k$.                  
(2) The experiments on large-scale data, like ImageNet.    
(3) Is the $ST_k$ loss complementary to current state-of-the-art long-tailed learning methods.    
(4) Ablation for the optimization manner of model parameters and $\lambda$.
limitations: limitations: Limitations are discussed.
paper_summary: paper_summary: This paper addresses the Top-K problem by introducing a new loss function. Building on the Average Top-K Loss, the authors incorporate a smoothed ReLU function to create the $ST_k$ loss, which is fully differentiable. Through experiments on various datasets, they demonstrate the effectiveness of their approach.
strength_weakness: strength_weakness: 1. The core weakness lies in the lack of significant innovation. The main technical contribution is the replacement of the ReLU-like operation in an existing optimization method for the Top-K problem with a smoothed ReLU function, controlled by a hyperparameter for smoothness. Similar concepts already exist in various ReLU function variants, as seen in related works such as [1,2].
2. The only large-scale experiment is a long-tailed classification task involving parameter-efficient fine-tuning of a CLIP model pretrained on ImageNet-21K. Given that this model already has substantial pre-trained knowledge, the results are less convincing. Evaluating the method by training a model from scratch on these tasks would better demonstrate its effectiveness in practical scenarios.
3. The paper does not include an ablation study on the hyperparameter $\delta$ to study its impact on the performance.

[1]. Biswas, Koushik, et al. "SMU: smooth activation function for deep networks using smoothing maximum technique." arXiv preprint arXiv:2111.04682 (2021).

[2]. Hendrycks, Dan, and Kevin Gimpel. "Gaussian error linear units (gelus)." arXiv preprint arXiv:1606.08415 (2016).
questions: questions: 1. In lines 20-21, the authors mention that "the cost of the ranking process can no longer be ignored." I am curious about how significant this ranking cost is in the context of the full forward pass of a deep learning model. Given that larger models require more computational resources for their intermediate layers, this ranking cost might be negligible in most cases. Additionally, there is no detailed comparison of time costs between the Average Top-K Loss $AT_k$ and the proposed $ST_k$ loss, apart from on the toy synthetic dataset.
limitations: limitations: Yes.

Comments: title: General Response after Author-Reviewer Discussion
Thanks for the responses from the authors
Response
comment: We are very grateful for the efforts and suggestions of each reviewer during the author-reviewer discussion period. This has greatly helped in improving the quality of this work. Below are the supplements and improvements we made based on the Reviewer feedback.

**1. We added 3 experiments temporarily.**

* We present a temporary ablation over the values of the smooth coefficient $\delta$ on CIFAR-100-LT. All other settings remain the same as those in the paper, with only $\delta$ being changed.

| $\delta$   | 0.1 | 0.05 | 0.01 | 0.005 |
|--------------|-------|--------|-------|--------|
| Accuracy | 88.8 | 89.4 | 89.8 | 89.8 |

* We conduct experiments on “how ST$_k$ complements other long-tailed algorithms.”
The backbone here is an Vision Transformer (ViT) pre-trained by MAE / CLIP [1][2]. "CS" here represents the cost-sensitive learning; while PEL was provided by [4], which is SOTA method on 2 long-tailed classification leaderboards. If the model was pre-trained by MAE, we add an Linear Classifier after the backbone; other-wise we follow the default PEFT in [3]. The batch size was set to 2048, training 30,000 steps. The results are as follows:

Pretrained By| CS | PEL | ST$_k$ | ImageNet-LT | CIFAR-100-LT |
|-------|-------|-------|-----------|-------------------|---------------------|
| MAE| –     | –      | –          | 65.701           | 78.572             |
| MAE| ✓     | –      | –         | 69.884           | 82.135             |
| MAE| ✓     | –      | ✓         | 70.140           | 83.048             |
| CLIP| –      | ✓     | –          | 78.296           | 89.103             |
| CLIP| –      | ✓     | ✓         | 79.148           | 89.833             |

  **[1]** He K, Chen X, Xie S, et al. Masked autoencoders are scalable vision learners[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022: 16000-16009.

  **[2]** Radford A, Kim J W, Hallacy C, et al. Learning transferable visual models from natural language supervision[C]//International conference on machine learning. PMLR, 2021: 8748-8763.

  **[3]** Shi J X, Wei T, Zhou Z, et al. Parameter-efficient long-tailed recognition[J].

* We experiment with the time cost of a few common sorting algorithms, $AT_k$ and $ST_k$ to calculate the ranking average. The specific scheme of the experiment is to find the Top-k (k=5) sum from 10,000 standard normally distributed samples. For $AT_k$ and $ST_k$, we iterate until the error is less than $10^{-2}$, and for each algorithm, we conduct 50 experiments and record the average time taken in the third column of the Table.

| Algorithm | Complexity       | Average Time(s)          |
|-----------|------------------|--------------------------|
| BubbleSort | $\mathcal{O}(n^2)$     | 20.42196 ± 3.7015        |
| HeapSort   | $\mathcal{O}(n \log(n))$ | 0.1243 ± 0.0446          |
| $AT_k$   | $\mathcal{O}(n + k)$     | 0.2167 ± 0.1528          |
| $ST_k$ (Ours) | $\mathcal{O}(n + k)$ | **0.0127 ± 0.0020**      |


**2. We realize that our current wording may causes some confusion.**

The approximation error between the smoothed loss function and the original loss function can uniformly bounded by $\delta/2, \forall (\lambda, w_{model})$ (see page 12, A.1), which means this is the uniform convergence, a recognized very strong condition.



**Thanks again for reviewing our work. We are grateful for the valuable feedback provided by each reviewer. It is a great honor to meet you at NeurIPS2024.**
Thanks again for reviewing our work. After carefully considering all your suggestions, we realized there is one more thing we need to clarify. 

We hold a different view on your concern that **"point-wise convergence is somehow weak."** The approximation error between the smoothed loss function and the original loss function can uniformly bounded by $\delta/2, \forall (\lambda, w_{model})$ (see page 12, A.1), which means this is the **uniform convergence**, a very strong condition.

We sincerely hope you may reconsider whether our response fully addresses your concern.

Best wishes,
Authors.
We carefully reconsidered your questions. We would like to further clarify our intentions.

**For your questions:**

> Q2. The proposed loss function is general and should be applicable to balanced data, like ImageNet. If not, the authors should provide more analysis on why the loss is specific to imbalanced data. How does the proposed method alleviate the overfitting issue of imbalanced data?

* The curves in Figure 3 clearly reflect that as the degree of imbalance increases, the Average Loss continuously misguides the model to shift the decision boundary toward the minority class.
* In contrast, AT$_k$ Loss prevent the model from shifting by ignoring those individual losses with smaller values.
* However,  the discontinuity affected the stability of the optimization process of $\lambda$. Therefore, we need to find a 
 objective function, 
    * with continuous gradient; 
    * at the same time sufficiently approximate the Average Top-k Loss $\frac{k \lambda}{n} + \frac{1}{n}\sum_{i=1}^n [\ell_i - \lambda]_+$.
* $ST_k$ happens to achieve this. Theoretically, the approximation error can be bounded by $\delta/2, \forall \lambda, w_{model}$. This is a **uniform convergence**, which is known as a very strong condition. (because $\delta \in \mathbb{R}$, we call it point-wise convergence, this may cause some misunderstanding)
* With the help of $ST_k$, models show consistent performance improvements, on synthetic datasets and real-world datasets.
Thanks for your insightful suggestions. Hope our explanations address your concerns. Based on the feedback, we are glad to inform you that we include 2 additional experiments.

**1.** We present a temporary ablation over the values of the smooth coefficient $\delta$ on CIFAR-100-LT. All other settings remain the same as those in the paper, with only $\delta$ being changed.

| $\delta$   | 0.1 | 0.05 | 0.01 | 0.005 |
|--------------|-------|--------|-------|--------|
| Accuracy | 88.8 | 89.4 | 89.8 | 89.8 |

**2.** We conduct experiments on “how ST$_k$ complements other long-tailed algorithms.”
The backbone here is an Vision Transformer (ViT) pre-trained by MAE / CLIP [1][2]. "CS" here represents the cost-sensitive learning; while PEL was provided by [4], which is SOTA method on 2 long-tailed classification leaderboards. If the model was pre-trained by MAE, we add an Linear Classifier after the backbone; other-wise we follow the default PEFT in [3]. The batch size was set to 2048, training 30,000 steps. The results are as follows:

Pretrained By| CS | PEL | ST$_k$ | ImageNet-LT | CIFAR-100-LT |
|-------|-------|-------|-----------|-------------------|---------------------|
| MAE| –     | –      | –          | 65.701           | 78.572             |
| MAE| ✓     | –      | –         | 69.884           | 82.135             |
| MAE| ✓     | –      | ✓         | 70.140           | 83.048             |
| CLIP| –      | ✓     | –          | 78.296           | 89.103             |
| CLIP| –      | ✓     | ✓         | 79.148           | 89.833             |

**[1]** He K, Chen X, Xie S, et al. Masked autoencoders are scalable vision learners[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022: 16000-16009.

**[2]** Radford A, Kim J W, Hallacy C, et al. Learning transferable visual models from natural language supervision[C]//International conference on machine learning. PMLR, 2021: 8748-8763.

**[3]** Shi J X, Wei T, Zhou Z, et al. Parameter-efficient long-tailed recognition[J].


**3.** The explanations we provided to Reviewers `SwHE` and `46Ey` (we appreciate the concerns they raised) may appropriately answer your question **"What is the theoretical advantage of ST$_k$ vs AT$_k$?"**

* The motivation here is that the discontinuity has affected the stability of the parameter $\lambda$ optimization process. 
Therefore, we need to find a smooth objective function to approximate: $\min_\lambda \frac{k \lambda}{n} + \frac{1}{n}\sum_{i=1}^n [\ell_i - \lambda]_+$

* Other ReLU-variants cannot sufficiently approximate ReLU as they are only similar in shape. Therefore, we design SReLU and proved that the approximation error can be bounded by $\delta / 2$ (see page 12, A.1). 

  $ \left[ \frac{k \lambda}{n} + \frac{1}{n}\sum_{i=1}^n [\ell_i - \lambda]_+ \right] $

  $- \left[ \frac{k \lambda}{n} + \frac{1}{n}\sum_{i=1}^n \frac{1}{2}\left[(\ell_i-\lambda) + \delta \left( \sqrt{\frac{(\ell_i-\lambda)^2}{\delta^2}+1}-1 \right)\right] \right]$

  $= \sum_{i=1}^n \frac{1}{2n}  \left[ \sqrt{(\ell_i-\lambda)^2} - \sqrt{(\ell_i-\lambda)^2 + \delta^2} + \delta \right]$

  $= \sum_{i=1}^n \frac{1}{2n}  \left[ \delta - \frac{\delta^2}{\sqrt{(\ell_i-\lambda)^2} + \sqrt{(\ell_i-\lambda)^2 + \delta^2}} \right]$

  $=  \sum_{i=1}^n \frac{1}{2n}  \left[ \delta \left( 1 - \frac{1}{\sqrt{\frac{(\ell_i-\lambda)^2}{\delta^2}} + \sqrt{\frac{(\ell_i-\lambda)^2}{\delta^2} + 1}} \right) \right] < \frac{\delta}{2}$

* The smoothed objective function shows jointly convex w.r.t $(w_{model}, \lambda)$, making ST$_k$ a special case of the non-linear multiple choice knapsack problems, with at most $q=2$ roots, which can be found in constant time. Thus, the problem can be solved in $\mathcal{O}(n*lnq)=\mathcal{O}(n)$ time as $q$ being fixed [4] (section 4, "non-linear case"). 

* Experiments have shown that, empirically, we don't even need to spend additional time optimizing $\lambda$ separately. Instead, by optimizing it synchronously with the model parameters $w_{model}$, we can still achieve performance improvements.

   **[4]** Zemel E. An O (n) algorithm for the linear multiple choice knapsack problem and related problems[J]. Information processing letters, 1984, 18(3): 123-128.

    **[5]** N. Megiddo, Linear programming in linear time when the dimension is fixed, J. ACM 31 (1984) 114–127.

    **[6]** Ogryczak, Wlodzimierz, and Arie Tamir. "Minimizing the sum of the k largest functions in linear time." Information Processing Letters 85.3 (2003): 117-122.

* By smoothing the $[\cdot]_+$ function, the optimization process becomes stable and efficient (see standard division in Table 2, 3 and 5, also time cost in Table 1).
We appreciate your recommendation. Best wishes.
Thanks for listing out all your confusions! Some of them can be explained theoretically, while others can be clarified through additional experiments. Here is the point-by-point response.

> Q1. As the authors claim that ST$_k$ can dynamically adjust the k for loss calculation, there should be some empirical or theoretical analysis to confirm the hypothesis.

* The $\lambda^* = \ell_{[k]}$ not only serves as a filter, excluding individual losses that smaller than $\ell_{[k]}$ (see $\min_\lambda k\lambda + \sum_{i=1}^n [\ell_i - \lambda]_+$). The $\lambda$ itself, also serves as a trainable parameter aiming to minimize the loss.

* By smoothing the $[\cdot]_+$ function, the optimization process becomes stable and efficient (see standard division in Table 2, 3 and 5, also time cost in Table 1).


> Q2. The proposed loss function is general and should be applicable to balanced data, like ImageNet. If not, the authors should provide more analysis on why the loss is specific to imbalanced data. How does the proposed method alleviate the overfitting issue of imbalanced data?

* In this paper, we use the ST$_k$ module to smooth the Average Top-k Loss (AT$_k$ Loss) to demonstrate that ST$_k$ can be applied in deep learning scenarios. We did not conduct experiments on balanced datasets because: AT$_k$ Loss was designed for imbalanced datasets [1], as this was proved by its creator, empirically and theoretically.

* The application of ST$_k$ does not require additional computational time/resources while consistently bringing performance improvements.

[1] Lyu S, Fan Y, Ying Y, et al. Average top-k aggregate loss for supervised learning[J]. IEEE transactions on pattern analysis and machine intelligence, 2020, 44(1): 76-86.


> Q3. The authors claim that the ST$_k$ complements other long-tailed algorithms. However, there is no empirical analysis to support it.

This is a very valuable suggestion, glad that you asked! The fact is that, we did conduct experiments on “how ST$_k$ complements other long-tailed algorithms.” However, due to the limited space, we did not include them in the present version. Here we would like to  provide a temporary experiment, and consider providing a detailed ablation study in our final version.

The backbone here is an Vision Transformer (ViT) pre-trained by MAE / CLIP [2][3]. "CS" here represents the cost-sensitive learning; while PEL was provided by [4], which is SOTA method on 2 long-tailed classification leaderboards. If the model was pre-trained by MAE, we add an Linear Classifier after the backbone; other-wise we follow the default PEFT in [4]. The batch size was set to 2048, training 30,000 steps. The results are as follows:

Pretrained By| CS | PEL | ST$_k$ | ImageNet-LT | CIFAR-100-LT |
|-------|-------|-------|-----------|-------------------|---------------------|
| MAE| –     | –      | –          | 65.701           | 78.572             |
| MAE| ✓     | –      | –         | 69.884           | 82.135             |
| MAE| ✓     | –      | ✓         | 70.140           | 83.048             |
| CLIP| –      | ✓     | –          | 78.296           | 89.103             |
| CLIP| –      | ✓     | ✓         | 79.148           | 89.833             |

[2] He K, Chen X, Xie S, et al. Masked autoencoders are scalable vision learners[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022: 16000-16009.

[3] Radford A, Kim J W, Hallacy C, et al. Learning transferable visual models from natural language supervision[C]//International conference on machine learning. PMLR, 2021: 8748-8763.

[4] Shi J X, Wei T, Zhou Z, et al. Parameter-efficient long-tailed recognition[J].


> Q4. Why is the ST$_k$ specific to the Adam optimizer?

* We claimed that the proposed loss can be optimized by the most common optimizers like SGD / Adam, only to stress the fact that no  additional time is needed to update $\lambda$ and the model parameters $w_{model}$ separately in two steps by using BCD (Block Coordinate Descent).

* Empirically, ST$_k$ is not specific to Adam, it also works with SGD, AdaGrad, and BCD. The Supplementary Material we uploaded include the source code for the experiments. If necessary, we would consider providing the results of different optimizers in our final version.


Hope these explanations address your concerns. If you have any confusion, please feel free to reach out to us anytime.
Thank you for addressing my concerns about the theoretical support. I would like to increase my final rating.
Thanks for raising this issue, and we are glad to address this for all reviewers.

**1.** The motivation here is that the discontinuity has affected the stability of the parameter $\lambda$ optimization process. 
Therefore, we need to find a smooth objective function to approximate:

$\min_\lambda \frac{k \lambda}{n} + \frac{1}{n}\sum_{i=1}^n [\ell_i - \lambda]_+$

**2.** Other ReLU-variants cannot sufficiently approximate ReLU as they are only similar in shape. Therefore, we design SReLU and proved that the approximation error can be bounded by $\delta / 2$ (see page 12, A.1). 

  $ \left[ \frac{k \lambda}{n} + \frac{1}{n}\sum_{i=1}^n [\ell_i - \lambda]_+ \right] $

  $- \left[ \frac{k \lambda}{n} + \frac{1}{n}\sum_{i=1}^n \frac{1}{2}\left[(\ell_i-\lambda) + \delta \left( \sqrt{\frac{(\ell_i-\lambda)^2}{\delta^2}+1}-1 \right)\right] \right]$

  $= \sum_{i=1}^n \frac{1}{2n}  \left[ \sqrt{(\ell_i-\lambda)^2} - \sqrt{(\ell_i-\lambda)^2 + \delta^2} + \delta \right]$

  $= \sum_{i=1}^n \frac{1}{2n}  \left[ \delta - \frac{\delta^2}{\sqrt{(\ell_i-\lambda)^2} + \sqrt{(\ell_i-\lambda)^2 + \delta^2}} \right]$

  $=  \sum_{i=1}^n \frac{1}{2n}  \left[ \delta \left( 1 - \frac{1}{\sqrt{\frac{(\ell_i-\lambda)^2}{\delta^2}} + \sqrt{\frac{(\ell_i-\lambda)^2}{\delta^2} + 1}} \right) \right] < \frac{\delta}{2}$

**3.** The smoothed objective function shows jointly convex w.r.t $(w_{model}, \lambda)$, making ST$_k$ a special case of the non-linear multiple choice knapsack problems, with at most $q=2$ roots, which can be found in constant time. Thus, the problem can be solved in $\mathcal{O}(n*lnq)=\mathcal{O}(n)$ time as $q$ being fixed [1] (section 4, "non-linear case"). 

[1] Zemel E. An O (n) algorithm for the linear multiple choice knapsack problem and related problems[J]. Information processing letters, 1984, 18(3): 123-128.

[2] N. Megiddo, Linear programming in linear time when the dimension is fixed, J. ACM 31 (1984) 114–127.

[3] Ogryczak, Wlodzimierz, and Arie Tamir. "Minimizing the sum of the k largest functions in linear time." Information Processing Letters 85.3 (2003): 117-122.

**4.** Experiments have shown that, empirically, we don't even need to spend additional time optimizing $\lambda$ separately. Instead, by optimizing it synchronously with the model parameters $w_{model}$, we can still achieve performance improvements.

Hope these explanations address your concerns. If you have any confusion, please feel free to reach out to us anytime.
Thanks for the responses from the authors. I still have the following confusion on the paper.

Q1. As the authors claim that $ST\_{k}$ can dynamically adjust the k for loss calculation, there should be some empirical or theoretical analysis to confirm the hypothesis.

Q2. The proposed loss function is general and should be applicable to balanced data, like ImageNet.
       If not, the authors should provide more analysis on why the loss is specific to imbalanced data. How does the proposed method alleviate the overfitting issue of imbalanced data? 

Q3. The authors claim that the $ST\_{k}$ complements other long-tailed algorithms. However, there is no empirical analysis to support it.

Q4. Why is the $ST\_{k}$ specific to the Adam optimizer?

If the concerns can't be addressed, I prefer to keep the initial rating.
Thanks for the authors' response. Some of my concerns have been addressed in the rebuttal. However, after reviewing comments from other reviewers, such as reviewer uCvC, I also share curiosity about the theoretical support behind the proposed method, which has not been directly addressed in the rebuttal. The SReLU approach, in particular, seems to be an empirical result without clear theoretical guidance.
Thanks again for reviewing our work. As the author-reviewer discussion phase wraps up, could you let us know if our responses have addressed your concerns? If so, would you reconsider your rating? If you still have concerns, please share so we can address them.
Thanks again for reviewing our work. As the author-reviewer discussion phase wraps up, could you let us know if our responses have addressed your concerns? If so, would you reconsider your rating? If you still have concerns, please share so we can address them.
We really appreciate your recommendation!

Here we present a temporary ablation over the values of the smooth coefficient $\delta$ on CIFAR-100-LT. All other settings remain the same as those in the paper, with only $\delta$ being changed.

| $\delta$   | 0.1 | 0.05 | 0.01 | 0.005 |
|--------------|-------|--------|-------|--------|
| Accuracy | 88.8 | 89.4 | 89.8 | 89.8 |
 
Integral experiments will be included in our final version. Thanks again for your suggestions.
I thank the authors for their response. My concerns have been answered, i i crease my recommendation.