Title: Bootstrap Your Uncertainty: Adaptive Robust Classification Driven by Optimal-Transport
Venue: NeurIPS 2025
Accepted: yes
Keywords: distributionally robust optimization

optimal transport
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-Transport

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 as 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. 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: 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 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 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-ofthe-art robustness.

1 Introduction

Distribution shifts between training and deployment environments present a fundamental challenge to 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 a popular framework that optimizes for worst-case performance across a set of potential test distributions, known as the ' uncertainty set '. Unlike empirical risk minimization (ERM), which directly minimizes the expected loss ℓ ( θ, x ) over the empirical training distribution P tr , DRO seeks robustness by solving the minimax problem: min θ sup Q ∈U ( P tr ) E Q [ ℓ ( θ, x )] . Here, U ( P tr ) represents the uncertainty set that encompasses distributions in the neighborhood of P tr , capturing a range of plausible test-time deviations.

A critical challenge in implementing effective DRO lies in designing an appropriate uncertainty set. Ideally, the uncertainty set should accurately capture potential distributional shifts while remaining amenable to efficient optimization. Optimal transport (OT) [61] provides an elegant mathematical framework for this purpose, offering both strong theoretical foundations and intuitive physical interpretations (see Sec. 2.1). These properties make OT distances suitable for constructing meaningful

uncertainty set in distribution space. In this work, we leverage the entropy-regularized OT distance (also known as, Sinkhorn distance [16], and the formal definition is shown in Def. 2) to define our uncertainty set. Let D ϵ denote the Sinkhorn distance and the DRO problem [62, 3] is formalized as:

Here δ > 0 serves as the radius to control the size of allowable distributional shift; P ( X×Y ) denotes the set of Borel probability distributions over the data-label space ' X ×Y '. The uncertainty set reflects prior knowledge about potential distributional shifts. As we will demonstrate in Sec. 2.2, the Sinkhorn distance D ϵ ( P tr , Q ) introduces a reference distribution ν that guides the uncertainty set's shape. This reference distribution determines the 'closeness' between distributions and the training data P tr , thereby serving as a prior knowledge about the 'geometry' of the uncertainty set.

Despite its popularity, OT-based DRO ( e.g., (1)) still faces two 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 as 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 ν , which can mitigate those issues to some extent. In particular, the distribution ν regularizes the transport plan and implicitly shapes the geometry of the uncertainty set U ( P tr , δ ) by controlling where and how probability mass can be transported [62]. However, constructing an appropriate reference distribution ν remains challenging. When training data is limited, ν is typically synthesized through heuristics (e.g., mixup [67] or noise injection [26]). But, this operations may introduce unrealistic distributions in U ( P tr , δ ) . Moreover, we are also confronted with a tricky dilemma for constructing the uncertainty set U ( P tr , δ ) :

Expanding the uncertainty 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, yielding a dynamic uncertainty set . While introducing this flexibility to the uncertainty set U ( P tr , δ ) 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 as follows,

We introduce a novel theoretical perspective that frames entropy-regularized DRO as a dynamic exploration-and-alignment 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. Furthermore, the OT-based structure of (1) enables seamless integration of semantic calibration and refinement into the training process.

Extensive experiments demonstrate that our method improves both robustness and accuracy across a 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 1. Due to limited training data, the learned decision boundary significantly deviates from the ideal boundary. The reference distribution ν , 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

2

2

0

0

-1

-1

-2

-2

2

-1

0

1

2

2

1

0

1

2

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 a Kullback-Leibler ball around the empirical measure [29, 46]. 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 [9, 10] and extensions to high-dimensional settings by Gao and Kleywegt [23]. Entropy-regularized OT, or Sinkhorn DRO, introduces a reference measure and entropy penalty to yield smoother uncertainty set and more efficient algorithms [64, 62]. The full discussion on related works is deferred to Appendix A.1.

2 Preliminaries

Consider the input-output space Ξ = X × Y , where X ⊆ R d is the data space and Y := [ K ] = { 1 , · · · , K } denotes the label space containing K distinct classes. Let P (Ξ) represent the set of Borel probability distributions supported on Ξ . For 1 ≤ i ≤ n , each sample p i = ( x i , y i ) is drawn from an underlying distribution P ∈ P (Ξ) . The empirical distribution P n = 1 n ∑ n i =1 δ p i is induced by the dataset { p 1 , . . . , p n } , where δ p i is the Dirac measure concentrated at p i . We endow Ξ with a transport cost function: C ( p i , p j ) = C X ( x i , x j ) + C Y ( y i , y j ) , where C X : X × X → R and C Y : 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 our neural network as f Θ ( · ) = h W ◦ g θ ( · ) , where g θ ( · ): X → Z serves as the encoder , parameterized by θ , which maps input data x ∈ X into a latent representation space Z ⊂ R d e . The function h W ( · ): Z → Y acts as a linear classification head that generates predictions in the output space Y . The classifier weight matrix W = [ w 1 , . . . , w K ] ∈ R K × d e , where each w i is the class-specific weight vector for the i -th class. The complete set of model parameters is Θ = { θ, W } . For notational simplicity, we sometimes abbreviate the notation as f ( · ) = h ◦ g ( · ) 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 a cost function, while IOT infers appropriate cost functions that best explains observed coupling patterns.

Let P and Q be probability distributions, and Π( P, Q ) denote the set of all joint distributions (called transport plans or couplings ) with marginals P and Q . Formally, any γ ∈ Π( P,Q ) satisfies ∫ γ ( p,q )d q = P ( p ) and ∫ γ ( p,q )d p = Q ( q ) . Let µ and ν be reference measures such that P ≪ µ and

Q ≪ ν (i.e., P and Q are absolutely continuous with respect to µ and ν , respectively). We begin with the entropy regularization.

Definition 1 (Relative Entropy of Transport Plans) . For any transport plan γ ∈ Π( P, Q ) , the relative entropy of γ w.r.t. product measure µ ⊗ ν is:

This entropy term requires supp ( P ) ⊆ supp ( µ ) and supp ( Q ) ⊆ supp ( ν ) to remain finite. A common choice [40] uses uniform reference measures ( d µ ( p )d ν ( q )= const), simplifying the entropy to:

Entropy-Regularized Optimal Transport Cuturi [16] 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 ( a.k.a. Entropy-regularized OT Distance)) . Given distributions P, Q , a cost function C : Ξ × Ξ → R + , and regularization strength ϵ> 0 , the Sinkhorn distance is:

Inverse Optimal Transport (IOT) IOT [57, 39, 13, 12] solves the inverse problem: inferring cost function that explains observed couplings. To capture this, we add the subscript θ to the previously defined notations ' γ ' and ' C ', i.e., γ θ and C θ , which denote the corresponding optimal transport plan and cost learned by the model parameterized by θ . In addition, we let ¯ γ ∈ Π( P,Q ) denote an observed coupling that reflects ground-truth semantic relationships (details in Sec.3.2). Then we have the following formulation for IOT:

Here, the KL divergence KL(¯ γ | γ θ ) quantifies how well C θ explains the observed matching pattern ¯ γ . The entropy term H ( γ ) is defined as in (3). This bilevel structure allows us to learn transport costs C θ 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 (1), highlighting its theoretical foundations and advantages over alternative approaches. Sinkhorn DRO extends Wasserstein DRO (WDRO) [9, 10, 23] through entropy regularization. When the transport cost C is a metric and ϵ = 0 , the formulation (1) 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 [19, 42, 56].

The introduction of entropy regularization in D ϵ brings benefits through two complementary mechanisms: Reference-guided uncertainty set: The relative entropy term H ( γ | µ ⊗ ν ) leverages reference distributions µ 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-based DRO (e.g., KL-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 ν , following Wang et al. [62], we set µ = P tr , the empirical training distribution. This choice is both practical and theoretically sound, as the relative entropy term differs only by a constant when P tr ≪ µ , leaving the optimization unaffected. In contrast, the reference measure ν plays a critical role: it defines the support of feasible worst-case distributions by enforcing Q ≪ ν . Through dual analysis, Wang et al. [62] derived the closed-form expression for the worst-case distribution. When the optimal Lagrange multiplier λ > 0 , the worst-case distribution Q λ has the density:

ref.dist.1

IOTloss

DynamicFilter

Model-derived

Coupling

Ground-truth

Coupling

Projector

tr

中()

Reweight

SampleCost

Encoder

feature

C(cir)

0

X

KL(I)

(cia)

Semantic

g()

map

Calibration

Classifier

f()

LabelCost

DynamicFilter

Co(ci,k)

KL()

（c,k)

The role of ν . The formula (6) shows that the worst-case distribution Q λ is essentially a reweighted version of ν , 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 ν , highlighting why the design of ν is crucial for effective distributional robustness. By dynamically refining the reference distribution ν , we can induce a dynamic uncertainty set that evolves with the learning dynamic to better reflect the underlying data geometry, which has not been explored previously.

3 Our Method: AdaDRO

To address the dilemma highlighted in Sec. 1, we propose a 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 a bilevel optimization whose upper-level and lower-level objectives are OT and IOT objectives, respectively (Sec. 3.1). This perspective inspires us to incorporate a more well-designed IOT for semantic calibration in Sec. 3.2; moreover, the geometric implication of OT coupling can serve as a feedback to refine the uncertainty set (Sec. 3.3). The overall framework is illustrated in Figure 1.

3.1 Sinkhorn DRO Revisited: A New Perspective from OT

In this section, we introduce a new perspective for Sinkhorn DRO from OT. First, we cast Sinkhorn DRO as a semi-relaxed optimal transport (OT) problem , where the benefit is that the worst-case distribution Q λ (defined in (6)) emerges naturally from the optimal transport plan. Additionally, we reinterpret the softmax cross-entropy loss as a form of inverse OT . Taken together, these insights enable the entire learning dynamics of DRO to be formulated through the OT/IOT objectives, so that it 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 ' γ ∈ Π( P, Q ) ' with ' γ ∈ Π( P ) := { γ | ∫ γ ( p, q )d q = P ( p ) } ',

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

the reference measure ν . 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 (1) with P tr and ν being the source and reference measures. If the transport cost in the semi-relaxed OT problem (7) is taken as C ′ ( p, q ) = 1 λ ( λC ( p, q ) -ℓ ( q )) , then the optimal transport plan γ ∗ ∈ Π( P tr ) has its target marginal equal to the worst-case distribution Q λ in (6) .

Classification as Inverse OT. Complementing the above view, we 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:

where ¯ γ ( x i , k ) encodes the coupling between sample x i and its true label y i , and C θ ( x i , k ) defines the model-learned transport cost to class k ∈ [ K ] .

As shown by Shi et al. [53], when using the cost C θ in the semi-relaxed OT formulation (7), the resulting optimal plan γ θ takes the form of a softmax distribution: γ θ ( x i ,k )= exp( ⟨ g ( x i ) ,w k ⟩ /ϵ ) ∑ K j =1 exp( ⟨ g ( x i ) ,w j ⟩ /ϵ ) , and the inverse OT objective, defined as KL(¯ γ | γ θ ) , becomes:

This expression exactly recovers the temperature-scaled cross-entropy loss commonly used in classification tasks.

Remark 1. This analysis shows that the DRO objective in (1) can be cast as a bilevel optimization [32] grounded in OT: Lower level (Exploration): exploring the worst-case distribution Q λ ( θ ) by a semirelaxed OT (as Lemma 3.1). Upper level (Alignment): updating parameters θ 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 ( p i , p j ) = C X ( x i , x j ) + C Y ( y i , y j ) . 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 to learn them directly from the data using IOT.

Calibrating the Sample-level Transport Cost C X . To ensure that transport costs capture taskrelevant semantic relationships, we propose learning C X θ through a principled IOT formulation. We define the parameterized transport cost as:

where c is a constant ensuring non-negativity, g θ ( x ) represents the embedding of sample x , and Ker( · , · ) is a kernel function (e.g., cosine or Gaussian) measuring similarity in the embedding space.

Given a sample x i drawn from the worst-case distribution Q λ , we generate a semantically equivalent counterpart x ′ i ∼ Aug ( x i ) using standard data augmentation techniques (e.g., cropping, rotation). These pairs ( x i , x ′ i ) define ground-truth semantic matching encoded as a coupling:

where I [ · ] is the indicator function. This coupling represents the semantically meaningful matches that our transport cost should respect.

To calibrate C X θ , we formulate the following IOT problem:

Similar to (7), Π( Q λ ) denotes the set of couplings whose first marginal is Q λ , i.e., Π( Q λ ) = { γ | ∑ j γ ( x i ,x ′ j )= Q λ ( x i ) } . This optimization aims to learn parameters θ such that the transport plan γ X θ , induced by the cost function C X θ , aligns with the ground-truth coupling ¯ γ X defined in (11).

Calibrating Label-level Transport Cost C Y . 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.1 offers a principled mechanism for semantic calibration in label space. To formalize this connection, we revisit the classification loss through the lens of inverse OT:

Here, the cost function is defined as C θ ( q i , k ) = c -⟨ g θ ( x i ) , w k ⟩ , which can be seen as a special case of a kernel-based cost: C θ ( q i , k ) = c -Ker( g θ ( x i ) , w k ) , where Ker( · , · ) is a linear kernel. A fundamental property of such kernels, known as the reproducing property [47], establishes that: ⟨ Ker( w k , · ) , Ker( w k ′ , · ) ⟩ = Ker( w k , w k ′ ) . This property reveals that standard cross-entropy training, 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-level transport cost as:

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 w k to encode label semantics has also appeared in prior work, such as [58], although they do not interpret this mechanism by inverse OT alignment.

3.3 Our Overall Algorithm

As discussed in Sec. 1, the reference distribution ν serves as a crucial prior that defines the support of the DRO uncertainty set. However, high-quality priors are rarely available: an overly broad ν may include implausible samples, while an overly restrictive one weakens robustness guarantees. The OT formulations in Sec.3.1-3.2 incorporate relaxed marginal constraints, reducing the need to forcefully match noisy or outlier samples [5, 14]. 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 ν , thereby refining the uncertainty set during training.

Refining ν via OT Feedback. We have introduced two Inverse OT frameworks for semantic calibration-feature-level IOT in (12) and label-level IOT in (13). 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 γ θ with the ground-truth coupling ¯ γ X (defined in (11)). For any matched pair ( p, q ) ∼ ¯ γ X , we define its matching confidence as γ X θ ( p, q ) , which quantifies how strongly the model aligns p with q . A global average confidence is then computed as τ 1 = E ( p,q ) ∼ ¯ γ X γ X θ ( p, q ) .

Because some targets may inherently be more difficult to match, particularly when the marginal distribution of the ground-truth coupling ¯ γ X is imbalanced, we define a target-specific local matching capacity for each q : τ 2 ( q ) = E p γ X θ ( p, q ) , which quantifies how much mass is transported to q across all source. A larger τ 2 ( q ) indicates that the target q is easier to match. Combining global and local signals, we define a dynamic threshold:

Each source p is retained in the filtered reference distribution ̂ ν only if it aligns well with at least one target q :

where χ normalizes ̂ ν to form a 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

Algorithm 1 OT Driven Adaptive Distributionally Robust Optimization (AdaDRO)

Require: Initial parameter θ 0 , step size α > 0 , total iterations T.

1: Train θ by the joint objective in (12) and (13) for N 0 epochs. ▷ Semantic calibrate by IOT

2: for t = 0 , 1 , . . . , T -1 do

3: ̂ ν t ← FilterReference( ν, θ t ) ▷ Filter via (16)

̂ 5: Estimate gradient grad t via RT-MLMC estimation for the IOT objectives (13) and (12) ▷ See Algorithm 2 in Appendix B.4

4: ̂ Q λ t ← SinkhornWorstCase( P tr , ν t , C θ t , λ ) ▷ Compute worst-case distribution via (6).

6: Update θ t +1 ← θ t -α t grad t

7: end for

sample-specific thresholds to account for varying target difficulties, overcoming a key limitation of static filtering rules. This design of our filtering mechanism is guided by the geometric intuition behind transport couplings, as 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-specific) 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-feature-level IOT in (12) and label-level IOT in (13)-to obtain a refined reference distribution ̂ ν . We then perform Sinkhorn distributionally robust optimization using ̂ ν . The full procedure of our OT-driven adaptive robust algorithm is summarized in Algorithm 1, and its schematic overview is provided in Figure 2. Theorem 3.2 (Convergence of AdaDRO) . Let the objective at iteration t to be defined as F t ( θ ) = E q ∼ ̂ Q λ t [ ℓ ( θ, q )] . Under mild smoothness assumptions, Algorithm 1 with decaying step size α t = α/ √ t +1 converges to a stationary point after at most T = O ( ϵ -4 log 1 ϵ ) iterations:

Here the total sample complexity O ( ϵ -4 log 2 1 ϵ ) .

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 λ t 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 C θ and the filtered reference distribution ̂ ν θ vary smoothly with θ , as shown in Proposition B.1. This ensures that changes to the uncertainty set remain controlled during training.

involves a composition of expectations.

For challenge (ii) , the objective F t ( θ ) = E q ∼ ̂ Q λ t [ ℓ ( θ, q )] Because the worst-case distribution ̂ Q λ t ( q ) , as we introduced in (6), is given by

Here the denominator of ϕ ( q ) is also an expectation over u ∼ ̂ ν t . This nested structure makes it impossible to construct an unbiased gradient estimator using standard Monte Carlo, and naive stochastic approximations often incur intractable cost [38, 62]. To address this, we adopt the multilevel Monte Carlo with Randomized Truncation (MLMC-RT) technique [28], a method has been used in robust optimization to trade off bias, variance, and computational cost efficiently. Related works such as [38] use MLMC for estimating gradients in worst-case distributional settings. In our case, we 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 ϵ -stationary point using a total of O ( ϵ -4 log 2 1 ϵ ) samples (as in Theorem 3.2)-only logarithmically worse than the known lower bound for general non-convex stochastic optimization [1]. This result demonstrates that robust optimization with dynamically adaptive uncertainty sets and semantic costs can also be performed efficiently at scale . Our full theoretical details are provided in Appendix B.

4 Experiments

We evaluate the proposed AdaDRO framework across a diverse set of benchmarks to assess its effectiveness under various types of distribution shifts.

4.1 Experimental Setup

Datasets. We evaluate on three widely studied distribution shift settings: Colored MNIST [2], which tests robustness under spurious correlations; Waterbirds [50], a real-world dataset with strong background-label correlation; CelebA [43], a benchmark for facial attribute recognition. We also evaluate on several long-tailed benchmark datasets: CIFAR-10-LT and CIFAR-100-LT [37].

Baselines. We compare AdaDRO with the following methods: KL-DRO [46], Wasserstein DRO ( WDRO ) [9], Sinkhorn DRO SDRO [62], a unified DRO method UniDRO [11] , semantic-aware robust methods such as GroupDRO [50], an invariant risk minimization method IRM [2].

Our models are implemented with PyTorch on a single NVIDIA RTX 6000 Ada GPU using the AdamW optimizer [45]. For AdaDRO, we use cosine similarity as the kernel in (10), and employ basic augmentations (flip, crop) for semantic calibration in Sec. 3.2 unless otherwise specified.

4.2 Experimental Results

AdaDRO is broadly applicable thanks to its flexible reference distribution design. While this flexibility enables adaptation to various settings, our method still relies on prior knowledge to construct the reference distribution, which may limit its applicability in some cases. Nevertheless, domain priors can often be effectively incorporated, as in many specialized robust learning approaches. For instance, in imbalanced settings, the reference distribution ν 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 1 and 2 show part of the results on Colored MNIST and Waterbirds. AdaDRO consistently achieves top worst-group accuracy, outperforming or matching specialized methods like GroupDRO.

Compared with other baselines, our method offers two main benefits. First, it performs a labelagnostic, sample-level semantic calibration as in (12). Second, it introduces a coupling-based filtering mechanism derived from the coupling introduced by (12) and (13). An intuitive example is the presence of noisy labels as shown in Table 3: 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 [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 a cleaner subset of the data. In contrast, methods that rely heavily on label supervision, such as ERM with augmentation, suffer from an 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.

A similar advantage is observed in Colored MNIST as in Table 4. 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.

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 of 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 2021YFA1000900, the NSFC through grants No. 62432016 and No. 62272432, and the Provincial NSF of Anhui through grant 2208085MF163. The authors would also like to thank the anonymous reviewers for their valuable comments and suggestions.

References

Yossi Arjevani, Yair Carmon, John C. Duchi, Dylan J. Foster, Nathan Srebro, and Blake E. Woodworth. Lower bounds for non-convex stochastic optimization. Math. Program. , 199 (1):165-214, 2023. doi: 10.1007/S10107-022-01822-7. URL https://doi.org/10.1007/ s10107-022-01822-7 .

Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint arXiv:1907.02893 , 2019.

Waïss Azizian, Franck Iutzeler, and Jérôme Malick. Regularization for wasserstein distributionally robust optimization. ESAIM: Control, Optimisation and Calculus of Variations , 29:33, 2023.

Aurélien Bellet, Amaury Habrard, and Marc Sebban. A survey on metric learning for feature vectors and structured data. arXiv preprint arXiv:1306.6709 , 2013.

Jean-David Benamou. Numerical resolution of an 'unbalanced' mass transport problem. ESAIM: Mathematical Modelling and Numerical Analysis , 37(5):851-868, 2003.

Jean-David Benamou, Guillaume Carlier, Marco Cuturi, Luca Nenna, and Gabriel Peyré. Iterative bregman projections for regularized transportation problems. SIAM Journal on Scientific Computing , 37(2):A1111-A1138, 2015.

Dimitris Bertsimas and Nihal Koduri. Data-driven optimization: A reproducing kernel hilbert space approach. Operations Research , 70(1):454-471, 2022.

Christopher M Bishop and Nasser M Nasrabadi. Pattern recognition and machine learning , volume 4. Springer, 2006.

Jose Blanchet and Karthyek Murthy. Quantifying distributional model risk via optimal transport. Mathematics of Operations Research , 44(2):565-600, 2019.

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.

Anh Tuan Bui, Trung Le, Quan Hung Tran, He Zhao, and Dinh Phung. A unified wasserstein distributional robustness framework for adversarial training. In International Conference on Learning Representations , 2022.

Zihao Chen, Chi-Heng Lin, Ran Liu, Jingyun Xiao, and Eva L Dyer. Your contrastive learning problem is secretly a distribution alignment problem. In The Thirty-eighth Annual Conference on Neural Information Processing Systems , 2024. URL https://openreview.net/forum? id=iNUKoLU8xb .

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://proceedings.mlr.press/v162/chiu22b.html .

Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard. Unbalanced optimal transport: Dynamic and kantorovich formulation, 2019. URL https://arxiv.org/ abs/1508.05216 .

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 Pattern Recognition , pages 9268-9277, 2019.

Marco Cuturi. Sinkhorn distances: Lightspeed computation of 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, Nevada, United States. , pages 2292-2300, 2013.

Marco Cuturi and David Avis. Ground metric learning. The Journal of Machine Learning Research , 15(1):533-564, 2014.

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.

Tongtong Fang, Nan Lu, Gang Niu, and Masashi Sugiyama. Rethinking importance weighting for deep learning under distribution shift. Advances in neural information processing systems , 33:11996-12007, 2020.

Tongtong Fang, Nan Lu, Gang Niu, and Masashi Sugiyama. Generalizing importance weighting to a universal solver for distribution shift problems. Advances in Neural Information Processing Systems , 36, 2024.

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.

Rui Gao, Xi Chen, and Anton J. Kleywegt. Wasserstein distributional robustness and regularization in statistical learning. CoRR , abs/1712.06050, 2017.

Rui Gao, Xi Chen, and Anton J Kleywegt. Wasserstein distributionally robust optimization and variation regularization. Operations Research , 72(3):1177-1191, 2024.

Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics , pages 16081617. PMLR, 2018.

Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization. arXiv preprint arXiv:2007.01434 , 2020.

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 .

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.

Yifan Hu, Xin Chen, and Niao He. On the bias-variance-cost tradeoff of stochastic optimization. Advances in Neural Information Processing Systems , 34:22119-22131, 2021.

Zhaolin Hu and L Jeff Hong. Kullback-Leibler divergence constrained distributionally robust optimization. Available at Optimization Online , 1(2):9, 2013.

Jiawei Huang and Hu Ding. An effective manifold-based optimization method for distributionally robust classification. In The Thirteenth International Conference on Learning Representations , 2025. URL https://openreview.net/forum?id=nzjSvVZBIp .

Ruomin Huang, Jiawei Huang, Wenjie Liu, and Hu Ding. Coresets for wasserstein distributionally robust optimization problems. Advances in Neural Information Processing Systems , 35: 26375-26388, 2022.

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 International 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://proceedings.mlr.press/v139/ji21c.html .

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.

Mahmut Kaya and Hasan ¸ Sakir Bilge. Deep metric learning: A survey. Symmetry , 11(9):1066, 2019.

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.

Masanari Kimura and Hideitsu Hino. A short survey on importance weighting for machine learning, 2024. URL https://arxiv.org/abs/2403.10175 .

Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton. Imagenet classification with deep convolutional neural networks. Advances in neural information processing systems , 25, 2012.

Daniel Levy, Yair Carmon, John C Duchi, and Aaron Sidford. Large-scale methods for distributionally robust optimization. Advances in Neural Information Processing Systems , 33: 8847-8860, 2020.

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: //jmlr.org/papers/v20/18-700.html .

Tianyi Lin, Nhat Ho, and Michael I Jordan. On the efficiency of entropic regularized algorithms for optimal transport. Journal of Machine Learning Research , 23(137):1-42, 2022.

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.

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.

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. URL https://doi.org/10.1109/ICCV.2015.425 .

Ziwei Liu, Zhongqi Miao, Xiaohang Zhan, Jiayun Wang, Boqing Gong, and Stella X Yu. Largescale long-tailed recognition in an open world. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 2537-2546, 2019.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations , 2018.

Hongseok Namkoong and John C Duchi. Stochastic gradient methods for distributionally robust optimization with f-divergences. Advances in neural information processing systems , 29, 2016.

Vern I Paulsen and Mrinal Raghupathi. An introduction to the theory of reproducing kernel Hilbert spaces , volume 152. Cambridge university press, 2016.

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.

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 http://proceedings. mlr.press/v119/pham20a.html .

Shiori Sagawa, Pang Wei Koh, Tatsunori B Hashimoto, and Percy Liang. Distributionally robust neural networks. In International Conference on Learning Representations , 2019.

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.

Dvir Samuel and Gal Chechik. Distributional robustness loss for long-tail learning. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 9495-9504, 2021.

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 on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 31408-31421. PMLR, 23-29 Jul 2023. URL https://proceedings.mlr.press/v202/shi23j.html .

Jun Shu, Qi Xie, Lixuan Yi, Qian Zhao, Sanping Zhou, Zongben Xu, and Deyu Meng. Metaweight-net: Learning an explicit mapping for sample weighting. Advances in neural information processing systems , 32, 2019.

Richard Sinkhorn. Diagonal equivalence to matrices with prescribed row and column sums. The American Mathematical Monthly , 74(4):402-405, 1967.

Agnieszka Słowik and Leon Bottou. On distributionally robust optimization and data rebalancing. In Gustau Camps-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 https://proceedings.mlr.press/v151/slowik22a.html .

Andrew M. Stuart and Marie-Therese Wolfram. Inverse optimal transport. SIAM Journal on Applied Mathematics , 80(1):599-619, 2020. doi: 10.1137/19M1261122. URL https: //doi.org/10.1137/19M1261122 .

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-47680. PMLR, 21-27 Jul 2024. URL https: //proceedings.mlr.press/v235/tan24f.html .

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.

Alexander B Tsybakov. Optimal aggregation of classifiers in statistical learning. The Annals of Statistics , 32(1):135-166, 2004.

Cédric Villani. Optimal Transport: Old and New , volume 338. Springer, 2009.

Jie Wang, Rui Gao, and Yao Xie. Sinkhorn distributionally robust optimization, 2023. URL https://arxiv.org/abs/2109.11926 .

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: Selfadaptive thresholding for semi-supervised learning. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=PDrUPTXJI_ A .

Zifan Wang, Yi Shen, Michael Zavlanos, and Karl H Johansson. Outlier-robust distributionally robust optimization via unbalanced optimal transport. Advances in Neural Information Processing Systems , 37:52189-52214, 2024.

Kilian Q Weinberger and Lawrence K Saul. Distance metric learning for large margin nearest neighbor classification. Journal of machine learning research , 10(2), 2009.

Eric Xing, Michael Jordan, Stuart J Russell, and Andrew Ng. Distance metric learning with application to clustering with side-information. Advances in neural information processing systems , 15, 2002.

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 https://openreview.net/forum?id=r1Ddp1-Rb .

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

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 be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it 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 are not discussed in the paper.

The authors are encouraged to create a 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 on 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 a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not 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 to address problems of privacy and fairness.

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 acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 a complete (and correct) proof?

Answer: [Yes]

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to 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 experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of 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 a 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 it 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 to 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 a 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 submissions 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 it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or 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 reproducibility. In the case of closed-source models, it 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 to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: included

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a 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: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to 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 a 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 as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental setting/details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 a 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

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

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, confidence intervals, or statistical significance tests, at least for the experiments that support the 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, or 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).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

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 a 96% CI, if the hypothesis of Normality of errors 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 of 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

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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 of compute required for each of the individual experimental runs as well as 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 didn't make it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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 circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in 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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to

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 to 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.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to 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

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

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 a 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or 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?

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 or dataset.

The authors should state which version of the asset is used and, if possible, include a 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 (e.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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

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 of the dataset/code/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 is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an 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 as 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 in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. 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 an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: does not involve crowdsourcing nor 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, or originality of the research, declaration is not required.

Answer: [NA]

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 ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

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 on estimating pairwise distance metrics from observed interactions, typically constraining solutions to specific metric classes such as Mahalanobis distances [66, 65, 18]. However, while metric learning directly optimizes distances, inverse optimal transport (IOT) [57, 39, 13, 48] employs a more sophisticated bilevel optimization structure that accounts for global coupling relationships between point sets, significantly increasing its complexity. Cuturi and Avis [17] conceptualized the inverse OT problem as a metric learning task, an approach later applied to generative models [24]. Taking a different direction, Frogner et al. [21] proposed constructing cost functions using external priors such as word embeddings, where semantic similarity between labels is 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.

Distributional robustness. Distributionally robust optimization (DRO) [38, 52] provides a principled 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 a predefined uncertainty set around the training distribution. A key ingredient of DRO is its dual formulation , which enables tractable reformulations and efficient algorithms [31]. This paradigm has been extended to group-wise robustness, as in Group DRO [51], which seeks to ensure balanced performance across predefined or latent subgroups-even when some groups are underrepresented. Recent methods such as [59] further refine this idea by incorporating logit adjustment to improve performance parity across groups without explicit group labels.

Other related robustness methods. Another prominent class of techniques for handling distribution 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 a meta-learning framework that learns a sample-weighting function using a small unbiased validation set, effectively addressing issues like class imbalance and label noise. Fang et al. [19] introduced an 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) [8, 7] serves as 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 high-dimensional 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 infinite-dimensional space while enabling computations through finite-dimensional kernel functions. We begin by defining the key components that constitute a Reproducing Kernel Hilbert Space.

Definition 3 (Reproducing Kernel) . Let H be a Hilbert space of functions defined on a non-empty set X . A function Ker : X × X → R is called a reproducing kernel for H if: (i) For all x ∈ X , Ker( · , x ) ∈ H . (ii) For all f ∈ H and x ∈ X , the reproducing property holds: f ( x ) = ⟨ f, Ker( · , x ) ⟩ H .

Definition 4 (Reproducing Kernel Hilbert Space) . A Hilbert space H of functions on X is called a Reproducing Kernel Hilbert Space if there exists a reproducing kernel Ker for H .

The reproducing property ensures that the evaluation of a function f ∈ H at a point x ∈ 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 (10) 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 ( x, y ) ≥ 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 ν . 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, as the regularization parameter ϵ → 0 , the dual objective of Sinkhorn DRO converges to that of Wasserstein DRO because the Sinkhorn distance is an entropy-regularized version of the Wasserstein distance. (2) Connection to KL DRO: When the reference distribution ν = P tr , and the transport cost function C ( · , · ) ≡ 0 , Sinkhorn DRO reduces to the KL-DRO problem. In 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.

A.4 Solutions for Entropy-Regularized Optimal Transport

For completeness, we summarize the closed-form solutions for entropy-regularized optimal transport (OT) problems under different levels of marginal constraint relaxation. These results are standard in the OT literature and can be found in [53]. Such relaxations provide analytical forms for the optimal transport plan γ , which significantly reduce computational complexity and enable efficient integration into modern learning algorithms.

Remark 2. Let Π( P, Q ) denote the set of couplings (joint probability measures) with specified marginals:

Relaxed Marginal Constraint: γ ∈ Π( P ) This semi-relaxed formulation preserves only the source marginal. The entropy-regularized OT problem is defined as:

where H ( γ | P ⊗ ν ) denotes the relative entropy (defined in (2)) with respect to the product measure of P and a reference measure ν . The optimal coupling has a closed-form expression:

This expression yields a softmax-like kernel around each point p , shaped by the cost C and regularization ϵ , and normalized with respect to the reference measure ν .

Fully Relaxed Constraint: γ ∈ Π(1) In this most relaxed case, only the total mass is constrained to be 1:

The optimal transport plan admits the following form:

This represents a global Gibbs distribution over the joint space X × Y and is often used for unconstrained matching.

Fully Constrained OT: γ ∈ Π( P, Q ) This is the standard OT formulation with both source and target marginals enforced. In this case, no closed-form solution exists, but iterative scaling via the Sinkhorn algorithm [55, 6] yields the unique entropy-regularized solution:

This procedure is guaranteed to converge to the unique solution γ ∗ ∈ Π( P, Q ) satisfying the marginal constraints.

These closed-form and iterative solutions illustrate how different constraint relaxations influence the structure of optimal couplings. In our framework (see Sec. 3.1 and Sec. 3.2), we make use of the semi-relaxed formulation γ ∈ Π( P ) due to its analytical tractability and its suitability for dynamic uncertainty set modeling under limited supervision.

A.5 Proof of Lemma 3.1

Proof. We begin by considering the semi-relaxed optimal transport (OT) problem where the coupling γ preserves only the source marginal P tr , with reference distribution ν , and a modified transport cost:

We seek the optimal coupling γ ∗ ∈ Π( P tr ) that minimizes the entropy-regularized OT objective:

Applying the standard solution for entropy-regularized semi-relaxed OT (see Appendix A.4), the optimal coupling has the form:

Substituting the definition of C ′ ( p, q ) , we obtain:

Taking the marginal over p ∼ P tr , the target marginal (i.e., the induced distribution on q ) becomes:

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 ¯ γ

In the main text, we introduced a filtering mechanism for the special case of one-to-one ground-truth matchings. Here, we present a generalized formulation applicable to many-to-many matchings, which naturally extends our approach while preserving its key properties.

Generalized Matching Quality Metrics. For a given source p , we define its matching quality as:

which measures the expected alignment between the model-derived coupling γ θ ( p, q ) and the groundtruth 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:

providing a dataset-wide measure of matching performance.

For each target q , we quantify its matching capacity by:

representing the total transport mass received by target q from all sources. A higher value of τ 2 ( q ) indicates that q is easier to match, consistent with our intuition in the main text.

Generalized Filtering Mechanism. Using these metrics, we compute a dynamic threshold for each target q :

which adjusts the global threshold τ 1 based on the relative matching capacity of target q .

For each source p , we compute a personalized threshold by taking the expectation over all its associated targets:

The generalized filtering rule retains source p in the filtered distribution if its matching quality exceeds its personalized threshold:

where χ normalizes ̂ ν to a 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 to One-to-One Case) . When the ground-truth coupling ¯ γ is one-to-one (i.e., for each source p , there exists exactly one target q p such that ¯ γ ( p, q p ) = 1 and ¯ γ ( p, q ) = 0 for all q = q p ), the generalized filtering rule in (40) reduces to the one-to-one filtering rule in (16) .

̸

Proof. For a one-to-one ground-truth coupling ¯ γ , each source p is associated with exactly one target q p where ¯ γ ( p, q p ) = 1 and ¯ γ ( p, q ) = 0 for all q = q p . Under this condition:

(i) The source matching quality simplifies to:

(ii) The source-specific threshold becomes:

(iii) Therefore, the filtering condition τ 0 ( p ) ≥ τ ( p ) becomes:

This is equivalent to the condition in (16): " ∃ q : ¯ γ ( p, q ) > 0 ∧ γ θ ( p, q ) ≥ τ ( q ) " since q p is the only target for which ¯ γ ( p, q ) > 0 . 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.

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 to the biased gradient estimation.

Weprove 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.1). Combined with our Lipschitz semantic transport cost, we obtain the continuity of the worst-case distribution ∥ ̂ Q λ t +1 -̂ Q λ t ∥ TV (Theorem B.4).

To handle nested gradient computation, we employ Multi-Level Monte Carlo with Random Truncation (MLMC-RT), showing that bias decays quadratically with approximation level ( O (2 -L ) ) while variance increases only logarithmically (Proposition B.5).

By carefully integrating these results with stochastic optimization theory, we prove that our algorithm with decaying stepsize α t = α/ ( √ t +1) achieves min t<T E ∥∇ F t ( θ t ) ∥ 2 ≤ ϵ in T = O ( ϵ -4 log 1 ϵ ) iterations, with total sample complexity O ( ϵ -4 log 2 1 ϵ ) (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 1 (Smoothness of the Model) . (i) The loss function ℓ ( θ, x ) is L θ -smooth in θ : ∥∇ ℓ ( θ 1 , x ) - ∇ ℓ ( θ 2 , x ) ∥ ≤ L θ ∥ θ 1 -θ 2 ∥ ∀ x ∈ X ; (ii) The induced semantic cost is Lipschitz continuous for any p, q ∈ P (Ξ) : ∥ C θ 1 ( p, q ) -C θ 2 ( p, q ) ∥ ≤ L C ∥ θ 1 -θ 2 ∥ .

Assumption 2 (Lipschitz Continuity of Matching Quality) . (i) The matching quality function γ θ ( p, q ) is L γ -Lipschitz continuous in θ :

(ii) The threshold function τ ( q ; θ ) is L τ -Lipschitz continuous in θ :

These assumptions abstract the key continuity properties required for our theoretical results. Assumption 1 ensures that both the model loss and the resulting semantic cost vary smoothly with respect to model parameters. Assumption 2 guarantees that the OT-based matching quality and the adaptive threshold used in filtering are stable under small changes in the model. In Section C, we 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 ¯ γ , where each source p is associated with one target q p such that ¯ γ ( p, q p ) > 0 , we assume that for any ε > 0 and θ , the probability mass of points in the ε -neighborhood of the decision boundary is bounded:

for some constant L B > 0 independent of θ and ε .

Assumption 4 (Non-degenerate Filtering) . For the ground-truth coupling ¯ γ , where each source p is associated with one target q p such that ¯ γ ( p, q p ) > 0 , we assume there exists δ > 0 such that for all θ in the parameter space, ν ( S θ ) ≥ δ , where S θ = { p : γ θ ( p, q p ) ≥ τ ( q p ; θ ) } is 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-analogous to the Tsybakov margin condition in statistical learning [60]-prevents pathological scenarios where infinitesimal parameter changes could cause discontinuous shifts in the filtered distribution. Non-Degenerate Filtering (Assumption 4) guarantees

the filtered distribution maintains a minimum effective support throughout training. Implementationally, this is 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 B.1), which is essential for the convergence guarantees of our adaptive robust optimization algorithm.

Algorithm Overview We analyze the convergence properties of our proposed iterative procedure (Algorithm 1) for updating parameters θ ∈ R d . At iteration t , let ̂ ν θ t represent the filtered reference distribution and C θ t denote the transport cost. The algorithm proceeds through the following steps at each iteration t :

Reference Distribution Filtering: Given θ t , we filter the reference distribution ν to obtain ̂ ν θ t by removing low-confidence samples according to:

Semantic Cost Adaptation: We update the transport cost function to C θ t , capturing semantic distances in both feature and label spaces.

Worst-Case Distribution Computation: Using ̂ ν θ t and C θ t , we compute the Sinkhorn DRO worst-case distribution ̂ Q λ t with fixed regularization parameter λ :

where Z θ t ( p ) = E u ∼ ̂ ν θ t ( e ℓ ( u ) -λC θ t ( p,u ) ) .

Gradient-Based Update: We perform a gradient step on the iteration-specific objective:

We will show that if ̂ ν θ t and C θ t vary "slowly" with respect to θ t , and F t is smooth in θ , then the sequence { θ t } converges to a stationary point under standard non-convex optimization conditions.

B.2 Continuity of Filtered Distributions

A key aspect of our 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 A.6.

̸

For a one-to-one ground-truth coupling ¯ γ , each source p is associated with exactly one target q p such that ¯ γ ( p, q p ) > 0 and ¯ γ ( p, q ) = 0 for all q = q p . The filtering operation is defined as:

where χ θ is a normalization constant ensuring ̂ ν θ is a probability distribution.

Proposition B.1 (Lipschitz Continuity of Filtered Distributions) . Under Assumptions 2-4, the filtering operation produces distributions that are Lipschitz continuous with respect to θ in the total variation distance:

holds for any θ 1 , θ 2 and some constant L TV > 0 .

Proof. We proceed in three steps:

(i) Bounding the symmetric difference. Let S 1 = { p | γ θ 1 ( p, q p ) ≥ τ ( q p ; θ 1 ) } and S 2 = { p | γ θ 2 ( p, q p ) ≥ τ ( q p ; θ 2 ) } .

For p ∈ S 1 \ S 2 , we have: γ θ 1 ( p, q p ) ≥ τ ( q p ; θ 1 ) and γ θ 2 ( p, q p ) < τ ( q p ; θ 2 )

Using the Lipschitz continuity from Assumption 2:

where L γ and L τ are the Lipschitz constants from Assumption 2.

For p ∈ S 1 \ S 2 , these inequalities imply:

Similarly, for p ∈ S 2 \ S 1 , we obtain:

Therefore, the symmetric difference S 1 △ S 2 = ( S 1 \ S 2 ) ∪ ( S 2 \ S 1 ) is contained within:

By Assumption 3, we have:

where L ν = L B · ( L γ + L τ ) and L B is the constant from Assumption 3.

(ii) Bounding normalization constants. The normalization constants are χ θ 1 = ν ( S 1 ) and χ θ 2 = ν ( S 2 ) , with difference:

By Assumption 4, both χ θ 1 and χ θ 2 are bounded below by δ > 0 , where δ is the constant from Assumption 4.

(iii) Bounding the total variation distance. The total variation distance between ̂ ν θ 1 and ̂ ν θ 2 is:

We decompose this into:

For the first term:

where we used ν ( A ∩ S 1 ) ≤ χ θ 1 and χ θ i ≥ δ .

For the second term:

Combining both terms:

Since this bound holds for any measurable set A , we have:

B.3 Stability Analysis of the Worst-case Distribution

To guarantee convergence of our algorithm, it is 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 it inherits the L θ -smoothness of the loss function with respect to model parameters.

Lemma B.2 (Smoothness of Sinkhorn DRO Objective) . Under Assumption 1, at any iteration t , the Sinkhorn DRO objective F t ( θ ) = E q ∼ ̂ Q λ t [ ℓ ( θ, q )] is also L θ -smooth.

Proof. By Assumption 1, the loss function ℓ ( θ, q ) is L θ -smooth in θ for all q :

For the objective function F t ( θ ) = E q ∼ ̂ Q λ t [ ℓ ( θ, q )] , where ̂ Q λ t is fixed for a given t , we have:

Therefore, the objective function F t ( θ ) inherits the L θ -smoothness from the loss function ℓ ( θ, q ) .

To analyze the stability of our algorithm, we need to understand how the worst-case distribution changes with model parameters. Let ̂ Q λ t be the worst-case distribution computed using parameter θ t and regularization λ :

where Z θ t ( p ) = E u ∼ ̂ ν θ t ( e ℓ ( θ t ,u ) -λC θ t ( p,u ) ) . The total variation (TV) distance provides a natural metric for quantifying the difference between probability distributions.

Proposition B.3. Let ( X , µ ) be a measure space. Let f, g : X → R be measurable functions with finite log-partition functions Z f = ∫ X e f ( x ) d µ ( x ) , Z g = ∫ X e g ( x ) d µ ( x ) < ∞ . Define probability densities d P ( x ) = e f ( x ) Z f d µ ( x ) , d Q ( x ) = e g ( x ) Z g d µ ( x ) . Then ∥ P -Q ∥ TV ≤ √ ∥ f -g ∥ ∞ .

Proof. We use Pinsker's inequality to relate the total variation distance to the KL divergence:

For the KL divergence between P and Q , we have:

Since | f ( x ) -g ( x ) | ≤ ∥ f -g ∥ ∞ for all x , we have ∫ X e f ( x ) Z f ( f ( x ) -g ( x ))d µ ( x ) ≤ ∥ f -g ∥ ∞ .

For the partition function ratio, since g ( x ) ≤ f ( x ) + ∥ f -g ∥ ∞ for all x , we have:

Therefore, | log Z g Z f | ≤ ∥ f -g ∥ ∞ . Combining these bounds:

Substituting into Pinsker's inequality: ∥ P -Q ∥ TV ≤ √ 1 2 · 2 ∥ f -g ∥ ∞ = √ ∥ f -g ∥ ∞

Thus, we have established ∥ P -Q ∥ TV ≤ √ ∥ f -g ∥ ∞ .

Theorem B.4 (Continuity of the Worst-Case Distribution) . Under Assumptions 1-4, for consecutive iterations t and t +1 :

where G = G ℓ + λL C + L TV combines the gradient bound G ℓ = sup q ∈X ∥∇ θ ℓ ( θ, q ) ∥ , the Lipschitz constant L C of the transport cost, and the Lipschitz constant L TV of the filtered distribution.

Proof. Recall that the worst-case distribution at iteration t is given by:

We introduce an intermediate distribution to separate the effect of parameter changes on conditional weights and reference distribution:

By the triangle inequality:

For the first term, since ̂ Q λ t +1 and ˜ Q λ t +1 differ only in their reference distributions:

where the last inequality follows from Proposition B.1.

For the second term, we analyze the change in conditional weights. For each fixed p , define the functions f t ( q ) = ℓ ( θ t , q ) -λC θ t ( p, q ) and f t +1 ( q ) = ℓ ( θ t +1 , q ) -λC θ t +1 ( p, q ) . By proposition B.3, we have:

We bound ∥ f t +1 -f t ∥ ∞ using Assumption 1:

For the loss function, by Taylor's theorem with the smoothness property:

For the transport cost, by Lipschitz continuity:

Combining these bounds:

Since this bound holds for all p , taking expectation over P tr preserves the inequality:

Combining the bounds (79) and (82) into (78):

where G = G ℓ + λL C + L TV .

This establishes that the worst-case distribution has a 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 λ . The quadratic term reflects the impact of the loss function's smoothness. This continuity property is crucial for establishing the convergence of our algorithm in 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 C θ and the filtered distribution ̂ ν θ with respect to model parameters θ in Sec. B.3. To tackle the second challenge, we 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 [28]. In our setting, we aim to optimize the following non-convex objective at each iteration t :

where Z θ t ( p ) = E q ∼ ̂ ν θ t [ e ℓ ( q ) -λC θ t ( p,q ) ] is a normalization factor. Because of the nested expectation and normalization term, directly obtaining unbiased gradients is intractable.

To overcome this difficulty, we introduce a sequence of approximations indexed by level l . Specifically, we approximate the Sinkhorn DRO objective F t ( θ ) via the following levell Monte Carlo approximation:

where:

̂ ν ( l ) θ t denotes a Monte Carlo approximation of ̂ ν θ t using 2 l independent and identically distributed (i.i.d.) samples from ̂ ν θ t , i.e., ̂ ν ( l ) θ t = { q 1 , . . . , q 2 l } ,

Z ( l ) θ t ( p ) = E u ∼ ̂ ν ( l ) θ t [ e ℓ ( u ) -λC θ t ( p,u ) ] is the corresponding normalization factor at level l .

Algorithm 2 MLMC-RT Gradient Estimation for Sinkhorn DRO

Require: Current parameter θ t , target accuracy ϵ , constants a = 2 , b = c = 1

1: Set the maximum level L = ⌈ log 2 ( ϵ -1 ) ⌉

2: Define sampling probabilities prob ( l ) = 2 -( l +1) / (1 -2 -( L +1) ) for all l ∈ { 0 , 1 , . . . , L }

3: Randomly sample a level l ∼ { prob (0) , prob (1) , . . . , prob ( L ) }

4: Draw 2 l i.i.d. samples from ̂ ν θ t to form ̂ ν ( l ) θ t

6: Compute G ( l ) t ( θ t , ζ ( l ) ) (from (86)) and G ( l -1) t ( θ t , ζ ( l -1) ) if l > 0

5: If l > 0 , reuse the first 2 l -1 samples to construct ̂ ν ( l -1) θ t

7: Return the gradient estimate:

MLMC-RT gradient idea. The gradient ∇ θ F t ( θ ) is approximated by estimating the gradient of F ( l ) t ( θ ) :

where ζ ( l ) denotes the random samples used at level l . As l → ∞ , G ( l ) t ( θ, ζ ( l ) ) → ∇ F t ( θ ) . 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 1-4 and bounded loss and transport cost, the sequence { F ( l ) t ( θ ) } ∞ l =0 satisfies:

Squared bias parameter a : ∥∇ F ( l ) t ( θ ) - ∇ F t ( θ ) ∥ 2 ≤ M a 2 -al with a = 2 (quadratic decay)

Variance parameter b : For the difference estimator H ( l ) t ( θ, ζ ( l ) ) := grad ( l ) t ( θ, ζ ( l ) ) -grad ( l -1) t ( θ, ζ ( l -1) ) , we have the variance Var( H ( l ) t ( θ, ζ ( l ) )) ≤ M b 2 -bl with b = 1 (linear decay)

Cost parameter c : The cost of computing H ( l ) t ( θ, ζ ( l ) ) is bounded by C ( l ) ≤ M c 2 cl with c = 1 (linear growth)

Here M a , M b and M c are all constants.

Proof. For the bias parameter a , note that ∇ F t ( θ ) involves expectation over ̂ ν θ , while ∇ F ( l ) t ( θ ) approximates this with 2 l i.i.d. samples. Under standard Monte Carlo convergence results for bounded smooth gradients, we have ∥∇ F ( l ) t ( θ ) -∇ F t ( θ ) ∥ = O (2 -l ) , so the squared bias is O (2 -2 l ) , implying a = 2 .

For variance, the difference estimator is:

where ζ ( l ) | 2 l -1 reuses the first 2 l -1 samples. This shared randomness ensures that the two estimates are highly correlated, and their difference has variance O (2 -l ) , hence b = 1 .

For cost, both G ( l ) t and G ( l -1) t require evaluating gradients for 2 l and 2 l -1 samples respectively. Since 2 l +2 l -1 = 1 . 5 · 2 l , the cost is O (2 l ) , implying c = 1 .

Remark 3. The efficiency of MLMC-RT depends crucially on three problem-specific parameters: (i) the bias decay rate a (how fast the approximation error decreases with level), (ii) the variance

decay rate b (how fast the variance of level differences decreases), (iii) and the cost growth rate c (how expensive each level is to evaluate). The above proposition establishes that, for our Sinkhorn DRO setting, the objective gradient satisfies the standard MLMC assumptions with parameters ( a, b, c ) = (2 , 1 , 1) . This result is key to applying MLMC-RT effectively and proving convergence with nearly optimal sample complexity.

B.4.2 MLMC-RT Gradient Estimation

The stochastic gradients grad t employed in our optimization procedure are computed using a MultiLevel Monte Carlo estimator with Randomized Truncation (MLMC-RT), following the methodology developed in Hu et al. [28].

Since computing exact expectations over ̂ Q λ t is intractable due to nested expectations and normalizations, MLMC-RT approximates the gradient via:

Randomly selecting a level l ∈ { 0 , 1 , . . . , L } according to a probability distribution { prob ( l ) } , where prob ( l ) ∝ 2 -( b + c ) l/ 2 ,

Computing a reweighted difference estimator based on samples at level l .

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 t satisfies E [ grad t ] = ∇ θ F ( L ) ( θ t ) rather than ∇ θ F ( θ t ) , hence it is a biased estimator of the true gradient. However, by carefully choosing L = ⌈ 1 a log(4 M a ϵ -1 ) ⌉ , we ensure that this bias remains controlled and does not impair the convergence of our algorithm.

The detailed procedure for MLMC-RT in our setting is as follows:

The MLMC-RT estimator enjoys the following theoretical guarantees:

Lemma B.6 (Properties of MLMC-RT Estimator) . Let grad t be generated according to Algorithm 2. Then:

Controlled Bias: ∥ E [ grad t ] -∇ F t ( θ t ) ∥ ≤ √ M a ϵ ,

Logarithmic Variance: E [ ∥ grad t -E [ grad t ] ∥ 2 ] ≤ K log( ϵ -1 ) for some constant K > 0 ,

Efficient Computation: The expected computation cost per iteration is O (log( ϵ -1 )) .

Proof. The controlled bias follows from truncating at level L , which ensures ∥∇ F ( L ) t ( θ ) -∇ F t ( θ ) ∥ 2 ≤ M a 2 -aL ≤ M a ϵ 2 , giving the bias bound ∥ E [ grad t ] -∇ F t ( θ t ) ∥ ≤ √ M a ϵ .

For the variance, with b = c = 1 and prob ( l ) ∝ 2 -l , we have:

We set prob ( l ) = 2 -l ∑ L k =0 2 -k = 2 -l 1 -2 -( L +1) (which is approximately 2 -l -1 for large L ). Substituting:

This is consistent with the results from [28], Table 1, for the MLMC-RT estimator with a = 2 , b = c = 1 .

We now analyze the computational cost of generating the MLMC-RT gradient estimate at each iteration. Let W t 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 t .

Since the MLMC-RT procedure selects a random level l according to the distribution { prob (0) , prob (1) , . . . , prob ( L ) } , and computing the estimator at level l requires a total cost W l proportional to 2 l , the expected computational cost per iteration is given by:

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 a logarithmic growth of variance, significantly better than the O ( ϵ -1 ) 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 our adaptive Sinkhorn DRO algorithm under MLMCRT gradient estimation.

Theorem B.7 (Convergence of AdaDRO) . Suppose Assumptions 1-4 hold, and the iterates are bounded as ∥ θ t ∥ ≤ M . Using the MLMC-RT gradient estimator with maximum level L = ⌈ log 2 ( ϵ -1 ) ⌉ and probability distribution prob ( l ) ∝ 2 -l , and setting the step size as α t = α √ t +1 for some constant α > 0 , we have that after T = Θ ( ϵ -4 log 1 ϵ ) iterations, Algorithm 1 satisfies

with a total sample complexity of O ( ϵ -4 log 2 1 ϵ ) .

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 L θ -smoothness of F t (Lemma B.2):

To account for the changing objective function, we add and subtract F t +1 ( θ t +1 ) to obtain:

Step 2: Bounding the adaptive drift. Let M ℓ := sup q | ℓ ( θ, q ) | , by Theorem B.4, we have:

Step 3: MLMC-RT bias-variance decomposition. For the key term E [ ⟨∇ F t ( θ t ) , grad t ⟩ ] , we apply the bias-variance decomposition:

From Lemma B.6, we have ∥ E [ grad t ] - ∇ F t ( θ t ) ∥ ≤ √ M a ϵ , by Cauchy-Schwarz inequality, yielding:

For the second-moment term, using the variance bound from Lemma B.6:

Step 4: Combining the bounds. Combining the estimates from Steps 1-3, and pluging (94), (95) and (96) into (93) gives

Substituting the bound for E [ ∥ grad t ∥ 2 ] from (96):

Bias cross-term. Using Young's inequality ab ≤ 1 4 a 2 + b 2 with a = √ α t ∥∇ F t ( θ t ) ∥ , b = √ M a α t ϵ ,

Adaptive-drift linear term. First apply Jensen: E ∥ grad t ∥ ≤ √ E ∥ grad t ∥ 2 ; then plug (96) and use Young's inequality again to obtain

Quadratic term from smoothness. For the quadratic terms involving ∥∇ F t ( θ t ) ∥ 2 in (98), choosing α small enough such that α t L θ (1 + M ℓ ) ≤ 1 4 for all t , we have

Combining all pieces (99)-(101) into (98) gives

Step 5: Telescoping. Summing (102) for t = 0 , . . . , T -1 and setting ∆ F := E [ F 0 ( θ 0 )] -inf θ F ( θ ) , we have:

Since α t = α √ t +1 , we know ∑ T -1 t =0 α t = O ( √ T ) and ∑ T -1 t =0 α 2 t = O (log T ) . Because ∑ T -1 t =0 α t ≥ α 2 √ T , we deduce:

Choosing T = Θ( ϵ -4 log ϵ -1 ) ensures log ϵ -1 log T T = O ( ϵ 4 log ϵ -1 ) , and log ϵ -1 T log T = O ( ϵ 4 ) . Hence

Each iteration costs O (log ϵ -1 ) samples (Lemma B.6), so the total sample complexity is

The sample complexity O ( ϵ -4 log 2 ( ϵ -1 )) matches the lower bound Ω( ϵ -4 ) [1] for general nonconvex stochastic optimization only by logarithmic factors.

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 C X and label transport cost C Y are Lipschitz continuous with respect to θ .

Proof. For the feature transport cost C X θ ( x i , x j ) = c -Ker( g θ ( x i ) , g θ ( x j )) , we assume:

(i) The encoder g θ is L g -Lipschitz continuous with respect to parameters θ :

(ii) The kernel function Ker( · , · ) is L K -Lipschitz continuous in both arguments:

These are reasonable assumptions for common kernel functions like RBF kernels or dot product kernels with bounded inputs.

For the feature transport cost:

Thus, C X θ is Lipschitz continuous with constant L C X = 2 L K L g .

For the label transport cost C Y kk ′ = c -Ker( o k , o k ′ ) where o k = w k (the classifier weights for class k ), a similar argument holds. Since w k is a subset of the parameters θ , we have:

Therefore:

Thus, C Y θ is Lipschitz continuous with constant L C Y = 2 L K .

The combined transport cost C θ ( p i , p j ) = C X θ ( x i , x j )+ C Y θ ( y i , y j ) is therefore Lipschitz continuous with constant L C = 2 L K ( L g +1) .

Lemma C.2 (Matching Quality Lipschitz Continuity) . Under suitable conditions on the encoder network and kernel function, the matching quality function γ θ ( p, q ) = γ θ ( p, q ) is L γ -Lipschitz continuous in θ .

Proof. We fix p and view γ θ ( p, q ) as a probability density over q :

Let f ( q ) = -C θ 1 ( p, q ) /ϵ , g ( q ) = -C θ 2 ( p, q ) /ϵ . Then, γ θ i ( p, · ) for i = 1 , 2 defines a probability density over ( Q , ν ) of the form d P i ( q ) = e f ( q ) Z f d ν ( q ) . By Proposition B.3:

From Lemma C.1, C θ is L C -Lipschitz in θ , so

This implies pointwise Lipschitz continuity:

Hence, γ θ ( p, q ) is L γ -Lipschitz in θ with L γ = √ L C /ϵ .

Lemma C.3 (Threshold Lipschitz Continuity) . Under the conditions established in Lemma C.2, the threshold function τ ( q ; θ ) is L τ -Lipschitz continuous in θ .

Proof. The threshold function is defined as:

where τ 1 ( θ ) = E ( p,q ) ∼ ¯ γ γ θ ( p, q ) and τ 2 ( q ; θ ) = E p γ θ ( p, q ) .

From Lemma C.2, γ θ ( p, q ) is L γ -Lipschitz. Therefore:

Similarly, τ 2 ( q ; θ ) is L γ -Lipschitz, and:

Let r ( q ; θ ) = τ 2 ( q ; θ ) max q ′ τ 2 ( q ′ ; θ ) . Under regularity conditions ensuring that τ 2 ( q ; θ ) is bounded and max q ′ τ 2 ( q ′ ; θ ) is bounded away from zero, we can establish that r ( q ; θ ) is L r -Lipschitz.

Now consider the full threshold function:

where L τ = L r + L γ = O ( 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.

D Experimental Results

We evaluate the proposed AdaDRO framework across a diverse set of benchmarks to 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], a real-world dataset with strong background-label correlation; CelebA [43], a benchmark for facial attribute recognition. We also evaluate on several long-tailed benchmark datasets: CIFAR-10-LT and CIFAR-100-LT [37].

Baselines. We compare AdaDRO with the following well-known methods: KL-DRO [46], Wasserstein DRO ( WDRO ) [9], Sinkhorn DRO SDRO [62], a unified DRO method UniDRO [11], semanticaware robust methods such as GroupDRO [50], an invariant risk minimization method IRM [2].

For fairness, we report the average accuracy and the distribution shift gap (drop in performance under shift). All models are implemented with PyTorch on a single NVIDIA RTX 6000 Ada GPU. The model is optimized using the AdamW optimizer [45]. For AdaDRO, we use cosine similarity as the kernel in (10), and employ basic augmentations (flip, crop) for semantic calibration in Sec. 3.2.

D.2 Robustness to Spurious Correlation and Group Shift.

AdaDRO is broadly applicable thanks to its flexible reference distribution design. In imbalanced settings, ν can be constructed via reweighting or resampling. In other tasks, domain priors can be easily 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 5. This demonstrates its ability to 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 6 show the results on Waterbirds.

Long-tail Generalization. We conduct our experiments on two widely used long-tailed benchmark datasets: CIFAR-10-LT and CIFAR-100-LT . These datasets are derived from the original CIFAR-10 and CIFAR-100 by artificially reducing the number of training examples in tail classes to simulate long-tailed distributions. The imbalance in these datasets is 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: IF=10 , IF=50 , and IF=100 , representing increasing levels of class imbalance.

Baselines. We compare AdaDRO with a diverse set of baselines tailored for long-tailed learning: Decouple [27]: 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-classify (often tail) examples; DRO-LT [52]: A distributionally robust optimization framework specifically adapted for long-tailed classification; SSL [35]: A contrastive learning-based method that enhances representations for tail classes via augmented positive samples; Resample [15]: Performs oversampling of minority classes to reduce imbalance; Reweight [15]: Adjusts sample weights in the loss function based on inverse class frequency; SinkhornDRO [62]: A recent OT-based DRO method that regularizes transport cost to mitigate over-pessimism in robustness.

Results and Insights. As shown in Table 7, our method consistently outperforms all baselines across different imbalance settings on both CIFAR-10-LT and CIFAR-100-LT. 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 baselines. This highlights the advantage of our adaptive distributional robustness formulation in focusing learning on underrepresented regions while maintaining calibrated semantic alignment. Table 9 details the quantitative comparison between our method and the strong baseline ERM+Aug across standard DomainBed benchmarks [25].

D.3 Efficiency Analysis.

Tables 10 and 11 report normalized total training time. While AdaDRO incurs moderate overhead due to semantic alignment and filtering, it remains tractable and justifiable given the performance gains.

D.4 Ablation Studies

We perform ablation studies to assess the contribution of key components in AdaDRO. Detailed results are shown in Tables 12-14.

(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.

(ii) Adaptive Filtering. We replace the filtered reference distribution ̂ ν t with a fixed one, either generated by Mixup [67] or uniform sampling. Without filtering, performance degrades due to irrelevant or noisy samples in ν .

(iii) Robustness to Uncertainty Radius. We compare AdaDRO with Sinkhorn DRO under varying uncertainty set radius λ . As shown in Table 14, AdaDRO maintains stable robustness even under aggressive uncertainty, whereas Sinkhorn DRO suffers from severe over-pessimism or instability.

D.5 Limitations

Our method relies on effective data augmentations to serve as a prior for 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 ¯ γ through semantic invariant data augmentation. In some domains, richer priors may be available. Effectively integrating such

heterogeneous sources of prior knowledge to build a more powerful coupling remains an open question.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nThis 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.\n\nQUESTIONS:\nQ1  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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nPros:\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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-03"}

comments: 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.

STRENGTHS AND WEAKNESSES:
Strengths:

The 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. 

The paper provides the theoretical convergence guarantees of the proposed method. 

Weaknesses:

The paper mentioned that: “This result demonstrates that robust optimization
with 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. 

It is not clear to the reviewer what the limitations of the proposed method are. 

In the experimental section of the main paper, it is not clear about the encoder g used in the paper.

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.
TITLE: Official Review
CLARITY: 2
IMPACT: 2
ORIGINALITY: 3
DATE: 2025-07-06

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nThe 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.\n\nQUESTIONS:\n- 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\n\nLIMITATIONS:\n\"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.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 1, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-06"}