Distributionally Robust Optimization with Data Geometry
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Distributionally Robust Optimization (DRO) serves as a robust alternative to empirical risk minimization (ERM), which optimizes the worst-case distribution in anuncertainty set, typically speciﬁed by distance metrics including f -divergence andthe Wasserstein distance.
The metrics deﬁned in the ostensible high dimensionalspace lead to exceedingly large uncertainty sets, resulting in the underperformanceof most existing DRO methods.
It has been well documented that high dimensionaldata approximately reside on low dimensional manifolds.
In this work, to furtherconstrain the uncertainty set, we incorporate data geometric properties into thedesign of distance metrics, resulting in our novel Geometric Wasserstein DRO(GDRO).
Empowered by Gradient Flow, we derive a generically applicable approximate algorithm for the optimization of GDRO, and prove the bounded errorrate of the approximation as well as the convergence of our algorithm.
We alsotheoretically characterize the edge cases when certain existing DRO methods aredegeneracy of GDRO.
Extensive experiments justify the superiority of our GDROto existing DRO methods in multiple settings with strong distributional shifts, andconﬁrm that the uncertainty set of GDRO adapts to data geometry.
Introduction
Machine learning algorithms with empirical risk minimization often suffer from poor generalizationperformance under distributional shifts in real applications due to the widespread latent heterogeneity, domain shifts, and data selection bias, etc .
It is demanded for machine learning algorithms toachieve uniformly good performances against potential distributional shifts, especially in high-stakeapplications.
Towards this goal, distributionally robust optimization (DRO) [17, 14, 19, 3, 11, 9],stemming from the literature of robust learning, has been proposed and developed in recent years.
It optimizes the worst-case distribution within an uncertainty set P ( P tr ) lying around the trainingdistribution P tr .
When the testing distribution P te is contained in P ( P tr ) , DRO could guarantee thegeneralization performance on P te .
In principle, the effectiveness of DRO heavily depends on the rationality of its uncertainty set P ( P tr )which is commonly formulated as a ball surrounding the training distribution endowed with a certaindistance metric.
An ideal uncertainty set should be constituted by all realistic distributions that maybe encountered in test environments.
However, the existing DRO methods adopting Wassersteindistance (i.e. WDRO methods [17, 19, 3, 11]) or f -divergence distance (i.e. f -DRO methods [14, 9])tend to generate over-ﬂexible uncertainty sets that incorporate unrealistic distributions far beyond theideal uncertainty set [13, 11].
As such unrealistic distributions must violate the underlying predictingmechanism, they are prone to be the worst-case and attract much optimization energy in the DROframework, making the learned model deviate from the true predicting mechanism.
[Figure]
Here we argue that the unrealistic distributions mentioned above originate from the distance metrics’inherent ignorance of data geometry, as illustrated in Figure 1.
The Euclidean-norm transportationcost measured by the L 2 -Wasserstein metric (Equation 2) leads to a straight-line transportation path asshown in Figure 1(a)
(red dotted line) which deviates from the data manifold (blue region).
Therefore,
WDRO methods tend to create unrealistic samples beyond the underlying data manifold, resultingin unrealistic distributions.
f -divergence can also be interpreted as a data geometry-independentmeasure of the transportation cost conﬁned in the support of P tr .
Taking χ 2 -divergence for example,the cost is a constant to transfer per unit of probability weights between samples, like a virtual tunnel(yellow dotted line in Figure 1(a)).
In such a case, the noisy samples (e.g. outliers or samples withlabel noises) are more prone to be the worst-case and thus gather much larger weights than normalsamples.
The distribution with such a support is obviously unrealistic.
To mitigate the problem, it is demanded to introduce a new distance metric incorporating datageometry to further constrain the uncertainty set and avoid the undesired cases.
As illustratedin Figure 1(a), considering the common assumption that data lie on a low-dimensional manifold[16, 20, 2], we expect the probability density transportation path (the blue dotted line) is restrictedwithin the data manifold (the blue region).
In this way, the resulted uncertainty set (i.e. the Geometric
Wasserstein Set as shown in Figure 1(b)) could inherently exclude the distributions beyond the datamanifold.
Furthermore, it is harder to gather probability weights on isolated noisy samples, whichalso mitigate the undesired cases in f -DRO.
In this work, we propose a novel Geometric Wasserstein DRO (GDRO) method by exploiting thediscrete Geometric Wasserstein distance [4] which measures the transportation cost of probabilitydensity along the geodesic in a metric space.
As the Geometric Wasserstein distance does notenjoy an analytical expression, we derive an approximate algorithm utilizing Gradient Flow inthe Finsler manifold endowed with Geometric Wasserstein Distance (in section 3.2).
We furthertheoretically specify an exponentially vanishing error rate of our approximation as well as a O (1 / √ T )convergence rate of our algorithm, and characterize the edge cases when GDRO will degenerate tof -DRO or Wasserstein DRO (in section 3.3 and 3.4).
Comprehensive experiments encompassingvarious distributional shifts, including sub-population shifts, domain shifts, and class difﬁculty shifts,validate the effectiveness of our proposed GDRO (in section 4).
We also observe a lower Dirichlet Energy (i.e. higher smoothness) of GDRO’s estimated worst-case distribution w.r.t the data manifoldcompared with existing DRO methods, justifying its adaptability to data geometry.
Preliminaries on Distributionally Robust Optimization
Notations.
X ∈ X denotes the covariates,
Y ∈ Y denotes the target, f θ ( · ) :
X → Y is the predictorparameterized by θ ∈ Θ .
P tr ( X, Y ) and P te ( X, Y ) abbreviated with P tr and P te respectivelyrepresent the joint training distribution and test distribution.
The random variable of data points isdenoted by Z =
( X, Y ) ∈ Z .
Distributionally Robust Optimization (DRO) is formulated as:
[Equation]
(1) where (cid:96) is a loss function and P ( P tr )
= { P : Dist ( P, P tr ) ≤ (cid:15) } characterizes the uncertaintyset surrounding the training distribution restricted by a radius (cid:15) in terms of Dist , a distance metric between probability distributions.
The constraints for the uncertainty set distinguish various DROmethods.
Most literature speciﬁes Dist as the Wasserstein distance [17, 19, 3] or f -divergence [14, 9].f -divergence DRO ( abbr.
f -DRO) f -divergence is deﬁned as D f ( P (cid:107) Q ) =
(cid:82) f ( dP/dQ ) dQ ,where f ( · ) is a convex function and f (1) = 0 .
Two typical instances of f -divergences are KLdivergence ( f ( t ) =
t log t ) and χ 2 -divergence ( f ( t ) =
( t − 1) 2 ).
As shown in Figure 1(a), fdivergence ignores the data geometry and only cares about the probability of each sample (onlydP, dQ occur).
Since it is well-accepted that data lie on a low-dimensional manifold and adjacentdata points have similar degrees of importance, data geometry information is crucial for a reasonableuncertainty set.
Such ignorance induces many undesired cases in the f -divergence ball.
For example,for heterogeneous data, one hopes to focus on some sub-population (e.g., put more weights on agroup of data), but without data geometric information, the distribution in the f -divergence ballis prone to only focus on some isolated samples with higher noises (as shown in Figure 2(a)).
Corresponding with our analysis, in Figure 2(b), we ﬁnd the worst-case distribution found by f -DRO(with KL-divergence) has lower smoothness (larger Dirichlet Energy) w.r.t.
data manifold.
Wasserstein DRO ( abbr.
WDRO)
Despite the common deﬁnition of Wasserstein distance by 89 Kantorovich, we shall introduce the Benamou-Brenier reformulation of L 2 -Wasserstein metric whichcould be readily generalized to Geometric Wasserstein Distance [21].
For probability measures P and 91 Q supported on X , consider a transport plan p ( t, x ) : ([0 , 1] , R n ) (cid:55)→ R continuous on [0 , 1] whichtransforms the measure P into Q over the time interval [0 , 1] .
Here p ( t, x ) abbreviated with p t ( x )denotes the probability density of x at time t and is an absolutely continuous probability measure.
Then the L 2 -Wasserstein metric could be deﬁned as:
[Equation]
where v : ([0 , 1] , R n ) (cid:55)→ R n 1 is a time-dependent vector ﬁeld satisfying the continuity functionin Equation 2.
Intuitively v is a velocity ﬁeld continuously transporting masses to convert thedensity distribution from P to Q along a curve in the Wasserstein space.
And the Wasserstein Distance measures the geodesic length which is shortest among all possible plans.
The length ofa curve in Wasserstein space is calculated in Equation 2 by integrating a total "kinetic energy" ofthe velocity ﬁeld over the transportation process.
WDRO is empowered with the ﬂexibility to createsamples beyond the training distribution, thus inspiring applications for adversarial robustness againstrestricted attacks [19].
It causes an over-pessimistic estimation under strong distributional shifts [11]where mislabeling prevails for perturbed samples as demonstrated by our experiments (Figure 2(a)).
Proposed Method
We take as input a training dataset D tr = { ( x i , y i ) } ni =1 and the corresponding manifold structuretypically approximated by a weighted ﬁnite graph G
= ( V, E, w ) , where V = { 1 , . . . , n } is thevertex set, E is the edge set and w = ( w ij ) i,j ∈ V are the edge weights.
In case the graph G 0 is notpredeﬁned, the k-nearest-neighbor (kNN) algorithm sufﬁces for a ﬁne approximation.
For large-scaledatasets, we use the NN-Descent method [8] for approximate kNN graph construction, which has alow empirical complexity of O ( n 1 . 14 ) and can be easily parallelized.
In the following of this section,we ﬁrst introduce the Geometric Wasserstein Distance and propose the overall objective of GDRO;then we derive an approximate algorithm for optimization utilizing Gradient Flow; ﬁnally, sometheoretical properties are proved and connections with existing DRO methods are demonstrated.
3.1 Discrete Geometric Wasserstein Distance GW G 0 and GDRO
We ﬁrstly introduce the Discrete Wasserstein distance, which extends the Benamou-Brenier formulation of optimal transport problem in Equation 2 to a metric space.
Since the manifold structure ismerely accessible by a discrete graphical representation, one has to further discretize the BenamouBrenier formula to describe the optimal transport on the graph.
The ﬁrst step is to deﬁne a discretevelocity ﬁeld and its discrete divergence, which we mainly follow the construction by Chow et al.
For a given weighted ﬁnite graph G 0 = ( V, E, w ) with n nodes, a velocity ﬁeld v = ( v ij ) i,j ∈ V ∈ 121 R n × n on G 0 is deﬁned to be a skew-symmetric matrix on the edge set E such that v ij = − v ji if( i, j ) ∈ E .
Given a probability set (simplex) supported on V : P ( G 0 ) =
{ ( p i ) ni =1 ∈ R n | (cid:80) ni
=1 p i =
1 , p i ≥ 0 , for any i ∈ V } and its interior denoted by P o ( G 0 ) , one could thereafter deﬁne thedivergence of pv as div G 0 ( pv ) :
= − ( (cid:80) j ∈ V :(
i,j ) ∈ E √ w ij v ij θ ij ( p ))
ni
=1 which is a vector in R n .
The divergence vector is supposed to lie in the tangent space of P
o ( G 0 ) , summing over all thein-ﬂuxes and out-ﬂuxes along edges of a certain node, with each edge transporting a probabilitydensity √ w ij v ij θ ij ( p ) .
θ ij is a predeﬁned "cross-sectional area" typically interpolated with theassociated nodes’ densities p
i , p j .
The direct approach is to take the arithmetic average such thatθ ij ( p ) =
( p i + p j ) / 2 .
However, to ensure the positiveness of p during optimization, we adopt theupwind interpolation: θ ij ( p )
=
I ( v ij > 0) p j
+ I ( v ij ≤ 0) p i .
Now we are ready to deﬁne Geometric
Wasserstein Distance as a generalization from Equation 2.
Deﬁnition 3.1 (Discrete Geometric Wasserstein Distance GW G 0 ( · , · ) [4]) .
Given a ﬁnite graph G 0 ,for any pair of distributions p 0 , p 1 ∈ P
o
( G 0 ) , deﬁne the Geometric Wasserstein Distance:
[Equation]
where the inﬁmum is taken over all velocity ﬁelds on G 0 , p is a continuously differentiable curvep ( t ) : [0 , 1] → P o ( G 0 ) , and θ ij ( p ) is a pre-deﬁned interpolation function between p i and p j .
Equation 3 resembles the Benamou-Brenier formulation of continuous L 2 -Wasserstein distance inthe sense that it integrates the transportation cost over an optimal transport process characterized bythe distribution curve p ( t ) , except that the probability density could only be transferred through theedges of a graph , thus ensuring the transportation path to stay within the manifold (as the blue dottedline shown in Figure 1(a)) and inducing a smoother estimate of worst-case probability distributionw.r.t the data structure since weights are exchanged just between neighbors.
Note that the metricspace endowed with Geometric Wasserstein Distance ( P o ( G 0 ) , GW G 0 ) is a Finsler structure [10].
Then we present the overall objective of Distributionally Robust Optimization with Geometric 144 Wasserstein Distance (GDRO).
Given the training dataset D tr = { ( x i , y i ) } ni =1 and its empiricalmarginal distribution
ˆ P tr = 1 n (cid:80) i δ ( x i ) , along with a manifold structure represented by graph G 0 ,we intend to obtain a distributionally robust predictor parameterized by θ ∗ such that for certain (cid:15) > 0 :
[Equation]
We add a minor entropy-regularization with a small β as proposed in the entropy-balancing literature[12] to avoid singular cases and ensure the convergence of our optimization in section 3.2.
Since Equation 4 degenerates to a f -DRO as β → ∞ , the entropy regularized objective is an interpolationbetween GDRO and the optimization-friendly f -DRO.
Owing to the Geometric Wasserstein Distance,the uncertainty set of GDRO excludes those distributions supported on points beyond the datamanifold and the Geometric Wassserstein Ball is directional in Wasserstein space as it stretches alongthe data structure, as depicted in Figure 1(b).
3.2 Optimization
In this subsection, we derive the optimization algorithm for GDRO.
For lack of an analyticalexpression of Geometric Wasserstein Distance, we give up providing an exact solution to Equation 4.
Instead, we propose to solve the minimax objective with alternate optimization.
For ﬁxed probabilityweights p , the parameter θ could be optimized via gradient descents for R n ( θ, p ) w.r.t.
θ in parameter
Input: Training Dataset D
= { ( x i , y i ) } ni =1 , learning rate α θ , gradient ﬂow iterations T , entropy term β , manifold representation G to Epochs do
Simulate gradient ﬂow for T time steps according to Equation 6 ∼ 7 to learn an approximate worst-case probability weight p T .
θ ( i +1)
← θ ( i ) − α θ ∇ θ ( (cid:80)
i p Ti (cid:96) i ( θ ))
end for space Θ .
The inner supremum problem can be approximately solved via gradient ascents for R n ( θ, p )w.r.t.
p in the Geometric Wasserstein space.
And the cost measured by GW 2 G 0 ( ˆ P tr , · ) could beapproximated with the length of the gradient ﬂow, which is a curve in ( P o ( G 0 ) , GW G 0 ) .
Here we clarify some notations.
p : [0 , T ] (cid:55)→ P o
( G 0 ) denotes the continuous gradient ﬂow, and 164 the probability weight of the i -th sample at time t is abbreviated as p i ( t ) .
The time-discretizedgradient ﬂow corresponding with the time step τ is denoted as ˆ p τ : [0 , T ] (cid:55)→ P o ( G 0 ) , and ˆ p τ ( t ) isabbreviated as ˆ p tτ .
For the optimization, we adopt the time-discretized deﬁnition of Gradient Flow 167 [21] for − R n ( θ, p ) in the Geometric Wasserstein space as: (with the time step τ )
[Equation]
When τ → 0 , the time-discretized gradient ﬂow ˆ p τ becomes the continuous one p .
Note that
Equation 5 describes the Gradient Flow as a steepest ascent curve locally optimizing for a maximalobjective within an inﬁnitesimal Geometric Wasserstein ball, and it coincides with the Lagrangianpenalty problem of Equation 4.
In theorem 3.1 we would prove that Equation 5 ﬁnds the exactsolution to a local GDRO at each time step.
Following Chow et al.
[4], the analytical solution to Equation 5 as τ → 0 could be derived as:
[Equation]
∈ E [Equation] (6) where p i denotes the time-dependent probability function of the i -th sample, and we take an upwind 175 interpolation of θ : θ ij ( p )
=
I ( v ij > 0) p j
+ I ( v ij ≤ 0) p i , so that the probability density transferred 176 on an edge equals the density from the origin node associated with the velocity ﬁeld.
The upwindinterpolation guarantees that the probability weight p stays positive along the Gradient Flow in 178 Equation 6.
Then we discretize equation 6 with Forward Euler Method:
[Equation]
where α is a learning rate.
For our algorithm, we control the maximum time step as t ≤ T in
Equation 7 to approximately restrict the radius of the Geometric Wasserstein ball.
We prove intheorem 3.2 that for the ﬁnal time step t =
T , the probability weights p ( T ) learned by Equation 7guarantees a global error rate e − CT from the worst-case risk R n ( θ, p ∗ ) constrained in an (cid:15)
( θ ) -radiusball where (cid:15) ( θ ) =
GW 2 G 0 ( ˆ P tr , p ( T )) and p ∗ =
arg sup p { R n ( θ, p ) : GW 2 G 0 ( ˆ P tr , p ) ≤ (cid:15) ( θ ) } .
Theresult is similar to conventions in WDRO [19], which gives up providing a prescribed radius of itsuncertainty set but turns to an approximation with a intermediate hyperparameter.
Pseudo-code of thewhole algorithm is shown in Algorithm 1.
The whole derivations are in Appendix.
3.3 Theoretical Properties
We prove the equivalence of our Gradient-Flow based algorithm to a local GDRO problem and thebound of its global error rate.
We ﬁrst provide the robustness guarantee for the Lagrangian penaltyproblem in Equation 5.
Theorem 3.1 (Local Robustness Guarantees of Lagrangian Penalty Problem) .
For any τ > 192 0 , t > 0 and given θ , denote the solution of Equation 5 as p ∗ ( θ ) =
arg sup p ∈ P ( G 0 )
R n ( θ, p ) −
τ GW 2 G 0 (ˆ p tτ ( θ ) , p ) .
Let (cid:15) τ ( θ ) =
GW [Equation] G 0 (ˆ p tτ ( θ ) , p ∗ ( θ )) , we have
[Equation]
Theorem 3.1 proves that at each time step our Lagrangian penalty problem is an exact solution to alocal GDRO within the (cid:15)
τ ( θ ) -radius Geometric Wasserstein ball.
It further shows that with τ → 0 in
Equation 5, our gradient ﬂow constantly ﬁnds the steepest descent direction.
Then we theoreticallyanalyze the global error rate brought by our approximate algorithm.
Theorem 3.2 (Global Error Rate Bound) .
Given the model parameter θ , denote the approximateworst-case by gradient descent in Equation 7 after time t as p t ( θ ) , and (cid:15) ( θ ) =
GW 2 G 0 ( ˆ P tr , p t ( θ ))denotes the distance between our approximation p t and the training distribution
ˆ P tr .
Then denotethe real worst-case distribution within the (cid:15) ( θ ) -radius discrete Geometric Wasserstein-ball as p ∗ ( θ ) ,that is,
[Equation]
Here we derive the bound w.r.t.
the error ratio of objective function R n ( θ, p ) (abbr. R ( p ) ).
For 204 θ ∈ Θ , there exists C > 0 such that
[Equation]
and when t → ∞ , Error Rate → 0 .
The value of C depends on (cid:96), β, n .
Theorem 3.2 theoretically characterizes ’how far’ our approximation p t is from the real worst-case p
∗in terms of the drop ratio of the objective function R ( p ) .
At last we derive the convergence rate ofour Algorithm 1.
Theorem 3.3 (Convergence of Algorithm 1) .
Denote the objective function for the predictor as:
[Equation]
which is assumed as L -smooth and R ( θ, p ) satisﬁes L p -smoothness such that (cid:107)∇ p R ( θ, p ) −∇
[Equation] p
[Equation] (cid:107) 2 .
(cid:15) ( θ ) follows the deﬁnition in Theorem 3.2.
Take a constant∆ F ≥ F ( θ (0) )
− inf θ F ( θ ) and set step size as α = (cid:112) ∆ F / ( LK ) .
For t ≥ T 0 where T 0 is aconstant, denote the upper bound of (cid:107) p t − p ∗ (cid:107) 22 as γ and train the model for K steps, we have:
[Equation]
We assume a common smoothness of the objective function as in [19].
As K → ∞ , ∇ θ F ( θ ( k ) ) willachieve a square-root convergence only if γ is controlled by the exponentially vanishing error rate in
Theorem 3.2.
And the accuracy parameter γ remains a ﬁxed effect on optimization accuracy
3.4 Connections with Conventional DRO Methods
In Theorem 3.4, we illustrate the connections of our GDRO with f -DRO.
Theorem 3.4 (Connection with f -DRO with KL-diveregence (KL-DRO).) .
Relax the discrete Geo- 220 metric Wasserstein-ball regularization (set (cid:15) → ∞ ) and set the graph G 0 to a fully-connected graph,and then the solution of GDRO is equivalent to the following form of KL-DRO:
[Equation]
where p ∗ ( θ ) = arg max p (cid:80) ni =1 p i (cid:96) ( f θ ( x i ) , y i ) − β (cid:80) ni =1 p
i log p i .
Remark (Connections with WDRO) .
Since conventional WDRO allows distributions to extendtraining support, our proposed GDRO is intrinsically different from WDRO.
Intuitively, for inﬁnitesamples, if the graph G 0 is set to a fully-connected graph with edge weights w ij = (cid:107) z
i − z j (cid:107) 2 and βis set to 0, our GDRO resembles support-restricted version of WDRO.
Experiments
In this section, we investigate the empirical performance of our proposed GDRO on differentsimulation and real-world datasets under varying kinds of distributional shifts , including subpopulation shifts, domain shifts, and label shifts.
As for baselines, we compare with empirical riskminimization (ERM), WDRO and two typical f -DRO methods, including KL-DRO ( f ( t ) =
t ln t )and χ 2 -DRO ( f ( t ) =
( t − 1) 2 ).
4.1 Simulation Data
As for the simulation data, we simulate domain shifts between training data and testing data viaanti-causal effect and selection bias mechanism respectively.
And we also investigate the inﬂuencebrought by label noises to demonstrate that our GDRO could to some extent resist label noise.
1. Toy Example: Domain Shift via Anti-Causal Effect
Firstly, inspired by [1], we induce the domain shifts via the anti-causal effect as S ∼ N (0 , 1) :
[Equation]
[Table]
where X = [ S, V ] T , S serves as a stable feature with an unchanged relationship with Y (thus itshould be used for prediction), but V is a spurious feature with changeable relationships with Y(thus one should avoid using it), and α S is set to 5 .
for all data.
For training data, we sample10000 points, and design different settings with varying minor group ratios r .
For testing data, wesimulate 6 domains with strong shifts by varying α V ∈ {− 3 , − 2 , − 1 , 1 , 2 , 3 } ( V = α V Y + N (0 , 1) )and calculate the Mean_Error, Std_Error and parameter estimation error Est_Error for each method.
From the results in Table 1, WDRO performs similarly to ERM, and two f -DRO methods (KL-DROand χ 2 -DRO) outperform ERM and WDRO.
Such results verify our analysis that the Wassersteinuncertainty set cannot be large due to its over-ﬂexibility, which greatly impairs its performance understrong domain shifts.
Our GDRO outperforms all baselines and achieves much lower prediction errorsand estimation errors in all settings, which shows our uncertainty set built on Geometric Wassersteindistance is much more practical and reasonable.
2. Selection Bias: Domain Shift via Selection Bias Mechanism
Secondly, a more complicated mechanism is designed via selection bias as X =
[ S, V ] T :
[Equation]
Similar to the toy example above, S are stable features while the relationships between V and Yare perturbed in different domains.
Speciﬁcally, a data point is selected with probability P ( x i , y i ) =|
r | − 5 ∗| y
i − sign ( r ) · v bi | , which induces the spurious correlation between a certain covariate V b ∈ V and
Y .
In training, we generate 10000 points, where the major group contains 95% data with r = 1 .
9and the minor group contains 5% data with r = − 1 .
In testing, we ﬁrst report the performancesof the two training groups, and then we further vary r ∈ {− 1 . 5 , − 1 . 7 , − 1 .
9 , − 2 .
3 , − 2 .
7 , − 3 . 0 } tosimulate more challenging domain shifts that cannot be obtained by interpolation between traininggroups.
The average results over ten runs are shown in Table 2 (simulation 2).
As for sub-populationshifts, GDRO achieves good tail performance at a slight sacriﬁce of the performance of the majorgroup.
Further, for more challenging domain shifts, GDRO signiﬁcantly outperforms all baselines inall testing distributions, even though the shifts are much stronger than the training.
Bias Experiments.
We report the root mean square errors
[Table]
[Table]
3. Label Noises: Add Label Noises
Since DRO methods are risk-aware, they are prone to be affected by label noises.
Though the effect oflabel noises cannot be eliminated due to the nature of DRO, our proposed GDRO could signiﬁcantlyimprove the resistance of DRO methods to a minor degree of label noises.
Based on the selection biasexperiment, we random sample 20 points and add label noises on them via ˜ Y
=
Y + Std ( Y ) wherestd
( Y ) denotes the standard derivation of the marginal distribution of Y .
From results shown in Table
2(simulation 3), both f -DRO methods are signiﬁcantly affected and perform similarly to ERM undersuch a minor degree of label noises.
And our GDRO is only slightly affected by the label noises.
Further, to demonstrate the difference between f -DRO, WDRO, and GDRO, for the label noiseexperiment, we visualize the learned worst-case distribution of three methods.
In the ﬁrst twoﬁgures in Figure 2(a), we draw the learned sample weights of KL-DRO and GDRO, where red pointsrepresent the noisy samples in training data, and the size of each point is proportional to its sampleweight.
We can see that KL-DRO puts heavy weights on the noisy points (the red nodes are muchlarger), while our GDRO only slightly increases them and the weights of data in the minor groupare raised, which results in the difference between their performances in the label noise experiment.
The reason for such phenomenon of KL-DRO is that it ignores the data geometry, allowing forsome isolated nodes with much heavier weights than surrounding nodes, which corresponds with our
Theorem 3.4.
Since our GDRO intrinsically naturally constrains weight learning on the data manifold,the learned weights are smooth w.r.t the data manifold.
In the third ﬁgure in Figure 2(a), the redpoints represent the samples, to which WDRO introduces label noises with the ratio larger than 50%,and the size of each point is proportional to its label noise ratio.
We can see that although WDROextends the support, the created samples are quite noisy and greatly harm the learning process.
(3) Toquantify this property, we measure the smoothness via Dirichlet Energy and plot the Dirichlet Energyw.r.t the relative entropy KL ( p (cid:107) ˆ P tr ) between the learned distribution and training distribution in
Figure 2(b), which shows that the learned weights of GDRO are much more smooth on the datamanifold than that of KL-DRO.
3.0 RelativeEntropy 1 e1 e1 e1 e1 e1 e1 eD i r i c h l e t
E n e r g y f
-DROGDRO with =1.0 GDRO with =0.1 GDRO with =0.
(b) Weight smoothness.
4.2 Real-World Data
Finally, towards a comprehensive comparison with existing DRO methods, we evaluate our methodon four real-world datasets with various kinds of distributional shifts, including sub-population shifts,domain shifts, class-wise shifts, and label noises.
All experimental details can be found in Appendix.1.
Retiring Adults: Sub-population Shifts The Retiring Adults dataset [7] is derived from US
Census surveys, where the sub-population shifts are natural since there are geographic variationsacross different states.
Our experiments involve three prediction tasks deﬁned in [7], including
Income Prediction, Public Health Insurance Coverage Prediction, and Commuting Time Prediction.
For each task, we randomly sample 2000 points from state A and 2000 · r points from state B , wherer < 1 and the second state is viewed as the minor group.
In Figure 3, we report the predictionaccuracy of the minor group for different minority ratios r in three tasks.
HIV-1:
Sub-population Shifts & Label Imbalance HIV-
Protease Cleavage Dataset
[5]involves a task to predict whether an octamer would be cleaved by HIV-1 protease, given a 160dimensional one-hot vector encoding the sequence of 8 amino acids composing the octamer.
Thedataset contains 4 splits, and following [15], we merge ’746’ and ’1625’ as subpopulations A andview ’Impens’ as B .
In training, subpopulations A and B are mixed at a ratio of 1 : r , where r ≤ 1is the minor group ratio.
Coupled with the sub-population shift, class labels are imbalanced with33% positive in sub-population A and 16% positive in B , which is more challenging.
In testing, were-balance the two classes and plot the overall accuracy w.r.t.
the minor group ratio in Figure 4.
Colored MNSIT: Domain Shifts & Label Noises Following Arjovsky et al.
[1], we conduct abinary classiﬁcation task on the MNIST dataset.
Firstly, a binary label Y is assigned to each imageaccording to its digit: Y = 0 for digit 0 ∼ 4 and Y = 1 for digit 5 ∼ 9 .
Secondly, we induce noisylabels ˜ Y by randomly ﬂipping the label Y with a probability of 0.2.
Then we induce domain shifts bysampling the color id C spuriously correlated with ˜ Y .
Speciﬁcally, we generate C by ﬂipping ˜ Y withprobability r , which can be viewed as the indicator of different domains.
In training, we randomlysample 5000 data points and set r = 0 .
85 and in testing, we set r = 0 , which induces strong domainshifts between training and testing.
Results are shown in Table 3.4.
Ionosphere Radar Classiﬁcation: Label Shifts Ionosphere Radar Dataset [6] consists ofreturn signals from the ionosphere of a phased array radar system in Google Bay, Labrador.
Theelectromagnetic signals were processed by an auto-correlation function to produce 34 continuousattributes.
The target is to determine whether the return signal indicates speciﬁc physical structuresin the ionosphere (good return) or not (bad return).
Due to the disparity between classes, ERM wasfound to achieve a much lower accuracy on bad returns than good ones [18].
DRO methods areexpected to achieve higher accuracy in the harder class.
Results are shown in Table 3.
Analysis From the results on real-world data, we ﬁnd that in most scenarios, WDRO only slightlyoutperforms ERM, and two f -DRO methods show signiﬁcant promotions to ERM.
Our proposed
GDRO outperforms all baselines signiﬁcantly in all scenarios.
(1) The discrepancy between WDRO( ≈ ERM) and f -DRO ( > ERM) indicates that the extending the distribution support in practice isnot as promising as might be expected, which we think is because creating new data points is nearlyimpossible in real scenarios.
(2) The signiﬁcant discrepancy between GDRO ( (cid:29) ERM) and WDRO( ≈ ERM) shows the superiority of restricting the distribution support, which enables DRO to providerobustness in a much larger uncertainty set without worrying about generating unrealistic samples.(3)
GDRO exhibits signiﬁcant advantages under strong distributional shifts (small minor ratio in
Figure 3 and Figure 4; strong shifts in Colored MNSIT in Table 3), which shows that by incorporatinggeometric properties, our uncertainty set mitigates the over-ﬂexibility problem and our GDRO canresist stronger distributional shifts.
Conclusion
In this work, we take the ﬁrst step to incorporate data geometry information to mitigate the overﬂexibility problem in DRO.
Theoretical analysis and empirical results validate the effectiveness ofour proposed GDRO.
We ﬁnd there are potentials to incorporate more complicated models like GNNsfor a better characterization of the data manifold, which we leave for the future work.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes] All proofs are placed
in Appendix.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e
g., data splits, hyperparameters,
how they
were chosen)?
[Yes] All details are placed in Appendix.
(c) Did you report error bars (e
g., with respect to the random seed after running exper-
iments multiple times)?
[N/A]
The error bar is quite small in our experiments and
therefore we omit it.
(d) Did you include the total amount of compute and the type of resources used (e.g., type
of GPUs, internal cluster, or cloud provider)?
[N/A]
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[Yes]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]