Distributionally Robust Optimization with Data Geometry
Authors: Jiashuo Liu
Jiayun Wu
Bo Li
Peng Cui
Venue: NeurIPS 2022 Accept (2022)
Keywords: Distributional Robustness
Over-flexibility Problem
Data Geometry

Abstract: Distributionally Robust Optimization (DRO) serves as a robust alternative to empirical risk minimization (ERM), which optimizes the worst-case distribution in an uncertainty set typically specified by distance metrics including $f$-divergence and the Wasserstein distance. The metrics defined in the ostensible high dimensional space lead to exceedingly large uncertainty sets, resulting in the underperformance of 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 properties into the design of distance metrics, obtaining our novel Geometric Wasserstein DRO (GDRO). Empowered by Gradient Flow, we derive a generically applicable approximate algorithm for the optimization of GDRO, and provide the bounded error rate of the approximation as well as the convergence rate of our algorithm. We also theoretically characterize the edge cases where certain existing DRO methods are the degeneracy of GDRO. Extensive experiments justify the superiority of our GDRO to existing DRO methods in multiple settings with strong distributional shifts, and confirm that the uncertainty set of GDRO adapts to data geometry.

Distributionally Robust Optimization with Data Geometry
Jiashuo Liu",, Jiayun Wu"", Bo Li² Peng Cuilit
Department of Computer Science & Technology, Tsinghua University, Beijing, China School of Economics and Management, Tsinghua University. Beijing, China hujiashuo77 hete7/@@maaa..ooommm. com,jiayun....... wu, work @ gmail..com liho liho@sem.tsinghua.ece edu.en, cuip crip@tsinnhua.ed..n
Abstract
Distributionally Robust Optimization (DRO) serves robust alternative to em- pirical risk minimization (ERM), which optimizes the worst-case distribution in an uncertainty set typically specified by distance metrics including f-divergence and the Wasserstein distance. The metrics defined the ostensible high dimensional space lead exceedingly large uncertainty sets, resulting in the underperformance most existing DRO methods. It has been well documented that high dimen- sional data approximately resides on low dimensional manifolds. In this work, to further constrain the uncertainty set, we incorporate data geometric properties into the design of distance metrics, obtaining our novel Geometric Wasserstein DRO (GDRO). Empowered by Gradient Flow, we derive a generically applicable approximate algorithm for the optimization of GDRO, and provide the bounded eтror rate of the approximation as well as the convergence rate of our algorithm. We also theoretically characterize the edge cases where certain existing DRO methods are the degeneracy of GDRO. Extensive experiments, justify the superiority of our GDRO existing DRO methods in multiple settings with strong distributional shifts, and confirm that the uncertainty set of GDRO adapts to data geometry.
Introduction
Machine learning algorithms with empirical risk minimization often suffer from poor generalization performance under distributional shifts real applications due the widespread latent heterogeneity, domain shifts, and data selection bias, etc. It is demanded for machine learning algorithms to achieve uniformly good performances against potential distributional shifts, especially in high-stake applications. Towards this goal, distributionally robust optimization (DRO) [27 23, 4 13 [1],
stemming from the literature robust learning, has been proposed and developed in recent years. It optimizes the worst-case distribution within an uncertainty set P(P.,) lying around the training distribution Pr.. When the testing distribution Pte contained in P(P DRO could guarantee the generalization performance on Pte.
In principle, the effectiveness of DRO heavily depends on the rationality of its uncertainty set P(Ptr) which is commonly formulated as a ball surrounding the training distribution endowed with certain distance metric. An ideal uncertainty set should be constituted by realistic distributions that may be encountered in test environments. However, existing DRO methods adopting the Wasserstein distance (i.e. WDRO methods 29] 29 (3) or f-divergence distance (i.e. f-DRO methods 23 tend to
generate over-flexible uncertainty sets that incorporate unrealistic distributions far beyond the ideal
Eqqual Contributions Corresponndig Author
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Data Manifold
Wasserstein Space
2|
Sample AwT
Regioss
Geonetrk Wasserstiaa Set
Sunple
Resssosble
dw
f-Divergeacck-3
Suppert
2
Wasserstein Set
Regiont
Geametric Watserselia
(a) Comparison of metrics on data manifold.
(b) Comparison of uncertainty sets.
Figure I: Toy illustrations. Figure (a) illustrates the transportation paths and corresponding costs of the distance metries. (b) depicts the uncertainty set of WDRO and GDRO compared with ideal uncertainty set in Wasserstein space, where each point denotes distribution.
As such unrealistic distributions must violate the underlying predicting mechanism, they are prone to be the worst-case and attract much optimization energy in the DRO framework, making the learned model deviate from the true predicting mechanism.
uncertainty set 15 13).
Here we argue that the unrealistic distributions mentioned above originate from the distance metrics' inherent ignorance of data geometry, as illustrated in Figure [T] The Euclidean- norm transportation cost measured by the La-Wasserstein metric leads to a straight-line transportation path as shown in Figure T(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, resulting in unrealistic distributions. f-divergence can also be interpreted as a data geometry- -inneerr-indeeendentt measure of the transportation cost confined in the support of P.... Taking x >=-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 l[a). In such case, the noisy samples (e.g. outliers or samples with label noises) are more prone to be the worst-case and thus gather much larger weights than normal samples. The resultant distribution is obviously unrealistic.
To mitigate the problem, it is imperative introduce a new distance metric incorporating data geometry to further constrain the uncertainty set and avoid the undesired cases. As illustrated (a) considering the common assumption that data lie on a low-dimensional manifold we expect the probability density transportation path (the blue dotted line) is restricted within the data manifold (the blue region). In this way, the uncertainty set (i.e. the Geometric Wasserstein Set shown in Figure [(b)) could inherently exclude the distributions beyond the data manifold. Furthermore, it is harder to gather probability weights on isolated noisy samples, which also mitigate the undesired cases in f-DRO.
Figure
12 30
In this work, we propose a novel Geometric Wasserstein DRO (GDRO) method by exploiting the discrete Geometric Wasserstein distance [6] which measures the transportation cost of probability density along the geodesic a metric space. As the Geometric Wasserstein distance does not enjoy analytical expression, we derive an approximate algorithm from the Gradient Flow in the Finsler manifold endowed with Geometric Wasserstein Distance (in section (.2). We further theoretically specify an exponentially vanishing error rate of our approximation as well as a ((1/VT) convergence rate of our algorithm, and characterize the edge cases where GDRO will degenerate to
f-DRO or Wasserstein DRO (in section 3.3 and 3.4p. Comprehensive experiments encompassing various distributional shifts, including cub-population shifts and class difficulty shifts, validate the effectiveness of our proposed GDRO (in section 4 We also observe lower Dirichlet Energy (i.e. higher smoothness) of GDRO's estimated worst-case distribution w.r.t the data manifold compared with existing DRO methods, justifying its adaptability to data geometry.
2 Preliminaries on Distributionally Robust Optimization
Notations. X E denotes the covariates, Y E y denotes the target, fo(-) : + y the predictor parameterized by 0 E O Px) and (x,(X, Y) abbreviated with Per and Pt. respectively represent the joint training distribution and test distribution. The random variable of data points denoted by Z = (X,Y)= Z.

---

Distributionally Robust Optimization (DRO) is formulated as:
arg min sup BP(((f((X), Y)), a5 PEP(P()
(I)
where e is a loss function, P(P.)) = (P Disi(P, P.r)) S }) characterizes the uncertainty set surrounding the training distribution restricted by radius e, and Dist is a distance metric between probability distributions. Most works specify the Dist metric as the f-divergence (23 or the
Wasserstein distance [27 29 26 127.
f f-divergence DRO (abbr: f-DRO) f-divergence is defined as D)(P)1Q) ((d) where f()) is a convex function and f(1) 0. Two typical instances of
f f-divergences are KL- divergence (f(t) t tlg and X2 x2-divergence (f(t) = (t- 1) (23] theoretically demonstrates the equivalence between x" ~DRO and the arriance-regularizes empirical risk minimization (ERM) problem, and [11] derives the optimization algorithm for family of S-DRO. However, as proven in
[5]. f-RRO faces the over-pessing weer-pessimism problem and ends up giving a classifier only fitting the given training distribution, which we attribute to the ignorance of data geometry. As shown in Figure [[a] f-divergence only cares about the probability of each sample (only dP, dQ occur). However, data geometry information is crucial for a reasonable uncertainty set, since it is well-accepted that data lie on a -in-dimensional manifold and adjacent data points have similar degrees of importance. For example, for heterogeneous data, while one hopes focus some sub-population (e.g.. put more weights on a group of data), without data geometric information, the distribution in the f-divergence ball is prone to only focus some isolated samples with higher noises (as shown in Figure (3a)). And in Figure 33b) we find the worst-case distribution of o f-DRO (with Kl-ddivergence) 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 the support of the training distribution, the uncertainty set built with Wasserstein distance allows for the extension of the support 22 29 4] [26 12 4 convert the original problem into a regularized ERM problem, but it is suitable only for limited class of loss functions and transportation cost functions. (29 proposes an approximate optimization method for Wasserstein DRO that could be applied to deep neural networks, which protects models from adversarial attacks. However, the flexibility of the Wasserstein ball also causes an over wver-pessimisti estimation under strong distributional shifts [3] where the created samples too noisy to obtain 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 geometric properties into the uncertainty set. Compared with traditional shape-constrained methods 17 [6] for multivariate extreme event amalysis that use the unimodality to constitute the uncertainty set, our proposed method characterizes the data manifold in data-driven way and incorporates it into the DRO framework intrinsically via the Geometric Wasserstein distance metric, which is also compatible with manifold learning and graph learning methods.
3 Proposed Method
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 and propose the overall objective of GDRO; then we derive an approximate algorithm for optimization utilizing Gradient Flow; finally, some theoretical properties are proved and connections with existing DRO methods are demonstrated.
3.1 Discrete Geometric Wasserstein Distance GWc, and GDRO
We firstly introduce the Discrete Geometric Wasserstein distance, which extends the Senamou-Brentee formulation of the optimal transport problem to metric space. The first step is to define a discrete velocity field and its discrete divergence, which we mainly follow the construction by Chow et al. [6].
Consider a given weighted finite graph Go (V, E, w) with nodes, where V = the vertex set, E the edge set and wr = (wrr) the weight of each edge. A velocity field v = (wn)) (.jev E Rax* on Go is defined to be skew-symmetric matrix on the edge set E such that Big Oyz if (i, E E. The probability set (simplex) (Go) supported on V
n.
1,p: 2 0, for any V} and its interior is defined as (Go) (p.)) E R" 1 Ei Pi
3

---

is denoted by P.(Go). Kij is a predefined "cross- sectional area" typically interpolated with the associated nodes' densities Pi. Pj. The direct approach to take the arithmetic average such that Kay(p) (p. + Pj)/2. However, to ensure the positiveness of p during optimization, we adopt the upwind interpolatione:. (p) = T(ris > 0)p; + U1, I O)p- One could thereafter define the product pe E " called flux function on Go- by pu := (i (P))()EE The divergence of po is dive. (pu) :- -
(EV.())) - W5 (p) p)){ which is vector in R". The divergence vector is supposed to lie in the tangent space of P (Go), summing over all the in-fluxes and out-fluxes along edges of a certain node, with each edge transporting a probability density Wusftighas Ray (p). Now are ready to define Geometric Wasserstein distance in Equation 2
Definition 3.1 (Discrete Geometric Wasserstein Distance (Wco(s)) ()). Given a finite graph Go.
for any pair of distributions p², pp' E PooGo),, define the Geometric Wasserstein Distance:
GWo (p" ,2') = inf [
(os(D)ofztt dp + =. Sr (2) (6,))E
Σ
0,p =
where v E Ruxn denotes the velocity field on Go. is a continuously differentiable curve p(t):
[0, + (G,),, and Kig (p) is a pre-defined interpolation function between Pt and py-
Intuitively is a velocity field continuously transporting masses to convert the density distribution from p" to p² along curve in the Wasserstein space [31]. Squation@meaasun length among all possible plans, which is calculated by integrating a total "kinetic energy" of the velocity field over the transportation process. Compared with the Benamou-Brentee continuous Lz-Wasserstein distance, it ensures that the transportation path stays within the manifold (as the blue dotted line shown in Figure I(a).. and it induces a smoother estimate of the worst-case probability distribution w.r.t the data structure since weights are exchanged just between neighbors.
the shortest (geodesic)
formulation of
Then we present the overall objective function of Distributionally Robust Optimization with Geometric Wasserstein distance (GDRO). Given the training dataset Der {( w));= and its empirical marginal distribution Per +/4 D 8xxz), along with a manifold structure represented by graph Go. intend to obtain a distributionally robust predictor parameterized by 0 such that for certain > 0:
0" arg mill BEE
sup P.OW PSS+
(0.
Σ p(((f((r)).))
B
ps log
We add a minor entropy- reguaarreuulaiizatiinn with a small 33 as proposed in the -alaociig literature 14 to avoid singular cases and ensure the convergence of our optimization in section 3.2 Owing to the Geometric Wasserstein distance, the uncertainty set of GDRO excludes those distributions supported on points beyond the data manifold and the Geometric Wassserstein Ball directional in
Wasserstein space as it stretches along the data structure, as depicted in Figure T(b)
How is Go estimated? To characterize the data manifold, the Go used in GDRO constructed as a k-nearest neighbor (KNN) graph from the training data only, as the KNN graph is shown to have a good approximation of the geodesic distance within local structures on the manifold (7. Note that 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 analytical form of the Geometric Wasserstein distance, we give up providing prescribed amount of robustness
Algorithm 1 Geometric Wasserstein Distributionally Robust Optimization (GDRO)
Input: Training Dataset D., ((x,, y ) learning rate aoo, gradient flow iterations T. entropy term B. manifold representation Go (leameed by kNN algorithm from Dzr).
Initialization: Sample weights initialized as (1/n...., 1/n)²"
Predictor's parameters initialized as ((0)
for i= to Epochs do
1. Simulate gradient flow for T time steps according to Equation 5 to learn an approximate worst-case prohability weight p".
2. g(++1) + g( Vo (E
end for

---

in Equation 3 and propose an alternate optimization algorithm as an approximation. For fixed probability weights p. the parameter could be optimized via gradient descents for R, (0.p) w.r.t 0 in parameter space ④ The inner supremum problem can be approximately solved via gradient ascents for (.. p) w.r.t p in the Geometric Wasserstein space (P.(G)) G.G.). And the cost
measured by GW2 Go (Ptr..) could be approximated with the length of the gradient flow, which is a curve (G.)))
Here we clarify some notations. p [0, + Po(Go) denotes the continuous gradient flow, and the probability weight of the i-th sample at time L is abbreviated as p.(1). The time-discretized gradient flow corresponding with the time step r is denoted as P. [0, T* + Po (Go), and P- (t) is abbreviated as o². For the optimization, we adopt the diserrtiiet definition of Gradient Flow 31 for -R, (0, p) in the Geometric Wasserstein space (9 (Go(, ( as: (with the time step T)
(t+ arg
max R., p) gwa (t).)). pf (GG))
(4)
2T
When 0, the time-discretized gradient flow 7r becomes the continuous one p. Note that Equation describes the Gradient Flow as a steepest ascent curve locally optimizing for a maximal objective within an infinitesimal Geometric Wasserstein ball, and it coincides with the Lagrangian penalty problem of Equation 3 In theorem 3.1 we would prove that Equation 4 finds the exact
solution to local GDRO each time step.
Following Chow er al. (6] the analytical solution to Equation 4 as + 0 could be derived as:
dp:
1i (2) 8
Σ H((E))
Miggae P; log
(5)
dt
where p; denotes the time-dependent probability function of the i-th sample, 4; denotes the loss of the i-th sample and we take an upwind interpolation of 6: K;;(p) = M[1s > 0)p, + T( S 0)p, so that the probability density transferred on an edge equals the density from the origin node associated with the velocity field. The upwind interpolation guarantees that the probability weight stays positive along the Gradient Flow in E Then we discretize equation 5 with Forward Euler Method:
pt a) ri(() + (ddp)()/dl
(6)
where is learning rate. For our algorithm, we control the maximum time step as / < T in
Equation 6 to approximately restrict the radius of the Geometric Wasserstein ball. We prove in theorem 3.2 that for the final time step L = T, the probability weights p(T) learned by Equation 6 CT guarantces a global error rate e -CT from the worst-case risk R, (0, constrained in an e(f)-radius ball where <0)) = GW2 Go (P.r. p(T) and p= = arg supp (R., (0.p) p) gW?, Go (P.r. p) I <(0)). The result is similar to conventions in WDRO [2)]. which gives up providing a prescribed radius of its uncertainty set but turns to an approximation with a intermediate hyperparameter. Pseudo-code of the whole algorithm is shown in Algorithm T The whole derivations are in Appendix.
3.3 Theoretical Properties
In this section, we prove the equivalence between our Gradient-Flow-bee algorithm and a local GDRO problem, and the bound of its global error rate as well as the convergence rate is derived. We first provide the robustness guarantee for the Lagrangian penalty problem in Equation 4
Theorem 3.1 (Local Robustness Guarantees of Lagrangian Penalty Problem). For any T >
0, t > 0 and given 0, denote the solution of Equation 4 as p" (0) peesuppee R, (,,0pp) p) Lw Go (0). (0) = GW3 (((()))p(()) we have
sup (0. gww co (p), (0). p.o (p) (0.).,)))))
sup
R. (0,
PE...
Theorem 3.1 proves that at each time step our Lagrangian penalty problem is equivalent to local GDRO within the --(diuu Geometric Wasserstein ball. It further shows that with + 0 in 4 our gradient flow constantly finds the steepest descent direction. Then we theoretically analyze the global error rate brought by our approximate algorithm.
Equation
Theorem 3.2 (Global Error Rate Bound). Given the model parameter 0, denote the approximate worst-case by gradient descent in Equation 6 after time t as p' (0), and c(0) = gWZ (Prr.7 p' (0)) denotes the distance between approximation p' and the training distribution Per. Then denote
5

---

the real worst-case distribution within the c(f)) discrete Geometric Wasserstein-ball as p (0)),
that is,
(0) arg
Enle pul. - - sup
B p. log
(8)
P.G (Pc-))))((0) i=1
Here we derive the bound w.E.t. the evvor ratio of objective function (0, (abbr. R(p)). For 0 E O there exists C > such that
Error Rate - (((oi)))((())))))))) -R(P.)) <E
(9)
and when - 00, Error Rate -+ The value of C depends on f, 13, n.
Theorem 3.2 theoretically characterizes "how far' approximation p' 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 of our Algorithm
Theorem 3.3 (Convergence Algorithm D. Denote the finctionn for the predictor as:
F(O) ewa, (pc)))
sup
(0,
(10)
which is assumed as L-smooth and (., satisfies Lp smssohhtsssss such that VV (0. V (,, p')ll2 S Lplip - -p'lle 12 e(0) follows the definition in Theorem 3.2 Take a constant AF 2 F(O(O)) ) inf, F(0) and set step size as a == =VDpi(Lie
For Z To where To is
constant, denote the upper bound of llp' - -p" 11/2 as and train the model for K steps, we have:
K
(1
(1 (1+ LDp/K)
1 2 LAr/K 1
2AF
(11)
K
VArK 2LAF
Here we make a common assumption on the smoothness of the objective function as in 29 As K > (oF((())) will achieve a square-root convergence only if y is controlled by the exponentially
vanishing error rate Theorem 3.2 And the accuracy parameter 7 remains a fixed 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 Geometric Wasserstein-bell regularization (set + 2o) and set the graph Go fully-connected graph, and then the solution of GDRO equivalent to the following form of KL-DRO:
min 8E0
sup P.D
pil(fo(zz). v.).
with 2(0) = (xx(F"(0)))
(12)
where p" (0) = arg max, En ((l(fo(z))..)) - B. -
log Pt.
Remark (Connections with WDRO). Since conventional WDRO allows distributions extend training support, our proposed GDRO is intrinsically different from WDRO. Intuitively, for infinite
samples, if the graph Go set to a fully-connected graph with edge weights way = || 2y|12 and is set to 0, our GDRO resembles support- version of WDRO.
4 Experiments
In this section, we investigate the empirical performance our proposed GDRO on different simula- tion and real-world datasets under various kinds of distributional shifts, including sub-population shifts and class difficulty shifts. As for baselines, we compare with empirical risk minimization (ERM), WDRO
[4] 29 and two typical fDROO methods [1], including KL-DRO (f(t) In and
-DRO (f(t) = 1)
1)?).
Implementation Details For all experiments, Go is constructed k-nearest neighbor graph from the training data only at the initialization step. Specifically, we adopt NN-Descent [10] to efficiently estimate the k-nearest neighbor graph for the large-scale dataset Colored MNIST while performing

---

k = 20
k = 100
43
Figure 2: Visualization of learned kNN graph with different l of the regression data, which is projected on the plane spanned by the unit vector of V axis and Os with projection matrix [0.. 01.4
0
01.4 0
an exact search for k-nearest neighbors in the other experiments. We adopt MSE as the empirical loss function 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. Besides, find that the two-stage optimization is enough for good performances, use in our experiments. Note that GDRO compatible with any parameterized models including deep models. The simulation of gradient flow in Equation o implemented by message propagation which scales linearly with sample size and enjoys parallelization by GPU.
mentioned in [19), and we
with DGL. package 32
4.1 Simulation Data
In this subsection, we use simulations to verify that our GDRO could deal with sub- cub-population shifts and to some extent resist the label noises. And also visualize the effects of the kNN algorithm as well as the sensitivity of GDRO to the parameter k.
Table I: Results on the Selection Bias Experiments. We report the root mean square errors.
Simulation regression data without noises
Trancmujor)
Trainiminor)
Parumeter
Bias Ratio
r 1.3
-=-33
-=-15
22
-2.1
3.0
Est Fimor
ERM
0.339
0.876
0.892
0.884
0.864
0.880
0.845
0.858
0.423
WDRO
0.339
0.877
0.894
0.885
0.865
0.882
0.844
0.890
0.424
-DRRO
0.411
0.744
0,757
0.741
0.733
0,742
0.714
0.755
0.367
KL-DRO
0.370
0.713
0.728
0.716
0.708
0.713
0.685
0.724
0.319
GDRO
0.492
0.508
0.489
0.501
0.483
0,486
0.496
0.033
Simulation 2: regression data with label noises
ERM
0.335
0.845
0.885
0.879
0.874
0.884
0.882
0.876
0.422
WDRO
0.335
0.896
0.887
0,880
0.875
0,886
0.884
0.877
0,423
2DDRO
0.375
0.866
0,855
0.856
0,843
0,860
0.854
0.845
0.408
KI-DRO
D.393
0.879
0.868
0.866
0.856
0.876
0.866
0.861
0.391
GDRO
0.537
0.553
0.549
0.539
0.555
0.058
Simulation 3: vary under Sumulation
GDRO (k - 5)
0.493
0.492
0.508
0.489
0.501
0.483
0.486
0.496
0.033
GDRO
0.507
0.502
0.504
0.036
GDRO (A= 100)
0.379
0.673
0,688
0.676
0.670
0.672
0.647
0.683
0.286
1. Regression: Sub-population Shifts via Selection Bias Mechanism
Data Generation The input features X = U, VJ" E I2 are comprised of stable features S RRS noisy features U R4 and the spurious feature V E R:
U,
(0, 21a) RR, U E
03 0.1 +N
(13)
1 ) ER,
~ Laplace(sign(r)) .Y, 1/
(14)
where As R i is the coefficient of the true model. |r| a factor for each sub-population. are stable features with the invariant relationship with Y. U are noisy features such that U 1 Y. And V is the spurious feature whose relationship with unstable and 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 spurious correlation is.
7

---

Simulation Setting In training, we generate 10000 points, where the major group contains 95% data with r 1.9 (i.e. strong positive spurious correlation) and the minor group contains 5% data with r = -1.3 (i.e. weak negative spurious correlation). As shown in Figure 2 the training data the union of two sub-spaces. In testing, we vary r É {-1.55 -1.7, -
2.3. -2.7, -3.0) to simulate stronger negative spurious correlations between V and Y. Notably, the testing data also lie on the same manifold as the training. We use the linear model and calculate the rouu-rean-s errors (RMSE) and the parameter estimation errors Est Error = 1 - 1 of different methods (0* = 00g, .................. 0)"). The results are shown in the Simulation 1 in Table
Simulation Setting 2 Then to test whether GDRO could resist label noises, we randomly sample 20 points and add label noises to them via Y Y + Std(Y) where std(Y) denotes the standard derivation of the marginal distribution of Y. The results are shown in the Simulation 2 in Table ]
And we visualize the learned worst-case distribution of three methods in Figure (3a and 3(b)
Analysis (1) From the results of Simulation / and Simulation 2 in Table [ GDRO outperforms all the baselines in terms of low prediction error on the minor group under different strengths of spurious correlations. (2) From Simulation 2 in Table [ compared with KL-DRO and x²-DRO, GDRO is only 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 GDRO focuses 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 between
Energy. In Figure 3(b) we plot the Dirichlet Energy w.r.t the relative entropy K the learned distribution P and training distribution Per, 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 label noises, since GDRO does not allow extremely high weights on the isolated points. (3) The third sub-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 of the KNN graph Go- we vary & (5, 20, 100} and test the performances of our GDRO under simulation setting 1. We also visualize the KNN graphs in Figure| 2 which show that KNN consistently manages to fit the data manifold well until k 100. And empirical results of Simulation 3 in Table D prove that with k < 100, GDRO performs stably better than the baselines with small and moderate k. except that smaller 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 data manifold. When k increases to an extremely large number h 100, the neighborhood of KNN diffuses and two manifolds start to merge on the graph, in which case GDRO could not distinguish between two aub-populations and its performance degrades as shown in the Table T 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 GDRO perform stably well for a large range of k.
ADOO
GDRO
WDRO
ED10 cDnd with 2.0
cDnD whth 0.1
Major Growp
Mader Croun
Mep Grean
CORD wbth 0.02
Minu Drowe
Miner Greap
Meer Dowe
Ralative Extropy
Modry Aaniiaa weightt
Labee
Holry Label
(a) Visualization of learned weights of fDROO and GDRO.
(b) Weight smoothness.
Figure 3: Explanatory studies of Simulation 2 for the regression data. Figure (a) visualizes the learned worst-case distribution of f-DRO, GDRO, and WDRO on KNN, and the size of each node is proportional to its sample weight or its label noise ratio. Figure (b) plots the Dirichlet Energy w.r.t the relative entropy, which measures the smoothness of learned weights given the same (pr)
2. Classification: Sub-population Shifts with High- digh-diimssional Manifold Data
Data Generation In this setting, data are high nighddimessional but with a low-dimensional structure. The data generation is similar to (25) and is typical classification setting in OOD generalization.
8

---

We introduce the spurious correlation between the label Y = [-1,-1] and the spurious attribute A = (+1,- -). We firstly generate low-dimensional data Xtow = [S, V E 10:00 as:
where A Y. with probability т, with probability 1-
S S~N((,,0))))) V ~~
(15)
Intuitively, r E [0. 1] tunes the proportions of sub-populations and controls the spurious correlation between A and Y. When r > 0.5, the spurious attribute A is positively correlated with Y; and when r < 0.5, the spurious correlation becomes negative, And larger |r 0.5| results in stronger spurious correlation between A and Y. Then convert the low-dimensional data high-dimensional space,
Xtow is multiplied by a column full rank matrix as:
Xhigh (H R 300
(16)
where H E E 30000 is full column rank, and we randomly choose H in each run.
Simulation Setting For both the training and testing data, we set o = 1.0 and o² = 0.3. We use linear models with cross-entropy loss for all methods. In training, we set r = 0.85 (A is positively correlated with Y). In testing, we design two environments with rı = 0.5 (A Y) and r2 = 0.0 (A is negatively correlated with Y) to introduce distributional shifts. Apart from the natural setting without label noises, we also test the performances under label noises. Specifically, we add 4% label noises in the training data by flipping the label Y. We run the experiments 10 times, each time with one random matrix H. We report the mean accuracy in Table 2
Analysis From the results in Table 2 our GDRO outperforms all baselines under the uub-population shifts, and it is not affected much by the label noises, which validates the effectiveness of our GDRO.
4.2 Real-Worle 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 uub-populations shifts and class difficulty shifts, and the others can be found in Appendix. We use MLPs with cross- loss in these experiments.
Sub-population Shifts & Label Noises Following Arjovsky et al. M. Colored MNIST a binary classification task constructed on the MNIST dataset. Firstly, a binary label Y is assigned to each image according to its digit: Y y= 0 for digit 0 ~ 4 and Y = for digit 5 ~ 9. Secondly, we induce noisy labels Y by randomly flipping the label Y with probability of 0.2. Then we sample the color id C' spuriously correlated with Y as C={ +Ý.
Colored MNIST:
probability 1 1-r,
-P,
with probubility
Intuitively, controls the spurious correlation between Y and C. When r < 0.5, C' is positively correlated with Y; and when > 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 points and set r = 0.85 (strong negative spurious correlation between C' and Y) and in testing. we set r 0 (strong positive spurious correlation), inducing strong shifts between training and testing. Results are shown in Table 3
Tonosphere Radar Classification: Class Difficulty Shifts Ionosphere Radar Dataset g consists of return signals from the ionosphere of phased array radar system in Google Bay, Labrador. The electromagnetic signals were processed by an auto-correlation function to produce 34 continuous attributes. The task is to predict whether the return signal indicates specific physical structures in the ionosphere (good return) or not (bad return). However, the prediction difficulty of two classes is quite different, and ERM was found to achieve a much lower accuracy on bad returns than good ones (28) In this experiment, both the training and testing sets consist of samples with balanced label distribution. But due to the disparity of class difficulty, the prediction accuracy of two classes is quite different, 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
From the results on real-world data, we find that all DRO methods (WDRO and fDDROs) show significant promotions to ERM, reflecting the reasonability of our experimental settings. And our proposed GDRO outperforms all baselines significantly when dealing with sub-population shifts and class difficulty shifts, which validates the effectiveness of our GDRO.
Analysis

---

Table 2: Results of the classification simu- lated experiment.
Table 3: Results of Colored MNIST data and lonosphere data.
No Label
Noises
4d
Label Noises
Method
Colored MNIST
lonoxpbere Hard Class Ace
1 0.5
r2
ft =
rg
Train Acc
Test Acc
Easy Class Ace
AUC Score
ERM
0.373
0.153
0.573
0.152
ERM
0.867
0.116
0.952
0.481
0.683
WDRO
0.576
0.159
0.576
0,157
WDRO
1.000
0.335
0.944
0.630
0.774
KL-DRO
0.654
0,340
0.625
0.269
X²DRRO
0.839
0.420
0.976
0.519
0.756
XDDRO
0.734
0.644
0.666
0.554
KLDBO
1,000
0.287
0,984
0.630
0.826
GDRO
0,768
0.767
0.760
0,703
GDRO
0.717
0.696
0.962
0,741
0.883
5
Related Work
Distributionally robust optimization (DRO) directly solves the OOD generalization problem by optimizing the worst-case
in pre-defined uncertainty set, which is often constrained by
moment or support conditions [9] 3]. shape constraints [22 17 16, [1, f-divergence [23] 11] and
Wasserstein distance 226 29 12).
set moment or support conditions for the distributions in the uncertainty set. As for shape constraints, one commonly used is unimodality, and [16] uses the orthounimodality to constitute the uncertainty set for DRO for multivariate extreme event analysis. As for f-divergence, [23] theoretically demonstrates that it is equivalent to the variance penalty, and derives the optimization algorithm from its dual reformulation. Compared with f-divergences which require the support of distributions in the uncertainty set is fixed, the uncertainty set built with Wasserstein distance contains distributions with different support and could provide robustness to unseen data. Despite the capacity of a Wasserstein uncertainty set, the optimization of Wasserstein DRO is quite hard. [26] 12 convert the original DRO problem into regularized ERM problem, but it is suitable only for a limited class of loss functions and transportation cost functions.
(29) proposes an approximate optimization method for Wasserstein DRO and could be applied to deep neural networks, which protects the models from adversarial attacks. Besides, DRO methods have also been used for structured data. [8] studies the DRO problem for data generated by a time-homogeneous, ergodic finite-state Markov chain.
Although DRO methods could guarantee the OOD generalization performances when the testing distribution is included in the uncertainty set, there are works [[5] [3 doubting their real effects in practice. In order to guarantee the OOD generalization ability, in real scenarios, the uncertainty set has to be overwhelmingly large to contain the potential testing distributions. Such overwhelmingly large set makes the learned model make decisions with fairly low confidence, and it also referred to as the over- over-pessiie pessimism problem. To mitigate such problem, [13] proposes to incorporate additional unlabeled data to further constrain the uncertainty set, and (20 learns the transportation cost function for WDRO with the help of multiple environment data.
6 Conclusion
Through this work, we take the first step to incorporate data geometry information to mitigate the over-flexbiility problem in DRO. In this work, we the k-nearest-neighhor graph characterize the data manifold, while our proposed method is compatible with any manifold learning or graph learning methods. And believe that more accurate estimated data structure with advanced manifold learning and graph learning algorithms will further boost the performance of GDRO, which we leave for future work.
Acknowledgements
This work was supported in part by National Key R&D Program of China (No. 2018AAA0102004 No. 20103340 National Natural Science Foundation of China (No. U1936219, 62141607), Beijing Academy of Artificial Intelligence (BAAI). Bo Li's research was supported by the National Natural Science Foundation of China (No.72171131)) the Tsinghua University Initiative Scientific Research Grant (No. 2019THZWJCII Technology and Innovation Major Project of the Ministry of Science and Technology of China under Grants 2020AAA0108400 and 2020AA We would like to thank Yuting Pan, Renzhe Xu, Hao Zou for helpful comments.
10

---

Checklist
1. 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]
2. 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.
3. 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 URL)? [No]
(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]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you eite 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 obtained from people whose data you're you're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]

---

References
[1] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk mini- mization, arXiv preprint arXiv: /90700299 2019.
[2] Mikhail Belkin and Partha Niyogi. Laplacian eigenmaps for dimensionality reduction and data 2003.
representation. Neural computation,
[3] Dimitris Bertsimas, Vishal Gupta, and Nathan Kallus. Data-driven robust optimization. Mathe- matical Programminng 167(():235-299) 2018.
[4] Ruidi Chen and loannis Paschalidis. robust learning approach for regression models based on distributionally robust optimization. Journal of Machine Learning Research, 19(13), 2018.
[5] Xi Chen, Simai He, Bo Jiang, Christopher Thomas Ryan, and Teng Zhang, The discrete moment problem with nonconvex shape constraints. Operations Research,
2021.
[6] Shui-Nee Chow, Wuchen Li, and Haomin Zhou, Entropy dissipation of fokker-planck equations on graphs. arXiv preprint arXiv:1701 170100441,, 2017.
[7] Emma Dann, Neil Henderson, Sarah A Teichmann, Michael D Morgan, and John C Marioni. Differential abundance testing on single-cell data using k-nearest neighbor graphs. Nature Biotechnology, 40(2):245-255 2022.
[8] lonosphere Radar Dataset. https://archive--etes cettaics.uci.edu//ee htt.sscc...cc////dtusss//oopppre
[9] Erick Delage and Yinyu Ye. Distributionally robust optimization under moment uncertainty with application to data-driven problems. Operations research, 58(3):595 612, 2010.
[10] Wei Dong, Moses Charikar, and Kai Li. Efficient k-nearest neighbor graph construction for generic similarity measures. In Sadagopan Srinivasan, Krithi Ramamritham, Arun Kumar, M. P. Ravindra, Elisa Bertino, and Ravi Kumar, editors, Proceedings of the 20th International Conference on World Wide Web, WWW 2011, Hyderabad, India, March 28 April I, 2011, pages 577-586. ACM, 2011.
[11] John C Duchi and Hongseok Namkoong. Learning models with uniform performance via distributionally robust optimization. The Annals of Statistics, 49(3):1378-1406, 2021.
[12] Peyman Mohajerin Esfahani and Daniel Kuhn. Data-driven distributionally robust optimization using the wasserstein metric: performance guarantees and tractable reformulations. Math. Program., (71(1-2):115-111 2018
[13] Charlie Frogner. Sebastian Claici, Edward Chien, and Justin Solomon. Incorporating unlabeled data into distributionally robust learning. arXiv preprint 1911v:191,.07729, 2019.
[14] Jens Hainmueller. Entropy balancing for causal effects: A multivariate reweighting method to produce balanced samples in observational studies. Political analysis, 20(1):25-46, 2012.
[15] Weihua Hu, Gang Niu, Issei Sato, and Masashi Sugiyama. Does distributionally robust super- vised learning give robust classifiers? In International Conference on Machine Learning, pages 2029-2037. PMLR, R, 2018.
[16] Henry Lam, Zhenyuan Liu, and Xinyu Zhang. Orthounimodal distributionally robust optimiza- tion: Representation, computation and multivariate extreme event applications. arXiv preprint arXiv:2111.07894, 2021.
[17] Henry Lam and Clementine Mottet. Tail analysis without parametric models: A worst-case perspective. Operations Research, 15(6):169-11111 2017.
[18] Mengmeng Li, Tobias Sutter, and Daniel Kuhn. Distributionally robust optimization with markovian data. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML. 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 1493-6503. PMLR, 2021.
12

---

[19] Evan Z Liu, Behzad Haghgoo, Annie S Chen, Aditi Raghunathan. Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. In International Conference on Machine Learning, pages 6781-6792. PMI.R, 2021.
[20] Jiashuo Liu, Zheyan Shen, Peng Cui, Linjun Zhou, Kun Kuang, Bo Li, and Yishi Lin. Stable adversarial learning under distributional shifts. In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pages 8662-8670, 2021.
[21] Leland McInnes, John Healy, and James Melville. Umap: Uniform manifold approximation and projection for dimension reduction. arXiv preprini
2018
[22] Clementine Mottet and Henry Lam. On optimization over tail distributions. arXiv preprint arXiv: 1rXiv:1711.00573, 2017.
[23] Hongseok Namkoong and John C Duchi. Variance-based regularization with convex objectives.
Advances neural information processing systems, 30, 2017.
[24] Sam T Roweis and Lawrence K Saul. Nonlinear dimensionality reduction by locally linear embedding, science, 290(5500):2323--3333 2000.
[25] Shiori Sagawa, Aditi Raghunathan, Pang Wei Koh, and Perey Liang. An investigation of why overparameterization exacerbates spurious correlations. In International Conference on Machine Learning, pages 8346-8356. PMLR, 2020.
[26] Soroosh Shafieezadeh-Abeeee Abadeh, Peyman Mohajerin Esfahani, and Daniel Kuhn. Distributionally robust logistic regression. In Corinna Cortes, Neil D. Lawrence, Daniel D. Lee, Masashi Sugiyama, and Roman Garnett, editors, Advances Neural Information Processing Systems 28: Annual Conference Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada, pages 1576-1584, 2015.
[27] Soroosh Shafieezadeh Abadeh, Peyman M Mohajerin Esfahani, and Daniel Kuhn. Distribu- tionally robust logistic regression. Advances in Neural Information Processing Systems, 28, 2015.
[28] Vincent G Sigillito, Simon Wing, Larrie Hutton, and Kile B Baker. Classification of radar returns from the ionosphere using neural networks. Johns Hopkins APL Technical Digest, 10(3):262-266, 1989.
[29] Aman Sinha, Hongseok Namkoong, Riccardo Volpi, and John Duchi. Certifying some dis- tributional robustness with principled adversarial training. arXiv preprint 17Xiv:1710.10571, 2017.
[30] Joshua B Tenenbaum, Vin de Silva, and John C Langford. A global geometric framework for nonlinear dimensionality reduction. science, 290(5500):23192253 2000.
[31] Cédric Villani. Topics in optimal transportation. 58, 2021.
[32] Minjie Wang, Da Zheng, Zihao Ye, Quan Gan, Mufei Li, Xiang Song. Jinjing Zhou, Chao Ma, Lingfan Yu, Yu Gai, Tianjun Xiao, Tong He, George Karypis, Jinyang Li, and Zheng Zhang. Deep graph library: A graph-centric, highly-performant package for graph neural networks. arXiv preprint arXiv: 190.03339 2019.
13


Reviews:
[{"review_id": "eSyFkO2mEPI", "reviewer": "Reviewer_XN8V", "summary": "This paper considers the data geometry in the distributionally robust optimization (DRO) problems and proposes a novel framework called Geometric Wasserstein DRO (GDRO) to achieve their goal. The authors also provide some theoretical analyses, such as approximated optimization error and convergence rate, to theoretically show the strengths of GDRO. Finally, the experimental results show the effectiveness of the proposed approach.", "questions": "My major concerns are listed below:\n- The authors should provide more discussions related to their work, including but not limited to the following concepts:\na)\tDistributionally Robust Optimization with Markovian Data.\nb)\tOrthounimodal Distributionally Robust Optimization.\nc)\tDistributionally robust shape and topology optimization.\n- Can the GDRO apply to large DNNs?\n- In Thm.3.1, $\\ell(\\theta)$ is assumed to be convex. What about the non-convex?\n\n==================================\n\nMy concerns about this work have been well addressed. I am happy to vote for its acceptance.", "limitations": "None\n", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The motivation and contributions are good. This work attacks an overlooked issue in the DRO community and proposes a reasonable method to alleviate this. This would raise much more research attention in this direction.\n- This paper presents some roughly decent theoretical guarantees, which support the effectiveness theoretically.\n\nWeaknesses:\n- This paper lacks comprehensive discussions about the related work. As we all know, DRO has attracted tremendous research interest in the machine learning community, and there are many studies about DRO.\n- The proposed GDRO may has somewhat limitations, since it is not easy to apply to the deep neural networks (DNNs) (Perhaps I am wrong, but at least the authors do not mention it).\n\nAs a whole, despite some limitations, I believe the proposed method is a qualified work and could bring some new insights into the DRO community. Therefore, I tend to accept this paper.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "a9oyQCNiri", "reviewer": "Reviewer_aH5S", "summary": "This work aims to solve the problem that DRO is too “pessimistic” (the uncertainty set is too large) and often leads to poor results in practice. The motivation is that “high dimensional data approximately reside on low dimensional manifolds” (lines 6-7), so this work tries to constrain the uncertainty set on this low dimensional manifold. To do this, this work (i) uses the NN-Descent method to estimate the low dimensional manifold; (ii) formulates the GDRO objective and optimizes it with alternating optimization and proves that it converges. The authors then conduct a series of experiments on synthetic and real datasets and claim that the proposed method is better than existing DRO methods.", "questions": "1. I suggest the authors provide some numerical results to demonstrate how well NN-Descent can estimate the low-dimensional manifold. This is very important, because if the manifold is not well estimated and the target distribution is outside the estimated manifold, then GDRO could completely fail. Moreover, since the authors are using kNN, studying the effect of k is also important. If the method works for some k but not for others, then the authors need to elaborate on how to select a proper k with the training samples alone.\n\n2. Could the authors clarify for which of the experiments is the graph G_0 directly provided to GDRO, and for which of them is the G_0 estimated by NN-Descent?\n\n3. In Figure 3, on the Retiring Adults dataset, GDRO seems to maintain the same high performance regardless of the minor ratio. This is not a good signal. When does the performance of GDRO start to drop? For instance, does GDRO still have a very high performance when the minor ratio is 0.01? What about when the ratio is 0? If GDRO always has such a high performance, then I believe that either G_0 provides too much information to GDRO (for instance, the target distribution could be directly obtained from G_0), or there is a bug in the code.\n", "limitations": "The limitations are not sufficiently addressed.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I really like the high-level idea of this paper. “DRO is too pessimistic” is a well-known open problem in this field, and this work tries to solve this problem by constraining the uncertainty set on a low dimensional manifold, which makes lots of sense. I also think that a lot of credit should be given to the authors for the theory part: The GDRO objective is nicely formulated, and optimizing this objective with alternating optimization makes sense. The convergence result is also very nice.\n\nThe major weaknesses of this work, however, come from the experiment section. Recall that the core motivation of this work is that “high dimensional data approximately reside on low dimensional manifolds” (lines 6-7), whereas there is a huge gap between the motivation and the experiments, which makes the experiments very confusing and unconvincing.\n\nTake the first experiment in Section 4.1 as an example. First of all, the experimental setting is very confusing. I suppose that the task is to infer Y from (S, V). I also couldn’t find the definition of alpha_V, so I suppose that it is used to define V. Thus, in this task the input data is 2-dimensional, and (S, V) seems to also reside on a 2-dimensional manifold, which is the union of a number of 1-dimensional curves (if alpha_V is in [a, b] for some a and b). I don’t think this could be called “high dimensional data approximately residing on low dimensional manifolds”. \n\nThen the authors claim that GDRO is much better than other DRO on this task. I don’t know how the graph G_0 is estimated. I suppose that the authors just simply provide G_0 to the algorithm because NN-Descent is only used “for large-scaled datasets” (lines 109-110). So it seems to me that the reason why GDRO is so good is that G_0 leaks some additional information about the target distribution to it but not to other methods, not because it leverages the fact that the data resides on a low dimensional manifold.\n\nOf course, it is nice that GDRO could utilize this additional leaked information from G_0. The question is: How to get this G_0 in practice? The authors propose to estimate G_0 with NN-Descent, but they don’t demonstrate how well NN-Descent can estimate G_0 on realistic tasks. If G_0 is not well estimated, and the target distribution is outside the estimated manifold, then I imagine that GDRO could completely fail.\n\nMoreover, most of the tasks in the experiments are not really high-dimensional (<50 dimensions), and all tasks seem to follow some simple, unrealistic structures, which make it easier for GDRO to achieve high performances. It is questionable whether these good performances are transferable to real-world applications with realistic distribution shifts.\n\nA valid experimental setting I would suggest the authors try is the following: The input data comes from a low dimensional manifold in a high dimensional space (at least 200 dimensions), but the structure of this manifold is unknown (for instance, introduce randomness into the manifold structure), so GDRO must first estimate G_0 by itself. This setting is closer to the authors’ motivation that “high dimensional data approximately reside on low dimensional manifolds”. Otherwise, it is always questionable whether the performance gain of GDRO comes from the information leakage from the provided G_0, rather than its ability to estimate and utilize the low-dimensional manifold.\n\nIn summary, I really like the high-level idea and the theory part of this paper, but the experiment section does require a lot of improvement. Currently, there is a huge gap between the authors’ motivation and the experiments, making the main conclusion of this paper highly debatable. For this reason, I recommend rejecting this paper for this time, and hope that the authors could resubmit after rewriting the experiment section.\n\n****** Post Rebuttal ******\nThe authors have revised the paper as suggested, so I would like to raise my rating to accept.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "DPIGA9XSRlf", "reviewer": "Reviewer_wQ3n", "summary": "This paper proposed a novel Geometric Wasserstein DRO (GDRO) method by exploiting the discrete Geometric Wasserstein distance. A generically applicable approximate algorithm is derived for model optimization. Extensive experiments on both simulation and real-world datasets demonstrate its effectiveness. ", "questions": "1. Since the manifold is constructed by the training set, is it still applicable for unseen distributions? Data from unseen distributions may fall out of the data manifold. \n2. Does the graph need to be updated at every iteration? If so, it would be time-consuming to estimate the manifold for large-scale datasets. \n", "limitations": "Yes.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Pros:\n1. The proposed method is well motivated and reasonable. This paper studied an important problem of DRO: the uncertainty set is too over-flexible such that it may include implausible worst-case distributions. To address this issue, the authors proposed to use Discrete Geometric Wasserstein distance to construct the uncertainty set, in order to constrain the uncertainty set within the data manifold. The method is somewhat novel and interesting. \n\n2. Both convergence rate and the bounded error rate are provided. And the superiority of the proposed method is also empirically demonstrated through experiments on both simulation and real-world datasets.\n\nCons:\n1. Data from unseen distributions may fall out of the manifold constructed by training data. In this case, simply constraining the uncertainty set may not be helpful for OOD generalization.\n\n2. Training efficiency. The authors use a graph to represent the manifold structure. It may be problematic for large-scale datasets since the graph needs to be estimated at every iteration.\n\n3. In the experiments, the authors only compare with ERM and DRO-based methods. It would be a bonus if some general methods for  OOD generalization can be included. \n  ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "CZRnOA9M12S", "reviewer": "Reviewer_PrzG", "summary": "In this paper, the authors study DRO with data geometry considered in the uncertainty set, making use of the so-called Geometric Wasserstein distance. The authors derive an approximate algorithm for the proposed GDRO and prove its convergence. Numerical experiments are performed to demonstrate the proposed GDRO framework over ERM and other DRO frameworks. ", "questions": "The mathematics of this article is very hard to follow due to abuse of notation. For example, $p$ has been used to represent a lot of different notions. Various claims are also stated without proof or sufficient explanations (e.g., line 151). The use of English language should also be largely improved. The theoretical results also appear to be very trivial. The description of the experiments must be improved—we don’t even know what loss functions are used in the experiments. The exact DRO problems solved in the experiments, especially for the real-world data ones, are completely unstated. The authors also have to rectify the following issues: \n\n- Typesetting or formatting issues (non-exhaustive): \n    - You have to add spaces before all parentheses and citations \n    - Add punctuations at the end of display style equations\n    - Format tables according to the instructions in the NeurIPS paper template—no vertical lines! \n    - Use italics instead of bold to emphasize words\n\n- Typos or writing issues (non-exhaustive):\n    - Line 42: $\\mathcal{X}_2$—do you mean $\\chi^2$? \n    - Line 60: $O(1/\\sqrt{T})$\n    - Line 109: large-scale~~d~~\n    - Line 128: ~~cross section~~ cross-sectional\n    - Line 129: ~~algorithmic~~ arithmetic\n    - Line 137: Brenier\n    - Line 146: $\\delta(x_i)$ vs $\\delta(x_i, y_i)$\n    - Line 159: ~~alternate~~ alternating\n    - Line 161: ascent~~s~~\n    - Line 193: $\\mathcal{R}_n$ vs $R_n$\n    - Line 200: $\\mathcal{GW}$ vs $GW$\n    - Line 204: what is $\\mathcal{R}(p)$? Not previously defined. \n    - Table 1: why adding underscores between “mean” and “error”? Use a dot for abbreviations instead of an underscore. \n", "limitations": "Relevant discussion of limitations might appear in the paper but is hard to find. \n", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths: \n- An interesting framework proposed using the Geometric Wasserstein distance\n- Extensive experiments\n\nWeaknesses: \n- The whole paper is quite hard to follow due to: \n    - the lack of self-containedness—some rather non-standard notions are not well-defined or well-explained, or very hard to understand\n    - unsatisfactory English writing (wrong grammar, choice of words, etc.)\n    - inaccuracy in mathematics exposition, the lack of mathematical rigor and lots of abuse of notation\n- Related works are not discussed in detail, e.g., DRO\n- A lot of typos, typesetting and formatting issues \n- The settings and descriptions of the experiments are unclear\n\n\n=============================\nPost-rebuttal: Thanks the authors for their effort. The revised version has addressed my concerns and I have raised my score. Although it is allowed in this NeurIPS submission cycle, I think the authors should have provided sufficient experimental details in their initial submission. A lot of new content have been added in the revised version of the paper, including supplementary material, say for the experimental details. I feel it is quite unfair to assess the merits of the paper according to this updated version, comparing to other paper submissions. It appears to me that the authors was submitting unfinished work on the submission deadline and take advantage of the rebuttal phase. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept