Noise Tolerance of Distributionally Robust Learning
Authors: Ramzi Dakhmouche
Ivan Lunati
Hossein Gorji
Venue: ICLR 2026 Poster (2026)
Keywords: Distributional Robustness
Wasserstein Distance
Deep Learning
Operator Learning
Primary area: learning on time series and dynamical systems
TL;DR: We propose a computationally efficient learning approach featuring robustness to global forms of noise.

Abstract: Given the importance of building robust machine learning models, considerable efforts have recently been put into developing training strategies that achieve robustness to outliers and adversarial attacks. Yet, a major aspect that remains an open problem is systematic robustness to global forms of noise such as those that come from measurements and quantization. Hence, we propose in this work an approach to train regression models from data with additive forms of noise, leveraging the Wasserstein distance as a loss function. Importantly, our approach is agnostic to the model structure, unlike the increasingly popular Wasserstein Distributionally Robust Learning paradigm (WDRL) which, we show, does not achieve improved robustness when the regression function is not convex or Lipschitz. We provide a theoretical analysis of the scaling of the regression functions in terms of the variance of the noise, for both formulations and show consistency of the proposed loss function. Lastly, we conclude with numerical experiments on physical PDE Benchmarks and electric grid data, demonstrating competitive performance with an order of magnitude reduction in computational cost.

Published as a conference paper ICLR 2026
NOISE TOLERANCE OF DISTRIBUTIONALLY ROBUST LEARNING
Ramzi Dakhmouche
Ivan Lunati
Institute of Malhematics, EPFL, Switzerland Computational Engineering Lab, Empua Switzerland
Computational Engineering Lab, Empa, Switzerland
ivan, lunat (empa.ch
ramzi 1. Kamoucch 2epfl.ch
Hossein Gorji
Computational Engincering Lab, Empa, Switzerland mchomaadoaseain. gordiee dempa.co
ABSTRACT
Given the importance of building robust machine learning models, considerable efforts have recently been put into developing training strategies that achieve robustness to outliers and adversarial attacks. Yet, major aspect that remains an open problem is systematic robustness to global forms of noise such as those that come from measurements and quantization. Hence, we propose in this work an approach to train regression models from data with additive forms of noise, leveraging the Wasserstein distance loss function. Importantly, our approach is agnostic to the model structure, unlike the increasingly popular Wasserstein Distributionally Robust Learning paradigm (WDRL) which, show, does not achieve improved robustness when the regression function is not convex or Lipschitz. We provide theoretical analysis of the scaling of the regression functions in terms of the variance of the noise, for both formulations and show consistency of the proposed loss function. Lastly, we conclude with numerical experiments on physical PDE benchmarks and electric grid data, demonstrating competitive performance with an order of magnitude reduction in computational cost.
INTRODUCTION
In real-world applications, collected data is often tainted with different forms of noise. Whether it is sensor noise engineering systems or measurement uncertainty biological experiments, such noise usually demands costly and time-consuming pre-procee pre-processing steps, before meaningful results can be extracted us- ing predictive machine learning algorithms. In order to streamline that process, different robust learning approaches have been proposed with a focus on robustness to outliers and adversarial attacks (Mohajerin Esfa- hani and Kuhn 2018 Steinhardt et al. 2018 Bai et al.) 2023 Levine and Feizi 2020). Most of such strategies rely on augmenting the data with adversarial examples (Goodfellow et al. 2014 Madry et al. 2018) or designing suitable loss regularization techniques (Dong al. 2020). However, for more global forms of noise, which are commonly encountered in practice, these approaches face both statistical and practical limitations. In the case of data augmentation, the limitations inherent to its design, while adversarial regularization often targets bounded perturbations, thereby overlooking standard noise models that arise in real-world settings. In contrast, the increasingly popular paradigm of Wasserstein Distributionally Robust Learning (WDRL) (Mohajerin Esfahani and Kuhn 20188] SShaiiezaadhh-Aadee et al.) 2209 Gao et al.) (024) represents more general framework that allows for arbitrary perturbations, and is more theoretically appealing while

---

Published as a conference paper ICLR 2026
leading competitive performance. Yet, there seems be a gap in the literature when comes robustness properties of WDRL with respect to global forms of noise, as pointed out by Hu et al. (2020) for instance. In this work, we address this question in a regression setting from multiple perspectives:
1. We study the global robustness properties of the popular WDRL formulation, through theoretical and numerical analysis its scaling terms of the variance of the noise.
2. Notably, we show that WDRL may fail improve the performance when the regression functions are neither Lipschitz nor convex.
3. To address this limitation, we propose a simple yet powerful robust learning approach that is agnostic to the structure of regression functions, enabling more expressive models. We further provide a theoretical analysis of its dependence on the variance of the noise.
4. We numerically demonstrate the performance our proposed approach through various physical problem benchmarks and electric grid usage time series data.
The rest of the paper is organized as follows. In Section we review the regression setting. We explore the shortcomings of WDRL, which has not been investigated before, in Section 3] Section inttooduces the novel regression approach, where we establish its main theoretical properties in Section We justify the improved scaling of the introduced regression method with respect to the empirical risk minimization and WDRL in Proposition 5.1 Section i) devoted to the numerical results, where the competitive performance of our method is demonstrated. Finally, Section discusses the limitations, outlook and concluding remarks.
1.1 RELATED WORKS
Denoising and Filtering. Extensive research has been conducted on denoising and filtering techniques, ranging from Kalman filtering (Ito and Xiong, 2000) and wavelet denoising (Sardy et al. 200T) to deep learning based methods (Jain and Scung| 2008 Xing and Egiazarian, 2021; Krull et al. 2019] Lehtinen et al. 2018). For a comprehensive overview in the context of image data modalities, see Elad ct al. (2023). However, most of these approaches require low noise data, focus Gaussian noise distributions require an explicit noise model. Additionally, they introduce costly pre-processing steps that must be performed prior the modeling. In contrast, we propose an approach that directly trains competitive models from noisy data, eliminating the need for extensive pre-processing.
Adversarial Defense. Early works introduced techniques to augment the training data with adver- sarial examples (Goodfellow al. 2014: Madry et al. 2018), leveraging the expressive power of neural networks to improve robustness. Building on this, several regularization techniques such as entropic regularization (Dong et al., 2020) and adversarial weight perturbation (Wu et al.) (020) have been proposed, further enhancing their performance. In parallel, certified robusiness approaches have focused quantifying the proportion of samples that remain robust to arbitrary perturbations within given bound (Tjeng et al. 2019; Raghunathan al. 2018; Dathathri et al., 2020). However, these techniques often lead to overly conservative models, which can degrade performance in the presence of global noise perturbations (Bai et al.
2023
Distributionally Robust Optimization. It is concerned with minimizing the worst-case loss over a given set of distributions (Mohajerin Esfahani and Kuhn 2018 Föllmer and Weber 2015 Blanchet and Murthy 2019). which is formally expressed as the minimax problem
inf sup Eg [2)(
#E QEP
where the supremum taken over a suitably chosen class of distributions Q E P. Recent focus (Shafeezaade Abadeh et al. 2015 Staib and Jegelka| 2017 Shafieezadeh-Abadel et al. 2019 Chen and
2

---

Published as a conference paper ICLR 2026
Paschalidis 2018 Bartl al, 2021, Gao et 2024 Phan al. 2023) has been given the formulation with Wasserstein ambiguity set P PP) which is the ball centered at the empirical distribution P with radius under the Wasserstein distance, leading to WDRL. See (Gao and Kleywegt 2023) for a discussion on theoretical advantages of this choice. WDRL has demonstrated remarkable performance in out-of-sample linear regression (Mohajerin Esfahani and Kuhn 2018 and classification (Shafieezadeh- Abadeh et al. 2019) tasks, well as in defending against adversarial attacks (Bai et al. 2023 [Bui et al. 2022) neural networks. For classification problems, another highly effective, although computationally costly, approach was proposed by Zhai et al. (202I). In contrast, consider robustness to unbounded forms of noise encountered in regression, which to the best of our knowledge, has not been much investigated for deep learning models (Liu et al.
2024
2 PROBLEM SETTING BACKGROUND
We denote by (2( the set of probability measures over Ω with finite variance. The Euclidean 2-norm is given by ll2- Consider the realizations (X,)isn from measure in P2((2). we employ ml(X;))soul 1/n Ei= 6x,, the shorthand for empirical measure. Consider the regression task of predicting response variables y/ from input features E X. Given class of regression functions {Se, 0 E Rd) and data samples (X;, Jig from an underlying distribution ((X,Y) E P2 Y) with ʃo + Y, the standard goal to find a model E IRd that minimizes the empirical risk
1
ASE E arg min
(x
&ER
N2
Wasserstein Distributionally Robust Learning Ben- Tal et 2013 Bartl et al. 202I). the aim instead is
minimize a stronger form of the empirical risk
ÖWNRT. E arg min sup (X Walni)
E (X,(()),,66)),,
(1)
BER
where á E (2 x y) is the empirical distribution over the training data, & pre-chosen loss function, and W2 the 2-Wasserstein distance given by
Wa(u, (t)
min rEll(a,)) (xxx))2 ||a -B1 drr (o.
with 11(, (i) the set of couplings between the probability measures and á. and a, B E P2(x x y).
this work, we focus on the setup where the response samples (Y; + oEi)ion are tainted with in- dependent and identically distributed noise with variance o², with the objective of training deep learning models that are the least sensitive to the noise level a.
3 DRAWBACKS OF WDRL
In order to compute the objective function of WDRL, certain structural assumptions on the regression functions fe. R2d as well as the loss function must be imposed. This is necessary since the original formulation ① involves solving an infinite dimensional optimization problem, which is generally intractable. For that matter, two main settings have been proposed:
Assuming that the function lg (x.y) + (( is finite maximum of concave functions;
3

---

Published as a conference paper ICLR 2026
(b) Assuming that the function Eo (z,y) + (x)) Lipschitz continuous.
In either cases, the objective function given in ① can be rewritten (Mohajerin Esfahani and Kuhn 2018 Shafieezadeh Abadeh et al. 2015) under the tractable form []
d2 d2(()))ssn (fo (fo().))sn) ==
sup (X,Y)~ Wa(
[xle(Y,
inf X20
Ad
{e(E) - fo(Ee)) -1 (a)
sup
Eill X|XX; - =2112)
TI
where the optimal solutions (0) and (5)))) are reached for all A E R Yet, to satisfy (a) or (b) a regression setting where the data distributions have unbounded domains, one typically needs to set 2(y, ly æ/| and enforce structural properties of convexity or Lipschitzness models, therefore, reducing their expressive power. A natural question that emerges is whether using the tractable expression of d2 as loss function, regardless of whether the equality holds, can improve the robustness of the neural network models. We provide negative answer to this question by exploring the behavior of d2 in training a convolutional neural operator (CNO) (Raonic et al. 2024) to solve the two- dimensional Slvier-Stokes equation. In particular, we estimate the operator that maps the initial condition
the neural network
(T = 0), represented as an image, to the final state (T 1). To this end, we train the WDRL regression model employing stochastic descent- ascent algorithm, exploring the model behavior as the noise level increases. We use the hyperparameters optimized by the authors who proposed the CNO architecture (Raonic et al. 2023). We obtain the results shown in Figs. T and 2 for both Gaussian and heavy-tail noise distributions, respectively. For the latter case, we use the standard Cauchy distribution, where a represents the scale parameter, Cauchy random variable does not have finite variance due to the heavy tails. We examine model performance via the mean absolute relative error (MAE). Our results indicate that under heavy-tail noise, WDRL training performs significantly worse than the standard MSE training. In the Gaussian noise setting, both lead to comparable results, without noticeable improvement from WDRL. This is in contrast with the novel regression approach introduced the follow-up section, whose performance on this setting is demonstrated in Section O Note that this limitation of WDRL has not been raised so far, to the best of our knowledge, mainly because previous works focused on image classification where the data domains are bounded, honouring the Lipschitz property.
4 WASSERSTEIN BATCH MATCHING
The key idea behind our approach is to relax the strict matching between features X; and their responses Y; + oe, for i E Lp. where I, denotes the index set a training batch. The motivation for this relaxation is that, in the presence of noise, the observed response Y, + σe; already deviates from the true response Y, Consequently, the batch elements are close enough, allowing features to match responses without a fixed correspondence yields more robust estimates while reducing the sensitivity of the loss function to the noise. We provide in the following a consistency result as well as an asymptotic analysis of the XBM loss function.
4.1 FORMULATION & CONSISTENCY
The formal way to implement this idea is to compute the Wasserstein distance between the empirical distributions of the predictions (X)())) and the responses (Yi)se,, leading to Wasserstein Batch
Blanchet and Murthy 2019) propose more general condition for the equality to hold, but leave the question of existence of optimizers, which is essential here, open.
4

---

Published as conference paper ICLR 2026
VDAL Regression Regression
Negression Regiession
175
250
Polabove ewar (50
125
Relative enter (N0
300
52
25
0.0
0.5
1.0
2.0
2.5
3.0
10
Nosse
Figure I: Test error evolution for Navier-Stokes op- erator learning with 30% corrupted training data with Cauchy noise
Figure 2: Test error evolution for Navier-Stokes op- crator learning with 30% corrupted training data with Gaussian noise
Matching (WBM) regression
ÓWRM E arg min BERd
instead of the Mean Squared Error (MSE) regression. We illustrate the idea in figure For our setting of empirical distributions, note that the Wasserstein distance reduces to
Wa(m m m[((
min ( (P,, M(XX))e,, PEC
where
i.jE,, i.je,, is the matrix of the pairwise norms between the predictions and the target values, and C the set of coupling matrices of di- mension #Ip. As a sanity check, we show in proposition (1) below that, asymptotically such matching scheme recovers any continuously differentiable bandlimited function, among its co-monotonic functions from its samples, in the noise-free regime. The proof is reported in Appendix (A).
fv g
Proposition 4.1. (Consistency)
= arg min fog
dcgd
Let f :R - be continuously differentiable and inte- grable function with compactly supported Fourier transform and let (f(* be its values sampled at ordered points (xa ( where ó is an unknown permutation preserving the batch partition. Then, given fixed batch size and an arbitrary amount of samples, f is completely characterized by minimiz- ing
Figure 3: Wasserstein Batch Matching il- lustration. The regression through cloud of points
batch (depicted by blue) tackled by finding optimal map, depicted by black line, between distributions of (X:); and (Y, + ae);, shown by red and green histograms.
min 9EG
P21
where (Tp. 1} is the finite collection of batch index and G the set of continuously differentiable and integrable bandlimited functions that are co-monotonic with
5

---

Published as conference paper ICLR 2026
Remark. (Complexity) From computational complexity perspective, training with WBM involves solving a linear program at each training step, which costs (ss), where s dim()) is the dimension of the response space V. However, this is independent of the strucutre of the regression functions. On the other hand, WDRL involves solving minimax problem which is in O(*²) when the function Eg is convex-concave. However, in the absence of this structure, the problem can become arbitrarily hard.
Remark. (Differentiability) The proposed WBM loss is differentiable with respect to the regression parameters Fdd by the envelope theorem (Bonnans and Shapiro 2013), which makes it suitable for training deep learning models.
5 NOISE SCALING ANALYSIS
To investigate the theoretical advantage of the WEM regression, we examine the robustness of the learned function with respect to the noise variance. First, provide sensitivity analysis of the scaling of the WDRL- and WBM- loss functions in terms of the variance of the noise, and contrast their scaling with that of the MSE loss. Next, investigate and discuss the consequences of this scaling on the learned regression functions.
5.1 NOISE SCALING OF Loss FUNCTIONS
An important feature of the introduced WBM loss is its scaling with respect to the noise which is derived in the following. The proof is reported in Appendix (B).
Proposition 5.1. (Noise Scaling) Assume the response variables (Y:)isn to be normalized and let o² denote the variance of the output noise. Then, for a E (0, 1), we have for the WBM loss
We (ml(X, + m((e)))),,((((((())))))).))
2 We (m[(X))ess))
=
Σ [0) - fo(,,) (V, fo(x))) Piyos, +0
JE
where (Pas)ag is a constant coupling matrix. For the WDRL loss, we have
1
d2 ((Y + ocs))es, (f((X)))ex)) - d2 (Yi)ee,, (fo(X)))e))) a.
Σ |el o(o).
where assumed the regression functtioo Se to be Lipschitz and continuously differentiable.
Remark. To put the previous bounds into perspective, consider the noise scaling of the MSE loss
dass ((Y, + dεi)tEt, (fo(X))) dass ((Y)) (fo(X)))ser)) = Σ (Y, - (xx)
2x
#Ip iElp
where we took y IR, for simplicity of presentation and dwse is the MSE loss, given by
duss ((Y:)sE (Xi)ier,) #Tp iE
Σ - -
Comparing the scaling of WBM loss with the one resulting from MSE, we observe that for small deviations between (Y))ee,, and (fo(x)))se)p) the prefactor of the linear term is smaller the former. This reduction
without loss of generality
6

---

Published as a conference paper ICLR 2026
arises because the WEM loss penalizes the prefactor through a cubic term when the deviation less than unity. This is not surprising, since WPM performs regression over the infimum of all possible couplings between the regression function and the response, leading to diminished dependency of the WEM loss compared to MSE. This theoretical insight is further confirmed by numerical experiments demonstrated in Section 6 In the following, we clarify how the derived scalings can have implications on the learned function so for 0 E IRd.
5.2 NOISE SCALING OF REGRESSION FUNCTIONS
To estimate the effect of the scaling of the loss on the learned parameters defining the estimated regression functions, we consider the invariant probability measures of the Markov chains defined by the Stochastic Gradient Descent (SGD) iterates, given in the MSE case by
(6X-) -x-
1 (x. (Y #I
0 E0.- #I, EEIP
0x+1 = Og VVa
Vo
This is motivated by the fact that SGD is standard way to solve the regression problem, when training deep learning models. Furthermore, this perspective has the advantage of taking into account other sources of randomness, encountered in model training. Yet, to keep the presentation concise, following constant step-size stochastic approximation literature (Diculeveut et al. 2020 Park et al. 2022 Vlatakis Vlatakis- Gkaragk-unka
al. 2024). we assume in the following the loss functions be strongly convex. This assumption can be
relaxed when considering decreasing step-sizes algorithms (Atchadé et al. 2017 Ben Arous et al.. 2022 Arous et al., 202T). this instance, this assumption translates to considering regularized versions of the loss functions with reproducing Hilbert kernel space hypotheses classes (Wainwright) (201)). The following results address regularized MSE and WEM, and are based on applying results by (Diculeveut al. 2020). We
postpone their proofs to Appendix (D) and EE respectively.
Corollary 5.2. Assume the MSE and WBM loss functions (denoted generically as le) to be strongly convex. Then, under suitable regularity assumptions (detailed in Appendix D). the averaged optimization iterates
0
0
k
converge to a unique stationary distribution Ty with expectation On = -ga a ( such that 0, - (V2 (V) A(0") V(O*) ) O11²)
where 0* is the optimum of the noise free respective problem, A (V @1+18 -1 and V(0*)= = [(V.lo
Remark. We see the effect of the first order coefficient of the loss functions in terms of the noise level V@8)). For the WDRL loss, the gradient iterates have a non-centered bias term making convergence not achievable in general. We derive bounds on the resulting learnt parameters in corollary E. (see Appendix E.
6 NUMERICAL RESULTS
We demonstrate the performance of WBM regression on two important practical problems: operator learning for PDEs and electric grid usage forecasting. Evaluation on test data is carried out using the mean absolute error (MAE) in all displayed figures, where compare models trained with the MSE those trained with WBM. We focus on this setting, since the ofate-of-thh--ar method for practical noise regimes WDRL was shown in section 3 to underperform ERM. For completeness though, report comparison of WBM divergence-based
7

---

Published as a conference paper ICLR 2026
methods CVaR DRC and Chi-Sq DRO (Duchi and Namkoong 2021). in figure As for GCDRO LLu et al. 2024), it is based on KNN graph construction, which is known to perform poorly on high-dimensional data RRadvaanovi et al,, 2009), which we focus on in this work. We explore the robustness properties of WBM both to standard Gaussian and heavy-tail Cauchy noise. Heavy-tail noise is present in many real-world applications such as vibration sensors for intelligent monitoring, power consumption sensors, and LIDAR systems. It comes from transient events, sudden extreme changes such as short circuits, atmospheric noise which exhibits heavy-tails. Additionally, we explore robustness to distribution shift properties, by training on (almost) noise-free data and testing noisy data. Such a use-case is encountered in practice, when a model is developped based on cleaned data before being deployed on real data. The results are averages over 13 samples. Code available code
6.1 LEARNING PDES
PDEs model a wide range of physical and engincering problems and feature a rich set of dynamical processes that illustrate the performance of machine learning models on a wide range of practical regression problems. For that matter, we demonstrate the performance WBM on two extensively used PDEs: the wave equation and the Navier-Stokes equation. More precisely, we consider the corresponding recently proposed benchmarks Raonic et al. (2024), where the task consists of learning operators mapping initial conditions (T' = 0), represented as images, the final state reached by the system, e.g., corresponding to (T' 1). The underlying images represent two-variable functions sampled at a given resolution. We train convolutional neural operators (Raonic et al., 2023), which have been proposed as featuring robustness properties notably to change in resolution. We set the hyperparameters optimized in (Raonic et al.) 2023), ,
and keep the same for WBM
WSN Regression
Regression
10
WBM Regression
MSE Hegression
100
Noan shoulato cIarr
Falohwe enon 156)
25
20
5
level
Noise
Figure 4: Test error evolution for Navier-Stokes op- erator learning; 30% corrupted test data with Cauchy noise
Figure 5: Test error evolution for wave equation oper- learning; 30% corrupted training data with Gaus- sian noise
training. except the batch size for which we explore different values. The convolutional neural operator architecture is based on mapping the sampled images back to function space using the Whitaker--Sannon interpolation formula (Raonic et al., 2023). We display the results Figs. 4 5 and 7 We note that, WBM regression consistently outperforms MSE regression. In particular, while both MSE and WDRL regressions indicate significant errors in Navier-Stoke
operator learning subject to the Cauchy noise, as shown in Fig. 6 the introduced WBM regression demonstrates notable robustness. Furthermore, regression outperforms divergence- based DRO methods, as shown in Fig. 7 while being significantly cheaper computationally. Indeed, note that WBM does require additional hyper-paramete tuning leading to least
note that WBM
8

---

Published as conference paper ICLR 2026
10-fold computational training gain compared CVaR and Chi-Sq DRO. We report additional numerical results and error bars in Appendix F
6.2 ELECTRIC GRID USAGE FORECASTING
Predicting electric load is an important and timely problem, especially given the increasing share of renewable energy sources in grid networks. We employ the recently proposed state-of-the-ar model TSMi xer (Chen et a to forecast electric transformer usage from the popular ETDataset (Zhou et al., 2021 Ilbert et al. 2024). TsMixer is based on mixing operations via stacking multi-layer perceptions. We train the model with input sequence length of 336 and prediction sequence length of 96. We utilize the hyperparameters proposed by the authors, except the number of training epochs which we reduce data. This is justified by the fact that we compare the model against itself trained with different loss functions. Hence, the comparison point can be chosen in a flexible way. We display the results in Figs. aandd where WBM outperforms MSE, especially in the case of Cauchy noise, suggesting its performance extends across high-dimensional learning problems.
single swap over the
WSH Regression
175
MGF Pagression
Regression
150
Mban abseluto enerr
Shean absokute entor
13-
125
100
Regressiion
Recression
CVait DRO Regression
0.0
0.5
1.0
1.5
2.0
2.5
30
10
Meiso leval
Noise level
Figure 6: Test error evolution for Navier-Stoker oper- ator learning: 30% corrupted training data with Cauchy noise
Figure Test error evolution for Navier- opera- tor learning: 30% corrupted training data with Gaussian noise
7 DISCUSSION
Limitations & Outlook. Although WBM demonstrates competitive performance, has certain limitations. First, as indicated by Proposition 4.1 WBM requires the underlying function to be sufficiently regular. As a consequence, it struggles operator learning tasks involving discontinuities, such as shock wave problems. Additionally, in the low-noise regime, WBM may exhibit slight underperformance. This can be attributed weak matching underlying WBM, which, while beneficial noisy settings, is suboptimal when perfect knowledge of the responses is available. That being said, most real-world data is noisy and there exists fundamental limit of learning algorithms when it comes to performing well both on clean and noisy data. Despite the limitations, WBM introduces novel perspective robust learning, paving the way for future research into robustness against unbounded, correlated, or structured forms of noise.
Conclusion. We considered the open problem of robustness to global forms of additive noise and proposed efficient learning approach WBM, overcoming the drawbacks of current robust learning methods.

---

Published as conference paper ICLR 2026
WEM Regression Regression
WUN Regression Negression
Moat absolute e17
Noat aboslute enter
25
3
5.0
12.5
17.5
20.0
0.0
0.5
2.0
1.5
2.0
2.5
3.0
3.5
4.0
Notce level
Figure 9: Test error evolution for electric time series forecasting; 30% corrupted test data with Gaussian noise
Figure 8: Test error evolution for electric series
forecasting: 30% corrupted test data with Cauchy noise
We investigated the scaling of the introduced regression along with other approaches with respect noise levels, offering a theoretical justification for the gains achieved by WBM. Furthermore, we demonstrated the practical performance of WBM via several numerical experiments involving PDE operators and electrical time series forecasting. We believe this work paves the way for robust learning methods that streamline the costly data pre-processing step, while advancing the development of reliable machine learning models.
ACKNOWLEDGMENTS
This work was supported by the Swiss National Science Foundation under grant No. 212876. We acknowledge computational resources from the Swiss National Supercomputing Centre CSCS.
REFERENCES
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. Online stochastic gradient descent on non-convex losses from high-dimensional inference. Journal of Machine Learning Research, 22(106):1-51, 2021.
Yves F Atchadé, Gersende Fort, and Eric Moulines. On perturbed proximal gradient algorithms. Journal of Machine Learning Research, 18(10):1-33, 2017.
Xingjian Bai, Guangyi He, Yifan Jiang, and Jan Obloj. Wasserstein distributional robustness of neural networks. Advances in Neural Information Processing Systems, 36, 2023.
Daniel Bartl, Samuel Drapeau, Jan Obłój, and Johannes Wiesel. Sensitivity analysis of wasserstein dis- tributionally robust optimization problems. Proceedings of the Royal Society A, 477(2256):0210117, 2021.
Gerard Ben Arous, Reza Gheissari, and Aukosh Jagannath. High-dimensional limit theorems for sgd: Effective dynamics and critical scaling. Advances in Neural Information Processing Systems, 35:25349-255 25362,, 2022.
10

---

Published as conference paper ICLR 2026
Aharon Ben-Tal, Dick Den Hertog, Anja De Waegenaere, Bertrand Melenberg, and Gijs Rennen. Robust solutions of optimization problems affected by uncertain probabilities. Management Science, 59(2): 341-357, 2013.
Jose Blanchet and Karthyek Murthy, Quantifying distributional model risk via optimal transport. Mathematics of Operations Research, 44(2):565-600, 2019.
Frédéric Bonnans and Alexander Shapiro. Perturbation analysis of optimization problems. Springer Science & Business Media, 2013.
Anh Tuan Bui, Trung Lc, Quan Hung Tran, He Zhao, and Dinh Phung. A unified wasserstein distributional robustness framework for adversarial training. In International Conference Learning Representations, 2022.
Ruidi Chen and loannis Ch Paschalidis. A robust learning approach for regression models based on distribu- tionally robust optimization. Journal of Machine Learning Research, 19(13)11-48, 2018.
Si-An Chen, Chun-Liang Li, Sercan O Arik, Nathanael Christian Yoder, and Tomas Pfister. Tsmixer: An all-mip architecture for time series forecast-ing. Transactions on Machine Learning Research.
Sumanth Dathathri, Krishnamurthy Dvijotham, Alexey Kurakin, Aditi Raghunathan, Jonathan Uesato, Rudy R Bunel, Shreya Shankar, Jacob Steinhardt, Ian Goodfellow, Percy S Liang, et al. Enabling certification of agginoc networks via memory-efficient semidefinite programming. Advances in Neural Information Processing Systems, 33:5318-555 1331, 2020.
Aymeric Dieuleveut, Alain Durmus, and Francis Bach. Bridging the gap between constant step size stochastic gradient descent and markov chains. The Annals of Statistics, 48(3):1348- 1382, 2020.
Yinpeng Dong, Zhijie Deng, Tianyu Pang. Jun Zhu, and Hang Su. Adversarial distributional training for robust deep learning. Advances in Neural Information Processing Systems, 33:8270-8288 2020.
John Duchi and Hongseok Namkoong. Learning models with uniform performance via distributionally robust optimization. The Annals of Statistics, 49(3):1378-1406 2021.
Michael Elad, Bahjat Kawar, and Gregory Vaksman. Image denoising: The deep learning revolution and beyond a survey paper. SIAM Journal on Imaging Sciences, 16(3):1994-164,, 2023.
Hans Föllmer and Stefan Weber. The axiomatic approach risk measures for capital determination. Annual Review of Financial Economics, 7(1):301-337, 2015. 2015.
Rui Gao and Anton Kleywegt. Distributionally robust stochastic optimization with wasserstein distance. Mathematics of Operations Research, 48(2):603 2023.
Rui Gao, Xi Chen, and Anton J Kleywegt. Wasserstein distributionally robust optimization and variation regularization. Operations Reseurch, 72(3):1 2024.
Matthias Gelbrich. On formula for the 12 wasserstein metric between measures on euclidean and hilbert spaces. Mathematische Nachrichten, 147(1):185-20) 1990.
lan Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. arXiv preprint arXiv: 1412:::4 2014.
Edward Hu, Adith Swaminathan, Hadi Salman, and Greg Yang. Improved image wasserstein attacks and defenses. arXiv preprint arXiv:2004. /2478, 2020.

---

Published as conference paper ICLR 2026
Romain Ilbert, Ambroise Odonnat, Vasilii Feofanov, Aladin Virmaux, Giuseppe Paolo, Themis Palpanas, and levgen Redko. Samformer: Unlocking the potential of transformers in time series forecasting with sharpness-aware minimization and channel-wise attention. In Forty-first International Conference on Machine Learning, 2024.
Kazufumi Ito and Kaiqi Xiong. Gaussian filters for nonlinear filtering problems. IEEE transactions on automatic control, (5(5):910-921) 2000.
Viren Jain and Sebastian Seung. Natural image denoising with convolutional networks. Advances in neural information processing systems, 21, 2008.
Alexander Krull, Tim-Oliver Buchholz, and Florian Jug. Noise2void-learninn denoising from single noisy images. In Proceedings of the IEEF/CVF conference on computer vision and pattern recognition, pages 2129-2137, 2019.
Jaakko Lehtinen, Jacob Munkherg, Jon Hasselgren, Samuli Laine, Tero Karras, Miik Aittala, and Timo Aila. Noise2noise: Learning image restoration without clean data. International Conference on Machine Learning, pages 2965-2974. PMLR, 2018.
Alexander Levine and Soheil Feizi. Wasserstein smoothing: Certified robustness against wasserstein adversar- ial attacks. International Conference on Artificial Intelligence and Statisties, pages 3938-3944 PMLR, 2020.
Jiashuo Liu, Jiayun Wu, Tianyu Wang. Hao Zou, Bo Li, and Peng Cui. Geometry-calibratee dro: Combating over-pessimism with free energy implications. In Forty-first International Conference Machine Learning, 2024.
Giulia Luise, Alessandro Rudi, Massimiliano Pontil, and Carlo Ciliberto. Differential properties of sinkhorn approximation for learning with wasserstein distance. Advances in Neural Information Processing Systems, 31, 2018.
Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. In International Conference on Learning Representations, 2018.
Peyman Mohajerin Esfahani and Daniel Kuhn. Data-driven distributionally robust optimization using the wasserstein metric: Performance guarantees and tractable reformulations. Mathematical Programming, 171(1):115-166. 2018.
Sejun Park, Umut Simsekli, and Murat A Erdogdu. Generalization bounds for stochastic gradient descent via localized s-covers. Advances in Neural Information Processing Systems, 35:2790-2800 2022.
Hoang Phan, Trung Le, Trung Phung, Anh Tuan Bui, Nhat Ho, and Dinh Phung. Global-local regularization via distributional robustness. In International Conference on Artificial Intelligence and Statistics, pages 7644-7664. PMLR, 2023.
Miloš Radovanović, Alexandros Nanopoulos, and Mirjana Ivanović. Nearest neighbors in nigh-dimensional data: The emergence and influence of hubs. In Proceedings of the 26th Annual International Conference on Machine Learning, pages 875-872,4 2009.
Aditi Raghunathan, Jacob Steinhardt, and Percy S Liang. Semidefinite relaxations for certifying robustness to adversarial examples. Advances in neural information processing systems, 31, 2018.
12

---

Published as a conference paper ICLR 2026
Bogdan Raonic, Roberto Molinaro, Tim De Ryck, Tobias Rohner, Francesca Bartolucci, Rima Alaifari, Siddhartha Mishra, and Emmanuel de Bézenac. Convolutional neural operators for robust and accurate learning of pdes. Advances Neural Information Processing Systems, 36, 2023.
Bogdan Raonic, Roberto Molinaro, Tim De Ryck, Tobias Rohner, Francesca Bartolucci, Rima Alaifari, Siddhartha Mishra, and Emmanuel de Bézenac. Convolutional neural operators for robust and accurate learning of pdes. Advances in Neural Information Processing Systems, 36, 2024.
Sylvain Sardy, Paul Tseng, and Andrew Bruce. Robust wavelet denoising. IEEE transactions on signal processing. 49(6):1146-11555 2001.
Soroosh Shafieezadeh Abadeh, Peyman M Mohajerin Esfahani, and Daniel Kuhn. Distributionally robust logistic regression. Advances in neural information processing systems, 28, 2015.
Soroosh Shafieezadeh-Abadeh. Daniel Kuhn, and Peyman Mohajerin Esfahani. Regularization via mass transportation. Journal of Machine Learning Research, 20(103):1-66 2019.
Richard Sinkhorn and Paul Knopp. Concerning nonnegative matrices and doubly stochastic matrices. Pacific Journal of Mathematics, 21(2):34-3348, 1967.
Matthew Staib and Stefanie Jegelka. Distributionally robust deep learning generalization of adversarial training. In NIPS workshop on Machine Learning and Computer Security, volume 3, page 4, 2017.
Jacob Steinhardt, Moses Charikar, and Gregory Valiant. Resilience: A criterion for learning in the presence of arbitrary outliers. In 9th Innovations Theoretical Computer Science Conference (ITCS 2018), volume 94, page 45. Schloss
Craittull-LiibniiZzeettttt fuer Informatik, 2018.
Vincent Tjeng, Kai Y Xiao, and Russ Tedrake. Evaluating robustness of neural networks with mixed integer programming. In International Conference on Learning Representations, 2019.
Emmanouil Vasileios Vlatakis
Vlatakis-Gkaragkouunn Gkaragkounis, Angeliki Giannou, Yudong Chen, and Qiaomin Xie. Stochastic methods in variational inequalities: Ergodicity, bias and refinements. International Conference on Artificial Intelligence and Slutistics, pages 4123-4131, PMLR, 2024.
Martin J Wainwright. Iligh-dimensional statistics: A non-asymptotic viewpoint, volume 48. Cambridge university press, 2019.
Dongxian Wu, Shu-Tao Xia, and Yisen Wang. Adversarial weight perturbation helps robust generalization. Advances in neural information processing systems, 33:2958-2969, 2020.
Wenzhu Xing and Karen Egiazarian. End-to-end learning for joint image demosaicing, denoising and super- resolution. In Proceedings of the IEEF/CVF conference computer vision and pattern recognition, pages 3507-3516, 2021.
Runtian Zhai, Chen Dan, Zico Kolter, and Pradeep Ravikumar. Doro: Distributional and outlier robust optimization. In International Conference on Machine Learning, pages 12345-12355, PMLR, 2021.
Haoyi Zhou, Shanghang Zhang, Jieqi Peng, Shuai Zhang. Jianxin Li, Hui Xiong, and Wancai Zhang. Informer: Beyond efficient transformer for long sequence time-series forecasting. In The Thirty-Fifth AAAI Conference on Artificial Intelligence, AAAI 2021, Virtual Conference, volume 35, pages 11106-11111 AAAI Press, 2021.
13

---

Published as conference paper ICLR 2026
A PROOF OF PROPOSITION 4.1: CONSISTENCY
Let's first assume that y R. Let f.g g G. Given that g is continuously differentiable, it has bounded variations on every compact, that is for all a, b IR such that a < b, we have
sup
(0(x+1)) - g(x:)) < +00,
PrEPr
where the supremum is taken over the set
ppr = { ... Unge) } pr a partition of [a, b] satisfying x; S xi+1 for S S Apr 1}}
This implies that there exists a partition of the feature space into intervals of lengths (6n) such that g monotonous on every interval. The same holds for f. Hence, we consider the partition formed by I; n Jj where ENN) and (Jz) are the chosen partitions for f and g respectively. We denote by (6n) the new lengths. Since, f is bandlimited, let the support of its Fourier transform be included in [-B, B] with B > 0. We can choose (.,), such that 6n B for all n N. Furthermore, we can sample each interval a number times equal to the prefixed batch size. Since, satisfies
min JEG
ml Wi (a
Wa(m(((((:)))e ((((()))))
P21
P21
0,
we know that minimizer g E 9 of
Wz (mll
g +
P21
must satisfy
vp 1
Wylml(((x.))iei),...ll m(.ls)))))
Furthermore, since / and g are co-monotonic, the Wasserstein matching recovers the true matching. Last, by the Shannon sampling theorem we conclude g equal f. As for generalization to the nulti-dimensional setting, the co-monotonicity property is expressed through V V 2 0. Intervals are replaced by hyperrect- angles (Cartesian products of intervals) and the matching argument carries on over hyperrectangles. Last, the bounded variation property carries on as an implication of Mean Value Inequality.
B PROOF OF PROPOSITION 5.1: NOISE SCALING
B.1 WASSERSTEIN BATCH MATCHING
For simplicity of the presentation, we take y R. The proof relies on establishing the scaling for the Sinkhorn regularized version WA of the Wasserstein distance, then concluding by taking the limit as - +00. This motivated by the fact that WA is strictly convex with respect to the coupling matrix P, hence it allows for richer characterization of the optimal coupling matrix p and its differentiability properties. First, recall that for À > 0,
1 ((P) with M(P)
a
Wa b) min (P, (aU(a.)
Σ :==1
Pas(log P.j - 1),
where M is the cost matrix, that depends on the data (X,, Y, + destign and a, b the discrete distributions under consideration. In our case, we consider uniform and b. That is, a; b, #T, for #lp. Denoting
14

---

Published as conference paper ICLR 2026
by Px the solution of the above optimization problem, we have by Sinkhorn's scaling theorem (Sinkhorn and Knopp 1967 Luise et al. 2018) the characterization
Px diag(c Aa "e Je- diag(e>0"),
where a² and 3 are the optimal Lagrange multipliers. We would like to compute the derivative of Px with respect to a to deduce its sensitivity to noise for small values of a, for egg. a E (0, But, first let us establish that + (x( is differentiable. Consider the Lagrangian
M-Myy-aa 10
a 8) = + Σ
A
By theorem in (Luise et al. 2018 it is smooth and strictly convex y- (a, 3). hence for every fixed σ >> O, there exists 7 (a) such that Z (o)) min, 7) As result, following the proof of theorem 2 in (Luise al. 2018), by the implicit function theorem, we deduce that PA continuously differentiable. Furthermore, denoting & C, we have
(a:7) = a- P21 (a-P11)
6 1
with
<(o;)((o)) =0.
By differentiation, get
VIv T
and consequently
Vor( (oa (0) = (oYY-fa(X,))) + 20rEs)
(2(Y, fo()))) + 22TE)).ssm)
Now, we can compute
dMy(o) (Pt(a)) + Σ Mij(a) d(Pi(o)) do
My(0))(X(0))s)- = Σ
do
do
#1=1
Σ 2(Y, M() 1 MMOG))GMMHMMMMMMMMMR)RRR F
2(Y M
1-=1
where we discarded the terms in (g)). Finally, since this holds for each (0, 1), and Px -(p) P(a) where ((a) is the optimal coupling in the non-regularized Wasserstein distance, which is continuous (see for e.g. (Bonnans and Shapiro 2013)), we conclude by Dini's theorem that the convergence is uniform, hence by taking the limit when À +00, of the estimate on the derivative, we get the result.
B.2 WASSERSTIEN DISTRIBUTIONALLY ROBUST LEARNING
Under the Lipschitz assumption, obtain the existence of optimizers for d2. Such optimizers (i) satisfy
dx oes)isn: (fo(X)))
15

---

Published as conference paper ICLR 2026
where d dY + dej)isn, (x) the Clarke subdifferential of d2 leveraging the differentia- bility properties of fo and the implicit function theorem, we use this inclusion to obtain the derivatives of the optimizers with respect to a. Note that the only non-smooth part comes from taking ((x,)) y) y/. Specifically, we
50 1 - 2 Y; + rs;)) 01 (.0 Vo - 22(E2 X;)] 0|
by differentiation we get
(0E E 2X -
+ 84) 0|
do fo( - da
d5 00 E
22 0 der
since 0| 1 -1= 1], and leveraging the Lagrangian to differentiate À with respect to we get
0 (ō 161 - Y11 - 11 - x,II) dd
+
21 11,11
do
Now, as we have the derivatives of the optimizers, we can differentiate
da oes)ign: (fo(X)))) X+8 Σ feist fa(E))) - A||L; -ill X||X, [5113
leading to the result.
C DISCUSSION ON ASYMPTOTIC SCALING
Further insights into the relationship between WBM and MSE regression be obtained by considering the asymptotic limit of regression problem in a batch of infinitely many samples. Specifically, we analyze the case where the feature-response pair (X, Y) X² random variables, satisfying the following conditions.
1. The laws of X and Y, ,x,y, are Gaussian measures with mean values mx,y and covariances Σx,y.
2. The covariance between X and Y is optimal, i.e., Cov(X,Y) = (E "r
1/2
3. The regression is performed over linear functions of 0 E TR.
this setting, is straight-forward to see that the two loss functions become identical. In particular, following Gelbrich bound (Gelbrich 1990), we have
||mx Omy 112 + Ey + f Ex -20r (MX Ey 22/22 1/2
Wa(py, Sogux) =
where fo fx is fosh-forward of px by fo, and therefore the MSE loss becomes
E 0
|mx - Omyy 112 Ey
f Ex 20tr Cov(X,,) Y) We(my, foHHx).
16

---

Published as conference paper ICLR 2026
D PROOF OF COROLLARY 5.2
First, let detail the assumptions. In addition to strong convexity, assume that the loss function is five times continuously differentiable with second to fifth bounded derivatives. Moreover, we assume for the covariance matrix C(0) BB((V [22, there exists s, k 2 0 such that for all 0 E TRd, have
(0) - M(1 + 10 0²11²).
max ve (1,2333)
Recall that for the MSE loss given by
duss (0' + des) (Jol (J)(E)))sE,,)
Σ
(Y: 0e; - So(x.))²
#I IEI
we have
Vaduss (OY+ os:)Et,, (To())o((()a)))) (fo(X)))iet)) Vo
Σ
+Vo -
(Y - 16xx)))
a -Yi-Jo(X))
#I,
#I, SELP
sEIp
along with
Σ (Y- fo(x)))
= Σ -Y)((x))))
V
#I p iELp
#Ip iELp
and thus
]
E Vo
EY: fo(Xi))
Σ (Y, fo().)) O.
#Ip
#Ip iel,
iET,
Similarly for the WBM loss, we get
Vo with E(g((Y)) (fe(X.)) (s:):)] = 0. Hence, by Proposition 2 and Theorem 4 in DDeuuleveut et al. 2020). we obtain the result for both loss functions.
E Wo (m[(Y) + (((()))))
(Jo(X.)), (c:))),
E PARAMETER ERROR BOUNDS
Corollary E.1. Assume the loss functions be strongly convex with constant &. > 0. Then, for o E (0, 1). the WBM noise-free estimate 8 satisfies
Diman
D2
Da
10
6 Pheerll
K
where t% is the noisy estimate and D Di D2 > 0 constant. In contrast, the noise-free MSE estimate satisfies
D
D2 0² Mae
10: Oiselll
+
Ki
K
where D is constant that depends (X,, Yt, EilEET 0r is the noisy estimate, and D2 0.
17

---

Published as conference paper ICLR 2026
Proof. First, denote the loss functions generically as lg. The proof is based on classical argument, and we
show it for W3M. Specifically, strong convexity implies for 0, 0"0" E Rd
lo 2 lo- + 110 1-
the minimizers
Log (g) S E0z (0) +
and
koz (0) S Eazze (0) +
for a > 0.
By combining these inequalities and the fact that
1e; (a) ( < D 0 + D2
Nest o²
one gets
x 110 Oise 11 S 2(D Dine o
which concludes the argument.
ADDITIONAL NUMERICAL RESULTS
We illustrate the evolution of the WEM loss for different Gaussian noise levels for the polynomial function (r 2) 4)(x 5)
as opposed to the MSE in figure 10 We note that the WBM loss changes more slowly with noise levels, suggesting it recovers regression functions that less impacted by noise. Additionally, report error bars
1.25
Ludiddian distance
wasserseen disancc
135
Enor
0.5
8.6
020
0.0
1.0
1.6 1.8
Figure 10: Loss evolution between function and its noisy version
over the test batches for the results presented section 6 as well as complementary numerical results for learning of the Navier-Stokes equation, Wave equation and Electricity Times Series. Note that for training on noisy the error bounds are much smaller due to the fact that the model training smoothens the noise, whereas for testing noisy data the difference can be significant among batches.
Batch Size Scaling
We report in figure 20 the relative test error for various batch sizes, under noise level 1. We note that WBM is not very sensitive to batch size values especially in the smaller range (up to 80).
18

---

Published as conference paper ICLR 2026
WTTM
VIIM
Nez ubooliee SnOE
100
Nozh asolitto ener
30
8C
20
20
Notso loval
No0ee Ievel
Figure 11: Test error evolution for Navier-Stokes operator learning 30% corrupted test data with Cauchy noise
Figure 12: Test error evolution for wave equation operator learning 30% corrupted training data with Gaussian noise
200
WBM
WSM
PRN
Chi-sq DAQ
Claft
150
Heen aboole STDE
Noin atselure omorr
125
25
0.0
0.5
3.0
2.0
2.5
Nosee knvid
Nosee lovell
Figure 13: Test error evolution for Navier-Stokes operator learning 30% corrupted training data with Cauchy noise,
Figure 14: Test error evolution for Navier-Stokes operator learning 30% corrupted training data with Gaussian noise
Wall-Clock Time We report in figure 21 the wall-clock time for WBM and the state-of-the-art x- DORO method across batch sizes, for the Navier-Stokes problem and 100 epoch of training. We note that WBM does not require higher time, since the cost is dominated by the dimension complexity remark of section 4. We also report wall-clock time across dimensions in figure 24 demonstrating the dominance of the dimension for time complexity. Furthermore, we report table ?? below the run time for each class of methods without hyper-paramedd
discussed the
tuning. Since WBM does not require any additional tuning beyond the standard deep learning reduction of computational cost of at least order of magnitude compared to the state-of-the-an which requires at least one additional
(such as batch size), it results in
19

---

Published as conference paper ICLR 2026
WIlN
C.1-sq
Noz abolibe STOE
Noch aboouut cITYr
40
2-
111
2.5
3.0
7.5
10.0
12.5
75.
17.5
20.0
hoisa leval
Momee
Figure 15: Test error evolution electric time series forecasting 30% corrupted test data with Cauchy noise
Figure 16: Test error evolution for electric time series forecasting 30% corrupted test data with Gaussian noise
100
WBM Regression
WBM Regression
Regression
Regression
CVaR DRO Regression
CVDit DRO Regression
Mean absolute arorr
Mean abooltte errt
150
125
100
25
0.0
0.5
1.5
2.0
2.5
3.0
3.5
4.0
10
Noise level
Noise level
Figure 17: Test error evolution for Navier-Stokes operator learning 30% corrupted training data with Cauchy noise
Figure 18: Test error evolution for Navier-Stokes operator learning 30% corrupted test data with Gaussian noise
Method Time (s)
WBM 2160.92
=DORO
WDRL 3476.27
2211.51
Table 1: Execution times of different methods.
20

---

Published as conference paper ICLR 2026
Regression MSE Regression
40
-00
erro
20-
Noise
Figure 19: Test error evolution for wave equation operator learning 30% corrupted training data with Cauchy noise
23
14-
130
120
Batch
Figure 20: Sensitivity of WBM to batch size
Multiplicative noise
We explore the effectiveness of WBM in case where the noise is multiplicative, for the Navier-Stokes problem. We report the relative MAE in figure 23 We observe that WBM maintains a relatively stable error across noise levels.
Correlated noise
Similarly, report the relative error of WBM for the Navier-Stokes problem, with strongly correlated noise. The noise is generated by sampling a Gaussian with covariance matrix the form + a11 T, where a control the correlation level and Σ follows Wishart distribution with diagonal scaling matrix. We note that WBM is less effective for correlated noise, which is an expected effect given that WBM is based on distributional matching.
21

---

Published as conference paper ICLR 2026
200
WBM
DORD
195
700
185
Time (3)
Time (3)
600
100
400
170
300
205
Batch size
Input Output dimension
Figure 21: Wall-Clock time across batch size
Figure 22: Wall-Clock time across dimensions
WBM
WBM
27.5
25.0
Reletive MANE (50
22.5
Relstve MAE (%)
82.5
80.0
37.5
35.0
12.5
30.5
1.0
1.2
1.4
1.6
1.8
2.0
2.2
2.4
1.0
12
1.4
1.5
3.8
2.0
2.2
2.4
Noise
correlation
Figure 23: Test error across multiplicative noise levels
Figure 24: Test error across noise corelation lev- els
22

---

Published as a conference paper ICLR 2026
G REPRODUCIBILITY
We provide version of the code used for the numerical experiments in the following link: code It is based on modifications of the publicly available code from (Raonic et al. (023) and (mbert et al. 2024). The experiments were performed NVIDIA Tesla T4 and Tesla P100 16GB Cores.
23


Reviews:
[{"review_id": "s4vFxFVmVT", "reviewer": "Reviewer_gAn2", "summary": "The paper studies robustness to global, additive label noise and argues that standard Wasserstein Distributionally Robust Learning does not improve robustness when the regression function is neither Lipschitz nor convex. It proposes Wasserstein Batch Matching: in each step, match the empirical distributions of predictions and responses within a batch via the 2-Wasserstein distance and train by minimizing that batchwise Wasserstein loss. The authors prove a consistency result for WBM in the noiseless case and analyze noise scaling of the loss and its effect on SGD iterates, showing milder sensitivity for WBM than MSE and linear-in-σ deterioration for WDRL under its assumptions. Experiments on operator learning and electric load forecasting find WBM competitive or better than MSE and divergence-based DRO baselines, especially under heavy-tailed noise, with lower computational burden than WDRL.", "questions": "- Could you check the format of the paper? margin and fontsize seem do not mach with the iclr template?\n- Missing reference on global regularization using WDR https://arxiv.org/abs/2203.00553", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- Clear formulation & intuition: Batchwise distribution matching gives a simple, architecture-agnostic robust objective.\n- Robustness demonstrated on wave/Navier–Stokes operator learning and ETDataset forecasting with TSMixer; WBM outperforms MSE and beats CVaR/Chi-Sq in cost-adjusted comparisons.\n- WBM involves a tractable per-step LP in the response dimension and avoids WDRL’s minimax complexity.", "weaknesses": "- The text says WBM solves an LP costing O(s) with s=dim(Y) and compares to WDRL’s O(s^3) when convex–concave; however, in practice the batch assignment (e.g., Hungarian/OT) can scale with batch size. Please specify the exact solver (e.g., Sinkhorn, network simplex) and report end-to-end wall-clock vs. batch size and s", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6ZaV0Kqj50", "reviewer": "Reviewer_uyo8", "summary": "This paper studies the problem of robustness to global additive noise such as measurement and quantization noise, which is often encountered in practical regression and physical modeling tasks. The authors introduce a new training framework named Wasserstein Batch Matching (WBM). It replaces instance-wise matching between predictions and responses with batch-level Wasserstein alignment between their empirical distributions.  \n\nTheoretically, the paper proves (i) **consistency** of WBM estimators (Proposition 4.1), and (ii) **favorable noise scaling properties** (Proposition 5.1), showing that WBM loss grows sublinearly with noise variance σ² compared to MSE and Wasserstein DRO (WDRL). Empirical results on **PDE operator learning (wave & Navier–Stokes)** and **electric grid forecasting (TSMixer)** confirm that WBM achieves significantly better robustness under Gaussian and especially Cauchy heavy-tailed noise, with lower computational cost than divergence-based DROs (CVaR-DRO, Chi²-DRO).", "questions": "1. The paper shows clear improvements on PDE operator learning and electric load forecasting. How effective is the proposed WBM method more generally? Could it achieve comparable gains on other types of regression or prediction tasks beyond these two domains?\n2. How sensitive is the method to design choices such as batch size, noise level, and the Sinkhorn regularization parameter?\n3.  The paper claims that WBM is model-agnostic and could, in principle, generalize beyond regression tasks. However, no experimental or theoretical evidence is provided to support this claim. Could the authors clarify whether WBM has been tested or formally analyzed for classification or structured prediction settings ？\n4. The appendix mentions that WBM achieves about a 10-fold computational gain mainly because it does not require hyperparameter tuning. Could the authors provide quantitative evidence or runtime measurements to support this statement, or clarify whether the comparison includes hyperparameter search time for other methods?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "+ Addresses a **real and underexplored problem** (robustness to global additive and heavy-tailed noise).\n+ Provides **rigorous theoretical analysis**, including a new _noise scaling law_ comparison with MSE and WDRL.\n+ Demonstrates **consistent empirical gains** on PDE and time-series tasks, matching theoretical expectations.\n+ Clear writing and strong logical flow from motivation → theory → experiments.\n+ Computation-friendly and easy to implement (batch-level matching without extra hyperparameters).", "weaknesses": "+ The theoretical contribution is relatively limited from the optimal transport perspective. The paper mainly adapts existing Wasserstein formulations to a batch-level matching context rather than developing new theoretical results or convergence guarantees.\n+ Given the moderate theoretical depth, stronger empirical support is needed to demonstrate the method’s effectiveness. However, the experiments are narrow in scope, focusing only on two case studies (PDE operator learning and electric load forecasting).\n+ The lack of broader evaluations across different domains such as tabular, visual, or structured regression reduces the generality of the claims. More systematic ablation studies on batch size, noise variance, and sample size would strengthen the experimental evidence.\n+ Claims about the model-agnostic property and computational efficiency of WBM are qualitative and are not supported by formal complexity analysis or quantitative runtime evaluation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3nWR1kNI31", "reviewer": "Reviewer_P1nf", "summary": "This paper addresses robustness to additive noise in regression, proposing a more distributional measure as an alternative to sample-wise Wasserstein robustness. The authors argue that standard Wasserstein robustness fails when regression functions are neither convex nor Lipschitz, and demonstrate that their \"batch matching\" method performs better with high-variance noise while being computationally cheaper.", "questions": "-- Is there any connection between your \"Wasserstein Batch Matching\" approach and the Central Limit Theorem? If so, I think it would be interesting to make this connection explicit.\n\n-- Why not compare to simpler robust losses like Huber loss or quantile regression?\n\n-- Strong convexity is assumed for SGD stationary distribution analysis (Cor. 5.2), then justified via RKHS regularization. This diverges from the deep-net setting used in experiments. Could you provide some discussion here?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "-- The focus on global noise in WBM is timely and distinct from adversarial/outlier robustness, while the authors provide a clear drawback to standard adversarial robustness measures such as sample-wise Wasserstein\n\n-- The theory seems genuinely novel (at least, it is to me) and interesting\n\n-- The scheme is model-agnostic and plugs into SGD, and the loss is differentiable via envelope theorem\n\n-- The experiments on PDE learning and time series demonstrate practical utility for the scheme\n\nOverall, I find the paper enlightening on an important limitation of standard sample-wise Wasserstein robustness, with an interesting practical fix as well as supporting analysis.", "weaknesses": "My main concern is that the paper is fuzzy on implementation details, which makes me a bit concerned for the reproducibility of the method, e.g.\n\n-- The WBM scheme's performance must heavily depend on batch size (larger batches = more flexible matching?), but this is barely discussed\n\n-- There is a 10 fold run improvement mentioned but no timing comparisons provided", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RuhUb5Hanx", "reviewer": "Reviewer_SL38", "summary": "This paper introduces WBM, a new regression loss to improve robustness to additive and global noise, addressing the limitations of Wasserstein Distributionally Robust Learning when regression functions are not Lipschitz or convex case. Theoretical results demonstrate that WBM loss aligns better with noise variance than WDRL. Numerical evaluation on PDE operator learning and electric grid forecasting further valid the robustness of proposed method.", "questions": "1.\tHow does WBM perform in non-convex neural networks beyond CNN operators?\n\n2.\tCan WBM expand to multiplicative or correlated noises?\n\n3.\tWhat will happen if the true regression function is discontinuous? \n\n4.\tHow does WBM performance depend on batch size? Is there any optimal range balancing between robustness and bias?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1.\tThe motivation of this paper is good. As WDRL lacks of robustness to global additive noise.\n\n2.\tThe paper provides rigorous proofs Wasserstein Batch Matching and noise scaling case.\n\n3.\tExperiments over diverse tasks validate the effectiveness of proposed WBM.", "weaknesses": "1.\tSome assumptions of theoretical results are not reasonable. For example, Proposition 4.1 requires bandlimited and continuously differentiable $f$ and co-monotonic of $g$, which may not hold for modern DNNs. Moreover, Corollary 5.2 relies on strong convexity and smoothness with second to even fifth order, that is not reasonable.\n\n2.\tNo error bars or confidence intervals. The number of runs only has 13 samples, which is small for modern deep learning.\n\n3.\tThe computational gains do not show practical analysis. Although a complexity gains achieved, the real-time runtime such as GPU time evaluation is missed.\n\n4.\tThere is limited comparison with other regression baselines. For instance, Noise2Noise.\n\n5.\tAppendix B used partial derivatives of Lagrange multipliers $\\alpha$, $\\beta$ with respect to $\\sigma$ without proving differentiability rigorously, only a sketch via implicit function is provided.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)