Coping with Label Shift via Distributionally Robust Optimisation
Authors: Jingzhao Zhang
Aditya Krishna Menon
Andreas Veit
Srinadh Bhojanapalli
Sanjiv Kumar
Suvrit Sra
Venue: ICLR 2021 Poster (2021)
Keywords: Label shift
distributional robust optimization

Abstract: The label shift problem refers to the supervised learning setting where the train and test label distributions do not match. Existing work addressing label shift usually assumes access to an unlabelled test sample. This sample may be used to estimate the test label distribution, and to then train a suitably re-weighted classifier. While approaches using this idea have proven effective, their scope is limited as it is not always feasible to access the target domain; further, they require repeated retraining if the model is to be deployed in multiple test environments. Can one instead learn a single classifier that is robust to arbitrary label shifts from a broad family? In this paper, we answer this question by proposing a model that minimises an objective based on distributionally robust optimisation (DRO). We then design and analyse a gradient descent-proximal mirror ascent algorithm tailored for large-scale problems to optimise the proposed objective.  Finally, through experiments on CIFAR-100 and ImageNet, we show that our technique can significantly improve performance over a number of baselines in settings where label shift is present.

Published as a conference paper at ICLR 2021
COPING WITH LABEL SHIFT VIA DISTRIBUTIONALLY ROBUST OPTIMISATION
Jingzhao Zhang
Massachusetts Institute of Technology jzhzhangémie edu
Aditya Krishna Menon & Andreas Veit & Srinadh Bhojanapalli & Sanjiv Kumar Google Research (adityakmenon, aveit, bsrinadh, sanjivk edu
Suvrit Sra
Massachusetts Institute of Technology suvrit@mit.edd
ABSTRACT
The label shift problem refers to the supervised learning setting where the train and test label distributions do not match. Existing work addressing label shift usually assumes access to an unlabelled test sample. This sample may be used to estimate the test label distribution, and to then train suitably re-weighted classifier. While approaches using this idea have proven effective, their scope is limited as it is not always feasible to access the target domain; further, they require repeated retraining if the model is to be deployed in multiple test environments. Can instead learn a single classifier that is robust to arbitrary label shifts from a broad family? In this paper, answer this question by proposing model that minimises based on distributionally robust optimisation (DRO). We then design and analyse a gradient descent- prreenma mirror ascent algorithm tailored for large-scale problems to optimise the proposed objective. Finally, through experiments on CIFAR-100 and ImageNet, we show that our technique can significantly improve performance over a number of baselines in settings where label shift is present.
objective
INTRODUCTION
Classical supervised learning involves learning model from training distribution that generalises well on test samples drawn from the same distribution. While the assumption of identical train and test distributions has given rise to useful methods, it is often violated many practical settings (Kouw Loog, 2018). The label shift problem is one such important setting, wherein the training distribution over the labels does not reflect what is observed during testing (Saerens et al., 2002). For example, consider the problem of object detection in self-driving cars: model trained in one city may see vastly different distribution of pedestrians and cars when deployed in a different city. Such shifts in label distribution can significantly degrade model performance. As concrete example, consider the performance of a ResNet-50 model on ImageNet. While the overall eгror rate is ~ 24%, Figure I reveals that certain classes suffer an error as high as 80%. Consequently, label shift that increases the prevalence of the more erroneous classes in the test set can significantly degrade performance.
Most existing work on label shift operates in the setting where one has an unlabelled test sample that can be used to estimate the shifted label probabilities (du Plessis & Sugiyama, 2014; Lipton et al., 2018: 8; Azizzadeneshelii et al., 2019). Subsequently, one can retrain a classifier using these probabilities in place of the training label probabilities. While such techniques have proven effective, it is not always feasible to access an unlabelled set. Further, one may wish to deploy a learned model in multiple test environments, each one of which has its own label distribution. For example, the label distribution for a vehicle detection camera may change continuously while driving across the city. Instead of simply deploying separate model for each scenario, deploying single model that is

---

Published as a conference paper at ICLR 2021
120
Ertr on chinced disbbute
0.5
Frequency claas
1000
0.7
0.6
60
0.5
40
0.4
20
Par-class validition
K2.
Figure 1: Distribution of per-class test errors of a ResNet-50 on ImageNet (left). While the average enror rate is 24%, classes achieve error as high 80%. An adversary thus significantly degrade test performance (right) by choosing Pto (1) with more weight these classes.
robust to shifts may be more efficient and practical. Hence, we address the following question in this work: can we learn a single classifier that robust a family of arbitrary shifts?
We answer the above question by modeling label shift via distributionally robust optimisation (DRO) (Shapiro et al., 2014; 4; Rahimian Mehrotra, 2019). DRO offers a convenient way of coping with distribution shift, and have lead to successful applications (e.g. Faury et al. (2020)). Intuitively, by seeking a model that performs well on all label distributions that are "close" to the training data label distribution, this task can be cast game between the learner and adversary, with the latter allowed to pick label distributions that maximise the learner's s loss. We remark that while adversarial perspectives have informed popular paradigms such GANs, these pursue fundamentally different objectives from DRO (see Appendix A for details).
Although several previous works have explored DRO for tackling the problem of example shift (e.g.. adversarial examples) (Namkoong & Duchi, 2016; 6; 2017: Duchi & Namkoong, 2018), an application of DRO to the label shift setting poses several challenges: (a) updating the adversary's 's distribution naively requires solving a nontrivial convex optimisation subproblem with limited tractability, and also needs careful parameter tuning: and (b) naively estimating gradients under the adversarial distribution on randomly sampled minibatch can lead to unstable behaviour (see $3.1). We overcome these challenges by proposing the first algorithm that successfully optimises a DRO objective for label shift on large scale dataset (i.e.. ImageNet). Our objective encourages robustness to arbitrary label distribution shifts within a Kl.-divergence ball of the empirical label distribution. Importantly, we show that this choice of robustness set admits an efficient and stable update step.
Summary of contributions
(1) We design a gradient descent- descent-proximal mirror ascent algorithm tailored for optimising large- scale problems with minimal computational overhead, and prove its theoretical convergence.
(2) With the proposed algorithm, we implement a practical procedure to successfully optimise the robust objective ImageNet scale for the label shift application.
(3) We show through experiments on ImageNet and CIFAR-100 that our technique significantly improves over baselines when the label distribution is adversarially varied.
2 BACKGROUND AND PROBLEM FORMULATION
In this section we formalise the label shift problem and motivate its formulation as adversarial optimisation problem. Consider a multiclass classification problem with distribution Ptr over instances X and labels y [L]. The goal is to learn classifier he: X - y parameterised by 0 E O, with the aim of ensuring good predictive performance on future samples drawn from Par More formally, the goal is minimise the objective ming () [((x, 0)], where L: X y O R4 is loss function. In practice, we only have access to finite sample S {(x,, p= ph. which motivates us to use the empirical distribution Pemp (I, y) = = 1(r = Zi, y:) in place of Per. Doing so, we arrive at the objective of minimising the empirical risk:
min E Ep... [f(x, 0)] = "" " E( y:. 0).
(1)
The assumption underlying the above formulation is that test samples are drawn from the same distribution Pur that is used during training. However, this assumption is violated in many practical settings. The problem of learning from a training distribution Ptr. while attempting to perform well on test distribution Pte * Ptr is referred domain adaptation (Ben-David al., 2007). In
2

---

Published as a conference paper at ICLR 2021
Label distribution
Reference
Train distribution
Standard ERM
Specified a-priori (e.g, balanced)
(Elkan, 2001; Xie & Manski, 1989; Cao et al., 2019)
Estimated test distribution
(du Plessis & Sugiyama, 2014; Lipton et al., 2018; Azizzadenesheli
ct al., 2019; Garg et al., 2020; Combes et al., 2020)
porst-performing class Worst k-performing classes
(Hashimoto ct al., 2018; Mohri et al., 2019; Sagawa et al., 2020)
(Fan et al., 2017; Williamson & Menon, 2019; Curi et al., 2019; Duchi ct al., 2020)
Adversarial shifts within KL-divereence
This paper
Table I: Summary of approaches to learning with a modified label distribution.
the special case of label shift, one posits that Pte (xr y) - Ptr(x y). but the label distribution Pie (1) F Pir(y) (Saerens al., 2002); i..., the test distribution satisfies (r. y) Pte (t) (2 v)- The label shift problem admits the following three distinct settings (see Table for a summaryy):
(1) Fixed label shift. Here, one assumes a-priori knowledge of (te One may then adjust the outputs of a probabilistic classifier post-hoe to improve test performance (Elkan, 2001). Even when the precise distribution is unknown, it is common to posit uniform Pie (t)) Minimising the resulting balanced error has been the subject of large body of work (He & Garcia, 2009), with recent developments including Cui al. (2019); Cao et al. (2019); Kang et al. (2020); Guo et al. (2020).
(2) Estimated label shift. Here, we assume that (to is unknown, but that we have access to an unlabelled test sample. This sample may be used estimate Pte (y)) e.g,, via kernel mean- matching (Zhang et al., 2013), minimisation of suitable KL divergence (du Plessis & Sugiyama, 2014), or using black-box classifier outputs (Lipton et al., 2018: Azizzadenesheli et al., 2019; Garg et al., 2020). One may then use these estimates to minimise suitably re- re-weighted empirical risk.
(3) Adversarial label shift. Here, we assume that (v) is unknown, and guard against suitably defined worst-case choice. Observe that an extreme case of label shift involves placing all probability mass on single y" E y. This choice can be problematic, as (1) may be rewritten as
min Σ
Pemp (V) vE(L]
Σ nly M
E(E,, (111)310)) yo 01
where ny is the number of training samples with label y. The empirical risk is thus a weighted average of the per-class losses. Observe that if some y" has a large per-class loss, then an adversary could degrade performance by choosing a Pte with (to()") being large. One means of guarding against such adversarial label shifts is to minimise the minimax risk MAdiiz-Rodriguez et al., 2007; Davenport et al., 2010; Hashimoto et al., 2018; Mohri et al., 2019; Sagawa et al., 2020)
f(x, Ms. 0
min max
Σ ((V)- DELL]
Σ 9L-9
# TEAL
Fly
where A denotes the simplex. In (2), we combine the per-label risks according the worst-case label distribution. In practice, focusing on the worst-case label distribution may be overly pessimistic. One may temper this by instead constraining the label distribution. A popular choice is enforce that || < 1 for suitable k, which corresponds to minimising the average of the top-k largest f per-class losses for integer (Williamson & Menon, 2019; Curi al., 2019; Duchi et al., 2020).
We focus on the adversarial label shift setting, as it meets the desiderata of training a single model that robust to multiple label distributions, and not requiring access to test samples. Adversarial robustness has been widely studied (see Appendix A for more related work), but its application to label shift is much less explored. Amongst techniques in this area, Mohri et al. (2019); Sagawa et al. (2020) are most closely related to our work. These works optimise the worst-case loss over subgroups induced by the labels. However, both works consider settings with a relatively small (S 10) number of subgroups; the resultant algorithms face many challenges when trained with many labels (see Section 4). We now detail how a suitably constrained DRO formulation, coupled with optimisation choices, can overcome this limitation.
3

---

Published as a conference paper at ICLR 2021
Algorithm 1 ADVSIIFTOO,, Ye, A, NNOpt, Pemp: 7/))
1: Initialise adversary distribution as (,,,, 7,).
2: for t . do
3: Sample mini-batch of b examples [(r,, yi Vi)
4: Evaluate stochastic gradient ge Σ relyi) Vollsi, Mc,
Pamp(yss
5: Update neural network parameters 0(+1 NNOpt(ge)
Update Lagrangian variable a = if (Te, p
27,2 if < (Te, Pomp).
7:
Evaluate adversarial gradient 9x (i) 2/=1 Pemp
(8,, mi) Vol(z,, 8c+1).
8: Update adversarial distribution M++1 (Ie Pomp) 21(4aa) exp (T/
3 ALGORITHM: DISTRIBUTIONALLY ROBUST KL-DIVERGENCE MINIMISATION
To address the adversarial label shift problem, propose replace the empirical risk (1) with min max E [e(2, y, 0)], == { A d(f.. S r):
(3)
0 TEP
where 3 is an uncertainty set containing perturbations of the empirical distribution pemp- This an instance of distributionally robust optimisation (DRO) (Shapiro et al., 2014), a framework where one minimises the worst-case expected loss over a family of distributions. In this work, we instantiate DRO with P being a parameterised family of distributions with varying marginal label distributions in Kl.-divergence, i.e., d(p, q) Rggr-g [- log p(y)/g()) (We use this divergence, as opposed to a generic f-divergence, as it affords closed-form updates; see $3.3.) Solving (3) thus directly addresses adversarial label shift, as it ensures our model performs well on arbitrary label distributions from P. Observe further that the existing minimax risk (2) is special case of (3) with +00.
Having stated our learning objective, we now turn to the issue of how to optimise it. One natural thought is to leverage strategies pursued in the literature example-level DRO using f-divergences. For example, Namkoong & Duchi (2016) propose an algorithm that alternately performs iterative gradient updates for model parameters and adversarial distribution æ, assuming access projection oracles, and the ability to sample from the adversarial distribution. However, there are challenges in applying such techniques on large-scale problems (e.g., ImageNet):
(1) directly sampling from π is challenging in most data loading pipelines for ImageNet.
(2) projecting onto the feasible set P requires solving constrained convex optimization problem at every iteration, which can incur non-trivial overhead (see Appendix E).
We now describe ADVSHIFT (Algorithm 1), our approach solve these problems. In nutshell, we iteratively update model parameters 9 and adversarial distributions л. In the former, we update exactly per ERM optimization (e.g., ADAM, SGD), which we denote NNOpt (neural network optimiser); in the latter, we introduce a Lagrange multiplier to avoid projection. Extra care is needed to obtain unbiased gradients and speed up adversarial convergence, as we detail.
3.1 ESTIMATING THE ADVERSARIAL MINIBATCH GRADIENT
For a fixed E A to estimate the parameter gradient E,, [Vol(z, y, 0. on a training sample 8 ((ri, Di) Di) we employ the importance weighting identity and write
E Vo
Eypemp
1
Vollz, v, 0)
Σ
a(yc)) Pomp (3rc)
Vollz,, Vi,
(e)
We may thus draw a minibatch as usual from 8, and apply suitable weighting to obtain unbiased gradi- ent estimates, A similar reweighting is necessary to compute the adversary gradients E,, [,,((, v, Making the adversarial update efficient requires further effort, as we now discuss.
3.2 REMOVING CONSTRAINTS BY LAGRANGIAN DUALITY
To efficiently update the adversary distribution я in (3), we would like to avoid the cost projecting onto P. To bypass this difficulty, we make the following observation based on Lagrangian duality.
4

---

Published as a conference paper at ICLR 2021
Proposition 1. Suppose é is bounded, and Pemp is not the boundary of the simplex. Then, Vr By" > such that for every No 2 72, the constrained objective solvable in unconstrained form:
argmax TEAL, KTAp-Pamp)Sr
Ex [e y. 0)] argmax Egg [[(x y, 0)] + min(0, 7e(r (L, (emp))
SEAL
Motivated by this, we may thus transform the objective (3) into:
min max E2 y, 0)] + min (0, % Pemp)))
(4)
8 TEAL
where > 0 is sufficiently large constant; in practice, this may be chosen by bisection search. The advantage of this formulation is that it admits an efficient update for æ, as we now discuss.
3.3 ADVERSARIAL DISTRIBUTION UPDATES
We now detail how
employ proximal mirror descent to efficiently update и. Observe that
we may decompose the adversary's 's (negated) objective into two terms: f(0, π) := -E [ y.
and h(T) = max (0, % (KL,, Pemp) r)), where h(x) is independent of the samples. Such decomposable objectives suggest using proximal updates (Combettes Pesquet, 2011):
1
M++1 = prox ( AV,I(0,, == argmin ((T) + 2X (1 x1² 2A(V (0, (5)
where À serves as the learning rate. The value of proximal descent relies on the ability to efficiently solve the minimisation problem in (5). Unfortunately, this does not hold as-is for our choice of l(T)), essentially due to a mismatch between the use of KL-divergence in h, and Euclidean distance || T| in (5). Motivated by the advantages of mirror descent over gradient descent on the simplex (Bubeck, 2014), we propose to replace the Euclidean distance with KL-divergence:
1T argmin (TT)) (KL(T,, mT) 2X(g;; m)),
(6)
TEAL
2X
where gt an unbiased estimator of V [0;, Tt). We have the following closed-form update. Lemma 2. Assume the optimal solution TT-1 to (6) satisfies [L Pemp) * т, and that all the classes appeared at least once in the empirical distribution, i.e. Vii, Pomp > O. Let 7 Tte if r < KT(T++1, 1,pmp), and 7 0
(Tt+1. then TT-1 permits a closed form solution
T++1 = (TT @ Pemp)) p
/(1+a))))) exp (nng))/C,
(7)
2y C (T O Pemp) 1/(++)) exp (л=9))l1 projects T++1 the simplex, and a @ is the element-wise product between two vectors a, b.
where The (++1/2X)(1a))
In Algorithm 1, we set 7 To if < KL (T Pemp) and 0 otherwise to appoximate the true y. Such approximation works well when r KL Te, Pemp) does not change sign frequently.
3.4 CONVERGENCE ANALYSIS
We provide below convergence analysis of our gradient descent-proximal mirror ascent method for nonconvex-conen stochastic saddle point problems. For the composite objective ming m f(0, π² + h(TT), and fixed learning rate no. we abstract the Algorithm update as:
0,1 = 0, - (0,) Te+1 = argmax h(T)) 2Λ
(KL m.) + 2X(g(T,),
(8)
x
where g(T), (@)) are stochastic gradients assumed to satisfy the following.
Assumption 1. The stochastic gradient g( with respect satisfies that for some > 0,
[0) = (0, x), and [[(lg(() [)](())) < o²
Assumption 2. The stochastic gradient (T)) with respect to æ satisfies that for some G >
B[g(=)) (0. T), and [[]lg(m)]]3 G².

---

Published as a conference paper at ICLR 2021
We make the following assumptions about the objective, similar to Lin et al. (2019; 2020):
Assumption 3. x(),z) + h(()) is L-smooth and L-pisschit: f0), and h(T) are concave in и.
Assumption 4. Every adversarial distribution iterate satisfies Pem - R for some R >> 0.
Assumption 3 and 4 may be enforced by adding constant e to the adversarial updates, which prevents Ti from approaching the boundary of the simplex. Assumption 2 in the label shift setting implies that the loss is upper and lower bounded. Such an assumption may be enforced by clipping the loss for computing the adversarial gradient, which can significantly speed up training (see Appendix ??). Furthermore, this is a standard assumption for analyzing nonconvex-concave problems (Lin et al., 2019). The assumption that the square Loc norm is bounded is weaker than L2 norm being bounded; such a relaxation results from using mirror rather than Euclidean update.
Given that the function F()) == maxre /(0, π) + l(()) nonconvex, our goal is to find a stationary point instead of approximating global optimum. Yet, due to the minimax formulation, the function F(0) may not necessarily be differentiable. Hence, we define convergence following some recent works (Davis & Drusvyatskiy, 2019; Lin et al., 2019; Thekumparampil et al., 2019) nonconvex- concave optimisation. First, Assumption 3 implics F(0) is L-weakly convex and I-Lipschitz (Lin al., 2019, Lemma 4.7). Hence, we define stationarity in the language of weakly convex functions. Definition 1. A point 9 is e-stationayy point of weakly convex function F if || (1/21 (0) where /22 (6) denotes the Moreau envelope Fi F1/2r. (6) min, F(w) + D| 0|1?.
With the above definition, we can establish convergence of the following update: Theorem 3 (informal). Under Assumptions 1-4, the update in (8) finds point 0 with B[IVF+ + (0)Il) S e in O(E iterations.
For a precise description of the theorem, please see Appendix H. The above result matches the best known rate in Lin et al. (2019) for optimising nonconvex-concave problem with stochastic gradients. To our knowledge, this is the first result that studies convergence of composite objectives with proximal methods under nonconvex-concave settings. By utilizing the proximal operator, solves the objective with extra ((T)) term without incurring additional complexity cost.
3.5 CLIPPING AND REGULARISING FOR FASTER CONVERGENCE
In addition to the proposed algorithm, we apply two additional techniques. We explain them here with motivations. First, we also observe that the adversarial's 's update could be very sensitive to the adversarial gradient 9k, i.c. label-wise loss in each minibatch, because the gradient appears in the exponential of the update. To avoid convergence degradation resulted from the noise in Jk. we clip the label-wise loss value 2. Second, we notice that the KI. divergence from any interior point of simplex to its boundary is infinity. Hence, updates near boundary can be highly unstable due to the nonsmooth KL loss. To cope with this, add constant term on the adversarial distribution to avoid the adversarial distribution reaching any of the vertices on the simplex. The term and clipping critical in both training and convergence analysis. We conduct an ablation of the sensitivity to these parameters in Figures 5 and 6. Note that the experiments show that even without these tricks, our proposed algorithm alone still outperform baselines.
3.6 DISCUSSION AND COMPARISON TO EXISTING ALGORITHMS
number of existing learning paradigms (e.g., fairness, adversarial training, and domain adaptation) have connections to the problem of adversarial label shift; see Appendix A for details.
We comment on some key differences between ADVSHIFT and related techniques in the literature. For the problem of minimising the worst-case loss (2) - which is equivalent to setting the radius r +oo in (3) Sagawa et al. (2020) propose an algorithm that assumes the ability to sample data from given group in order to evaluate adversarial gradients. Such sampling is cumbersome to implement in most ImageNet data loading pipelines. Mohri et al. (2019) propose a way to evaluate gradients using importance sampling, and then apply projected gradient descent-ascent This method suffers from instability owing to sampling (upon which we improve with proximal updates), and incurs non-trivial computational overhead due to the projection step. We will illustrate these

---

Published as a conference paper at ICLR 2021
Toi readton
VM Sooaoo orrer rod. aoddec
A
Dalonced
Flod
Ruet
ADESHI 2.11
AdKINT E11
divegence threshid
Sivergence treshals
thesselt
(a) Ablation train.
(b) Ablation validation.
(c) Comparison train.
(d) Comparison validation.
Comparison of performance on ImageNet under adversarial label distributions. For each method, we vary the KI. divergence threshold т, and for each report the maximal validation error induced by the adversarial shift within the threshold. Subplots (a) (b) compare the performance of ADVSIIITT trained with different DRO radius r against the default ERM training. We subtract the baseline error of ERM from all values for easy visualization. Absolute values can be found in Figure in the Appendix. Combined with (c), (d), we see that ADVSHIFT can reduce the adversarial validation error by over 2.5% compared to the BASELINE method and is consistently superior to the AGNOSTIC, BALANCED and FIXED methods. Figure 3(c) illustrates adversarial distributions for varying thresholds r.
Figure 2:
problems in our subsequent experiments (see results for AGNOSTIC in $4). Finally, for an uncertainty set I based on the CVaR, Curi et al. (2019) provide an algorithm that updates weights using ExP3, This approach relies on a determinantal point process, which has a poor cemension-ddeendeen
4 EXPERIMENTAL RESULTS
We now present a series of experiments to evaluate the performance of the proposed ADVSHIFT algorithm and how it compares to related approaches from the literature. We first explain our experiment setups and evaluation methods. We then present the results on ImageNet dataset, and show that under the adversarial validation setting, our proposed algorithm significantly outperforms other methods discussed in Table 1. Similar results on CIFAR-10 100 are shown in the Appendix.
4.1 EXPERIMENTAL SETUP
To evaluate the proposed method, use the standard image classification setup of training a ResNet- 50 on ImageNet using SGD with momentum as the neural network optimiser. All algorithms are run for 90 epochs, and are found to take almost the same clock time. Note that ImageNet has a largely balanced training label distributions, and perfectly balanced validation label distributions.
We assess the performance of models under adversarial label shift follows. First, we train model on the training set and compute its error distribution on the validation set. Next, we pick a threshold the allowable KL divergence between the train and target distribution and find the adversarial distribution within this threshold which achieves the worst-possible validation error. Finally, we compute the validation performance under this distribution. Note that T 0 corresponds to the train distribution, while 7 = +oo corresponds to the worst-case label distribution (see Figure 1).
We evaluate the following methods, each corresponding one in Table 1: (i) standard empirical risk minimisation (BASELINE) (ii) balanced empirical risk minimisation (BALANCED) (iii) agnostic federated learning algorithm of Mohri et al. (2019), which minimises the worst-case loss (AGNOSTIC) (iv) our proposed KL-divergence based algorithm, for various choices of adversarial radius (AD- VSHIFT) (v) training with ADVSHIFT with a fixed adversarial distribution extracted from Figure 3(c) (FIXED). This corresponds to the estimated test distribution row in Table I with an ideal estimator.
4.2 RESULTS AND DISCUSSION
Figure 2 shows the train and validation performance on ImageNet. Each curve represents the average and standard deviation across 10 independent trials. To better illustrate the differences amongst methods, we plot the difference in error to the BASELINE method. (See Figure in the Appendix for unnormalised plots.) Subfigures (a) and (b) compre the performance of ADVSHIFT for various choices of radius to the ERM baseline; (c) and (d) compare SDIFH the remaining methods.
7

---

Published as a conference paper at ICLR 2021
1.0
10:1
distanoo:
0.8
0.0
Advesarial meght
distance:
0.0
9.6
10-1
2044
0.4
10-
10-
0.2
Baseline
Baseline
Advshiit
60
200
600
800 1000
Training
Trairing
Label soted validation
(a) Train.
(b) Validation.
(c) Adversarial distributions.
Figure 3: Subplots (a) (b) show violin plots of the distribution of errors for both the BASELINE and our ADVSHIFT methods over the course of training. On the training set, ADVSHIFT significantly reduces the worst-case error, evidenced by lower upper endpoints of the distribution. On the validation set, the reduction is consistent, albeit less pronounced owing to a generalisation gap. Subplot (c) illustrates adversarie distributions at KL distances of 1, 2 and 3 for model trained with BASELINE. Even at 7 = 1, the adversarial distribution is highly concentrated on only a few hard labels.
Epoch
Epoch 18
Epoch 54
Epoch 90
1000
200
Frequency
150
150
150
600
100
400
50
50
50
200
0.0005
0,001
0.0015
0.002
0.0005
0.003
0.0015
0.002
0.0005
0.001
0,0015
0.002
0.0005
0.001
0.0015 0.002
Adversari weight
Adversarial weight
Adversarial weight
Adversarial weight
Figure 4: Evolution of learned adversarial distribution (T) across training epochs. Starting off from uniform distribution over labels, the adversary quickly infers the relative difficulty of small fraction of labels, assigning nearly 2 the weight them compared to the average. This distribution remains largely stable in subsequent iterations, getting gradually more concentrated as training converges.
Hyperparameters for each method are separately tuned. FIXED 1, 2, 3 corresponds to training with each of the three adversarial distributions in Figure 3(c). We see that:
the reduction offered by ADVSHIFT consistently superior to that afforded by the AGNOSTIC, BALANCED and FIXED methods. On the training set, observe significant (~ 8%) reduction in performance for large KL divergence thresholds. On the validation set, the gains are less pronounced (~ 2.5%), indicating some degradation due generalisation gap.
while ADVSHIFT consistently improves above the baseline across adversarial radii, we observe best performance for r = 0.1. Smaller values of r lead to smaller improvements, while training becomes increasingly unstable for larger radii. Please see the discussion the last section.
during training, AGNOSTIC either learns the adversarial distribution too slowly (such that behaves like ERM), This highlights the importance of the proximal mirror ascent updates in our algorithm.
too large learning rate for the adversary (such that the training fails).
Illustration of distributions at fixed KL thresholds. Figure 3(c) visualises the adversarial distribu- tions corresponding to a few values of the KL threshold т. At a threshold of r = 3, the adversarial distribution is concentrated on only a few hard labels. Consequently, the resulting performance on such distributions is highly reflective of the worst-case distribution that can happen in reality.
Training with a fixed adversarial distribution. Suppose we take the final adversarial distributions shown in Figure 3(c), and then employ them as fixed distributions during training; this corresponds to the specified a-priori and estimated validation distribution approaches in Table 1. Does the resulting model similarly reduce the error on hard classes? Surprisingly. Figure 2(d) indicates this is not so, and performance is in fact significantly worse on the "casy" classes. Employing a fixed adversarial distribution may thus lead to underfitting, which has intuitive explanation: the model must struggle to fit difficult patterns from early stages of training. Similar issues with importance weighting in conjunction with neural networks have been reported in Byrd & Lipton (2019).
Evolution of error distributions. To dissect the evolution of performance during training, Figure shows violin plots of the distribution of errors for both the BASEL INE and our ADVSHIFT methods
8

---

Published as a conference paper at ICLR 2021
reduction
Acvdlin dlp 6.1
AdvShift dip
AdvShift Olo
Advohil die 10
Advshi It. tip 10
2
KL. divergence threshold
divergence threshold
(a) ImageNet train.
(b) ImageNet validation.
Figure 5: Ablation of loss clipping threshold. We see that when the clipping threshold is either too large or too small, validation performance of the model tends to suffer.
Train reducction
Baseline Adrihilt
Ves ialidooonnn enrer reduction
Basdling
Advshil
AdrShin: 001
AdvShit 0.001
-5
Adv/Shit 0.03
Ad/Shirt
AdvStii 0.01
AdVShiit
-20
KL divergence threshold
KL. divergence threshold
(a) ImageNet train.
(b) ImageNet validation.
Ablation of gradient stabilisation parameter which is constant added to the gradient updates to prevent iterates from reaching the vertices of the simplex. We see that without any gradient stabilisation, the model's performance rapidly degrades the adversarial radius increases. Conversely, performance also suffers when the stablisation is too high.
Figure 6:
after fixed training epochs. We observe that on the training set, ADVSIIFT significantly reduces the worst-case error, evidenced by the upper endpoints of the distribution being reduced. Note also that, as expected, the adversarial algorithm is slower to reduce the error on the "easy" classes early in training, evidenced by the lower endpoints of the distribution initially taking higher values. On the validation set, the reduction is consistent, albeit less pronounced owing to generalisation gap.
Evolution of learned adversarial weights. To understand the evolution of the adversarial distribu- tion across training epochs, Figure plots the histogram of adversary weights at fixed training epochs. Starting off from a uniform distribution, the adversary is seen to quickly infer the relative difficulty of small fraction of labels, assigning 2x the weight on them compared to the average. In subsequent iterations the distribution becomes more concentrated, and gradually reduces the largest weights.
Ablation of clipping threshold and gradient stabiliser.
Figures and show an ablation of the choice loss clipping threshold, and the gradient stabiliser We see that when the clipping threshold is either too large or too small, validation performance of the model tends to suffer (albeit still better than the baseline). Similarly, we see that without any gradient stabilisation, the model's ''s performance rapidly degrades as the adversarial radius increases. Conversely, performance also suffers when the stablisation is too high.
In summary, our experiments show that our proposed DRO formulation can be effectively solved with ADV SHIFF, and results in a model that is robust to adversarial label shift.
5 DISCUSSION AND FUTURE WORK
We proposed ADVSHIFT, an algorithm for coping with label shift based on distributionally robust optimisation, and illustrated its effectiveness of real-world datasets. Despite this, our approach does not solve the problem fully. First, Figure 2ba)(b) shows that the generalization gap increases the perturbation radius increases. Understanding why there is correlation between hard examples and bad generalization could improve robustness. Second, Figure 2(a) shows that even the train the algorithm threshold r does not translate to the model's ' from the interplay of model expressivity and data distribution, whose future study is of interest.
level of robustness. We conjecture this results

---

Published as a conference paper at ICLR 2021
REFERENCES
Rocío Alaiz-Rodriguez. Alicia Guerrero-Curieses and Jesús Cid-Sueiro. Minimax regret classifier for imprecise class distributions. Journal of Machine Learning Research, 8:103-130, May 2007.
ISSN 1532-4435.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization, 2019.
Kamyar Azizzadenesheli, Anqi Liu, Fanny Yang, and Animashree Anandkumar. Regularized learning for domain adaptation under label shifts. In International Conference on Learning Representations, 2019.
Shai Ben-David, John Blitzer, Koby Crammer, and Fernando Pereira. Analysis of representations for domain adaptation. In Advances in neural information processing systems, pp. 137-144, 2007.
Sébastien Bubeck. Convex optimization: Algorithmss and complexity, 2014.
Jonathon Byrd and Zachary Chase Lipton. What is the effect of importance weighting in deep learning? In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings the 36th Inter- national Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, volume 97 of Proceedings of Machine Learning Research, pp. 872 281. PMLR, 2019.
Kaidi Cao, Colin Wei, Adrien Gaidon, Nikos Arechiga, and Tengyu Ma. Learning imbalanced datasets with aabel-disribution-avvaa margin loss. In Advances in Neural Information Processing Systems, 2019.
Remi Tachet des Combes, Han Zhao, Yu-Xiang Wang, and Geoff Gordon. Domain adaptation with conditional distribution matching and generalized label shift. arXiv preprint arXiv:2003.0447 2020.
Patrick Combettes and Christoope Pesquet. Proximal splitting methods in signal processing. In Fixed-poine algorithms for inverse problems in science and engineering, pp. 185-212. Springer, 2011.
Yin Cui, Menglin Jia, Tsung-Y Yi Lin, Yang Song, and Serge Belongie. Class-balanced loss based on effective number of samples. In CVPR, 2019.
Sebastian Curi, Kfir. Y. Levy, Stefanie Jegelka, and Andreas Krause. Adaptive sampling for stochastic risk-averse learning, 2019.
M. Davenport, R. G. Baraniuk, and C. D. Scott. Tuning support vector machines for minimax and neyman-pearson classification. IEEE Transactions on Pattern Analysis and Machine Intelligence, 18880)188881188 2010.
Damek Davis and Dmitriy Drusvyatskiy, Stochastic model-based minimization of weakly convex functions. SIAM Journal on Optimization, 29(1)2007-239, 2019.
Michele Donini, Luca Oneto, Shai Ben-David, John Shawe-Taylor, and Massimiliano Pontil. Em- pirical risk minimization under fairness constraints. In Proceedings of the 32nd International Conference Neural Information Processing Systems, NIPS' 18, pp. 2796-2806 Red Hook, NY, USA, 2018. Curran Associates Inc.
Marthinus Christoffel du Plessis and Masashi Sugiyama. Semi-supervised learning of class balance under class dlas--prior change by distribution matching. Neural Newworks, 50:1 10-119,,0001 2014.
J. Duchi and H. Namkoong. Learning models with uniform performance via distributionally robust optimization. arXiv preprint arXi:: 2018
John Duchi, Tatsunori Hashimoto, and Hongseok Namkoong. Distributionally robust losses for latent covariate mixtures, 2020.
Charles Elkan. The foundations of cost-sensitive learning. In Proceedings of the 17th International Joins Conference on Artificial Intelligence Volume 2, IJCAI'O1, pp. 973-978, San Francisco. CA, USA, 2001. Morgan Kaufmann Publishers Inc. ISBN 155860815 25..
10

---

Published as a conference paper at ICLR 2021
Yanbo Fan, Siwei Lyu, Yiming Ying, and Baogang Hu. Learning with average top-k loss. In Advances in neural information processing systems, pp. 497-505, 2017.
Louis Faury, Ugo Tanielian, Elvis Dohmatob, Elena Smirnova, and Flavian Vasile. Distributionally
robust counterfactual risk minimization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 3850-3857, 2020.
Saurabh Garg, Yifan Wu, Sivaraman Balakrishnan, and Zachary C. Lipton. A unified view of label shift estimation, 2020.
Ian J. Goodfellow, Jean Pouget-Abiis Abadie, Mchdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron C. Courville, and Yoshua Bengio, Generative adversarial nets. In Zoubin Ghahramani, Max Welling, Corinna Cortes, Neil D. Lawrence, and Kilian Q. Weinberger (eds.). Advances in Neural Information Processing Systems 27: Annual Conference on Neural Information Processing Systems 2014, December 8-13 2014, Montreal, Quebec, Canada, pp. 2672-2688) 2014.
Jiaxian Guo, Mingming Gong, Tongliang Liu, Kun Zhang, and Dacheng Tao. Ltf: label transfor- mation framework for correcting label shift. In International Conference on Machine Learning, pp. 3843-3853. PMLR, 2020.
Tatsunori Hashimoto, Megha Srivastava, Hongseok Namkoong, and Percy Liang. Fairness without demographics in repeated loss minimization. In Jennifer Dy and Andreas Krause (eds.), Interna- tional Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1929-1998,, Stockholmsmässann Stockholm Sweden, 10-15 Jul 2018. PMLR.
Haibo He and Edwardo A. Garcia. Learning from imbalanced data. IEEE Transactions on Knowledge and Data Engineering, (2)3:126311284 2009.
Weihua Hu, Gang Niu, Issei Sato, and Masashi Sugiyama. Does distributionally robust supervised learning give robust classifiers? In Jennifer G. Dy and Andreas Krause (eds.). Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 2042-2044.. PMLR, 2018.
Bingyi Kang. Saining Xie, Marcus Rohrbach, Zhicheng Yan, Albert Gordo, Jiashi Feng, and Yan- nís Kalantidis. Decoupling representation and classifier for long-tailed recognition. In Eighth International Conference Learning Representations (ICLR), 2020.
Wouter M. Kouw and Marco Loog. An introduction to domain adaptation and transfer learning. CoRR, abs/ 1ass//81..1111 1806, 2018 URL http: http::///axxvv/oa/ee wwtp/////////oooooo///.../...ss/ 1111111111. 11806.
Daniel Levy, Yair Carmon, John C Duchi, and Aaron Sidford. Large-scale methods for distributionally robust optimization. Advances in Neural Information Processing Systems, 33, 2020.
Tianyi Lin, Chi Jin, and Michael I. Jordan. On gradient descent ascent for nonconvex-coneave minimax problems, 2019.
Tianyi Lin, Chi Jin, and Michael. I. Jordan. Near-optimal algorithms for minimax optimization, 2020.
Zachary Lipton, Yu-Xiang Wang, and Alexander Smola. Detecting and correcting for label shift with black box predictors. In Jennifer Dy and Andreas Krause (eds.), International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 3122-3130, Stockholmsmässan, Stockholm Sweden, 10-15 15 Jul 2018. PMLR.
Mehryar Mohri, Gary Sivek, and Ananda Theertha Suresh. Agnostic federated learning. In Interna- tional Conference on Machine Learning, 2019.
Krikamol Muandet, David Balduzzi, and Bernhard Schölkopf. Domain generalization via invariant feature representation. In Sanjoy Dasgupta and David McAllester (eds.), International Conference on Machine Learning. Proceedings of Machine Learning Research, pp. 10-18, Atlanta, Georgia, USA, 17-19 Jun 2013. PMLR.
H. Namkoong and J. Duchi. Stochastic gradient methods for distributionally robust optimization with f-divergences. In Advances in Neural Information Processing Systems (NeurIPS), 2016.

---

Published as a conference paper at ICLR 2021
Hongseok Namkoong and John C Duchi. Variance-based regularization with convex objectives. In Guyon, U. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Gareett (eds.), Advances in Neural Information Processing Systems 30, pp. 2971-2980. Curran Associates, Inc., 2017.
Hamed Rahimian and Sanjay Mehrotra. Distributionally robust optimization: A review, 2019.
Marco Saerens, Patrice Latinne, and Christine Decaestecker. Adjusting the outputs of a classifier to new priori probabilities: A simple procedure. Neural Computation, 14(1):21-41 2002.
S. Sagawa, P. W. Koh, T. B. Hashimoto, and Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. In International Conference on Learning Representations (ICLR), 2020.
Alexander Shapiro, Darinka Dentcheva, and Andrzej Ruszczyúski. Lectures on Stochastic Program- ming: Modeling and Theory, Second Edition. Society for Industrial and Applied Mathematics, Philadelphia, PA, 2014. doi: 10.1137/1.978161 1973433.
Aman Sinha, Hongseok Namkoong, and John Duchi. Certifying some distributional robustness with principled adversarial training. arXiv preprint arXiv: 10571, 2017.
Aman Sinha, Hongseok Namkoong, and John Duchi. Certifiable distributional robustness with principled adversarial training. In International Conference on Learning Representations, 2018.
Amos J Storkey and Masashi Sugiyama. Mixture regression for covariate shift. In B. Schölkopf, J. C. Platt, and T. Hoffman (eds.), Advances in Neural Information Processing Systems 19, pp. 1337-1344. MIT Press, 2007.
Kiran Koshy Thekumparampil. Prateek Jain, Praneeth Netrapalli, and Sewoong Oh. Efficient algorithms for smooth minimax optimization, 2019.
Robert C. Williamson and Aditya Krishna Menon. Faimess risk measures. In Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, pp. 6786-6797 2019.
Yu Xie and Charles Manski. The logit model and response-based samples. Sociological Methods & Reseaach, 17(3):283-302, 1989.
Muhammad Bilal Zafar, Isabel Valera, Manuel Gomez Rodriguez, and Krishna P. Gummadi. Fair- ness beyond disparate treatment disparate impact: Learning classification without disparate mistreatment. In Proceedings of the 26th International Conference World Wide Web, WWW '17, pp. 1171-1180, Republic and Canton Geneva, CHE, 2017. International World Wide Web Conferences Steering Committee. ISBN 9781450349130.
Kun Zhang, Bernhard Schölkopf, Krikamol Muandet, and Zhikun Wang. Domain adaptation under target and conditional shift. In Proceedings of the 30th International Conference on International Conference on Machine Learning Volume 28, ICML' 13, pp. III-81--III-8 JMLR.org. 2013.
12

---

Published as a conference paper at ICLR 2021
A RELATED PROBLEMS
Example-level DRO. Existing work on DRO has largely focussed on the setting where I encom- passes shifts in the instance space (Namkoong & Duchi, 2016; 2017; Sinha et al., 2018; Duchi & Namkoong, 2018; Levy et al., 2020). This notion of robustness has a natural link with adversarial training (Sinha al., 2017), and involves a more challenging problem, as requires parameterising the adversary's 's distribution. Hu et al. (2018) illustrate the potential pitfalls of DRO, owing to mismatch between surrogate and 0-1 losses. They also propose to encode an uncertainty set based on latent label distribution shift (Storkey & Sugiyama, 2007), which requires domain knowledge. The techniques in example-level DRO are mostly designed for small scale dataset with SVM models, as these techniques require sampling according adversarial distribution, which can be very unstable if implemented with importance sampling only. It also requires maintaining a vector proportional to the number of labels and indexing each sample during training to match up the sample index, which is not available in most dataloading pipelines.
Fairness. Adversarial label shift may be related algorithmic fairness. Abstractly, this concerns the mitigation of systematic bias in predictions on sensitive subgroups (e.g., country of origin). One fairness criteria posits that the per-subgroup errors should be equal (Zafar et al., 2017; Donini et al., 2018), an ideal that may be targetted by minimising the worst-subgroup error (Mohri et al., 2019; Sagawa al., 2020). When the subgroups correspond labels, ensuring this notion of fairness is tantamount to guarding against an adversary that can place all mass on the worst performing label.
GANs. GANs (Goodfellow et al., 2014) involve solving a min-max objective that bears some similarity to the DRO formulation (3), but is fundamentally different in details: while DRO considers reweighting of samples according to a fixed family, GANs involve parameterised adversarial family, with the training objective augmented with an additional penalty.
Domain adaptation. Label shift can be viewed as a special case of domain adaptation, where Ptr and Phe can systematically differ. Typically, one assumes access to small sample from Pte. which may be used to estimate importance weights (Combes et al., 2020), or samples from multiple domains, which may be used to estimate generic comainnagnostto representation (Muandet et al., 2013). In causal inference, there has been interest in similar classes of models (Arjovsky et al., 2019).
B ALGORITHM IMPLEMENTATION DETAILS
We introduce some additional details in our implementation of ADVSHIFT. First, as observed in Section 3.1, our algorithm requires knowing the empirical label distribution. As the exact value not always available, we estimate the empirical label distribution online for all the experiments presented later in Section 4 using an exponential moving average, Pemp = 3 Pemp + (1 3) Pbatch: where Phanch is the label distribution in the minibatch. We set 3 = 0.999. The number is set such that the exponential moving average has a half-life roughly equal to the number of iterations in one epoch of ImageNet training using our setup.
In all the experiments, we set 2%a = 1 in Algorithm for simplicity. For learning the adversarial distribution, we only tune the adversarial learning rate T/
C ADDITIONAL EXPERIMENTAL RESULTS
We present here additional experimental results, including:
for ImageNet, an illustration of the lack of correlation between a label's frequency in the training set, and its validation error. (Figure 7)
unnormalised versions of the results on ImageNet shown in the body, where we do not subtract the baseline performance from each of the curves; this gives sense of the absolute performance numbers obtained by each method. (Figure 8)
an ablation of the loss clipping threshold and gradient stabiliser as introduced above. (Figure 5,6)
results on CIFAR-1000 to complement those for ImageNet. (Figure 9,10)
13

---

Published as a conference paper at ICLR 2021
0.8
Validation
0.2
0.0
0.0006
0.0007 0,0008 0. 0.0009 0.0010 Label frequency
Figure Illustration that training label frequency does not strongly correlate with test error. Observe that several classes with high error appear frequently in the training set.
70
Agnostic 0.
Agrostic 0.05
60
Balanced
Validation error
Balanced
Advšhitt 0.1
Advšhitt 0.1
50
50
30
20
divergence threshald
divergance threshold
(a) ImageNet train.
(b) ImageNet validation.
Figure 8: Comparison of performance of various methods on ImageNet under adversarial label distributions. For each plot, we vary a KL divergence threshold т, and for given construct the label distribution which results in maximal test error for the baseline model. We then compute test error under this distribution. Note that the case 7 = 0 corresponds to using the train distribution, while = +oo corresponds to using the worst-case label distribution, which is concentrated on the worst-performing label. Our proposed ADVSHIFT can reduce the adversarial test error by over 2.5% over the baseline method.
C.I BALANCED LABELS
+ BALANCED PERFORMANCE
Figure shows that training label frequency does not strongly correlate with test error. Observe that several classes with high error appear frequently in the training set. Indeed, the three classes with highest error casette player, maillot, and water jug all appear an equal number of times in the training set.
C.2 UNNORMALISED PLOTS ON IMAGENET
Figure presents plots of the unnormalised performance the various methods compared in the body. Here, rather than subtract the performance of the baseline, show the absolute accuracy of each method as the adversarial radius is varied. Evidently, the baseline and AGNOSTIC models tend to suffer in their validation error as the adversarial radius increases.
C.3 RESULTS ON CIFAR-100
Figure shows results on CIFAR-100, where we train various methods using CIFAR-ResNet--1 as the underlying architecture, Here, we see consistent and sizable improvement from ADVSHIFT over the baseline method. On this dataset, AGNOSTIC fares better, and eventually matches the performance of ADVSHIFT with large adversarial radius. This is in keeping with the intended use-case AGNOSTIC, i.c., minimising the worst-case loss. Figure 10 supplements these plots with unnormalised versions, illustrate the absolute performance differences.
14

---

Published as a conference paper at ICLR 2021
2.5
Validation error retuction
Agnostic 0.001
0.0
Agnostk: 0.001
-5
Agnostic
Agnostic
+2.5
Agnostk 0.01
Agnesti:
-10
Advšhitt 0.01
-5.0
Advšhit 0.01
AdvShift
-15
-7.
-20
-10.0
-25
divergence thresholo
divergence threshols
(a) CIFAR-100 train.
(b) CIFAR-100 validation.
Figure 9: Comparison performance of various methods on CIFAR-100.
Baseline
Baseline
Agnostic
Agnostic
Aghostic 0.01
Aghostic 0.01
Agnostic 0.1
60
Agnostis 0.1
80
AdvShift 0.1
50
AdvShiit 0.1
30
AdVShift
AdvShift
20
3 KL divergence threshold
1
KL divergence threshold
(a) CIFAR-100 100 train.
(b) CIFAR- 100 validation.
Figure 10: Comparison of performance of various methods on CIFAR-100 (unnormalised),
We remark here that choice a CIFARRResNet-- results in an underparameterised model, which does not perfectly fit the training data. In the overparameterised case, there challenges with employing DRO, as noted by Sagawa et al. (2020). Addressing these challenges settings where the training data balanced remains interesting open question.
D CONSTRAINED DRO DOES NOT PERMIT BOLTZMAN SOLUTION
We start with a simple example with three label classes {a, b, c} with class losses / = (1,2.4)
respectively. We assume an uniform empirical distribution, i.c. Pemp {1/3, 1/3, 1/3). We consider two different problems. The first is to find the optimal solution to regularised objective,
p argmin, pT <KL(p, (pm.pp
This problem is well known (e.g. see 2.7.2 of lecture to permit solution of form p(x) explz/t Ereeta exp for some t.
contrast, we show that distributions the above form does not solve the constrained version the problem. In particular, consider the following optimisation problem:
max p
F
such that K L(, Pemp) S
If the solution is of form p(x) = Ex((ludee) / then we know for 4, # lor
expl,/t
(Log(Pa) log(p=))/ (po) log(pe)) (la lc)/(l le). We solve the above problem with convex optimizer using 0.01 and found Pa = 0.283, 0.322, Pe = 0.395.
(log(pa) (og(p:))/((og//e log(p.)) 1.64 + (L. 1.) (1. 1.5.
Note that the above example shows that all solutions the contrained problem can be written a Boltzman distribution, i.c. p(x) = Extetasse expl,/t Yet, this does not contradict with results
exply/t
15

---

Published as a conference paper at ICLR 2021
[e.g. Lemma 4, Faury et al. (2020)] that claim there is Boltzman distribution whose function value matches the optimal value of the constraint problem. Mathematically, we can have p * p' but Ep]([)]] E. (xx))
E PROJECTING AN ADVERSARIAL DISTRIBUTION
The projection operator in our setting aims project distribution p into the P {g: KL P) - r by solving the following problem:
min 11 pll
such that
Σ 4i log(g:/p)) Sr
Σ 9: 1,
Vi, 2 0,
where Qi: denotes the ieh component of q: and denotes number of classes. Given that our implementation based on Tensorflow, we use the "trust-region constrained algorithm" provided by SciPy for easy integration with our python-based training procedure. However, even after extensive tuning, solving each problem up to 1% relative constraint error requires more than minute when 1000 (the number of labels in ImageNet). This means that if we train ResNer50 on ImageNet for 100k iterations, we need to spend 100k minutes on projection operation, which is not affordable.
F PROOF OF PROPOSITION
Proof. We only need to show that for large enough Te. any minimiserp' p' of the unconstrained problem satisfies that KL(p", PPmp)) S Since the distance from boundary of the simplex to any interior point is too, we can safely assume that the point lies within the relative interior of the simplex. To prove the proposition, denote c inf V (p Pemp)ll p E A s.t. KL Pemp) r). By strict convexity of KI. divergence and the fact that r >> 0, we know c 0. Denote the upper bound of loss as M, then when Yo > M/c,, we know that KL. Pemp)
= O ş Oz (E, [(x U, 0)] mi, Yo(r KL(T, Pemp))mp)) However, since p minimises the objective if and only if 0 Oz (E,, [6(x, y, 0)] + min Jo(r KL(T, (L(TTPemp))) }), we have that any minimiser must satisfy KL(T, (Pmpp) S r.
O
G PROOF OF LEMMA 2
Recall that we want to find πK+1 that minimises the following objective,
1
argmin ((T)) + (KL(TTTTA)) + 2X(gu, T))
Th+1 =
TEA
2X
argmin max { 0, 1 az
Oe KL Pemp))
+ (KI(TT, TT) - 7/96, T)),
TEA
a.
where 7 1/(2T. + 1/A), a. 27d.. Denote vli) the in component of vector Notice that the simplex can be written as a constraint that E; (1) = Vi, x(i) 2 0. Based on this constraint, we first write (6)'s Lagrangian dual
7(1) 1) + max 0, ae (KL(T, Pemp) - -
L(a,b, Σ (a.T (i) + b M
ar
(KI TT) 17 T))
1 ac
where (1) denotes the in component of т. If ne > 0 component-wiss then the optimal x cannot lie on the boundary (i.e. Vi, (i)) > 0), which results in KT,(T, mx) co. By Lagrangian duality and
16

---

Published as a conference paper at ICLR 2021
complementary slackness, we know that for if (L( Pemp)
a b, = b+ ae लπ(()
0
(i) + 1+
log(<(3)/TT(()) - -1 + 1.
It ac
On the other hand, if (T.(T, Pemp) <
0 d θπ(1) bz) b log(z(i)/T<(i)) (1) + 1.
1 ae
We discuss the case when (L(T,, Pemp) < r, and the other case follows similarly, Rearrange the optimality condition of Lagrangian multiplier, get
a (emm/(/))mmm()) + I+ log(T(i)/<<(1) - (1) -
Since b is a constant for all coordinates,
t(i) (Pemp(1)2 (T)) exp
(19((1) Itoc
The result follows by noting that Σ; T(i) = 1.
H PROOF OF THEOREM
For completeness, we define several terms used in optimisation. A function S(0) is Lasl--ipssoo -ipscchhtz if for 0,0',
[f(0) f(00) < (|00 0'1l- function f(0) is L smooth for all 0, 0',
[V) f(")l) s L|le 8'1|- function f(0) is L L-weklly convex if (0) 1 110 convex.
Then we can state the formal theorem below.
Theorem (formal version of Theorem 3). Under Assumptions 1-4, the update in (8) generates
sequence of points 01, ..,....8 Or with the following property:
2 (F1/22(00) -Fi/2)) + +G + G2
2 + (12 + 02)1/2)
T
TI/I
+ (h- /(TTO))/T
Proof. For convenience, denote Φ(0, π) f00 h(T), F(O) maxp $(0,p). We start by following the standard SGD proof. Denote as the stochastic gradient evaluated at step / - 1 with respect to 0. Denote = prox 2/2 == arg min, (F(w) + 2L 01 Conditioned on f,-1- we have
E||| [|||0 0bll9]= = =100,-1 1--1112 + 2n ((; - Q1-1, go)l + + E 1|
=|00,- 112 2118 (0,-1 O.-1, (0 Te. 1))
+ 2 +0 (9)
where the first equality follows by 02 = 0,- - - V & (0,00) गe. TT- Next. we observe that
(0.1 0,-1, Vo 1)) s 200, Tr 40 Te-1)
L
1100-
0,111 (10)
$ F(0.-1) - @(0 me-1) +
-
L
11001 -0,111²
L
< F(O,.)) - 00,0-1,, mm-1
11 -,-
2
17

---

Published as a conference paper at ICLR 2021
The first line follows by convexity and s-sootthnssss. The second line by definition of F. The third line by definition of Q. Next, by definition of Moreau envelop,
F1/21 - F(0r-1) + T110,-1 0ell
Take expectation on both sides and we get (+)))
(11)
S F() + B| - -
-((0-1) + -(110-1 - 11- + 2n (0 - 0, Vo TT- + 1112 + o²²))
S F11/21(0-1)) 1) + ( (0)-1 11) 200 T TT- - 110-1 11-1 + L (12 + o²)
The second line substitutes in (9). The third line follows by convexity and L smoothness. Denote that A, = 1( (02-,1, mT-1) 2 (0,-11, 1c-1 0001-1, 1-)) We can sum over / and take expectation recursively to get,
Σ B/I/VA//22 (0.)l 2L Bl1|0, 0c-111
(12)
- Lipe
(Fi (Fi/21 - F1)21) + + D Tho + d²))
where F2 ming F2 (0). The first equality follows by the definition of Moreau envelope. The second inequality follows by rearranging (11).
Next, we aim to bound the accumulated error Σ, A..
Recall that the update for the is follows for E[gs] (0,
T == argmin (-2X x) 21h(=) + (T)
Applying Lemma 5 with L(TT) 22( TT) + 22)((T), we get
((T)))))) (gx, " (0.)) + KT.(TT, mx)
2 - (g Tx+1) h(TT+I))))))))))))))))))(()))( + (L(TR+I,TH) + KL(T", TK+1)
Rearrange and take expectation we get
2X(EE((99 ([- T" (0.))) - [[(((T++1)) + (O.)))
< -E[KL TT+ mx) + KL Tx+ KL(=", T&) 2X Te T8+1)))
S -E[KL (T.. x [T 1]+1)]
Tx. 1/ 2X 21 11 + +| - T 11//2
The second inequality follows by the fact that di-divereence is strongly convex with respect to L₁ norm and Mauchy-Schwand inequality. We further observe that
-E( Th x (0.))) = (OV, Tx), F - = (0.)) 2 - mx) + f(0,, x (0.)) -J(Okk TR) + f(O,, x" (0x)) [(0,, T" (0x)) + [(0, "" (0.)) 2 -f(Kk, TK) + f(2,, n (0x)) - f(0,, T (0x)) + f(0,, 7 (0x))
-f(0,, x" (0.)) + f(0,, x" x0{)))
2 -[(0k, Tx) + [(0,, x (0a)) - 27|10, -0.00 0xll
The first inequality follows by concavity. The third line follows by f(e,, n² (fx)) f(0,, (0.)).
The last inequality follows by Lipschitzness. Similarly,
-h(Ok, x h:0,, x" (0x)) -h(Ok, =x) hi(Ok, T" ((xx)) 2 0x1
We can take iterative expectation and get sum over s + 1, 1,..., 8 B to get
Σ k=s+1
B[- TR) + f(0e, T" (0w)) M(Tx)) + h (0x))
s+B
S -hh(T,) + E[h((T++))) + 41 Σ Σ
E[10 - 01+1111 + ABG²
A
1
[E]KL(T"(0,) To) - EILL T [[K), T.))
2X
18

---

Published as a conference paper at ICLR 2021
Note that 4 -f(0k, T&) f(0k. ((x)) h(TT)) + h(x=/(xx0000 (6x)), hence
++
B| - - ( 2ng ABG²
k=s+1
1 (B|KL( (0.), mb)] E[KL(TT(0.)) T.]])
2
By further sum over all blocks and divide by total number of iterations T, we get
T/B
s+B
T
Σ
Σ
E|AN|
< (-h(m)) + B (- BG + XG2
b=
k=ba+1
+ 2XB
(O.), E[KL(T (0))
Substitute the above inequality into (12) and we get
1
2
T Σ
E [[
(F1))n(00) + Bn + XG
TLmo
[E To)) +)
2AB
we set 178 3/44 B 1/2 X 1/4 then we see that
1 Σ 21
2
R + (12 + 0222122)
(Fi/2n Fī,21) 2G + G2 2
TI/4
(-h(mT) + +/)
Lemma S. For any differentiable convex function L, if = argmin,ea + KL(x, zo)), then for any x' E A, we have
((x') (L(a', 2 {(xx) KL(x", xo) + (x(xx²xx))
Proof. This Lemma is well-known, but we include a proof for completeness. By optimality of x and convexity of A, we know that
(Vl(=") Vo(x") vo(z)),=-=")=
where 6(z) E, x; log( and the Bregman divergence defined according to ó is KL-divergence. Then
(xx') > ((x' (V(()"), x²)
2 +(x")+ + (Vo(z)) - Vo(a*), ), - x²)
=((x)) (Vo(v)), x0' zo) 0055)(()))))))))))))) (0x0))
+VV((zo) x' mp) 6(x') o(zn) (Vo(x"), x²) )000) 6(x*)
=((x") KL(x=" ro) KL(,', 20) + KL x²)
19


Reviews:
[{"review_id": "Q5gWi6AKUPN", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Weighting each class differently during model training is a common technique to deal with label imbalance and label shift. In this paper, the authors propose AdvShift, a method for learning these weights in a data- and model-dependent way through distributionally robust optimization (DRO). \n\nI found the paper especially well-written and clear, from the problem setup and contextualization within related work, to the description and intuition behind the method and the results, as well as the candid discussion of potential limitations of the method. Thank you to the authors for the enjoyable read. \n\nOverall, the method and results seemed promising. However, I had two general areas of concern that prevented me from giving a higher rating; the first with the comparison to prior work, and the second about the experimental details.\n\n**Prior work**\nThe authors identify Mohri et al. (2019) and Sagawa et al. (2020) as their closest work, so I primarily read up on and referred to those papers. \n\n[1a] Compared to Mohri et al. (2019), the differences here seems to be that (a) Mohri et al. use a chi-squared regularizer vs. the KL divergence term used in AdvShift; (b) Mohri et al. use (computationally expensive) projections whereas AdvShift uses proximal updates; and (c) AdvShift also includes a couple of optimization tricks, loss clipping and gradient stabilization, as detailed in the appendix. The argument that the authors make is that (b) is the important ingredient; that is, it is a question of optimization. Is this a fair characterization?\n\n[1b] If so, the tricks used in (c) seem to be critical to the performance of AdvShift, so they might perhaps explain the gap between Mohri et al. and AdvShift. The authors write that the proximal mirror ascent updates solve the optimization instability (compared to the projection operator), but do not elaborate. \n\n[1c] Moreover, if the superiority of the method hinges on optimization, then the phenomenon discussed in the last paragraph is especially striking: it seems like AdvShift is not optimizing properly with regard to the algorithm threshold r, as the authors point out. This suggests that there could be fundamentally something wrong with the algorithm and its optimization.\n\n[2] Compared to Sagawa et al. (2020), the authors write that the method in Sagawa et al. requires the ability to sample data from a given group. This does not appear to be true: Sagawa et al. describe their method as using standard minibatching and SGD. The authors also write that Sagawa et al. only test their method on a small number of labels, which is true, so it seems plausible that AdvShift will perform better on a large number of labels, but this would need to be empirically determined.\n\n[3] Both Mohri et al. and Sagawa et al.’s methods were originally tested on a small number of labels. How does AdvShift compare to these when there are a small number of labels? In the Appendix, it appears that Mohri et al. performs comparably to AdvShift on CIFAR-100. Is the pitch that one should use AdvShift when there are a large number of labels, but that it does not matter otherwise?\n\n**Experimental details**\n[4] Fig 2(a) and 2(c): Why is there no difference in training error at tau = 0 even when training for a very conservative model (e.g., AdvShift 1.0, or Fixed 3)? This would make sense if train errors are close to 0 but they are not: Fig 6 in the appendix shows that the training errors are >20% at tau = 0. The models used seem very underparameterized, presumably as a result of hyperparameter tuning; do they actually beat a normal ResNet with standard hyperparameters in the adversarial setting? I couldn’t find hyperparameter tuning details in the appendix.\n\n[5] Figs 2 and 3: Are the authors reporting test or validation accuracies? The figure captions all say validation, but the text all says test. \n\n[6] Fig 2(d): I am confused by why the fixed baseline does so badly. To clarify, for each value of tau, does each model have its own adversarial shift? Specifically, for the fixed model, is it still being evaluated against the fixed distribution? If not, could we plot it as an oracle? If it is, that seems very strange to me that optimizing for exactly that distribution would get higher train error. The authors cite Byrd and Lipton (2019) but I wasn’t able to figure out how that paper addressed this case, since Byrd and Lipton seem to be operating in the separable data setting. This suggests that there are some optimization issues?\n\n**Minor**\n[7] It is reasonable to focus on the adversarial setting, but the adversarial test error is still exceedingly high despite some improvements from AdvShift, so from the point of view of providing guarantees, this is only marginally better. I’m curious if AdvShift helps on non-worst-case settings. Other work on label imbalance, for example, https://arxiv.org/abs/1901.05555 show that their reweighting techniques can improve accuracy on the standard test set for ImageNet. \n\nTypos: \nS4.1, “has a largely balanced training label distributions”\nS4.2, “compre”\n\n**Update**\nThank you to the authors for the detailed response. Overall, I'll leave my rating as a 6. The paper is clear and proposes a promising optimization method, with moderate comparisons to prior approaches and reasonable results on ImageNet (but not CIFAR-100). The impact of these results are somewhat diluted by the overall absolute low scores in the adversarial setting, but as a general DRO optimization method, the result is interesting. Thank you for all the discussion.\n\nComments to the latest author response (posting this as an edit since public comments are now disabled):\n\n[1] Thanks, the updated version is clearer.\n\n[2] I agree with this characterization (for example, the unconstrained perturbation might cause optimization instability). The paper still seems to contain the language about the Sagawa paper requiring a custom sampler. Overall, I agree that there are reasons to believe that the proposed optimization method is more stable, but I think the authors should be clear that they only compared to Mohri et al. experimentally (in your response to Reviewer 3, you mentioned that you extensively discussed and compared to both Mohri and Sagawa).\n\n[6] Thanks. I also agree with this. I had two points: first, the differences between your proposed algorithm and the standard algorithm might be more stark when the training is not balanced. Second, in imbalanced settings, it is more likely that the test distribution will be skewed (e.g., rare animals in iNaturalist).\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "J5w_HiRyg3c", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper attacks the issue of mismatch in distribution of labels between train and test samples. The authors propose a DRO-based approach which amounts to solving a modified ERM problem. Compared to classical approaches, the proposed method doesn’t entail fitting many different models: just a single model. The method builds on recent progress on solving nonconvex-concave games for approximate stationary points. The resulting algorithm is a mirror-descent scheme with explicit convergence rates, under Setting : train and test label distribution do not match\n\nThe general approach is:\nLearn a classifier robust to arbitrary label shifts (from a family) -> doing this by using DRO approach because it allows them to train a model that performs well on all label distributions sufficiently close to the training data label distribution.\nThe paper proposes the DRO approach with KL-divergence for the label shift problem, an analysis of the proposed algorithm (optimization techniques and convergence analysis) and experiments on Imagenet (ResNet50)\n\nGood points\n      - It paper is well-written and asy to follow paper. The problem and the algorithm are clearly described and well analyzed. The experiments are also quite well done. \nThe label shift pb is clearly explained and described\nThe analysis of the proposed algorithm (more or less all section 3) is quite thorough.\nThe experiment section is clearly described (and the additional experiments in Appendix are useful).\n\nBad points:\nThe experiments are not really convincing in the sense that the improvement is either really limited on the test set by comparisons to the baseline (on Imagenet), or similar to that of the Agnostic method (on CIFAR-100, Appendix)\nIt would have been nice to have a more detailed explanation of the complexity of the algorithm in practice, and more explanation on the choices of hyper-parameters (Appendix B). There are many hyper-parameters that were set to a fixed value, and only one parameter is tuned (how?), so how/why do you choose these fixed values? \n\nRecommendations:\n      -  Contrary to what the authors say, inner problem in (4) does indeed have a closed-form solution. It is the Fenchel-Legendre transform of KL divergence, which equals log-sum-exp(...). Also, this max is attained at Boltzman distribution which is proportional to exp(loss / temperature). See Lemma 4 of Faury et al. (AAAI 2020) Distributionally Robust Counterfactual Risk Minimization. BTW, its would be interesting to contrast this result with the result of Lemma 2, I. Equation (7),  of the current manuscript.\n\nErrors:\n      -  Equation (2) and previous equation (unlabeled): yi should be replaced with y\n      - First equation  (unlabeled) in section 3.1: missing pi(y) in importance weight \n\nQuestions\nIt is mentioned that the Agnostic method faces many challenges during training, but these challenges are not explained afterwards (computation time? Difficult parameters tuning? Etc.)\n\nI’ll happily revise my score upward if my concerns are addressed.", "extra_scores": {}, "extra_text": {}}, {"review_id": "AGBY9pKRgKs", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper tackles label shift in supervised learning via distributionally robust optimization. The main idea is to train by solving a min-max problem, where the max problem searches for the worst-case label shift in an Kullback-Leibler divergence ambiguity set. The KL ambiguity set will generate some form of adversarial reweighing of the sample training points, which gives us hope that the learned parameters will perform better in the test data (with shift). The paper proposes to solve the Lagrangian version instead of the constrained version of the DRO problem, and proposes a gradient descent-ascent type of algorithm.\n\nThe main contributions of this paper is condensed in Section 3:\n- From the mathematical viewpoint, the DRO problem (3) and the mini batch gradient estimator are not new.\n- In Section 3.2, the authors proposed the Lagrangian problem. However, the difficulty of this problem lies in choosing the \"right\" value of gamma_c. Indeed, finding gamma_c is not trivial, and simply saying that a bisection search can be employed (as suggested by the authors) is not sufficient. To use bisection search, the authors should argue that the optimal value of problem (4) is at least convex in gamma_c, and this result has not been established. \n- The authors relax the 2-norm to a KL-divergence in equation (6). This relaxation makes the complaint (2) in page 4 become invalid. Indeed, complaint (2) says that projection is difficult under 2-norm, however, we can also relax to the projection using KL-divergence to alleviate this difficulty.\n- Problem (3) and (4) are equivalent only if gamma_c is optimally tuned. The algorithm analyzed in Section 3.4 is has guaranteed only for problem (4). It is still not clear what guarantee we can have on the original problem (3), which is the main problem of interest.\n\nGetting back to the main contributions that are listed in page 2: I can only partially agree with contribution (2). However, I do not think that contribution (1) about the formulation and contribution (3) about the numerical results can really be justified as the main contributions of this paper.\n\nMinor comment:\n1. In algorithm 1, line 1 should be \\pi_0. In line 2, the for loop should start with t = 0\n2. I don't understand why p_emp should not be on the boundary of the simplex. If p_emp is on the boundary, we can simply drop the samples (x_i, y_i) with p_emp(y_i) = 0, and that should not affect the problem.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: 
This paper considers the problem of training neural networks to be robust to label shifts. To do so, it proposes to use a distributionally robust optimization (DRO): instead of minimizing the expected error with respect to the empirical data distribution, the worse case expected error is minimized over a KL-divergence "ball" of distributions centered at the empirical distribution with a given radius. The main contribution of the paper is an efficient algorithm for achieving this optimization, that avoids the need to project onto the uncertainty set or to sample in non-standard ways from the training set. The paper provides evidence on the ImageNet data set and the ResNet-50 architecture that the proposed AdvShift algorithm outperforms reasonable baselines.

Reviewers raised concerns about the novelty of the algorithm, and claims the paper makes about the infeasibility of the sampling required by one of the competing baselines, and the need for the Lagrangian parameter used in the algorithm to be well-tuned. The rebuttals addresses the concerns suitably; in particular, the novelty of the algorithm lies it it being the first DRO-based solution for the label shift problem and its efficiency obtained by using KL uncertainty set and the Lagrangian formulation of the problem, which allows a closed form solution. 

Due to the strong empirical and theoretical performance of the proposed AdvShift algorithm, it is recommended that this paper be accepted.
Decision: Accept (Poster)