D
IFFERENTIABLE T OP - k C LASSIFICATION L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
The top- k classiﬁcation accuracy is one of the core metrics in machine learning.
Here, k is conventionally a positive integer, such as 1 or 5 .
In this work, we relax this assumption and propose to draw k from a probability distribution for training.
Combining this with recent advances in differentiable sorting and ranking, we propose a new family of differentiable top-
k cross-entropy classiﬁcation losses.
We ﬁnd that relaxing k does not only produce better top- 5 accuracies, but also makes models more robust, which leads to top- 1 accuracy improvements.
When ﬁnetuning publicly available ImageNet models, we achieve a new state-of-the-art on
ImageNet for publicly available models with an 88 .
37% top- 1 and a 98 .
68% topaccuracy.
I NTRODUCTION
Classiﬁcation is one of the core disciplines in machine learning and computer vision.
With the advent of classiﬁcation problems with hundreds or even thousands of classes, the top- k classiﬁcation accuracy has established itself as an important task, i.e., an algorithm can suggest k classes and one of them has to be the correct class.
Usually, models are trained to optimize the top- 1 accuracy and top- 5 etc. are used for evaluation only.
Some works (Lapin et al., 2016; Berrada et al., 2018) have challenged this idea and proposed top-
k losses, such as a smooth top- 5 margin loss.
These methods have demonstrated superior robustness over the established top- 1 softmax cross-entropy in presence of additional label noise (Berrada et al., 2018).
In standard classiﬁcation settings, however, these methods have so far not shown improvements over the established top- 1 softmax cross-entropy.
In this work, instead of selecting a single top- k metric such as top- 1 or top- 5 for deﬁning the loss, we propose to specify k to be drawn from a probability distribution P K , which may or may not depend on the conﬁdence of speciﬁc data points or on the class label.
Examples for distributions
P K are [ . 5 , 0 , 0 , 0 , .
5] ( 50% top- 1 and 50% top- 5 ), [ . 1 , 0 , 0 , 0 , .
9] ( 10% top- 1 and 90% top- 5 ), and [ . 2 , .
2] ( 20% top- k for each k from 1 to 5 ).
Note that, when k is drawn from a distribution, this is done sampling-free as we can compute the expectation value in closed form.
Conventionally, given scores returned by a neural network, softmax produces a probability distribution over the top- 1 rank.
Recent advances in differentiable sorting and ranking (Grover et al., 2019;
Prillo & Eisenschlos, 2020; Cuturi et al., 2019; Petersen et al., 2021) provide methods for generalizing this to probability distributions over all ranks represented by a matrix P .
Based on differentiable ranking, multiple differentiable top- k operators have recently been proposed.
They found applications in differentiable k -nearest neighbor, differentiable beam search, attention mechanisms, and differentiable image patch selection (Cordonnier et al., 2021).
In these areas, integrating differentiable top- k improved results considerably by creating a more natural end-to-end learning setting.
However, to date, none of the differentiable top- k operators have been employed as neural network losses for top- k classiﬁcation learning with k > 1 .
Building on differentiable sorting and ranking methods, we propose a new family of differentiable top- k classiﬁcation losses where k is drawn from a probability distribution.
We ﬁnd that our top- k losses improve not only top- k accuracies, but also top- 1 accuracy on multiple learning tasks.
We empirically evaluate our method using four differentiable sorting and ranking methods on the
CIFAR-100 (Krizhevsky et al., 2009), ImageNet-1K (Deng et al., 2009), and the ImageNet-21K-P (Ridnik et al., 2021) data sets.
Using CIFAR-100, we demonstrate the capabilities of our losses to train models from scratch.
On ImageNet-1K, we demonstrate that our losses are capable of ﬁne-tuning state-of-the-art models and achieve a new state-of-the-art for publicly available mod- els on both top- 1 and top- 5 accuracy.
We benchmark our method on multiple recent models and demonstrate that our proposed method consistently outperforms the baselines for the best two differentiable sorting and ranking methods.
With ImageNet-21K-P, where many classes overlap (but only one is the ground truth), we demonstrate that our losses are scalable to more than 10 000 classes and achieve improvements of over 1% with only last layer ﬁne-tuning.
Overall, while the performance improvements on ﬁne-tuning are rather limited (because we retrain only the classiﬁcation head), they are consistent and can be achieved without the large cost of training from scratch.
The absolute 0 . 2% improvement that we achieve on the ResNeXt-101 32x48d
WSL top- 5 accuracy corresponds to an error reduction by approximately 10% , and can be achieved at much less than the computational cost of (re-)training the full model in the ﬁrst place.
We summarize our contributions as follows:
• We derive a novel family of top- k cross-entropy losses and relax the assumption of a ﬁxed k . •
We ﬁnd that they improve both top- k and top- 1 accuracy.
• We demonstrate its scalability to more than 10 000 classes.
• We propose splitter selection nets, which require fewer layers than existing selection nets.
• We achieve new state-of-the-art results (for publicly available models) on ImageNet1K.
B ACKGROUND : D IFFERENTIABLE S ORTING AND R ANKING
We brieﬂy review NeuralSort, SoftSort, Optimal Transport Sort, and Differentiable Sorting Networks.
We omit the fast differentiable sorting and ranking method (Blondel et al., 2020b) as it does not provide relaxed permutation matrices / probability scores, but rather only sorted / ranked vectors.
2. N EURAL S ORT & S OFT S ORT
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 sufﬁcient to build on pairwise differences of elements of the vectors to be sorted instead of the cumulative sums.
They ﬁnd that SoftSort performs approximately equivalent in their experiments to NeuralSort.
2. O PTIMAL T RANSPORT / S INKHORN S ORT
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 the Sinkhorn sort to top- k classiﬁcation learning.
2. D IFFERENTIABLE S ORTING N ETWORKS
Petersen et al.
(2021) 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 of the sorting network in each layer of the sorting network by a logistic distribution, i.e., instead of min and max they use softmin and softmax .
Similar to the methods above, this method produces a relaxed permutation matrix, which allows us to apply it to top- k classiﬁcation learning.
Note that sorting networks are a classic algorithmic concept (Knuth, 1998a) are not neural networks nor refer to differentiable sorting.
Differentiable sorting networks are one of multiple differentiable sorting and ranking methods.
[Figure]
a i r p l a n e p a n d a h u
s k y rankm a mm a l g o l d fi s h is the probability for a class to be among the top-k where k is drawn from P K
T OP - k L EARNING
In this section, we start by introducing our objective, elaborate its exact formulation, and then build on differentiable sorting principles to efﬁciently approximate the objective.
A visual overview over the loss architecture is also given in Figure 1.
The goal of top- k learning is to extend the learning criterion from only accepting exact (top- 1 ) predictions to accepting k predictions among which the correct class has to be.
In its general form, for top- k learning, k may differ for each application, class, data point, or a combination thereof.
For example, one practitioner may want to rank 5 predictions and assign a score that depends on the rank of the true class among these ranked predictions, while, on the other hand, another practitioner may want to obtain 5 predictions but does not care about their order.
Yet another practitioner in image classiﬁcation may want to enforce a top- 1 accuracy on images from the “person” super-class, but resign to a top- 3 accuracy for the “animal” super-class, as it may have more ambiguities in classlabels.
We model this by a random variable K , following a distribution P K that describes the relative importance of different values k .
The discrete distribution P K is either a marginalized distribution for a given setting (such as the uniform distribution), or a conditional distribution for each class, data point, etc.
This allows the practitioner to specify a marginalized / conditional distribution k ∼ P K .
This generalizes the ideas of conventional top- 1 supervision (usually softmax cross-entropy) and top- k supervision for a k like k = 5 (usually based on surrogate top- k margin/hinge losses like (Lapin et al., 2016; Berrada et al., 2018)) and uniﬁes them.
The objective of top- k learning is maximizing the probability of accepted predictions of the model f Θ on data X, y ∼ D given marginal distribution P K (or conditional P K | X,y if it depends on the class y and/or data point X ).
In the following, P k,y is the predicted probability of y being the k th-best prediction for data point X .
[Equation]
To evaluate the probability of y to be the top- 1 prediction, we can simply use softmax y ( f Θ ( X )) .
However, k > 1 requires more consideration.
Here, we require probability scores P k,c for the k th prediction over classes c ∈ C , where (cid:80) nc =1 P k,c = 1 (i.e., it P is row stochastic) and ideally additionally (cid:80) nk
=1 P k,c = 1 (i.e., it P is also column stochastic and thus doubly stochastic.)
With this, we can optimize our model by minimizing the following loss
[Equation]
which is the cross entropy over the probabilities that the true class is among the top- k class for each possible k .
Note that (cid:80) nk =1 P K ( k )
= 1 .
To compute P k,c , we require a function mapping from a vector of real-valued scores to an (ideally) doubly stochastic matrix P .
The most suitable for this are the differentiable relaxations of the sorting and ranking functions, which produce differentiable permutation matrices P , which we introduced in Section 2.
We build on these approximations to propose instances of top- k learning losses and extend differentiable sorting networks to differentiable top-
k networks, as just ﬁnding the top- k scores is computationally cheaper than sorting all elements and reduces the approximation error.
3. T OP - k P ROBABILITY M ATRICES
The discussed differentiable sorting algorithms produce relaxed permutation matrices of size n × n .
However, for top- k classiﬁcation learning, we require only the top k rows for the number k of topranked classes to consider.
Here, k is the largest k that is considered for the objective, i.e., where
P K ( k ) > 0 . As n (cid:29) k , producing a k × n matrix instead of a n × n matrix is much faster.
For NeuralSort and SoftSort , it is possible to simply compute only the top rows, as the algorithm is deﬁned 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.
For differentiable sorting networks , it is (via a bi-directional evaluation) possible to reduce the cost from O ( n 2 log 2 ( n )) to 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 × n ) .
[Equation]
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 × 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 × n matrices and sparse n × n matrices instead of dense n × n and sparse n × n matrices.
With this, we reduce the asymptotic complexity from O ( n 2 log 2 ( n )) to O ( nk log 2 ( n )) .
3.1. D IFFERENTIABLE T OP - k N ETWORKS
As only the top- k rows of a relaxed permutation matrix are required for top- k classiﬁcation learning, it is possible to improve the efﬁciency 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, 1998b).
We propose splitter selection networks (SSN), a novel class of selection networks that requires only O (log n ) layers (instead of the O (log 2 n ) layers for sorting networks) which makes top- k supervision with differentiable top- k networks more efﬁcient 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: ﬁrst, we propose a novel kind of selection networks that needs fewer layers, and second, we relax those similarly to differentiable sorting networks.
3. I MPLEMENTATION D ETAILS
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 (ImageNet21K-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., 2021).
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 ﬁnd 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 (for the top- 1 component) and the top- k cross-entropy (for the top- k ≥ 2 components) as follows:
[Equation]
[Table]  R ELATED W ORK
We structure the related work into three broad sections: works that derive and apply differentiable top- k operators, works that use ranking and top- k training objectives in general, and works that present classic selection networks.
4. D IFFERENTIABLE T OP - k O PERATORS
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.
(2020b), and Petersen et al.
(2021) 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 (Blondel et al., 2020a) to derive a differentiable top- k operator, which they use for differentiable image patch selection.
Lee et al.
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 ).
4. R ANKING AND T OP - k T RAINING O BJECTIVES
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 classiﬁcation 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.
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 classiﬁcation.
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 classiﬁcation accuracy.
Speciﬁcally, 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 classiﬁcation accuracy in unmodiﬁed 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 classiﬁcation, 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.
4. S ELECTION N ETWORKS
Previous selection networks have been proposed by, i.a., (Wah & Chen, 1984; Zazon-Ivry & Codish, 2012; Karpi ´ nski & Piotr ´ ow, 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 (cid:28) n ), and additionally we continuously relax them.
E XPERIMENTS 5. S ETUP
We evaluate the proposed top- k classiﬁcation loss for four differentiable ranking operators on
CIFAR-100, ImageNet-1K, as well as the winter 2021 edition of ImageNet-21K-P. CIFAR-100 may be considered a small-scale data set with only 100 classes.
We use CIFAR-100 to train a ResNet model (He et al., 2016) from scratch and show the impact of the proposed loss function on the topand 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 beneﬁt of allowing more settings, making our work easily reproducible, and allowing to perform multiple runs on different seeds to improve the statistical signiﬁcance of the results.
For ImageNet-1K, we use two publicly available state-ofthe-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 EfﬁcientNet-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 048dimensional embeddings and for the Noisy Student EfﬁcientNet-L2, we extract 5 504 -dimensional embeddings of ImageNet-1K and ﬁne-tune on them.
We apply the proposed loss in combination with various available differentiable sorting and ranking approaches, namely NeuralSort, SoftSort, SinkhornSort, and DiffSortNets.
The most important hyperparameter for each method is the temperature (or reparameterized as the inverse temperature).
To determine the optimal temperature for each approach via grid search at a resolution of factor 2 .
For training, we use the Adam optimizer (Kingma & Ba, 2015).
For training on CIFARfrom 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 .
For
ImageNet-21K-P, we train for up to 40 epochs at a batch size of 500 and a learning rate of 10 − 4 . [Table]
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)
5. C OMPARISON TO THE S TATE - OF - THE -A RT
[Figure] We compare the proposed results to current stateof-the-art methods in Table 1.
We focus on methods that are publicly available and build upon two of the best performing models, namely Noisy Student
EfﬁcientNet-L2 (Xie et al., 2020a), and ResNeXt101 32x48d WSL (Mahajan et al., 2018).
Using both backbones, we achieve improvements on both metrics, and when ﬁne-tuning on the Noisy Student
EfﬁcientNet-L2, we achieve a new state-of-the-art for publicly available models.
In Figure 2, we demonstrate our improvements on the four model sizes of ResNeXt-101 WSL (32x8d, 32x16d, 32x32d, 32x48d).
Our method improves the model in all settings.
Signiﬁcance Tests. To evaluate the signiﬁcance of the results, we perform a t-test (with signiﬁcance level of 0 . 01 ). We ﬁnd that our model is signiﬁcantly better than the original model on both topand top- 5 accuracy metrics. Comparing to the observed accuracies of the baseline ( 88 .
33 | 98 . 65 ),
DiffSortNets are signiﬁcantly better (p= 0 .
00001 | 0 . 00005 ).
Comparing to the reported accuracies of the baseline ( 88 . 35 | 98 . 65 ), DiffSortNets are also signiﬁcantly better (p= 0 .
00087 | 0 . 00005 ).
5. I MPACT OF THE D ISTRIBUTION P K AND D IFFERENTIABLE S ORTING M ETHODS
We start by demonstrating the impact of P K , which is the distribution from which we draw k .
Let us ﬁrst consider the case where k is 5 with probability α and 1 with probability 1 − α , i.e., P K = [Figure]
[1 − α, 0 , 0 , 0 , α ] .
In Figure 3 (left), we demonstrate the impact that changing α , i.e., transitioning from a pure top- 1 loss to a pure top- 5 loss, has on ﬁne-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 topposition vanishes (or is only indirectly given by being among the top- 5 .)
While the top- 5 accuracy in this plot is largest for a pure top- 5 loss, this generally only applies to the Sinkhorn algorithm and overall training is stable if a pure top- 5 is avoided.
This can also be seen on the left of Table 2.
In Table 2, we consider more additional settings with all diff.
ranking methods.
Speciﬁcally, 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] , 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 ﬁnd the main insight of this paper: the best performance cannot be achieved by optimizing top- k for only a single k , but instead relaxing this constraint improves performance for all metrics. [Table]
[Table] [Table] Comparing the differentiable ranking methods, we can ﬁnd 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., 2021) 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 3 we observe a similar pattern.
Here, again, using the top- k loss alone is not enough to signiﬁcantly increase accuracy, but combining both, top- 1 and top- k loss helps to improve accuracy on both reported metrics.
While NeuralSort struggles with a good loss for this large-scale ranking problem and stays below the softmax baseline,
DiffSortNets (Petersen et al., 2021) provide the best top- 1 and top- 5 accuracy with 40 .
22% and 70 .
88% , respectively.
We note that we do not claim that all settings (especially all differentiable sorting methods) improve the classiﬁcation performance.
Instead, we include all methods and also additional settings to demonstrate the capabilities and limitations of each differentiable sorting method.
It is notable that SinkhornSort overall achieves the 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 ﬁne-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 efﬁcient selection networks.
5. D IFFERENTIABLE R ANKING S ET S IZE m
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 3 (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 speciﬁcally tune m , and that better performance can be achieved by ﬁne-tuning m , as displayed in the plot.
5. T RAINING FROM S CRATCH
Finally, we demonstrate that the proposed loss can also be used to train a network from scratch.
As reference baseline, we train a ResNet18 from scratch on CIFAR-100, considering the same settings as above.
The results are shown in Table 4.
Again, we ﬁnd that training with top- k alone slightly improves the top- 5 but not the top- 1 accuracy, whereas a weighted combination of both losses can signiﬁcantly improve the performance on both metrics.
Notably, here, [ . 2 , . 2 , .
2] with DiffSortNets yields the best results on top- 1 accuracy.
In Supplementary Material A, an extension to learning with top- 10 and top- 20 components on
ImageNet-21K-P can be found.
C ONCLUSION
In this work, we presented a novel loss, which relaxes the assumption of using a ﬁxed k for top- k classiﬁcation learning.
For this, we leveraged recent differentiable sorting and ranking operators.
We performed an array of experiments to explore different top- k classiﬁcation learning settings and achieved a state-of-the-art on ImageNet for publicly available models.
R EPRODUCIBILITY S TATEMENT
The source code for the experiments will be made available upon publication.
We only use publicly available resources, and each experiment can be reproduced using small-scale hardware.
We used (wherever possible) multiple seeds for our experiments.
A E XTENSION TO T OP -10 AND T OP -
We further extend the training settings, measuring the impact of top- 10 and top- 20 components on the large-scale ImageNet-21K-P dataset.
The results are diplayed in Table
[Table]
where we report top- 1 , top- 5 , top-
[Table]
and top- 20 accuracy for all conﬁgurations.
Again, we observe that 50% top- 1 and 50% top- k produces the overall best performance and that training with top- 5 yields the best top- 1 , top-
[Table]
and top- 10 accuracy.
We observe that the performance decays for top- 20 components because (even among 10 450 classes) there are virtually no top- 20 ambiguities
[Table]
and artifacts of differentiable sorting methods can cause adverse effects.
Note that top- 10 ambiguities do exist in ImageNet-21K-P, e.g., there are 11 class hierarchy levels Ridnik et al.
[Table]
B S PLITTER S ELECTION N ETWORKS
Similar to a sorting network, a selection network is generally a comparator network and hence it consists of wires (or lanes) carrying values and comparators (or conditional swap devices) connecting pairs of wires.
A comparator swaps the values on the wires it connects if they are not in a desired order.
However, in contrast to a sorting network, which sorts all the values carried by its wires, a ( k, n ) selection network, which has n wires, moves the k ≤ n largest (or, alternatively, the k smallest) values to a speciﬁc set of wires (Knuth, 1998b), most conveniently consecutive wires on one side of the wire array.
Note that the notion of a selection network usually does not require that the selected values are sorted.
However, in our context it is preferable that they are, so that P K can easily be applied, and the selection networks discussed below all have this property.
Clearly, any sorting network could be used as a selection network, namely by focusing only on the top k (or bottom k ) wires.
However, especially if k is small compared to n , it is possible to construct selection networks with smaller size (i.e. fewer comparators) and often lower depth (i.e. a smaller number of layers, where a layer is a set of comparators that can be executed in parallel).
A core idea of constructing selection networks was proposed in (Wah & Chen, 1984), based on the odd-even merge and bitonic sorting networks (Batcher, 1968): partition the n wires into subsets of at least k wires (preferably 2 (cid:100) log 2 ( k ) (cid:101) wires per subset) and sort each subset with odd-even mergesort.
Then merge the (sorted) top k elements of each subsets with bitonic merge, thus halving the number of (sorted) subsets.
Repeat merging pairs of (sorted) subsets until only a single (sorted) subset remains, the top k elements of which are the desired selection.
This approach requires 1 2 (cid:100) log 2 ( k ) (cid:101) ( (cid:100) log 2 ( k ) (cid:101) + 1) + ( (cid:100) [Table] [Table]
Improvements to this basic scheme were developed in (Zazon-Ivry & Codish, 2012; Karpi ´ nski &
Piotr ´ ow, 2015) and either rely entirely on odd-even merge (Batcher, 1968) or entirely on pairwise sorting networks (Parberry, 1992).
Especially selection networks based on pairwise sorting networks have advantages in terms of the size of the resulting network (i.e. number of needed comparators).
However, these improvements do not change the depth of the networks, that is, the number of layers, which is most important in the context considered here.
Our own selection network construction draws on this work by focussing on a speciﬁc ingredient of pairwise sorting networks, namely a so-called splitter (which happens to be identical to a single bitonic merge layer, but for our purposes it is more comprehensible to refer to it as a splitter).
A splitter for a list of m wires having indices [ (cid:96) 0 , . . .
, (cid:96) m − 1 ] has comparators connecting wires (cid:96) i and (cid:96) i + s where s = (cid:100) log 2 ( m ) (cid:101) − 1 for i ∈ { 0 , . .
, m − 1 − 2 (cid:100) log 2 ( m ) (cid:101) ) − 1 } .
A pairwise sorting network starts with what we call a splitter cascade .
That is, an initial splitter partitions the input wires into subsets of (roughly) equal size.
Each subset is split recursively until wire singletons result (Zazon-Ivry & Codish, 2012).
An example of such a splitter cascade is shown in Figure 4 for 8 wires and in purple color for 16 wires in Figure 5 (arrows point to where the larger value is desired).
After a splitter cascade, the value carried by wire (cid:96) i has a minimum rank of r = 2 b ( i ) − 1 , where b ( i ) counts the number of set bits in the binary number representation of i .
This minimum rank results from the transitivity of the swap operations in the splitter cascade, as is illustrated in Figure 4 for 8 wires:
By following upward paths (in splitters to the left) through the splitter cascade, one can ﬁnd
[Table]
splitter selection n sort                                                                                        for each wire (cid:96) i exactly r = 2 b ( i ) − 1 wires with smaller indices that must carry values no less than the value carried by wire (cid:96) i .
This yields the minimum ranks shown in Figure 4 on the right.
The core idea of our selection network construction is to use splitter cascades to increase the minimum ranks of (the values carried by) wires.
If such a minimum rank exceeds k (or equals k , since we work with zero-based ranks and hence are interested in ranks { 0 , . .
, k − 1 } ), a wire can be discarded, since its value is certainly not among the top k .
On the other hand, if there is only one wire with minimum rank 0, the top 1 value has been determined.
More generally, if all minimum ranks no greater than some value r occur for one wire only, the top r + 1 values have been determined.
We exploit this as follows:
Initially all wires are assigned a minimum rank of 0, since at the beginning we do not know anything about the values they carry.
We then repeat the following construction: traversing the values r = k − 1 , . . .
, 0 descendingly, we collect for each r all wires with minimum rank r and apply a splitter cascade to them (provided there are at least two such wires).
Suppose the wires collected for a minimum rank r have indices [ (cid:96) 0 , . . .
, (cid:96) m ( r ) − 1 ] .
After the splitter cascade we can update the minimum rank of wire (cid:96) i to r + 2 b ( i ) − 1 , because before the splitter cascade there is no known relationship between wires with the same minimum rank, while the splitter cascade establishes relationships between them, increasing their ranks by 2 b ( i ) − 1 .
The procedure of traversing the minimum ranks k − 1 , . .
, 0 descendingly, collecting wires with the same minimum rank and applying splitter cascades to them is repeated until all minimum ranks 0 , . . .
, k − 1 occur only once.
As an example, Figure 5 shows a (5 , 16) selection network constructed is this manner, in which the minimum ranks of the wires are indicated after certain layers as well as when certain wires can be discarded (red crosses) and when certain top ranks are determined (green check marks).
Comparators belonging to the same splitter cascade are shown in the same color.
While selection networks resulting from adaptations of sorting networks (see above) have the advantage that they guarantee that their number of layers is never greater than that of a full sorting network, our approach may produce networks with more layers.
However, if k is sufﬁciently small compared to n (in particular, if k ≤ log 2 ( n ) ), our approach can produce selection networks with considerably fewer layers, as is demonstrated in Table 6.
Since in the context we consider here we can expect k ≤ log 2 ( n ) , splitter-based selection networks are often superior.