Bootstrap Your Uncertainty: Adaptive Robust Classification Driven by Optimal-Transport
Authors: Jiawei Huang
Minming Li
Hu Ding
Venue: NeurIPS 2025 poster ()
Keywords: distributionally robust optimization
optimal transport
Primary area: optimization

Abstract: Deep learning models often struggle with distribution shifts between training and deployment environments. Distributionally Robust Optimization (DRO) offers a promising framework by optimizing worst-case performance over a set of candidate distributions, which is called as the \emph{uncertainty set}. However, the efficacy of DRO heavily depends on the design of uncertainty set, and existing methods often perform suboptimally due to inappropriate and inflexible uncertainty sets. In this work, we first propose a novel perspective that casts entropy-regularized Wasserstein DRO as a dynamic process of distributional exploration and semantic alignment, both driven by optimal transport (OT). This unified viewpoint yields two key new techniques: \emph{semantic calibration}, which bootstraps semantically meaningful transport costs via inverse OT, and \emph{adaptive refinement}, which adjusts uncertainty set using OT-driven feedback. Together, these components form an exploration-and-feedback system, where the transport costs and uncertainty set evolve jointly during training, enabling the model to better adapt to potential distribution shifts. Moreover, we provide an in-depth analysis  on this adaptive process and prove the theoretical convergence guarantee. Finally, we present our experimental results across diverse distribution shift scenarios, which demonstrate that our approach significantly outperforms existing methods, achieving state-of-the-art robustness.

Bootstrap Your Uncertainty: Adaptive Robust Classification Driven by Optimal-Transpon
Jiawei Huang!. Minming Li
Hu Ding
School of Computer Science and Technology. University of Science and Technology of China DDpaatmentt of Computer Science, City University of Hong Kong {hjw0330, huding) &mail. ustc. edu. cn, minming www.ciit..eedu.... edu.
Abstract
Deep learning models often struggle with distribution shifts between training and deployment environments. Distributionally Robust Optimization (DRO) offers a promising framework by optimizing worst-case performance over a set of candidate distributions, referred to the uncertainty set. However, the efficacy of DRO heavily depends on the design of the uncertainty set, and existing methods often perform suboptimally due to an inappropriate or inflexible uncertainty set. this work, we first propose a novel perspective that casts entropy-regulame Wasserstein DRO as dynamic process of distributional exploration and semantic alignment, both driven by optimal transport (OT). This unified viewpoint yields two key new techniques: semantic calibration, which bootstraps semantically meaningful transport costs via inverse OT, and adaptive refinement, which adjusts uncertainty set using OT-driven feedback. Together, these components form an suploration-and-feeddaad system, where the transport costs and uncertainty set evolve jointly during training, enabling the model to better adapt to potential distribution shifts. Moreover, provide an in-depth analysis of this adaptive process and prove theoretical guarantees of convergence. Finally, we present our experimental results across diverse distribution shift scenarios, which demonstrate that our approach significantly outperforms existing methods, achieving state-of- the-art robustness.
Introduction
Distribution shifts between training and deployment environments present fundamental challenge machine learning systems, often leading to significant performance degradation in real-world applications. This challenge has motivated robust learning algorithms that ensure performance under distributional uncertainty. Among these approaches, distributionally robust optimization (DRO) [22 46) has emerged as popular framework that optimizes for worst-case performance across of potential test distributions, known the "uncertainty ser".. Unlike empirical risk minimization (ERM), which directly minimizes the expected loss {(0, x)) over the empirical training distribution Pur. DRO seeks robustness by solving the minimax problem: ming supQeu(P..) Eq (((0, x)). Here, ((.rr) represents the uncertainty set that encompasses distributions in the neighborhood of Pur. capturing a range of plausible test-time deviations.
critical challenge implementing effective DRO lies in designing an appropriate uncertainty set. Ideally, the uncertainty set should accurately capture potential distributional shifts while remaining amenable efficient optimization. Optimal transport (OT) 61 provides an elegant mathematical framework for this purpose, offering both strong theoretical foundations and intuitive physical inter- pretations (see Sec. 2.1 These properties make OT distances suitable for constructing meaningful
"Corresponding author.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

uncertainty set in distribution space. In this work, we leverage the entropy-regularized OT distance (also known as, Sinkhorn distance [[6], and the formal definition is shown in Def. 2 to define our uncertainty set. Let De denote the Sinkhorn distance and the DRO problem 62 is formalized as:
min sup Bo( where U(Pu. {Q E P(XX)) D. Q) 53. Q..(P...))
(I)
Here 8 >> O serves the radius to control the size of allowable distributional shift; P(X'X)) denotes the set of Borel probability distributions over the data-label space "X' x ) The uncertainty set reflects prior knowledge about potential distributional shifts. As we will demonstrate in Sec. 2.2 the Sinkhorn distance De (Pr.. Q) introduces a reference distribution 1/ that guides the uncertainty set's shape. This reference distribution determines the "closeness" between distributions and the training data Ptr, thereby serving as prior knowledge about the "geometry" of the uncertainty set.
Despite its popularity, OT-based DRO (e.g., ① still faces major practical limitations 42 First, the uncertainty set can be overly conservative, encompassing unrealistic worst-case distributions that harm practical performance. Second, conventional transport cost is not very effective for encoding semantic information, because the input training data is often in lack of explicit expression for
semantic similarities in some practical areas such computer vision [30] For instance, pixel-wise
distance is a poor proxy for image-level similarity 42 As mentioned in above, the Sinkhorn distance
introduce a reference distribution v/ which can mitigate those issues to some extent. In particular, the
distribution y regularizes the transport plan and implicitly shapes the geometry of the uncertainty set (Pu.. 8) by controlling where and how probability mass can be transported [6]]. However, constructing an appropriate reference distribution v remains challenging. When training data is limited, v is typically synthesized through heuristics (e.g.. mixup [67] or noise injection (26)). But, this operations may introduce unrealistic distributions in (Ptt,. 5). Moreaver, we are also confronted with tricky dilemma for constructing the uncertainty set (P....
Expanding the uncertainly set to capture meaningful shifts risks incorporating unrealistic perturbations, while overly restricting it may exclude plausible distributional variations.
Our core idea is to design an adaptive DRO framework in which both the semantic distances and uncertainty set evolve during training, yiclding dynamic uncertainty set. While introducing this flexibility to the uncertainty set U(Pu.. is appealing. it also presents significant challenges for implementation and analysis. Our proposed framework is grounded in optimal transport, offering a principled geometric interpretation that guides the algorithm design and analysis. To our best knowledge, our proposed framework is fundamentally different with existing DRO methods, from both the formulation-level and algorithmic-design level (most of them rely on static uncertainty set). We name our approach "AdaDRO" and summarize our contributions follows,
We introduce novel theoretical perspective that frames cettropyrreulaaiizee DRO dynamic explorationann--liinnee process driven by OT. This perspective gives rise to two core components: semantic calibration, which bootstraps transport costs via inverse OT, and adaptive refinement, which filters the uncertainty set using OT- based feedback.
We establish theoretical convergence guarantees despite the evolving nature of the transport costs and uncertainty set, ensuring robust optimization stability throughout training. Further- more, the OT-based structure of (T) enables seamless integration of semantic calibration and refinement into the training process.
Extensive experiments demonstrate that our method improves both robustness and accuracy across range of distribution shift scenarios, thereby broadening the applicability of DRO.
To illustrate the benefit of our approach, consider the 2D binary classification task shown in Figure Due to limited training data, the learned decision boundary significantly deviates from the ideal boundary. The reference distribution V, constructed from Gaussian distributions around training points, may include samples that fall outside the true data manifold. Such low-quality samples can corrupt the uncertainty set and degrade classifier quality. By identifying and filtering out such samples, we obtain a boundary (blue curve) that better reflects the structure of the data distribution and more closely approximates the true decision boundary.

---

(a) ERM
(b) DRO and Ours
Figure 1: Binary classification "double moon" dataset. (a) SVM (RBF-kernel) decision boundary (red) trained on 20 samples (10 per class). (b) Comparison of Sinkhorn DRO (green) using Gaussian reference (small dots show 600 samples) versus our adaptive DRO (blue). Our method better captures the true data manifold by filtering low-confidence regions in v.
1.1 Other Related Works
Distributionally robust optimization (DRO) methods seek model parameters that minimize worst-case expected loss over an uncertainty set of distributions around the training data. Early work in this area used ő-divergence uncertainty set, yielding tractable reformulations via duality; for instance, KL-DRO optimizes against distributions within Bullbackk-eibblee ball around the empirical measure Wasserstein DRO (WDRO) instead constructs uncertainty set based on optimal transport (OT) costs, providing geometric robustness guarantees; seminal results include finite-sample bounds and dual formulations by Blanchet and Murthy ( [0] and extensions to high-dimensional settings by Gao and Kleywegt [24] Entropy-regularizer OT, or Sinkhorn DRO, introduces reference measure
29 46
and entropy penalty to yield smoother uncertainty set and more efficient algorithms 664 62
The full
discussion related works is deferred to Appendix A.1
2 Preliminaries
Consider the input-output space E X' V, where X' C Rd is the data space and y [K] (1..... . K} denotes the label space containing K distinct classes. Let P(E) represent the set of Borel probability distributions supported on E. For 1 < S. n, each sample p; (xx vi) is drawn from an underlying distribution P E PE)). The empirical distribution I,, 1 nn 5p, is induced by the dataset (Pi,..., Pn.). where dp. the Dirac measure concentrated at p.. We endow with transport cost function: C (pt, P2) C (3)) C* (yis, Uj), where C X X R and C² y y + R quantify the transport costs between features and labels, respectively. This transport cost provides the foundation for defining the optimal transportation distance in (4).
We define neural network f.) /ww go(-), where go(+): X - Z serves the encoder, parameterized by 0, which maps input data E X into a latent representation space Z c R&4, The function hw (w):)) Z y acts as linear classification head that generates predictions in the output space The classifier weight matrix W [w,,... wk] E kxd, where each w; the class-specific weight vector for the i-th class. The complete set of model parameters is θ = {0, W}. For notational simplicity, sometimes abbreviate the notation as f(-) h 9(-) and denote the classification loss function as ((-).
2.1
OT and Inverse-OT
Optimal Transport (OT) and Inverse Optimal Transport (IOT) offer complementary perspectives: OT seeks optimal couplings between distributions given cost function, while IOT infers appropriate cost functions that best explains observed coupling patterns.
Let P and Q be probability distributions, and IIP, Q) denote the set of all joint distributions (called transport plans or couplings) with marginals P and Q. Formally, any 7 e (PP,Q) satisfies ((p..)dg=PPl P(p) and dp. = Q(g). Let μ and v be reference measures such that P < p and
3

---

Q& V (i.e., P and Q are absolutely continuous with respect to μ and v, respectively). We begin with the entropy regularization.
Definition (Relative Entropy of Transport Plans). For any transport plan y y Q), the relative entropy of 7 w.r.t. product measure @ v is:
H(r μ @ v) - f (a(wsv) (p. g)) dr((p.g)
(2)
This entropy term requires supp( P) supp(u) and supp(Q) (upp(()) remain finite. A common choice 40 uses uniform reference measures (du(p)dvv(q) = const), simplifying the entropy to:
H (7) (mpa)>>> [log (d.(p.))) 1],
Optimal Transport Cuturi [66 introduces an entropic penalty to the classical OT problem, encouraging smoother couplings between distributions while minimizing transport cost. This leads to the formulation of the Sinkhorn distance:
Definition 2 (Sinkhorn Distance (akka. retroppy-reuularized OT Distance)). Given distributions P, Q. a cost function C' =x R. and regularization strength >> O, the Sinkhorn distance is:
Q) inf TEIT() (m-a)~
{
E LLBB,ICC (C(p.:))++H((lleeee)ee =H((luuv)
12 solves the inverse problem: inferring cost function that explains observed couplings. To capture this, we add the subscript 0 to the previously defined notations "y" and "C"", i.e., To and Cg. which denote the corresponding optimal transport plan and cost learned by the model parameterized by 9. In addition, we let 5 II(..Q) denote an observed coupling that reflects ground-trume semantic relationships (details in Sec 3.2). Then we have the following formulation for IOT:
Inverse Optimal Transport (IOT) 1OT 57
ming KT.(7110), s.t. 700 arg min E [Colo.9) + +H))) 7) (p.q)~>
(5)
Here, the KL divergence (L( quantifies how well C. explains the observed matching pattern 7. The entropy term (7)) is defined in This bilevel structure allows to learn transport costs Cg such that the resulting optimal transport plan aligns with observed empirical semantic couplings.
2.2 Sinkhorn DRO
We now analyze the Sinkhorn DRO formulation in ① highlighting its theoretical foundations and advantages over alternative approaches. Sinkhorn DRO extends Wasserstein DRO (WDRO) I 23 through entropy regularization. When the transport cost C' is metric and 0, the formulation (I) reduces to standard WDRO. However, WDRO typically yields discrete, adversarial worst-case distribution, making it well-suited for adversarial robustness but less effective for modeling natural distribution shifts, which often appear as continuous, structured deformations
[6].
The introduction of entropy regularization in D. brings benefits through two complementary mecha- nisms: Reference-gule guided uncertainty set: The relative entropy term H (T/μ @ v)) leverages reference distributions yu and to shape the transport plan, guiding it toward more plausible distribution shifts rather than arbitrary worst-case scenarios; Geometric Awareness: In contrast to divergence-bard based DRO (e.g., KI.-DRO (29). Sinkhorn DRO incorporates transport costs C'p, q) that reflects pairwise sample similarities, preserving crucial geometric information within the data space.
Although the Sinkhorn distance in 4 depends on two reference measures μ and v, following Wang al. [62], we set μ Pt.. the empirical training distribution. This choice is both practical and theoretically sound, as the relative entropy term differs only by constant when Ptr < leaving the optimization unaffected. In contrast, the reference measure plays critical role: it defines the support of feasible worst-case distributions by enforcing Q & p. Through dual analysis, Wang et al. (22 derived the closed-form expression for the worst-case distribution. When the optimal Lagrange multiplier X > 0, the worst-case distribution Q^ has the density:
(((())- (C(
dQ*(a) Ep- EE (e(((u)-AC(p.))) dv(g).
(6)

---

ref. V
Dynamic Filter
Nadeldertred Couping
Groundtruth Coupling
Projector (C)
Reweight
Sample Cost
Encoder feature
KLCH
Semantic Calibration
g(-)
map
Classifler f(-)
Label Caliuui)
Dynamic Filter
KL(H)
Figure 2: Overview of the AdaDRO framework. Starting from the empirical distribution Piz, construct a reference distribution v and dynamically filter out low-confidence samples to obtain the refined set i7. Semantic calibration is performed by aligning model-derived transport couplings with ground-truth couplings using inverse optimal transport (IOT) losses. This guides the construction of semantically meaningful and adaptive uncertainty sets.
The role of v. The formula (6) shows that the worst-case distribution Q* is essentially a reweighted version v, with weights determined by the loss function é and transport cost C. The absolute continuity condition Q < ensures that the support of Q* remains within that of v. highlighting why the design of v crucial for effective distributional robustness. By dynamically refining the reference distribution v, dynamic to better reflect the underlying data geometry, which has not been explored previously.
induce dynamic uncertainty set that evolves with the learning
Our Method: AdaDRO
To address the dilemma highlighted in Sec. M propose self-adaptive framework that jointly learns semantic costs and refines uncertainty set based on OT.
High-level idea We reformulate the Sinkhorn DRO objective as bilevel optimization whose upper-level and lower-level objectives are OT and IOT objectives, respectively (Sec. 3.T). This perspective inspires us incorporate more well-designed IOT for semantic calibration Sec. 3.2, moreover, the geometric implication of OT coupling can serve feedback to refine the uncertainty set (Sec. 3.3).
The overall framework illustrated in Figure]
3.1 Sinkhorn DRO Revisited: A New Perspective from OT
In this section, we introduce new perspective for Sinkhorn DRO from OT. First, we cast Sinkhorn DRO semi-relaxed optimal transport (OT) problem, where the benefit that the worst-case distribution Q^ (defined in (6)) emerges naturally from the optimal transport plan. Additionally, we reinterpret the softmax cross-entropy loss form of inverse OT. Taken together, these insights enable the entire learning dynamics of DRO to be formulated through the OT/IOT objectives, that paves the way for building our framework and theoretical analysis in the following subsections.
Distribution Exploration via OT. The optimal coupling of OT in Def. 2 represents the most efficient way to transport a source distribution P to the target distribution Q. By relaxing the target marginal constraint, i.e., replace the "y E 11(P, Q)" with "y E I1(P) = {r (p, g)dq = P(p)
D.(P,)) = inf (Bip-)),, [C(p, g)] - +H( P p)r) -ETI(P)
(7)
This relaxation considers couplings that preserve only the source marginal, allowing the distribution P to "explore" the target space. The exploration is guided by the transport cost function C and
5

---

the reference measure v. Under an appropriately chosen cost, this exploration process recovers the Sinkhorn DRO worst-case distribution Qº as its target marginal.
Lemma 3.1. Consider the DRO objective in ① with Pir and v being the source and refer- ence measures. If the transport cost in the semi-relaxed OT problem ( is taken as C"(p,q) q) (AC(p, -)(()) then the optimal transport plan 7 ElI(Prr) has its larget marginal equal to the worst-case distribution Q^ in (6)
Classification as Inverse OT. Complementing the above view, interpret classification through the lens of inverse OT. In this formulation, the goal is to learn a transport cost (implicitly parameterized by the model) such that the resulting optimal coupling aligns with the known label structure. To make this concrete, we define the ground-truth coupling and learned cost as:
(rs k) == Q* (xx))
Co(oe, k) := c (g0(x;), wk))
(8)
where (0, k) encodes the coupling between sample xy and its true label yı and Co(2s, k) defines the model-learned transport cost to class k E [K].
As shown by Shi et al. [53, when using the cost Cg in the semi-relaxed OT formulation the
resulting optimal plan "/o takes the form of softmax distribution: 700 (00 دー exp((o(pt).xx))//))
exp((g(u,),
and the inverse OT objective, defined (7 ) becomes:
exp (ap((()))
KL (5hee = =E log Σ Re exp((ata:::mm)//))) + Const.
This expression exactly recovers the semperature-scalle cross-entropy loss commonly used in classi- fication tasks.
Remark 1. This analysis shows that the DRO objective in 1 can be cast a bilevel optimization 32 grounded OT: Lower level (Exploration); exploring the worst-case distribution (*) by semi- relaxed OT (as Lemma 3.1. Upper level (Alignment): updating parameters A by solving an inverse OT problem that aligns the model's induced transport plan with ground-truth semantic labels. This OT-based perspective not only provides theoretical clarity but also suggests natural mechanisms for
adaptive algorithmic design, as explored in the following subsections.
3.2 Semantic Calibration via Inverse OT
Recall from Sec. 2 that we defined our overall transport cost as the sum of feature and label components: C Ppt, = C* (x;, + C³ (y/, Vj)- The effectiveness of DRO critically depends on how well these costs capture semantic relationships in the data. Rather than manually specifying these costs, we propose a principled approach learn them directly from the data using IOT.
Calibrating the Sample-level Transport Cost C* To ensure that transport costs capture task- relevant semantic relationships, we propose learning Cg. through a principled IOT formulation. We define the parameterized transport cost as:
C5 (I, 22) = c Ker (go(g)) go (z;)))
(10)
where is constant ensuring non-negativity, (a(x)) represents the embedding of sample z, and Ker(-::) is a kernel function (e.g., cosine or Gaussian) measuring similarity in the embedding space.
Given sample x; drawn from the worst-case distribution Q^, we generate a semantically equivalent counterpart af ~ Aug(æ;) using standard data augmentation techniques (e.g.. cropping, rotation). These pairs (#t, define ground-truth semantic matching encoded as a coupling:
xx (I,, x1)) [|j = i],
(11)
where 1[-] is the indicator function. This coupling represents the semantically meaningful matches that our transport cost should respect.
To calibrate C%K formulate the following 1OT problem:
min KL hx"), s.t. 2 arg min E IC + >ET((AA) (R.Z))
(12)
Similar to (T())) denotes the set of couplings whose first marginal is Q*, i.e., ITQ*)= (7) Σ, 7 (ayaa Q (z:)). This optimization aims to learn parameters such that the transport plan 70 induced by the cost function C... aligns with the ground- coupling 5 defined in (1

---

Calibrating Label-level Transport Cost C. To define meaningful transport costs between class labels, we need to quantify semantic relationships among them. Unlike data points, class labels lack explicit representations in the embedding space. Fortunately, the "Classification as Inverse OT" formulation from Sec. 3.TJoferss principled mechanism for semantic calibration in label space. To formalize this connection, we revisit the classification loss through the lens of inverse OT:
min KL KL (5l70), s.t. No -0--ap arg min E (Co(s, [o]])))]= + +e ((:
(13)
(EL((QA)
Here, the cost function defined as Colgi, k) = c- (go (x;), wa), which can be seen as special case of kernel-based cost: Co(g, k) = c (er(go(z;); wx), where Ker(-.-) is linear kernel. fundamental property of such kernlss, known as the reproducing property 447] establishes that: (Ker(wg,)) -), Ker ( ))> = Ker(wk, wp)). This property reveals that standard cross-entropy train- ing, which is equivalent to minimizing an inverse OT objective as shown in Sec. 3.1 implicitly
calibrates the semantic structure in the label space. Leveraging this insight, we define the label-Ievel transport cost as:
(k; (k,k)) - c- Ker(ww, ww),, Vk, k' [K]
(14)
where c is a constant ensuring non-negativity and Ker(..)) is the linear kernel. Classes with similar weight vectors will have lower transport costs between them, reflecting their semantic proximity. A similar idea of using classifier weight vectors wy to encode label semantics has also appeared in prior work, such (88, although they do not interpret this mechanism by inverse OT alignment.
3.3 Our Overall Algorithm
As discussed in Sec. [ the reference distribution v serves as a crucial prior that defines the support of the DRO uncertainty set. However, high-quality priors are rarely available: overly broad may include implausible samples, while an overly restrictive one weakens robustness guarantees. The OT formulations in Sec incorporate relaxed marginal constraints, reducing the need to forcefully match noisy or outlier samples
[5] [4]. This flexibility enables the transport coupling to
reflect the intrinsic structure of the data 49 64 Motivated by this observation. we introduce an OT-based feedback mechanism that adaptively filters the reference distribution v, thereby refining the uncertainty during training.
Refining v' via OT Feedback. We have introduced two Inverse OT frameworks for semantic Cerruttion-fattre---ee IOT in 112 and label-level IOT in (3). A key question now arises: how can we leverage the resulting couplings to dynamically refine the uncertainty set?
To address this, we propose a confidence-based filtering mechanism that evaluates the alignment strength between source-target pairs based on the inferred coupling. We illustrate this using the feature-level IOT framework in (12). The key idea is to compare the model-induced coupling "o with the ground-truth coupling 7t (defined in III ). For any matched pair (p, q) 5 we define its matching confidence as " (p, 4). which quantifies how strongly the model aligns p with q. A global average confidence is then computed as = (p))))))))) (p.
Because some targets may inherently be more difficult match, particularly when the marginal distribution of the ground-truth coupling 7x is imbalanced, we define target-specific local matching capacity for each g: (2)g) -E, (p, q). which quantifies how much mass is transported to q across all source. A larger (2) indicates that the target q is easier to match. Combining global and local signals, we define a dynamic threshold:
T(4) =
(15)
(g')
Each source is retained in the filtered reference distribution D only if it aligns well with at least one target q:
(p)) ((p) (p.q)>ann (oag)oo (p)))0000(0)))270 (p.)22(())
(16)
x
where X normalizes & to form valid distribution.
This filtering mechanism offers several advantages: (1) It establishes a feedback loop between model
training dynamics and uncertainty set refinement, enabling adaptive robust optimization; (2) It uses
7

---

Algorithm 1 OT Driven Adaptive Distributionally Robust Optimization (AdaDRO)
Require: Initial parameter Oo. step size a > O, total iterations T.
Train by the joint objective in [12 and (33 for No epochs. b Semantic calibrate by IOT
2: for t = 0, o......,
do
3: Dt f- FilterReference(v... 02)
Filter via (16)
4: Q f SinkhornWorst
Pu., Dr, Co.: A) D Compute worst-case distribution via (6)
5: Estimate gradient grad, via RT-MLMC estimation for the IOT objectives (13) and 12 See
Algorithm 2 Appendix
6: Update 0c+1 +- 02 at grad,
7: end for
sample-specific thresholds to account for varying target difficulties, overcoming key limitation of static filtering rules. This design of our filtering mechanism is guided by the geometric intuition behind transport couplings, previously discussed. We also draw conceptual connections to the well-known thresholding strategy introduced in FreeMatch [63], which adaptively filters samples in semi-supervised learning by combining global and local (class-specifie) thresholds. Such methods eliminate the reliance on fixed, manually tuned thresholds and leverage data-driven signals to inform selection.
Overall Algorithm. We apply the filtering rule in both semantic calibration tasks tasks- feature-Ievel IOT in (12) and label-level IOT in 13 to obtain a refined reference distribution D. We then perform Sinkhorn distributionally robust optimization using i7. The full procedure of our OT-driven adaptive
robust algorithm is summarized in Algorithm [ and its schematic overview is provided in Figure
Theorem 3.2 (Convergence of AdaDRO), Let the objective at iteration to be defined as F((0)
E @ [0(0,g), g)). Under mild smoothness assumptions, Algorithm 7 with decaying step size ar =
s/ +1 converges to a stationary point after at most T = O log 2) iterations:
min JE [IV < 2. DSKT
(17)
Here the total sample complexity O (6- log" 2
Theorem 3.2 establishes the convergence of our adaptive Sinkhorn DRO algorithm to a stationary point, despite two major challenges: (i) the dynamically evolving uncertainty set Q² that depends on both the model and the filtered reference distribution, and (ii) the presence of nested expectations in the objective, which renders unbiased gradient estimation intractable.
To handle (i), we ensure that both the transport cost Co and the filtered reference distribution Do vary smoothly with 6, as shown in Proposition B.I This ensures that changes to the uncertainty set remain controlled during training.
For challenge (ii), the objective F.(0) E g ()( involves a composition of expectations. Because the worst-case distribution (g), as we introduced in Ob. is given by
( (()))(()) D:(4), and ((g) Epp
(l()) A0Ce-)/AA)))
(18)
Here the denominator of (0q) also an expectation over u ~ Dr. This nested structure makes it impossible to construct unbiased gradient estimator using standard Monte Carlo, and naive stochastic approximations often incur intractable cost 62]. To address this, we adopt the multilevel Monte Carlo with Randomized Truncation (MLMC-RT technique (28). method has been used in robust optimization to trade off bias, variance, and computational cost efficiently. Related works such 38 use MLMC for estimating gradients in worst-case distributional settings. In our case, adapt MLMC-RT to the setting of dynamically changing uncertainty set and show that it remains theoretically sound.
Together, these properties enable our AdaDRO to converge to an e-stationary point using a total of O {) samples (as in Theorem 3.2 only logarithmically worse than the known lower bound for general non-convex stochastic optimization II. This result demonstrates that robust optimization with dynamically adaptive uncertainty sets and semantic costs can also be performed efficiently at scale. Our full thcoretical details provided in Appendix B

---

Experiments
We evaluate the proposed AdaDRO framework across diverse set of benchmarks to assess its effectiveness under various types of distribution shifts.
4.1 Experimental Setup
Datasets. We evaluate three widely studied distribution shift settings: Colored MNIST (2 which tests robustness under spurious correlations; Waterbirds [50], real-world dataset with strong background-lan correlation; CelebA 4 benchmark for facial attribute recognition. We also
evaluate on several long-tailed benchmark datasets: CIFAR-IO-LY and CIFAR-100-L1 37
Baselines. We compare AdaDRO with the following methods: KL-DRO 46
Wasserstein DRO
(WDRO) 9 Sinkhorn DRO SDRO 662 a unified DRO method UniDRO
semantic-aware invariant risk minimization method IRM (2)
robust methods such as GroupDRO 50
Our models are implemented with PyTTrch on a single NVIDIA RTX 6000 Ada GPU using the AdamW W optimizer [45]. For AdaDRÓ, we use cosine similarity as the kernel in (TO), and employ basic augmentations (flip, crop) for semantic calibration Sec. 3.2 unless otherwise specified.
4.2 Experimental Results
AdaDRO is broadly applicable thanks to its flexible reference distribution design. While this flexibilityy enables adaptation to various settings, our method still relies on prior knowledge construct the reference distribution, which may limit its applicability in some cases. Nevertheless, domain priors can often be effectively incorporated, as many specialized robust learning approaches. For instance, in imbalanced settings, the reference distribution v can be constructed through reweighting or resampling. We evaluate AdaDRO across diverse scenarios, including spurious correlation, group shift, and long-tail generalization. Additional experimental details are in Appendix D
Tables and show part of the results on Colored MNIST and Waterbirds. AdaDRO consistently achieves top worst-group accuracy, outperforming or matching specialized methods like GroupDRO.
Table 1: Accuracy (%) on Colored MNIST under high spurious correlation. AdaDRO achieves the highest robustness across methods.
Method
Avg Accuracy
Worst-group Acc
Gap
IRM
79.5
75.3
4.2
GroupDRO 150
78.4
73.0
5.4
KL-DRO 46
74.9
68.5
6.4
WDRO 9
75.8
69.1
6.7
SDRO 62
79.0
72.6
6.4
UniDRO |||
67.1
47.8
20.7
AdaDRO (ours)
87.2
79.9
7.3
Table 2: Accuracy (%) on Waterbirds. AdaDRO significantly improves the SDRO.
Method
Avg Accuracy Worst-group Ace
Gap
IRM (2)
82.6
73.2
9.4
GroupDRO 50
86.3
80.5
5.8
KL-DRO 46
81.0
71.5
9.5
SDRO 62
85.0
75.4
9.6
AdaDRO (ours)
87.1
80.5
6.6
Compared with other baselines, our method offers two main benefits. First, it performs label-
agnostic, sample-level semantic calibration as in (12). Second, it introduces a coupling-based filtering mechanism derived from the coupling introduced by (12) and (3)). An intuitive example is the
presence of noisy labels shown in Table 3 when portion of labels are flipped, sample-level

---

calibration [12 can still reliably capture semantics. This is because the process is unsupervised, similar to contrastive learning [53], and thus remains robust to label noise. Building on this, the coupling-based filtering mechanism can discards samples with wrong label, enabling the cross-entropy to train on cleaner subset of the data. In contrast, methods that rely heavily on label supervision, such as ERM with augmentation, suffer from inherent "garbage in, garbage out" problem. No augmentation can correct mislabeled examples, and these label errors inevitably propagate through the learning process, compromising model robustness.
similar advantage observed in Colored MNIST in Table Even when augmentation occasion- ally alters the spurious attribute (e.g.. color), standard ERM still learns entangled representations, as the label remains strongly correlated with color. In contrast, our label-agnostic calibration encourages the model to identify truly invariant features ttose that remain stable under augmentation -naturally prioritizing shape over color in the absence of a strong supervisory signal.
Table Accuracy (%) on CIFAR-100 with 50% label noise. Half of the training labels are randomly corrupted. We report test accuracy under different regularization strengths X.
Accuracy (%)
A=0.5
>=1.0
2=2.0
A=3.0
A=5.0
AdaDRO (Ours)
52.5 = 0.6
54.2 0.5
57.2 + 0.8
55.6 + 1.1
53.2 + 0.8
SDRO 62
42.5
0.3
41.2
44.1 = 0.3
44.8 + 0.5
43.9 + 0.3
Table 4: Accuracy (%) Colored MNIST benchmark under strong spurious correlations. ERM
Aug denotes empirical risk minimization with carefully designed data augmentation.
Method
ERM Aug
KL-DRO [46
IRM
SDRO 62
AdaDRO (Ours)
Accuracy (%)
43.5 = 0.6
41.2 + 0.5 45.1 + 1.0
44.8 + 0.6
57.2 + 0.8
5
Conclusion
We proposed a novel OT-based framework for adaptive robust optimization, where transport costs and uncertainty set co-evolve through inverse optimal transport. By combining semantic calibration with OT-based feedback, our method forms a closed-loop process that adapts to distribution shifts during training. The experiments show consistent gains over strong DRO baselines, and our framework opens promising directions for enhancing the flexibility of DRO methods. Moreover, our framework offers a new perspective for enhancing the flexibility DRO through the lens of OT, paving the way for incorporating richer OT tools into robust learning.
Acknowledgments and Disclosure of Funding
This work was supported in part by the National Key R&D program of China through grant 2021 YFA 1000900, the NSFC through grants No. 62432016 and No. 62272432, and the Provincial NSF of Anhui through grant 2208085MF163. The authors would also like thank the anonymous reviewers for their valuable comments and suggestions.
References
[1] Yossi Arjevani, Yair Carmon, John C. Duchi, Dylan Foster, Nathan Srebro, and Blake Woodworth. Lower bounds for non-convex stochastic optimization. Math. Program., 199 (1):165-214, 2023. doi: 00.1007/110007002--112222 URL httptt//p///d/ waciooi.oog//0.. 1007/
810107-022-01822-1
[2] Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk mini- mization. arXiv preprint arXiv:1907.02899 2019.
[3] Waiss Azizian, Franck lutzeler, and Jérôme Malick. Regularization for wasserstein distribu- tionally robust optimization. ESAIM: Control, Optimisation and Calculus of Variations, 29:33, 2023.
10

---

[4] Aurélien Bellet, Amaury Habrard, and Mare Sebban. A survey on metric learning for feature vectors and structured data. arXiv preprint arXiv:1306.6706 2013.
[5] Jean-David Benamou. Numerical resolution of an "unbalanced" mass transport problem. ESAIM: Mathematical Modelling and Numerical Analysis, 37(5):851-868, 2003.
[6] Jean-David Benamou, Guillaume Carlier, Marco Cuturi, Luca Nenna, and Gabriel Peyré. Itera- tive bregman projections for regularized transportation problems. SIAM Journal on Scientific Computing, 37(2): -1388 2015.
[7] Dimitris Bertsimas and Nihal Koduri. Data-driven optimization: A reproducing kernel hilbert space approach. Operations Research, 70(1):454-471 2022.
[8] Christopher M Bishop and Nasser M Nasrabadi. Pattern recognition and machine learning, volume 4. Springer, 2006.
[9] Jose Blanchet and Karthyek Murthy. Quantifying distributional model risk via optimal transport. Mathematics of Operations Research, 44(2):565- 600, 2019.
[10] Jose Blanchet, Yang Kang, and Karthyek Murthy. Robust Wasserstein profile inference and applications to machine learning. Journal of Applied Probability, 56(3):830- 857, 2019.
[1I] Anh Tuan Bui, Trung Le, Quan Hung Tran, He Zhao, and Dinh Phung. unified wasserstein distributional robustness framework for adversarial training. In International Conference on Learning Representationss 2022.
[12] Zihao Chen, Chi-Heng Lin, Ran Liu, Jingyun Xiao, and Eva L Dyer. Your contrastive learning problem is secretly distribution alignment problem. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL BRLhhtpss: /tttps:///penrevien wwwpriiwwwweeffrrreee net/forum?
1d=1NUKoLU8xb
[13] Wei-Ting Chiu, Pei Wang, and Patrick Shafto. Discrete probabilistic inverse optimal transport. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 3925-3946. PMLR, 17-23 Jul 2022. URL https://procceeddngg wtw.eeiiiinnnnn....ee..eeeeeeeee wwtpppeeee/e///eeveeeer chiu22 html
[14] Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François- Xavier Vialard. Unbalanced optimal transport: Dynamic and kantorovich formulation, 2019. URL https://tart/ org/
Abs/1508.05216
[15] Yin Cui, Menglin Jia, Tsung- Yi Lin, Yang Song, and Serge Belongie. Class-balanced loss based on effective number of samples. In Proceedings of the IEEE/CVF Conference on Computer Vision and Puttern Recognition, pages 2267-9277, 2019.
[16] Marco Cuturi. Sinkhorn distances: Lightspeed computation optimal transport. In Advances in Neural Information Processing Systems 26: 27th Annual Conference on Neural Information Processing Systems 2013. Proceedings of a Meeting Held December 5-8, 2013, Lake Tahoe, 2013.
Nevada, United States., pages
[17] Marco Cuturi and David Avis. Ground metric learning. The Journal of Machine Learning Research, 15())5333-564, 2014.
[18] Jason V Davis and Inderjit S Dhillon. Structured metric learning for high dimensional problems. In Proceedings of the 14th ACM SIGKDD international conference on Knowledge discovery and data mining, pages 195-203, 2008
[19] 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:1 10096-12007, 2020.
[20] Tongtong Fang, Nan Lu, Gang Niu, and Masashi Sugiyama. Generalizing importance weighting to universal solver for distribution shift problems. Advances in Neural Information Processing Systems, 36, 2024.
11

---

[21] Charlie Frogner, Chiyuan Zhang, Hossein Mobahi, Mauricio Araya, and Tomaso A Poggio. Learning with a wasserstein loss. Advances in neural information processing systems, 28, 2015.
[22] Rui Gao, Xi Chen, and Anton J. Kleywegt. Wasserstein distributional robustness and regulariza- tion in statistical learning. CoRR, abs/1712.06055. 2017.
[23] Rui Gao, Xi Chen, and Anton J Kleywegt. Wasserstein distribuuionally robust optimization and variation regularization. Operations Research, 72(3):1177-119,, 2024.
[24] Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics, pages 1608- 1617. PMLR, 2018.
[25] Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. arXiv preprint arXiv:2007.01434, 2020.
[26] Dan Hendrycks, Norman Mu, Ekin Dogus Cubuk, Barret Zoph, Justin Gilmer, and Balaji Lakshminarayanan. Augmix: A simple data processing method to improve robustness and uncertainty. In International Conference on Learning Representations.
[27] Youngkyu Hong, Seungju Han, Kwanghee Choi, Seokjun Seo, Beomsu Kim, and Buru Chang. Disentangling label distribution for long-tailed visual recognition. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 6626-6636 2021.
[28] Yifan Hu, Xin Chen, and Niao He. On the bias-variance-costt tradeoff of stochastic optimization. Advances in Neural Information Processing Systems, 34:22119 221321,, 2021.
[29] Zhaolin Hu and L Jeff Hong. Kullback-Leibler divergence constrained distributionally robust optimization. Available at Optimization Online, 1(2):9, 2013.
[30] Jiawei Huang and Hu Ding. An effective manifold-based optimization method for distri- butionally robust classification. In The Thirteenth International Conference on Learning Representations, 2025. URL
[31] Ruomin Huang, Jiawei Huang, Wenjie Liu, and Hu Ding. Coresets for wasserstein distribution- ally robust optimization problems. Advances in Neural Information Processing Systems, 35: 26375-26388 2022.
[32] Kaiyi Ji, Junjie Yang, and Yingbin Liang. Bilevel optimization: Convergence analysis and enhanced design. In Marina Meila and Tong Zhang, editors, Proceedings of the 38th Interna- tional Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pages 4882-4892 PMLR, 2021. URL http://pproceedings mlr, aress/v139/ji21e html
[33] Bingyi Kang, Saining Xie, Marcus Rohrbach, Zhicheng Yan, Albert Gordo, Jiashi Feng, and Yannis Kalantidis. Decoupling representation and classifier for long-tailed recognition. International Conference on Learning Representations, 2020.
[34] Mahmut Kaya and Hasan Şakir Bilge. Deep metric learning: A survey. Symmetry, 11(9):1066.
2019.
[35] Prannay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishnan. Supervised contrastive learning. Advances in neural information processing systems, 33:18661-18673, 2020.
[36] Masanari Kimura and Hideitsu Hino. A short survey on importance weighting for machine learning, 2024. URL https: org/ 2btp//3.///222.32/// 10175
[37] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems, 25, 2012.
[38] Daniel Levy, Yair Carmon, John C Duchi, and Aaron Sidford. Large-scale methods for distributionally robust optimization. Advances Neural Information Processing Systems, 33: 8847-8860, 2020.
12

---

[39] Ruilin Li, Xiaojing Ye, Haomin Zhou, and Hongyuan Zha. Learning to match via inverse optimal transport. Journal of Machine Learning Research, 20(80):1-37, 2019. URL http:
7/jmlr.orz wwtpeerrs//////---// htmll
[40] Tianyi Lin, Nhat Ho, and Michael Jordan. On the efficiency of entropic regularized algorithms for optimal transport. Journal of Machine Learning Research, 23(137):1-42, 2022.
[41] Tsung- Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, and Piotr Dollár. Focal loss for dense object detection. In Proceedings of the IEEE International Conference on Computer Vision, pages 2980-2988, 2017.
[42] Jiashuo Liu, Jiayun Wu, Bo Li, and Peng Cui. Distributionally robust optimization with data geometry. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
[43] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In 2015 IEEE International Conference on Computer Vision, ICCV 2015, Santiago, Chile, December 7-13, 2015, pages 3730-3738. IEEE Computer Society, 2015, doi: 10.1109/ICCV. 2015.425. htthttps: //oii org 1wogg//....0 10. 1109/ /CB///CC./..// 2015.022520 425
[44] Ziwei Liu, Zhongqi Miao, Xiaohang Zhan, Jiayun Wang. Boqing Gong, and Stella X Yu. Large- scale long-tailed recognition in an open world. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 2537-2546, 2019.
[45] Uya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2018
[46] Hongseok Namkoong and John C Duchi. Stochastic gradient methods for distributionally robust optimization with f-divergences. Advances in neural information processing systems, 29, 2016.
[47] Vern I Paulsen and Mrinal Raghupathi. An introduction to the theory of reproducing kernel Hilbert spaces, volume 152. Cambridge university press, 2016.
[48] Mikhail Persiianov, Arip Asadulaev, Nikita Andreev, Nikita Starodubcev, Dmitry Baranchuk, Anastasis Kratsios, Evgeny Burnaev, and Alexander Korotin. Inverse entropic optimal transport solves semi-supervised learning via data likelihood maximization. arXiv preprint arXiv:2410.02628, 2024.
[49] Khiem Pham, Khang Le, Nhat Ho, Tung Pham, and Hung Bui. On unbalanced optimal transport: An analysis of sinkhorn algorithm. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 7673-7682. PMLR, 2020. URL httt:::://roceeedingg mlr. prress/v19/phaee html
[50] Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust neural networks. In International Conference on Learning Representations, 2019.
[51] Shiori Sagawa*, Pang Wei Koh*, Tatsunori B. Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. In International Conference on Learning Representations, 2020.
[52] Dvir Samuel and Gal Chechik. Distributional robustness loss for long-tail learning. In Pro- ceedings of the IEEE/CVF International Conference on Computer Vision, pages 9495-9504, 2021.
[53] Liangliang Shi, Gu Zhang, Haoyu Zhen, Jintao Fan, and Junchi Yan. Understanding and generalizing contrastive learning from the inverse optimal transport perspective. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference Machine Learning. volume 202 of Proceedings of Machine Learning Research, pages 31408-11421. PMI.R 23-29 Jul 2023. URL WRLLhttps:tm
mlr. wwwpsrresssssee. w20sss//222/a//e shi .ttll
13

---

[54] 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.
[55] Richard Sinkhorn. Diagonal equivalence to matrices with prescribed row and column sums. The American Mathematical Monthly, 74(4):402-40) 1967.
[56] Agnieszka Stowik and Leon Bottou. On distributionally robust optimization and data rebal- ancing. In Gustau Camps-Vall Valls, Francisco J. R. Ruiz, and Isabel Valera, editors, Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151 of Proceedings of Machine Learning Research, pages 1283-1297. PMLR, 28-30 Mar 2022. URL mlr.
html
[57] Andrew M. Stuart and Therese Wolfram. Inverse optimal transport. SIAM Journal on Applied Mathematics. 80(1):599 619, 2020. doi: 10.1137/19M126112. 122. URL https
//doi.org// 1tt1//1..117/111111111111
[58] Zhiquan Tan, Kaipeng Zheng, and Weiran Huang. OTMatch: Improving semi-supervised learning with optimal transport. In Ruslan Salakhutdinov, Zico Kolter, Katherine Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp, editors, Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 47667 77667-47680. PMLR, 21-27 Jul 2024. URL https //proceedings.mil mlr, press/ viless//235/tatt//
html
[59] Christos Tsirigotis, Joao Monteiro, Pau Rodriguez, David Vazquez, and Aaron C Courville. Group robust classification without any group information. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 56553-56575 Curran Associates, Inc., 2023.
[60] Alexander Tsybakov. Optimal aggregation of classifiers statistical learning. The Annals of Statistics, 32(1):135-166 2004.
[61] Cédric Villani. Optimal Transport: Old and New, volume 338. Springer, 2009.
[62] Jie Wang. Rui Gao, and Yao Xie. Sinkhorn distributionally robust optimization, 2023. URL https: //arxi:///xxvv. orrg/iioorr//oooo abs/2109. 11926
[63] Yidong Wang, Hao Chen, Qiang Heng, Wenxin Hou, Yue Fan, Zhen Wu, Jindong Wang, Marios Savvides, Takahiro Shinozaki, Bhiksha Raj, Bernt Schiele, and Xing Xie. Freematch: Self- adaptive thresholding for semi cemi-supervised learning. In The Eleventh International Conference on Learning Representations, 2023. URL https : A
htap::///openae //tp///vepeeeeeee////eeeeee mete/Eoorum?P--PPXUPXXIIII
[64] Zifan Wang, Yi Shen, Michael Zavlanos, and Karl H Johansson. Outlier-robust distribution- ally robust optimization via unbalanced optimal transport. Advances Neural Information Processing Systems, 37:52189-52214 2024.
[65] Kilian Q Weinberger and Lawrence K Saul. Distance metric learning for large margin nearest neighbor classification. Journal of machine learning research, 10(2), 2009.
[66] Eric Xing, Michael Jordan, Stuart Russell, and Andrew Ng. Distance metric learning with application to clustering with side-information. Advances in neural information processing systems, 15, 2002.
[67] Hongyi Zhang, Moustapha Cissé, Yann N. Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018 URL httpps://openreviee net/ wewt/oorumiddrrDDooD
14

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification:
the abstract and introduction summarize our contributions.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: included
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question:
For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
15

---

Justification:
All assumptions are stated and complete proofs are included in the appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: included
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
16

---

Answer: [Yes]
Justification: included
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ ctpppp.ccppulli//guiii wwwiicccguie///eeeeeeeeis CoodeSubeiiie ssi wswonnolli onnolicy)))) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Anonymized code and data are included in the supplemental material. Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:[No
Justification: it would be too computationally expensive.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
17

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: included
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g.. preliminary or failed experiments that NeurIPS Code of Ethics https: http::::////iis/.eeeee ctt.rrisscc..c.aececceooooooe httpp///ccc//ccccccccieeoeeeleee
didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the
Answer:
[Yes]
Justification: we confirm our work complies with all stated principles.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is a special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: does not involve social impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to
18

---

generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our model and data do not pose high misuse risks requiring safeguards.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: the third-party datasets and code used are properly cited with licenses.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, snperswwihcoode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's is creators.
13. New assets
Question: Are new assets introduced in the paper well documented and the documentation provided alongside the assets?
19

---

Answer:
[NA]
Justification: does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/codd/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: does not involve human subjects
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
20

---

Justification:
No LLMs were used as a part of the core methodology.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy http:eae//eerrppoocccc/oooeeeeee for what should or should not be described.
21

---

A Omitted Materials
A.1 Other Related Works
Metric learning. The inverse optimal transport problem shares conceptual and methodological similarities with metric learning, thoroughly examined by 4 34]. Metric learning focuses
estimating pairwise distance metrics from observed interactions, typically constraining solutions to specific metric classes such Mahalanobis distances [66
65
18] However, while metric
learning directly optimizes distances, inverse optimal transport (IOT) [57
48 employs more
sophisticated bilevel optimization structure that accounts for global coupling relationships between point sets, significantly increasing its complexity. Cuturi and Avis [77] conceptualized the inverse OT problem as a metric learning task, an approach later applied to generative models
[24]. Taking
different direction, Frogner et al. [21] proposed constructing cost functions using external priors such as word embeddings, where semantic similarity between labels derived from Euclidean distances in pre-trained word2vec spaces (as demonstrated with the Flickr dataset). Though intuitive, this approach introduces external dependencies and may not adequately capture task-specific semantic relationships.
provides prin- Distributional robustness. Distributionally robust optimization (DRO) 38 52
cipled framework for learning models that remain reliable under test- time distribution shifts. The central idea is to optimize model performance against the worst-case distribution within predefined uncertainty set around the training distribution. A key ingredient of DRO is its dual formulation, which enables tractable reformulations and efficient algorithms (3T) This paradigm has been extended to group-wise robustness, as in Group DRO [5I], which seeks to ensure balanced performance across predefined or latent subgroups-ewe as [59] further refine this idea by incorporating logit adjustment to improve performance parity across groups without explicit group labels.
when some groups are underrepresented. Recent methods such
Other related robustness methods. Another prominent class of techniques for handling distri- bution shift is importance weighting (IW) [36), which reweights training samples to better match the test distribution. Several recent advances have extended IW to broader settings. For instance, Shu et al. (54) proposed meta-learning framework that learns sample-weighting function using small unbiased validation set, effectively addressing issues like class imbalance and label noise. Fang al. [19 introduced iterative method that jointly estimates weights and performs classification in an end-to-end fashion. More recently, Fang et al. (20) proposed Generalized Importance Weighting (GIW), which can handle complex shifts, including scenarios where the training and test distributions differ in support.
A.2 Reproducing Kernel Hilbert Space (RKHS)
Reproducing Kernel Hilbert Space (RKHS) g 7 serves a cornerstone in modern machine learning
and functional analysis, providing a rigorous mathematical framework for kernel-based methods. This paper outlines the fundamental concepts of RKHS, presents its formal definition, and highlights its significance in computational applications.
In many areas of machine learning and statistics, the representation and manipulation of data in digh-dimmesional spaces are crucial for understanding complex patterns. Reproducing Kernel Hilbert Space (RKHS) provides an elegant framework that combines the computational efficiency of kernels with the theoretical rigor of Hilbert spaces. RKHS is particularly notable for its role in kernel methods, such as support vector machines and Gaussian processes.
The power of RKHS lies in its ability to map data into a potentially enabling computations through dimenssooan kernel functions. We begin by defining the key components that constitute Reproducing Kernel Hilbert Space.
space while
Definition 3 (Reproducing Kernel). Let H be a Hilbert space of functions defined on a non- set X. A function Ker X x X + IR is called reproducing kernel for H. if: (i) For all x e X. Ker(,, E H. (ii) For all f E H and x E X, the reproducing property holds: f(x) = (f, Ker(,,)))H+
Definition& 4 (Reproducing Kernel Hilbert Space). A Hilbert space H of functions on X' called a Reproducing Kernel Hilbert Space if there exists reproducing kernel Ker for H.
22

---

The reproducing property ensures that the evaluation of a function f E H. at point E X' can be expressed as an inner product in H. This property underpins the computational efficiency of kernel-based methods. In our setting, we leverage this property to construct semantic transport costs (see (T0 and (14) that reflect similarities in latent and label spaces, while enabling efficient gradient-based optimization.
A.3 Differences and Connections with other DRO methods
Differences and Connections with Wasserstein DRO and KL DRO (1) Flexibility of the Cost Function: WDRO requires the transport cost function C(. -) to be a metric, whereas Sinkhorn DRO only imposes the weaker condition of non-negativity C(I,1) y) 2 0. This makes Sinkhorn DRO more flexible. In WDRO without regularization, the worst-case distribution is typically discrete, while in Sinkhorn DRO, the worst-case distribution Q* is absolutely continuous with respect to the reference measure v. The reference measure // can be interpreted as prior knowledge about the structure of the worst-case distribution, which enhances both the interpretability and practical applicability of the model. Notably, the regularization parameter c + O, the dual objective of Sinkhorn DRO converges to that of Wasserstein DRO because the Sinkhorn distance is an antropy-regulaninn version of the Wasserstein distance. (2) Connection to KL DRO: When the reference distribution and the transport cost function C(.,) = O, Sinkhorn DRO reduces to the KL-DRO problem. this case, the entropy regularization term of the Sinkhorn distance simplifies to the KL divergence, which measures the cross-entropy difference between the worst-case distribution and the input distribution.
Pu-
A.4 Solutions for cntropy-RReulaare Optimal Transport
For completeness, we summarize the closed-form solutions for entropy regrop--rggulrrizedd optimal transport (OT) problems under different levels of marginal constraint relaxation. These results standard in the OT literature and can be found in (53). Such relaxations provide analytical forms for the optimal transport plan 7, which significantly reduce computational complexity and enable efficient integration into modern learning algorithms.
Remark 2. Let IIP, Q) denote the of couplings (joint probability measures) with specified marginals:
(I) = {r II [ 9) dg = P(p), dp IIP)) = {r [ dq = P(p)
(19)
(20)
((1) (- ( dpdy 1}
(21)
This semi-relaxed formulation preserves only the source marginal. The retropy-regularrieed OT problem is defined as:
Relaxed Marginal Constraint: IT(P)
= arg min *En(P)
[y 4) + +cH(7 P))
(22)
where H(y P&v) denotes the relative entropy (defined in (2) with respect to the product measure of P and reference measure The optimal coupling has closed-foree expression:
exp(- C(p.9))
> g) J exp( Glo.e) du(u)
P(p).
(23)
This expression yields a softmax-like kernel around each point p, shaped by the cost C' and regular- ization 6, and normalized with respect to the reference measure v.
Fully Relaxed Constraint: 7 E (T(1) In this most relaxed case, only the total mass is constrained to be I:
7" = arg min 71)(1)
E(p (C(p, 4) + cH())-
(24)
The optimal transport plan admits the following form:
exp( c6..)
where Z = IS exp (- C(u..) dudv.
>" (p. q)
(25)
Z
This represents global Gibbs distribution over the joint space X y and is often used for uncon- strained matching.
23

---

Fully Constrained OT: y E QP)))Q) This is the standard OT formulation with both source and target marginals enforeed. In this case, no dosed-form solution exists, but iterative scaling via the Sinkhorn algorithm 55 6 yields the unique entropy-regularized solution:
(p. exp( C(p.9)
(26)
(p.g) P(p): (row normalization)
(27)
[T(0))p)') )
(p.
J d' Q(g), (column normalization)
(p.g)
(28)
This procedure is guaranteed to converge to the unique solution y E (I(P.QQ satisfying the marginal constraints.
These closed-fome and iterative solutions illustrate how different constraint relaxations influence the structure of optimal couplings. our framework (see Sec. S.1. semi-relaxed formulation y E II(P) due to its analytical tractability and its suitability for dynamic uncertainty set modeling under limited supervision.
35.]and and Sec. 3.2), we make use of the
A.5 Proof of Lemma 3.1
Proof. We begin by considering the semi- relaxee optimal transport (OT) problem where the coupling y preserves only the source marginal Pur, with reference distribution v, and modified transport cost:
Cp,, ②) = (AC(p, - (CC(),))e(qu))) f(g))
(29)
We seek the optimal coupling >* II(Ptr) that minimizes the entropy-regularized OT objective:
arg min PEll(Pa)
((p.a)>,(C'pp.. + P.. @ 0v)
(30)
Applying the standard solution for entropy-reeullaizze semi-relaxed OT (see Appendix A.4 the optimal coupling has the form:
expl ('(9)))
7 (p.9)
Pir(p).
(31)
[exp
C'(p.w) du(u)
Substituting the definition of C' 9). we obtain:
exp( ((C()))))))
(0 4) Tepp [xppp_ (A(p)) (p)
(32)
(CC2)))(()
(((l)))--0(pp)))))))l)a (p).
(33)
Taking the marginal over p ~ Pur, the target marginal (i.e., the induced distribution on g) becomes:
() / 7 (p, Ep- [
(((()) ((()))ACG--))((000)
(34)
This matches exactly the closed-form expression for the worst-case distribution Q* in 6 proving the claim.
A.6 General Case for the Ground-truth Matching 7
In the main text, we introduced filtering mechanism for the special case of one-to-one ground-truth matchings. Here, we present generalized formulation applicable to many-to-many matchings, which naturally extends approach while preserving its key properties.
24

---

Generalized Matching Quality Metrics. For given source p. we define its matching quality as: (p(p) E (T(..) q) (o 9)),
(35)
which measures the expected alignment between the model-derived coupling (p, q) and the ground- truth coupling (p. q) across all targets associated with p. This metric provides a source-specific measure of how well the learned transport map matches the desired coupling structure.
The global average matching quality is computed across all sources as:
71 = (p-r)(o(p)]
(36)
providing date measure of matching performance. For each target q, we quantify its matching capacity by:
(2 = Ep- (p, g)],
(37)
representing the total transport mass received by target q from all sources. A higher value of (2 indicates that g is easier to match, consistent with our intuition in the main text.
Generalized Filtering Mechanism. Using these metrics, we compute dynamic threshold for each larget g:
T2(2)
(q) = T1 max, T2(
(38)
which adjusts the global threshold 5 based the relative matching capacity of target g.
For each source p. we compute a personalized threshold by taking the expectation over its associated targets:
T(p) = (70 T())]
(39)
The generalized filtering rule retains source p in the filtered distribution if its matching quality exceeds its personalized threshold:
(pp) v(p) [/T((p) 2 T(p)
(40)
X
where x normalizes 2 valid distribution.
This generalized mechanism preserves the key benefits of our approach: (1) it establishes an adaptive feedback loop between model training and uncertainty set refinement, and (2) it accounts for varying difficulties across different targets and source-target associations.
Lemma A.1 (Reduction One-to-One one Case). When the ground-truth coupling & is one-to-one (i.e., for each source there exists exactly one target qp such that (p, 4p) = and 7(p.4) = for all q & qp), the generalized filtering rule in (40) reduces to the one-to-one filtering rule 16
Proof. For a one-to-one ground truth coupling 7, each source p is associated with exactly one target 4p where (p, qp) and (p, q) for all q # 9p- Under this condition:
(i) The source matching quality simplifies to:
Tp(p) o q) ()(p 70 4p)
(41)
(ii) The source-specific threshold becomes:
(p) EB ( [()) = T(p))
(42)
(iii) Therefore, the filtering condition To(p) 2 T(p) becomes:
70(P. 9p) 2 T(p))
(43)
This is equivalent to the condition in 16 "3: 7p,.) g) > 0 A (p g) 2 r(1)" since qp is the only target for which 7(, q) > Thus, the generalized filtering mechanism correctly reduces to the specific case presented in the main text when the ground-truth coupling is one-to-one.
D
25

---

B Theoretical Foundations
This section establishes the convergence of our adaptive Sinkhorn DRO approach. We address two key technical challenges: (1) the dynamic nature of the uncertainty set during optimization, and (2) the nested expectations in our objective leads the biased gradient estimation.
We prove convergence under four mild regularity assumptions that govern the behavior of the reference
distribution filtering mechanism. First, we establish that the filtered reference distribution is Lipschitz continuous with respect to model parameters (Proposition B.I). Combined with our Lipschitz semantic transport cost, we obtain the continuity of the worst-case distribution || Qillvv (Theorem B.4h.
To handle nested gradient computation, we employ Multi-Level Monte Carlo with Random Truncation (MLMC-RT), showing that bias decays quadratically with approximation level (0(2-4) while variance increases only logarithmically (Proposition B.5
By carefully integrating these results with stochastic optimization theory, prove that algorithm with decaying stepsize a, = a/(Vt+ achieves min:<T E|] (..) 12 in T' O( log ) iterations, with total sample complexity O log" +) (Theorem B.7 This complexity nearly matches lower bounds for non-convex stochastic optimization, differing only by logarithmic factors.
B.1 Key Assumptions
We first formalize the assumptions that underpin our theoretical analysis.
Assumption (Smoothness of the Model).
(i) The loss function E(0, x) is Lg-smooth
[]V =f) < Lo |0_ 02ll Ve X; (ii) The induced semantic cost is Lipschitz continuous for any q E P(E): CCo (p. Co2 4)l S Lc 02ll-
Assumption 2 (Lipschitz Continuity of Matching Quality). (i) The matching quality function 7o (p, 4) is L, Lipschitz continuous in 0:
Iro, (p. q) - ro 4)l S 1, - 02ll
(44)
(ii) The threshold function (4; 0) is L, Lipschit continuous in 0: |r 01) ((: 02)1 - L- 02ll
(45)
These assumptions abstract the key continuity properties required for our theoretical results. Assump- tion T ensures that both the model loss and the resulting semantic cost vary smoothly with respect to model parameters. Assumption 2 guaranttess that the OT-based matching quality and the adaptive threshold used in filtering are stable under small changes in the model. In Section CC verify that these conditions hold in our AdaDRO framework, specifically for the semantic cost function and the OT-based feedback mechanism that determines the adaptive threshold.
Assumption 3 (Boundary Regularity). For the ground-truth coupling 7. where each source is associated with one target Up such that 7(p 4p) >> O, we assume that for any & >> 0 and 0, the probability mass of points in the e e-neighborhood of the decision boundary is bounded:
(tp hpop,, gp) T (4pi 0)1 < s}) S LBs
(46)
for some constant LB > independent of 9 and E. Assumption 4 (Non-degenerate Filtering). For the ground-truth coupling 7, where each source p is associated with one target 9p such that 7(, 9p) > 0, we assume there exists 5 >> O such that for
all in the parameter space, v(So) 2 6, where So = (p 70(p. 9p) 2 T(p;: 0)) the set of points accepted by the filtering operation.
Interpretation of Regularity Assumptions The validity of our convergence analysis relies on two interpretable regularity conditions governing the filtering mechanism: Boundary Regularity (Assumption 3) establishes that the reference distribution's density decays linearly near decision boundaries. This smoothness condition condition-analogonn analogous to the Tsybakov margin condition in statistical learning 600 prevents pathological scenarios where infinitesimal parameter changes could cause discontinuous shifts in the filtered distribution. Non- -on-Deeenee Filtering (Assumption guarantees
26

---

the filtered distribution maintains minimum effective support throughout training. Implementation ally, this achieved through adaptive threshold calibration or entropy constraints in the transport plan. Theoretically, it prevents algorithmic instability from excessive sample rejection.
Generally speaking, these assumptions ensure that the filtering operation produces distributions that vary continuously with respect to model parameters (Proposition [.I). which is essential for the convergence guarantees of adaptive robust optimization algorithm.
Algorithm Overview We analyze the convergence properties of our proposed iterative procedure (Algorithm [ for updating parameters IRd. At iteration let Vo, represent the filtered reference distribution and Co, denote the transport cost. The algorithm proceeds through the following steps at each iteration
Given 0, we filter the reference distribution v to obtain Dg, by removing low-confidence samples according to:
Reference Distribution Filtering:
i9g (A) v(AnSa.)
where Sq. {p Te, (p. q) 2 T(q; (p. 3}, Xe. = vSe.) (47)
2. Semantic Cost Adaptation: We update the transport cost function to Co capturing semantic distances in both feature and label spaces.
Worst-Case Distribution Computation: Using Pe. and Cs., we compute the Sinkhorn DRO worst-case distribution Q7 with fixed regularization parameter A:
Te (((g)-ACo, (p.9) dop, (1)
(Q) Bup
(48)
Zo, (p)
where (p) = Eu-P, (M (n.u)).
4. Gradient-Based Update: We perform gradient step the iteration-specific objective:
(0) = E Φ-
[((0. q)]
(49)
01+1 02 aigrad, where grad, ss VgFi(0.)
(50)
We will show that if Do, and Co, vary "slowly" with respect to fr, and F, smooth in 0, then the sequence {6,} converges to stationary point under standard non non-convex optimization conditions.
B.2 Continuity of Filtered Distributions
key aspect of approach is the filtering operation applied to the reference distribution. We establish that this operation produces distributions that vary continuously with respect to model parameters. The following analysis focuses on the one-to-one ground-truth coupling case described in the main paper, though similar results can be established for the general case in Appendix AA.6]
For one-to-one ground-truth coupling 7, each source p is associated with exactly one target Sp such that (p, 4p) and (p, q) = for all # 9p- The filtering operation is defined as:
v(p) [[ qp) 2 (pp: 0)
Do(p)
(51)
X
where is a normalization constant ensuring Do is a probability distribution.
Proposition B.1 (Lipschitz Continuity of Filtered Distributions). Under Assumptions 2 the filtering operation produces distributions that Lipschitz continuous with respect to in the total variation distance:
100
To, lrv
Lov [|0: -02ll
(52)
holds for any 01 02 and some constant Lrv > 0.
Proof. We proceed in three steps:
(i) Bounding the symmetric difference. Let S₁ {pl 70, (p. 2 ((pp;))) and S2 = {p)
Ye, (p. 9p) (go: 02)}.
27

---

For p E S1 \ S2, we have: Ye, (p. 2 (4p: 01) and Yo, (p, qp) < (qpp: 02)
Using the Lipschitz continuity from Assumption 2
70 (P, (p) 2 70, (p, {p) - Ln110 02
(53)
T9pi 02) I {pp: 01) -- 82ll
(54)
where T.y and by are the Lipschitz constants from Assumption 2
For S1 \ S2, these inequalities imply:
S To, (p. 9p) - (qp: < (Ly + L,)1 - 82ll
(55)
Similarly, for E S2 \ S1, we obtain:
- (L, + 1|0 0211 < 70, (p. 9p) 7 (qp: 01) <
(56)
Therefore, the symmetric difference S1 452 (S) \ (2) U (S2 \ S1) is contained within:
{p: hes (p. 9p) (qp: 01)1 < (Ly L, +=)1 02113
(57)
By Assumption 3 we have:
v(5 I LB (I, + L,) 021 = L, 110 02ll
(58)
where L,, = LB (In + [.) and L.B is the constant from Assumption
(ii) Bounding normalization constants. The normalization constants are X0, v(Si) and Xea =
v(2),, with difference:
Ixe, xox (ell=(((S) v(S2) s ((8:48)) < L, |0 02ll
(59)
By Assumption 4 both Xe, and Xa, are bounded below by 8 > 0, where d is the constant from Assumption 4
(iii) Bounding the total variation distance. The total variation distance between Do and Vo, is:
1|a. Dig 1TV SUPAEF
((ASi))
((AnSs)
(60)
X01
X#2
We decompose this into:
v(AnS:)
v(An82)
<
((AnS))
v(AnS))
v(AnS:)
v(AnS))
(61)
X*1
X02
X6;
X02
X#2
X02
For the first term:
M(ANSL)
p(AnS))
=v(Ans S1)
1
= v(AAns))
Xx0: X0al
X=1
X0z
X01
X8a
X0,
1xo1
X62
Ly 110
02ll
(62)
X0g
where we used (i) S Xo, and Xo, 2 8.
For the second term:
[u(An S1) ~(A n S2)l
v(AnS)) v(AnSs)
X62 X0g
X0=
(((ANS))AAA S2))
v(An (SIAS))
X02
X02
((S
Lu 110 021l ④
L
(63)
X02
Combining both terms:
v(AnS))
v(AnS))
< 3 1 82ll
(64)
X4g
Since this bound holds for any measurable set A, we have: 2L, 10 02 LTv 110 021l
| Porl <
(65)
0
where LTV 2hx 24n (hn-(L,+b-))
28

---

B.3 Stability Analysis of the Worst-case Distribution
To guarantee convergence of our algorithm, it essential to characterize how the objective function evolves in response to changes in the filtered distribution and the transport cost. The following lemma establishes the smoothness of the Sinkhorn DRO objective, showing that inherits the Lg Lsssmoothsess of the loss function with respect to model parameters.
Lemma B.2 (Smoothness of Sinkhorn DRO Objective). Under Assumption ! at any iteration the Sinkhorn DRO objective F.(0) = y [8(0, q)) is also La-smooth.
Proof. By Assumption | the loss function E(O. 9) La-smooth in for all qF
[|Vol(0r.q) - Vol(O2, 4) S Loll0, - 02ll
(66)
For the objective function Fr(0) E3, -+ [e(0, ()), where @ is fixed for a given we have:
[VF((()) (12(0))
Ve - 2 [e(0), 9)] Voll
q [((02.))]] 9)ll
= || و- Vol(O), 9)] - Φ Vol(02,.0)
s E Φ+ Vol((1. 9) Vol(
s E 0; |0 02ll
Loll01 021
(67)
Therefore, the objective function (0) inherits the La-smoothness from the loss function ((0, q).
To analyze the stability of our algorithm, we need to understand how the worst-case distribution changes with model parameters. Let Q₂ be the worst-case distribution computed using parameter 0, and regularization A:
(9(.a)) (p-a)
dQ)(9) = Epp
dūg, (q)
(68)
Zo. (p)
where Ze, (p) = Euwin, (0) ACo, (n.w)). The total variation (TV) distance provides a natural metric for quantifying the difference between probability distributions.
Proposition B.3. Let (K, be a measure space. Let f.g X + R be measurable functions with
finite log- pe---aatiiin functions Z₁ = Ix (((()))(())) Z, fx (ala)dd(e < c0. Define probability densities dz(x)) = 22 du(x), dQ(x)= = cdlr) 2, du(x). Then | Qlizv S Vllf glloe.
((,)
Proof. We use Pinsker's inequality to relate the total variation distance to the KL divergence:
PP- - Qlivv S V (L(P)Q)
(69)
For the KL divergence between P and Q, we have:
KL(PIQ) = / P(x) log
P(x) dx
ela))
ef(0) /Z, du(x)
log e0(x) /Zg
(70)
Q(x)
fx Zj
/ ed(m)
Zg )dp(z)
x Z, (f(x) -g(x) + loog
(71)
Z₁
/ Z5 (f(x) g(x)))d((x) + log Zg
((x)
(72)
Z
Since [f(x) g(x)) < IIf - glooo for all x, we have fx ed(e)
2, (f(x) g(x))da(x)) & [lf - glloo.
For the partition function ratio, since g(x)= S fxx + |lf al|xo for all x. we have:
Un / (ett)d(( / e(0)) dạ(x)) = ells- 2,
(73)
x
29

---

Zz S llf gllcc Combining these bounds: Therefore, logg 2,
KL(PIQ) & Ilf gll0o + Il - glloo = =2|1f - -glooo
(74)
Substituting into Pinsker's inequality: ||P Q1iv V 211 001 -
Thus, have established ||P Qlitv S VIIS- glooo-
Theorem B.4 (Continuity of the Worst-Case Distribution), Under Assumptions for consecutive iterations and +
Le 2
- - VillT S G/|0++1 - 0zll + 110++1 - 01
(75)
where G = GG+ + LTv combines the gradient bound G2 sup VV 9) the Lipschitz constant Le of the transport cost, and the Lipschitz constant LTv of the filtered distribution.
Proof. Recall that the worst-case distribution at iteration t is given by:
e (()) ((g) (P-9)
di (g) Ep- Pie
dia, (4)
(76)
(p)
We introduce an intermediate distribution to separate the effect of parameter changes conditional weights and reference distribution:
[ ((1))ACo,-1) (p.9)
dQ)+1(0) EpnPu
dos, (q)
(77)
Z0.11 (p)
By the triangle inequality:
l 27 2rv S 1122 Lilllr + 1122+
Qillev
(78)
For the first term, since @ and Qi11 differ only in their reference distributions:
l - 1lrv I ll De. 1lrv & LTV - 0ell
(79)
where the last inequality follows from Proposition B.I
For the second term, we analyze the change in conditional weights. For each fixed p. define the functions ft(g) = E(Oc, q) - ACo, (p. q) and fe-1(g) ((0++1, -XCoet (p. q). By proposition B.3 we have:
el+((()
elt(+)
I [lf+1 - fell
Zo. (p)
Zo, (p)
TV
We bound || felloo using Assumption [
1. For the loss function, by Taylor's theorem with the smoothness property:
10(0++1.g) E(O.. q)1 < (V,e( 0)1 - -Oel - Lg
||0++1 0.11
2
< G2 10+1 - 0ell + Lg
110 0.112
(80)
2
For the transport cost, by Lipschitz continuity:
ICo.. (p. g) (p. 9)l S Lcl10++1 0.11
Combining these bounds:
[lfe+1 felboo sup ff)+1(9) fe(g)l
Iap 2
< sup (Gell0e41 0ell +
|0:41 - -bal2 + ALool0++1 - -bel)
Lg
= (Ge + ALC) - 0,1l +
11 - -
(81)
2
30

---

Since this bound holds for all taking expectation over Pır preserves the inequality:
Q - Qillav I (G2 + XLc)1l0+11 - Bell + T-0
l102+1 - 02l1
(82)
2
Combining the bounds (9) and
into 78
Le ||0+11 0.l1
1 Qillr & Lrv 10+1 0el + (Ge + ALc)) --ll 2
La 110
(G ALC + LTV 01001 0z11 +
0111
Lg
= =G - 0ell 10++1 0.11
(83)
where G= G5 ALC + LTV.
This establishes that the worst-case distribution has continuous dependence on model parameters, accounting for changes in both conditional weights and reference distribution.
This result shows that small changes in model parameters lead to controlled changes in the worst-case distribution, with a continuity constant that depends linearly on the regularization parameter A. The quadratic term reflects the impact of the loss function's 's smoothness. This continuity property is crucial for establishing the convergence of our algorithm next subsection.
B.4
Convergence Analysis
Now we present our main convergence result. Compared to standard optimization settings, our algorithm must overcome two key challenges: (i) the dynamic nature of both transport costs and uncertainty sets; (ii) biased gradient estimates arising from the nested expectations in the Sinkhorn DRO objective.
For the first challenge, we have established the continuity of both the transport cost Ca and the filtered distribution Ba with respect to model parameters 0 in Sec. [33 To tackle the second challenge, employ the Multi-Level Monte Carlo (MLMC) gradient estimation technique, specifically the Random Truncation variant (MLMC-RT), which achieves an effective trade-off between bias, variance, and computational cost.
B.4.1 Multi-Level Approximation for Sinkhorn DRO
We adopt the Multilevel Monte Carlo with Randomized Truncation (MLMC-RT) framework setting, aim to optimize the following non-convex objective at each iteration t:
28 In
(0) Exp. Rigeo
((0.4)
(84)
Zo. (p)
where Ze, (p) = EEg (e (p-4)] is normalization factor. Because of the nested expectation and normalization term, directly obtaining unbiased gradients is intractable.
To overcome this difficulty, we introduce sequence of approximations indexed by level Specif- ically, approximate the Sinkhorn DRO objective F.(0) via the following level-( Monte Carlo approximation:
((()) ACo, (C-.)
F. F() Ep P.. E
Σ 2
((0,)))
(85)
gED!!
7 (p)
(1))
where:
denotes a Monte Carlo approximation of Dp. using 2 independent and identically distributed (i.i.d.) samples from Dy i.e., D!!)) (q.,...,ppp
1 (p) E EEa [ ACo, (n.u) is the corresponding normalization factor level I.
31

---

Algorithm 2 MLMC-RT Gradient Estimation for Sinkhorn DRO
Require: Current parameter O., target accuracy 6, constants 2, = cc
1: Set the maximum level L [log (e (e]1) 1)1
2: Define sampling probabilities prob (1) 2 2-(11) -2- for all (0,,,....,1 1
3: Randomly sample level I {prob(") (0) prob (2) ....,probl())
4: Draw 2' i..dd samples from Do, to form o!
d,
5: If > 0, reuse the first 2/-1 samples to construct D% ((-1)
6: Compute G. (0,, ( (from (66) and ((t- (0c, (1-1)) if 0
7: Return the gradient estimate:
C ( (6,,
if = 0, otherwise.
prob(o)
grad,
G G (BL, - (0,,
prob
MLMC-RT gradient idea.
The gradient VoFi(0) is approximated by estimating the gradient of
Fl (0):
(AC)
(= Ep E graa!!
Vol(o. .0)
(86)
((0(pp) 0, (p)
where ((9 denotes the random samples used level As - 0o, G. (( (0. (F)
In this framework, the gradient estimator is constructed by combining estimates across multiple approximation levels, where higher levels yield lower bias but incur greater computational cost. We now verify that this family of approximations satisfies the key MLMC conditions:
Proposition B.5 (MLMC Parameter for Sinkhorn DRO), At each iteration t, under Assumptions
and bounded loss and transport cost, the sequence (P!!) (6) satisfies:
Squared bias parameter a: |VP (0) (0))1 S M,2-ol with a (quadratic decay)
2. Variance parameter b: For the difference estimator H. (0, (()) = grad (0, (()) gradk (1- ((,,()))) we have the variance Var(H, (1)) s M,2 with (linear decay)
3. Cost parameter e: The cost of computing H! (0, ()) bounded by C()) S M.2d with C= 1 (linear growth)
Here Ma My and Me are all constants.
Proof. For the bias parameter a, note that (()) involves expectation over De. while VF (0) approximates this with 2 i.i.d. samples. Under standard Monte Carlo convergence results for bounded smooth gradients, we have /VF!) (0) (F) O(2-1), so the squared bias is O(2-21), implying
For variance, the difference estimator is:
H - (0, (0) Gi (il-!) (0, (0)) (2-1),
G.
(87)
where (()) 21 reuses the first 2 samples. This shared randomness ensures that the two estimates are highly correlated, and their difference has variance O(2-1), hence 1.
For cost, both G and G((-1) require evaluating gradients for 2 and 2 samples respectively. Since + 2/-11 = 1.5 2!, the cost is (22)), implying e 1.
Remark 3. The efficiency of MLMC-RT depends crucially on three problem-specine parameters: (i) the bias decay rate (how fast the approximation error decreases with level), (ii) the variance
32

---

decay rate (how fast the variance of level differences decreases), (iii) and the cost growth rate (how expensive each level is evaluate). The above proposition establishes that, for our Sinkhorn DRO setting, the objective gradient satisfies the standard MI.CC assumptions with parameters (a, b, c) = (2, 1,). This result is key to applying ML.MC-RT effectively and proving convergence with nearly optimal sample complexity.
B.4.2 MLMC-RT Gradient Estimation
The stochastic gradients grad, employed in our optimization procedure computed using Multi- Level Monte Carlo estimator with Randomized Truncation (MLMC-RT), following the methodology developed in Hu al. (28).
Since computing exact expectations over Q7 is intractable due to nested expectations and normaliza- tions, MLMC-RT approximates the gradient via:
Randomly selecting level / (prob")), where prob( 2 2-(++))1/2
L} according to probability distribution
Computing a reweighted difference estimator based on samples level
This randomized truncation technique effectively balances bias and variance, while significantly reducing computational cost compared to fixed-level or naive methods.
The resulting estimator grad, satisfies E[grad,] (=)( rather than o ((,,), hence it is biased estimator of the true gradient. However, by carefully choosing L I 1og(44,, -1), ensure that this bias remains controlled and does not impair the convergence of our algorithm.
The detailed procedure for MLMC-RT in setting is as follows:
The MLMC-RT estimator enjoys the following theoretical guarantees:
Lemma B.6 (Properties of MLMC-RT Estimator). Let grad, be generated according to Algorithm Then:
1. Controlled Bias: |E]grad]] FF & VMas
2. Logarithmic Variance: IEll gradg Blgrad,]] 3| < K log(E-)) some constant K > 0,
3. Efficient Computation: The expected computation cost per iteration is O(log(E-)),
Proof. The controlled bias follows from truncating at level L, which ensures [VE (0)
F. < M,2-aL Mae2, giving the bias bound [[E[grad]] (F,( MMe.
For the variance, with b c and (ro 2-1, we have:
Var(He (i)
Billgrad, BBgradd]]l = prob (1)
(88)
(prob(1)a
1=0
M,2-bi
prob
(89)
(prob(U)2
We set prob(!) (1)
2-1
2-
(which approximately 2-1-1 for large L.). Substituting:
Ln 2
1-2 ior
2
M,2
E Blllgrad, EEgrad, <
(L+1)
(2-1/(1 - 2 (1+11))
Ma(1 (4+11)
Σ
1= M1 2- +1)
1=1
O(L) = O(lo((e-))
(90)
33

---

This is consistent with the results from 228 Table 1, for the MLMC-RT estimator with a = 2. b
We now analyze the computational cost of generating the MLMC-RT gradient estimate at each iteration. Let W denote the random variable representing the computational cost incurred at iteration t, corresponding to the number of samples used to compute the stochastic gradient grad,
Since the MLMC-RT procedure selects random level according to the distribution (prob("), prob(!) ..... ...,,pob(()) prob("), proportional to 2', the expected computational cost per iteration is given by:
and computing the estimator at level requires total cost Wi
B[W.] Σ FroboWe S Me Σ
pr
(where and W S M22')
2-
2-1
= Me Σ prob (1) = Me
Σ
2 2 (L+11
(since prob
1 - 1-2 2-1(++1)
1=0
1=0
(M(L+))
= 1 2-2-1(+1) = ((log(E-)))
(91)
Thus, the expected per-iteration computational cost grows only logarithmically with the desired precision, which highlights the efficiency of the MLMC-RT approach.
Remark 4 (Comparison with Standard SGD). The MLMC-RT estimator achieves logarithmic growth of variance, significantly better than the O variance growth under naive stochastic gradient descent (SGD). This improvement critically reduces the overall sample complexity while maintaining bias control.
B.4.3 Main Convergence Theorem
We now establish the convergence guarantees of adaptive Sinkhorn DRO algorithm under MLMC- RT gradient estimation.
Theorem B.7 (Convergence of AdaDRO). Suppose Assumptions 7 hold, and the iterates are bounded as ||0el] < M. Using the MLMC-RT gradient estimator with maximum level L = [o)g and probability distribution prob(t) (1) x 2-1, and setting the step size a V4I for some constant a > 0, we have that after T (E (=loogo !) iterations, Algorithum Sariaammsatii satisfies
min E TCT
FI 11 =,,
with total sample complexity of log !)
Proof. Our proof integrates both the adaptive dynamics of our algorithm and the properties of MLMC-RT gradient estimation.
Step 1: Descent lemma with adaptive objectives. By the La-smoothness of Fr (Lemma B.2
Løa
10(0++1) S (,,).) (v
g|
(92)
2
To account for the changing objective function, we add and subtract Fe+ (0++1) to obtain:
Fet1(044) F(,O,) ar an (V:))))) grad,) + Loog
2 |grad,1 2 (F1+1(0041) F((0+11) (93)
Step Bounding the adaptive drift. Let M₁ == sup, (e(0, 9)1. by Theorem B.4 we have:
(F41(0+41) - FF((++1) & M2 - Q lrv
Ta |0 - 0dll
< M Gl - 0ell +
2
Mo
= =Maaayllrrddll
|grad,1?
(94)
2
34

---

Step 3: MLMC-RT bias-variance decomposition. For the key term B((VFr(0:), (rradd we apply the birs-variance decomposition:
E [VV()))))
+ Mfgrad,]] ())(0.))
From Lemma B.6 yielding:
we have
E [lgrad,]] - VF:((:)) - VMac, by
inequality,
(V) grad,)) = (VF,))))({+((((((,))EEEgaaa VF.((.))
z
||lll-||Bgggarr [[[gradd]] V (0.)l
(F.())))1 VMae [V)())))
(95)
For the recond-momedd term, using the variance bound from Lemma B.6
E Elllgrdd,1 || = |[[grad,]|| + gllggrd, Elg Elgrad,11(
s (IIV))(0))) + e2 + K log(E-1) < ( + 2M + K log(c
(96)
Step 4: Combining the bounds. Combining the estimates from Steps 1-3, and pluging (94).
995
and 99 into (93) gives
P (0++
[.))) - 0r E V [= + Ma (F,,).
Loof T Igrad, |2] GMa,,, Fllggrad,1] + La M
B Bllgrad, (97)
2
Substituting the bound for B ggrad || from 96
IF| Fi+1 (0c-1)) < F(()( - ar E VE Maage (1)
+ Lo M1)
(2)) + 2M + Kloge )) + GMa H[lgradd11]- (98)
Bias cross-term. Using Young's inequality ab < fa² + b² with (V VMa Mare [VF. S 0t Tel ( + Maate2
(99)
Adaptive-drift linear term. First apply Jensen: E Elgad,, || s grad ?: then plug (96) and use Young's 's inequality again to obtain
GMaa Bl|grad,
< GMaa (V2RV + 22,,² + Klog -
01
[E]V + 20, (GMM3 + 2M, + K Klgge
(100)
Quadratic term from smoothness. For the quadratic terms involving [[E(())) in choosing small enough such that a; Lo (1 + M₂) < for all f, we have
@t (VF)2))) Lg(1 M, 2 (..
(101)
Combining all pieces
999
101
into
98
gives
E [[F++1(0++1] < E
(())))
04 E
F, (0))) + (aae?) Oaglogg Oa log =!),
4
(102)
Step 5: Telescoping. Summing 102) for O..... - and setting AF = E (F)( info F(0).
we have:
T-1
T-1
T-1
T-1
Σ
o. B/ s AF + O
ax)
Alogg
Σ 1=0
ax) Ologg e
Σ a?)
t=0
1=U
35

---

Since a, V++I we know Σt= a2 = =(VT) and Ei= ai = (log T). Because Σ ar
2 VT, we deduce:
E ( & 0
(
AF VT
+ 0 (
O
loggelyy TVT
+ O
log c- vT
(103)
T
VT
1=
Choosing T = loge log
ensures
c log log ( (oge=)), and loge = O(e*). Hence TloTT
min 0steT
EI]EE(0))) s
Each iteration costs Olloge samples (Lemma B.6). so the total sample complexity is
Ore 1)
The sample complexity O -lg-+ (e-1)) matches the lower bound [?(=*) for general non- convex stochastic optimization only by logarithmic factors.
log"
C Verification of Technical Assumptions
To complete our theoretical analysis, we verify the key technical assumptions required for our convergence guarantees.
Lemma C.1 (Lipschitz Continuity of Semantic Costs). Under mild conditions on the encoder network and kernel function, both feature transport cost O and label transport cost C³ are Lipschitz continuous with respect to 0.
Proof. For the feature transport cost Cg (z), = Ker(gg(az).go(1) we assume:
(i) The encoder go is Lipsichs continuous with respect to parameters 0:
196, (z) (x))l 1 021l Vx E X
(104)
(ii) The kernel function Ker(-,)) is LK -ipschitz continuous in both arguments: [Kr((1, w1) Ker (2,, w22l - Lx - 22ll + ||w wzll)
(105)
These are reasonable assumptions for common kernel functions like RBF kernels or dot product kernels with bounded inputs.
For the feature transport cost:
IC (x,, 2;) Cit (2,,, Ker(ge, (2;), go, (x,)) Kor(go, (z,); go. (x,))) LK(llge, (xx) 90, (I;) 1199, (I₂) ge, (I1)
< LK (LgL 021 + 1p1101 - 02ll)
2L Lgg01100 021
(106)
Thus, Og is Lipschitz continuous with constant Lox 2L
For the label transport cost RR. Ker(oR, ow) where k), similar argument holds. Since wn subset of the parameters we have:
(the classifier weights for class
||1 w2x < ||0 021l
(107)
Therefore:
(k,,(,K)) (k,)'))) [er(WIR, wix) Ker(wzx Waxr)
S LL w2x |l + || w2r I|)
s 20 02l
(108)
36

---

Thus, Oj is Lipschitz continuous with constant Lov = 2L
The combined transport cost Cp(p,, = Og (I 2ÿ) 0g (yt, Mj) is therefore Lipschitz continuous with constant Lc = 22 (Lgg 1).
Lemma C.2 (Matching Quality Lipschitz Continuity). Under suitable conditions the encoder network and kernel function, the matching quality function To(P, = Te 9) Lng-Lipsciit continuous in 0.
Proof. We fix p and view (o q) as a probability density over qi
Tolp. 9) exp(-Co(p.a)/o) v(g)
(109)
Zo(p)
Let f(()) = -Cn, (p.
g)e. g4) -02 (p. q)/c. Then, Ya, (v, .) for i = 1,2 defines probability
density over (2, v) the form dl:(g) = 2l55)
dv(g). By Proposition B.3
1:0, (p. yos (p, Vlrv S Vllf - ol10 / Col(p,)) Co. olo
(110)
From Lemma C.I
Ce Li-Lsschiiz in 0, so
Le 100 - 02l
1he, (p..) yo, (p. .lTv $
(111)
This implies pointwise Lipschitz continuity:
Le
170, (p. 4) - 70, (p. S (p,)) ~00, .lrvv S Hence, 7o (p, q) Lipssipschitz in 9 with I,, = VTc.fc.
10: //
02ll
(112)
Lemma C.3 (Threshold Lipschitz Continuity). Under the conditions established in Lemma threshold function T(q; 8) is L, Lipschitt continuous in 0.
C.2 the
Proof. The threshold function is defined as:
<(:: 0) = (0) T2( 0)
(113)
max =2
where (0) = E(pa) yo (p, 9) and T2 0) = (p 9)
From Lemma
C.2
7o (p. q) is L,, Lipschitz. Therefore:
(01) Ti (02) [ (re, (p, g) yo, (p. s (pp 11,0, (p. yo, (p.9))
I 1 021
(114)
Similarly, 72(9; 0) is Lipssiip,h and: max (2(4 (1) max T2(9':00) & max [2((': -220''=2)) & L, 021l
(115)
Let 0) max, r2lg':0)' Under regularity conditions ensuring that (2( is bounded and max 72(9': 0) is bounded away from zero, we can establish that r(g: 0) is -ipschith
=2(g:0)
Now consider the full threshold function:
[r(q: 01) T(::22)
In (())r((;))))))))())))))))))))))) (1 (02)r(g; 02)
$ |n (0:))lr(q r(q: 02) (2)g:1)))11 (02)
1, 101 - 02ll
(116)
where L, = = -
These Lipschitz continuity properties ensure that small changes in model parameters lead to bounded changes in both the matching quality function and the threshold function, which is essential for the stability and convergence of our adaptive robust optimization algorithm.
37

---

D Experimental Results
We evaluate the proposed AdaDRO framework across diverse set of benchmarks assess its effectiveness under various types of distribution shifts. Here is the code
D.1 Experimental Setup
Datasets. We evaluate on three widely studied distribution shift settings: Colored MNIST 2
which tests robustness under spurious correlations; Waterbirds [50], real-world dataset with strong background-labed correlation; CelebA 43 a benchmark for facial attribute recognition. We also
evaluate several long-tailed benchmark datasets: CIFAR-IO-LT and CIFAR-1O0-LT
(37
Baselines. We compare AdaDRO with the following well well-known methods: KL-DRO [46]. Wasser
stein DRO (WDRO) (9), Sinkhorn DRO SDRO [62). unified DRO method UniDRO [II. semantic- aware robust methods such as GroupDRO [50]. an invariant risk minimization method IRM (2]
For fairess,, we report the average accuracy and the distribution shift gap (drop in performance under shift). All models are implemented with Py Torch on single NVIDIA RTX 6000 Ada GPU. The model is optimized using the AdamW optimizer [45]. For AdaDRO, we use cosine similarity as the kernel (TO), and employ basic augmentations (flip, crop) for semantic calibration Sec. 3.2
D.2 Robustness to Spurious Correlation and Group Shift.
AdaDRO broadly applicable thanks to flexible reference distribution design. In imbalanced settings, can be constructed via reweighting resampling. In other tasks, domain priors can be casily integrated, similar to many specialized robust methods. We evaluate AdaDRO across diverse scenarios, including spurious correlation, group shift, and long-tail generalization.
Spurious Correlation. On Colored MNIST, AdaDRO outperforms Sinkhorn DRO and IRM by 7.3% and 4.6% (on worst-group) respectively in worst-group accuracy, as shown in Table This demonstrates its ability filter misleading correlations during training. Notably, when the correlation between color and label is strong (bias = 0.9), all other DRO methods suffer a significant drop except AdaDRO, which effectively mitigates the over-pessimism nature of DRO methods.
Real-world Group Shift. On Waterbirds, AdaDRO significantly improves the worst-group accuracy (+5.1%) over Sinkhorn DRO and matches GroupDRO, which is specialized for group-shift scenarios. Table show the results on Waterbirds.
Table Accuracy (%) on Colored MNIST under high spurious correlation. AdaDRO achieves the highest robustness across methods.
Method
Avg Accuracy
Worst-group Ace
Gap
IRM
79.5
75.3
4.2
GroupDRO 50
78.4
73.0
5.4
KL-DRO 46
74.9
68.5
6.4
WDRO 9
75.8
69.1
6.7
SDRO 62
79.0
72.6
6.4
UniDRO
67.1
47.8
20.7
AdaDRO (ours)
87.2
79,9
7.3
Table 6: Accuracy (%) Waterbirds. AdaDRO is competitive with GroupDRO while outperforming
others.
Method
Avg Aecuracy
Worst-group Acc
Gap
IRM 2
82.6
73.2
9.4
GroupDRO (50]
86.3
80.5
5.8
KL-DRO 46
81.0
71.5
9.5
SDRO 62
85,0
75.4
9.6
AdaDRO (ours)
87.1
80.5
7.0
38

---

Long-tail Generalization. We conduct our experiments on two widely used long-tailed tailed benchmark datasets: CIFAR-IO-LT and CIFAR-100-LA These datasets are derived from the original CIFAR-10 10 and CIFAR-100 100 by artificially reducing the number of training examples in tail classes to simulate long-tailed distributions. The imbalance in these datasets quantified by the imbalance factor (IF), defined as the ratio between the number of samples in the most frequent class and that in the least frequent class. We evaluate our model under three imbalance levels: I0=10, IF=50, and I00, representing increasing levels of class imbalance.
Baselines. We compare AdaDRO with diverse set of baselines tailored for long-tailed learning: Decouple [7]: Decouples representation learning from classifier training, using balanced sampling in the classifier phase: Focal Loss [41]: Reweights the loss to focus training on hard-to-on card-to-classify
(often tail) examples; DRO-LT [52] A distributionally robust optimization framework specifically
35 A contrastive learning- based method that enhances representations for tail classes via augmented positive samples; Resample [[5] Performs oversam- pling of minority classes to reduce imbalance; Reweight [15] Adjusts sample weights in the loss function based on inverse class frequency; SinkhoraDRO [62]: A recent OT-based DRO method that regularizes transport cost to mitigate over-pessimism in robustness.
adapted for long-tailed classification; SSL
Results and Insights. As shown in Table our method consistently outperforms all baselines across different imbalance settings both CIFAR-IO-LT and CIFAR-100-HE AdaDRO achieves the highest accuracy under all imbalance factors, especially under the most severe imbalance (IF=100), where it shows substantial gains over both traditional and robust basclines. This highlights the advantage of our adaptive distributional robustness formulation in focusing learning underrepresented regions while maintaining calibrated semantic alignment. ddtalls the quantitative comparison between our method and the strong baseline ERM+Aug across standard DomainBed benchmarks [25
Table 7: Accuracies ResNet32 on long-tailed CIFAR-10 and CIFAR-100 datasets with different imbalance factors (10, 50, and 100).
CIFAR-10-LT
CIFAR-100-LIT
Dataset
100
100
Decouple
2
70.4
76.2
86.4
41.2
46.8
57.9
Focal Loss
70.3
76.7
86.6
38.4
44.3
55.7
DRO-I.T
52
73.7
77.2
86.9
45.4
55.3
61.2
SSL 35
67.3
75.4
86.5
37.5
44.0
56.7
Resample
16
66.5
74.8
86.4
33.4
43.9
55.1
DR 12
65.7
74.5
84.7
38.2
49.3
56.3
AdaDRO
74.2
78.2
87.9
46.4
56.6
64.8
Table 8: Performance on ImageNet-LT. We report accuracy on three class splits: many-shot (more than 100 images), medium-shot (20-100 images), and few-shot (fewer than 20 images), which follows the settings in 33 44
Method
Many-shot
Medium-shot
Few-shot
ERM+Aug
41.8 + 0.2
24.7 + 0.4
11.5 + 0.6
Class-Balanced Loss
40.8 0.3
31.1 1 0.3
20.8 1 0.5
DRO-LT
42.2 + 0.2
32.6
22.5 0.4
Sinkhorn DRO (SDRO)
40.1 - 0.3
2..5
21.4 0.5
AdaDRO (Ours)
3.34
27.3 27.3000 0.4
AdaDRO w/o Adaptive Filtering
41.5 + 0.2
0.3
21.7 0.4
39

---

Table 9: Performance comparison across different domain generalization benchmarks.
Algorithm
CMNIST
RMNIST
VLCS
PACS
Office-Home
TerraInc
ERM+Aug
52.0 0.1
98.0 + 0.0
77.4 0.3
85.7 7 0.5
67.5 + 0.5
47.2 + 0.4
Ours
63.8 = 0.3
98.0 = 0.1
77.8
86.2 = 0.6
67.1 + 0.3
50.2 + 0.5
D.3 Efficiency Analysis.
Tables and due semantic alignment and filtering, remains tractable and justifiable given the performance gains.
and [1 report normalized total training time. While AdaDRO incurs moderate overhead
Table 10: Normalized total training time (relative to IRM) Colored MNIST and Waterbirds.
Method
Colored MNIST (Norm.) Waterbirds (Norm.)
GroupDRO
1.00
1.00
IRM
1.12
1.06
KL-DRO
1.16
1.14
WDRO
6.24
6.22
SDRO
1.25
1.27
UniDRO
6.03
5,49
AdaDRO (ours)
1.31
1.34
Table 11: Normalized total training time (relative to Decouple) on 11FAR-I0-LT and CIFAR-I00-LIT
Method
CIFAR-I0-LT (Norm.)
CIFAR-100-LT (Norm.)
DRO-LT
1.00
1.00
Focal Loss
0.84
0.79
Decouple
1.07
1.03
SSL
1.05
1,02
Resample
0.95
0.93
SinkhornDRO
1.12
1.19
AdaDRO (ours)
1.28
1.26
D.4 Ablation Studies
We perform ablation studies to assess the contribution of key components in AdaDRO. Detailed results are shown in [2bles[[2]11
(i) Semantic Calibration. We disable the semantic calibration objective by replacing the inverse OT formulation in (12) with a simple Euclidean distance cost. This leads to significant drops in both clean and worst-group accuracy, confirming the importance of semantically aligned transport.
(īi) Adaptive Filtering. We replace the filtered reference distribution i; with a fixed one, either generated by Mixup (67] uniform sampling. Without filtering, performance degrades due irrelevant or noisy samples in v.
(iii) Robustness to Uncertainty Radius. We compare AdaDRO with Sinkhorn DRO under varying uncertainty set radius X. As shown in Table [4] AdaDRO maintains stable robustness even under aggressive uncertainty, whereas Sinkhomn DRO suffers from severe over-pessimism or instability.
D.5 Limitations
Our method relies on effective data augmentations to serve as prior constructing the distributional uncertainty set. However, the efficacy of this prior is task-dependent and lacks a universal standard. Our method also relies on constructing the ground-truth coupling 7 through semantic invariant data augmentation. In some domains, richer priors may be available. Effectively integrating such
40

---

heterogeneous sources of prior knowledge to build a more powerful coupling remains an open question.
Table 12: Impact of removing semantic calibration. Performance significantly drops without IOT- based feature alignment.
Variant
Clean Acc
Worst-group Acc
Full AdaDRO
84.2
79.9
w/o Semantic Calibration
80.6
73.7
Table 13: Adaptive filtering improves robustness by focusing the uncertainty meaningful support.
Variant
Clean Ace
Worst-group Acc
AdaDRO (adaptive)
84.2
79.9
Fixed (Mixup)
82.1
75.0
Fixed v (Uniform)
81.4
74.3
Table 14: Robustness comparison of AdaDRO and Sinkhorn DRO under varying on Colored MNIST. Higher À corresponds to smaller uncertainty sets. AdaDRO maintains stable performance, while Sinkhorn DRO suffers significantly as À increases.
AdaDRO
Sinkhorn DRO
A
Worst-grous group Ace (%)
Average Acc (%)
Worst-group Acc (%)
Average Ace (%)
0.5
71.0
81.2
62.1
74.3
1.0
76.4
84.0
69.1
75.8
2.0
79,9
87.2
72.6
78.6
5.0
78.7
85.5
72.3
79.0
41


Reviews:
[{"review_id": "OHdzn9wXJk", "reviewer": "Reviewer_Z7dt", "summary": "The authors introduce a novel DRO framework where (semi-relaxed and entropy-regularized) Optimal Transport is utilized together with (semi-relaxed) Inverse Optimal Transport during the DRO optimization to find a robust model/classifier. It is argued that the ability of such models to withstand various realistic distribution shifts, while showing high source accuracy, depends critically on the support \\nu of the computed (uncertainty) distributions Q. The authors then, inspired by the proposed Inverse Optimal Transport which refines the computed couplings of the semi-relaxed OT step, propose an adaptive filtering operation to refine this reference distribution \\nu and hence shape the geometry of the computed Q_{\\lambda} distributions. Results are shown, over various DRO benchmarks, where it is shown that the proposed method can improve over state-of-the-art approaches, such as Sinkhorn DRO.", "questions": "- It would be good to expand on the Figure 1 and illustrate why and how the proposed method can find better / more meaningful / realistic uncertainty sets (through semantic calibration / adaptive filtering etc.)\n- It would be clearer in the discussion that follows eq. (1) to mention that \\nu forms the support of the distributions Q.\n- \"This choice is both practical and theoretically sound, as the relative entropy term then differs only by a constant when Ptr ≪ μ leaving the optimization unaffected.\"\nThis sentence was not clear to me.\n- The paper becomes sometimes quite confusing to read and to understand in section 3, partly due to the fact that that the authors do not separate the problem formulation from the contributions/to be introduced method. For instance the 'Distribution Exploration via OT' paragraph can be streamlined with section 2 where the authors introduce the derivation of eq. (6) from Wang et al. [58].  The following Lemma 3.1 should also be mentioned to follow from [58].  Also, the 'classification as inverse OT' part in section 3.1 is not clear as suddenly the OT problem changes from matching p(x) and q(x) to p(x) and y.\n- line 157: lambda depends on delta, which is not mentioned.\n- It would make section 3 clearer if Remark 1 is explained more explicitly: show formally how (1) can be cast as OT + IOT mentioned in section 3.1\n- It is not clear what is the benefit of emphasizing 'semantic calibration' with 'inverse OT' over plain source risk minimization using cross-entropy. As one trains the model f_theta on the distributions Q, the parameters theta changes the OT costs and there doesn't seem to be any additional benefit of formulating this problem as an 'inverse OT'. For instance the 'calibration' of costs via IOT in (11) and (12) can I think be equally achieved by ERM that includes augmentations to learn a better model g_augmented, which would then improve the semantic quality of costs C(g_augmented(x_i), g_augmented(x_j)) over C(g(x_i), g(x_j)).\n- \"Leveraging this insight, we define the label-level transport cost as...\" The preceding discussion for motivating (14) is not clear for me.\n- Given the heuristic nature of the introduced filtering operations for \\nu in (15), (16) I have a hard time understanding the reasoning behind them. Since this part seems to form the core of the paper's contributions, e.g. over [60]: \"Outlier-robust distributionally robust optimization via unbalanced optimal transport\", it would make sense I think to explain, motivate and analyze them, as well as the connections to other thresholding strategies such as FreeMatch, in much more detail. When do these heuristics work well and when would they not for instance?\n- In the discussion that follows Thm 3.2, it would be nice to mention similar results for other DRO algorithms (e.g., ones compared to in the experiments). Is Thm 3.2 an extension of such results for dynamically evolving uncertainty sets?\n- It is not clear why Algorithm 1 needs step 1 [train \\theta by the joint objective in (12) and (13) for N0 epochs] before the for loop.\n  \nExperiments\n- Since they include augmentations, the authors should compare against ERM with augmentations included. The other methods compared against should then also include augmentations for fairness.\n- I suspect that the proposed method may not beat the 'ERM with augmentation' baseline across various distribution shift scenarios. See the \"In search of lost domain generalization\" paper and the DomainBed platform for why I think so. Would it be possible to try more algorithms or datasets from the DomainBed repository?\n- GIW is not evaluated in the experiments, nor should it be: in DRO one does not assume unlabeled test/target distribution to be available and domain adaptation methods such as GIW require such data to be available. However it is mentioned as if it is compared against.\n- The authors should describe the distribution shift in each scenario: e.g., in colored MNIST, the 'worst group' do not contain the colors correlated with the labels, I believe. What is the difficulty in the Waterbirds dataset? Do the spurious correlations (background to label correlation) change in the 'worst-group' distribution?\n- In general the experiments section, even when taken together with Appendix D, is quite inadequate to appreciate the contributions. The authors should streamline parts of Section 3 with Section 2 to create more space in the main text, so that they can significantly widen the\ndiscussion of the experiments section.\n- Hyperparameters set during the experiments should be mentioned in particular, e.g., what is the value of the uncertainty set radius \\delta\n  in the experiments? Is it varying or constant?\n- \"Notably, when the correlation between color and label is strong (bias = 0.9), all other DRO methods suffer a significant drop except AdaDRO, which effectively mitigate the over-pessimism nature of DRO\" Which experiment is this? Show the quantitative effect.\n- Report error bars, at least in the appendix.\n- \"We conduct ablation experiments to isolate the effect of AdaDRO’s technical components, please refer to the supplementary materials\"\nPlease include all relevant material in the Appendix.\n- \"Full experimental details and additional results, including the visualization of the dynamic \\nu_t sensitivity analyses, and runtime statistics, are available in the supplementary materials.\"\nPlease include some visualizations and other analyses/statistics in the appendix / experiments section also. Ideally, there shouldn't be any need to include a supplementary material besides the appendix.\n- How do the different imbalance factors effect the CIFAR-10/100-LT experiments?\n- Which experiment/dataset does Table 6 correspond to?\n- In Table 8, \"Higher λ corresponds to smaller uncertainty sets. AdaDRO maintains stable performance, while Sinkhorn DRO suffers significantly as λ increases\"\nHigher \\lambda should correspond to bigger uncertainty sets (bigger radius \\delta) I believe. The conclusion also seems unwarranted: according to the results, there is similar variance between the two methods.\n\nMinor comments:\n- W matrix is R^{d_e \\cross K} instead\n- note that eq. (2) is the KL-divergence between \\gamma and \\mu \\cross\n  \\nu whereas eq. (3) is the \\emph{negative} entropy H plus some\n  constant. Hence it would be more appropriate to use KL-notation for\n  eq. (2) and show that it is equal to the negative of Shannon entropy\n  H(.) up to a constant.  Minimizing the KL-divergence in this case\n  maximizes the Shannon entropy and vice versa.  Regularizing with the\n  KL-divergence of a coupling likewise inflates the entropy of the\n  coupling, depending on the regularization constant.\n- not a big fan of the p,q notation for (x,y) positions as they're easily\n  confused with the upper case notation for the distributions\n- line 151: add space after 'following'\n- 'Geometric Awareness' is not a benefit of the entropy regularization\n  in D_{\\epsilon} but rather a 'benefit' of the OT-formulation.\n- The inputs xi and yi generally live in different vector spaces (with\n  diff. dimensions), so the coupling gamma_bar in (8) is not standard\n  for Wasserstein distances (Gromov-Wasserstein OT considers points\n  that can live in different spaces.) You'd need to introduce the\n  operation of the model f_theta: x -> y_hat with parameters theta\n  inside Q^{\\lambda} to make it compatible with standard OT\n  formulation.\n- N0 is also a required hyperparameter of Alg1", "limitations": "\"If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.\" \nIt would be good to briefly mention the positive societal impacts.", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 1, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The authors identified the underlying dilemma of the DRO problem quite correctly and proposed an interesting solution to solve this important issue, which prevents the wider deployment of DRO methods in practice (as the proposed solutions are either too conservative and have poor source accuracy or are unable to capture the relevant uncertainty set). However, the proposed solution is a heuristic that is not explained or motivated carefully. I also found some parts of section 3 to be unclear in some aspects (as explained below in detail), which prevented me from appreciating the proposed method and the resulting Algorithm 1. Equally importantly, the experiments sections was quite limited and the results were not explained carefully. A much more detailed experiments section with more results (over more distribution shift scenarios) and ablations, would make a more convincing case. Please find below some detailed comments concerning both these issues. \n\nAs of now I cannot advocate for acceptance of this paper in its current form, but if the authors can answer my concerns adequately I'd be happy to raise my score.", "was_revised": true, "final_justification": "I thank the authors for their carefully prepared rebuttal. After going through the response, I have no major concerns and advocate for the acceptance for this paper. I have raised my overall score (and confidence) accordingly. \n\nA few issues still remain, although they are not as critical:\n\n- ERM with augmentation is still competitive or beating the proposed algorithm in the DomainBed benchmark that is posted by the authors, except for the Colored MNIST experiments. Given the response (to Q1-Q2) where the authors claim that their method would be better due to (i) label noise, (ii) label-agnostic calibration, I infer that (i) label-noise is probably not prevalent in a lot of these 'curated' datasets (is that so?), and that (ii) label-agnostic calibration indeed helps to avoid spurious correlations in a way that ERM with augmentation cannot. I would like the authors to discuss these experiments and these (or similar) conclusions in more detail.\n- The heuristic for the nu-filtering should be explained better and integrated with Figure 1 (which the authors kindly agreed to do). It would be quite nice to also discuss potential problems with this heuristic.\n- It is not clear from the main text what is the precise novelty of Theorem 3.2 and what is the related work. The contributions of the paper towards this result should be written and highlighted clearly in the main text.", "rebuttal": "We thank the reviewer for the thoughtful comments, and address your concerns below.\n\n**Q1 - Q6:**\n\nWe sincerely thank the reviewer for their detailed and constructive feedback on the manuscript's presentation. We agree that the clarity and structure can be improved.  In particular, we will enhance Figure 1 by including additional visualizations to better illustrate how our semantic calibration and adaptive filtering contribute to the formation of more meaningful uncertainty sets. We will also revise the text to explain Remark 1 more explicitly. Furthermore, we will incorporate additional details as suggested. \n\n**Q3:** \"This choice is both practical and theoretically sound, as ...\n\n$H(\\gamma | \\mu \\otimes \\nu) = E_{(p,q)\\sim\\gamma} [ \\log(\\frac{d\\gamma(p,q)}{dP_{tr}(p)d\\nu(q)}) + \\log(\\frac{dP_{tr}(p)}{d\\mu(p)}) ]  = E_{\\gamma} [ \\log(\\frac{d\\gamma(p,q)}{dP_{tr}(p)d\\nu(q)}) ] + E_{p\\sim P_{tr}}[\\log(\\frac{dP_{tr}(p)}{d\\mu(p)}) ]$. Thus \nany choice of $\\mu$ satisfies $P_{tr} \\ll \\mu$ yields an equivalent entropy term up to an additive constant, and does not affect the optimization. \n\n\n**Q4:** ... the 'classification as inverse OT' part... the OT problem changes from matching $p(x)$ and $q(x)$ to $p(x)$ and $y$.\n\nOur contribution in Lemma 3.1 is to provide a key reinterpretation for DRO formulated in (1): we show that the worst-case distribution (derived in [1]) can be understood as the target marginal of a specific semi-relaxed OT problem. This perspective forms the foundation for our subsequent components—semantic calibration via inverse OT (IOT) and adaptive filtering—both of which exploit the geometric structure of OT. \nRegarding the discussion of ‘classification as inverse OT’, we adapt the IOT framework between the source distribution of training data $p(x)$, and the target distribution over labels $p(y)$. This is a necessary step for modeling classification as a transport-based matching. \n    \n\n**Q7:** ... the benefit of emphasizing 'semantic calibration' with 'inverse OT' ...\n\nAs described in lines 102 and 209, the transport cost between two samples is decomposed into two parts: a sample-level cost and a label-level cost, \n$C(p_i, p_j) = C^X(x_i, x_j) + C^Y(y_i, y_j)$.\nOur semantic calibration framework explicitly addresses both components using the IOT objectives (12) and (13), respectively. We reinterpret the cross-entropy loss through the lens of IOT and show that it coincides with our objective in (13). Importantly, this perspective does not complicate the formulation, and provides a deeper understanding of its function within our IOT framework. \n\nIn contrast, standard ERM does not provide well-founded mechanism for learning the sample-level cost $C^X(x_i, x_j)$, and thus cannot guarantee that the learned feature space reflects meaningful semantic geometry—even with data augmentation. \nTo address this, we introduce the IOT objective (12), which explicitly learns $C^X$ by recovering latent semantic distances from observed matching pairs. This approach helps to extract $C^X$, enabling the semantic alignment that cannot be achieved by ERM alone. \n    \n\n**Q8:** ... preceding discussion for motivating (14) is not clear for me.\n\nWe will revise this section to improve clarity. The high-level idea is as follows: standard cross-entropy training is equivalent to an instance of IOT, where the training objective effectively calibrates the distance defined in (8) between the input embedding $g_\\theta(x_i)$ and its corresponding class weight vector $w_{y_i}$.\nThis process naturally yields a set of weight vectors $\\{w_k\\}_{k=1}^K$ that carry semantic meaning; classes that are semantically similar will have similar weight vectors. Within the IOT framework, these weights induce a transport cost over the label space $C^Y$. \nWe will present this idea more directly in the revision.\n    \n\n**Q9:** ... I have a hard time understanding the reasoning...\n\nThe key idea of the filtering is to use the model's own confidence (i.e., the coupling $\\gamma_\\theta$) as a self-adaptive signal. If the model has difficult to align a sample $p$ with its ground-truth match $q$ (i.e., $\\gamma_\\theta(p,q)$ is low), it suggests that $p$ may be an outlier or noisy sample that is harming the uncertainty set.\n**Adaptive Threshold $\\tau(q)$:** The adaptive threshold is designed to balance global and local confidence. The term $\\tau_1$ is the global average confidence, setting a baseline. $\\tau_2(q)$ measures the \"matchability\" of a target $q$. The ratio \"$\\tau_2(q) / \\max(\\tau_2)$\" adjusts the global threshold locally. For easy targets (high $\\tau_2$), the threshold is higher; for hard targets (low $\\tau_2$), the threshold is lower to avoid incorrectly filtering. The popular FreeMatch [2] used a similar self-adaptive thresholding, and was later also reinterpreted through the lens of OT [3].  \n\n\n**Q10:** ... mention similar results... Is Thm 3.2 an extension of such results ...?\n\nThank you for prompting us to highlight this point.\nIn addition to handling the dynamically evolving reference distribution $\\hat{\\nu}$, our analysis addresses a more general non-convex setting. As far as we know, the most relevant prior result is provided in [1]. While [1] focuses on convex settings, our analysis extends to non-convex regimes with dynamically evolving uncertainty sets, using different techniques. \n    \n\n**Q11:**  ... why Algorithm 1 needs step 1 ...\n\nStep 1 acts as a warm-up phase to ensure stable and reliable signals for use in (12) and (13).\n\n\n**Experiments**\n\n**Q1 Q2: ... 'ERM with augmentation' baselines ...** \n\nThe refereed paper show that ERM+augmentation is a strong baseline. \nWe believe our method is unique in two aspects: 1, its label-agnostic sample-level semantic calibration (13); 2, the filtering mechanism based on the coupling introduced by (12) (13). \n\nAn intuitive example is the presence of noisy labels: when a portion of labels are flipped, sample-level calibration (12) can still reliably capture semantics. This is because the process is unsupervised, similar to contrastive learning [4], and thus remains robust to label noise. Building on this, the coupling-based filtering mechanism can discards samples with wrong label, enabling the cross-entropy to train on a cleaner subset of the data.\nIn contrast, methods that depend on label supervision, including ERM with augmentation, face an inherent \"garbage in, garbage out\" issue. No  augmentation can correct mislabeled examples, and these label errors inevitably propagate through the learning process, compromising model robustness. \n\nA similar advantage is observed in Colored MNIST. Even when augmentation occasionally alters the spurious attribute (e.g., color), standard ERM still learns entangled representations, as the label remains strongly correlated with color. In contrast, our label-agnostic calibration encourages the model to identify truly invariant features—those that remain stable under augmentation—naturally prioritizing shape over color in the absence of a strong supervisory signal. \n\nFor the space limit, we only present results on the DomainBed benchmark.\n\n| Algorithm | CMNIST         | RMNIST         | VLCS           | PACS           | Office-Home    | TerraInc       | DomainNet      |\n| --------- | -------------- | -------------- | -------------- | -------------- | -------------- | -------------- | -------------- |\n| ERM+Aug   | 52.0 ± 0.1     | **98.0 ± 0.0** | 77.4 ± 0.3     | 85.7 ± 0.5     | **67.5 ± 0.5** | 47.2 ± 0.4     | **41.2 ± 0.2** |\n| Ours      | **63.8 ± 0.3** | **98.0 ± 0.1** | **77.8 ± 0.3** | **86.2 ± 0.6** | 67.1 ± 0.3     | **50.2 ± 0.5** | 39.7 ± 0.3     |\n\n**Q3-Q13:** We thank the reviewer for sufficient suggestions.\nWe will clarify these points in the revised experimental section:\n(Q3) GIW was mentioned for context and we will make it clear that it was not used as a baseline; (Q4) For Waterbirds, the spurious correlation is between the background (water vs. land) and the label (waterbird vs. landbird). The 'worst-group' distribution contains samples where this correlation is broken (e.g., waterbirds on land); (Q6) The uncertainty radius $\\delta$ is a hyperparameter, which is controlled by the dual variable $\\lambda$ in practice. We will add its search space and the selected value for each experiment to the appendix; (Q12) Table 6 corresponds to the Waterbirds dataset. We will add a caption.\n\nWe agree with these excellent suggestions for improving the paper's readability and completeness. In the revision, we will:\n(Q5) Streamline Sections 2 and 3 to create more space for a detailed discussion of the experiments; (Q8, Q9, Q10) Integrate all supplementary materials, including ablation studies, visualizations, error bars, and full experimental details, directly into the main paper's appendix to create a single, self-contained document; (Q11) Add a discussion on the effect of imbalance factors in the CIFAR-LT experiments.\n\n**Response to Minor Comments:** Thank you for your valuable comments. We have corrected the typos and made several revisions to improve the overall writing quality of the manuscript.\n\nRegarding the comment on the inputs $\\textit{x}_i$ and labels $\\textit{k}$ residing in different spaces, our formulation does not compute a  Wasserstein distance, which is defined in a metric space. Instead, our approach leverages the more general OT framework, which permits flexible ground distance function that  do not necessarily satisfy the properties of a metric. As long as a ground distance function is defined, an OT distance can be computed.  \n\n[1] Sinkhorn distributionally robust optimization. 2023\n\n[2] Freematch: Self-adaptive thresholding for semi-supervised learning. ICLR'23\n\n[3] OTMatch: Improving Semi-Supervised Learning with Optimal Transport. ICML'24\n\n[4] Understanding and Generalizing Contrastive Learning from the Inverse Optimal Transport Perspective. ICML'23"}, {"review_id": "QKZ7U8KHZy", "reviewer": "Reviewer_JCUz", "summary": "The efficacy of Distributionally Robust Optimization (DRO) heavily depends on the design of the uncertainty set. To deal with this, the paper proposed a novel optimal transport-based framework for adaptive robust optimization, where transport costs and uncertainty set co-evolve through inverse optimal transport.", "questions": "In the paper, the initial reference distribution $\\nu$ is constructed using data augmentation. The reviewer wonders if data augmentation is also used in the compared methods.", "limitations": "While the authors answered 'Yes' to the question 'Does the paper discuss the limitations of the work performed by the authors?', the reviewer did not find a section that discusses the limitations of the work.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\nThe paper proposed a novel optimal transport-based framework for Distributionally Robust Optimization. The framework consists of two core components, i.e., semantic calibration, which bootstraps transport costs, and adaptive refinement, which filters the uncertainty set, and the transport costs and uncertainty set co-evolve during the learning process. \n\nThe paper provides the theoretical convergence guarantees of the proposed method. \n\nWeaknesses:\n\nThe paper mentioned that: “This result demonstrates that robust optimization\nwith dynamically adaptive uncertainty sets and semantic costs can also be performed efficiently at scale.”. However, there are no large-scale experiments (e.g., experiments on the ImageNet dataset) to confirm this. \n\nIt is not clear to the reviewer what the limitations of the proposed method are. \n\nIn the experimental section of the main paper, it is not clear about the encoder g used in the paper.", "was_revised": true, "final_justification": "After reading the authors' response, I decide to maintain my rating.", "rebuttal": "We thank the reviewer for the thoughtful comments, and address your concerns below.\n\n**Response to Weaknesses**\n\n**W1:** ... no large-scale experiments (e.g., experiments on the ImageNet dataset) to confirm this. \n\nWe will add more experiments. \nRegarding scalability, as shown in the Sec 2.4 of Supplement, AdaDRO's runtime is only 1.3x that of standard baselines, while robust methods like WDRO are 5-6x slower. This efficiency stems from our use of a closed-form solution to the optimal transport (OT) problem, eliminating the need for costly iterative solvers. \n\nImageNet-LT\n\n| Method                        | Many-shot      | Medium-shot    | Few-shot       |\n| ----------------------------- | -------------- | -------------- | -------------- |\n| ERM+Aug                       | 41.8 ± 0.2     | 24.7 ± 0.4     | 11.5 ± 0.6     |\n| Class-Balanced Loss           | 40.8 ± 0.3     | 31.1 ± 0.3     | 20.8 ± 0.5     |\n| DRO-LT                        | **42.2 ± 0.2** | 32.6 ± 0.3     | 22.5 ± 0.4     |\n| Sinkhorn DRO (SDRO)           | 40.1 ± 0.3     | 29.5 ± 0.4     | 21.4 ± 0.5     |\n| **AdaDRO (Ours)**             | 42.1 ± 0.2     | **34.4 ± 0.3** | **27.3 ± 0.4** |\n| AdaDRO w/o Adaptive Filtering | 41.5 ± 0.2     | 29.8 ± 0.3     | 21.7 ± 0.4     |\n\n\n\n**W2:**  ...  the limitations of the proposed method.\n\n\nWe will include a discussion of the limitations in our revised version. A key aspect we will elaborate on is the construction of the ``ground-truth'' coupling, $\\bar{\\gamma}$, which is central to our semantic calibration via Inverse OT.\n\nOur approach use a practical strategy for constructing $\\bar{\\gamma}$ by leveraging data augmentation, which implicitly defines a semantic correspondence between a sample and its augmented variant. However, we acknowledge that this represents a relatively simple means of encoding semantic priors. A central challenge lies in how to optimally incorporate all available prior knowledge to define a task-specific observed coupling. In some domains, richer priors may be available—such as explicit class hierarchies, textual descriptions, or auxiliary information from other modalities. Effectively integrating such diverse sources to construct a more powerful coupling remains a non-trivial and open research question. \n\n\n\n**W3:** In the experimental section of the main paper, it is not clear about the encoder $g$ used in the paper.\n\n\nThe encoder $g_\\theta$ is not an additional module but refers to the standard feature extractor backbone of the neural network used in our experiments. Correspondingly, the function $h_W$ represents the final linear classification layer (or \"head\") that maps the feature representation from $g_\\theta$ to the logits. For instance, in experiments using a ResNet-32 backbone, $g_\\theta$ includes all layers except the final linear classifier, which is represented by $h_W$.\n\n\nThis standard architectural decomposition ensures that our method does not introduce any additional parameters or architectural modifications beyond the baseline models, ensuring a fair comparison. We will include these implementation details explicitly in the revised manuscript’s experimental setup section to improve clarity and reproducibility. \n\n\n**Response to Questions**\n\n**Q1:** In the paper, the initial reference distribution is constructed using data augmentation. The reviewer wonders if data augmentation is also used in the compared methods. \n\nTo ensure a fair and rigorous comparison, all baseline methods were trained using the data augmentation protocol (e.g., random crops, flips) as our proposed method. However, different methods may prefer different augmentation strengths, and we report the best performance achieved for each. We will add a clarifying sentence to the experimental setup section to make this explicit and avoid ambiguity."}, {"review_id": "qbYmXx6NjW", "reviewer": "Reviewer_sEAy", "summary": "This paper proposes a new framework that reformulates entropy-regularized Wasserstein DRO as a dynamic optimal transport process. The key contributions are: (1) semantic calibration that learns meaningful transport costs via inverse optimal transport (IOT), and (2) adaptive refinement that dynamically filters the uncertainty set using OT-based feedback. The method addresses limitations of existing DRO approaches that rely on static, potentially inappropriate uncertainty sets and manually designed transport costs.", "questions": "1. How does the OT computation scale with high-dimensional data and large sample sizes?\n\n2. Can you provide more detailed hyperparameter sensitivity analysis across different datasets?\n\n3. What is the actual wall-clock time comparison with baseline methods?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Pros:\n\n1. The reformulation of Sinkhorn DRO as a bilevel OT/IOT optimization problem provides valuable geometric intuition and theoretical clarity. Rigorous convergence analysis despite the challenging dynamic uncertainty sets and nested expectations.\n\n2. Semantic calibration via IOT is novel and addresses the long-standing issue of manually designing transport costs. The adaptive filtering mechanism combining global and local thresholds is more principled than fixed approaches.\n\n3. Consistent improvements across multiple standard benchmarks (Colored MNIST, Waterbirds, CelebA, etc.). Clear ablation studies demonstrating the contribution of each component.\n\nCons:\n\n1. MLMC-RT gradient estimation, while theoretically efficient, adds significant implementation complexity. Multiple OT problems need to be solved at each iteration, potentially creating computational bottlenecks.\n\n2. Assumptions 3-4 regarding boundary regularity and non-degenerate filtering may not hold in practice. Lipschitz continuity requirements for kernels and encoders could limit applicability.\n\n3. Multiple hyperparameters (λ, ε, filtering thresholds) require careful tuning. High implementation barrier compared to simpler baseline methods.", "was_revised": true, "final_justification": "The proposed method is novel and provides valuable geometric intuition. The rebuttal is clear and solid.  All my major concerns have been addressed.", "rebuttal": "We thank the reviewer for the thoughtful comments, and address your concerns below.\n\n**Response to Weaknesses**\n\n**W1:** MLMC-RT gradient estimation, while theoretically efficient, adds significant implementation complexity. Multiple OT problems need to be solved at each iteration, potentially creating computational bottlenecks.\n\nWe would like to clarify two points regarding computational complexity:\n\n**MLMC-RT:** MLMC-RT is introduced not as a source of algorithmic complexity, but as a tool to make gradient estimation for our nested objective tractable and efficient. A naive Single-Level Monte Carlo (SLMC) estimator suffers from a fundamental trade-off: achieving low bias requires a large number of samples in the inner loop—on the order of $O(1/\\epsilon)$ per iteration (line 913). In contrast, MLMC-RT achieves the same estimation accuracy with significantly lower total complexity $O(\\log 1/\\epsilon)$ (line 948), making it exponentially more sample-efficient for any given target accuracy. Even when both MLMC-RT and SLMC use the same number of samples in the inner loop, their average per-iteration costs are comparable. \n\n\n**OT:** While entropy-regularized OT is often solved with the Sinkhorn's algorithm, our formulation allows for a **closed-form** for the transport plan (detailed in Appendix A.4). This is because we adopt a semi-relaxed OT formulation that only enforces the source marginal constraint ($\\gamma \\in \\Pi(P)$). This specific relaxation circumvents the need for iteratively balancing two marginals and instead admits a direct analytical solution.  In essence, we only need to compute the cost matrix $C_\\theta$ and then perform a few element-wise operations (like exp and log-sum-exp), which are highly optimized on modern hardware like GPUs. This means we bypass the need for any iterative OT solver, and the additional computational overhead per step is minimal. \n\n\n **W2:** Assumptions 3-4 regarding boundary regularity and non-degenerate filtering may not hold in practice. Lipschitz continuity requirements for kernels and encoders could limit applicability.\n\n\n1, **Assumption 3 (Boundary Regularity)** is a standard smoothness condition, similar to the Tsybakov margin condition in statistical learning (line 714). It helps prevent pathological cases and supports the stability of the filtering mechanism. This assumption is both common and reasonable for well-behaved data distributions, especially when data augmentation techniques are employed.  \n2, **Assumption 4 (Non-degenerate Filtering)** can be ensured through appropriate algorithmic design. In our case, the adaptive thresholding mechanism (lines 717–719, 264–267) can prevent excessive sample rejection and maintain sufficient support in the filtered distribution, thereby avoiding potential algorithmic instability. \n3, The **Lipschitz continuity** requirements for encoders and kernels are also common and practical. For instance, we analyzed in Appendix C.1 that this property holds for common choices like the cosine similarity used in our experiments (line 321, 961). \n\n\n**W3:** Multiple hyperparameters ($\\lambda$, $\\epsilon$, filtering thresholds) require careful tuning. High implementation barrier compared to simpler baseline methods. \n\n**Response to Questions**\n\n**Q1:** How does the OT computation scale with high-dimensional data and large sample sizes? \n\nPlease see our response to W1 for the  concerns on scalibility. \n\n**Q2:** Can you provide more detailed hyperparameter sensitivity analysis across different datasets?\n\n\nThe table below corresponds to the Colored MNIST benchmark, where we evaluate how various domain generalization methods perform under strong spurious correlations. To evaluate sensitivity to IOT calibration, we replace standard augmentations with multi-crop [1], which provides semantically invariant samples at multiple scales for ground-truth pairing. \n\n| Method           | ERM + Aug  | KL-DRO     | IRM        | Sinkhorn DRO | **AdaDRO (Ours)** | AdaDRO w/ multi crop. [1] |\n| ---------------- | ---------- | ---------- | ---------- | ------------ | ----------------- | -------------------------------- |\n| **Accuracy (%)** | 43.5 ± 0.6 | 41.2 ± 0.5 | 45.1 ± 1.0 | 44.8 ± 0.6   | **57.2 ± 0.8**    | 58.0 ± 0.7                       |\n\nThe following table presents results on CIFAR-100 with 50% label noise, where half of the training labels are randomly corrupted. This setting evaluates the robustness of different methods under severe label noise. We report accuracy across different values of the regularization parameter $\\lambda$.\n| Accuracy (%) | $\\lambda=0.5$ | $\\lambda=1.0$ | $\\lambda=2.0$  | $\\lambda=3.0$  | $\\lambda=5.0$ |\n| ------------ | ------------- | ------------- | -------------- | -------------- | ------------- |\n| AdaDRO       | 52.5 ± 0.6    | 54.2 ± 0.5    | **57.2 ± 0.8** | 55.6 ± 1.1     | 53.2 ± 0.8    |\n| Sinkhorn DRO | 42.5 ± 0.3    | 41.2 ± 0.7    | 44.1 ± 0.3     | **44.8 ± 0.5** | 43.9 ± 0.3    |\n\nDue to time constraints, we will include additional content in the revision.\n\n\n\n**Q3:** What is the actual wall-clock time comparison with baseline methods?\n\nWe provide a detailed wall-clock time comparison in Section 2.4 of our supplementary material, with key results summarized in Tables 4 and 5 therein. The empirical results demonstrate that our method, AdaDRO, is highly competitive and incurs only a minimal computational overhead compared to standard baselines, while being significantly more efficient than some DRO methods. \nSpecifically, on the Colored MNIST and Waterbirds datasets, the training time of AdaDRO is approximately 1.3x that of standard methods like GroupDRO and IRM. This is much faster than other robust optimization methods like WDRO and UniDRO, which are 5-6 times slower. Similarly, on the CIFAR-10-LT and CIFAR-100-LT benchmarks, AdaDRO's runtime is comparable with the efficient baselines like Decouple and Focal Loss, and SinkhornDRO. \n\n[1] Caron, Mathilde, et al. Emerging properties in self-supervised vision transformers. ICCV'21"}, {"review_id": "UVOUumrjoi", "reviewer": "Reviewer_YPiL", "summary": "This paper aims to improve the worst case performance in distribution shift via reframing the DRO problem as a bilevel OT problem, and optimizing the uncertainty set. Their method successfully bootstrap the semantic meaning through the “semantic calibration” and adaptive refinement of the undercaterin set based on the OT feedback. They further validate their results on the long-tailed benchmark datasets including CIFAR-10-LT and CIFAR-100-LT.", "questions": "Q1  You propose bootstrapping transport costs by solving an inverse‐OT problem (Eq. 5). How sensitive is this calibration step to the choice of the “observed coupling” γ̄? Have you evaluated alternative sources of semantic couplings (e.g., feature‐space clustering) and their impact on convergence?\n\n\nQ2: The uncertainty set evolves during training—can you provide theoretical or empirical guarantees on its “diameter” (e.g., how δ effectively grows or shrinks) over epochs? In particular, under what conditions might the set collapse?\n\n\nQ3: Entropy‐regularized OT and its inverse both require Sinkhorn iterations. What is the alignment plan of your two OT loops per minibatch, and how does it scale with n? Have you compared against approximations (e.g., sliced‐Wasserstein) in large‐scale settings?\n\n\nQ4: Your experiments focus on synthetic “double moon” and Gaussian‐based shifts. How does AdaDRO perform on real‐world image‐classification shifts (e.g., CIFAR → CIFAR‐C or domain‐adaptation benchmarks)? Can you characterize which types of shifts (e.g., covariate vs. label shift) benefit most from adaptive versus static uncertainty sets?\n\n\nQ5  You learn the feature-level cost via IOT using augmented pairs (Eq. 12). How sensitive is the learned cost to the choice and strength of augmentations? I\n\n\nQ6: You adopt multilevel Monte Carlo with randomized truncation for nested expectations (Sec. 3.3). What is the per-iteration time overhead of MLMC-RT compared to a single-level Monte Carlo?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "This paper demonstrates several strengths across the dimensions of originality, quality, clarity, and significance:\n\n\n\n\n\nQuality:\n\n\nThis paper makes a good contribution by reframing entropy‐regularized DRO as a bilevel optimal‐transport problem, where a semi‐relaxed OT “explores” worst‐case distributions and an inverse‐OT “aligns” semantic costs, enabling both feature- and label-level calibration. Each theory of paper is supported by rigorous proof.\n\n\n\n\n\nClarity:\n\n\nThe paper is generally very clear and well‐organized to show the development from OT preliminaries, bilevel formulation, semantic calibration, adaptive refinement, definitions and lemmas are stated precisely. The only concern here is lack the background for the OT or the IOT applications, making this connection ambiguous.\n\n\n\n\nSignificance:\n\n\nThis work advances DRO by unifying worst-case distribution search and model alignment under a single OT framework. Its novel OT-driven feedback mechanism adaptively refines the uncertainty set during training, and proves the convergence. Empirically, AdaDRO delivers strong, consistent improvements over DRO and semantic-aware baselines across spurious-correlation, group-shift, and long-tail benchmarks. \n\n\n\n\n\nOriginality:\n\nThe paper’s originality lies in its unified OT‐based reframing of entropy‐regularized DRO as a bilevel problem, coupling semi‐relaxed OT exploration with inverse‐OT alignment. Moreover, casting softmax cross‐entropy as an inverse‐OT objective and integrating MLMC‐RT for nested expectation gradients under evolving distributions further distinguishes the work, marking a novel synthesis of transport theory and adaptive robust optimization.\n\n\n\n\n\nWeakness:\n\nW1: Unclear about the computational costing of this method. Could you summarize the time and space complexity to run these algorithms?\n\n\nW2: The manuscript currently omits the background literature to illustrate the inverse optimal transport (IOT) and feature‐level alignment works that directly underpin its semantic calibration and adaptive refinement. To address this, the authors should explicitly cite foundational IOT frameworks, including:\n\nLi, Ruilin, et al. \"Learning to match via inverse optimal transport.\" Journal of machine learning research 20.80 (2019): 1-37.\n\nChiu, Wei-Ting, Pei Wang, and Patrick Shafto. \"Discrete probabilistic inverse optimal transport.\" International Conference on Machine Learning. PMLR, 2022.\n\nChen, Zihao, et al. \"Your contrastive learning problem is secretly a distribution alignment problem.\" Advances in Neural Information Processing Systems 37 (2024): 91597-91617.\n\nPersiianov, Mikhail, et al. \"Inverse Entropic Optimal Transport Solves Semi-supervised Learning via Data Likelihood Maximization.\" arXiv preprint arXiv:2410.02628 (2024).\n\nW3: Not sure about the hyperparameter sensitivity, the key components—regularization strength ϵ, filter radius δ, kernel choice/bandwidth, and thresholding scheme—require careful tuning, yet ablations on their robustness are limited.\n\n\nW4: It seems that the semantic calibration relies on data augmentations; if augmentations produce semantically invalid pairs, the learned cost may misalign, but the paper lacks analysis of failure modes.", "was_revised": true, "final_justification": "The author addressed my concerns. I think it is a good paper.", "rebuttal": "We thank the reviewer for the thoughtful comments, and address your concerns below.\n\n**Response to Weaknesses**\n\n**W1:** **Could you summarize the time and space complexity to run these algorithms?**\n\nWhile entropy-regularized OT is often solved with the Sinkhorn's algorithm, our formulation allows for a **closed-form** for the transport plan (as seen in Appendix A.4). In essence, we only need to compute the cost matrix $C_\\theta$ and then perform a few element-wise operations (like exp and log-sum-exp), which are highly optimized on modern hardware like GPUs. This means we bypass the need for any iterative OT solver, and the additional computational overhead per step is minimal.\n\n\n**W2 W3:** The manuscript currently omits the background literature to illustrate the inverse optimal transport (IOT) and feature‐level alignment works...; ... hyperparameter sensitivity.... \n\nWe will address the related works you mentioned in our revised version. Additionally, we will include more ablation studies; the corresponding tables are omitted here due to space limitations. \n\n\n**W4:** It seems that the semantic calibration relies on data augmentations; if augmentations produce semantically invalid pairs...the paper lacks analysis of failure modes.\n\nWe will discuss this limitation. If the augmentations are too aggressive and create semantically invalid pairs (e.g., cropping a \"cat\" image to show only grass), the learned transport cost could be misled. In our experiments, we followed standard practices and used only mild data augmentations (such as random flipping and cropping), which are widely adopted in semi-supervised and self-supervised learning, and have been shown to effectively preserve the semantics of the samples. \n\n\n**Response to Questions**\n\n**Q1 and Q5:** Sensitivity to the choice of observed coupling $\\bar{\\gamma}$. Sensitivity to the choice and strength of augmentations.  \n\nIn our current work, we define $\\bar{\\gamma}$ using a simple one-to-one matching between a sample and its augmented counterpart (11).  We have not yet explored alternative sources for $\\bar{\\gamma}$. However, we hypothesize that using more sophisticated couplings could be a promising direction for future work. It might provide a richer, many-to-many semantic signal. \n\nThe learned cost is not sensitive to the augmentations.  Data augmentation is a widely used technique to construct *semantically invariant samples*, especially in contrastive learning. Recent works have explored multi-level augmentations such as **multi-crop** [1], which creates stronger invariance by including different spatial scales. To study the sensitivity of our method to the augmentation strategy, we replaced our standard augmentations with multi-crop augmentation.\n\nEmpirically, we observe that the performance of our method remains stable under this change. In some cases, multi-crop leads to slight improvements, suggesting that our method is robust to the choice of augmentation. For instance, when applied to the *Colored MNIST* task with a spurious correlation of 0.8, we observe a 0.2% increase in accuracy. On *CIFAR-100* with 50% label noise, multi-crop yields a 0.8% improvement. Due to space constraints, we omit detailed tables.\n\n\n\n\n**Q2:** Guarantees on the diameter of the evolving uncertainty set.\n\nThe dynamic uncertainty set does not refer to a change in radius $\\delta$, but to the adaptive adjustment of the uncertainty set's geometry and support. This is achieved through two core mechanisms: 1, dynamically refining the reference distribution $\\nu$ into $\\hat{\\nu}$ via OT feedback (16); 2, calibrating the semantic transport cost $C_\\theta$ via Inverse OT (Sec. 3.2). To prevent collapse of the uncertainty set, we introduce Assumption 4 (Non-degenerate Filtering), which ensures that the samples retained by our filtering mechanism maintain a minimum effective support throughout training. Our algorithmic design satisfies this assumption. \n\nEmpirically, while precisely tracking diameter evolution in high-dimensional space is challenging, our toy example in Fig. 1 intuitively demonstrates the process’s effectiveness. The method successfully identifies and removes low-quality outliers lying off the true data manifold, thereby producing a cleaner decision boundary. Furthermore, our ablation studies (Table 7) and experiments in Table 8 validate that this adaptive process both prevents performance collapse and significantly improves robustness. We plan to include additional visualizations in the revised version to better illustrate the dynamics of this process. \n\n**Q3:** Alignment of OT loops and scalability.\n\nOur framework entirely avoids iterative OT solvers. This is achieved because both our forward and inverse OT problems are formulated using a **Semi-Relaxed Marginal Constraint**, as detailed in Appendix A.4. In this setting, we only enforce the source marginal of the transport plan, which allows the optimal coupling $\\gamma^*$ to have a closed-form analytical solution (23). \n\nTherefore, our alignment plan per minibatch is highly efficient and non-iterative: We first compute the cost matrix $C_\\theta$. We then directly obtain the transport plan $\\gamma^*$ by applying the closed-form solution from (23), which only involves highly-parallelizable, element-wise operations (e.g., `exp` and a normalization) that are heavily optimized on GPUs. Both the exploration (OT) and alignment (inverse OT) components of our framework leverage this formulation, meaning neither requires an iterative loop like Sinkhorn. The computational complexity per step is dominated by the $O(nm)$ cost matrix computation, which is standard for many OT-based methods and is far more efficient than the $O(nm \\times L)$ cost of a full Sinkhorn algorithm, where $L$ is the number of iterations.\n\nRegarding the comparison against approximations like sliced-Wasserstein: these methods are indeed powerful alternatives for the general OT problem, which lacks a closed-form solution and requires iteration. Our approach (Semi-Relaxed OT) achieves efficient, non-iterative computation. We agree that exploring such approximations is an interesting future direction, particularly when extending our framework to fully constrained transport scenarios or other variants like Gromov OT. We will include this point in our discussion.\n\n\n**Q4:** Performance on real-world image classification shifts.\n\nWe believe our method is unique in two aspects: 1, its label-agnostic sample-level semantic calibration (13); 2, the filtering mechanism based on the coupling introduced by (12) (13). \n\nAn intuitive example is the presence of noisy labels: when a portion of labels are flipped, sample-level calibration (12) can still reliably capture semantics. This is because the process is unsupervised, similar to contrastive learning [4], and thus remains robust to label noise. Building on this, the coupling-based filtering mechanism can discards samples with wrong label, enabling the cross-entropy to train on a cleaner subset of the data.\nIn contrast, methods that depend on label supervision, including ERM with augmentation, face an inherent \"garbage in, garbage out\" issue. No  augmentation can correct mislabeled examples, and these label errors inevitably propagate through the learning process, compromising model robustness. \n\nA similar advantage is observed in Colored MNIST. Even when augmentation occasionally alters the spurious attribute (e.g., color), standard ERM still learns entangled representations, as the label remains strongly correlated with color. In contrast, our label-agnostic calibration encourages the model to identify truly invariant features—those that remain stable under augmentation—naturally prioritizing shape over color in the absence of a strong supervisory signal.  \n\n\n\n| Method           | ERM + Aug  | KL-DRO     | IRM        | Sinkhorn DRO | **AdaDRO (Ours)** | AdaDRO w/o Semantic Calib. (Ablation) |\n| ---------------- | ---------- | ---------- | ---------- | ------------ | ----------------- | ------------------------------------- |\n| **Accuracy (%)** | 43.5 ± 1.2 | 41.2 ± 1.5 | 45.1 ± 1.0 | 44.8 ± 1.1   | **57.2 ± 0.8**    | 46.5 ± 0.9                            |\n\n**Q6:** ... overhead of MLMC-RT compared to a single-level Monte Carlo?\n\n\nTo provide a complete comparison regarding the per-iteration time overhead, we compare MLMC-RT against Single-Level Monte Carlo (SLMC) in two scenarios: \n\n1, When we fix the number of inner samples for both SLMC and MLMC-RT, their average per-iteration cost is similar. However, SLMC suffers from significant and uncontrollable estimation bias due to its fixed and limited sampling. In contrast, MLMC-RT maintains low cost by prioritizing low-complexity levels (with sampling probability roughly proportional to $2^{-l}$, making its average cost comparable to that of a naive SLMC), but still achieves a much lower bias thanks to its multi-level correction. This leads to substantially more accurate estimates without increasing average computational cost. \n\n2, To achieve the same low estimation bias as MLMC-RT (for a target accuracy $\\epsilon$), SLMC must use a large number of inner samples, on the order of $N = O(1/\\epsilon)$, making its per-iteration cost prohibitively high. In contrast, MLMC-RT achieves the same bias accuracy with far fewer samples, requiring only $N = O(\\log^2(1/\\epsilon))$, and is therefore exponentially more efficient under the same bias target. \n\nIn summary, MLMC-RT offers the best of both worlds: its average per-iteration cost is comparable to that of a cheap but biased SLMC estimator, while its accuracy rivals that of an expensive high-fidelity SLMC. We will add a paragraph to the appendix elaborating on this trade-off, and include the wall-clock time comparison (as noted in W1) to provide concrete empirical support. \n\n[1] Caron, Mathilde, et al. Emerging properties in self-supervised vision transformers. ICCV'21"}]

Author rebuttal:

Meta-review: This paper develops a new DRO framework where Optimal Transport is utilized together with (semi-relaxed) Inverse Optimal Transport during the DRO optimization to find a robust model/classifier. The proposed approach addresses the current limitation of DRO approach, which is the design of the uncertainty set. Instead, the authors develop a novel perspective that casts entropy-regularized Wasserstein DRO as a dynamic process of distributional exploration and semantic alignment. Theoretical convergence guarantees were also provided by the authors. All the reviewers are positive about this manuscript and appreciate its novelty and significance for improved robustness over existing benchmarks. Some minor concerns, although not critical, still exist. The authors may address them in a future version of this paper.
Decision: Accept (poster)