Differentiable Top-$k$ Classification Learning

To make the sorting operation differentiable, Grover et al. (2019) proposed relaxing permutation matrices to unimodal row-stochastic matrices. For this, they use the softmax of pairwise differences of (cumulative) sums of the top elements. They prove that this, for the temperature parameter approaching $0$, is the correct permutation matrix, and propose a variety of deep learning differentiable sorting benchmark tasks. They propose a deterministic softmax-based variant, as well as a Gumbel-Softmax variant of their algorithm. Note that NeuralSort is not based on sorting networks.

Prillo & Eisenschlos (2020) build on this idea but simplify the formulation and provide SoftSort, a faster alternative to NeuralSort. They show that it is sufficient to build on pairwise differences of elements of the vectors to be sorted instead of the cumulative sums. They find that SoftSort performs approximately equivalent in their experiments to NeuralSort.

Cuturi et al. (2019) propose an entropy regularized optimal transport formulation of the sorting operation. They solve this by applying the Sinkhorn algorithm (Cuturi, 2013) and produce gradients via automatic differentiation rather than the implicit function theorem, which resolves the need of solving a linear equation system. As the Sinkhorn algorithm produces a relaxed permutation matrix, we can also apply Sinkhorn sort to top-$k$ classification learning.

Petersen et al. (2021a) propose differentiable sorting networks, a continuous relaxation of sorting networks. Sorting networks are a kind of sorting algorithm that consist of wires carrying the values and comparators, which swap the values on two wires if they are not in the desired order. Sorting networks can be made differentiable by perturbing the values on the wires in each layer of the sorting network by a logistic distribution, i.e., instead of $\min$ and $\max$ they use $\operatorname{softmin}$ and $\operatorname{softmax}$. Similar to the methods above, this method produces a relaxed permutation matrix, which allows us to apply it to top-$k$ classification learning. The method has also been improved by enforcing monotonicity and bounding the approximation error (Petersen et al., 2022). Note that sorting networks are a classic algorithmic concept (Knuth, 1998), are not neural networks nor refer to differentiable sorting. Differentiable sorting networks are one of multiple differentiable sorting and ranking methods.

The discussed differentiable sorting algorithms produce relaxed permutation matrices of size $n\times n$. However, for top-$k$ classification learning, we require only the top $k$ rows for the number $k$ of top-ranked classes to consider. Here, $k$ is the largest $k$ that is considered for the objective, i.e., where $P_{K}(k)>0$. As $n\gg k$, producing a $k\times n$ matrix instead of a $n\times n$ matrix is much faster.

For NeuralSort and SoftSort, it is possible to simply compute only the top rows, as the algorithm is defined row-wise.

For the differentiable Sinkhorn sorting algorithm, it is not directly possible to improve the runtime, as in each Sinkhorn iteration the full matrix is required. Xie et al. (2020b) proposed a Sinkhorn-based differentiable top-$k$ operator, which computes a $2\times n$ matrix where the first row corresponds to the top-$k$ elements and the second row correspond to the remaining elements. However, this formulation does not produce $\bf{P}$ and does not distinguish between the placements of the top-$k$ elements among each other, and thus we use the SinkhornSort algorithm by Cuturi et al. (2019).

For differentiable sorting networks, it is (via a bi-directional evaluation) possible to reduce the cost from $\mathcal{O}(n^{2}\log^{2}(n))$ to $\mathcal{O}(nk\log^{2}(n))$. Here, it is important to note the shape and order of multiplications for obtaining $P$. As we only need those elements, which are (after the last layer of the sorting network) at the top $k$ ranks that we want to consider, we can omit all remaining rows of the permutation matrix of the last layer (layer $t$) and thus it is only of size $(k\times n)$.

Note that during execution of the sorting network, $P$ is conventionally computed from layer $1$ to layer $t$, i.e., from right to left. If we computed it in this order, we would only save a tiny fraction of the computational cost and only during the last layer. Thus, we propose to execute the differentiable sorting network, save the values that populate the (sparse) $n\times n$ layer-wise permutation matrices, and compute $P$ in a second pass from the back to the front, i.e., from layer $t$ to layer $1$, or from left to right in Equation 3. This allows executing $t$ dense-sparse matrix multiplications with dense $k\times n$ matrices and sparse $n\times n$ matrices instead of dense $n\times n$ and sparse $n\times n$ matrices. With this, we reduce the asymptotic complexity from $\mathcal{O}(n^{2}\log^{2}(n))$ to $\mathcal{O}(nk\log^{2}(n))$.

As only the top-$k$ rows of a relaxed permutation matrix are required for top-$k$ classification learning, it is possible to improve the efficiency of computing the top-$k$ probability distribution via differentiable sorting networks by reducing the number of differentiable layers and comparators. Thus, we propose differentiable top-$k$ networks, which relax selection networks in analogy to how differentiable sorting networks relax sorting networks. Selection networks are networks that select only the top-$k$ out of $n$ elements (Knuth, 1998). We propose splitter selection networks (SSN), a novel class of selection networks that requires only $\mathcal{O}(\log n)$ layers (instead of the $\mathcal{O}(\log^{2}n)$ layers for sorting networks) which makes top-$k$ supervision with differentiable top-$k$ networks more efficient and reduces the error (which is introduced in each layer.) SSNs follow the idea that the input is split into locally sorted sublists and then all wires that are not candidates to be among the global top-$k$ can be eliminated. For example, for $n=1024,k=5$, SSNs require only $22$ layers, while the best previous selection network requires $34$ layers and full sorting (with a bitonic network) requires even $55$ layers. For $n=10450,k=5$ (i.e., for ImageNet-21K-P), SNNs require $27$ layers, the best previous requires $50$ layers, and full sorting requires $105$ layers. In addition, the layers of SSNs are less computationally expensive than those of the bitonic sorting network. Details on SSNs, as well as their full construction, can be found in Supplementary Material B. Concluding, the contribution of differentiable top-$k$ networks is two-fold: first, we propose a novel kind of selection networks that needs fewer layers, and second, we relax those similarly to differentiable sorting networks.

Despite those performance improvements, evaluating the differentiable ranking operators still requires a considerable amount of computational effort for large numbers of classes. Especially if the number $n$ of elements to be ranked is $n=1\,000$ (ImageNet-1K) or even $n>10\,000$ (ImageNet-21K-P), the differentiable ranking operators can dominate the overall computational costs. In addition, for large numbers $n$ of elements to be ranked, the performance of differentiable ranking operators decreases as differentially ranking more elements naturally introduces larger errors (Grover et al., 2019; Prillo & Eisenschlos, 2020; Cuturi et al., 2019; Petersen et al., 2021a). Thus, we reduce the number of outputs to be ranked differentially by only considering those classes (for each input) that have a score among the top-$m$ scores. For this, we make sure that the ground truth class is among those top-$m$ scores, by replacing the lowest of the top-$m$ scores by the ground truth class, if necessary. For $n=1000$, we choose $m=16$, and for $n>10\,000$, we choose $m=50$. We find that this greatly improves training performance.

Because the differentiable ranking operators are (by their nature of being differentiable) only approximations to the hard ranking operator, they each have their characteristics and inconsistencies. Thus, for training models from scratch, we replace the top-$1$ component of the loss by the regular softmax, which has a better and more consistent behavior. This guides the other loss if the differentiable ranking operator behaves inconsistently. To avoid the top-$k$ components affecting the guiding softmax component and avoid probabilities greater than $1$ in $p$, we can separate the cross-entropy into a mixture of the softmax cross-entropy ($\mathrm{sm}$, for the top-$1$ component) and the top-$k$ cross-entropy ($\mathrm{top}-k$, for the top-$k\geq 2$ components) as follows:

Grover et al. (2019) include an experiment where they use the NeuralSort differentiable top-$k$ operator for $k$NN learning. Cuturi et al. (2019), Blondel et al. (2020), and Petersen et al. (2021a) each apply their differentiable sorting and ranking methods to top-$k$ supervision with $k=1$.

Xie et al. (2020b) propose a differentiable top-$k$ operator based on optimal transport and the Sinkhorn algorithm (Cuturi, 2013). They apply their method to $k$-nearest-neighbor learning ($k$NN), differential beam search with sorted soft top-$k$, and top-$k$ attention for machine translation. Cordonnier et al. (2021) use perturbed optimizers (Berthet et al., 2020) to derive a differentiable top-$k$ operator, which they use for differentiable image patch selection. Lee et al. (2021) propose using NeuralSort for a differentiable top-$k$ operator to produce differentiable ranking metrics for recommender systems. Goyal et al. (2018) propose a continuous top-$k$ operator for differentiable beam search. Pietruszka et al. (2020) propose the differentiable successive halving top-$k$ operator to approximate the normalized Chamfer Cosine Similarity ($nCCS@k$).

Fan et al. (2017) propose the “average top-$k$” loss, an aggregate loss that averages over the $k$ largest individual losses of a training data set. They apply this aggregate loss to SVMs for classification tasks. Note that this is not a differentiable top-$k$ loss in the sense of this work. Instead, the top-$k$ is not differentiable and used for deciding which data points’ losses are aggregated into the loss.

Lapin et al. (2015, 2016) propose relaxed top-$k$ surrogate error functions for multiclass SVMs. Inspired by learning-to-rank losses, they propose top-$k$ calibration, a top-$k$ hinge loss, a top-$k$ entropy loss, as well as a truncated top-$k$ entropy loss. They apply their method to multiclass SVMs and learn via stochastic dual coordinate ascent (SDCA).

Berrada et al. (2018) build on these ideas and propose smooth loss functions for deep top-$k$ classification. Their surrogate top-$k$ loss achieves good performance on the CIFAR-100 and ImageNet1K tasks. While their method does not improve performance on the raw data sets in comparison to the strong Softmax Cross-Entropy baseline, in settings of label noise and data set subsets, they improve classification accuracy. Specifically, with label noise of $20\%$ or more on CIFAR-100, they improve top-$1$ and top-$5$ accuracy and for subsets of ImageNet1K of up to $50\%$ they improve top-$5$ accuracy. This work is closest to ours in the sense that our goal is to improve learning of neural networks. However, in contrast to (Berrada et al., 2018), our method improves classification accuracy in unmodified settings. In our experiments, for the special case of $k$ being a concrete integer and not being drawn from a distribution, we provide comparisons to the smooth top-$k$ surrogate loss.

Yang & Koyejo (2020) provide a theoretical analysis of top-$k$ surrogate losses as well as produce a new surrogate top-$k$ loss, which they evaluate in synthetic data experiments.

A related idea is set-valued classification, where a set of labels is predicted. We refer to Chzhen et al. (2021) for an extensive overview. We note that our goal is not to predict a set of labels, but instead we return a score for each class corresponding to a ranking, where only one class can correspond to the ground truth.

Previous selection networks have been proposed by, i.a., (Wah & Chen, 1984; Zazon-Ivry & Codish, 2012; Karpiński & Piotrów, 2015). All of these are based on classic divide-and-conquer sorting networks, which recursively sort subsequences and merge them. In selection networks, during merging, only the top-$k$ elements are merged instead of the full (sorted) subsequences. In comparison to those earlier works, we propose a new class of selection networks, which achieve tighter bounds (for $k\ll n$), and relax them.

We evaluate the proposed top-$k$ classification loss for four differentiable ranking operators on CIFAR-100, ImageNet-1K, as well as the winter 2021 edition of ImageNet-21K-P. We use CIFAR-100, which can be considered a small-scale data set with only $100$ classes, to train a ResNet18 model (He et al., 2016) from scratch and show the impact of the proposed loss function on the top-$1$ and top-$5$ accuracy. In comparison, ImageNet-1K and ImageNet-21K-P provide rather large-scale data sets with $1\,000$ and $10\,450$ classes, respectively. To avoid the unreasonable carbon-footprint of training many models from scratch, we decided to exclusively use publicly available backbones for all ImageNet experiments. This has the additional benefit of allowing more settings, making our work easily reproducible, and allowing to perform multiple runs on different seeds to improve the statistical significance of the results. For ImageNet-1K, we use two publicly available state-of-the-art architectures as backbones: First, the (four) ResNeXt-101 WSL architectures by Mahajan et al. (2018), which were pretrained in a weakly-supervised fashion on a billion-scale data set from Instagram. Second, the Noisy Student EfficientNet-L2 (Xie et al., 2020a), which was pretrained on the unlabeled JFT-300M data set (Sun et al., 2017). For ResNeXt-101 WSL, we extract $2\,048$-dimensional embeddings and for the Noisy Student EfficientNet-L2, we extract $5\,504$-dimensional embeddings of ImageNet-1K and fine-tune on them.

We apply the proposed loss in combination with various available differentiable sorting and ranking approaches, namely NeuralSort, SoftSort, SinkhornSort, and DiffSortNets. To determine the optimal temperature for each differentiable sorting method, we perform a grid search at a resolution of factor $2$. For training, we use the Adam optimizer (Kingma & Ba, 2015). For training on CIFAR-100 from scratch, we train for up to $200$ epochs with a batch size of $100$ at a learning rate of $10^{-3}$. For ImageNet-1K, we train for up to $100$ epochs at a batch size of $500$ and a learning rate of $10^{-4.5}$. For ImageNet-21K-P, we train for up to $40$ epochs at a batch size of $500$ and a learning rate of $10^{-4}$. We use early stopping and found that these settings lead to convergence in all settings. As baselines, we use the respective original models, softmax cross-entropy, as well as learning with the smooth surrogate top-$k$ loss (Berrada et al., 2018).

We start by demonstrating that the proposed loss can be used to train a network from scratch. As a reference baseline, we train a ResNet18 from scratch on CIFAR-100. In Table 1, we compare the baselines (i.e., top-$1$ softmax, the smooth top-$5$ loss (Berrada et al., 2018), as well as “pure” top-$5$ losses using four differentiable sorting and ranking methods) with our top-$k$ loss with $k\sim[.2,.2,.2,.2,.2]$.

We find that training with top-$5$ alone—in some cases—slightly improves the top-$5$ but has a substantially worse top-$1$ accuracy. Here, we note that the smooth top-$5$ loss (Berrada et al., 2018), top-$5$ Sinkhorn (Cuturi et al., 2019), and top-$5$ DiffSort (Petersen et al., 2021a) are able to achieve good performance. Notably, Sinkhorn (Cuturi et al., 2019) outperforms the softmax baseline on the top-$5$ metric, while NeuralSort and SoftSort are less stable and yield worse results especially on top-$1$ accuracy.

By using our loss that corresponds to drawing $k$ from $[.2,.2,.2,.2,.2]$, we can achieve substantially improved results, especially also on the top-$1$ accuracy metric. Using the DiffSortNets yields the best results on the top-$1$ accuracy and Sinkhorn yields the best results on the top-$5$ accuracy. Note that, here, also NeuralSort and SoftSort achieve good results in this setting, which can be attributed to our loss with $k\sim[.2,.2,.2,.2,.2]$ being more robust to inconsistencies and outliers in the used differentiable sorting method. Interestingly, top-$5$ SinkhornSort achieves the best performance on the top-$5$ metric, which suggests that SinkhornSort is a very robust differentiable sorting method as it does not require additional top-$k$ components. Nevertheless, it is advisable to include other top-$k$ components as the model trained purely on top-$5$ exhibits poor top-$1$ performance.

In this section, we discuss the results for fine-tuning on ImageNet-1K and ImageNet-21K-P. In Table 2, we find a very similar behavior to training from scratch on CIFAR-100. Specifically, we find that training accuracies improve by drawing $k$ from a distribution. An exception is (again) SinkhornSort, where focussing only on top-$5$ yields the best top-$5$ accuracy on ImageNet-1K, but the respective model exhibits poor top-$1$ accuracy. Overall, we find that drawing $k$ from a distribution improves performance in all cases.

To demonstrate that the improvements also translate to different backbones, we show the improvements on all four model sizes of ResNeXt-101 WSL (32x8d, 32x16d, 32x32d, 32x48d) in Figure 3. Also, here, our method improves the model in all settings.

We start by demonstrating the impact of $P_{K}$, which is the distribution from which we draw $k$. Let us first consider the case where $k$ is $5$ with probability $\alpha$ and $1$ with probability $1-\alpha$, i.e., $P_{K}=[1-\alpha,0,0,0,\alpha]$. In Figure 2 (left), we demonstrate the impact that changing $\alpha$, i.e., transitioning from a pure top-$1$ loss to a pure top-$5$ loss, has on fine-tuning ResNeXt-101 WSL with our loss using the SinkhornSort algorithm. Increasing the weight of the top-$5$ component does not only increase the top-$5$ accuracy but also improves the top-$1$ accuracy up to around $60\%$ top-$5$; when using only $k=5$, the top-$1$ accuracy drastically decays as the incentive for the true class to be at the top-$1$ position vanishes (or is only indirectly given by being among the top-$5$.) While the top-$5$ accuracy in this plot is best for a pure top-$5$ loss, this generally only applies to the Sinkhorn algorithm and overall training is more stable if a pure top-$5$ is avoided. This can also be seen in Tables 1 and 2.

In Tables 3 and 4, we consider more additional settings with all differentiable ranking methods. Specifically, we compare four notable settings: $[.5,0,0,0,.5]$, i.e., equally weighted top-$1$ and top-$5$; $[.25,0,0,0,.75]$ and $[.1,0,0,0,.9]$, i.e., top-$5$ has larger weights; $[.2,.2,.2,.2,.2]$, i.e., the case of having an equal weight of $0.2$ for top-$1$ to top-$5$. The $[.5,0,0,0,.5]$ setting is a rather canonical setting which usually performs well on both metrics, while the others tend to favor top-$5$. In the $[.5,0,0,0,.5]$ setting, all sorting methods improve upon the softmax baseline on both top-$1$ and top-$5$ accuracy. When increasing the weight of the top-$5$ component, the top-$5$ generally improves while top-$1$ decays.

Here we find a core insight of this paper: the best performance cannot be achieved by optimizing top-$k$ for only a single $k$, but instead, drawing $k$ from a distribution improves performance on all metrics.

Comparing the differentiable ranking methods, we can find the overall trend that SoftSort outperforms NeuralSort, and that SinkhornSort as well as DiffSortNets perform best. We can see that some sorting algorithms are more sensitive to the overall $P_{K}$ than others: Whereas SinkhornSort (Cuturi et al., 2019) and DiffSortNets (Petersen et al., 2021a) continuously outperform the softmax baseline, NeuralSort (Grover et al., 2019) and SoftSort (Prillo & Eisenschlos, 2020) tend to collapse when over-weighting the top-$5$ components.

Comparing the performance on the medium-scale ImageNet-1K to the larger ImageNet-21K-P in Table 2, we observe a similar pattern. Here, again, using the top-$k$ component alone is not enough to significantly increase accuracy, but combining top-$1$ and top-$k$ components helps to improve accuracy on both reported metrics. While NeuralSort struggles in this large-scale ranking problem and stays below the softmax baseline, DiffSortNets (Petersen et al., 2021a) provide the best top-$1$ and top-$5$ accuracy with $40.22\%$ and $70.88\%$, respectively.

In Supplementary Material A, an extension to learning with top-$10$ and top-$20$ components can be found.

We note that we do not claim that all settings (especially all differentiable sorting methods) improve the classification performance on all metrics. Instead, we include all methods and also additional settings to demonstrate the capabilities and limitations of each differentiable sorting method.

Overall, it is notable that SinkhornSort achieves the overall most robust training behavior, while also being by far the slowest sorting method and thus potentially slowing down training drastically, especially when the task is only fine-tuning. SinkhornSort tends to require more Sinkhorn iterations towards the end of training. DiffSortNets are considerably faster, especially, it is possible to only compute the top-$k$ probability matrices and because of our advances for more efficient selection networks.

We consider how accuracy is affected by varying the number of scores $m$ to be differentially ranked. Generally, the runtime of differentiable top-$k$ operators depends between linearly and cubic on $m$; thus it is important to choose an adequate value for $m$. The choice of $m$ between $10$ and $40$ has only a moderate impact on the accuracy as can be seen in Figure 2 (right). However, when setting $m$ to large values such as $1\,000$ or larger, we observe that the differentiable sorting methods tend to become unstable. We note that we did not specifically tune $m$, and that better performance can be achieved by fine-tuning $m$, as displayed in the plot.

We compare the proposed results to current state-of-the-art methods in Table 5. We focus on methods that are publicly available and build upon two of the best performing models, namely Noisy Student EfficientNet-L2 (Xie et al., 2020a), and ResNeXt-101 32x48d WSL (Mahajan et al., 2018). Using both backbones, we achieve improvements on both metrics, and when fine-tuning on the Noisy Student EfficientNet-L2, we achieve a new state-of-the-art for publicly available models.

This work was supported by the IBM-MIT Watson AI Lab, the DFG in the Cluster of Excellence EXC 2117 “Centre for the Advanced Study of Collective Behaviour” (Project-ID 390829875), and the Land Salzburg within the WISS 2025 project IDA-Lab (20102-F1901166-KZP and 20204-WISS/225/197-2019).