C OPING WITH LABEL SHIFT VIA DISTRIBUTIONALLY ROBUST OPTIMISATION
Anonymous authors
Paper under double-blind review
A BSTRACT
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 classiﬁer.
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 classiﬁer 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 CIFARand ImageNet, we show that our technique can signiﬁcantly improve performance over a number of baselines in settings where label shift is present.
I NTRODUCTION
Classical supervised learning involves learning a model from a 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 in many practical settings (Kouw &
Loog, 2018).
The label shift problem is one such important setting, wherein the training distribution over the labels does not reﬂect what is observed during testing (Saerens et al., 2002).
For example, consider the problem of object detection in self-driving cars: a model trained in one city may see a vastly different distribution of pedestrians and cars when deployed in a different city.
Such shifts in label distribution can signiﬁcantly degrade model performance.
As a concrete example, consider the performance of a ResNet-50 model on ImageNet.
While the overall error rate is ∼ 24% , Figurereveals that certain classes suffer an error as high as ∼ 80% .
Consequently, a label shift that increases the prevalence of the more erroneous classes in the test set can signiﬁcantly 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;
Azizzadenesheli et al., 2019).
Subsequently, one can retrain a classiﬁer 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 a separate model for each scenario, deploying a single model that is robust to shifts may be more efﬁcient and practical.
Hence, we address the following question in this work: can we learn a single classiﬁer that is robust to a family of arbitrary shifts?
We answer the above question by modeling label shift via distributionally robust optimisation (DRO) (Shapiro et al., 2014; Rahimian & Mehrotra, 2019).
DRO offers a convenient way of coping with distribution shift, by seeking a model that performs well on all label distributions that are “close” to the training data label distribution.
Intuitively, this task can be cast as a game between the learner and an adversary , with the latter allowed to pick label distributions that maximise the learner’s loss.
We remark that while adversarial perspectives have informed popular paradigms such as GANs, these pursue fundamentally different objectives from DRO (see Appendix A for details).
[Figure]
F r e q u e n c y
o
f c l a ss    4 KL( p tr , p te ) 0.
E rr
o r o n s h
i f t
e d d i s t r i b u t i o n
Label distribution
[Figure]
Standard ERM
Speciﬁed a-priori (e.g., balanced) (Elkan, 2001; Xie & Manski, 1989; Cao et al., 2019)
[Table]
(du Plessis & Sugiyama, 2014; Lipton et al., 2018; Azizzadenesheli et al., 2019; Garg et al., 2020; Combes et al., 2020)
[Table]
(Hashimoto et al., 2018; Mohri et al., 2019; Sagawa et al., 2020)
[Table]
(Fan et al., 2017; Williamson & Menon, 2019; Curi et al., 2019; Duchi et al., 2020)
[Figure] [Figure]
This paper
Although several previous works have explored DRO for tackling the problem of example shift (e.g., adversarial examples) (Namkoong & Duchi, 2016; 2017; Duchi & Namkoong, 2018), an application of DRO to the label shift setting poses several challenges: (a) updating the adversary’s distribution na ¨ ıvely requires solving a nontrivial convex optimisation subproblem with limited tractability, and also needs careful parameter tuning; and (b) na ¨ ıvely estimating gradients under the adversarial distribution on a randomly sampled minibatch can lead to unstable behaviour (see § 3.1).
We overcome these challenges by proposing the ﬁrst algorithm that successfully optimises a DRO objective for label shift on a 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 efﬁcient and stable update step.
Summary of contributions
(1) We formulate a robust optimisation objective for the label shift problem based on KL-divergences.
(2) We design a gradient descent-proximal mirror ascent algorithm tailored for optimising large-scale problems with minimal computational overhead, and prove its theoretical convergence.
(3) We show through experiments on ImageNet and CIFAR-100 that our technique signiﬁcantly improves over baselines when the label distribution is adversarially varied.
B ACKGROUND AND PROBLEM FORMULATION
In this section we formalise the label shift problem and motivate its formulation as an adversarial optimisation problem.
Consider a multiclass classiﬁcation problem with distribution p tr over instances X and labels Y =
[ L ] .
The goal is to learn a classiﬁer h θ : X → Y parameterised by θ ∈ Θ , with the aim of ensuring good predictive performance on future samples drawn from p tr .
More formally, the goal is to minimise the objective min θ E ( x,y ) ∼ p tr [ (cid:96) ( x, y, θ )] , where (cid:96) :
X × Y × Θ → R + is a loss function.
In practice, we only have access to a ﬁnite sample S =
{ ( x i , y i ) } ni
=1 ∼ p n tr , which motivates us to use the empirical distribution p emp ( x, y ) =
1 n (cid:80) ni =1 1 ( x = x i , y = y i ) in place of p tr .
Doing so, we arrive at the objective of minimising the empirical risk:
[Equation]
The assumption underlying the above formulation is that test samples are drawn from the same distribution p tr that is used during training.
However, this assumption is violated in many practical settings.
The problem of learning from a training distribution p tr , while attempting to perform well on a test distribution p te (cid:54) =
p tr is referred to as domain adaptation (Ben-David et al., 2007).
In the special case of label shift , one posits that p te ( x | y ) =
p tr ( x | y ) , but the label distribution p te ( y ) (cid:54) =
p tr ( y ) (Saerens et al., 2002); i.e., the test distribution satisﬁes p te ( x, y ) =
p te ( y ) p tr ( x | y ) .
The label shift problem admits the following three distinct settings (see Table 1 for a summary): (1) Fixed label shift .
Here, one assumes a-priori knowledge of p te ( y ) .
One may then adjust the outputs of a probabilistic classiﬁer post-hoc to improve test performance (Elkan, 2001).
Even when the precise distribution is unknown, it is common to posit a uniform p te ( y ) .
Minimisation of the resulting balanced error has been the subject of a large body of work (He & Garcia, 2009), with recent developments including Cui et al.
; Cao et al.
; Kang et al.
(2) Estimated label shift .
Here, we assume that p te ( y ) is unknown, but that we have access to an unlabelled test sample.
This sample may be used to estimate p te ( y ) , e.g., via kernel meanmatching (Zhang et al., 2013), minimisation of a suitable KL divergence (du Plessis & Sugiyama, 2014), or using black-box classiﬁer outputs (Lipton et al., 2018;
Azizzadenesheli et al., 2019; Garg et al., 2020).
One may then use these estimates to minimise a suitably re-weighted empirical risk.
(3) Adversarial label shift .
Here, we assume that p te ( y ) is unknown, and guard against a suitably deﬁned worst-case choice.
Observe that an extreme case of label shift involves placing all probability mass on a single y ∗ ∈ Y .
This choice can be problematic, as (1) may be rewritten as
[Equation]
where n y 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 ∗ ∈ Y has a large per-class loss, then an adversary could degrade performance by choosing a p te with p te ( y ∗ ) being large.
One means of guarding against such adversarial label shifts is to minimise the minimax risk (Alaiz-Rodr ´ ıguez et al., 2007; Davenport et al., 2010; Hashimoto et al., 2018;
Mohri et al., 2019; Sagawa et al., 2020)
[Equation]
where ∆ L denotes the simplex.
In (2), we combine the per-label risks according to 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 to enforce that (cid:107) π (cid:107) ∞ ≤ 1 k for suitable k , which corresponds to minimising the average of the top- k largest per-class losses for integer k
(Williamson & Menon, 2019; Curi et 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 is 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.
; 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 ( ≤ 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.
A LGORITHM : DISTRIBUTIONALLY ROBUST KL- DIVERGENCE MINIMISATION
To address the adversarial label shift problem, we propose to replace the empirical risk (1) with
[Equation]
where P is an uncertainty set containing perturbations of the empirical distribution p emp .
This is 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
[Table]
DRO with P being a parameterised family of distributions with varying marginal label distributions in KL-divergence, i.e., d ( p, q ) =
E y ∼ q [ − log p ( y ) /q
( y )] .
(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 a special case of (3) with r = + ∞ .
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 on 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 to 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 a constrained convex optimization problem at every iteration, which can incur non-trivial overhead (see Appendix E).
We now describe A DV S HIFT (Algorithm 1), our approach to solve these problems.
In a nutshell, we iteratively update model parameters θ and adversarial distributions π .
In the former, we update exactly as per ERM optimization (e.g., ADAM, SGD), which we denote as 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 now detail.
3. E STIMATING THE ADVERSARIAL MINIBATCH GRADIENT
For a ﬁxed π ∈ ∆ L , to estimate the parameter gradient E π [ ∇ θ (cid:96) ( x, y, θ )] on a training sample S = { ( x i , y i ) } ni =1 , we employ the importance weighting identity and write
[Equation]
n
[Equation]
We may thus draw a minibatch as usual from S , and apply suitable weighting to obtain unbiased gradient estimates.
A similar reweighting is necessary to compute the adversary gradients E
[Equation]
, θ )] .
Making the adversarial update efﬁcient requires further effort, as we now discuss.
3. R EMOVING CONSTRAINTS BY L AGRANGIAN DUALITY
To efﬁciently update the adversary distribution π in (3), we would like to avoid the cost of projecting onto P .
To bypass this difﬁculty, we make the following observation based on Lagrangian duality.
Proposition 1.
Suppose (cid:96) is bounded, and p emp is not on the boundary of the simplex.
Then, ∀ r > 0 , ∃ γ ∗ > 0 such that for every γ c ≥ γ ∗ , the constrained objective is solvable in unconstrained form:
[Equation]
Motivated by this, we may thus transform the objective (3) into:
[Equation]
(4) where γ c > 0 is a sufﬁciently large constant; in practice, this may be chosen by a bisection search.
The advantage of this formulation is that it admits an efﬁcient update for π , as we now discuss.
3. A DVERSARIAL DISTRIBUTION UPDATES
We now detail how we can employ proximal mirror descent to efﬁciently update π .
Observe that we may decompose the adversary’s (negated) objective into two terms: f ( θ, π ) :
= − E π [ (cid:96) ( x, y, θ )] and h ( π ) :
= max { 0 , γ c (KL( π, p emp ) − r ) } , where h ( π ) is independent of the samples.
Such decomposable objectives suggest using proximal updates (Combettes & Pesquet, 2011):
[Equation]
where λ serves as the learning rate.
The value of proximal descent relies on the ability to efﬁciently solve the minimisation problem in (5).
Unfortunately, this does not hold as-is for our choice of h ( π ) , essentially due to a mismatch between the use of KL-divergence in h , and Euclidean distance (cid:107)
π t − π (cid:107) 2 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:
[Equation] [Equation]
where g t is an unbiased estimator of ∇ π f ( θ t , π t ) .
We have the following closed-form update.
Lemma 2.
Assume the optimal solution π t +1 to (6) satisﬁes KL( π t +1 , p emp ) (cid:54) =
r , and that all the classes appeared at least once in the empirical distribution, i.e. ∀ i, p i emp > 0 .
Let γ = γ c if r <
KL( π t +1 , p emp ) , and γ = 0 if r > KL( π t +1 , p emp ) , then π t +1 permits a closed form solution
[Equation]
where η π = 1 (
[Equation]
+1 /
[Equation]
)(1+ α ) , α = 2 γλ , C = (cid:107) ( π t (cid:12) p α emp ) 1 / (1+ α ) exp ( η π g t ) (cid:107) 1 projects π t +1 onto the simplex, and a (cid:12) b is the element-wise product between two vectors a, b .
In Algorithm 1, we set γ = γ c if r < KL( π t , p emp ) and 0 otherwise to appoximate the true γ .
Such approximation works well when r − KL( π t , p emp ) does not change sign frequently.
3. C ONVERGENCE A NALYSIS
We provide below a convergence analysis of our gradient descent-proximal mirror ascent method for nonconvex-concave stochastic saddle point problems.
For the composite objective min θ max π ∈ ∆
L f ( θ, π ) + h ( π ) , and ﬁxed learning rate η θ , we abstract the Algorithm 1 update as:
[Equation]
[Equation]
where g ( π ) , g ( θ ) are stochastic gradients assumed to satisfy the following.
Assumption 1.
The stochastic gradient g ( θ ) with respect to θ satisﬁes that for some σ > 0 ,
E [ g ( θ )]
= ∇ θ f
( θ, π ) , and E [ (cid:107)
[Equation]
Assumption 2.
The stochastic gradient g ( π ) with respect to π satisﬁes that for some G > 0 ,
[Equation]
We make the following assumptions about the objective, similar to Lin et al.
Assumption 3.
f ( θ, π ) + h ( π ) is L − smooth and l − Lipschitz; f ( θ, π ) and h ( π ) are concave in π .
Assumption 4.
Every adversarial distribution iterate π t satisﬁes KL( π t , p emp ) ≤ R for some R > 0 .
Assumption 3 and 4 may be enforced by adding a constant (cid:15) to the adversarial updates, which prevents π t 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 signiﬁcantly speed up training (see Appendix C.3).
Furthermore, this is a standard assumption for analyzing nonconvex-concave problems (Lin et al., 2019).
The assumption that the square L ∞ norm is bounded is weaker than L 2 norm being bounded; such a relaxation results from using mirror rather than Euclidean update.
Given that the function F ( θ ) := max π ∈ ∆ f ( θ, π )+ h ( π ) is nonconvex, our goal is to ﬁnd a stationary point instead of approximating global optimum.
Yet, due to the minimax formulation, the function
F ( θ ) may not necessarily be differentiable.
Hence, we deﬁne convergence following some recent works (Davis & Drusvyatskiy, 2019;
Lin et al., 2019; Thekumparampil et al., 2019) on nonconvexconcave optimisation.
First, Assumption 3 implies F ( θ ) is L − weakly convex and l -Lipschitz (Lin et al., 2019, Lemma 4.7).
Hence, we deﬁne stationarity in the language of weakly convex functions.
Deﬁnition 1.
A point θ is an (cid:15) − stationary point of a weakly convex function F if (cid:107)∇ F 1 / 2 L ( θ ) (cid:107) ≤ (cid:15) , where F 1 / 2 L ( θ ) denotes the Moreau envelope F 1 / 2 L ( θ )
=
min w F ( w ) + L (cid:107) w − θ (cid:107) 2 .
With the above deﬁnition, we can establish convergence of the following update:
Theorem 3 (informal) .
Under Assumptions 1–4, the update in (8) ﬁnds a point θ with E [ (cid:107)∇ F 12 L ( θ ) (cid:107) ] ≤ (cid:15) in O ( (cid:15) − 8 ) iterations.
For a precise description of the theorem, please see Appendix H.
The above result matches the best known rate in Lin et al.
for optimising nonconvex-concave problem with stochastic gradients.
To our knowledge, this is the ﬁrst result that studies convergence of composite objectives with proximal methods under nonconvex-concave settings.
By utilizing the proximal operator, it solves the objective with an extra h ( π ) term without incurring additional complexity cost.
3. D ISCUSSION AND COMPARISON TO EXISTING ALGORITHMS
A 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 A DV S HIFT and related techniques in the literature.
For the problem of minimising the worst-case loss (2) — which is equivalent to setting the radius r = + ∞ in (3) — Sagawa et al.
(2020) propose an algorithm that assumes the ability to sample data from a 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 a non-trivial computational overhead due to the projection step.
We will illustrate these problems in our subsequent experiments (see results for A GNOSTIC in § 4).
Finally, for an uncertainty set P based on the CVaR, Curi et al.
(2019) provide an algorithm that updates weights using E XP 3.
This approach relies on a determinantal point process, which has a poor dimension-dependence.
E XPERIMENTAL RESULTS
We now present a series of experiments to evaluate the performance of the proposed A DV S HIFT algorithm and how it compares to related approaches from the literature.
We ﬁrst 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 signiﬁcantly outperforms other methods discussed in Table 1.
Similar results on CIFAR-100 are shown in the Appendix.
4. E XPERIMENTAL SETUP
To evaluate the proposed method, we use the standard image classiﬁcation setup of training a ResNet50 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 as follows.
First, we train a model on the training set, and compute its error distribution on the validation set.
Next, we pick some threshold τ on the allowable KL divergence between the train and target distribution.
We then ﬁnd the adversarial distribution within this threshold which achieves the worst-possible validation error.
Finally, we compute the validation performance under this distribution.
Note that τ = 0 corresponds to using the train distribution, while τ = + ∞ corresponds to using the worst-case label distribution (see Figure 1).
[Figure]
T r a i n e rr o r r e
d
u c t
i o n
[Figure]
4 KL divergence [Figure]
V a li d a t i o n e rr o r r e
d
u c t
i o n ( [Figure] Ablation validation.
4 KL divergence [Figure]
T r a i n e rr o r r e
d
u c t
i o n (
[Figure]
Comparison train.
[Figure]
[Figure]
Even at τ = 1 , the adversarial distribution is highly concentrated on only a few hard labels.
We evaluate the following methods, each corresponding to one row 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 r (A DV S HIFT ) (v) training with A DV S HIFT with a ﬁxed adversarial distribution extracted from Figure 3(c) (F IXED ).
This corresponds to the estimated validation distribution row in Table 1 with an ideal estimator.
4. R ESULTS 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 6 in the Appendix for unnormalised plots.)
Subﬁgures (a) and (b) compre the performance of A DV S HIFT for various choices of radius r to the ERM baseline; (c) and (d) compare A DV S HIFT to the remaining methods.
Hyperparameters for each method are separately tuned.
F IXED 1, 2, 3 corresponds to training with each of the three adversarial distributions in Figure 3(c).
We see that:
• the reduction offered by A DV S HIFT is consistently superior to that afforded by the AGNOSTIC , B ALANCED and F IXED methods.
On the training set, we observe signiﬁcant ( ∼ 8% ) reduction in performance for large KL divergence thresholds.
On the validation set, the gains are less pronounced ( ∼ 2 . 5% ), indicating some degradation due to a generalisation gap.
[Figure] • while A DV S HIFT consistently improves above the baseline across adversarial radii, we observe best performance for r = 0 .
Smaller values of r lead to smaller improvements, while training becomes increasingly unstable for larger radii.
Please see the discussion in the last section.
• during training, AGNOSTIC either learns the adversarial distribution too slowly (such that it behaves like ERM), or uses too large a learning rate for the adversary (such that the training fails).
This highlights the importance of the proximal mirror ascent updates in our algorithm.
Illustration of distributions at ﬁxed KL thresholds .
Figure 3(c) visualises the adversarial distributions corresponding to a few values of the KL threshold τ .
At a threshold of τ = 3 , the adversarial distribution is concentrated on only a few hard labels.
Consequently, the resulting performance on such distributions is highly reﬂective of the worst-case distribution that can happen in reality.
Training with a ﬁxed adversarial distribution .
Suppose we take the ﬁnal adversarial distributions shown in Figure 3(c), and then employ them as ﬁxed distributions during training; this corresponds to the speciﬁed 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 signiﬁcantly worse on the “easy” classes.
Employing a ﬁxed adversarial distribution may thus lead to underﬁtting, which has an intuitive explanation: the model must struggle to ﬁt difﬁcult 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, Figureshows violin plots of the distribution of errors for both the BASELINE and our A DV S HIFT methods after ﬁxed training epochs.
We observe that on the training set, A DV S HIFT signiﬁcantly 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 a generalisation gap.
Evolution of learned adversarial weights .
To understand the evolution of the adversarial distribution across training epochs, Figure 4 plots the histogram of adversary weights at ﬁxed training epochs.
Starting off from a uniform distribution, the adversary is seen to quickly infer the relative difﬁculty of a small fraction of labels, assigning ∼ 2 × the weight on them compared to the average.
In subsequent iterations the distribution becomes more concentrated, and gradually reduces the largest weights.
In summary, our experiments show that our proposed DRO formulation can be effectively solved with A DV S HIFT , and results in a model that is robust to adversarial label shift.
D ISCUSSION AND FUTURE WORK
We proposed A DV S HIFT , 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 2(a)(b) shows that the generalization gap increases as the perturbation radius increases.
Understanding why there is a correlation between hard examples and bad generalization could improve robustness.
Second, Figure 2(a) shows that even on the train set, the algorithm threshold r does not translate to the model’s level of robustness.
We conjecture this results from the interplay of model expressivity and data distribution, whose future study is of interest.
A R ELATED PROBLEMS
Example-level DRO .
Existing work on DRO has largely focussed on the setting where P encompasses 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 et al., 2017), and involves a more challenging problem, as it requires parameterising the adversary’s distribution.
Hu et al.
(2018) illustrate the potential pitfalls of DRO, owing to a 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 to 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 to 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 et al., 2020).
When the subgroups correspond to 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 ﬁxed family, GANs involve a 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 p tr and p te can systematically differ.
Typically, one assumes access to a small sample from p te , which may be used to estimate importance weights (Combes et al., 2020), or samples from multiple domains, which may be used to estimate a generic domain-agnostic representation (Muandet et al., 2013).
In causal inference, there has been interest in similar classes of models (Arjovsky et al., 2019).
B A LGORITHM IMPLEMENTATION DETAILS
We introduce some additional details in our implementation of A DV S HIFT .
First, as observed in
Section 3.1, our algorithm requires knowing the empirical label distribution.
As the exact value is not always available, we estimate the empirical label distribution online for all the experiments presented later in Section 4 using an exponential moving average, p emp = β · p emp
+ (1 − β ) · p batch , where p batch is the label distribution in the minibatch.
We set β = 0 .
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 addition to estimating the empirical distribution, we also observe that the adversarial’s update could be very sensitive to the adversarial gradient g k , i.e. 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 g k , we clip the label-wise loss at value 2 .
Furthermore, we add a constant (cid:15) term on the adversarial distribution to avoid the adversarial distribution reaching any of the vertices on the simplex.
The (cid:15) term and clipping is critical in both training and convergence analysis.
We conduct an ablation of the sensitivity to these parameters in Appendix C.3.
In all the experiments, we set 2 γ c λ = 1 in Algorithm 1 for simplicity.
For learning the adversarial distribution, we only tune the adversarial learning rate η π .
C A DDITIONAL 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 5)
[Figure]
• 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 a sense of the absolute performance numbers obtained by each method.
(Figure 6) • an ablation of the loss clipping threshold and gradient stabiliser (cid:15) as introduced above.
(Figure 7,8) • results on CIFAR-100, to complement those for ImageNet.
(Figure 9,10)
C. B ALANCED LABELS (cid:54) = ⇒ BALANCED PERFORMANCE
Figure 5 shows that training label frequency does not strongly correlate with test error.
Observe that several classes with a 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. U NNORMALISED PLOTS ON I MAGE N ET
Figure 6 presents plots of the unnormalised performance of the various methods compared in the body.
Here, rather than subtract the performance of the baseline, we show the absolute accuracy of each method as the adversarial radius is varied.
Evidently, the baseline and A GNOSTIC models tend to suffer in their validation error as the adversarial radius increases.
C. A BLATION OF CLIPPING THRESHOLD AND GRADIENT STABILISER
Figures 7 and 8 show an ablation of the choice of loss clipping threshold, and the gradient stabiliser (cid:15) .
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 performance rapidly degrades as the adversarial radius increases.
Conversely, performance also suffers when the stablisation is too high.
C. R ESULTS ON CIFAR-
Figure 9 shows results on CIFAR-100, where we train various methods using a CIFAR-ResNetas the underlying architecture, Here, we see a consistent and sizable improvement from A DV S HIFT over the baseline method.
On this dataset, A GNOSTIC fares better, and eventually matches the performance of A DV S HIFT with a large adversarial radius.
This is in keeping with the intended use-case of A GNOSTIC , i.e., minimising the worst-case loss.
Figure 10 supplements these plots with unnormalised versions, to illustrate the absolute performance differences.
We remark here that the choice of a CIFAR-ResNet-18 results in an underparameterised model, which does not perfectly ﬁt the training data.
In the overparameterised case, there are challenges with employing DRO, as noted by Sagawa et al.
Addressing these challenges in settings where the training data is balanced remains an interesting open question.
[Figure] [Figure]
[Figure]
n e rr o r r e
d
u c t i o n
[Figure]
We start with a simple example with three label classes { a, b, c } with class losses l = { 1 , 2 , 4 } respectively. We assume an uniform empirical distribution, i e. p emp = We consider two different problems.
The ﬁrst is to ﬁnd the optimal solution to regularised objective, p = argmin p p (cid:62) l + γ KL( p, p emp ) . [Figure] [Figure] This problem is well known (e
[Figure]
g. see 2.7.2 of lecture ) to permit a solution of form p ( x ) = exp l
[Figure]
for some t .
In contrast, we show that distributions of the above form does not solve the constrained version of the problem.
In particular, we consider the following optimisation problem:
[Equation]
such that KL ( p, p emp ) ≤ r
If the solution is of form p ( x ) =
[Equation]
, then we know for l b
[Equation]
l
[Equation]
We solve the above problem with a convex optimizer using r = 0 . 01 and found p a = 0 .
283 , p b = 0 .
322 , p c = 0 .
[Equation]
Note that the above example shows that not all solutions of the contrained problem can be written a Boltzman distribution, i.e. p ( x ) =
[Equation]
Yet, this does not contradict with results [e.g. Lemma 4, Faury et al.
that claim there is a Boltzman distribution whose function value matches the optimal value of the constraint problem.
Mathematically, we can have p (cid:54) = p (cid:48) but E p [ l ( x )] =
E p (cid:48)
[ l ( x )] .
E P ROJECTING AN ADVERSARIAL DISTRIBUTION
The projection operator in our setting aims to project a distribution p into the set P = { q : KL( q, ˆ p ) ≤ r } by solving the following problem:
[Equation]
such that (cid:88) i
[Equation]
[Equation]
[Equation]
where q i , p i denotes the i th component of q, p and n denotes number of classes.
Given that our implementation is based on Tensorﬂow, 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 1 minute when n = 1000 (the number of labels in ImageNet).
This means that if we train ResNet50 on ImageNet for 100k iterations, we need to spend 100 k minutes on projection operation, which is not affordable.
F P ROOF OF P ROPOSITION
Proof.
We only need to show that for large enough γ c , any minimiser p ∗ of the unconstrained problem satisﬁes that KL( p ∗ , p emp ) ≤ r .
Since the distance from boundary of the simplex to any interior point is + ∞ , we can safely assume that the point lies within the relative interior of the simplex.
To prove the proposition, denote c = inf {(cid:107)∇ p KL( p, p emp ) (cid:107) | p ∈ ∆ L , s.t. KL( p, p emp ) > r } .
By strict convexity of KL divergence and the fact that r > 0 , we know c > 0 .
Denote the upper bound of loss as M , then when γ c > M/c , we know that KL( p, p emp ) >
r = ⇒
(cid:54)∈ ∂ π ( E π [ (cid:96) ( x, y, θ )]
+ min { 0 , γ c ( r − KL( π, p emp )) } ) .
However, since p minimises the objective if and only if 0 ∈ ∂ π ( E π [ (cid:96) ( x, y, θ )]
+ min { 0 , γ c ( r − KL( π, p emp )) } ) , we have that any minimiser p must satisfy KL( π, p emp )) ≤ r .
G P ROOF OF L EMMA
Recall that we want to ﬁnd π k +1 that minimises the following objective,
[Equation]
where η = 1 / (2 γ c + 1 /λ ) , α c = 2 γ c λ .
Denote v ( i ) as the i th component of vector v .
Notice that the simplex can be written as a constraint that (cid:80)
i π ( i ) = 1; ∀ i, π ( i ) ≥ 0 .
Based on this constraint, we ﬁrst write (6)’s
Lagrangian dual
L ( a, b, π ) =
(cid:88)
[Equation]
where π ( i ) denotes the i th component of π .
If π k > 0 component-wise, then the optimal π cannot lie on the boundary (i.e. ∀ i, π ( i ) > 0 ), which results in KL( π, π k ) = ∞ .
By Lagrangian duality and complementary slackness, we know that for if KL( π, p emp ) >
r , 0 =
[Equation]
On the other hand, if KL( π, p emp ) < r
[Equation]
We discuss the case when KL( π, p emp ) < r , and the other case follows similarly.
Rearrange the optimality condition of Lagrangian multiplier, we get
[Equation]
Since b is a constant for all coordinates,
[Equation]
The result follows by noting that (cid:80)
i π ( i ) = 1 .
H P ROOF OF T HEOREM
For completeness, we deﬁne several terms used in optimisation.
A function f ( θ ) is l − Lipschitz if for all θ, θ (cid:48) , | f ( θ ) − f ( θ (cid:48) )
| ≤ l (cid:107) θ − θ (cid:48) (cid:107) .
A function f ( θ ) is L − smooth if for all θ, θ (cid:48) , (
cid:107)∇ f ( θ ) − ∇ f
( θ (cid:48) )
(cid:107) ≤ L (cid:107) θ − θ (cid:48)
(cid:107) .
A function f ( θ ) is L − weakly convex if f ( θ ) + L 2 (cid:107) θ (cid:107) 2 is convex.
Then we can state the formal theorem below.
Theorem 4 (formal version of Theorem 3) .
Under Assumptions 1–4, the update in (8) generates a sequence of points θ 1 , ..., θ T with the following property:  T (cid:88)
[Equation]
Proof.
For convenience [Equation] denote Φ( θ
[Equation]
π ) = f ( θ, π ) + h ( π ) [Equation] F ( θ ) =
max p Φ( θ, p ) .
We start by following the standard SGD proof.
Denote g θ as the stochastic gradient evaluated at step t − 1 with respect to θ .
Denote ˆ θ = prox F/ 2 L ( θ ) :
= arg min w { F ( w ) + 2 L (cid:107) w − θ (cid:107) 2 } .
Conditioned on θ t − 1 , we have
[Equation]
where the ﬁrst equality follows by θ t
= θ t − 1 − η θ ∇ θ Φ( θ t − [Equation] π t − 1 ) .
Next
[Equation]
we observe that
[Equation]
The ﬁrst line follows by convexity and L − smoothness.
The second line by deﬁnition of F .
The third line by deﬁnition of ˆ θ .
Next, by deﬁnition of Moreau envelop,
[Equation]
Take expectation on both sides and we get
[Equation]
The second line substitutes in (9).
The third line follows by convexity and L − smoothness.
Denote that ∆ t :=
F (ˆ
[Equation]
−
[Equation]
π
[Equation] [Equation]
− Φ
[Equation]
θ
[Equation]
− 1 , π
[Equation]
− 1 ) .
We can sum over t and take expectation recursively to get,
[Equation]
where F ∗ 1 / 2 L = min θ F 1 / 2 L ( θ )
The ﬁrst equality follows by the deﬁnition of Moreau envelope.
The second inequality follows by rearranging (11).
Next, we aim to bound the accumulated error (cid:80) t
[Equation] Recall that the update for the π is as follows for E [ g π ] = ∇ π f ( θ, π ) ,
[Equation]
Applying Lemma 5 with L ( π ) = − 2 λ (cid:104)
g π
[Equation]
π (cid:105) + − 2 λh ( π )
[Equation]
we get − h ( π ∗ (
[Equation] )) − [Equation]
Rearrange and take expectation we get
[Equation]
The second inequality follows by the fact that KL − divergence is strongly convex with respect to Lnorm and Cauchy-Schwartz inequality.
We further observe that
[Equation]
The ﬁrst inequality follows by concavity.
The third line follows by f ( θ s , π ∗ ( θ k ))
≤
f ( θ s , π ∗ ( θ s )) .
The last inequality follows by Lipschitzness.
Similarly,
[Equation]
We can take iterative expectation and get sum over k
= s
+ 1 , ..., s + B to get s +
[Equation]
k [Equation] +
[Equation]
Note that ∆ k = − f
( θ k , π k )
+ f ( θ k , π ∗ ( θ k ))
− h ( π k ) + h ( π ∗ ( θ k )) , hence
[Equation]
By further sum over all blocks and divide by total number of iterations T , we get
[Equation]
Substitute the above inequality into (12) and we get
T
[Equation]
If we set η θ = T − 3 / 4 , B = T 1 / 2 , λ = T − 1 / 4 , then we see that T
[Equation]
Lemma 5.
For any differentiable convex function L , if x ∗ = argmin x ∈ ∆ { L ( x ) + KL( x, x 0 ) } , then for any x (cid:48) ∈ ∆ , we have
[Equation]
Proof.
This Lemma is well-known, but we include a proof for completeness.
By optimality of x ∗ and convexity of ∆ , we know that
[Equation]
where φ ( x ) = (cid:80)
i x i log( x i ) , and the Bregman divergence deﬁned according to φ is KL-divergence.
Then
[Equation]