Distributionally Robust Optimization with Data Geometry
Anonymous Author(s)
AffiliationAddress 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 specified by distance metrics including f -divergence andthe Wasserstein distance.
The metrics defined 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 dimensional data approximately resides on low dimensional manifolds.
In this work,to further constrain the uncertainty set, we incorporate data geometric propertiesinto the design of distance metrics, obtaining our novel Geometric Wasserstein DRO (GDRO).
Empowered by Gradient Flow, we derive a generically applicableapproximate algorithm for the optimization of GDRO, and provide the boundederror rate of the approximation as well as the convergence rate of our algorithm.
Wealso theoretically characterize the edge cases where certain existing DRO methodsare the degeneracy of GDRO.
Extensive experiments justify the superiority of our GDRO to existing DRO methods in multiple settings with strong distributionalshifts, and confirm 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) [20, 16, 22, 3, 10, 8],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, existing DRO methods adopting the Wassersteindistance (i.e. WDRO methods [20, 22, 3, 10]) or f -divergence distance (i.e. f -DRO methods [16, 8])tend to generate over-flexible uncertainty sets that incorporate unrealistic distributions far beyond theideal uncertainty set [12, 10].
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 leads to a straight-line transportation path as shownin 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 confined 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 resultant distribution is obviously unrealistic.
To mitigate the problem, it is imperative 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[17, 23, 2], we expect the probability density transportation path (the blue dotted line) is restrictedwithin the data manifold (the blue region).
In this way, the 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 from the 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 where GDRO will degenerate tof -DRO or Wasserstein DRO (in section 3.3 and 3.4).
Comprehensive experiments encompassingvarious distributional shifts, including sub-population shifts and class difficulty shifts, validate theeffectiveness 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 manifold comparedwith 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 ℓ is a loss function, P ( P tr )
= { P : Dist ( P, P tr ) ≤ ϵ } characterizes the uncertainty setsurrounding the training distribution restricted by a radius ϵ , and Dist is a distance metric between probability distributions.
Most works specify the Dist metric as the f -divergence [16, 8] or the Wasserstein distance [20, 22, 3, 19, 9].f -divergence DRO ( abbr.
f -DRO) f -divergence is defined as D f ( P ∥ 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 ).
[16] theoretically demonstratesthe equivalence between χ 2 -DRO and the variance-regularized empirical risk minimization (ERM)problem, and [8] derives the optimization algorithm for a family of f -DRO.
However, as proven in[12], f -DRO faces the over-pessimism problem and ends up giving a classifier only fitting the giventraining distribution, which we attribute to the ignorance of data geometry.
As shown in Figure 1(a),f -divergence only cares about the probability of each sample (only dP, dQ occur).
However, datageometry information is crucial for a reasonable uncertainty set, since it is well-accepted that datalie on a low-dimensional manifold and adjacent data points have similar degrees of importance.
Forexample, for heterogeneous data, while one hopes to focus on some sub-population (e.g., put moreweights on a group of data), without data geometric information, the distribution in the f -divergenceball is prone to only focus on some isolated samples with higher noises (as shown in Figure 3(a)).
And in Figure 3(b), we find the worst-case distribution of f -DRO (with KL-divergence) is not smooth(with larger Dirichlet Energy) w.r.t.
the data manifold.
Wasserstein DRO ( abbr.
WDRO)
Compared with the f -divergence ball that does not extend thesupport of the training distribution, the uncertainty set built with Wasserstein distance allows for theextension of the support [20, 22, 3].
[19, 9, 3] convert the original problem into a regularized ERMproblem, but it is suitable only for a limited class of loss functions and transportation cost functions.[22] proposes an approximate optimization method for Wasserstein DRO that could be applied todeep neural networks, which protects models from adversarial attacks.
However, the flexibility of the
Wasserstein ball also causes an over-pessimistic estimation under strong distributional shifts [10],where the created samples are too noisy to obtain a confident model.
As demonstrated in Figure 3(a), WDRO adds much more noises to the data and thus hurts the generalization performances in practice. Therefore, to mitigate the over-pessimism problem of DRO, we propose to incorporate the geometricproperties into the uncertainty set.
Compared with traditional shape-constrained methods [14, 13]for multivariate extreme event analysis that use the unimodality to constitute the uncertainty set, ourproposed method characterizes the data manifold in a data-driven way and incorporates it into the DRO framework intrinsically via the Geometric Wasserstein distance metric, which is also compatiblewith manifold learning and graph learning methods.
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 this work, we propose Distributionally Robust Optimization with Geometric Wasserstein distance(GDRO).
In the following of this section, we first introduce the Geometric Wasserstein distance andpropose the overall objective of GDRO; then we derive an approximate algorithm for optimizationutilizing Gradient Flow; finally, some theoretical properties are proved and connections with existing DRO methods are demonstrated.
3.1 Discrete Geometric Wasserstein Distance GW G 0 and GDRO
We firstly introduce the Discrete Geometric Wasserstein distance, which extends the Benamou-Brenierformulation of the optimal transport problem to a metric space.
The first step is to define a discretevelocity field and its discrete divergence, which we mainly follow the construction by Chow et al.
For a given weighted finite graph G 0 = ( V, E, w ) with n nodes, where V = { 1 , 2 , . . . , n } isthe vertex set, E is the edge set and w = ( w ij ) i,j ∈ V is the weight of each edge. A velocityfield v = ( v ij ) i,j ∈ V ∈ R n × n on G 0 is defined to be a skew-symmetric matrix on the edge set E such that v ij = − v ji if ( i, j ) ∈ E .
The probability set (simplex) P ( G 0 ) supported on Vis defined as P ( G 0 ) =
{ ( p i ) n i =1 ∈ R n | (cid:80) n i =1 p i = 1 , p i ≥ 0 , for any i ∈ V } and its interioris denoted by P o ( G 0 ) . κ ij is a predefined "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 adoptthe upwind interpolation : κ ij ( p ) = I ( v ij > 0) p j + I ( v ij ≤ 0) p i .
One could thereafter define theproduct pv ∈ R n × n , called flux function on G 0 , by pv := ( v ij κ ij ( p )) ( i,j ) ∈ E . The divergence of pv is div G 0 ( pv ) :
= − ( (cid:80) [Equation]
V :( i,j ) [Equation] E √ w ij v ij κ ij ( p ))
ni
=1 which is a vector in R n .
The divergence vectoris supposed to lie in the tangent space of P
o ( G 0 ) , summing over all the in-fluxes and out-fluxesalong edges of a certain node, with each edge transporting a probability density √ w ij v ij κ ij ( p ) .
Nowwe are ready to define Geometric Wasserstein distance in Equation 2.
Definition 3.1 (Discrete Geometric Wasserstein Distance GW G 0 ( · , · ) [4]) .
Given a finite graph G 0 ,for any pair of distributions p 0 , p 1 ∈ P
o
( G 0 ) , define the Geometric Wasserstein Distance:
[Equation]
where v ∈ R n × n denotes the velocity field on G 0 , p is a continuously differentiable curve p ( t ) : [0 , 1] → P o ( G 0 ) , and κ ij ( p ) is a pre-defined interpolation function between p i and p j .
Intuitively v is a velocity field continuously transporting masses to convert the density distributionfrom p 0 to p 1 along a curve in the Wasserstein space [24]. Equation 2 measures the shortest (geodesic)length among all possible plans, which is calculated by integrating a total "kinetic energy" of thevelocity field over the transportation process. Compared with the Benamou-Brenier formulation ofcontinuous L 2 -Wasserstein distance, it ensures that the transportation path stays within the manifold(as the blue dotted line shown in Figure 1(a)), and it induces a smoother estimate of the worst-caseprobability distribution w.r.t the data structure since weights are exchanged just between neighbors.
Then we present the overall objective function of Distributionally Robust Optimization with Geometric 143 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 ϵ > 0 :
[Equation]
We add a minor entropy-regularization with a small β as proposed in the entropy-balancing literature[11] to avoid singular cases and ensure the convergence of our optimization in section 3.2.
Owingto the Geometric Wasserstein distance, the uncertainty set of GDRO excludes those distributionssupported on points beyond the data manifold and the Geometric Wassserstein Ball is directional in Wasserstein space as it stretches along the data structure, as depicted in Figure 1(b). How is G 0 estimated?
To characterize the data manifold, the G 0 used in GDRO is constructed as ak-nearest neighbor (kNN) graph from the training data only , as the kNN graph is shown to have agood approximation of the geodesic distance within local structures on the manifold [15, 5]. Notethat our GDRO is compatible with any manifold learning and graph learning methods .
3.2 Optimization
In this subsection, we derive the optimization algorithm for GDRO.
Due to the lack of an analyticalform of the Geometric Wasserstein distance, we give up providing a prescribed amount ϵ of robustnessin Equation 3 and propose an alternate optimization algorithm as an approximation.
For fixedprobability weights p , the parameter θ could be optimized via gradient descents for R n ( θ, p ) w.r.t.θ in parameter space Θ .
The inner supremum problem can be approximately solved via gradientascents for R n ( θ, p ) w.r.t. p in the Geometric Wasserstein space ( P o ( G 0 ) , GW G 0 ) . And the costmeasured by GW 2 G 0 ( ˆ P tr , · ) could be approximated with the length of the gradient flow, which is acurve in ( P o ( G 0 ) , GW G 0 ) .
Input: Training Dataset D tr = { ( x i , y i ) } ni =1 , learning rate α θ , gradient flow iterations T , entropy term β , manifold representation G 0 (learned by kNN algorithm from D tr ).
Initialization: Sample weights initialized initialized as θ (0) . for i = 0 to Epochs do
θ ( i +1)
← θ ( i ) − α θ ∇ θ ( (cid:80)
i p Ti ℓ i ( θ ))
end
Here we clarify some notations.
p : [0 , T ] (cid:55)→ P o
( G 0 ) denotes the continuous gradient flow , and 166 the probability weight of the i -th sample at time t is abbreviated as p i ( t ) .
The time-discretizedgradient flow 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 definition of Gradient Flow 169 [24] for −R n ( θ, p ) in the Geometric Wasserstein space ( P o ( G 0 ) , GW G 0 ) as: (with the time step τ )
[Equation]
When τ → 0 , the time-discretized gradient flow ˆ p τ becomes the continuous one p .
Note that
Equation 4 describes the Gradient Flow as a steepest ascent curve locally optimizing for a maximalobjective within an infinitesimal Geometric Wasserstein ball, and it coincides with the Lagrangianpenalty problem of Equation 3.
In theorem 3.1 we would prove that Equation 4 finds the exactsolution to a local GDRO at each time step.
Following Chow et al.
[4], the analytical solution to Equation 4 as τ → 0 could be derived as:
[Equation]
where p i denotes the time-dependent probability function of the i -th sample, l i denote the loss of the 178 i -th sample and we take an upwind interpolation of κ : κ ij ( p ) =
I ( v ij > 0) p j
+ I ( v ij ≤ 0) p i , so thatthe probability density transferred on an edge equals the density from the origin node associated with 180 the velocity field.
The upwind interpolation guarantees that the probability weight p stays positivealong the Gradient Flow in Equation 5.
Then we discretize equation 5 with Forward Euler Method:
[Equation]
where α is a learning rate.
For our algorithm, we control the maximum time step as t ≤ T in
Equation 6 to approximately restrict the radius of the Geometric Wasserstein ball.
We prove intheorem 3.2 that for the final time step t =
T , the probability weights p ( T ) learned by Equation 6guarantees a global error rate e − CT from the worst-case risk R n ( θ, p ∗ ) constrained in an ϵ ( θ )
-radiusball where ϵ ( θ ) =
GW 2 G 0 ( ˆ P tr , p ( T )) and p ∗ =
arg sup p {R n ( θ, p ) : GW 2 G 0 ( ˆ P tr , p ) ≤ ϵ ( θ ) } .
Theresult is similar to conventions in WDRO [22], 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
In this section, we prove the equivalence between our Gradient-Flow-based algorithm and a local GDRO problem, and the bound of its global error rate as well as the convergence rate is derived.
Wefirst provide the robustness guarantee for the Lagrangian penalty problem in Equation 4.
Theorem 3.1 (Local Robustness Guarantees of Lagrangian Penalty Problem) .
For any τ > 195 0 , t > 0 and given θ , denote the solution of Equation 4 as p ∗ ( θ ) =
arg sup p ∈ P ( G 0 )
R n ( θ, p ) −
τ GW 2 G 0 (ˆ p tτ ( θ ) , p ) .
Let ϵ τ ( θ ) =
GW 2 G 0 (ˆ p tτ ( θ ) , p ∗ ( θ )) , we have
[Equation]
Theorem 3.1 proves that at each time step our Lagrangian penalty problem is equivalent to a local GDRO within the ϵ τ ( θ ) -radius
Geometric Wasserstein ball.
It further shows that with τ → 0 in
Equation 4, our gradient flow constantly finds 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 6 after time t as p t ( θ ) , and ϵ ( θ )
= 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 ϵ ( θ ) -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 207 θ ∈ Θ , there exists C > 0 such that Error
[Equation]
[Equation]
and when t → ∞ , Error Rate → 0 .
The value of C depends on ℓ, β, 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 n ( θ, p ) satisfies L p -smoothness such that ∥∇ p R n ( θ, p ) −∇ p R n ( θ, p ′ )
∥ 2 ≤ L p ∥
p − p ′
∥ 2 .
ϵ ( θ ) follows the definition 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 ∥ p t − p ∗ ∥ 22 as γ and train the model for K steps, we have:
[Equation]
Here we make a common assumption on the smoothness of the objective function as in [22].
As K →∞ , ∇ θ F ( θ ( k ) ) will achieve a square-root convergence only if γ is controlled by the exponentiallyvanishing error rate in Theorem 3.2.
And the accuracy parameter γ remains a fixed effect onoptimization 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 224 Geometric Wasserstein-ball regularization (set ϵ → ∞ ) and set the graph G 0 to a fully-connectedgraph, 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 ℓ ( 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 infinitesamples, if the graph G 0 is set to a fully-connected graph with edge weights w ij = ∥
z i − z j ∥ 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 different simulation and real-world datasets under various kinds of distributional shifts, including sub-populationshifts and class difficulty shifts .
As for baselines, we compare with empirical risk minimization(ERM), WDRO [3, 22] and two typical f -DRO methods [8], including KL-DRO ( f ( t ) =
t ln t ) andχ 2 -DRO ( f ( t ) =
( t − 1) 2 ).
Implementation Details For all experiments, G 0 is constructed as a k-nearest neighbor graph fromthe training data only at the initialization step. Specifically, we adopt NN-Descent [7] to efficientlyestimate the k-nearest neighbor graph for the large-scale dataset Colored MNIST while performingan exact search for k-nearest neighbors in the other experiments. We adopt MSE as the empirical lossfunction for regression tasks and cross-entropy for classification tasks. We use MLPs for the Colored
MNIST and Ionosphere datasets, and linear models in the other experiments. Note that GDRO iscompatible with any parameterized models including deep models . The simulation of gradient flowin Equation 6 is implemented by message propagation with DGL package [25], which scales linearlywith sample size and enjoys parallelization by GPU. 4.1 Simulation Data
In this subsection, we use simulations to verify that our GDRO could deal with sub-population shiftsand to some extent resist the label noises.
And we also visualize the effects of the kNN algorithm aswell as the sensitivity of GDRO to the parameter k . 0 1 ,0 1 ,1 θ TS 0 1 , (cid:21) .
[Table] [Table]
KL-DRO
[Table]
GDRO
[Table]
[Table] [Table] under Simulation GDRO ( k = 5 )
[-[ 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 . 
-][Table] [Table] [Table] [Table]
[Table]
Data Generation The input features X
= [ S, U, V ] T ∈ R 10 are comprised of stable features S ∈ R 5 , noisy features U ∈ R 4 and the spurious feature V ∈ R :
[Equation]
V ∼ Laplace ( sign ( r ) · Y,
[Equation]
(14) where θ S ∈ R 5 is the coefficient of the true model. | r | > 1 is a factor for each sub-population. Sare stable features with the invariant relationship with Y .
U are noisy features such that U ⊥ Y . And V is the spurious feature whose relationship with Y is unstable and is controlled by the factor r .
Intuitively, sign ( r ) controls whether the spurious correlation between V and Y is positive or negative.
And | r | controls the strength of the spurious correlation: the larger | r | is, the stronger the spuriouscorrelation is.
Simulation Setting 1 In training, we generate 10000 points, where the major group contains 95%data with r = 1 .
(i.e. strong positive spurious correlation) and the minor group contains 5% datawith r = − 1 . 3 (i.e. weak negative spurious correlation).
As shown in Figure 2, the training datais the union of two sub-spaces. In testing, we vary r ∈ {− 1 . 5 , − 1 . 7 , − 1 .
9 , − 2 .
3 , − 2 .
7 , − 3 . 0 } tosimulate stronger negative spurious correlations between V and Y . Notably, the testing data also lieon the same manifold as the training.
We use the linear model and calculate the root-mean-squareerrors (RMSE) and the parameter estimation errors Est Error =
∥ ˆ θ − θ ∗ ∥ 2 of different methods( θ ∗ =
[ θ S , 0 , . . .
, 0] T ). The results are shown in the Simulation 1 in Table 1.
Simulation Setting 2 Then to test whether GDRO could resist label noises, we randomly sample20 points and add label noises to them via ˜ Y
=
Y + Std ( Y ) where std ( Y ) denotes the standardderivation of the marginal distribution of Y .
The results are shown in the Simulation 2 in Table 1.
And we visualize the learned worst-case distribution of three methods in Figure 3(a) and 3(b). Analysis (1) From the results of Simulation 1 and Simulation 2 in Table 1, GDRO outperforms allthe baselines in terms of low prediction error on the minor group under different strengths of spuriouscorrelations.
(2) From Simulation 2 in Table 1, compared with KL-DRO and χ 2 -DRO, GDRO isonly slightly affected by the label noises.
Also, from Figure 3(a), compared with GDRO, KL-DRO puts much heavier weights on the noisy points (red points of f -DRO are much larger). And GDROfocuses more on the minor group (blue points), which results in their different performances under Simulation 2 .
Further, to investigate this phenomenon, we quantify the smoothness via Dirichlet
Energy.
In Figure 3(b), we plot the Dirichlet Energy w.r.t the relative entropy KL ( ˆ P ∥
ˆ P tr ) betweenthe learned distribution ˆ P and training distribution ˆ P tr , which proves that the learned weights of GDRO are much smoother w.r.t. the data manifold.
And this property helps GDRO to resist the labelnoises, since GDRO does not allow extremely high weights on the isolated points. (3) The thirdsub-figure in Figure 3(a) verifies our analysis on WDRO that it introduces much more label noises(red points).
Discussion on kNN To test whether GDRO is sensitive to the parameter k of the kNN graph G 0 ,we vary k ∈ { 5 , 20 , 100 } and test the performances of our GDRO under simulation setting 1 . Wealso visualize the kNN graphs in Figure 2, which show that kNN consistently manages to fit thedata manifold well until k = 100 . And empirical results of Simulation 3 in Table 1 prove that withk < 100 , GDRO performs stably better than the baselines with small and moderate k , except thatsmaller k leads to slower convergence since sparse graphs restrain the flow of probability weights.
Still, we present an extreme failure case where KNN achieves a poor approximation of the datamanifold. When k increases to an extremely large number as k = 100 , the neighborhood of kNNdiffuses and two manifolds start to merge on the graph, in which case GDRO could not distinguishbetween two sub-populations and its performance degrades as shown in the Table 1. Actually, in Theorem 3.4 of this paper, we have proved that with an infinitely large k , GDRO could be reduced to
KL-DRO, which completely ignores data geometry. Still, we have to clarify that kNN and GDROperform stably well for a large range of k . [Figure] [Figure]
Data Generation In this setting, data are high-dimensional but with a low-dimensional structure. The data generation is similar to [18] and is a typical classification setting in OOD generalization.
We introduce the spurious correlation between the label Y =
{ +1 , − 1 } and the spurious attribute
A = { +1 , − 1 } .
We firstly generate low-dimensional data X low =
[ S, V ] T ∈ R 10 as: S ∼ N ( Y 1 , σ 2 s I 5 ) , V ∼ N ( A 1 , σ 2 v I 5 ) , where A = (cid:26) Y , with probability r, − Y , with probability 1 − r.
Intuitively, r ∈ [0 , 1] tunes the proportions of sub-populations and controls the spurious correlation 303 between A and Y .
When r > 0 . 5 , the spurious attribute A is positively correlated with Y ; and whenr < 0 . 5 , the spurious correlation becomes negative. And larger | r − 0 .
5 | results in stronger spuriouscorrelation between A and Y . Then to convert the low-dimensional data to high-dimensional space, 306 X low is multiplied by a column full rank matrix H as: [Equation] where H ∈ R 300 × 10 is full column rank, and we randomly choose H in each run.
Simulation Setting For both the training and testing data, we set σ 2 s = 1 . 0 and σ 2 v = 0 .
We uselinear models with cross-entropy loss for all methods.
{+In training, we set r = 0 . ( A is positivelycorrelated with Y ). In testing, we design two environments with r 1 = 0 . 5 ( A ⊥ Y ) and r 2 = 0 . 
+}
( A is negatively correlated with Y ) to introduce distributional shifts.
Apart from the natural settingwithout label noises, we also test the performances under label noises. Specifically, we add 4% label [Table] [Table]
noises in the training data by flipping the label Y .
We run the experiments 10 times, each time withone random matrix H .
The results over 10 runs are shown in Table 2.
Analysis From the results in [-three tasks. HIV-1: Sub-population Shifts & Label Imbalance HIV-
-]Table 2, our GDRO outperforms all baselines under the sub-populationshifts, and it is not affected much by the label noises, which validates the effectiveness of our GDRO.
4.2 Real-World Data
We evaluate our method on four real-world datasets.
Due to space limits, we place two of them here,with various kinds of distributions, including sub-populations shifts and class difficulty shifts, and theothers can be found in Appendix. We use MLPs with cross-entropy loss in these experiments.
Colored MNIST: Sub-population Shifts & Label Noises Following Arjovsky et al.
[1], Colored MNIST is a binary classification task constructed on the MNIST dataset.
Firstly, a binary label Yis assigned to each image according to its digit: Y = 0 for digit 0 ∼ 4 and Y = 1 for digit 5 ∼ 9 .
Secondly, we induce noisy labels ˜ Y by randomly flipping the label Y with a probability of 0.2.
Thenwe sample the color id C spuriously correlated with ˜ Y as C = (cid:26) +˜ Y , with probability 1 − r, − ˜ Y , with probability r. Intuitively, r controls the spurious correlation between Y and C . When r < 0 .
5 , C is positivelycorrelated with Y ; and when r > 0 . 5 , the spurious correlation becomes negative. And | r − 0 . 5 |controls the strength of the spurious correlation.
In training , we randomly sample 5000 data pointsand set r = 0 .
( strong negative spurious correlation between C and Y ) and in testing , we set r = {+
 ( strong positive spurious correlation), inducing+} strong shifts between training and testing.
Resultsare shown in Table 3.
Ionosphere Radar Classification: Class Difficulty Shifts Ionosphere Radar Dataset [6] consistsof return 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 task is to predict whether the return signal indicates specific physical structures inthe ionosphere (good return) or not (bad return).
However, the prediction difficulty of two classes isquite different, and ERM was found to achieve a much lower accuracy on bad returns than good ones[21].
In this experiment, both the training and testing sets consist of samples with balanced labeldistribution. But due to the disparity of class difficulty, the prediction accuracy of two classes is quitedifferent, while DRO methods are expected to achieve similar prediction accuracy for both classes. Therefore, in testing , we report the testing accuracy for the easy class and the hard class respectively,as well as the AUC score of the testing set.
Results are shown in Table 3.
Analysis From the results on real-world data, we find that all DRO methods (WDRO and f -DROs)show significant promotions to ERM, reflecting the reasonability of our experimental settings.
Andour proposed GDRO outperforms all baselines significantly when dealing with sub-population shiftsand class difficulty shifts, which validates the effectiveness of our GDRO.
Conclusion
Through this work, we take the first step to incorporate data geometry information to mitigate theover-flexibility problem in DRO.
In this work, we use the k-nearest-neighbor graph to characterize thedata manifold, while our proposed method is compatible with any manifold learning or graph learningmethods. And we believe that a more accurate estimated data structure with advanced manifoldlearning and graph learning algorithms will further boost the performance of GDRO, which we leavefor future work.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect 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 theywere 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., typeof 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 identifiable
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]