An Effective Manifold-based Optimization Method for Distributionally Robust Classification
Authors: Jiawei Huang
Hu Ding
Venue: ICLR 2025 Poster (2025)
Keywords: robustness
optimization
representation learning
Primary area: optimization
TL;DR: We aim to approximate (part of) the data manifold's tangent space through representation learning and use it to improve the performance of distributionally robust optimization. Theoretical analysis and empirical results confirm its effectiveness.

Abstract: How to promote the robustness of existing deep learning models is a  challenging problem for many practical classification tasks. Recently, Distributionally Robust Optimization (DRO) methods have shown promising potential to tackle this problem. These methods aim to construct reliable models by minimizing the worst-case risk within a local region (called ''uncertainty set'') around the empirical data distribution. However, conventional DRO methods tend to be overly pessimistic, leading to certain discrepancy between the real data distribution and the uncertainty set, which can degrade the classification performance. To address this issue, we propose a manifold-based DRO method that takes the geometric structure of training data  into account for constructing the uncertainty set. Specifically, our method employs a carefully designed ''game'' that integrates contrastive learning with Jacobian regularization to capture the manifold structure, enabling us to solve DRO problems constrained by the data manifold. By utilizing a novel idea for approximating geodesic distance on manifolds, we also provide the theoretical guarantees for its robustness. Moreover, our proposed method is easy to implement in practice. We conduct a set of experiments on several popular benchmark datasets, where the results demonstrate our advantages in terms of accuracy and robustness.

Published as a conference paper at ICLR 2025
AN EFFECTIVE MANIFOLD-BASED OPTIMIZATION METHOD FOR DISTRIBUTIONALLY ROBUST CLASSIFI-
CATION
Jiawei Huang1.2, Hu Ding'*
1, School of Computer Science and Technology, University of Science and Technology of China 2. Department of Computer Science, City University of Hong Kong hjw0330 @mail.mall.......lll
en, hudingeuste. edu.
ABSTRACT
How to promote the robustmess of existing deep learning models is a challeng- ing problem for many practical classification tasks. Recently, Distributionally Robust Optimization (DRO) methods have shown promising potential to tackle this problem. These methods aim to construct reliable models by minimizing the worst-case risk within local region (called "uncertainty set") around the em- pirical data distribution. However, conventional DRO methods tend to be overly pessimistic, leading to certain discrepancy between the real data distribution and the uncertainty set, which can degrade the classification performance. To address this issue, we propose manifold-based DRO method that takes the geometric structure of training data into account for constructing the uncertainty set. Specif- ically, our method employs a carefully designed "game" that integrates contrastive learning with Jacobian regularization to capture the manifold structure, enabling to solve DRO problems constrained by the data manifold. By utilizing a novel idea for approximating geodesic distance on manifolds, we also provide the the- oretical guarantees for robustness. Moreover, our proposed method easy to implement in practice. We conduct set of experiments several popular bench- mark datasets, where the results demonstrate our advantages in terms of accuracy and robustness.
INTRODUCTION
Neural network-based image classification methods have achieved significant success (He et al., 2016). but the data shift issue is still major obstacle that seriously affects their practical perfor- mance (Wiles et al. 2022 Fang al. 2020 Taori al., 2020 Data distribution shift refers the situation that the training and testing data distributions have certain degree of discrepancy, which lead decrease in the generalization of the model TTng et al. 2021 Zhou al. 2021 Lu let al. 2020 Zhang et al., 2022). Real-world machine learning tasks often encounter such issues due to various reasons, such as the changes of data sources (Koh et al.. 2021 Lin et al. 2021 or the
process for data collection (Dragoi et al. 2022 Gareía et al.) 2015).
If we have sufficient prior knowledge to assume that, the real data distribution falls within an area surrounding the training data (we denote this area as the "uncertainty set"), ) it could be possible to enhance the performance by considering all possible data distributions within such an uncertainty set. Namely, the goal is to promote the worst-case performance in this uncertainty set. This idea is called "distributionally robust optimization (DRO)" (Rahimian & Mchrotra, 2019). which has been used in many machine learning models to address distribution shift (Sagawa et al. 2020: Staib 2019 Samuel & Chechik, 202T). In particular, the DRO methods can also be imple-
& Jegelka mented neural networks through the techniques like gradient flow (Wang al. 2022). iterative optimization (Sinha et al.) 2018), and surrogate loss (Samuel & Chechik][021
The uncertainty set of DRO is typically defined the collection of distributions that lie within bounded distance from the empirical distribution of the training data. Common choices for mea-
"Corresponding author.

---

Published as a conference paper at ICLR 2025
suring this distance include f-divergences (Duchi & Namkoong 2021 and the Wasserstein dis- tance (Gao et al.) 2024; Li et al. 2019a). The Wasserstein distance is a popular metric for compar- ing data distributions, which effectively captures pointwise distance information between samples, making it well-suited for the tasks like pattern matching, style transfer, and handling distribution shift <Li et al. 2019b Ling & Okada 2007). A formal definition of Wasserstein distance is provided
in Section 3 The DRO methods relying on Wasserstein distance are called "Wasserstein-bass"
DRO (WDRO)" (Mohajerin Esfahani al. 2018).
DRO offers significant advantages in terms of generalization and robustness across various do- mains (Gao et al., 2024, Huang et al., 2022 Wang et al.) 2022). However, a major challenge in implementing DRO lies in the design of an appropriate uncertainty set. An ideal uncertainty set should effectively capture the underlying distributional uncertainty while can be easily integrated into the formulation for optimization. Consider WDRO as an example, where we denote the train- ing data by Pe.. The conventional WDRO approaches often define the following uncertainty set
u(Prr: 8) (Q P(RR) (IR)) P..) 80,
where P(IRd) is the set of all distributions in IR., W(:-) indicates the Wasserstein distance function and 5 > 0 is a pre-specified range for uncertainty. This formulation implies that the uncertainty set u(P....) 5) includes all distributions whose Wasserstein distance from P.r is at most 6. However, we are often confronted with dilemma of setting 8. On the one hand, 5 should be set large enough to ensure robustness. On the other hand, large ō can also lead reduced accuracy on real data, as the uncertainty set (Puz, 5) could encompass unrealistic distributions that deviate significantly from the true generating process.
Our intuition for addressing the above issue comes from the fact that realistic (e.g., images) typically locate nearby a now-dimensional manifold (Roweis & Saul 2000, Tenen- baum et al.) 2000, Yang et al. (2044). In addition, we can utilize the powerful representational capabilities of neural networks to capture the information from the underlying manifold (Chung let al. [ Cohen et al. 2020). Under these observations, it is reasonable to assume that the data is supported on a manifold M, and we consider to construct new uncertainty set as follows,
data
(go = {Q E P(M))GW(Q, P.,) S 8)},
(2)
where (W(:-) represents the Geodesic Wasserstein distance (formally defined in Section 3). and P(M) is the set of distributions supported on the manifold M. This uncertainty set Ugu [Pr Pt.. 8)
defined in 22 enjoys two important benefits compared with (dPt.,,) defined in Eq. II. First, it restricts the distribution to locate nearby the data manifold, preventing the DRO algorithm from being bothered by irrelevant points in the ambient space Rd Please see an illustrative case in Fig- ure (aab. Second, we consider the geodesic distance on the manifold, which offers a more precise measure of the semantic distance between samples for classification (Criminisi et al. 2008 Wang ct al. 2017). To illustrate this, consider a toy example shown in Figure (bb): the two samples (images) "A" and "B" lying on the manifold M are not semantically similar, although they have shorter pairwise Euclidean distance (compared with their geodesic distance).
However, solving the WDRO problem with the uncertainty set 2y (Piz. d) presents significant chal- lenges. The first challenge is due to the lack of neat representation for the manifold in neural net- works; second, it is hard to add explicit constraints to the objective function to satisfy the uncertainty set in WDRO. In this paper, we aim to develop an effective and -os-too-impleennt approach to tackle these challenges, and our contributions are summarized below:
1. A novel WDRO model with manifold guided constraint. Our proposed model can effee- tively identify the directions of semantic variation in the data manifold, where the key idea relies on a carefully designed game between Contrastive Learning (CL) (Chen et al., (020) and Jacobian regularization in the training process. Intuitively, CL is used to encourage the data to be updated along semantic-variant tangent directions of the manifold, while the Jacobian regularization is for suppressing the updating direction orthogonal the mani- fold. Our training process equipped with such a game can help us to efficiently realize the WDRO within the uncertainty set defined in 2).
2. An efficient algorithm for solving the WDRO model. We aware of several WDRO algorithms for neural networks, such as (Sinha et al. 2018 Bui et al.] (022), but they are unable to be directly extended the manifold constraints, to our best knowledge. To
2

---

Published as a conference paper at ICLR 2025
Points out of maniiold Points menifold
(b)
Figure 1: (a) Real images typically reside on or nearby a manifold in the aih-dimensioonl space, whereas points far away from the manifold could be meaningless, such random noise (as repre- sented by the orange point). (b) 3D visualization for a part of the MNIST dataset (Deng 2012) by t-SNE (Van der Maaten & Hinton 2008). The black curve and dashed line respectively represent the geodesic and Euclidean distances between two images A (digit "1") and B (digit "5"). Though their geodesic distance is large, their Euclidean distance relatively short. Obviously, the geodesic distance can reveal their semantic difference more accurately.
efficiently capture the manifold information, utilize the principal singular vectors for approximating the desired tangent directions of the manifold. Another critical idea in our algorithm utilizing the accumulated optimization trajectory for estimating the geodesic distance. Thanks to the strong concavity of the DRO's dual form, we can guarantee the robustness with this estimation theory. Our algorithm is easy to implement, and the experimental results illustrate its effectiveness for classification tasks on several benchmark datasets.
1.1 RELATIONS TO OTHER DRO METHODS.
The works of (Staib & Jegelka 2017: Sinha et al. 2018 Bui et al. 2022) are closely related to ours, as they also use the dual form of WDRO to search for worst-case perturbations. However, these approaches cannot be directly extended to manifold constraints. While they are effective against adversarial attacks, their overly pessimistic nature can decrease performance on clean data. Another approach is called "Group DRO" (Sagawa et al. 2020), which aims optimize performance across different demographic subgroups, addressing hidden biases and ensuring fairness. The method of f- divergence DRO (Namkoong & Duchi 2016) Duchi & Namkoong 2021), considers the uncertainty set within a certain /ddivergence from the nominal distribution. Several studies tried to incorpo- rate geometric information into the uncertainty set. Liu a constructed graph as the geodesic distance to guide distribution density transportation. Qiao & Peng (202)) calculated group centrality based on the affinity between data groups, serving as good nominal distribution for group DRO. A full introduction of related works is placed in Appendix A.2 A.2
(2022) used the shortest path over
PRELIMINARIES
Notations. "P, represent the input set of n labeled data items (x1,91). (Xn: 1n). For each iis S n, x; belongs to the input space X' T and yi belongs the lahel space y {1, ... C). with C' being the number of classes. In this paper. interchangeably the i-th sample point in Pa.. The Euclidean and Frobenius norms are denoted by || and || respectively. We assume Pt. lies collection of distributions supported on the manifold M. For any point x E M, (x)() denotes the tangent space of M x, which comprises all the tangent vectors x.
p; and (x,, y/ to denote
low-dimensional manifold M in Rd P(M) denotes the
Our neural network is denoted as I h g, where g(-) : X + is the encoder, h(-) - y the classification head and & is the feature space. The network I is parameterized by vector
3

---

Published as a conference paper at ICLR 2025
8 E G where θ is the parameter space. We denote the empirical loss
C(0, e(0,p.),
where E(0, Pt) represents the induced loss of the prediction T p;). We are aware of several existing works that focus on the learned manifolds in "feature space" (LEE et al. 2022 Ma et al. 2023 But it is worth noting that the manifold studied in our current work is about the "data space". " To emphasize this distinction, refer to the manifolds where and g(x) reside the data manifold and the representation manifold, respectively.
Roughly speaking, a manifold is space where each local area around a point is homeomorphic to lower dimensional Euclidean space. We refer readers to the recent textbook (Boumal 2023 for the formal definition of manifold. Given two points on a manifold, their geodesic is the curve between them on the manifold that locally minimizes the path length, and the length of the shortest curve is called the geodesic distance, which is denoted as d-(-.-). We also need the following two assump- tions for M, which are both commonly adopted for optimization problems on manifold (Boumal 2023
Assumption M a complete manifold embedded in an Euclidean space.
Due to the space limit, we leave more explanations on the rationale of Assumption Iin Section A.1 Before stating the second assumption, we introduce two important definitions.
Definition (Exponential map (Zhang & Sra, 2016) Given a point p E M and vector E (P) let 7 : [) + M be the geodesic trajectory starting from satisfying 7(0) = p and y'(0) = with the uniform velocity ||v|]. We use q := Exp,(v) to denote the endpoint 7(1), which called the exponential map at (p, v).
The following Definition 2 generalizes the concepts of u-strongly convexity and f-smoothness from Euclidean space to manifold.
Definition (Geodesically j-strongly convex and f-smooth) We say a function f f -M+RI is geodesically p-strongly convex if for any given x E M.v E (x), E 1]. the inequality (Expxx (tv)) 2 f(x) + + f(x), v) + [ holds. Furthermore, we say geodesically 8-smooth if for any given E M, (x)x), E [0, 1]. the inequality f (Expxx S f(x) +
((VM)(x) + 12 || holds.
In this definition, "P.M (f(x))" represents the Riemannian gradient of f at x, tangent vector in (x4(x) indicating the steepest ascent direction in function value. Under Assumption 1 it equals the orthogonal projection of the "classical" gradient to the tangent spaces (x))) (Boumal 2023). i.e. V-M = Projx (Vf(x)), where the Proj (-) is the orthogonal projection operator that projects a vector the tangent space (x) and (ff(x) is the gradient in Euclidean space Rd
Assumption 2 For any E O, the loss function E(f,x) x) is geodesically 3-smooth w.r.t. any M.
3 OUR FORMULATION: MANIFOLD-BASHE WDRO
We name our model using the uncertainty set (2) as Manifold-based WDRO (MWDRO), which relies on new concept called "geodesic Wasserstein distance".. Given two distributions P and Q. the standard Wasserstein distance measures the optimal transportation cost between them (Peyré et al.) 2019 Villani 2009). Following that, we define the geodesic Wasserstein distance, where we use the geodesic distance to measure the pairwise distance between points on the manifold M:
1/2
J inf GW(Q, P) = (TEI(( JMX. d2 g) (q. (a.
(4)
where the I1, P) represents the collection of all possible joint distributions combining Q and P as their marginal distributions. GW(Q, P) indicates the minimum cost required to transport one distribution into another under the geodesic distance dg.
4

---

Published as a conference paper at ICLR 2025
As mentioned before, our objective is to achieve the distributionally robust optimization over the un- certainty set Ugw (Pr.. 5), as shown in Eq. To realize this goal, we optimize E θ by minimizing the following MWDRO loss to replace the vanilla empirical loss ((0,Ptr)
(D Por) su (Pi..)
{ (((c@(((),)))))
(5)
The above formulation is intractable because needs consider continuous distributions within the uncertainty set under Wasserstein distance; as discussed in (Gao & Kleywegt (203)) such a WDRO model involving continuous distributions an infinite-dimensionall optimization problem. Although a number of theoretical works have been proposed for developing tractable formulations for DRO problems (Esfahani & Kuhn 2013 Shafieezadeh Abadeh al., 2015 Mohajerin Estahani et al. 2018 Postek et al. [016]] Lee & Mehrotra| 2015 Hanasusanto & Kuhn 2018), these methods are not applicable to our case due to the manifold constraint. To solve the problem (5. we adopt the strongly duality property proposed in (Gao & Kleywegt 2023 Theorem 1) to obtain the dual form (as [roposition[]eelee below), and then design our approach with several new ideas relying on this dual form. Note that their original article considers the general metric space rather than manifold. Because the geodesic distance also forms a metric on the manifold, we can derive the strong dual reformulation for Eq.55) (see Section
Proposition 1 For given 9 and the data distribution shift threshold 5 0, the "worst-case" loss in (5) can be reformulated as
(b, min {vō² + Ep,,,a(0,p;. v20
(6)
where (,(O,pr, == sup [((0.g) 9) -0a0--va2(00
gEM
Remark 1 (g.16) is the dual form of 55. that is, the optimal value of Eq. 16 is also the optimal value of Eq. (5) benefit of such reformulation is that we can ignore the complicated uncertainty set Ugu 5). 6). Instead, we only add surrogate loss E,,,, (f,) Pt, v) to the Eq which yields a more succinct formulation for optimizing the problem. However, there are still several significant challenges remaining unsolved, e.g. the feasible domain is confined to the manifold M. We will elaborate on the details for solving this new formulation in Section 4
4 OUR APPROACH FOR SOLVING MWDRO
For ease of understanding, we present an overview our approach Section 4.1 showing how to incorporate the (approximate) manifold constraint into the training process. Next, we offer a detailed analysis on the senifold-guided game in Section 4..2 explaining why this game can effectively encodes the manifold's tangent into model. Finally, we outline our MWDRO algorithm in Section 4.3 together with the theoretical quality guarantees.
4.1 THE OVERVIEW OF OUR APPROACH
We adopt the dual formulation from Proposition as the objective function in our model. However, a major challenge for solving the objective function arises from the surrogate loss in Eq. as it is difficult to find the supremum within manifold (which is non-convex region in high dimensional space). To solve this optimization problem on the manifold, we require the Riemannian gradient, which the orthogonal projection of the "classical" gradient onto the tangent spaces (Boumal, 2023), i.e., (x) = Proix (VJ(x) as shown in our preliminaries. The central problem, there- fore, reduces to recovering the tangent space (x) for any data point x E M. However, obtaining the tangent space on the data manifold is not easy job, the manifold lacks closed -fored repre- sentation. Our idea for addressing this issue comes from the fact that neural networks usually have powerful representation abilities, especially for encoding manifold information (Bronstein et al. 2021
Hu et al. 2023). This motivates us to ask the following question:
Can neural networks also be used to extract the tangent space of data manifold?
High-level idea. We should emphasize that the learned representation for manifold from neural networks is not sufficient for extracting tangent space, demonstrated experimentally in Figure 11
5

---

Published as a conference paper at ICLR 2025
(Appendix G A key part of approach is based a Contrastive Learning (CL) loss (Chen et al. (020) and a Jacobian regularization term. Roughly speaking, a CL loss is used for amplitying the differences between samples with distinct semantics in the feature space; on the other hand, a Jacobian regularization term tries to suppress the differences in the feature space for preserving the manifold constraint (if do not have such suppression, the updating direction could seriously deviate from the manifold). The combination of these two components enables the encoder g(-) to remain sensitive variations along the manifold while insensitive to other directions (a fine- ene-grained analysis is provided in Section 4.2) As result, the Jacobian matrix (x) aligns its principal singular vectors with (part of) the manifold's tan- gent space. This alignment implicitly encodes the data manifold's structure and facilitates effective manifold optimization for MWDRO, as detailed Section 4.3
cumifold-guides game for minimizing
Formally, we introduce the loss for our manifold-guide game:
L (0) LINONNCE + X 1I
where the term LInfONCE is the objective function of CL (that will be explained latter) and the term Jg [x(x)= ;= 0xg(8.x] is the Jacobian matrix of the encoder g(-) w.r.t. the sample x. The LInENNE and the Jacobian regularization terms are balanced by the coefficient X > O.
dx
For completeness, we provide more details on the popular self-supervised method CL. For each x; Pur we denote its two randomly generated copies as x/ x² via some random semantic semantic-preserving augmentation operation (such as cropping rotating the image). Let Ptr {
x², and Pir = {ar xx" Given any S # < n, we say the couple (ar/. ar.") is a "positive pair", and the couples (ark. mr)) and (ark, xx") are "negative pairs". The main idea of CL. is align represen- tations of positive pairs closer, while diversify representations of negative pairs (Tian et al. 2020
Wang & Isola, 2020; Xue et al. (0022). The Information Noise- Contr--Crttrtivvee Estimation (InfoNCE)
loss is a commonly used formulation in CL models (Chen et al. 2020
He et al. 2020):
LIAONCE Eigjn] [<:()
exp (sim (g(a²) g (x.) /T) UPE "" exp (sim (g (x() g(x))
where lal (xx) llg
T > 0 is scalar temperature hyperparameter, and sim(u, v) = u T //llu|/llul is the cosine simi- larity. Intuitively, the InfoNCE loss tends to bring g (x/) and g (z!") to be closer, and meanwhile repulse g (x() and g (x).
4.2 FINE-GRAINED ANALYSIS ON THE MANIFOLD-GUIDED GAME
In this section, we explain why our proposed method in Section 4.1 can effectively encode the langent information of the data manifold into our model. More specifically, what is the ad of combining the CL and Jacobian regularization term in our training method, as proposed in Eq.
To begin our analysis, we need to define some notations. For each data sample x E M,, we define the region containing all the augmentations of x (such as cropping or rotating in CL) on M as the Aemantic-Invarian region "Msr We also denote M : R M the Out-of-Manifold region. Consider the singular value decomposition (SVD) of Jg(x)
Uxx)E(x)V (x), where the rank of Jg Jz(x) is r O, the diagonal matrix Exx) diag ( (x), (x)),,, and the right singular vectors of (x) contains (x). (x)} (these r vectors form an orthonormal basis). Given a threshold To O, we define the subspace spanned by the principal singular vectors of the Jacobian as:
Tappr (x, TO) =
Σ
a.v((x) a, E [R], where = (i ( (x 2 TO. 1i.1iis).
(10)
ier
In [q. we truncate the small singular values less than To-
Why the game helps to extract tangent information? Due to the loss defined in Eq. the train- ing process of CL encourages the extracted features of æ! xx" E <M)r(x) to be close, while keeping the features of negative pairs distinct. Simultaneously. tion suppresses the feature variations for perturbations of across all directions (including both the directions tangent and normal to the manifold) (Ross & Doshi- Velez 2018). Therefore, we can re- gard these two objectives two players in a game: the CI. loss preserves feature variations along
is known that the Jacobian regulariza-

---

Published as a conference paper at ICLR 2025
the tangents apart from the Cemantic-Invariant region Ms)(x), to defend against the sup- pression from the Jacobian regularization. In other words, the namifold-guides game between CL and Jacobian regularization encourages the space spanned by the Jacobian's 's principal singular vec- tors (truncated by T)), i.e., (xppr((; To). to approximately align with directions within the tangent space that induce semantic changes. Therefore, T Tappr(X To) serves as an effective approximation of subset of the tangent T (x), specifically capturing directions of significant semantic variation within (x)... As we will see later in Section 4.3 T Taappr (x, TO) can be utilized approximate the computation of the Riemannian gradient.
Experimental verification. In the above analysis, we assume that Tappr (x, To) is a subspace of the tangent (x). We explain the rationale of this assumption here. According to the property of manifold, a tangent plane characterizes the vicinity of a point on the manifold by a local linear approximation. Similarly, the Jacobian matrix (g(x), which contains the partial derivatives of the features w.r.t. the input data, provides linear approximation that captures how the feature g(-) responds to input variations. Under our proposed manifold-guided game, Tappr (x, To) represents subspace of the tangent (x)() corresponding to semantic changes. Based on the low-dimensional manifold assumption, the dimension of T Tappe (x, To) should be significantly smaller than d, indi- cating a low-rank structure of the Jacobian. In Figure 2 we plot the cumulative percentage of the singular values for the Jacobian matrix of the models trained on CIFAR-10/100 We can that even a small number of top singular vectors can approximate the Jacobian Jg(x) quite well. For instance. in CIFAR-10 the cumulative percentage of the top five singular vectors already take over 95% of the sum of all singular values. This concentration of singular values enables us to compute T Tappe (x, To) with greater accuracy and efficiency (see Appendix
To see the benefit of using our designed game, we compare it with other models in Figure 2 Our model's cumulative percentage curve shows the steepest increase the early stage. This suggests that we have effectively suppressed the encoder g(-)'s 's sensitivity to data variations across most di- rections. We can reasonably assume that the space spanned by the remaining primary singular vec- tors aligns with directions of semantic variation within the data manifold's 's tangent. This alignment facilitates the solving of MWDRO, as detailed in Section 4.3
Cumulaativ Percentage
CIFAR-10
Cumulative Percentage
CIFAR-100
1.0
1.0
0.8
0.8-
0.6
occ+CL00=0.1)
=ac+
0.4
ac+CL(h-0.01)
Jace ac+CL00g-..0
ac
0.2
DropOute12
0.2
0.0
0.0
60
80
100
120
20
40
80 100 120
Number of Singular Values
Number of Singular Values
Figure 2: Cumulative percentage of variance of the Jacobian for different models trained on CIFAR- 10 and CIFAR-100 using ResNet- 18, respectively. The cumulative percentage of variance for the (())
first A singular values is calculated as Eist (x) The red line represents our model as described in (8. The blue line represents the Auto-Encoder model with Jacobian regularization (Jac+AE), and the green line represents the CL. model with Dropout and 1.2 regularizer (Dropout+I.2).
Remark 2 We also highlight several additional benefits of our proposed manifold- guided game from a representation learning perspective. By analyzing the characteristics of the encoder g()), we show that it is possible to further enhance our approach for capturing a smoother manifold. Due to the space limit, please refer to Appendix D for the detailed discussion.
4.3 THE ALGORITIM VIA THE MANIFOLD-GUIDER GAME
In this section, provide our algorithm to solve the MWDRO problem by minimizing (q. based on our analysis in Section 4.2) We minimize the MWDRO loss (q, with the constrant from
the manifold-guided game Eq. 8), keep the manifold's tangents being encoded into our model
7

---

Published as a conference paper at ICLR 2025
effectively in the training process, i.e.,
ming Ln, Ptr), s.t. (0) c0
(11)
where co > 0 is a pre-specified threshold. We convert the problem 11 into an unconstrained
optimization problem by Lagrange duality (Boyd et al. 2004) with X2 > 0 controlling the magnitude of regularization:
ming (D, Pir) + (2)
(12)
Sketch of Algorithm 1 Our MWDRO approach contains two stages. In the first stage, we train the encoder g(-) by minimizing the sanifold-guider loss in Eq. The goal is to ensure g(+) can yield sufficiently good approximate tangent T Tappr (x, To) for each data sample which is later used for Riemannian optimization to solve the surrogate loss. In the second stage, we optimize Eq. (2) by stochastic gradient descent algorithm. A key part of Algorithm iis to compute the "worst-case' distribution in each iteration, which is implemented by Algorithm2 Roughly speaking, Algorithm 2 employs the Riemannian gradient ascent to approach the worst-case distribution. It is worth noting heavily relies on the tangent plane provided by the encoder
that the performance of Algorithm 2 g()), and this is also the reason that add the constraint " "L (0) co" in problem 111
Algorithm Distributional Robust in Data Manifold M
Input: training set Per {(xi, yi) )= constant step size sequence {az > 03/20
Output: the final parameter 0N+11
1: Train g(-) to minimize L (0) in (8. 1 The first stage. Training for No epochs.
*/
2: 1 ...., do
3: Take a sample from Ptr uniformly at random, which is denoted as Piy. and compute new point p., + Evolve (.g()) w): / {y/ ............. returned by the Evolve function
(Algorithm 2 is an approximation of the "worst case" distribution. $/
4:
65+1 f- O ar Vo (E(0), p..) X2 (fa (pi,) + A (p(,)((( Objective Eq. 12) */
5:
end for/* The second stage. */
To deal with the objective function D. (0, Put) = min {vō² + Ep, (0, p.. v)} in Eq. 12 we do not explicitly set the radius 8 for the uncertainty set Ug (Puz, in Eq. Instead, we select a V > O to solve the empirical surrogate loss Ep., Ie, Pc, vi] rather than prescribing the robust range d; this approach, initially proposed by (Sinaa et al. (2018), has been shown to ensure certified distributional robustness. However, we cannot directly apply their method due to the lack of algorithmic tools for dealing with the manifold constraint. Thanks to our constructed approximate tangent T "aappe(x (x, To) in Section 4.2 we can approximate the "worst case" distribution by performing gradient ascent on the manifold. This process is detailed in Algorithm 2 In the algorithm, we fix V in Eq.(7), transforming it into geodesically strongly concave maximization problem. We iteratively approach the optimal "q" through Riemannian gradient ascent on M: generating a sequence 0g0h 2, v..tt..... ... .g*t1 " ..... Let ügt. denote the unit vector representing the initial direction of the geodesic from point g to point g on the manifold. Following optimization theory on manifold (Boumal 2023). we update with a step size a
follows:
++1 Exp,. (aV-M (e(0, va (g')))))
= -xpp, (aProij (VE(O, (0) 2owd ( g)
(13)
Approximation for the geodesic distance. A challenge to solve 13. is how to estimate the geodesic distance d2 "","))))))))),"')) We propose an effective approximation technique for the geodesic dis- tance (q" q²). Specifically, we approximate dg(d", q²) using the accumulated optimization trajec- tory length pt == Σ(= ||a. grad'll in the Step 4 of Algorithm 2 Consequently, we can efficiently update the worst-case distribution as the Step and of Algorithm [ Under mild smoothness as- sumptions, which typically hold for neural networks, we establish the quality of this approximation in Lemma [
Lemma 1 Under Assumption Mand we fix constant v 3 to set up Algorithm 2 (recall B is the smoothness parameter in Assumption D Let q" denote the optimal solution. Then, Algorithm 2 converges linearly; and for any t > 0, the following inequality holds: pt' S xt < vk(VR.
V - Tr (q ".") where 2+8 v-B
8

---

Published as a conference paper at ICLR 2025
Algorithm 2 Evolve(E.g(),.)
Input: A sample y, a constant and the encoder g(-).
Output: y²+1
for = 1, 2, .., until convergence do
1. initialize the "potential" for q: pt" 0:/* Using "potential" for approximating the geodesic
distance. */
2. grad' Proj, (V,1(0, q¹))
grad"); ): ) Conduct the gradient ascent on the manifold. */ 4. pt pt' + ||a grad"|12:
3. g Exp,t gt (a
end for
Lemma Timplies that we can utilize the accumulated step size to approximate the geodesic distance. This approximation plays key role in Theorem [] which establishes the robustness guarantee under our geodesic distance approximation. The proofs are deferred to Appendix
Theorem Suppose Assumption and Tand@hold We select an v > 3 in the surrogate loss (Eq. 7 and define as in Lemma L Using the accumulated step size approximate the geodesic dis- tance, let f be the optimal solution of the dual formulation under this approximation. Then Ep, (, Pi, v) d" ming E.,,, (0,,,(0,p0 Pi,, v). where e" x² (2 V
Remark As an example, if we set 2 173, the result of Theorem 7 indicates that E,,,0 v) S 5.1 ming Ep,, e, P.. It is also worth noting that a smaller 3 implies larger range for setting the parameter WDRO methods on the smoothness of the loss function has also been discussed in previous works. For instance, Sinha et (2018) showed that B is rypically small for neural networks. In Section we further explore the value of 3 and reveal an interesting result: Jacobian regularization is also beneficial for smaller B.
require > 3 in Theorem 7 The dependence of
Algorithm 2 solves the surrogate loss using Riemannian gradient ascent. As introduced in Section
and Section 4.1 the Riemannian gradient can be formulated as: V-Me(B, x) = Projx (V (x)x) where Projx denotes the projection of the Euclidean gradient onto the tangent space TM (x). How- ever, this exact projection is not directly computable due to the lack of analytical formula for the tangent T (x). To overcome this, as shown Section 4.2 we approximate the tangent information using T a'appe (x, To) by the manifold-guider guided game. Specifically, We approximate the Riemannian gra-
dient as V.M ((8, x) := Proj (VE(6, x)), where the (roj)ppr(() the orthogonal projection op- erator that projects the Euclidean gradient (x((0, x) to the plane Tpppe (x, To). In our problem, the loss function E(f, x) is composition of functions involving the encoder g(-). This structure enables us to provide theoretical guarantees regarding the quality of the approximation Projay (VE(0, x))- Informally, under the assumption that T appe (x, TD) (x)(x) (we analyzed the rationale behind this
assumption in Section 4.20. the error of our approximation
bounded <Me(B. x) M x)
by O(Tu) (due to space limit, defer the formal statement Theorem in Appendix
We replace the exponential map operation (x),,()) in Step 3 of Algorithm with a retraction operation, i.e., q=+1 = ' a grad', which is widely used in manifold-based algorithms (Bécigneul & Ganeal (2019). Further implementation de-
Other implementation details for the Algorithm 2]
tails, including efficient strategies for applying Jacobian regularization and computing the subspace T Tappe (x, To), are discussed in Appendix
5 EXPERIMENTS
We conduct series of experiments across several scenarios evaluate proposed method. Specifically, we test it on three typical distribution shift tasks: dealing with noisy data, attacked data and imbalanced data. All experiments are implemented with PyTorch single NVIDIA RTX 6000 Ada. In our experimental setup, we set the hyperparameters \₁ 0.01 and X = 1 in Eq.18) and 12) respectively to configure the algorithm, unless otherwise specified.

---

Published as a conference paper at ICLR 2025
CIFAR-10
CIFAR-100
0.9
MIDDR
MWWDR
0.8
DropOut+12
DropOut+12
Accuracy
WDRO
WDRO
0.6
DRD
KL-DRO
RS
RS
0.4
0.4
0.3
0.3
0.2
0.1
0.0
0.2
0.4
0.6
0.8
0.0
0.1
0.2
0.3
0.4
Cnoise
Onoise
(b)
Tiny-lmageNet 200
BuS
MWDRO
MWDRO
0.7
DropOut+L2
Cccuracy
0.6
WDRO
Accuracy
WDRD
KL-DRO
0.7
0.4
RS
RS
0.3
0.6
0.1
0.5
0.0
0.6
Choise
Onaise
(c)
(d)
Figure 3: Robustness evaluation under Gaussian white noise perturbations. The x-axis represents the variance of the Gaussian noise, ranging from 0 (the input data has been normalized). The results indicate that MWDRO enhances robustness across four datasets compared to baselines.
Basic Evaluation on Noisy Data. We test our model MWDRO against natural image corruptions by adding Gaussian white noise to images. We use CIFAR-10/100 (Krizhevsky et al. 2012), Tiny- ImageNet-200 (Krizhevsky et al. 2017), and a medical imaging dataset BuS (Mo et al. (023) for image classification. Our model is compared with DropOut+L2 (He al.) 2016), WDRO (Sinhal ct al. 2018), KL-DRO (Levy et al. 2020), and Randomized Smoothing (Cohen et al. 2019). Fig-
demonstrates that MWDRO method significantly enhances robustness compared other methods.
Further Evaluation on Distributional Shift Scenarios. We conduct additional evaluations on two typical distributional shift scenarios, detailed in Appendix F Specifically, we conduct assess- ments of our model's resilience against adversarial attacks using the CIFAR-10, CIFAR-100 100, and Tiny-ImageNet-200 datasets. Additionally, we evaluate the effectiveness of our proposed method on long-tailed benchmark datasets test performance under naturally imbalanced data distributions.
compare the running time with the standard WDRO (which does not consider the data manifold) and other adversarial training algorithms. While our algorithm requires more time, achieve better performance accuracy. We conduct other abla- tion studies in Appendix F.6 analyze the effectiveness of the combination of the manifold-guided game and the MWDRO method.
Other Experiments. In Appendix F.5
CONCLUSION AND FUTURE WORK
This paper introduces novel manifold-based WDRO method to enhance the robustness of image classification. We design a game that integrates contrastive learning with Jacobian regularization to capture the manifold structure, allowing to solve DRO problems constrained by the data manifold. Experiments on several popular benchmark datasets demonstrate the method's 's advantages in terms of accuracy and robustness. Additionally, framework for other scenarios that require access to the tangent of the data manifold, extending beyond the MWDRO problem. This is interesting avenue for future research.
manifold-batd model could potentially serve
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
The research of this work was supported in part by the National Key Research and Development Program of China (No. (0210000))))0) the National Natural Science Foundation of China (No. 62272432, No. 62432016), and the Natural Science Foundation of Anhui Province (No. 2208085MF163), H. Ding is the corresponding author.
REFERENCES
Philip Bachman, R Devon Hjelm, and William Buchwalter, Learning representations by maximizing mutual information across views. Advances in neural information processing systems, 32, 2019.
Randall Balestriero and Yann LeCun. Contrastive and non-contrastive self-supervised learning re- cover global and local spectral embedding methods. Advances in Neural Information Processing Systems, 35:26671 10:6611280 15226671-26685, 2022.
Gary Bécigneul and Octavian- Eugen Ganea. Riemannian adaptive optimization methods. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019. URL
net. www.teeet.eetteteeeeee
r rleigi09K7
Nicolas Boumal. An introduction to optimization on smooth manifolds. Cambridge University Press, 2023.
Stephen Boyd, Stephen Boyd, and Lieven Vandenberghe. Convex optimization. Cambridge uni- versity press, 2004.
Wieland Brendel, Jonas Rauber, Matthias Kümmerer, Ivan Ustyuzhaninov, and Matthias Bethge. Accurate, reliable and fast robustness evaluation. Advances in neural information processing systems, 32, 2019.
Michael M. Bronstein, Joan Bruna, Taco Cohen, and Petar Velickovic. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. CoRR, abs/2104. 2021.
Anh Tuan Bui, Trung Le, Quan Hung Tran, He Zhao, and Dinh Phung. unified wasserstein distri- butional robustness framework for adversarial training. In International Conference on Learning Representations, 2022. URL https: http:::::///peneeeee...... net wtwp:/ewwwweee/ooodddddddd--e
Jonathon Byrd and Zachary Lipton. What is the effect of importance weighting in deep learning?
In International conference machine learning, pp. 872-881 PMLR, 2019.
Nicholas Carlini and David Wagner. Towards evaluating the robustness of neural networks. In 2017 ieee symposium on security and privacy (sp), pp. 39-57. Iece, 2017.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference machine learning, pp. 1597-1607. PMLR, 2020.
Sue Chung, Daniel D. Lee, and Haim Sompolinsky. Linear readout object manifolds. Phys- ical Review E, 93(6), June 2016. ISSN 2470-0053.
Jeremy Cohen, Elan Rosenfeld, and Zico Kolter. Certified adversarial robustness via randomized smoothing. In international conference on machine learning. pp. 1310-1320. PMLR, 2019.
Uri Cohen, Sue Yeon Chung, Daniel D Lee, and Haim Sompolinsky. Separability and geometry of object manifolds in deep neural networks. Nature Communications, 11(1):1-13, 2020.
Antonio
Criminisi, Toby Sharp, and Andrew Blake. Geos: Geodesic image segmentation. In Com- puter Vision-ECCV 2008: 10th European Conference on Computer Vision, Marseille, France, October 12-18, 2008, Proceedings, Part 10, pp. 99-112. Springer, 2008
Francesco Croce and Matthias Hein. Reliable evaluation of adversarial robustness with an ensemble of diverse parameter-free attacks. In International conference on machine learning, pp. 2206- 2216. PMLR, 2020.
11

---

Published as a conference paper at ICLR 2025
Yin Cui, Menglin Jia, Tsung- Yi Lin, Yang Song, and Serge Belongie. blas--allaned loss based on effective number samples. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp.
2019.
Sebastian Damrich, Niklas Böhm, Fred A Hamprecht, and Dmitry Kobak. From t-SNE to UMAP with contrastive learning. In The Eleventh International Conference Learning Representations, ew. net/ fo forum? 3=
2023. URL
Li Deng. The mnist database of handwritten digit images for machine learning research. 1EEE Signal Processing Magazine, 29(6):141-142 2012.
Marius Dragoi, Elena Burceanu, Emanuela Haller, Andrei Manolache, and Florin Brad. Anoshift:
A distribution shift benchmark for unsupervised anomaly detection. In S. Koyejo, S. Mohamed,
A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 32854-32867. Curran Associates, Inc., 2022.
John C Duchi and Hongseok Namkoong. Learning models with uniform performance via distribu- tionally robust optimization. The Annals of Statistics, 49(3):137--144 2021.
Peyman Esfahani and Daniel Kuhn. Data-driven distributionally robust optimization using the Wasserstein metric: Performance guarantees and tractable reformulations. Mathematical Pro- gramming, 171, 05 2015.
Tongtong Fang, Nan Lu, Gang Niu, and Masashi Sugiyama. Rethinking importance weighting for deep learning under distribution shift. Advances in neural information processing systems, 33: 11996-12007, 2020.
Tongtong Fang, Nan Lu, Gang Niu, and Masashi Sugiyama. Generalizing importance weight- ing to universal solver for distribution shift problems.
Alice Oh, Tristan Nau- mann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Informa- tion Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16, 2023, 2023. URL (G229ee0017700322355555222eeeee2
ni nips lss html
Rui Gao and Anton Kleywegt. Distributionally robust stochastic optimization with wasserstein distance. Mathematics of Operations Research, 48(2):603-655, 2023.
Rui Gao, Xi Chen, and Anton Kleywegt. Wasserstein distributionally robust optimization and variation regularization. Operations Research, 12(3):1777-1191, 2024. 2024.
Salvador García, Julián Luengo, and Francisco Herrera. Data preprocessing in data mining, vol- 72. Springer, 2015.
Nathan Halko, Per-Gunnar Martinsson, and Joel Tropp. Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions. SIAM review, 53 (2):217-288, 2011.
Grani A Hanasusanto and Daniel Kuhn. Conic programming reformulations of two-stage distri- butionally robust linear programs over Wasserstein balls. Operations Research, 66(3):849-869, 2018.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEFE conference on computer vision and patlern recognition, pp- 770-778, 2016.
Kaiming He. Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEF/CVF conference on computer vision and patlern recognition, pp. 9729-9738. 2020.
Judy Hoffman, Daniel A Roberts, and Sho Yaida. Robust learning with jacobian regularization. arXiv preprint arXiv: 1908 0200990..9...... 2019.
12

---

Published as a conference paper at ICLR 2025
Youngkyu Hong, Seungju Han, Kwanghee Choi, Seokjun Seo, Beomsu Kim, and Buru Chang, Dis- entangling label distribution for long-tailes visual recognition. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 6626-6636, 2021.
Tianyang Hu, Zhilii Liu, Fengwei Zhou, Wenjia Wang, and Weiran Huang. Your contrastive leamn ing is secretly doing stochastic neighbor embedding. In International Conference on Learning Representations, 2023.
Ruomin Huang, Jiawei Huang, Wenjie Liu, and Hu Ding. Coresets for wasserstein distributionally robust optimization problems. Advances in Neural Information Processing Systems, 35:26375- 26388, 2022.
Weiran Huang, Mingyang Yi, Xuyang Zhao, and Zihao Jiang. Towards the generalization of con- trastive self-supervised learning. In International Conference on Learning Representations, 2023.
Daniel Jakubovitz and Raja Giryes. Improving dnn robustness to adversarial attacks using jacobian regularization. In Proceedings of the European conference on computer vision (ECCV), pp. 514 529, 2018.
Prannay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishman. Supervised contrastive learning. Advances in neural information processing systems, 33:18661-18673, 2020.
Masanari Kimura and Hideitsu Hino A short survey on importance weighting for machine learning, 2024. URL Uttpttps:: aawpiv.ivvvv//a////// org/abb/2403.110
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Bal- subramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, et al. Wilds: A benchmark of in-the- wild distribution shifts. In International Conference on Machine Learning, pp. 5637-5664. PMLR, 2021.
Alex Krizhevsky, lya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convo- lutional neural networks. Advances in neural information processing systems, 25, 2012.
Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convo- lutional neural networks. Communications of the ACM, 60(6):84 90, 2017.
Alexey Kurakin, lan J Goodfellow, and Samy Bengio. Adversarial examples in the physical world. In Artificial intelligence safety and security, pp. 99-112. Chapman and Hall/CRC, 2018.
Changhyeok Lee and Sanjay Mehrotra. distributtionally-oouus approach for finding support vector machines. Available from Optimization Online, 2015.
Yonghyeon LEE, Sangwoong Yoon, MinJun Son, and Frank C. Park. Regularized autoencoders for isometric representation learning. In International Conference on Learning Representations, 2022.
Daniel Levy, Yair Carmon, John C Duchi, and Aaron Sidford. Large-scale methods for distribution- ally robust optimization. Advances in Neural Information Processing Systems, 33:8847-8860, 2020.
Jajjnn Li, Sen Huang, and Anthony Man-Cho So. A first-order algorithmic framework for distri- butionally robust logistic regression. Advances in Neural Information Precessing Systems, 32, 2019a.
Pan Li, Lei Zhao, Duanqing Xu, and Dongming Lu. Optimal transport of deep feature for image style transfer. In Proceedings of the 2019 4th International Conference on Multimedia Systems and Signal Processing, ICMSSP '19, pp. 167-171, New York, NY, USA, 2019b. Association for Computing Machinery. ISBN 9781450371711.
Tsung- Yi Lin, Priya Goyal, Ross Girshick. Kaiming He, and Piotr Dollár. Focal loss for dense object detection. Proceedings of the IEEE international conference on computer vision, pp. 2980-2988, 2017.
13

---

Published as a conference paper at ICLR 2025
Zhiqiu Lin, Jia Shi, Deepak Pathak, and Deva Ramanan. The CLEAR benchmark: Continual learn- ing on real-world imagery. In Joaquin Vanschoren and Sai- Kit Yeung (eds.), Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks 1, NeurIPS Datasets and Benchmarks 2021, December 2021, virtual, 2021.
Haibin Ling and Kazunori Okada. An efficient earth mover's s distance algorithm for robust histogram comparison. IEEE transactions on pattern analysis and machine intelligence, 29(5):840-833, 2007.
Jiashuo Liu, Jiayun Wu, Bo Li, and Peng Cui. Distributionally robust optimization with data geom- etry. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Shangyun Lu, Bradley Nott, Aaron Olson, Alberto Todeschini, Hossein Vahabi, Yair Carmon, and Ludwig Schmidt. Harder or different? closer look distribution shift in dataset reproduction. In ICML Workshop on Uncertainty and Robustness in Deep Learning, volume 5, pp. 15, 2020.
Yanbiao Ma, Licheng Jiao, Fang Liu, Shuyuan Yang, Xu Liu, and Lingling Li. Curvature-balanced feature manifold learning for long-tailed classification. In Proceedings of the IEEE/CVE confer- ence on computer vision and pultern recognition, pp. 15824-15835. 2023.
Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu, To- wards deep learning models resistant to adversarial attacks. In International Conference on Learn ing Representations, 2018. URL http:: /ttpp//:////eeeevieweee net/ htt/://orommmiimmmm 20orum2id==J211E2Aa
James Martens, Ilya Sutskever, and Kevin Swersky. Estimating the hessian by back-propagating curvature. In Proceedings of the 29th International Coference International Conference Machine Learning, pp. 963-970, 2012.
Yuhao Mo. Chu Han, Yu Liu, Min Liu, Zhenwei Shi, Jiatai Lin, Bingehao Zhao, Chunwang Huang, Bingjiang Qiu, Yanfen Cui, al. Hover-trans: Anatomy-awan hover-transfore for roi-free breast cancer diagnosis in ultrasound images. IEEE Transactions on Medicul Imaging, 2023.
Peyman Mohajerin Esfahani, Soroosh Shafeerzadeh-Abadhh, Grani A Hanasusanto, and Daniel Kuhn. Data-driven inverse optimization with imperfect information. Mathematical Program- ming, 167(1):191-234, 2018.
Hongseok Namkoong and John C Duchi. Stochastic gradient methods for distributionally robust optimization with f-divergences. Advances in neural information processing systems, 29, 2016.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances neural information processing systems, 32, 2019.
Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends@ in Machine Learning, 11(--6):355-677, 2019.
Krzysztof Postek, Dick den Hertog, and Bertrand Melenberg. Computationally tractable counter- parts of distributionally robust constraints on risk measures. SIAM Review, 58(4))603-650, 2016.
Fengchun Qiao and Xi Peng. Topology-aware robust optimization for out-of-distribution general- ization. ICLR, 2023.
Hamed Rahimian and Sanjay Mehrotra. Distributionally robust optimization: A review. arXiv preprint arXiv:1908.0566 2019.
Salah Rifai, Yann N Dauphin, Pascal Vincent, Yoshua Bengio, and Xavier Muller. The manifold tangent classifier. Advances neural information processing systems, 24, 2011a.
Salah Rifai, Grégoire Mesnil, Pascal Vincent, Xavier Muller, Yoshua Bengio, Yann Dauphin, and Xavier Glorot. Higher order contractive auto-encoder. In Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2011, Athens, Greece, September 5-9, 2011, I, Proceedings, Part Il 22, pp. 645-660. Springer, 201 1b.
14

---

Published as a conference paper at ICLR 2025
Salah Rifai, Pascal Vincent, Xavier Muller, Xavier Glorot, and Yoshua Bengio. Contractive auto- encoders: Explicit invariance during feature extraction. In Proceedings of the 28th international conference on international conference machine learning, pp. 833 840, 2011c.
Andrew Ross and Finale Doshi Velez. Improving the adversarial robustness and interpretability of deep neural networks by regularizing their input gradients. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018.
Sam T Roweis and Lawrence K Saul. Nonlinear dimensionality reduction by locally linear embed- ding. science, 290(5500):2333-22222 2000.
Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust neural networks. In International Conference on Learning Representations, 2020.
Dvir Samuel and Gal Chechik. Distributional robustness loss for long-tail learning. In Proceedings of the JEEE/CVF International Conference on Computer Vision, pp. 9495-9504, 2021.
Soroosh Shafieezadeh Abadeh, Peyman M Mohajerin Esfahani, and Daniel Kuhn. Distributionally robust logistic regression. Advances in Neural Information Processing Systems, 28, 2015.
Jun Shu, Qi Xie, Lixuan Yi, Qian Zhao, Sanping Zhou, Zongben Xu, and Deyu Meng. Meta- weight-net: Learning an explicit mapping for sample weighting. Advances in neural information processing systems, 32, 2019.
Aman Sinha, Hongseok Namkoong, and John C. Duchi. Certifying some distributional robustness with principled adversarial training. In 6th International Conference on Learning Representa- tions, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceed- ings. OpenReview.net, 2018
Alessandro Sordoni, Nouha Dziri, Hannes Schulz, Geoff Gordon, Philip Bachman, and Remi Tachet Des Combes. Decomposed mutual information estimation for contrastive representation learning. In International Conference on Machine Learning, pp. 9859-98899 PMLR, 2021.
Matthew Staib and Stefanie Jegelka. Distributionally robust deep learning as generalization of adversarial training. In NIPS workshop on Machine Learning and Computer Security, volume 3, pp. 4, 2017.
Matthew Staib and Stefanie Jegelka. Distributionally robust optimization and generalization in ker- nel methods. In H. Wallach, Larochelle, A. Beygelzimer, F. d'Alché-Buu Fox, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Zhiqiang Tang, Yunhe Gao, Yi Zhu, Zhi Zhang, Mu Li, and Dimitris N. Metaxas. Crossnorm and selfnorm for generalization under distribution shifts. In Proceedings of the IEEE/CVF Interna- tional Conference on Computer Vision (ICCV), pp. 52-61, October 2021.
Rohan Taori, Achal Dave, Vaishaal Shankar, Nicholas Carlini, Benjamin Recht, and Ludwig Schmidt. Measuring robustness to natural distribution shifts in image classification. Advances in Neural Information Processing Systems, 33:18583-18599 2020.
Joshua B Tenenbaum, Vin de Silva, and John C Langford. A global geometric framework for nonlincar dimensionality reduction. science, 290(5500):219-22222 2000.
Yonglong Tian, Chen Sun, Ben Poole, Dilip Krishnan, Cordelia Schmid, and Phillip Isola. What makes for good views for contrastive learning? In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), Advances in Neural Information Processing Systems, volume 33, pp- 6827-6839. Curran Associates, Inc., 2020.
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(11), 2008
Cédric Villani. Optimal transport: old and new, volume 338. Springer, 2009.
15

---

Published as a conference paper at ICLR 2025
Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through align- ment and uniformity on the hypersphere. In International Conference on Machine Learning, pp- 9939. PMLR, 2020.
Yifei Wang, Wen Li, Dengxin Dai, and Luc Van Gool. Deep domain adaptation by geodesic distance minimization. In Proceedings of the IEEE International Conference on Computer Vision (ICCV) Workshops, Oct 2017.
Zhenyi Wang, Li Shen, Le Fang. Qiuling Suo, Tiehang Duan, and Mingchen Gao. Improving task- free continual learning by distributionally robust memory evolution. In International Conference on Machine Learning, pp. 22985-22998, PMLR, 2022.
Olivia Wiles, Sven Gowal, Florian Stimberg, Sylvestre-Alvie Alviie Rebuffi, Ira Ktena, Krishnamurthy Dj Dvijotham, and Ali Taylan Cemgil. fine-grained analysis on distribution shift. In International Conference on Learning Representations, 2022.
Yihao Xue, Kyle Whitecross, and Baharan Mirzasoleiman. Investigating why contrastive learning benefits robustness against label noise. In International Conference on Machine Learning, pp- 24851-24871 PMLR. 2022.
Meimei Yang, Qiao Liu, Xinkai Sun, Na Shi, and Hui Xue. Towards kemeliizing the classifier for hyperbolic data. Frontiers of Computer Science, 18(1):18130) 2024.
Chongzhi Zhang, Mingyuan Zhang, Shanghang Zhang, Daisheng Jin, Qiang Zhou, Zhongang Cai, Haiyu Zhao, Xianglong Liu, and Ziwei Liu. Delving deep into the generalization of vision trans- formers under distribution shifts. In Proceedings the IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), pp. 7277-7286, June 2022.
Hongyang Zhang, Yaodong Yu, Jiantao Jiao, Eric Xing, Laurent El Ghaoui, and Michael Jordan. Theoretically principled trade-off between robustness and accuracy. In International conference on machine learning, pp. 7472-7482. PMLR, 2019.
Hongyi Zhang and Suvrit Sra. First-order methods for geodesically convex optimization. In Con- ference learning theory, pp. 1617-1638. PMLR, 2016.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. In International Conference on Learning Representations, 2018.
Yang Zhao and Hao Zhang. Neighborhood region smoothing regularization for finding flat minima in deep neural networks. In Proceedings of the Asian Conference on Computer Vision, pp. 4697- 4710, 2022.
Chunting Zhou, Xuezhe Ma, Paul Michel, and Graham Neubig. Examining and combating spurious features under distribution shift. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research. pp. 12857-12867. PMLR, 18-24 Jul 2021.
A APPENDIX
A.1 BACKGROUNDS
In this section, we provide some background on the manifold. For more detailed information, we refer readers to the comprehensive textbook by Boumal (202)).
Assumption T states that the manifold is a complete manifold embedded Euclidean space. This assumption ensures that the manifold's geometric properties well-suited for optimization tasks. Roughly speaking, a complete manifold implies that every geodesic can be extended infinitely, pre- venting the optimization process from encountering boundaries or undefined regions that could dis- rupt convergence.
Moreover, assume that the manifold M is embedded in Euclidean space Rd meaning it inherits a well-defined metric from the ambient space Rd. This allows us to define geodesic distances, per- form projections, and compute gradients in manner that is consistent with the underlying geometry
16

---

Published as a conference paper at ICLR 2025
of the ambient Euclidean space. Such an embedded manifold also simplifies many computational tasks, such as the calculation of the Riemannian gradient. Specifically, the Riemannian gradient V. (f(x)) can be computed by the orthogonal projection of the "classical" gradient to the tangent spaces T (x)... i.e, V-M f(x) = Projx (V (x)), where the Proiz (-) is the orthogonal projection operator that projects vector to the tangent space M and (ff(x) is the gradient in Euclidean space R².
For the mathematical rigor, we briefly verify that the prerequisites of Gao & Kleywegt 20233 The- orem 1) are satisfied within our framework, as defined in (5) by examining the properties of the manifold MM. By equipping M with the geodesic distance, we can verify that it Polish space (i.e., a separable and complete metric space). Also, the space of Borel probability measures on M is well-defined. Additionally, the mild Assumption in Appendix [C Censurss the continuity of our loss function ( p). These properties allow us to safely apply Gao & Kleywegt (2023)'s 's strong duality results within our manifold- based WDRO framework, as formalized in Proposition
A.2 OTHER RELATED WORKS
More DRO methods. DRO demonstrate demonstrate significant advantages in terms of general-
ization and robustness across multiple domains (Gao et al. 20244 Huang et al. 2022 Wang et al.
2022). The works by (Staib & Jegelka 2017 Sinha et al. 2018 Bui et al. 2022) are closely related to ours, as they also use the dual form of Wasserstein Distributionally Robust Optimization (DRO) to scarch for worst-case perturbations. However, these approaches cannot extend directly to mani- fold constraints; while they are effective against adversarial attacks, their overly pessimistic nature clean data. Another line of works considers the Group DRO (Sagawal (et al. 2020). which aims to optimize performance across different demographic subgroups, address- ing hidden biases and ensuring fairness. f-divergence DRO (Namkoong & Duchi 2016 Duchi & Namkoong 2021). on the other hand, considers an uncertainty set within a certain j-divergence from the nominal distribution. Unlike Wasserstein distance, f f-divergence measures the dissimilar- ity in terms of densities at corresponding points, and becomes infinite when two distributions have different support sets. This intrinsic difference allows conventional WDRO to extend the training support, while f-divergence cannot (Liu et al.] (022). Several studies attempt to incorporate geo- metric information into the uncertainty set. Qiao & Peng (2023) calculate group centrality based on the affinity between data groups, serving as good nominal distribution for group DRO, Liu et al. (022) employ a k-nearest neighbor graph structure to capture the geometric structure of the dataset, where shortest paths approximate geodesic distances to guide the transportation of probability mass. Their Geometric Wasserstein DRO framework constrains the distribution support set to observed data points. In contrast, our approach learns the data manifold structure through neural networks. This continuous, differentiable representation offers several advantages: it captures smooth geomet- ric variations in the data manifold rather than relying on discrete graph approximations; furthermore, it naturally accommodates out-of-sample points, which enables more flexible adaptation to unseen data.
can decrease performance
Importance weighting (IW) has emerged as a powerful fam- ily of methods for addressing distribution shift in machine learning (Byrd & Lipton 2019 Kimura & Hino, 2024). Several advancements have extended IW techniques for tackling various challenges. For instance, Shu et al. (2099) developed a meta-learning framework that learns an explicit weight- ing function to handle biased training data, such as label noise and class imbalance, by leveraging a small unbiased validation set. [Fan et al. (2020) proposed an end-to-end framework that seamlessly integrates weight estimation and weighted classification through an iterative process. More recently, Fang et al. (2023) introduced generalized importance weighting (GIW), which extends traditional IW methods to handle various types of distribution shift, including challenging scenarios where the training and test distributions have different support sets. Jacobian regularization has been used to enhance robustness against small perturbations, such as noisy or adversarial data (Jakubovitz & Ciryes 2018 Hoffman et al., 2019). Jakubovitz & Giryes (2018) empirically demonstrated that enhances robustness with minimal impact on the network's accuracy on clean data. Hoffman et al. (2019) proposed a computationally efficient implementation of Jacobian regularization. Further- more, Rifai et al. (2011a c) applied Jacobian regularization to autoencoders to learn representations that capture the local directions of variation dictated by the data, thereby improving the accuracy of classification tasks.
Other related robustness methods.
17

---

Published as a conference paper at ICLR 2025
Contrastive learning is powerful self-superrised learning paradigm that has significantly advanced representation learning (Wang & Isola 2020: Xue et al, Chen et al.) 2020; He et al. 2020). Khosla et al. (2020) introduced the supervised contrastive loss, achieving superior performance compared to cross-entropy on several tasks. However, the
Contrastive learning methods.
2022
theoretical understanding of contrastive learning (CL.) remains limited. Some works interpret the
InfoNCE loss from a mutual information perspective (Bachman et al. 2019] (Soddoi et al.) (021),
showing it maximizes mutual information between different views of the data. Other studies (Wang & Isola, 2020 Huang et al., 2023) analyze the geometry of the embedding space, revealing that InfoNCE comprises two parts: one ensuring aligmment and the other preventing representation col-
lapse. Additionally, CL's feature embedding has been linked to spectral clustering (Balestriero &
LeCun 2022) and manifold learning (Hu et al. 2023
B MANIFOLD LEARNING PERSPECTIVE OF CONTRASTIVE LEARNING
In this section, we explore our proposed manifold-guided game through the lens of manifold learning and representation learning. By introducing important properties from contrastive learning, we ana- lyze how our model effectively captures the structure of the data manifold, providing deeper insights into our approach. In Appendix[B.1 we introduce several properties commonly discussed in the con- trastive learning literature. These properties will aid in analyzing the effects of our manifold-guidee approach, which we consider to be a complement to Section 4.2) Furthermore, in Appendix [B.2] we show that the Jacobian regularization can enhance contrastive learning by preserving neighbour relationships within the input data, which benefits classification tasks. Building upon this observa- tion, we refine our model to learn smoother manifolds, further improving its ability to capture the structure of the data manifold.
B. PROPERTIES OF THE CONTRASTIVE LEARNING
number of manifold learning methods, such as 1-SNE (Van der Maaten & Hinton 2008), UMAP (Damrich et al. 2023), and Isomap (Balestriero & LeCun 2022), are based the neighbour embedding techniques. These methods aim to preserve essential structural information from high- dimensional data in low-dimensional embedding space. This is achieved by learning mapping that encodes data points from high-dimensional space into a low-dimensional space, ensuring that data points that are neighbors in the original space remain closely situated in the embedding space. Specifically, given set of points (xi.... J and an encoder g(-). we define the input/embedding affinity matrices A, B Raxx (i) input space affinity: Ai,j denote the affinity between the sam- ples x; and X; in the original high- nigh-dimensional space; (ii) embedding affinity: B is the affinity computed in the embedding space between g(x,) and g(x,). Then the affinity- preetttiig entropy is given by Laffinity (g) = Σizy Aij log (Brr), which quantifies how well the embedding preserves the pairwise affinities of the original data.
Definition 3 (Neighbor-preserving property) Let A denote the impul space affinity matrix based on a k-nearest neighbor graph G constructed from the inpul data. Specifically, A 1 ((i.j) E G) //G where is the indicator function and [G] the total number of edges in the graph. The embedding affinity matrix B is computed using the Cauchy kernel: B,j
1(x+) An encoder g(-) is called neighbor-presseringg if it minimizes the affinity- preserving entropy (g)
The affinity matrices described in Definition 3 are, in fact, the same as those used in the Stochas- tic Neighbor Embedding (SNE) (Hu et al. 2023) method. Similarly, contrastive learning can be interpreted as special case that also minimizes the affinity-peseerving entropy, but with different choices for the affinity matrices. We define these affinities as:
[1, if ;r; and my are positive pairs (0, otherwise,
Ai.i =
(14)
exp (sim (f (xx) f (2,)) /T)
Bid = Ener exp (sim (f (xx). .f (xx)) /T)
(15)
18

---

Published as a conference paper at ICLR 2025
where sim(.:-) denotes a similarity measure (e.g.. cosine similarity), and is temperature pa- rameter. With these choices, the affinity-preserving entropy Latlinity (g) = - Eife Ai, log (B,,) reduces to the InfoNCE loss in Eq.()).
By the choice of the input/embedding affinity in Eq 14 15), experimental and theoretical results in prior works (Wang & Isola 2020 Hu et al., 2023) show that minimizing the InfoNCE loss leads to encoders that approximately achieve the properties of perfect alignment and perfect uniformity (see Definition 4. Specifically, the global minimizer of the contrastive loss requires that features of positive pairs align (ain)g())))))))))) g(x())) tends to 1) and those of negative pairs be distant as possible.
Definition 4 (Perfect Alignment/Uniformic WWng & Isola, 2020 (i) Perfect Alignment. An encoder g()) achieves perfectly alignment over dataset Pue g( g(x²) almost surely for all æ',, x E aug(:;), Va; E Pari (ii) Perfect Uniformity, An encoder g(+) achieves perfect unifor- mity over dataset Per if the embeddings g ((x)) Vx, E Par are maximally separated as |Prll points on the sphere Sa. E R where the de is the dimension of the embedding space.
We now re-examine our proposed manifold-guice game from the perspectives of alignment and uniformity. The repulsive force between negative pairs (i.e., pushing apart g(x() and g(x) in the InfoNCE loss (9. helps preserve maximal information of the data manifold. Simultaneously. the attractive force between the positive pairs (i.e., pulling together g(x() and g(x()), combined with Jacobian regularization, minimizes sensitivity orthogonal to the manifold or directed towards the semantic-invariant region Msr(x).
Based on this analysis, we summarize the desired properties of the encoder g(.) as follows: (I) Sensitivity. The encoder g(x) should change most when the sample moves along the semantic variant directions within the data manifold; (II) Insensitivity. The encoder g(x) should change minimally when moves off the manifold (i.e., M or towards the Cermantic-Ivvarian region (i.e., (x).
B.2 TOWARDS SMOOTHER MANIFOLD AND BETTER REPRESENTATION
In this section, we demonstrate that the Jacobian regularization can enhance contrastive learning by better preserving neighbour relationships within the input data. We have known that the alignment and uniformity properties (see Definition 4 enable the game between CL and Jacobian regulariza- tion shape representations that capture the tangents of the data manifold. However, considering only the contrastive loss may not adequately preserve the neighbor structure. This is because any permutation of the mapping - g (x) for E Per can achieve perfect alignment and uniformity, thus reaching the global minima of Eq. 9 Figure (4) illustrates this issue with two encoders, gı (-) and its permutation g2(-). Both achieve perfect alignment and uniformity, but g2(-) preserves the neighbor relationship of the input data, clustering similar samples (e.g.. birds) in the embedding space. This clustering is beneficial for downstream classification tasks due to the linear separabil- ity of the features (Wang & Isola, 2020) (as shown in Figure 4b embedding space by a linear classifier).
we can classify the birds in the
To analyze the quality of neighbor structure preserving. Hu et al 2023 propose a complexity measure defined as:
[[((x) x x' 2
1|2
C(g) = En
(16)
where x, x' are negative pairs as in Eq 9 This complexity measure captures the degree to which the mapping g(.) preserves the neighbor structure of the input data in the context of contrastive learning. For example, in Figure (4) C(ga) < C(g.), which implies g2 preserves the neighbor structure better than g1. Building on their findings, the Jacobian regularization can be considered a differential version of this complexity measure, C(g). Our empirical results confirm that Jaco- bian regularization reduces C(g), thereby enhancing the wewpbborhooo--peseeviie properties of the model.
However, the Jacobian provides only a local approximation to C(g) as defined in Eq. Eq. Intu- itively, enforcing smoothness on the Jacobian Jg(x) improves neighborhood preservation by ensur- ing that the encoder g(.) changes gradually without abrupt variations. Specifically, this smoothness
19

---

Published as a conference paper at ICLR 2025
Eneer classi fior
(a) Emhedding space g-)))
(b) Embedding space ga(:)
Figure 4: Both encoders g1 (-) and g-(-) achieve perfect alignment and uniformity, reaching the em- pirical minima of the loss Eq 99. The encoder g-(.) in Figure (4a) is permutation of the encoder mapping g2(-) in Figure 4b). g2 g2(") preserves the affinity of the input dataset Par, with similar sam- ples (birds, cats and ships) well clustered in the embedding space. We prefer g-(-) as the clustered features beneficial for downstream classification tasks due to their linear separability (Wang & Isola 2020
beneficial for accurately capturing the directions of semantic variation within the data manifold. As is well known, second-order derivatives (i.e., the Hessian), measure the smoothness of the Jaco- bian (Rifai al.] 20TTb; Zhao Zhang, 2022). However, computing the Hessian matrix com- putationally expensive. To address this, can the mixup data augmentation technique (Zhang) itt al. 2018) to enhance our manifold-guice -ggiided game, as mentioned in Remark Specifically, we cre- samples by mixing sample x; Per with one of its neighbor xj, forming Ax; + A)xx, and incorporate these new samples into the training process. This mixup technique helps stabilize the Jacobian when x moves along the data manifold, favoring smoother representation.
We construct the mixed dataset using the mixup (Zhang et al. 2018) technique and denote it as pmix : {ar,, ar// A) Jac P'r xrl E P Aij * 0, Uni D)}. where A is the input affinity matrix in Definition and Un, 1) is the uniform distribution over the interval (0, 1).
The modified InfoNCE loss with mixup is defined as:
sim ((xa (Axr Xaa
LInfONCE
E panix
log
(17)
Exerg (a sim(x ax) x())))
where the notations follow Eq. 9.. Then the complete objective function becomes:
L. (0)
A
oF TEP REPLUPPUPPE (x)
E
(18)
The mixup-enhaneee loss encourages the model to produce consistent embeddings for interpolated inputs, helping to smooth the representation and stabilize the Jacobian without the costly Hessian regularization.
We compute the C(g). as defined in Eq.(16. ()..F. for well-trained models on the CIFAR- 10 and CIFAR- 100 datasets. The results, presented in Table show that Jacobian regularization, particularly when combined with the mixup technique, significantly enhances the neighborhoods preservation in fea- ture embeddings. By stabilizing the Jacobian Jg as x traverses the data manifold, our model encodes the data manifold more smoothly. This yields two key advantages: first, it facilitates better extraction of tangents on the data manifold; second, the clustered embeddings are beneficial for downstream classification tasks, illustrated in Figure
20

---

Published as a conference paper at ICLR 2025
Table 1: We compute C(g) (as defined in (g)(6) for various well-trained models on CIFAR- 10 and CIFAR-100 100 datasets. The methods compared include: Contrastive Learning (CL, Eq. Contrastive Learning with Jacobian regularization (CL+Jac, (q. and Contrastive Learning with both Jacobian regularization and mixup (CL+Jac+mixup. Eq.(17)) We normalize all the results by dividing them over the results of CL. Lower values indicate better neighbour preservation.
Objectives
CIFAR-10
CIFAR-100
CL
1.00
1.00
CL+Jac
0.29
0.40
CL+Jac+mixup
0.10
0.11
C APPROXIMATION OF THE RIEMANNIAN GRADIENT
As discussed in Section we utilize the strong duality formulation of the MWDRO problem presented in Proposition as our objective function. The main challenge in optimizing this function arises from the surrogate loss defined in Eg.(7) To compute this surrogate loss f, pt. 1/) ==
SuP [E(0,q) - v & (g. p:)]. we use the first-order optimization method (Zhang & Sra 2016 The update rule shown Eq. [3). For clarity, we rewrite here:
g+tl= = E (aV-M (e(0, (d
= Exp,, (aProj (VE(O,'')) 20vdd( ()(gt.))
The above updating steps require the Riemannian gradient, which is the orthogonal projection of the "classical" gradient to the tangent spaces (Boumal 2023). i.e. fa')(g)) = Proj, (VI(q")).TT central problem, then, is how to recover the tangent space (m so as to approximate the Rie- mannian gradient.
In Section 4.2. we design the manifold-guided game for training the feature encoder g(-). We analyze the singular value decomposition (SVD) of the Jacobian matrix Jg(x) U()) (x), where U(x) dd, V (x)
[x), ... [x)] [xx) = diag(01 (x), ... (x)x)) We define the space spanned by the principal singular vectors of the Jacobian matrix in Eq. 10.. i.e., Tappr To) == span (xx)) ox(x) 2 To, i [r]].
R and
As discussed in our main paper, we aim to approximate the Riemannian gradient V-M(O,,'') q') by projecting the Euclidean gradient Ve(O, g²) onto the subspace Tappe (g', tO). We denote this opera- tion as Projapp: (VE(0. (VE(C, q²)). In what follows, we analyze the rationale behind this approximation. We begin with assumptions we require.
The key idea is utilizing Tappr (X, To) to approximate subset of the tangent TM(x), capturing directions of significant semantic variation within T (x).. We formalize this as Assumption 3
Assumption For some f 2 T Tappr (x, f) subspace of the tangent plane TM for any E M.
Now, let us examine the loss function (0,,)). We can view it composition of functions involving g(x). The loss E(-) is usually continuous w.r.t. g(x), as detailed in Assumption 4
Assumption The loss function E(O, x) differentiable with respect to g(x). Moreover, E(O, x) is T--Lipschitz with respect to g(x). This means that for any two inputs x1 and x2, the inequality |8(0,x)) x1) f(0, x2) < Lllg(x)) - g(xa)]l holds for any parameter 0.
Remark In classification tasks, a common choice is E(0,x) CE CE (S(Sottaax(g(()) where CE(-) is the cross-entropp loss. It is well-known that V E()))((( {(f,))) x) = Softmax(g(x)) y. where y is the one-hot label corresponding to x. Therefore, |IV g(x) ((0. x) S 2, implying that the cross- entropy loss is 2-Lipschitz with respect to g(x).
Theorem (Approximation of Riemannian Gradient) Under Assumptions and for any & S ME(E, 9¹) with the
To < max; G. Projep (VE(e, g¹)) approximates the Riemannian gradient following error bound: V E(0, q¹) Projsppr((V ))))2 $ L To-
21

---

Published as a conference paper at ICLR 2025
Proof. By the chain rule, we express the classical gradient
de
ve(O,9") = 0g
J UEVT
(19)
0g
Let w = [w],... _ = 8f U. Then the gradient can be expressed as: VE(O,g') = wEV
Eis1 W4 The approximate projection operator with threshold To is defined as:
Projaper (V((0,'))
Σ
W,N,V
(20)
2:0:270
The Riemannian gradient obtained by projecting the classical gradient onto the tangent space at q¹. This projection be decomposed follows:
e(0,0') = Projg
wio.v.) Σ
Proig (wroiv))
Σ
W,O.V, +
Σ
Projee (WiO.V,
(21)
2:
20) follows directly from the definition of the approximate tangent space Tappr (x, f).
By Assumption 3 we have Tappr (x, (x, C (M((),, which together with the definition of T apppr (oppr((x (x, f)
establishes 21.
The difference is:
V-Me(O,g') Proj (VI(0,)'))
Σ
Woo +
Σ is
Projet (wioov.)
(22)
Taking the norm and using the orthogonality ve:
[|0,M((0,'') Projuppe (Ve(O,
Σ
|wee 2o3
(23)
By Assumption 4 we have 11 l2 L.
Since U is orthonormal matrix, |U||
thus
at U 8g
$ L. Therefore we have
Σ
1u <
T
Σ
wy < TO
w 1 S L
(24)
1=1
Taking the square root:
M( (V((e,'') $ L.
(25)
This theorem provides a justification for approximating the Riemannian gradient by projecting onto
T Tappe (o)
D THEORETICAL ASPECTS OF THE DISTRIBUTIONAL ROBUST
D.1 JACOBIAN TERM AND THE SMOOTHNESS ASSUMPTION
In this section, we demonstrate that the Jacobian regularization term serves as an upper bound for
3. which characterizes the smoothness of our model Assumption 2 Since the data manifold M embedded in the ambient Euclidean space dd without loss of generality, we focus on discussing smoothness within d
22

---

Published as a conference paper at ICLR 2025
We derive the second derivative of (g())) with respect to the imput sample The smallest value 3 can take is the largest eigenvalue of the Hessian at any point x E X. In our setting, 02t
3 supxex Amax ixx
8-8
0l
00
0
0e
0Jg dg
de
Jg
g +
0x2
0x
0x
Ox
Ög
8x
0g
0x
0e
0eg gi
= J H.Jg g
+
0g. 0x
de
J HiJ R
+
H
(26)
0g
The notations used are:
J Jacobian of g, dimension de d.
He: Hessian matrix of the loss E(-) with respect to g(+), dimension de x de.
Hg 0J
D'g 2xx (x): Hessian tensor of g(-), dimension de x d x d.
|| |la: spectral norm of matrix, i.e. the largest singular value.
Using the sub-multiplicative property of the norm ||2, we have
Amx (J.
H
[ x
HH
||gll2
(27)
Ol
di
mmux
H
S
[[] (1,
(28)
0g
0g
Assuming the second partial derivatives of g and are continuous, then H2 and Hg (i, :) for all i [de] are symmetric. Further, not hard verify that J g Applying Weyl's inequality for eigenvalues, we get:
H,Jg and de ög
Hg are also symmetric.
026 Ox2
Amax (J) H:Jg) Amax ( Oe
Amax
Hg
0g
de
He 12
|
[[| (i, :12]
(29)
0g
Note that
J la < Jل For the Hessian term, following the approximation from (Jakubovitz &
Giryes Thus
2018
Martens et al. 2012). we have ( :l) se [ (1.;)))l = Jg
J (i.
06. 0g
ae [[]] (1, :)1T
de
ar
OR 0g
0g. Σ
[ig(),))
|0 J 1
(30)
0g
Finally we have
02
Amax
S
Hell +
0g o [,)||
(31)
0x2
Consequently, we can conclude that the Jacobian regularization term [x) provides an upper bound for B.
D.2 PROOF OF LEMMA
We can verify that the following proposition holds according to Definition [Boumal 2023) for more detailed information.
Please refer to
23

---

Published as a conference paper at ICLR 2025
Proposition? 2 Under Assumption 2 for a given p E M and a constant > B, the function
E(O.p') - vdg (p', p) is geodesically (v -)-strongly concave and geodesically (v + B)-smooth.
Proof.[of Lemma [ Algorithm solves the surrogate loss as defined in Eq. 7 For brevity, we denote E(O, q) as Elg). The surrogate loss is given by:
(s(a = sup [((g) (a
(32)
9EM
It can be shown that the function E(g) vd (²) is geodesically (v B)-strongly concave and geodesically (v 3m-smooth (v 8 0). Thus its condition number, defined as the ratio of the smoothness parameter to the strong concavity parameter, satisfies
シータ
We introduce a sequence of auxiliary functions for the t-th iteration of Algorithm 2 follows,
4g(g) = {(q) Vt dg (q".g).
(33)
where is the value computed as n- (p=')²
~पा (with v/ v). It is not hard note that 14 2 for all t 2 0. Therefore, the function (e(g) is geodesically (1/ - B) B--ttongyyy concave and geodesically
(1, + 3m-smooth, with 24 > > B > 0. As result, its condition number satisfies KG r+8 b- K.
21 48
2
The gradient computed in step 2 of Algorithm 2 (g')) - (v,1(0, (') 2upt'.
dg.. where igy.g denotes the unit vector representing the initial direction of the geodesic from point qº to point 9 on the manifold. Let g² denote the optimal solution of maximizing 4(()). According
to Algorithm 2 the update step is given by: g²+1 Exp,, (zM)x where a E [0, 243 is
the step size. We select a
Given that 4. is (vz + -mmsoohh, we have:
1 2(11 + 3)
(a("t*) 11 s u(("'t) (14")
(34)
Furthermore, by (Boumal 2023
Lemma 11.28), we obtain: (v-Mi(l'')1 2 2(1 - ()((((²**
(,(")), substituting this into Eq. 34) yields:
(²(g"t) 1a(1 s
(1+(a"*) (2 (1 (1- (-1(1"))
(35)
Kit
By recursively applying this inequality, we obtain linear convergence:
(e(a"t)) ( < " ro(9²)
-o("))
(36)
Utilizing the geodesically (v B) -tstong concavity, we have:
2 - - -4(
dg(att1 , <
3
s 2
1 ((o((00 -
(1
3
K
_ - vo(99))
< B
(1
(37)
K
By the 8 + smoothness, we have " ) S 40(9)) + B+V dg(g" a 2222 Then
dx a(' S V 1- 1
k. dat
(38)
Let c = - E 1). Then, we can rewrite as: dz(a', S c (rdzga'.9)
By triangle inequality, for any i 0, we have:
dg (g' +++) < de(a' ( + dglg"t1 .grht)
< c' Vedglg (."")) d+ Vndgla ,9"),
24

---

Published as a conference paper at ICLR 2025
which simplifies to:
ag(d',ttt)) S (1 + (1 ,9").
(39)
Summing over all steps along the trajectory, we obtain:
+ (1+
Σ 1=0
dgla' + s (1+CVBR
dd .-)
Vxdgle"
(40)
1-c
=00
Taking / + oc, we have:
8
1+c Vndy (1'
pt 00
Σ
d +1 S
c
1=0
Since H VK(VK + VR 1)², finally we have:
pt' S ValVK+ + 1)- .
(41)
D.3 PROOF OF THEOREM 1
As mentioned in our main paper, we do not explicitly set the radius for the uncertainty set Ugu (Pur, 8) in 2q. Instead, we select
> 0 to solve the empirical surrogate loss Ep, Ie.. ps. v)] rather than prescribing the robust range 8; this approach, originally proposed by Sinha et al. (0118) 9. has been shown to ensure certified distributional robustness. Then the objec- tive function 2 (R
0, Pt.) = min fur² + Px, (0, Pi, v)} } is simplified to Ep, (0, Pi, v). The following proposition provides the relationships between ô (the radius in Eq.(5) and (the dual variable in Proposition T..
Proposition 3 (Relationships between 5 and v) For a given 0, consider the dual funnctio Eq. O
(i) For 1/2 VI O, we denote 81, such that v1, V2 are the minimizer of (0, (0, Pa)), (D Por) respectively. Then have ōn 2 82.
a EEPER, ea(lpri))
(ii) With Assumption fix constant B and let 52 (0, P.t) v En, (.,
Then
Proof.
Proof (i)
We first establish the monotonic relationship between and 8. This result is intuitive, but for com- pleteness, we provide a rigorous proof here. From the strong duality property proposed by Gao & Kleywegt 2023.
Theorem 1),
Cb, P..) inf +20
v8² + sup (000:/00)0000(.P01
= min {u82 Ep, (.. #20
x)}
(42)
Besides, from the Proposition of (Sinha et al. 2018). we have
sup [EQ vQ P)}
(43)
For any 122 2 vn 0, Since v1, V22 are the minimizer of Lb (0, Pa), Lb (0, P..) respectively. Let
Q1 E arg sup (Eq q)] -GG(Q P.)). Q2 E argsup (Eg) - SW Pu)]. we Q
Q
have &n (QW(Q, P.,) Q Qi) and 82 (9W(Q, P.)) Q Q2).
25

---

Published as a conference paper at ICLR 2025
It is not hard to know sup (Egl((0 q)) (QQ 2 sup (Eole(0, -
(44)
Q
Considering that larger radius ó results in larger loss Dn (0, Pwr), it follows that 51 82.
Proof of
Now derive the precise relationship between and 8. From (Gao & Kleywegt 2023 Lemma (ii)). we know that (0, pr, v)) is convex and non non-increasing in v. Recall that in Eq. we define (8, Pa, v) as (0, Pt, v)) = sup [e(0, q) - vag (q. p.)]. Building on this formulation, we further
dEM
define the following quantities:
D p.). ==
lim sup (dg (q: p.) E(0, q.) 510
-vd
-a qt p.) 2 (.(0, v) 6},
(45)
D(0,u.p)) = lim inf { (g: ps): [(0, q:) va (e pı) 2 L, (0, Ps. D) 6}. 610
(46)
Furthermore, (Gao & Kleywegt 2023 (0, Ps, v) with respect to v. Specifically:
lemma 3 (iv)) shows that D (0,,,, is subderivative of
0ls(0,p,, Ōv
D² D',,, 5
< (,)
(47)
Now, define riinn (0.Pt.1) min {dg(9 p.) E(0, q:) vag (pr, 9:) Lo, ps, w)} as the mini- 2.M
mum distance between Pt and all the q's 's' that attain the supremum of E(0, qq)) - =d (Pi, q:) in M. (0, Ps: v)} as the
Similarly, define max (f., Pt, v/
max (dg(%. Pt) E(0,
maximum distance between p; and all the gi's that attain the infimum of E(A, (.) - vaglin k)) in M. Thus, have the following inequality:
-T -maax (O. -D'(,. (l.(0.pr.D) < (D,, S -m (0, (..Pi.))
(48)
dv
Note that ((0,P:) is geodesically B-smooth with respect to Pt. and e(0, vd (qt, Pt) is at least geodesically (v -)tstrongly concave and (v + 3)-smooth. Therefore, r (,, pi, v) 2 (0,Pi, v), which we denotes as r2,00p in the following. Finally, we obtain
0l,(0.P.. EEp, ()).p..))
8² -EE,,
(49)
dv
Lemma 2 Under Assumption 7 let the expected surrogate loss be denoted by Lo Pu, Ep. E0 Pi, v). The corresponding estimated value, computed with the approximate geodesic dis-
tance in Algorithm 2 is denoted as (6, Ptr. By fixing v 3 and setting v' CıU, where k(VR+ -K-1)* 1)a, we have the following relationship:
C. Pur, v') I (1 Pte, S (.,0P)...).
(50)
Proof.
Consider the surrogate loss (, Pi.v) sup [E(0,g) (g in Eq. 7). It is clear that 9EM
(0, P.. v) > (. Pt, v/) for S 1/ as shown in Proposition 3 Now. consider the approximated
surrogate loss defined as 2, pi, v) = sup () - vdz (q. p.)). where dg p.) an approxi-
cEM mated distance as stated by Lemma Therefore, we have
(, p.. < L.(0, pr. & (,
By taking the expectation over the dataset Ptr. we complete the proof.
26

---

Published as a conference paper at ICLR 2025
Proof. of Theorem Let = n(VR VK- and denote = civ. Recall that ô the solution of ming Ep,, (0, Pi, v) using the approximate geodesic distance, i.e. min Co Pu, = (0, Ptr, v). By Lemma 2 it follows that
0
Par, min (0, Par, P/), S C. (0, P/r, S min L,, Par, 1/) S min
(51)
From Proposition
3 ii), we know that 80,10.0..)) -p² (0, Ps, v). Recall that (0, Ps, v) = 2( q:) ud q:) (.. Pt. 21} represents the minimum distance between
min (g. p.) #EM
and all the gi 's that attain the supremum of e(0, g:) vay 9i) in M. Furthermore, since the function {(f, 9:) vdg (%: pi) is at least geodesically (2 B) -tstonyyyy concave and (v -m--mmothh for a given pi, there exists an S. = [e((O.p)) 2 satisfies A S r(0, S _ Thus, we
can express (0, Pt. //) as:
(/.0 l, Pi. +00) 1"d (0. -(0.pt) +
[-" - t)dt. (52)
t)dt /"d ((1
S 2
S7 1'
S
di
1+ 3
1+ 00
v +
3
Similarly, for (0, Di, v), we have:
= ((0,p.) + /-"
(0, S (
S
2
S?
S?
1-3
dt
1-3
1
V
3
8
We have (0, p;. v) S (0,pi))
St + w-B and (0, po; v') 2 ((0,p:) +
S? V+3
As ((0,p:) 2 0, we
obtain the following bound:
v' + 3
Ep. (..
E p. l, (, (0.P).
(53)
3
Combining this with Eq. we have
v/ + 3
l0 pi. v) <
Ep, l. (0,0p.
c + B v min 2. Per. v),
c min Ep, l (0, p.. v)
(54)
which completes the proof.
E DETAILS FOR IMPLEMENTATION
In Algorithm[] D for the sake of simplicity, we omit two potentially time-consuming aspects: Jacobian regularization and the SVD for computing the subspace Tappe (x))). For the former issue, Jaco- bian regularization can be efficiently implemented using random projection (Hoffman et al.) 2019)) thereby circumventing the necessity training time. Regarding the latter issue, we utilize efficient randomized singular value decom- position (SVD) method (Halko et al. 2011). which is particularly well suited to our setting. This is because, under approach, the top singular vectors of the Jacobian matrix are highly concentrated as shown in Figure 2 [] which will benefit both the efficiency and effectiveness the algorithm. Further details are provided below.
computing the full Jacobian and only slightly increasing
27

---

Published as a conference paper at ICLR 2025
Randomized SVD method The randomized SVD algorithm approximates the range of Jg Rd.ed by projecting it onto lower-dimentt limensiimesiooall subspace using random Gaussian matrix 52 Idddd, where the dg is the target dimension. Let Y = JgSt E Rd,Xd, denote this projection. We compute the SVD of the smaller matrix Y, yielding Y = UEV. The approximate right singular vectors of gg given by V .VV. Please refer (Halko al. 2011) for more details.
Proposition 4 (Halko et al. 2011) Let Ig be de X d (d > de) matrix with singular values Ti 2 2 ... 2 Ode- We compute the approximate SVD by the randomized method described above, and setting the target dimension to dg 2d + Then the expected approximation error in
the Frobenius norm, |]g Jg VVV V.VVile [|., is bounded by:
1/2
v6 2
de
Jg Jg S
(55)
The term (2) o?) 1/2
represents the square root of the sum of the squared singular values beyond the d,-th component. The Jacobian matrix Jg. derived from g trained using our manifold- guided game approach, shows a rapid decay in its singular values, as demonstrated in Figures 2 and [ As result, the majority of the singular values are concentrated in the top few ones. This concentration allows for an efficient and accurate computation of the subspace T apppr (x, To) defined (Q. To see this more concretely, consider an image task with an input dimension 224 * 224 and an output dimension de 128 (typically used in contrastive learning). In this case, performing the singular value decomposition (SVD) on Jg requires approximately 0.7 GFLOPs (giga floating-poit point operations). For comparison, the cost of forward propagation in Resnet-18 is 1.8 GFLOPs.
=
Moreover, leverage advanced automatic differentiation tools available in the deep learning frame- work PyTorch (Paszke et al. 2019). Specifically, we avoid the explicit computation of Jg by using Jacobian-vector and vecobrnJacobian product operations (jvp and vip). These operations elimi- nate the need to store the full Jacobian matrix and further improve computational efficiency. Addi- tionally, we update the approximated tangent information every few iterations in Algorithm In Appendix F.5 we compare the running time with standard WDRO (which does not consider the data manifold) and other adversarial training algorithms. While our algorithm requires more time (around 1.5 times), it achieves better performance on accuracy.
F EXPERIMENTS
We conduct series of experiments across several scenarios to evaluate our proposed method. Specifically, we test it three typical distribution shift tasks: dealing with noisy data, attacked data and imbalanced data. All experiments ware implemented with PyTorch on single NVIDIA RTX 6000 Ada, and each instance is repeated by 5 times. In Section F.I. we present the dynamic update mechnism for the hyperparameter In Section F2 and F.3 show that method out- performs current WDRO methods and certain SOTA detense techniques for clean and contaminated data. In Section F4 conduct data augmentation based MWDRO address the over-fitting issue and improve generalization on limited data. In our experimental setup, we set the hyperparam- eters A 0.01 and λ2 in Eq. and Eq. respectively configure the algorithm, unless otherwise specified.
F.1 SELECTING STRATEGY FOR THE LAGRANGIAN PARAMETER IN EQ.
In this section, we briefly introduce the dynamic adjustment approach for the Lagrangian parameter in 6) proposed in Uni-DR (Bui et al. 2022). We have modified some notations for scenarios, and the readers may check the paper for more details on this method.
In Algorithm 3 let v, denote the value of iteration j. The parameter π) represents the learning rate for updating v, determining how quickly v adjusts in response to the perturbation discrepancy; 5 is the predefined maximum allowable perturbation radius in Eq. (2); and dg dg(gi qi) is the geodesic
28

---

Published as a conference paper at ICLR 2025
distance between the adversarial example and the original input qis which be approximated by the accumulated steps as shown in Lemma | Finally, Ei=1 dg dg(9; : qi) represents the average (population) perturbation cost over a batch of samples.
When computing perturbed samples, V is dynamically adjusted, with the adjustment magnitude controlled by the parameter 7/2- Intuitively, if the average geodesic distance between adversarial and benign samples is smaller than a specified threshold, v decreases in the next iteration to allow adversarial samples di to deviate further from the benign samples Ji, and vice versa.
Algorithm 3 Algorithm | Jusin dynamic v
Input: training set Par {x;; Mi ())) constant v, step size sequence {az > 0}
Output: the final parameter (N++
1: Train g(:) to minimize L (0) in E
2: f,, ...., do
3: For each sample Pt
Ptr. compute the corresponding adversarial point pt
Evolve(fpi): g(-), v;);
4: Updating V1+1 according (5.
5: gi+1 + 0 ay E Vo ((P) p.) X2 (La (pi) + (p:))((p)))
6: end for
The dynamic adjustment of is governed by the following update rule:
MJ+1 = - Tl (8 1 M de(a, 9.)
8
(56)
We still need to select an initial Po to setup this dynamic adjustment strategy. As illustrated in (Bui et al. 2022), the performance of WDOR is relatively insensitive the choice of vo when the dynamic strategy is employed. Following the recommendation in (Bui et al. 2022 we set 1// = 0.5 across all experiments.
F.2 EVALUATION ON NOISY DATA
To assess the robustness of our model against natural corruptions, introduce perturbations to each original image using Gaussian white noise. The resulting corrupted image is denoted as x' = 6, N (0,0 T) We adjust the noise levels with different values for anoise-
Datasets. CIFAR-10 and CIFAR-100 Krizhevsky et al. 2012) are two popular datasets for image classification, each of which consists of 60000 color images of 32 32 pixels, the former one divided into 10 classes, with 6000 images per class and the latter one contains 100 classes, with 600 images per class. Tmny-ImageNet-222
(Krizhevsky al. (017) is subset of the ImageNet dataset,
comprising 100000 images across 200 classes, each with 500 images. The images are resized to 64 64 pixels and include both RGB and grayscale channels. BuS (Mo al., (223) is medical image dataset that contains 1,119 malignant and 886 benign benign breast ultrasound images. We employ this medical image dataset as noise in medical images is a critical issue.
Compared methods. We compare our proposed MWDRO with four baselines: DropOut+L2 (He jet al. (0166) as the standard method for image classification, a practical implementation of WDRO proposed by (Sinaa et al. 2018), and the KL-DRO by (Levy et al. 2020). The randomized smooth- ing method RS proposed by (Cohen et al. 2019).
Figure 5 demonstrates that our MWDRO method significantly enhances robustness compared to other methods, particularly as the noise level increases.
29

---

Published as a conference paper at ICLR 2025
CIFAR-10
CIFAR-100
0.9
MWDRO
MWDRR
0.8
DropOut+12
DropOut412
WDRO
WDRO
0.6
KL-DRO
KL-DRO
RS
RS
0.4
0.4
0.3
0.3
0.2
0.1
0.0
0.2
0.4
0.6
0.8
0.0
0.1
0.2
0.3
0.4
Onoise
Onaise
Tiny-ImageNet 200
BuS
MWDDR
MWDRO
0.7
DropOut+L2
DropOut
Accuracy
0.6
WDRO
Accuracy
WDRD
KL-DRO
0.7
0.4
RS
RS
0.3
0.6
0.1
0.5
0.0
Onoise
Onaise
(c)
(d)
Figure 5: Evaluation of robustness against Gaussian white noise perturbations. The x-axis represents the variance of the additive Gaussian noise, ranging from 0 to I (the input data has been normalized). The results demonstrate that our MWDRO method significantly improves robustness across datasets.
F.3 EVALUATION ON ADVERSARIAL ATTACK
In this section, we evaluate the robustness of our proposed method. MWDRO, against adversarial attacks. We conduct experiments ImageNet-200. These datasets are standard benchmarks for assessing the effectiveness of adversarial defenses.
widely-used datasets: CIFAR-I0, CIFAR-I00 and Tiny-
We utilize variety of adversarial attack methods to evaluate the performance of our model. The attacks include Fast Gradient Sign Method (FGSM) (Madry et al. (2018), Projected Gradient Descent (PGD) (Kurakin et al., 2018), AutoAttack (AA) (Croce & Hein 2020), the Jacobian-based Saliency
Map Attack (JSMA) (Brendel et al. 2019) and C&W (L2) by (Carlini & Wagner 2017). Complete
results presented in Table 2 Table3 and Table
We compare our proposed MWDRO with four baselines: a practical implementation of WDRO proposed by (Sinha et al.) 2018), unified WDRO framework Uni-DR (Bui et al. 2022) and the popular adversarial training methods PGD-AT (Madry et al. 2018) and TRADES (Zhang et al,, 2019). For both PGD-AT and TRADES algorithms, match their corresponding attack scenarios as shown in Table 3 and And their adversarial attacks executed with 40 iterations and step size of 0.005.
the defense perturbation bounds to
Tabbe 2: Comparisons of clean and adversarial accuracies against different attacks for the Tiny- ImageNet-200. Best scores are highlighted in boldface.
MWDRO (8 = 0.01)
MWDRO 0.005)
Tmny-mmageNet-200
WDRO
TRADES
PGD-AT
Clean accuracy
28.47
39.48
40.23
43.82
55.42
PGD (s 0.01)
8.37
14.83
15.26
17.62
15.18
AA (E 0.01)
6.13
11.23
12.38
13.48
10.73
C&W 0.01)
9.81
12.46
16.32
17.26
13.34
30

---

Published as a conference paper at ICLR 2025
Table 3: Comparisons of clean data accuracy and adversarial accuracies against different attacks for the CIFAR-10 dataset. The table displays the performance of four different baselines and our proposed method MWDRO. Each method is evaluated on clean data and under various adversarial attacks (FGSM, PGD, JSMA, AA, and C&W () with different perturbation strengths. Best scores are highlighted in boldface.
Uni-DRO (6 = 0.03)
MWDRO (6 = 0.03)
MWDRO (Ø = 0.06)
CIFAR-10
WDRO
TRADES
PGD-AT
Clean accuracy
83.10
87.49
85.48
85.41
88.45
85.71
FGSM (s 0.03)
57.64
73.18
73.35
72.44
77.51
76.92
FGSM (s 0.06)
48.38
53.33
52.65
51.58
68.28
68.11
PGD (8 0.03)
39.37
45.40
47.65
44.81
48.28
49.32
PGD (E = 0.06)
12.38
19.59
20.02
14.23
23.58
27.65
ISMA (s 0.03)
54.29
74.51
74.39
73.85
79.56
74.53
JSMA (E 0.06)
40.01
66.27
67.41
64.25
74.47
73.90
AA (E 0.03)
41.38
45,74
45,85
43.05
47.73
45.28
AA (8 0.06)
11.11
14.78
14.02
13.69
14.34
15.89
C&W (A 0.01)
46.38
54.74
55.85
52.05
57.73
55.91
Table 4: Comparisons of clean data accuracy and adversarial accuracies against different attacks for the CIFAR-100 dataset. The table displays the performance of four baselines and our proposed method MWDRO. Each method is evaluated on clean data and under various adversarial attacks (FGSM, PGD, JSMA, AA, and C& W) with different perturbation strengths. Best scores are high- lighted in boldface.
Uni-DRO (Ó = 0.01)
MWDRO (6 = 0.01)
MWDRO (Ō = 0.02)
CIFAR-100
WDRO
TRADES
PGD-AT
Clean accuracy
66.28
67.94
67.67
67.64
69.58
67.25
FGSM (s 0.01)
38.71
48.05
49.54
45.97
49.96
48.41
FGSM (= 0.02)
20.26
35.82
36.47
29.33
35.97
37.01
PGD (s 0.01)
41.36
44.55
45.59
42.64
46.42
46.21
PGD (s = 0.02)
17.95
31.49
32.22
22.24
33.71
32.10
ISMA (E = 0.01)
34.15
46.51
47.71
45.37
49.24
38.91
JSMA (E 0.02)
28.27
36,84
39.53
33.30
45.72
33.64
AA (E 0.01)
20.38
44.74
44.85
40.32
47.52
41.74
AA (E 0.02)
9.69
14.76
15.01
14.69
15.14
15.87
C&W (A 0.01)
20.38
28.74
29.85
28.05
28.93
26.44
F.4 EVALUATION ON LONG-TAILED DATA
In this section, we evaluate the effectiveness of our proposed method, MWDRO, in handling long- tailed data distributions. Long-taile datasets present challenge due the significant imbalance between the frequency of common and rare classes. We conduct our experiments on two long- tailed benchmark datasets: CIFAR-10-LT and CIFAR-I00-LT. The imbalance in these datasets quantified by the imbalance factor (IF), which represents the ratio of the number of samples in the most frequent class to those in the least frequent class. We test our method on three versions of these datasets with imbalance factors of 10, 50, and 100. compare MWDRO with several baseline methods designed to address long-tailes distributions: Decouple (Hong et al.
2021): Disentangle
the label distribution from the model prediction; Reweight (Cui et al. 2019): Reweight the loss;
Resample (Cui et al. 2019): Over- sampling minor classes; Focal Loss (Lin et al.
2017): Weighting
difficult samples; SSI (Khosla et al. 2020 representation learning method: DRO-LT (Smuell
31

---

Published as a conference paper at ICLR 2025
& Chechik 2021): A DRO method designed for long-tail learning. We chose three long-tailed versions with imbalance factors (IF) of 10, 50, and 100 for training and the results shows that in highly imbalanced dataset, our MWDRO method combing with basic seamplingrr-weighhee strategy achieve best accuracy.
Table 5: Accuracies of ResNet32 on long-tailed CIFAR-IO and CIFAR- 100 datasets. The table shows the performance of various methods across datasets with different imbalance factors (10%, 50%, and 100%). Baseline methods include Decouple, Reweight, Resample, Focal Loss, DRO- LT, and SSL. Our proposed method, MWDRO, and its combination with re-sapling/re-wwighinnng strategies are highlighted. The best accuracy for each dataset and imbalance factor combination is highlighted in boldface.
CIFAR-1O-LT
CIFAR- 100-LT
Dataset
100%
50%
10%
100%
50%
10%
Decouple Hong et al. 2021
70.4
76.2
86.4
41.2
46.8
57.9
Focal Loss LLin et al. 2017
70.3
76.7
86.6
38.4
44.3
55.7
DRO-LT Samuel Chechik,
73.7
77.2
86.9
45.4
55.3
61.2
SSL Khosla al. 2020
67.3
75.4
86.5
37.5
44.0
56.7
Reweight Cui et al.,
2019
70.5
74.8
86.4
34.0
43.9
57.1
+MWDRO
74.5
76.9
86.7
46.7
56.2
60.5
Resample Cui et al., 2019
66.5
74.8
86.4
33.4
43.9
55.1
+MWDRO
72.8
77.5
86.2
43.5
52.6
58.2
F.5 TIME COMPARISON
In Algorithm D and Algorithm 2 for the sake of simplicity, we omit two potentially cime-consuming aspects: Jacobian regularization and the SVD for computing the subspace T Tappr (x, 7o). Details on efficient implementation are provided in Appendix H Table 6 presents the experimental results on running times, normalized relative the running time of WDRO. The experimental environments are consistent with those in Appendix F3
Table 6: Comparison of normalized running time for different methods across datasets. We train each method with 200 epochs. All running times normalized by dividing them by the running time of WDRO.
Dataset
WDRO
TRADES
Uni-DRO
PGD-AT
MWDRO
CIFAR-10
1.0
0.826
1.481
1.264
1.471
CIFAR-100
1.0
0.947
1.489
1.172
1.434
1.0
1.191
1.148
1.460
1.521
BuS
1.0
1.071
1.364
1.169
1.521
In the above Table 6 we compare computational running times across various methods using 200 epochs, following standard practices in prior work such as (Bui et al.] 2022). As expected, our method requires longer computational time than WDRO because it maintains similar workflow while incorporating additional data manifold information. As shown in Table our method requires approximately 50% more computation time per epoch compared classical WDRO.
However, our method exhibits faster convergence - natural consequence of working with a more focused uncertainty set (a subset of WDRO's uncertainty set). Leveraging this characteristic, implemented early stopping, where training halts when the robust performance on the validation set shows no improvement for consecutive epochs.
32

---

Published as a conference paper at ICLR 2025
Tables 7 and present the results across various datasets. With early stopping enabled, method achieves comparable total training times while delivering superior accuracy and robustness, demon- strating effective balance between computational efficiency and model performance.
Table 7: Comparison of clean accuracy and adversarial robustness across different methods var- ious datasets. "Clean" represents accuracy on the original test set, while "PGD" shows accuracy under adversarial attack. denotes the perturbation bound for PGD attacks.
Dataset
Evaluation
WDRO
TRADES
Uni-DRO
PGD-AT
MWDRO
Clean
80.42
86.79
85.51
82.18
88.21
CIFAR-10
PGD (e 0.03)
41.58
47.35
49.64
45.52
52.27
Clean
64.38
67.32
66.92
41.47
68.17
CIFAR-100
PGD (c 0.01)
41.79
45.72
47.20
43.91
49.54
Clean
30.09
39.11
39.37
36.87
41.36
Tiny-ImageNet-2000
PGD (c = 0.01)
7.41
14.67
14.92
14.46
14.80
Clean
74.71
77.92
78.29
75.92
81.74
BuS
PGD (e 0.01)
58.28
66.73
68.53
63.99
69.48
Table 8: Comparison of normalized running time for different methods across datasets. Early stop- ping criteria were applied consistently across all methods. All running times are normalized by dividing them by the running time of WDRO.
Dataset
WDRO
TRADES
Uni-DRO
PGD-AT
MWDRO
CIFAR-10
1.0
0.994
1.070
1.126
1.115
CIFAR- 100
1.0
1.15
1.139
1.142
1.108
mmny-lmageNet2200
1.0
1.098
1.121
1.174
1.123
BuS
1.0
1.071
1.094
1.217
1.118
F.6
OTHER ABLATION STUDIES FOR THE MWDRO.
We conduct the ablation studies to analyze the effectiveness of the manifold-guided game and the MWDRO method. As shown in Table models trained using the manifold-guides game without DRO (i.c., regularized only by Eq. demonstrate significant vulnerability to adversarial attacks.
Table 9: We report the Top-1 accuracy on various datasets for both the MWDRO model and models trained without the DRO method, which are regularized solely by the manifold-guided game in Eq. The evaluation includes three types of attacks: PGD (c 0.01), AA (c 0.01). and C&W (A 0.01) except for CIFAR-10, except for CIFAR- 10, where 0.03 used for both PGD and AA attacks.
Cifar10
Cifar
Tiny liny-ImageNet-200
PGD
AA
C&W
PGD
AA
C&W
PGD
AA
C&W
Only with (q.
43.29
42.82
48.16
24.71
28.69
20.17
7.07
7.83
10.14
MWDRO
48.28
47.73
57.73
46.42
47.52
28.93
15.47
11.64
14.73
F.7 OTHER EXPERIMENTAL RESULTS
In the first stage of Algorithm the encoder g(-) is trained for No epochs following the objective of the manifold-guides game, defined in (q. This training phase is crucial for obtaining sufficiently accurate tangent information for each data sample x. This stage ensures g(-) can yield sufficiently good approximate tangent T Tappe Tn) for each data sample x, which is later used
33

---

Published as a conference paper at ICLR 2025
for the gradient ascent to solve the surrogate loss in Eq. Figure illustrates the impact of the hyperparameter No on the accuracy our model across various datasets.
Effect of No
Effect of No.
Robust Accuracy
1.0
Clean Accuracy
Tinytnage
Tryyimage
0.8-
0.8
0.6-
0.6
0.4
0.2
20
40
60
0
20
40
60
Start Epoch No
Start Epoch No
(a) Accuracy PGD (e 0.01) attacked datasets
(b) Accuracy clean data
Figure 6:
The effects of the hyperparameter No on model accuracy across different datasets. We evaluate the model's performance CIFAR-IO, CIFAR-100, Tinylmage-200, and BuS datasets, varying the starting epochs No for the second stage of Algorithm [ Figure [6a] and [6b are the achieved accuracy both attacked (PGD, 0.01) and clean datasets, respectively,
Effect of No
0.9
0.8
Accuracy
0.7
0.6
0.5
0.4
0.3
CIFAR-100
Tiny ImageNet
0.2
BuS
0.1
40
Start Epoch No
Figure 7: The effects of the hyperparameter No on accuracy across different datasets. We evaluate the model's performance CIFAR-10, CIFAR-100, Tinylmage-200, and BuS datasets with various starting epochs No for the second stage of Algorithm Gaussian noise with a variance of x² = 0.2 added to all datasets test the robustness.
Figure shows the cumulative variance percentage of the Jacobian for various models trained Tiny-mmageNet-21 and BuS datasets. Our model exhibits the most rapid initial increase in cu- mulative variance percentage, indicating that fewer singular vectors are needed approximate the tangent space effectively.
34

---

Published as a conference paper at ICLR 2025
Cumulative Percentage
CIFAR-10
Cumulative Percentage
CIFAR-100
1.0
1.0
0.8
0.8
0.6
Jac+CL(A==0,
Jac+
0.4
JaC+CL(A;00.00
0.4
=ac+AE()==0.)
Jac+AE(A;=0.1)
0.2
Drop0ut+L2
0.2
DropOut+L2
CL
0.0
0.0
60
100
120
20
160
120
Number of Singular Values
Number of Singular Values
(a)
(b)
Tiny ImageNet-200
Cumulative Percentage
BuS
Cumulative Percentage
1.0
1.0
0.8
0.8
0.6
Jac+ -c<+C[(;=0011
Jac+ =
0.4
Jac+ BCC+C(A==0,01
0.4
Jac+ 0000)
Jac -0
AEC)AE(A)-
0.2
DropOut+L22
0.2
DropOut+L2
CL
40
60
100
20
60
80
100
120
Number of Singular Values
Number of Singular Values
(c)
(d)
Figure 8: Cumulative percentage of variance of the Jacobian for different models trained on Tiny- ImageNet-200 and BuS datasets using ResNet-18 18, respectively. The cumulative percentage of vari- ance for first singular values is calculated as described in (8). The blue line represents the Auto-Encoder model with Jacobian regularization (Jac+AE), the green line represents the CL model with Dropout and L2 regularizer (DropOut+L.2), the yellow line represents the CL model without Jacobian regularization (CL).
Σ )
--11 [x The red line represents our model as
G VISUALIZATIONS
This section presents comprehensive visualizations demonstrating our manifold-guided approach's 's mechanisms and effectiveness. We showcase the method's ability to learn meaningful geometric structures through various experiments, including analysis of manifold dynamics toy datasets, comparison of learned tangent spaces across different models, and visualization of perturbation trajectories the MNIST dataset. Our results illustrate how the proposed approach successfully constrains perturbations to follow the intrinsic data manifold structure.
Figure provides schematic illustration of our methodology. The encoder g(æ;) should change most when the sample æ moves along the semantic variant directions within the data manifold; while changing minimally when moves off the manifold (i.e., M-), or towards the Semantic-Invarian region (i.e.. (x)( This results in distinctive "spiked" distribution of the Jacobian matrix's s singular values (as shown in Figure where the top singular vectors naturally correspond to the semantic variant directions in the tangent.
Figure T0 illustrats the procedure for obtaining the perturbed samples described in Algorithm 2 for 2D toy manifold dataset. The blue dots are generated based on a quadratic function x2 x?. Specifically, we sample points along X2 -x and add Gaussian noise with variance 8 0.01. To train the encoder g, employ eration is simplified to adding small noise. The tangent at any given point x₃ is approximated using the largest singular vector of the Jacobian J ((x1)x))) For further simplification, point has a gradient direction of (0, 1) when optimizing the surrogate loss in Eq. equation] 7 imply- ing that perturbation directions for all points are the unit vector (0, 1) depicted in Figure 10(b).
manifold-guides game framework, where the augmentation op-
assume each
35

---

Published as a conference paper at ICLR 2025
Data manifold
Demantic-Invarian space
Train samples
CL Augmentation
Semantic-Varian direction
Direction to leave manifold
Figure 9: This picture depicts the mechanism of our proposed manifold-guided game.
We select (S1 to S7 marked in Figure 10 points from X2 = -x, as the starting points and observe the trajectory with without the tangent constraint. Figure [0'a) illustrates that manifold- guided game extracts local tangent information, effectively constraining points to move approxi- mately along the manifold rather than other directions. Figure T0(b) shows that the samples move freely along the given gradient direction.
Figure | provides qualitative comparison of tangent spaces learned by different models: (i) dis- plays original image from the CIFAR- 10 dataset. The subsequent rows show extracted tangent vec- tors from three models: (ii) the standard ResNet- (iii) model trained with smoothness penalty term (0.==) manifold guided game. This comparison highlights that our method more effectively captures mean- ingful directions of variation within the data manifold, consistent with the analysis in Section 4.2
c) f(x)11 (A takes 0.01), and (iv) a model trained with our proposed
and Appendix
Figure [Ojillsstates the perturbation trajectories the MNIST dataset. (a) depicts an original sam- ple from the MNIST dataset. (b) and (c) depict perturbed samples after 50 iterations, generated using the LeNet model trained with standard training and our manifold-guide guided game, respectively. (d) provides t-SNE visualization (perplexity = 20) MNIST digits, along with the perturbed sam- ples from each iteration. Different colors represent distinct digit classes (0-9), while the green and black lines illustrate the perturbation trajectories. While the t-SNE does not perfectly represent the underlying geometric structure of the dataset, it offers valuable insights. Specifically, perturbations without manifold constraints can push data points to meaningless regions. In contrast, our manifold- guided game effectively restricts the exploration space, resulting in smaller and more meaningful uncertainty set for Wasserstein Distributionally Robust Optimization (WDRO). This demonstrates that our approach enhances the robustness of the model by aligning perturbations with the intrinsic data geometry.
36

---

Published as a conference paper at ICLR 2025
Trajectories w/ Manifold Constrint
Training Data Trajectory
5, 5j:(-4, - 16)
Trajectory 2
52:(-3,
Trajectory 3
5-:(-1,
Trajectory 4
54:10.07
Trajectory
X2
5;:(1. - 1)
Trajectory
5s:13. 9)
Trajectory 7 5,:(4.
X1
(a)
Trajectories w/o Manifold Constrint
Training Trajectory
5y: 16)
Trajectory
$2:(-3, 9)
Trajectory
(-:(-1, -
Trajectory
So:
Trajectory
5s: - 1)
X2
Trajectory
$6: -9)
Trajectory 5:: -
X1
(b)
Figure 10: Comparison of trajectory optimization with and without manifold constraints. We apply manifold-guided game leam the underlying manifold structure, then simulate the Algorithm for solving the perturbed samples with or without the manifold constraint. (a) shows trajectories constrained follow the learned manifold structure. (b) shows that perturbation directions for all points are simplified as the unit vector (0, 1). Here S, are the starting points.
37

---

Published as a conference paper at ICLR 2025
(i)
(10)
(M)
(a)
(1)
(i))
(ii)
(iv)
(b)
Figure 11: Visualization of the tangents learned by different ResNet-18 models. (i) displays original images from the CIFAR-10 dataset. The subsequent rows show extracted tangent vectors from three models: (ii) the standard ResNet-18; (iii) a model trained with a smoothness penalty term (1E=~A(0.02)) [x fx) I|2 (A takes 0.01); and (iv) model trained with our proposed manifold guided game. These visualizations reveal the models' learned local geometric structure.
38

---

Published as a conference paper at ICLR 2025
5
5
(a)
(b)
(c)
8
trajectory trajectory
-100
-100
-50
50 100
t-SNE dimension
(d)
Figure 12: Visualization of perturbation trajectories the MNIST dataset. (a) An original sample from the MNIST dataset. (b) and (c) show the corresponding perturbed samples after 50 iterations, generated using LeNet model trained with standard methods and our manifold-guide spectively. (d) A t-SNE visualization (perplexity 20) of MNIST digits, along with the perturbed samples mentioned above each iteration. Different colors represent distinct digit classes (0-9), while the green and black lines illustrate the perturbation trajectories.
game, re-
39


Reviews:
[{"review_id": "b69d5F2QFS", "reviewer": "Reviewer_v52n", "summary": "The authors propose a manifold-based DRO method that incorporates the geometric structure of training data to construct the uncertainty set. By integrating contrastive learning with Jacobian regularization, the method effectively captures the data manifold, providing theoretical guarantees for robustness. Experimental results on benchmark datasets demonstrate improved accuracy and robustness compared to conventional DRO methods.", "questions": "See Weaknesses above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed manifold-based DRO method effectively captures the geometric structure of training data, leading to better robustness and accuracy compared to traditional DRO methods. This is achieved by integrating contrastive learning with Jacobian regularization, which helps in maintaining the data manifold structure.\n2. The method provides theoretical guarantees for robustness by utilizing a novel approach to approximate geodesic distances on manifolds. This ensures that the model remains reliable even when faced with distributional shifts in the data.\n3. Despite the complexity of incorporating manifold constraints, the proposed method is easy to implement and computationally feasible. The authors demonstrate this through experiments on popular benchmark datasets, showing that the method achieves superior performance without significant increases in computational cost.", "weaknesses": "I believe the motivation of this paper is reasonable and the method is effective. However, I would like to ask the authors whether the effectiveness of the algorithm depends on specific data? What kind of data characteristics would give MWDRO a greater advantage?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gtRB3Vh1jc", "reviewer": "Reviewer_qtHQ", "summary": "This paper proposes performing distributional robustness (DR) over data manifolds. It develops the dual form of DR over data manifolds and    uses the top singular principal vectors of the Jacobian matrix to characterize the tangent space of the data manifolds.", "questions": "Please address my questions in Weaknesses.", "limitations": "", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The idea of developing distributional robust on data manifolds is intuitive and interesting.\n- The paper is well-written.", "weaknesses": "- It is unclear why we can use the linear span of top singular vectors of the Jacobian matrix of g can characterize the tangent space of data manifolds. The current explanations do not really convince me. Moreover, the experiments in Figure 2 only demonstrate that there are some dominant singular vectors and I cannot see how it explains why $T_{appr}(x, \\tau_0)$ is a subspace of the tangent space.\n- It is also unclear to me why the game in Section 4.2 helps extract tangent information. Evidently, the CL loss encourages the features of data examples in $M_{SI}(x)$ to be close, whereas the Jacobian regularization suppresses the feature variations for perturbations of x across all directions. They seem to share the same purpose and nature.\n- It is unclear to me why we can use the $pt^t$ in Alg 2 to approximate the geodesic distance from $q^t$ to $q^0$. \n- It is unclear how to do Exp operator to retract to the data manifolds. This seems impossible because we only assume that data lie in manifolds but we do not have anymore information of these manifolds. The Exp operator appears closed-form just for some simple and specific manifolds.\n-  The proposed method is very computationally expensive, even more expensive than adversarial training. This is because we need to compute a trajectory for each data example which requires computing Jacobian matrix and also the gradients to data examples. The analysis on the computational complexity is necessary.\n- The experiments do not demonstrate the benefit of performing DR on data manifolds. It would be better if the authors keep the same CL loss and Jacobian regularization, while replacing their manifold DR by standard DR on data space for a comparison. Moreover, the authors should provide the visualization of $q^1$,...,$q^t$ on the trajectory to see if we can make them on the data manifolds.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "j9qpizWLYn", "reviewer": "Reviewer_Q211", "summary": "This paper introduces a manifold-based Wasserstein Distributionally Robust Optimization (WDRO) method aimed at improving the robustness of deep learning models. Furthermore, to tackle the challenges posed by the data manifold structure, the authors propose a game that integrates contrastive learning with Jacobian regularization.", "questions": "Q1. In Figure 3, the accuracy of each model is already distinguishable even when the noise size is 0. Could you provide further discussion on this observation?\nQ2. Is the use of contrastive learning essential in the design of the manifold-guided game? Are there alternative optimization objectives that could be utilized instead?\nQ3. As noted in W3, could you further analyze which phase of training contributes to the increased time required for MWDRO?\nQ4. Is it possible to further optimize the design of the algorithm to reduce the time to loss? For example, by utilizing simplified geometric computations or low-rank approximations.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "S1. The paper is well-organized and presents its ideas clearly.\nS2. The proposed methods demonstrate superior performance compared to state-of-the-art algorithms.\nS3. The design of the manifold-guided game is a novel concept, leveraging neural networks to encode manifold information, particularly given the absence of a closed-form representation for the data manifold.", "weaknesses": "W1. The motivation for this work could be strengthened. In the Introduction, the authors mention that selecting an appropriate threshold for uncertainty is challenging and introduce a new uncertainty set by assuming that data is supported on a manifold. However, constraining the uncertainty set to manifolds also raises the same issue of threshold selection.\nW2. The explanation of how the game aids in extracting tangent information is not intuitive enough. Including a figure to illustrate this concept would be beneficial.\nW3. As the experimental results show, MWDRO takes more time compared to other algorithms, and it would be desirable to have a further discussion on this part of the reason.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Fa7hPDJQtd", "reviewer": "Reviewer_P9ds", "summary": "This paper proposed a manifold-based distributionally robust optimization method to promote the robustness of existing deep learning models. Specifically, it designed a game that trades off between CL and Jacobian regularization to solve the DRO problem constrained by the data manifold. Both theoretical and empirical results show the robustness of the proposed method.", "questions": "1. Could you please discuss the effect of $\\lambda$ and how to select the optimal value?\n2. Could you verify the effectiveness of your proposed method on different model architectures such as ConvNext or Swin?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The proposed method is simple and effective. \n2. This paper provides a comprehensive analysis of their proposed method.", "weaknesses": "1. This paper should be further polished and re-organized. The introduction and methodology section should be more concise, and some contents should be moved to a more relevant part. I provide the details as follows: 1) the introduction of WDRO from line 49 to 71 can be more brief. 2) I can understand that the authors discussed the most relevant literature in Section 1.1, but the structure is a little strange. I suggest considering this part as a discussion part such as \"Relation to xxx\".  3) In the methodology, I also suggest the authors provide the literature or experiment support for the sentence that \"we should emphasize that the learned representation for manifold from neural networks is not sufficient for extracting tangent space\". 4) Eq (2) as a significant part of the methodology is mentioned in Section 1, which makes the methodology separate. I suggest the comparison between WDRO and MWDRO in the introduction being high-level, while its details can be put into Section 3 as the motivation. \n2. Some significant ablation studies are missing. For example, the effect of $\\lambda$ and model architectures has not been discussed. I suggest the authors test a range of values of $\\lambda$ to test its effect. As for the model architectures, the authors only provided ResNet18, which is insufficient. I suggest the authors try some state-of-the-art models such as ConvNext and Swin.\n3.  The authors should clarify the advantages of the proposed method compared with [1] in Related Works and Experiments.\n\n[1] Liu J, Wu J, Li B, et al. Distributionally robust optimization with data geometry[J]. Advances in neural information processing systems, 2022, 35: 33689-33701.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Xz3SE82kad", "reviewer": "Reviewer_m5cH", "summary": "This paper presents a novel manifold-based Distributionally Robust Optimization (DRO) method aimed at enhancing the robustness of deep learning models against distributional shifts. By integrating contrastive learning with Jacobian regularization, the proposed approach captures the manifold structure of the training data to construct a more accurate uncertainty set, thereby improving classification performance under various distributional changes.", "questions": "Please see the comments in Weakness.\n\nI would be happy to raise my score if these issues are well addressed.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1.\tWriting: This work is presented with good writing style, where the summarized problems with detailed explanations make it easy for readers to understand the problem addressed in this article. However, there exist minor spelling and grammatical errors. The example in Fig.1 helps to understand the connections between geometric representation and semantic tasks.\n\n2.\tNovelty: It incorporates geometric structure into the construction of the uncertainty set, potentially leading to more realistic distributional assumptions. Contrastive learning with Jacobian regularization to encode manifold information is novel and seems to consume more geometric information. \n\n3.\tTheory: Theorems 1&2 based on geometric scheme provides guarantees on approximation and gradient estimation, offering strong dual reformulations and approximation techniques for geodesic distances, which are crucial for establishing the method's reliability. \n\nThe proof seems solid but I have not carefully checked the whole Appendix. \n\n4.\tExperiments: Several scenarios and recent baselines are considered, implying improvements in accuracy and robustness under various distributional shifts.", "weaknesses": "1.\tApproximation. There exists a gap between optimal solutions from original and dual problems. I’m not sure if Theorem 1 states the distances between the original and dual objects? \n\n2.\tSelection on v. The bound derived in Theorem 1 heavily relies on v. I would appreciate it if the authors could give more detailed illustrations on the “dynamic mechanism” on setting v empirically?\n\n3.\tBroader baselines and empirical settings. \nFor example, the settings for “Noisy Data” are kind of simple. What’s the variance of the added Gaussian white noise? It is suggested to follow the empirical settings in [1] to widen the difference between the training and testing sets, see Table 1 in [1]. \n\nMore DRO approaches [1-3] for learning from noisy data are suggested to be included.\nMoreover, it is kindly suggested to add introductions to these baselines in the supplementary file for better readability.\n\n[1] https://arxiv.org/abs/2305.14690\n[2] https://arxiv.org/abs/1902.07379\n[3] https://arxiv.org/abs/2006.04662\n\n4.\tIllustrations on manifolds. The paper could benefit from a more in-depth analysis of the algorithm's scalability, especially regarding its performance with large datasets or specific manifold structures. \nSynthetic experiments on toy examples with already known geometric structure, e.g., the Swiss or Torus, could help to visualize and estimate the investigated manifolds.\n\n\n5.\tComputational Cost. While the theoretical underpinnings are well-developed, the paper may not provide a comprehensive assessment of the computational efficiency and practicality of the proposed method in real-world applications. Like the computational complexity analysis or empirical time/memory cost.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)