Bootstrap Your Uncertainty: Adaptive Robust Classification Driven by Optimal-Transport

We thank the reviewer for the thoughtful comments, and address your concerns below.

Response to Weaknesses

W1: Could you summarize the time and space complexity to run these algorithms?

While entropy-regularized OT is often solved with the Sinkhorn's algorithm, our formulation allows for a closed-form for the transport plan (as seen in Appendix A.4). In essence, we only need to compute the cost matrix $C_\theta$ and then perform a few element-wise operations (like exp and log-sum-exp), which are highly optimized on modern hardware like GPUs. This means we bypass the need for any iterative OT solver, and the additional computational overhead per step is minimal.

W2 W3: The manuscript currently omits the background literature to illustrate the inverse optimal transport (IOT) and feature‐level alignment works...; ... hyperparameter sensitivity....

We will address the related works you mentioned in our revised version. Additionally, we will include more ablation studies; the corresponding tables are omitted here due to space limitations.

W4: It seems that the semantic calibration relies on data augmentations; if augmentations produce semantically invalid pairs...the paper lacks analysis of failure modes.

We will discuss this limitation. If the augmentations are too aggressive and create semantically invalid pairs (e.g., cropping a "cat" image to show only grass), the learned transport cost could be misled. In our experiments, we followed standard practices and used only mild data augmentations (such as random flipping and cropping), which are widely adopted in semi-supervised and self-supervised learning, and have been shown to effectively preserve the semantics of the samples.

Response to Questions

Q1 and Q5: Sensitivity to the choice of observed coupling $\bar{\gamma}$. Sensitivity to the choice and strength of augmentations.

In our current work, we define $\bar{\gamma}$ using a simple one-to-one matching between a sample and its augmented counterpart (11). We have not yet explored alternative sources for $\bar{\gamma}$. However, we hypothesize that using more sophisticated couplings could be a promising direction for future work. It might provide a richer, many-to-many semantic signal.

The learned cost is not sensitive to the augmentations. Data augmentation is a widely used technique to construct semantically invariant samples, especially in contrastive learning. Recent works have explored multi-level augmentations such as multi-crop [1], which creates stronger invariance by including different spatial scales. To study the sensitivity of our method to the augmentation strategy, we replaced our standard augmentations with multi-crop augmentation.

Empirically, we observe that the performance of our method remains stable under this change. In some cases, multi-crop leads to slight improvements, suggesting that our method is robust to the choice of augmentation. For instance, when applied to the Colored MNIST task with a spurious correlation of 0.8, we observe a 0.2% increase in accuracy. On CIFAR-100 with 50% label noise, multi-crop yields a 0.8% improvement. Due to space constraints, we omit detailed tables.

Q2: Guarantees on the diameter of the evolving uncertainty set.

The dynamic uncertainty set does not refer to a change in radius $\delta$, but to the adaptive adjustment of the uncertainty set's geometry and support. This is achieved through two core mechanisms: 1, dynamically refining the reference distribution $\nu$ into $\hat{\nu}$ via OT feedback (16); 2, calibrating the semantic transport cost $C_\theta$ via Inverse OT (Sec. 3.2). To prevent collapse of the uncertainty set, we introduce Assumption 4 (Non-degenerate Filtering), which ensures that the samples retained by our filtering mechanism maintain a minimum effective support throughout training. Our algorithmic design satisfies this assumption.

Empirically, while precisely tracking diameter evolution in high-dimensional space is challenging, our toy example in Fig. 1 intuitively demonstrates the process’s effectiveness. The method successfully identifies and removes low-quality outliers lying off the true data manifold, thereby producing a cleaner decision boundary. Furthermore, our ablation studies (Table 7) and experiments in Table 8 validate that this adaptive process both prevents performance collapse and significantly improves robustness. We plan to include additional visualizations in the revised version to better illustrate the dynamics of this process.

Q3: Alignment of OT loops and scalability.

Our framework entirely avoids iterative OT solvers. This is achieved because both our forward and inverse OT problems are formulated using a Semi-Relaxed Marginal Constraint, as detailed in Appendix A.4. In this setting, we only enforce the source marginal of the transport plan, which allows the optimal coupling $\gamma^*$ to have a closed-form analytical solution (23).

Therefore, our alignment plan per minibatch is highly efficient and non-iterative: We first compute the cost matrix $C_\theta$. We then directly obtain the transport plan $\gamma^*$ by applying the closed-form solution from (23), which only involves highly-parallelizable, element-wise operations (e.g., exp and a normalization) that are heavily optimized on GPUs. Both the exploration (OT) and alignment (inverse OT) components of our framework leverage this formulation, meaning neither requires an iterative loop like Sinkhorn. The computational complexity per step is dominated by the $O(nm)$ cost matrix computation, which is standard for many OT-based methods and is far more efficient than the $O(nm \times L)$ cost of a full Sinkhorn algorithm, where $L$ is the number of iterations.

Regarding the comparison against approximations like sliced-Wasserstein: these methods are indeed powerful alternatives for the general OT problem, which lacks a closed-form solution and requires iteration. Our approach (Semi-Relaxed OT) achieves efficient, non-iterative computation. We agree that exploring such approximations is an interesting future direction, particularly when extending our framework to fully constrained transport scenarios or other variants like Gromov OT. We will include this point in our discussion.

Q4: Performance on real-world image classification shifts.

We believe our method is unique in two aspects: 1, its label-agnostic sample-level semantic calibration (13); 2, the filtering mechanism based on the coupling introduced by (12) (13).

An intuitive example is the presence of noisy labels: when a portion of labels are flipped, sample-level calibration (12) can still reliably capture semantics. This is because the process is unsupervised, similar to contrastive learning [4], and thus remains robust to label noise. Building on this, the coupling-based filtering mechanism can discards samples with wrong label, enabling the cross-entropy to train on a cleaner subset of the data. In contrast, methods that depend on label supervision, including ERM with augmentation, face an inherent "garbage in, garbage out" issue. No augmentation can correct mislabeled examples, and these label errors inevitably propagate through the learning process, compromising model robustness.

A similar advantage is observed in Colored MNIST. Even when augmentation occasionally alters the spurious attribute (e.g., color), standard ERM still learns entangled representations, as the label remains strongly correlated with color. In contrast, our label-agnostic calibration encourages the model to identify truly invariant features—those that remain stable under augmentation—naturally prioritizing shape over color in the absence of a strong supervisory signal.

Q6: ... overhead of MLMC-RT compared to a single-level Monte Carlo?

To provide a complete comparison regarding the per-iteration time overhead, we compare MLMC-RT against Single-Level Monte Carlo (SLMC) in two scenarios:

1, When we fix the number of inner samples for both SLMC and MLMC-RT, their average per-iteration cost is similar. However, SLMC suffers from significant and uncontrollable estimation bias due to its fixed and limited sampling. In contrast, MLMC-RT maintains low cost by prioritizing low-complexity levels (with sampling probability roughly proportional to $2^{-l}$, making its average cost comparable to that of a naive SLMC), but still achieves a much lower bias thanks to its multi-level correction. This leads to substantially more accurate estimates without increasing average computational cost.

2, To achieve the same low estimation bias as MLMC-RT (for a target accuracy $\epsilon$), SLMC must use a large number of inner samples, on the order of $N = O(1/\epsilon)$, making its per-iteration cost prohibitively high. In contrast, MLMC-RT achieves the same bias accuracy with far fewer samples, requiring only $N = O(\log^2(1/\epsilon))$, and is therefore exponentially more efficient under the same bias target.

In summary, MLMC-RT offers the best of both worlds: its average per-iteration cost is comparable to that of a cheap but biased SLMC estimator, while its accuracy rivals that of an expensive high-fidelity SLMC. We will add a paragraph to the appendix elaborating on this trade-off, and include the wall-clock time comparison (as noted in W1) to provide concrete empirical support.

[1] Caron, Mathilde, et al. Emerging properties in self-supervised vision transformers. ICCV'21

Thank you again for your thoughtful comments. If you have any further comments, we will gladly make every effort to address them.

We thank the reviewer for the thoughtful comments, and address your concerns below.

Response to Weaknesses

W1: MLMC-RT gradient estimation, while theoretically efficient, adds significant implementation complexity. Multiple OT problems need to be solved at each iteration, potentially creating computational bottlenecks.

We would like to clarify two points regarding computational complexity:

MLMC-RT: MLMC-RT is introduced not as a source of algorithmic complexity, but as a tool to make gradient estimation for our nested objective tractable and efficient. A naive Single-Level Monte Carlo (SLMC) estimator suffers from a fundamental trade-off: achieving low bias requires a large number of samples in the inner loop—on the order of $O(1/\epsilon)$ per iteration (line 913). In contrast, MLMC-RT achieves the same estimation accuracy with significantly lower total complexity $O(\log 1/\epsilon)$ (line 948), making it exponentially more sample-efficient for any given target accuracy. Even when both MLMC-RT and SLMC use the same number of samples in the inner loop, their average per-iteration costs are comparable.

OT: While entropy-regularized OT is often solved with the Sinkhorn's algorithm, our formulation allows for a closed-form for the transport plan (detailed in Appendix A.4). This is because we adopt a semi-relaxed OT formulation that only enforces the source marginal constraint ($\gamma \in \Pi(P)$). This specific relaxation circumvents the need for iteratively balancing two marginals and instead admits a direct analytical solution. In essence, we only need to compute the cost matrix $C_\theta$ and then perform a few element-wise operations (like exp and log-sum-exp), which are highly optimized on modern hardware like GPUs. This means we bypass the need for any iterative OT solver, and the additional computational overhead per step is minimal.

W2: Assumptions 3-4 regarding boundary regularity and non-degenerate filtering may not hold in practice. Lipschitz continuity requirements for kernels and encoders could limit applicability.

1, Assumption 3 (Boundary Regularity) is a standard smoothness condition, similar to the Tsybakov margin condition in statistical learning (line 714). It helps prevent pathological cases and supports the stability of the filtering mechanism. This assumption is both common and reasonable for well-behaved data distributions, especially when data augmentation techniques are employed.
2, Assumption 4 (Non-degenerate Filtering) can be ensured through appropriate algorithmic design. In our case, the adaptive thresholding mechanism (lines 717–719, 264–267) can prevent excessive sample rejection and maintain sufficient support in the filtered distribution, thereby avoiding potential algorithmic instability. 3, The Lipschitz continuity requirements for encoders and kernels are also common and practical. For instance, we analyzed in Appendix C.1 that this property holds for common choices like the cosine similarity used in our experiments (line 321, 961).

W3: Multiple hyperparameters ($\lambda$, $\epsilon$, filtering thresholds) require careful tuning. High implementation barrier compared to simpler baseline methods.

Response to Questions

Q1: How does the OT computation scale with high-dimensional data and large sample sizes?

Please see our response to W1 for the concerns on scalibility.

Q2: Can you provide more detailed hyperparameter sensitivity analysis across different datasets?

The table below corresponds to the Colored MNIST benchmark, where we evaluate how various domain generalization methods perform under strong spurious correlations. To evaluate sensitivity to IOT calibration, we replace standard augmentations with multi-crop [1], which provides semantically invariant samples at multiple scales for ground-truth pairing.

The following table presents results on CIFAR-100 with 50% label noise, where half of the training labels are randomly corrupted. This setting evaluates the robustness of different methods under severe label noise. We report accuracy across different values of the regularization parameter $\lambda$.

Due to time constraints, we will include additional content in the revision.

Q3: What is the actual wall-clock time comparison with baseline methods?

We provide a detailed wall-clock time comparison in Section 2.4 of our supplementary material, with key results summarized in Tables 4 and 5 therein. The empirical results demonstrate that our method, AdaDRO, is highly competitive and incurs only a minimal computational overhead compared to standard baselines, while being significantly more efficient than some DRO methods. Specifically, on the Colored MNIST and Waterbirds datasets, the training time of AdaDRO is approximately 1.3x that of standard methods like GroupDRO and IRM. This is much faster than other robust optimization methods like WDRO and UniDRO, which are 5-6 times slower. Similarly, on the CIFAR-10-LT and CIFAR-100-LT benchmarks, AdaDRO's runtime is comparable with the efficient baselines like Decouple and Focal Loss, and SinkhornDRO.

[1] Caron, Mathilde, et al. Emerging properties in self-supervised vision transformers. ICCV'21

Thank you again for your thoughtful comments. If you have any further comments, we will gladly make every effort to address them.

We thank the reviewer for the thoughtful comments, and address your concerns below.

Response to Weaknesses

W1: ... no large-scale experiments (e.g., experiments on the ImageNet dataset) to confirm this.

We will add more experiments. Regarding scalability, as shown in the Sec 2.4 of Supplement, AdaDRO's runtime is only 1.3x that of standard baselines, while robust methods like WDRO are 5-6x slower. This efficiency stems from our use of a closed-form solution to the optimal transport (OT) problem, eliminating the need for costly iterative solvers.

ImageNet-LT

W2: ... the limitations of the proposed method.

We will include a discussion of the limitations in our revised version. A key aspect we will elaborate on is the construction of the ``ground-truth'' coupling, $\bar{\gamma}$, which is central to our semantic calibration via Inverse OT.

Our approach use a practical strategy for constructing $\bar{\gamma}$ by leveraging data augmentation, which implicitly defines a semantic correspondence between a sample and its augmented variant. However, we acknowledge that this represents a relatively simple means of encoding semantic priors. A central challenge lies in how to optimally incorporate all available prior knowledge to define a task-specific observed coupling. In some domains, richer priors may be available—such as explicit class hierarchies, textual descriptions, or auxiliary information from other modalities. Effectively integrating such diverse sources to construct a more powerful coupling remains a non-trivial and open research question.

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

The encoder $g_\theta$ is not an additional module but refers to the standard feature extractor backbone of the neural network used in our experiments. Correspondingly, the function $h_W$ represents the final linear classification layer (or "head") that maps the feature representation from $g_\theta$ to the logits. For instance, in experiments using a ResNet-32 backbone, $g_\theta$ includes all layers except the final linear classifier, which is represented by $h_W$.

This standard architectural decomposition ensures that our method does not introduce any additional parameters or architectural modifications beyond the baseline models, ensuring a fair comparison. We will include these implementation details explicitly in the revised manuscript’s experimental setup section to improve clarity and reproducibility.

Response to Questions

Q1: In the paper, the initial reference distribution is constructed using data augmentation. The reviewer wonders if data augmentation is also used in the compared methods.

To ensure a fair and rigorous comparison, all baseline methods were trained using the data augmentation protocol (e.g., random crops, flips) as our proposed method. However, different methods may prefer different augmentation strengths, and we report the best performance achieved for each. We will add a clarifying sentence to the experimental setup section to make this explicit and avoid ambiguity.

Thank you again for your thoughtful comments. If you have any further comments, we will gladly make every effort to address them.

We thank the reviewer for the thoughtful comments, and address your concerns below.

Q1 - Q6:

We sincerely thank the reviewer for their detailed and constructive feedback on the manuscript's presentation. We agree that the clarity and structure can be improved. In particular, we will enhance Figure 1 by including additional visualizations to better illustrate how our semantic calibration and adaptive filtering contribute to the formation of more meaningful uncertainty sets. We will also revise the text to explain Remark 1 more explicitly. Furthermore, we will incorporate additional details as suggested.

Q3: "This choice is both practical and theoretically sound, as ...

$H(\gamma | \mu \otimes \nu) = E_{(p,q)\sim\gamma} [ \log(\frac{d\gamma(p,q)}{dP_{tr}(p)d\nu(q)}) + \log(\frac{dP_{tr}(p)}{d\mu(p)}) ] = E_{\gamma} [ \log(\frac{d\gamma(p,q)}{dP_{tr}(p)d\nu(q)}) ] + E_{p\sim P_{tr}}[\log(\frac{dP_{tr}(p)}{d\mu(p)}) ]$. Thus any choice of $\mu$ satisfies $P_{tr} \ll \mu$ yields an equivalent entropy term up to an additive constant, and does not affect the optimization.

Q4: ... the 'classification as inverse OT' part... the OT problem changes from matching $p(x)$ and $q(x)$ to $p(x)$ and $y$.

Our contribution in Lemma 3.1 is to provide a key reinterpretation for DRO formulated in (1): we show that the worst-case distribution (derived in [1]) can be understood as the target marginal of a specific semi-relaxed OT problem. This perspective forms the foundation for our subsequent components—semantic calibration via inverse OT (IOT) and adaptive filtering—both of which exploit the geometric structure of OT. Regarding the discussion of ‘classification as inverse OT’, we adapt the IOT framework between the source distribution of training data $p(x)$, and the target distribution over labels $p(y)$. This is a necessary step for modeling classification as a transport-based matching.

Q7: ... the benefit of emphasizing 'semantic calibration' with 'inverse OT' ...

As described in lines 102 and 209, the transport cost between two samples is decomposed into two parts: a sample-level cost and a label-level cost, $C(p_i, p_j) = C^X(x_i, x_j) + C^Y(y_i, y_j)$. Our semantic calibration framework explicitly addresses both components using the IOT objectives (12) and (13), respectively. We reinterpret the cross-entropy loss through the lens of IOT and show that it coincides with our objective in (13). Importantly, this perspective does not complicate the formulation, and provides a deeper understanding of its function within our IOT framework.

In contrast, standard ERM does not provide well-founded mechanism for learning the sample-level cost $C^X(x_i, x_j)$, and thus cannot guarantee that the learned feature space reflects meaningful semantic geometry—even with data augmentation. To address this, we introduce the IOT objective (12), which explicitly learns $C^X$ by recovering latent semantic distances from observed matching pairs. This approach helps to extract $C^X$, enabling the semantic alignment that cannot be achieved by ERM alone.

Q8: ... preceding discussion for motivating (14) is not clear for me.

We will revise this section to improve clarity. The high-level idea is as follows: standard cross-entropy training is equivalent to an instance of IOT, where the training objective effectively calibrates the distance defined in (8) between the input embedding $g_\theta(x_i)$ and its corresponding class weight vector $w_{y_i}$. This process naturally yields a set of weight vectors $\{w_k\}_{k=1}^K$ that carry semantic meaning; classes that are semantically similar will have similar weight vectors. Within the IOT framework, these weights induce a transport cost over the label space $C^Y$. We will present this idea more directly in the revision.

Q9: ... I have a hard time understanding the reasoning...

The key idea of the filtering is to use the model's own confidence (i.e., the coupling $\gamma_\theta$) as a self-adaptive signal. If the model has difficult to align a sample $p$ with its ground-truth match $q$ (i.e., $\gamma_\theta(p,q)$ is low), it suggests that $p$ may be an outlier or noisy sample that is harming the uncertainty set. Adaptive Threshold $\tau(q)$: The adaptive threshold is designed to balance global and local confidence. The term $\tau_1$ is the global average confidence, setting a baseline. $\tau_2(q)$ measures the "matchability" of a target $q$. The ratio "$\tau_2(q) / \max(\tau_2)$" adjusts the global threshold locally. For easy targets (high $\tau_2$), the threshold is higher; for hard targets (low $\tau_2$), the threshold is lower to avoid incorrectly filtering. The popular FreeMatch [2] used a similar self-adaptive thresholding, and was later also reinterpreted through the lens of OT [3].

Q10: ... mention similar results... Is Thm 3.2 an extension of such results ...?

Thank you for prompting us to highlight this point. In addition to handling the dynamically evolving reference distribution $\hat{\nu}$, our analysis addresses a more general non-convex setting. As far as we know, the most relevant prior result is provided in [1]. While [1] focuses on convex settings, our analysis extends to non-convex regimes with dynamically evolving uncertainty sets, using different techniques.

Q11: ... why Algorithm 1 needs step 1 ...

Step 1 acts as a warm-up phase to ensure stable and reliable signals for use in (12) and (13).

Experiments

Q1 Q2: ... 'ERM with augmentation' baselines ...

The refereed paper show that ERM+augmentation is a strong baseline. We believe our method is unique in two aspects: 1, its label-agnostic sample-level semantic calibration (13); 2, the filtering mechanism based on the coupling introduced by (12) (13).

An intuitive example is the presence of noisy labels: when a portion of labels are flipped, sample-level calibration (12) can still reliably capture semantics. This is because the process is unsupervised, similar to contrastive learning [4], and thus remains robust to label noise. Building on this, the coupling-based filtering mechanism can discards samples with wrong label, enabling the cross-entropy to train on a cleaner subset of the data. In contrast, methods that depend on label supervision, including ERM with augmentation, face an inherent "garbage in, garbage out" issue. No augmentation can correct mislabeled examples, and these label errors inevitably propagate through the learning process, compromising model robustness.

A similar advantage is observed in Colored MNIST. Even when augmentation occasionally alters the spurious attribute (e.g., color), standard ERM still learns entangled representations, as the label remains strongly correlated with color. In contrast, our label-agnostic calibration encourages the model to identify truly invariant features—those that remain stable under augmentation—naturally prioritizing shape over color in the absence of a strong supervisory signal.

For the space limit, we only present results on the DomainBed benchmark.

Q3-Q13: We thank the reviewer for sufficient suggestions. We will clarify these points in the revised experimental section: (Q3) GIW was mentioned for context and we will make it clear that it was not used as a baseline; (Q4) For Waterbirds, the spurious correlation is between the background (water vs. land) and the label (waterbird vs. landbird). The 'worst-group' distribution contains samples where this correlation is broken (e.g., waterbirds on land); (Q6) The uncertainty radius $\delta$ is a hyperparameter, which is controlled by the dual variable $\lambda$ in practice. We will add its search space and the selected value for each experiment to the appendix; (Q12) Table 6 corresponds to the Waterbirds dataset. We will add a caption.

We agree with these excellent suggestions for improving the paper's readability and completeness. In the revision, we will: (Q5) Streamline Sections 2 and 3 to create more space for a detailed discussion of the experiments; (Q8, Q9, Q10) Integrate all supplementary materials, including ablation studies, visualizations, error bars, and full experimental details, directly into the main paper's appendix to create a single, self-contained document; (Q11) Add a discussion on the effect of imbalance factors in the CIFAR-LT experiments.

Response to Minor Comments: Thank you for your valuable comments. We have corrected the typos and made several revisions to improve the overall writing quality of the manuscript.

Regarding the comment on the inputs $*x*_i$ and labels $*k*$ residing in different spaces, our formulation does not compute a Wasserstein distance, which is defined in a metric space. Instead, our approach leverages the more general OT framework, which permits flexible ground distance function that do not necessarily satisfy the properties of a metric. As long as a ground distance function is defined, an OT distance can be computed.

[1] Sinkhorn distributionally robust optimization. 2023

[2] Freematch: Self-adaptive thresholding for semi-supervised learning. ICLR'23

[3] OTMatch: Improving Semi-Supervised Learning with Optimal Transport. ICML'24

[4] Understanding and Generalizing Contrastive Learning from the Inverse Optimal Transport Perspective. ICML'23

Dear reviewer

Please engage in the discussion with the authors at your earliest convenience.

Thanks

AC