{'fraction_non_alphanumeric': 0.05642199393197424, 'fraction_numerical': 0.01724596795656571, 'mean_word_length': 4.131093814010652, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 5, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Doubly-stochastic mining for heterogeneous retrieval

Modern retrieval problems are characterised by training sets with potentially billions of labels, and heterogeneous data distributions across subpopulations (e.g., users of a retrieval system may be from different countries), each of which poses a challenge. The first challenge concerns scalability: with a large number of labels, standard losses are difficult to optimise even on a single example. The second challenge concerns uniformity: one ideally wants good performance on each subpopulation. While several solutions have been proposed to address the first challenge, the second challenge has received relatively less attention. In this paper, we propose doubly-stochastic mining (S2M ), a stochastic optimization technique that addresses both challenges. In each iteration of S2M, we compute a per-example loss based on a subset of hardest labels, and then compute the minibatch loss based on the hardest examples. We show theoretically and empirically that by focusing on the hardest examples, S2M ensures that all data subpopulations are modelled well.


Introduction
Information retrieval concerns finding documents that are most relevant for a given query, and is a canonical real-world use case for machine learning [Manning et al., 2008]. The simplest incarnation of retrieval models involves learning a real-valued scoring function that ranks, for each example, the set of possible labels it may be matched to. A core challenge is scalability: there may be billions of examples (e.g., user queries) and labels (e.g., videos in a recommendation system), each of whose scores naïvely needs to be updated at every training iteration. Effective means of addressing both problems have been widely studied [Mikolov et al., 2013, Jean et al., 2015.
A distinct challenge is heterogeneity: the distribution over examples is often a mixture of diverse subpopulations (e.g., queries may arise from geographically disparate user bases). Naïve training on such data may lead to models that perform disproportionately well on one subpopulation at the expense of others; e.g., if queries originate from multiple countries, the retrieval model may only perform well on queries from the dominant country. Such behaviour is clearly undesirable.
Unfortunately, the heterogeneity problem plagues (to our knowledge) all state-of-the-art retrieval models. Indeed, while there has been considerable progress on improving the efficacy of retrieval systems [Prabhu et al., 2018, Guo et al., 2019, any such method reliant on optimising examples drawn uniformly at random from the training set (e.g., via SGD) will inherently be biased towards the "dominant" subpopulations. This raises the natural question: Can one mitigate such bias without introducing significant computational overhead?
In SGD, we compute the average loss over each instance, which in turn averages over each of the labels. In SNM, we only retain labels with the k-highest contributions to the per-example losses; here, k = 3. In top-k SGD, we only retain examples with top-k highest loss; here k = 2. In S 2 M, one combines both of these, allowing tractable training under a large number of labels, as well as good performance on subpopulations.
(i) we propose doubly-stochastic mining (S 2 M), an algorithm that optimises a loss computed over the hardest examples and labels in a randomly drawn minibatch.
(ii) we establish that S 2 M controls the retrieval loss over latent subpopulations, drawing on a connection between our objective and the conditional value-at-risk (CVaR).
(iii) we show that empirically, S 2 M yields improved retrieval performance on subpopulations compared to state-of-the-art retrieval methods.
In more detail, given a minibatch of examples, S 2 M randomly draws a sample of labels, from which we compute a per-example loss based on the hardest labels in this subset; we then average the hardest per-example losses to form the minibatch loss. Algorithmically, this is a combination of two distinct proposals in the literature (see Figure 1): the stochastic negative mining (SNM ) algorithm of , and q-SGD [Kawaguchi and Lu, 2019], which focus on the hardest labels and examples, respectively.
While combining these two algorithms is conceptually straightforward, it is far less obvious that the result can cope with the heterogeneity problem. Contribution (ii) exploits two non-apparent linksthat of the S 2 M objective to the conditional value-at-risk (CVaR) [Rockafellar and Uryasev, 2000], and of the CVaR to a worst-case loss over subpopulations -to establish that by focusing on the hardest examples, S 2 M ensures good performance on latent data subpopulations. Our analysis integrates recent results from both the fairness [Mohri et al., 2019] and the operations research [Cherukuri, 2019] literature.
The idea of focusing on the hardest examples has more broadly been explored in several different contexts. For example, minimising the maximal loss forms the basis for the hard-margin SVM [Vapnik, 1999], and has more broadly been studied in Shalev-Shwartz and Wexler [2016]. However, these works do not consider jointly sampling instances and labels, and are thus not attuned to the retrieval setting of interest. Active learning also concerns finding "hard" examples to label [Dagan and Engelson, 1995], but in an interactive setting distinct to our passive retrieval setting.

Background and notation
We review the multiclass retrieval setting of interest 1 , and approaches to address the underlying challenges.

Multiclass retrieval
In multiclass classification, we observe instances x ∈ X with associated labels y ∈ Y = [K] . = {1, 2, . . . , K} drawn from some distribution D over X × Y. We aim to learn a scorer f : X → R K which can order the labels by their relevance for an instance. In a retrieval setting, our goal is to minimise the retrieval risk, which for an integer r is where ret(r) (i, f ) . = i / ∈ top r (f ) , and top r (f ) denotes the r indices with highest score under f . When r = 1, this reduces to the misclassification risk L 01 (f ; D).
In practice, directly minimising L ret(r) (f ; D) is computationally prohibitive owing to the nondifferentiability of the loss ret(r) . This can be alleviated by minimising a suitable surrogate loss Examples of include the softmax cross-entropy (y, f (x)) = −f y (x) + log y e f y (x) , and binary ordered weighted losses (BOWLs)  (y, f (x)) = φ(f y (x)) + avg top(k) ({φ(−f y (x))} y =y ) , where φ : R → R + is a margin loss (e.g., hinge φ(z) = [1 − z] + ), k is an integer, and avg top(k) denotes the average of the top-k largest elements of a vector. When k = 1, this is the Cramer-Singer loss [Crammer and Singer, 2002], This provides a tight bound on L ret(1) (f ; D). When k = K − 1, this is the averaged loss [Zhang, 2004b], For a training set S ∼ D N and surrogate loss , one may seek to minimise the empirical risk

Challenges in multiclass retrieval
Multiclass retrieval poses several challenges. First, the number of samples N can potentially be in the order of billions, thus making computing L avg (f ; S) prohibitive [Bottou and Bousquet, 2007]. Second, the number of labels K may also be in the order of billions, thus making even computing a single (y i , f (x i )) prohibitive [Agrawal et al., 2013, Yu et al., 2014, Bhatia et al., 2015, Jain et al., 2016, Babbar and Schölkopf, 2017, Prabhu et al., 2018, Jain et al., 2019; e.g., the softmax cross-entropy requires computing log y ∈[L] e f y (xi) , which is expensive for large K.
Third, the distribution D typically comprises heterogeneous subpopulations; e.g., the users of the retrieval system may be from different countries. Formally, suppose D = p∈[P ] ν p · D p for P distributions {D p } p∈[P ] with mixture weights ν ∈ ∆ P , where ∆ P denotes the simplex. Our goal is to ensure good performance for each D p .
We now review proposals to deal with each challenge.

Stochastic gradient descent and negative mining
When N is large, one may perform minibatch stochastic gradient descent (SGD) [Robbins and Monro, 1951] by sampling a minibatch S mb = {(x j , y j )} N mb j=1 ⊆ S, and performing descent based on the average minibatch loss: When K is large, however, it is challenging to use SGD since even computing the minibatch loss may be prohibitive, let alone optimising it. As noted above, to compute the loss on even a single example, the softmax cross-entropy as well as the BOWL losses in (2) require the scores of all K labels, which may not be feasible.
One approach to cope with this problem is to sample the labels, and use a resulting stochastic approximation to the loss [Bengio and Senecal, 2008, Jean et al., 2015, Grave et al., 2017. In stochastic negative mining (SNM) , in addition to the minibatch of random examples, for each example one draws a random subset of "negative" labelsȲ ⊆ Y − {y} of sizē K K. One can then treatȲ as the entire label space, and compute a new BOWL snm (y, f (x);Ȳ) = φ(f y (x)) + avg top(k) ({φ(−f y (x))} y ∈Ȳ ).
As a simple example, for k = 1, where, by contrast to (3), the maximum is only over the labels inȲ. A key property of SNM is that one only needs to update O(K) labels' parameters, which can be a significant saving compared to updating the parameters for all K labels.
Observe that snm is stochastic, owing toȲ being random. The expected loss is itself a BOWL, which is calibrated [Zhang, 2004b] under mild conditions.

Maximum empirical loss
Both SGD and SNM take minimising (5) to be their basic goal. This involves minimising the average loss over the training set. One may however replace the average with other summaries, such as the maximum: Shalev-Shwartz and Wexler [2016] established that the minimiser of (9) guarantees good performance on subpopulations. Intuitively, (9) encourages predicting well on all examples, including those from a rare subpopulation. The authors also proposed a means of stochastically optimising the objective, but this requires O(N ) auxiliary variables, which can be prohibitive in large-scale settings.

Method
Loss on minibatch Large N ? Large K? Heterogenity?
SGD [Robbins and Monro, 1951] avg  avg φ(f y (x)) + avg top(k) {φ(−f y )} y ∈Ȳ × Top-k SGD [Kawaguchi and Lu, 2019] Table 2: Comparison of losses employed in stochastic gradient descent (SGD), stochastic negative mining (SNM), average top-k SGD (top-k SGD), and our proposed doubly stochastic mining (S 2 M). Here, S mb denotes a random minibatch of (instance, label) pairs (x, y), and f (x) ∈ R K the model predictions for each of K possible classes. We assume the use of a base convex loss φ : R → R + used to construct a multiclass loss per (2). In SGD, one computes the average loss over the minibatch, which implicitly averages over all labels. In SNM, one draws a random batch of "negative" labelsȲ ⊆ Y − {y}, and only operates on the top-k highest scoring labels, denoted by the avg top(k) (·) operation. In top-k SGD, one only operates on the top-k highest per-example losses, but for each per-example loss averages over all labels. In S 2 M, one combines both of these, with the aim of allowing tractable training under a large K, while ensuring good performance on heterogeneous data subpopulations.
3 Doubly-stochastic mining (S 2 M) We now present S 2 M, our doubly-stochastic mining algorithm, which addresses all three challenges discussed above: S 2 M scales to settings with a large number of examples, labels, and performs well on heterogeneous subpopulations.

Doubly-stochastic mining
Reviewing existing multiclass retrieval methods, we see that they fail to meet one or more of the challenges in §2.2 (see Table 2 for a summary); e.g., while SNM handles large N and K, it does not adapt to heterogeneous distributions.
A natural question is whether we can extend SNM to rectify this. Following §2.4, we take inspiration from Shalev-Shwartz and Wexler [2016] and move from average to the maximal empirical loss. Indeed, given a minibatch S mb = {(x j , y j )} N mb j=1 , the analogue of (9) is: Interestingly, this can be seen as an example level counterpart to the SNM loss of (8): in the latter, we compute a per-example loss snm (y, f (x)) via the maximal loss over all labels inȲ, while in (10), we compute the minibatch loss L max (f ; S mb ) via the maximal loss over all examples in the minibatch.
To further highlight this, suppose we replace the maximum with the average of the k largest losses: This is in contrast to (7), which obtains a per-example loss by averaging over the k labels with the highest contribution to the loss value. By combining the SNM loss of (7) with the top-k minibatch loss of (11), we have a doubly stochastic procedure which handles both large K and heterogeneous example distributions. Intuitively, this should behave similarly to the maximum loss in (10) in controlling subpopulation performance while being more noise-robust.
Formally, our doubly-stochastic mining (S 2 M) algorithm 2 with a training set S ∼ D N is summarised in Algorithm 1. In a nutshell, the procedure is as follows: first, following SGD, we draw a random minibatch S mb ⊆ S. Next, following SNM, for a given example (x, y) ∈ S mb , we draw a random sample ofK K labelsȲ ⊆ Y − {y}, and compute per-example loss snm (y, f (x)), per (7); this will only focus on the hardest k labels withinK. Given { snm (y, f (x);Ȳ) : (x, y) ∈ S mb }, we average the top-k loss values, and use this as our minibatch loss, per (11).
Algorithm 1 Doubly-stochastic mining (S 2 M) , training steps T , minibatch size N mb , label sample sizeK, k, k ∈ N+ 2: for t = 1, 2, . . . , T do 3: We emphasise that S 2 M involves two forms of top-k averaging, with different purposes. In constructing the SNM loss snm , we only average over the top-k highest scoring labels. This is motivated by yielding both computational tractability, and a tight bound to the retrieval loss. In constructing the final minibatch loss, we only average over the top-k highest per-sample losses. This is motivated by controlling the loss over subpopulations; we formalise this in §4.
Further,K and k may be picked following : to ensure good retrieval performance, one can pickK based on computational considerations, and k = r for the retrieval threshold r in (1). The choice of k is more subtle: the discussion following Proposition 5 in §4 reveals that k should correspond to the number of samples from the rarest subpopulation one expects to see in S. In practice, this can be specified as a tuning parameter.

Special cases of S 2 M
S 2 M generalises several existing approaches: (a) whenK = K and k = N mb , we recover SGD.
(b) whenK K and k = N mb , we recover SNM.
(c) whenK = K and k N mb , we recover average top-k SGD [Kawaguchi and Lu, 2019].
Our interest is in focussing on both the hardest per-example losses (k N mb ), and the hardest labels within each such loss (k K K). For point (c) above, for the small K setting, Fan et al. [2017], Kawaguchi and Lu [2019] considered minimising the average top-k loss per (11), with the latter specifically focussed on the stochastic setting. However, these works did not consider sampling over labels, nor theoretically justify the ability of top-k losses to handle heterogeneous example distributions (as we shall do in §4).
Compared to SNM, S 2 M replaces the average of all per-example losses with the average of the top-k per-example losses. For a minibatch size of N mb , this step becomes O(N mb · log N mb ) rather than O(N mb ). Since N mb N is typically a small constant, the overhead is minimal.

Expected loss under S 2 M
Recall that in SGD, one works with the stochastic average minibatch loss L avg (f ; S mb ) per (6). A key property of this loss is that its expectation is precisely the quantity we wish to optimise, namely, the empirical risk L avg (f ; S).
In S 2 M, our minibatch loss (11) has two sources of stochasticity, as we draw both a minibatch and a label sample. Further, the minibatch loss only keeps the top-ranked loss values, which makes its expected behaviour more subtle. As with SNM and average top-k SGD, this expected loss involves the order-weighted average [Usunier et al., 2009] of a vector, owa(z; θ) denotes the ith largest element of the vector z. We have the following.
Lemma 1. Pick any φ : R → R + and k , with induced SNM loss snm per (7). Then, the expected S 2 M loss (cf. (11)) over the draw of minibatch S mb and label setȲ is Concretely, suppose k = 1, i.e., we compute the maximal loss over the minibatch. By Lemma 1, the expected loss is not the maximum loss i (9), but rather an order-weighted average of the per-example losses. We thus place a low, but non-zero, weight on smaller per-example losses.
We thus see that S 2 M involves a doubly order-weighted average: the expected loss emphasises the individual examples with high loss, where the per-example loss in turn emphasises the individual labels with high loss. Reddi et al.
[2019] establish that the latter guarantees a tight bound on the retrieval performance. In §4, we show that the former ensures good performance on subpopulations.

Consistency of average top-k for retrieval
A minimal requirement to impose on the expected loss in Lemma 1 is that it is consistent [Zhang, 2004a, Bartlett et al., 2006 for retrieval, i.e., that its minimisation also ensures that the retrieval risk is minimised. Here, we establish this when is the softmax cross-entropy, provided that the label distribution P Y|X satisfies certain regularity conditions and k is chosen suitably as the number of sample N increases.
Proposition 2. Pick the softmax cross-entropy loss . Let α N . Then, for any α > 0, there exists an integer r(α) depending on P Y|X , such that the average top-k (N ) loss with is consistent for top-r(α) retrieval.

Analysis of S 2 M under heterogeneity
We now establish that minimising the average top-k loss controls performance over heterogeneous subpopulations of the data. Thus, in retrieval settings, we ensure that good retrieval over different data subpopulations.
Formally, following §2.2, suppose the underlying data distribution comprises P distinct subpopulations {D p } p∈[P ] ; i.e., D = p∈[P ] ν * p · D p for some unknown mixture weights ν * ∈ ∆ P . Our aim is to control the following maximal retrieval loss across different subpopulations: where L ret(r) is the retrieval loss for integer r (cf. (1)). We make two related comments here. First, any D will admit an infinitude of mixture representations involving different D p 's. Second, performance on a given D p ought to depend on how well-represented it is in D, i.e., the magnitude of ν * p . This intuition will be manifest in our subsequent bounds.
For simplicity, we focus on the full batch setting, i.e., N mb = N . We also consider the use of a generic multiclass loss , which could be the expected SNM loss from Lemma 1. Our analysis easily generalises to the minibatch setting, using Lemma 1 and the techniques of Reddi et al.

Warm-up: known subpopulation membership
As a warm-up, we illustrate how we can control (12) if we know in advance which subpopulation each training example in S belongs to. For p ∈ [P ], let S p . = {i ∈ [N ] : (x i , y i ) ∼ D p } with N p := |S p | be the indices of the samples drawn from D p , the distribution of the pth subpopulation. Given a surrogate loss , e.g., the SNM loss of (7), the empirical loss for each subpopulation is To achieve good performance on each subpopulation, we may consider the agnostic empirical loss from Mohri et al. [2019], which considers the worst-case performance over all mixtures of subpopulations: for a mixture set Λ ⊆ ∆ P , Observe that if we pick Λ = ∆ P , then the maximum in (14) is exactly (12). Equipped with this, one may then appeal to generalisation bounds for the agnostic empirical loss [Mohri et al., 2019, Theorem 2] to conclude that with high probability over the draw of the training sample, where C F,Λtop,N depends on a weighted Rademacher complexity of the function class F, and covering number of the mixture set Λ top ; see Mohri et al. [2019] for details.
We now move to the main focus of this paper, the more common and and challenging setting where the subpopulation membership for the training samples is unknown. In this case, we cannot even compute the empirical loss over each subpopulation (13), let alone minimise the agnostic empirical loss (14) to ensure good performance over each subpopulation. Fortunately, we can establish that the proposed average top-k minimisation controls the loss in (12). Our analysis proceeds as follows: (i) first, we relate the top-k risk that S 2 M minimises to the conditional value at risk (CVaR) (ii) next, we further relate the population version of the CVaR to the maximal subpopulation loss in (12) (iii) finally, we derive a generalisation bound for the CVaR, which shows that the maximal retrieval loss can be controlled by minimising the empirical top-k risk.

From top-k to maximal retrieval loss via CVaR
We begin by introducing the conditional value at risk (CVaR) at level α ∈ (0, 1] of a random variable U, defined as [Rockafellar and Uryasev, 2000, Theorem 1] where [x] + . = max{x, 0}. Given iid samples u i ∼ U, we may construct an empirical estimate for α = k N . Here, CVaR α (Û) is a biased estimate of CVaR α (U), since the expectation of the infimum is not the same as the infimum of the expectation [Brown, 2007].
Interestingly, (16) is equivalent to the average of the top-k largest elements of {u i } N i=1 [Rockafellar and Uryasev, 2002, Proposition 8], [Fan et al., 2017, Lemma 2]. This gives a means of re-expressing the average top-k loss in terms of the CVaR. Specifically, consider the random variable of loss values L f . = (Y, f (X)) for (X, Y) ∼ D and given scorer f .
gives us an empirical estimate of CVaR α (L f ).
This connection between the empirical average top-k risk and CVar has twofold advantage. First, the representation in (15) enables us to derive generalisation bounds for the retrieval problem. Second, appealing to a distinct representation for the CVaR evinces how it controls subpopulation performance [Ben-Tal and Teboulle, 2007, Example 4.2]: where the uncertainty set B(P, α) = {Q ∈ ∆ U | (∃R ∈ ∆ U ) P = α · Q + (1 − α) · R} with ∆ U denoting the set of all distributions on the sample space U. The CVaR is thus the maximal expectation under any minority subpopulation of the original distribution with mass at least α. A similar connection was made in [Duchi et al., 2019, Eq. (2)].
The representation in (18) can be seen as a generalisation of the agnostic loss in (14), wherein there are subpopulation dependent mixing weights, and the empirical subpopulation losses are replaced with average top-k counterparts.
We may thus relate the top-k risk to a maximal risk over heterogeneous subpopulations. Our next step is to convert this to a generalisation bound for empirical top-k minimisation.

Generalisation bounds for CVaR and top-k
From the above, the key to deriving a generalisation bound is to study the behaviour of the CVaR. Specifically, how well does the empirical quantity CVaR α ( L f ) approximate the population CVaR on L f , i.e., CVaR α (L f )? The first concern is that CVaR α ( L f ) is not an unbiased estimate of CVaR α (L f ). Fortunately, this bias asymptotically vanishes.
Lemma 4. Pick any α ∈ (0, 1). Fix a function class F ⊂ R X and bounded loss : We may thus seek a generalisation bound for CVaR α ( L f ) in terms of its population counterpart, with an additional bias term that asymptotically vanishes. While concentration bounds for the CVaR are well-studied [Brown, 2007, Wang andGao, 2010], these do not suffice for our purposes since we seek a uniform bound over all f ∈ F.
The following result presents two such bounds. The first bound relies on a result of Cherukuri [2019], and depends on the covering number of F. The second bound is an application of Rademacher bounds to the formulation in (15).
An immediate consequence of Proposition 5 is that minimising the empirical top-k risk for k = αN controls the maximum subpopulation risk : using (15), (17), and (18), where the uncertainty set B(D, α) comprises all subpopulations of D occupying at least α mass. The LHS is exactly the maximal retrieval loss in (12), assuming D = p∈[P ] ν * p · D p with ν * p ≥ α, for all p ∈ [P ]. This is intuitive: as the subpopulation becomes rarer (α → 0), it is harder to guarantee good performance using a finite number of samples (M 1 , M 2 → ∞).

Relation to existing work
The two key steps in our analysis are to relate the top-k loss to the CVaR, and to then relate the CVaR to a worst-case subpopulation loss. The implication of combining these observations for retrieval has not previously been noted.
The issue of controlling the risk across subpopulations has been studied in prior work. Shalev-Shwartz and Wexler [2016] showed that minimising the maximal loss over all examples can guarantee good performance on data sub-populations. However, this analysis was for separable binary problems. By contrast, our analysis is for non-separable multiclass problems, and uses the more robust top-k loss. Fan et al. [2017] proposed to minimise the average of the top-k largest per-example losses. However, their analysis only established consistency in binary classification settings, and did not establish guarantees on different subpopulations.
In the fairness literature, a line of work has explored the use of the maximum or average top-k per-subgroup losses [Alabi et al., 2018, Mohri et al., 2019, Williamson and Menon, 2019. These assume the subgroup memberships are known, akin to §4.1, which is not the case in our setup.
When subgroup membership is unknown, Duchi and Namkoong [2018], Hashimoto et al. [2018], Oren et al. [2019], Duchi et al. [2019] propose robust optimisation procedures. These can be seen as directly working with variational representations akin to (15), which requires fundamentally altering the training procedure (e.g., introducing an auxiliary parameter to be iteratively optimised over); further, such techniques are not attuned to the retrieval setting, where K is large. By contrast, the minibatch top-k minimisation in S 2 M requires only a minor modification to the state-of-the-art SNM retrieval method .

Experiments
We now present experiments confirming our central claim that S 2 M can ensure good performance on heterogeneous data subpopulations, while being efficient to optimise in settings with a large number  Table 3: Test set accuracy on MNIST and CIFAR-10, where the training set is downsampled so that 5 classes appear 1% of the time. We assess average top-k minibatch SGD for varying k , on a minibatch size of N mb = 64. The best performance on the tail subpopulation defined by the downsampled classes is achieved by choosing k < 64.
of labels. We will illustrate this on both a controlled setting where the subpopulations are artificially created, and a setting with natural subpopulations.
At the outset, we emphasise that to our knowledge, no existing retrieval method addresses the hetereogenity problem; state-of-the-art retrieval methods [Jain et al., 2016, Prabhu et al., 2018, Guo et al., 2019 focus on aggregate performance, which we shall demonstrate can lead to a bias towards a dominant subpopulation. Further, while we do not claim to improve upon the state-of-the-art in terms of standard retrieval performance -involving metrics agnostic to heterogeneous subpopulations -we show that S 2 M remains competitive on such metrics as well.
Results on small-scale data. We begin with an experiment on two benchmark datasets, MNIST and CIFAR-10, each of which has 10 classes. On each dataset, we artificially create pronounced subpopulations as follows: (a) we split the classes into two groups, Head & Tail (b) on the training set, we downsample the Tail classes such that the ratio of Head : Tail classes is 100 : 1.
The data associated with the resulting Head and Tail classes form two imbalanced subpopulations of the data distribution. We do not modify the test data in any way, and so there is no such population imbalance in the test set.
We consider standard neural architectures for these datasets: for MNIST, we train a LeNet [Lecun et al., 1998] with minibatch size N mb = 64 for 100K iterations, and for CIFAR-10, we train a ResNet-50 [He et al., 2016] with batch size N mb = 64 for 600K iterations. On each dataset, we train S 2 M for k ∈ {1, 16, 32, 64}, with softmax cross-entropy as our base loss. Since the number of labels in both datasets is only K = 10, we do not perform sampling over labels, i.e., k = K. Recall from §3.2 that this corresponds to performing average top-k SGD and that when k = 64, the method is identical to standard minibatch SGD. Table 3 shows the performance of all methods on both the downsampled and unaltered classes. We see that on both datasets, using the average top-k loss for k < 64 results in the best performance on tail, i.e., the rare subpopulation corresponding to the downsampled classes. Further, on the CIFAR-10 dataset, we even see slight gains on modelling head, i.e., the subpopulation associated with the unaltered classes. This indicates that even within the head, we are able to better model more difficult examples.
Results on large-scale data. We next consider large-scale classification on AmazonCat-13k, a standard benchmark for extreme multilabel classification. Following Reddi et al.  Table 4: Test set recall on AmazonCat-13k, where the evaluation is split into three subpopulations (Head, Torso, Tail) based on label popularity. We assess S 2 M for varying values of k , on a minibatch size of N mb = 2048. Note that SNM corresponds to choosing k = 2048. During each optimization step, we select the k = 64 hardest negative labels from a set ofK = 4096 randomly sampled negative labels. The best performance at high levels of recall is achieved by choosing k < 2048.
Unlike the previous experiment, on this dataset we consider the naturally occurring subpopulations induced by the inherent label imbalance. We thus do not modify the training set in any way; rather, at test time, we simply evaluate performance per-subpopulation. These subpopulations are defined based on the label frequency as follows: (a) for each label y ∈ [K], we compute the fraction π y of times it appears in the training set (b) we let q, q be the 66% and 33% quantile of the π y 's (c) we say a test sample (x, y) belongs to Head if π y > q, Torso if q ≥ π y > q , and Tail otherwise.
At test time, we then separately evaluate performance on each of these three subgroups. Observe that by this construction, the subpopulations have overlapping support: indeed, an instance x may be associated with both rare and frequent labels, in which case it will belong to multiple subgroups.
To measure performance on a subpopulation, we compute recall@r for varying values of r ∈ {5, 10, 25, 50}, which is simply the negation of the retrieval risk (cf. (1)). This measures the number of (x, y) pairs where y is amongst the top-r scored labels from the model. Since each instance can have many labels and frequent labels tend to get higher scores compared to rare labels, we expect higher recall@r for Head labels, especially for low values of r.
For minibatch size N mb = 2048, we compare the performance of S 2 M for k ∈ {256, 512, 2048}. Note that k = 2048 refers to the baseline of SNM, which averages the loss over all examples whithin a batch. Following Reddi et al.
[2019], we train a fully connected neural network with a single hidden layer of width 512 employing linear activations. We use the cosine contrastive loss [Hadsell et al., 2006] as our base loss . During each step of training, for each instance, we select the k = 64 hardest negative labels from a set ofK = 4096 negative labels sampled uniformly. Second, while we already observe increased recall of head labels in the top-5 results, with increasing r, recall further improves over the baseline across subpopulations of increasingly rare labels. For r ≥ 25, performance improves across all subpopulations. This indicates that in addition to overall improved recall, S 2 M learns a consistent ranking across labels, i.e., head labels are consistently ranked within the top labels, followed by torso and then tail labels.
Overall, S 2 M achieves good overall recall performance as well as across different data subpopulations, while being efficient to optimise with a large number of labels.

A Proofs
Proof of Lemma 1. Let N mb ⊆ [N ] denote the indices of the N mb samples selected in the minibatch S mb . Note that where (i) follows from Kawaguchi and Lu [2019]. The precise form of the weights vector θ ∈ R N + is per Kawaguchi and Lu [2019]. Compared to SGD, the expected loss is a weighted average of the loss over the samples. Further, the weighting places emphasis on samples with highest loss. Next, by following , we have where ϕ ∈ R K−1 + is a weight vector. Here, the weighting places emphasis on the labels with highest loss. (The precise form of the weights is per .) Now, Lemma 1 follows by combining (20) and (21).
Proof of Proposition 2. We follow the proof of Fan et al. [2017], which was done for the binary classification setting. By exploiting the connection between the average top-k risk and the CVaR (15), we have that where α = k N . Let (f * , t * ) be the minimizer of the RHS in (22). Further, letting f * be the minimiser of the risk with respect to loss , i.e., f * = argmin f : Therefore, we have where L * = min f : X→R E (X,Y)∼D [ (Y, f (X)] denotes the Bayes-risk (minimal loss) with respect to .
Equipped with this, let us consider the following problem subject to y g y (x) = 1 and g y (x) ≥ 0 ∀y, where g y (x) = e fy (x) j e f j (x) and τ = e t . In what follows, we suppress the dependence on x to obtain the following problem. minimize y p y · [− log τ g y ] + subject to y g y = 1 and g y ≥ 0 ∀y, The problem can be rewritten as Note that Let Y 1 = {y : g y = 1 τ } and Y 2 = {y : g y < 1 τ }. Assuming that K · 1 τ > 1, we claim that Y 1 ∪ Y 2 = Y, i.e., at the optimal solution we don't have g y > 1 τ . This easily follows from the observation that assigning more that 1 τ mass to g y does not further reduce the objective, however, it can increase the contribution of g y , for y = y, to the objective. From KKT condition, an optimal solution satisfies ∂W ∂g y = 0 ∀y or Now, since {g y } forms a valid distribution, we have or Thus, we have, Recall that, for y ∈ Y 2 , we have g * y < 1 τ . Thus, y ∈ Y 2 if Let τ * = e t * . Consequently, the minimiser of the average top-k risk collapses the probabilities of the largest |Y 1 | labels into 1 τ * . Further, these largest elements will have probability larger than the other labels. Consequently, the minimiser will be will be top-τ * consistent.
To get some intuition, recall that τ * = e t * for the optimal threshold t * . Thus, in light of (23), τ * ≤ e L * /α . When L * ∼ 0 (i.e., the problem is noise-free), and α ∼ 1 (i.e., we average most of the labels), this indicates that the loss will be top-1 consistent, as is intuitive.
Proof of Lemma 3. Consider the following set of weightings on the training samples: i.e., all k -sparse vectors with equal weight on the nonzero components. The empirical average top-k loss as defined in (11) can then be rewritten as Given τ ∈ Λ top(k ),N , for each subpopulation, we can further consider only those samples which have nonzero associated weight: S τ,p = {i ∈ S p : τ i > 0} ⊆ S p , p ∈ [P ]. Now, we can rewrite (35) as follows: Note that the maximisation over τ is equivalent to selecting λ p . = |Sτ,p| k fraction of the k hardest examples that define the average top-k loss from the p-th subpopulation. We may thus write Proof of Lemma 4. Fix any f ∈ F. From the definition of the empirical and population CVaR (16), (15), we have wheret * , t * are the optimal values of t in (16), (15),. Following Rockafellar and Uryasev [2000], these correspond to the empirical and population quantiles of the losses, which we denote byq α and q α respectively. Thus, A simple case-analysis reveals that |[ (y i , f (x i )) − q α ] + − [ (y i , f (x i )) −q α ] + | ≤ |q α −q α |. Thus, the RHS can be bounded by 1 + 1 α · E [|q α −q α |]. Now observe that the empirical quantileq α converges in distribution to a Gaussian with mean q α and variance O 1 N [Arnold et al., 1992, Theorem 8.5.1]. Since the mean of the absolute value of a zero-mean Gaussian with variance σ 2 is O(σ), the result follows.
Proof of Proposition 5. We provide both a covering number bound based on Cherukuri [2019], and a Rademacher bound based on Boyd et al. [2012].
Covering number bound. Note that for any fixed f ∈ F, Let the loss function : [K] × R K → R satisfy the Lipschitz condition: Under this assumption, it follows from Cherukuri [2019, Proposition IV.9] that P sup where γ( ) = O(diam(F) · ( · α) −1 ), and β( ) = O(α · 2 · B −2 ), where diam(F) is the diameter of the function class, and B = ∞ is the maximum possible CVaR value. Now, for a large enough constant µ, by setting we obtain that Thus, it follows from (36), (39), and (40) that with probability at least 1 − δ, Rademacher bound. Define L α (f, t; D) We thus have CVaR α (L f ) = L α (f, t * ; D) where t * is the (1 − α)th quantile of the loss values. We similarly have the empirical version L α (f, t; S) where again CVaR α ( L f ) = L α (f,t * ; S), fort * the empirical quantile of the loss values.