Elliptic Loss Regularization
Authors: Ali Hasan
Haoming Yang
Yuting Ng
Vahid Tarokh
Venue: ICLR 2025 Poster (2025)
Keywords: regularizer
loss landscape
diffusion
elliptic
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We propose a sampling based regularizer that constrains the loss as a function of inputs to satisfy an elliptic operator, giving favorable theoretical and empirical properties.

Abstract: Regularizing neural networks is important for anticipating model behavior in regions of the data space that are not well represented. In this work, we propose a regularization technique for enforcing a level of smoothness in the mapping between the input space and the loss. We specify the level of regularity by requiring that the loss of the network satisfies an elliptic operator over the data domain. To do this, we modify the usual empirical risk minimization objective such that we instead minimize a new objective that satisfies an elliptic operator over points within the domain. This allows us to use existing theory on elliptic operators to anticipate the behavior of the error for points outside the training set. We propose a tractable computational method that approximates the behavior of the elliptic operator while being computationally efficient. Finally, we analyze the properties of the proposed regularization to understand the performance on common problems of distribution shift and group imbalance. Numerical experiments empirically confirm the promise of the proposed regularization technique.

Published as a conference paper at ICLR 2025
ELLIPTIC LOsS REGULARIZATION
Ali Hasan 1.2.;, Haoming Yang? Yuting Ng', Vahid Tarokh²
Machine Learning Research, Morgan Stanley
2 Department of Electrical and Computer Engineering, Duke University
a ali. hasan, hacming yang, outing.ng, vahid,tarokh @duke. edu,
ABSTRACT
Regularizing neural networks is important for anticipating model behavior in regions of the data space that are not well represented. In this work, we propose a regularization technique for enforcing a level of smoothness in the mapping between the data input space and the loss value. We specify the level of regularity by requiring that the loss of the network satisfies an elliptic operator over the data domain. To do this, modify the usual empirical risk minimization objective such that we instead minimize a new objective that satisfies an elliptic operator over points within the domain. This allows to existing theory on elliptic operators to anticipate the behavior of the error for points outside the training set. We propose tractable computational method that approximates the behavior of the elliptic operator while being computationally efficient. Finally, we analyze the properties of the proposed regularization understand the performance on common problems of distribution shift and group imbalance. Numerical experiments confirm the promise the proposed regularization technique.
INTRODUCTION
Designing effective losses is longstanding goal for training neural networks. Standard techniques such as Empirical Risk Minimization (FRM) are commonly used for training, but these may overfit to finite samples of data and do not provide any explicit regularization for regions outside the support of the training data. This can be an issue in the case overparameterized neural networks, which minimize the loss up to arbitrary accuracy for data points in the training set but have unknown behavior outside the points in the training set. For real-world deployment scenarios, this can be especially concerning where data shifts and imbalances may be present. In response this limitation, various techniques have been proposed to regularize the loss landscape as a function of input samples of classifier, e.g. (Hermández-Gard) and König, 2018; Wang et al., 2021; Balestriero al., 2022).
ERM
Elliptic
Foundary
Pintericr
Prounday Maximum Principle
Vintetion
No Maximum Principle
Figure 1: Loss surface of the two moons dataset. The scattered points illustrate the training (blue boundary) and testing (orange interior) samples. The surface represents the loss values of well- trained classifier that classifies samples to their respective moon. Zooming in on the interior of the ERM loss surface, the training loss exceeds the loss of boundary (circled area) whereas elliptic regularization bounds the loss via the maximum principle of elliptic PDEs.
In this paper, we consider a viewpoint for regularization that stems from the properties of elliptic differential operators where we impose that the loss satisfies an elliptic partial differential equation with boundary data given by the loss. By imposing a loss that satisfies this operator, we can interpret
cqual contribution

---

Published as a conference paper at ICLR 2025
the expected loss through the lens of the powerful theory developed for elliptic operators and control the parameters of the operator to enforce desirable properties of the learned function approximator.
Why the PDE Perspective? PDEs are useful for characterizing the joint rates in change of different variables. Imposing constraints on how the loss changes as a function of perturbations in the input space is helpful in understanding the behavior of the learned function approximator under different transformations. Additionally, minimizing the loss under the PDE constraints allows us to use the tools from PDEs qualitatively describe properties of the loss function under the estimated function; an example is given in figure where we show how the maximun principle of elliptic operators is used to bound the loss of the estimator. This provides valuable insights into the robustness of the model for regions outside of the training data. To that end, we are mainly interested in the qualitative behavior of the operator and its implications in two major problems related to generalization in machine learning:
Data Shifts: How does the operator influence the loss under shifts in the data distribution? Data Imbalance: How does the operator influence the loss for underrepresented groups?
More explicitly, suppose we have data pairings of (X, y) with features X and target variables y. Data shifts are data that are collected under different conditions but represent the same phenomena, i.e. (T(x), y) for different transformations T applied to X. Data imbalance is the existence of subgroups o within the dataset that may be more or less frequently represented. By carefully choosing the parameters of the clliptic operator, under these regimes.
show how behavior of the approximator can be anticipated
To summarize, by grounding our regularization in the rich theory of elliptic PDEs, we provide principled approach for understanding the generalization properties of neural networks, specifically:
1. We describe new regularization that corresponds to the solution of an elliptic PDE;
We theoretically characterize the practical properties of this regularization through PDE theory;
We introduce an efficient computational approach that endows the properties of the elliptic regularization.
2 BACKGROUND
Our goal is to investigate loss that promotes functions that address the two above problems while minimally affecting performance related to our target tasks. We focus on problems of estimating a mapping between features and target variables, such as those in classification and regression tasks. To impose the desired regularity on our mapping, we can consider a few different techniques. On one hand, we can directly restrict the hypothesis class of functions from which we estimate our mapping to ensure those contain the relevant properties. On the other hand, it may be difficult to a priori know which class of functions is sufficient for the approximation task. We instead consider an approach that regularizes the loss over the data space without explicitly constraining the function class, which implemented through specific data augmentation.
Related Work The connection between regularization and data augmentation has been studied in a number of settings (Balestriero al., 2022; Geiping et al., 2022). A particular instance of augmentations is the mixup algorithm (Zhang et al., 2017). This algorithm has empirically shown to create more robust classifiers that are better suited perturbations of the data, provide level of uncertainty estimation with respect to new samples, and improve performance for classes with few samples (Zhang al., 2017). These properties lead to many extensions of mixup: one of which is mixupE. (Zou et al., 2023), which imposes an explicit directional derivative regularization to improve the mixup algorithm. The class of mixup algorithms has been theoretically analyzed from regularization perspective (Carratino et al., 2022; Zhang al., 2020). Various statements have been proved such as in Zhang et al. (2020) which showed that the loss at the interpolated points is an upper bound for & adversarial training loss. The regularization and generalization effect of mixup led to development of mixup based robust optimization methods such as UMIX (Han et al., 2022), which combines mixup and weighting algorithm to achieve robust training in classification;
2

---

Published as a conference paper at ICLR 2025
and c-mixup (Yao et al., 2022a), which samples based on the distribution of regression label distance for adapting mixup to the robust regression task.
Motivated by the thcoretical connection between regularization and data augmentation, propose PDE based regularization that can be implemented through a specific type of random data augmen- tation. This new regularization, which refer to elliptic regularization, is highly flexible and applies to arbitrary nonnegative loss functions in both regression and classification tasks.
Elliptic regularization focuses on improving optimization under the data shifts and data imbalance scenarios which are widely studied. Algorithms such as DRO (Sagawa et al., 2019) and DORO (Zhai al., 2021) were introduced to optimize over long-tail distributions and be robust against worst-case shifts in the data distribution. JTT (Liu et al., 2021) 1) and UMIX (Han et al., 2022) are two popular multi-stage methods that weigh data during training based on the earlier results. While most of the above algorithms deal with classification tasks, c-mixup has emerged as robust optimization method over regression (Yao et al., 2022a). These algorithms, including ours, are group-oblivious where the underlying subpopulation that causes data shifts and imbalance is unknown. We discuss group-informed methods Appendix F.4.
3 PROBLEM SETUP
Let's first define region which we will consider our domain. Given observations ((X,y) E X } with X' C Rd y C IR* we will define D = Dx Dy C X y subset of R x R* with nonzero Lebesgue measure. For purposes of analysis we will take D as the convex hull of X x y denoted as Cxxx- Additionally, denote the empirical measure of points within training set X x y
as xxxy = x Σ [X'30l 6(x.w): where | denotes the cardinality of set. We want to minimize the risk e D, x D,, + R, over all the points in the training set for a mapping fo Dx + D, with parameters O. We will disregard any implicit regularization associated with the optimization procedure associated with the minimization problem (i.e. the associated regularization when minimizing such problems using gradient descent). With these in mind, we will consider different flavors of the risk minimization problem and analyze the corresponding regularization from imposing additional terms.
Let us compare two frameworks for minimizing a risk f, the first performs empirical risk mini- mization (ERM) and the second minimizes the risk after applying class of transformations Tx. T parameterized by which we will refer as transformed empirical risk minimization (TRM):
ming Eixxs [((fo(X)))) (ERM)
ming Eg 1 (((fo(Tx X), vyy)) (TRM)
From equation ERM we can see that this optimization, which takes place over discrete set of points, relates to the issues we are trying to address. Since x is a set of measure zero, its unclear what the interpolating behavior of is for any subset of X' how ((fo(X) y) will behave for any subset of (X,9)= E D. Equation TRM on the other hand may cover more of 'D depending on how Ø is chosen.
Instead, we consider loss that is defined over all of 'D by modifying expectation in equation ERM in number of different ways. One way is explicitly regularizing the class of functions fa such that the behavior of {(Sol ()),)) over all D can be anticipated (for example, making So Lipschitz would constrain the growth between two points). However, explicitly regularizing fo can be difficult since constraining the function class when fe a neural network may require defining specific architectures. Alternatively, we can borrow techniques often used in computer vision problems (e.g. (Wang et al., 2021; Yang et al., 2023b)) by applying transformations or augmentations T T with parameters ø to X, y such that a new objective in equation TRM is minimized for some distribution on Φ. This also has the benefit that now can sample transformations rather than directly constrain the class of functions. The question now remains: what class of transformations would be useful to regularize (xfo(X), y) over DD
4 ELLIPTIC LOSS LANDSCAPES
Our answer to this question involves modifying the ERM problem construct loss that satisfies the elliptic operator. We introduce all components of the operator in three different parts. To first develop intuition behind the loss, describe elliptic operator that has only Laplacian terms that we wish
3

---

Published as a conference paper at ICLR 2025
to constrain. Using this example, then describe how this operator can be imposed through its stochastic representation given by the Feynman- -ac theorem, providing both a connection to the data augmentation regularization initially described and technique for implementation. Finally, using this stochastic representation, we describe the full method that includes low order terms through an importance sampling framework.
We define the loss landscape as a function u ((,,y) D + R.'. Our goal is to prescribe the function u(X,2) y) with a specified level of regularity over the data space in way that also imbues ((fx y) with regularity and thereby obtain the desired properties listed above. We propose doing this by solving new minimization problem defined by the following equations:
min e u),
(X,y) E D
= (X,),
(X,y) E
u(X, 3) ((f
(X,y) 0D
where V is taken with respect to (X, 1). a > is a coefficient related to the Brownian diffusion which will be later discussed, and 0D represents the boundary of the domain. The key aspects of this formulation are: 1. the constraint in equation 1, corresponding to the elliptic operator, ensures certain level of regularity within the domain beyond the training data; 2. the boundary condition in equation 2 connects the loss landscape u to the neural network loss (Xfo(X). v). The regularity imposed by the PDE in equation describes the behavior for regions of the space away from the observations. Specifically, equation represents the steady state of the heat equation, which diffuses the boundary data from equation 2. The diffusive behavior of this equation defines the regularity of the loss based on the closest points in the observation set for regions undefined in the observation set. Additionally, the equation provides a direct connection to an Itō diffusion process which we will discuss in the next section.
Interpreting the loss landscape The loss landscape u can be understood as the expected loss under new data point (X, E D. u satisfies an elliptic PDE with boundary data given by the points in the training set. Since the points in the training set provide all the information we have about the particular phenomena we are representing, the information at these points is diffused according to equation 1.
4.1 SOLVING THE REGULARIZED PROBLEM
Computing equation 1 requires solving PDE over high dimensional space, which can be com- putationally infeasible. Fortunately, solutions to PDEs of the type in equation 1 have a Monte Carlo representation which allows scalable computation of solutions. This is formalized through the Feynman-Kaa Kac formula, which describes the relationship between expectations of stochastic processes and PDEs (Øksendal, 2003; Pardoux and Raşcanu, 2014). For self-containment, we provide additional descriptions of the Feynman-Kaa formula as it relates to the problem we are solving in Appendix D. Using this form, the solution to equation I can be written as the following expectation
ux y) == min E (e(fo()))) U+) 20 I., yo y)
where Te. V/ satisfies d(x;,, 3t) = a dW, where = inf > 0 (xr, E X V). > 0 the diffusion coefficient, and W. is standard Brownian motion, and sample path distribution q. Since 5xxy Lebesgue measure zero, we will almost surely never hit a single point in the set, leading to a T being infinite. To circumvent this, we place a ball of size > 0 around each point and define new set (X' x x), = {z. y E D (('' s/) E X x y [[ v) - -<x,8')l || < E}. The hitting time is then defined under this new set as (c) = inf > 0 (zt, (e). E (X' x y)& U 0Cxxy',, }, and we compute
the expectation with respect to these hitting times. To relate this to equation TRM, sample paths of (Ia, V1) are a particular type of transformation of the data within the domain.
Using a finite number of samples, we write the empirical expectation as:
N ΣΣ e (1olo20).01) s.t. x6 200 j=1
min #
(1) = ,))
Note that this is different from the loss landscape definition which is given function of training iterations and considers perturbations in the parameter space.
4

---

Published as a conference paper at ICLR 2025
y1) (X' ). Sample paths of (Ize U/)) generated using standard Mally-Ma method. Qualitatively, as e increases, the behavior of the solution becomes piecewise constant each point in the training set. As E decreases, the solution becomes more diffusive where each data point becomes a point source.
where TE (6)
inf(L 0
4.2
GENERALIZATIONS OF THE OPERATOR
The form of the PDE described in equation 1 is motivated by the diffusive properties the second order derivative, which has the property that the loss values on the boundaries diffuse at rate proportional to a. Including only second order derivatives is not strictly necessary. however. The equation can be easily generalized to include lower order terms by including an advection term which acts as a deterministic term in the sense that information is being propagated according to an expected trajectory rather than diffused. The connection between the stochastic process (az, yt) and the operator easily extends to this case. Through Girsanov's 's theorem, we can define the Radon-Nikodym derivative dg between two path measures < to redefine the optimization criterion as
dp
Up(a, v) min Fyg ((Jo(E)))) V/)
3o s. yo
dq
which has the effect of reweighting some of the sample paths of (æe, vo). This results in u satisfying new PDE with lower order terms given by the structure of dp in dg
d+
8 u 227
=
b Vn
(5)
which we assume is parameterized by a drift function 6 D - gd-k This reweighting term can then be interpreted with respect the problems we are trying solve. For example, regarding the problem of data imbalance, sample paths corresponding to the underrepresented data groups can be reweighed greater influence the loss. This provides a connection to existing loss functions, such as the focal loss in Lin et al. (2017), where points within the training distribution are weighted according to their loss values if we choose b to be function of ((fo(X) y,y)
PRACTICAL CONSIDERATIONS OF THE REGULARIZER
There are few important practical notes regarding applying elliptic regularization. We will first describe how efficiently implement this loss and discuss choice for the importance factor in equation 4 while maintaining the desired properties of the operator. We then theoretically illustrate the main properties of inter- est and how they are relevant to the problems presented above: The maximum principle which allows us to bound the error on the interior of the domain in Propo- sition 1; and, the qualitative behavior the the loss landscape in regression cases of distribution shifts and class imbalance in Propositions 2 and 3.
X(2)
X
X
(4)
X
5.1 NUMERICAL. PROCEDURE
Illustration of the loss values over a domain with 4 points on the bound- ary. The expected loss at point X* is com- posed of losses at balls around X x(),
Figure 2:
In practice. it may be difficult to compute the first hit- ting time of the boundary within a reasonable amount of time. For example, computing the paths up to the first hitting time requires sampling for an undefined amount of time. For practical implementation, a sampling method that approximates the behavior of the first hitting time but does not require sampling unconstrained sample paths. We require that the approximation a) runs in a finite amount of time and b) maintains the important properties of the PDE. To do this, we first compute the pairwise distances between all points within batch according
1...4. 4. Black paths represent sample paths starting at X
consider
5

---

Published as a conference paper at ICLR 2025
to some distance metric. Then, for each starting point (æ. y), sample endpoint according to discrete distribution P over endpoints with mass inversely proportional to the distances computed. Equation 3 would ordinarily be solved over all (X D (e.g. by sampling uniformly over the space). We instead consider sampling over paths that connect data points using the minimum distance. Figure illustrates an example of this where four points on the boundary are used to define the loss at X The loss is minimized at all points along the paths between the source point X and the target points X( i = 1...4. 4. This is done by sampling Brownian bridges between the pairings and optimize the following loss:
JPM BB [%. e(fo(X.). va)ds
min
(6)
XB -- . as a Brownian bridge sample path where denotes starting points and - denotes end points and IT is the set of points in the support of d(xxy) x 6(xx))- This expectation involves solving the problem equation 3 for all points along the Brownian bridge between X- and X_. yy This minimizes the loss over all paths between the data points where the paths satisfy the distribution corresponding to the PDE between the endpoints. This loss, through an application of Dynkin's 's formula, corresponds to exactly solving equation 3 for points along the Brownian bridge paths, which we formally describe in Lemma in Appendix B. Since the support of the bridges includes the domain, all points within the domain should eventually be sampled. Additionally, the starting point of the bridge can be an arbitrary point within the domain; we use the data points for convenience, but the theoretical properties still hold with this sampling technique. The connection to the transformation based regularization described in equation TRM is made more explicit ' in equation TRM can be seen the Brownian bridge transformation on the data points. Note that there exist some numerical error with this approximation, since it requires discretizing a continuous process, we refer Graham and Talay (2013, Chapter 7.2) for additional details. We describe additional implementations of the bridge regularizer and how they relate to the original elliptic operator in Appendix E.
for all where we denote BB
5.2 BOUNDING THE LOSS
One property we would like to guarantee on our loss is, for any new point within D, what should we expect the loss to be? For example, under distribution shifts we may want to predict what the expected loss is. We do this through characterizing the loss landscape as a function of the observed data. Since we impose that u must satisfy a PDE, we can use the properties of the solution of the PDE to bound the solution in the interior of the domain. In particular, since the PDE is elliptic, the maximum principle is satisfied. This allows us to bound the interior of the domain by the values on the boundary (which are our training points). We formalize this in the following proposition:
Proposition (Bounding the Loss for an Interior Point). Consider any point X, E and suppose the function pairs u, Sa solves equation 1. Then, the expected loss u at X, y satisfies the following inequality:
min ((fo(X). v) S v(X,v) < X.yXXxy
X.gEXxy (X),)))
Proof. The proof follows from the fact that satisfies an elliptic PDE that satisfies both the maximum and minimum principle. Therefore, all extreme points lie on the boundary of the domain which correspond to the points in the training set.
This has an important implication insofar as we can anticipate the behavior of our loss landscape for new, unseen points by guarantceing the loss is bounded by the training set points. This condition holds for any point within the domain DD
5.3 EFFECTS ON DISTRIBUTION SHIFTS AND DATA IMBALANCE
Finally, we discuss how to interpret the elliptic operator for the problems of interest involving distribution shift and class imbalance which motivate this work. We focus on describing the behavior for the regression case where the loss given by the L distance and the function f is a two-layer ReLU neural network and study the loss landscape u( u(X,y y) satisfying the following conditions:
vVxu(X,)) 3) = 0, x(x,y)= (f(x) -
X,y E D
X, y E 0D
6

---

Published as a conference paper at ICLR 2025
with f(X) W2 (W) and where Wo R and W1 R
Applied to data shifts Suppose we define a set of transformations T : X' - X parameterized by s PCC & for some space of parameters & and require that these transformations only affect the input features X but do not affect the target variable y. As examples, we can think of medical imaging data collected at different hospitals under different operating conditions but of the same disease and of the same modality with the target variable being the class of disease in the image. Certain parameter values of &p may be sampled more in the dataset than others (e.g. data from larger hospitals) leading to greater uncertainty on samples from the underrepresented parameter values. We will denote an example subset with few sample points x As such, the subset of X associated with (X)),, y E $, will have small cardinality. Then, the expected hitting time of point within the space is longer since fewer points are within the vicinity of point within this subset.
In the following proposition, we formalize this idea and study the changes in the loss landscape when an affine transformation is applied to the features:
Proposition (Expected Error Under Affine Transformations), Consider the loss landscape satisfying equation 7 and consider the class of affine transformations T where T(X) = ArX + br e T for AT E Raxdd and br E TRd and (X, y) ED + X x Y. Suppose the true error is given by (J(X) - y)² and is C-Lipschitz. Then, the loss landscupe satisfies
w(T(X),,) S 2W, Wo F CIAKIA|+ + 2s),
where: = : If(X) v). A AT X + bt X, and TT = infr>o V1 OD Xo Ar br, 1/0 = y).
Proposition 2 allows to anticipate the behavior of our loss under affine distribution shifts of our data. By applying this regularization, we can guarantee that the expected loss is no greater than a function of the parameters of the neural network used for regression.
Applied to data imbalance With regards to imbalanced data in the classification setting, consider a local structure on the points within each class. We will partition y into two classes, Vv and V% for the over-represented and under-represented classes respectively. We can show that when the loss landscape satisfying the operator in equation 7, the expected loss is greater for regions of low data density versus for regions of high data density.
Following a similar argument as in Proposition 2, we note that the hitting times for Vv, is greater than the ones in 3%, with certain probability. The proposition is based on a probabilistic interpretation of how many components of the weight matrices in the neural network are greater than or equal to zero.
Proposition 3 (Expected Error in Regions of Low Density), Consider the loss landscape u u(X, y) satisfying equation 7. Let q(e) := P(V2(f(X)
y)² 2cW, Wo (W; Wo) be a continuous function, then an e* E (0, 1) such that q(c") =1 1 e" exists. Choose the smallest e*_ Let (Xv. and (Xx,ya) be two points where Y^ is in a class that is well represented such that for Tyv infesol(Xe, Mt) E 'D
Xo = Xv, y0 = yv3 and Tyn = infrso(((ee, yt) E D Xo = Xx, yo = yn } the relation TA 2 Tyx holds. Suppose also that (f(Xv) - yv)² = (f(Xx) - yn Then, u(Xv, vv) u(Xx, VA) with probability 1 e..
Proposition 3 states that within regions of low data density, the loss landscape to be greater than in regions of high density, which is important cases involving data imbalance.
EXPERIMENTS
We present empirical results over a comprehensive set of tasks related the proposed regularization scheme. We first empirically evaluate the bound derived from Proposition 1 and verify that the proposed regularization retains the necessary loss within the domain of interest. Next, we benchmark the elliptic regularization against another popular regularization scheme known mixup (Zhang et al., 2017) and its variants on balanced classification and in-distribution regression. We then investigate the benefits of elliptic training classification and regression with imbalanced group populations, imbalanced domains, and subdomain shifts. Finally, we experiment with elliptic training on imbalanced classification. Additional comparisons and ablation studies are provided in
7

---

Published as a conference paper at ICLR 2025
Table 1: Classification error on CIFAR-IO, CIFAR- 100, best method is bolded while the second best is underlined
Table 2: Classification accuracy on Tiny- Imagenet 200, best method is bolded while the second best is underlined
CIFAR-10
CIFAR-100
Top-1 (%)
Top-5(%)
ERM
4.71
24.68
ERM
54.97
72.71
mixup
4.53
23.03
mixup
54.65
72.53
PreActRN50
PreActRN18
mixupE
3.53
20.23
mixupE
02.21
82.09
Elliptic
4.15
18.95
Elliptic
65.03
84.98
ERM
4.21 20.07
23.20=004
ERM
57.25
72.58
PreActRN101
mixup
4.43
23,05
PreActRN34
mixup
57.79+0.4
76.15+0.4
mixupE
3.35
18.86
mixupE
65.37
83.77+0.4
Elliptic
3.5640.2
19.03
Elliptic
66.67 +1.7
85.71
ERM
2.24
22.20
ERM 55.91
73.50.10.6
mixup
3.03
19.38
mixup
54.8640.6
73.11
Wide-RN28
nixupE 2.94
17.12
PreActRN50
mixupE 67.22++0.4
85.1440.4
Elliptic
3.09
17.71
Elliptic
70.28
88.14
Appendix F with detailed data descriptions for all datasets in Appendix G. All hyperparameter settings described in Appendix H. We note that although the implementation described Section 5.1
specifies computing pairwise distances between data, we empirically show in Appendix F that this distance computation has minor effect on performance of the method.
6.1 EMPIRICAL STUDY OF PROPOSITION 1
Using the two moons dataset neural network by solving the PDE through equation 6 enforces the bound in Proposition 1. We first train 2-layer, 4-hidden unit multi- layer perceptron using the different op- timization criteria of ERM, elliptic, and mixup with the generated boundary data (blue Figure 3). We then compute the loss on the within-boundary data (orange in Figure 3). We sampled 10000 points on the boundary to train ERM and mixup while using 1000 boundary data with 10 discretized timesteps for the elliptic loss to sample Brownian bridges. We show in Fig- ure 3 that only the neural network trained using equation maintains the interior loss within the bounds of the boundary loss as specified in Proposition 1.
example, empirically illustrate that optimizing the
Boundary (x,y am))
Intedor E)
Ellipiic
Mixup
Figure 3: Comparison of loss for data on boundary vs within the boundary of the data space. The dashed line indicates the max loss of the boundary data.
6.2 EXPERIMENTS ON REGRESSION AND CLASSIFICATION
We now consider how the proposed elliptic regularization performs on standard regression and classification tasks in machine learning. While are interested in understanding the regimes in which the regularization scheme improves the performance for the cases of group imbalance and distribution shift, are also interested in seeing how the performance is in "normal" regimes without these issues. To do this, we evaluate elliptic training against suitable baselines of vanilla mixup (Zhang al., 2017) and related -tate-of-the-art mixup algorithms for classification (Zou et al., 2023) and regression (Yao et al., 2022a).
We then consider the performance when applying the Brownian bridge algorithm described in Sec- tion 5.1 without additional importance weighting. We showcase that using the elliptic regularization achieves comparable results the state-of-the-ar mixup algorithms for these balanced, in-distribution experiments. We benchmark on CIFAR-10 and CIFAR-100 100 with classification error rate in Table I;
Two endpoints and 8 diffusion steps.
8

---

Published as a conference paper at ICLR 2025
we also evaluate on tiny-Imagenn 200 and report the top-1 and top-5 accuracy in Table 2. A number of different architectures are used in this experiment including PreActResNet of various depths (He et al., 2016) and with 5 random seeds. The results suggest that elliptic training generally improves performance over the vanilla mixup algorithm and improves upon mixupE the more complicated Tiny-Imagene 200. The results also suggest that the elliptic regularization may be most helpful in datasets with larger number of classes. This may be because the interpolation between classes will more likely lie on the interior of the simplex rather than on discrete corners, as is the case for datasets with fewer classes. For regression, we benchmarked the ellip- tic regularization against c-mixup algorithm on two real world tabular datasets: Airfoil Self-Noise (Airfoil) (Brooks et al., 2014) and NO2 (Kooperberg, 1997); Airfoil contains acrodynamic features of airfoil blade and re- lated acoustic statistics and NO2 studies the concentration of NO2 partieles given traffic
10 (Zagoruyko and Komodakis, 2016). The experiment was repeated
Table 3: Average Regression (RMSE) on Regression
Algorithm
AirFoil
NO2
C-mixup (Yao et al., 2022a) Elliptic
2.88
0.524
2.85 40.14
0.523 +0.008
volume and meteorological variables. Average RMSE is reported in Table 3, and the standard deviations are calculated over 10 repetitions.
6.3 ROBUSTNESS TO GROUP IMBALANCE AND DISTRIBUTION SHIFT
We now explicitly apply the elliptic regularization to the loss function to evaluate the performance on data with imbalanced subpopulations and under distribution shifts. We continue to evaluate on both classification and regression tasks. For the robust classification task, we benchmark against two types of algorithms: 1) single-stage training algorithm including Focal Loss (Lin al., 2017). CVaR-DRO and DDRD (Levy et al., 2020), CVaR-DORO and =²-DORO (Zhai et al., 2021); and 2) two-stage training algorithm including JTT (Liu et al., 2021) and UMIX (Han et al., 2022). These are group- oblivious algorithms that train without subpopulation/domaan information; full comparison with group-informed algorithms on classification tasks is provided in appendix F.4. We benchmark elliptic training with importance weighting drift b- (((f(x), y) in equation 5 on WaterBirds (Koh et al., 2021) and CelebA (Sagawa et al., 2019) to evaluate performance on robustness for group imbalance. We also evaluate the Camelyon17 (Bandi et al., 2018) dataset examine robustness under domain shifts. To remain consistent with previous literature, we report average and worst group accuracy for Waterbirds and CelebA across random seeds and report 10 seed average accuracy for Camelyon 17. The classification results are presented in Table 4. The elliptic regularization greatly improves upon the one-stage algorithms and is comparable to the two-stage algorithm Tust-Train--e Twice (Liu et al., 2021), which requires two separate training sessions of the model.
Table 4: Robust classification accuracy under imbalance subpopulations and domain shift; best methods are bolded and best one-stage methods are underlined.
Algorithm
WaterBirds
CelebA
Camelyon17
Avg(%)
Worst(%)
Avg(%)
Worst(%)
Avg(%)
Focal Loss (Lin et al., 2017)
87.0+0.5
73.1
88.4
72. 13.1
68.1
CVaR-DRO (Levy et al., 2020)
90.3
77.2+2.2
86.8
76.9
70.5
CVaR-DORO (Zhai et al., 2021)
91.5
77.0+2.8
89.640.4
75.6+4.2
67.3
X (Levy al., 2020)
88.341.s
74.0+1.8
87.7
78.4+3.4
68.0+6.7
x" 2DDRRO (Zhai et al., 2021)
89.5
76.043.1
87.020.6
75.623.4
68.027.5
Elliptic +
92.0+a.a
84.141.1
91.3
77.4+4.6
77.9
Two-stage: JTT (Liu et al., 2021)
93.64
86.0
88.0
81.1+NA
69.146.4
Two-stage: UMIX (Han et al., 2022)
93.0 +0.5
90.0+1.1
90.1
85.3
75.1
For the regression task, we compare against the c-mixup (Yao et al., 2022a) algorithm. We experiment on SkillCraft and Crime to examine domain shift and RCF-MNIST for sub-domain shift. Communities and Crime (Crime) (Redmond, 2009) and SkillCraftI Master Table (SkillCraft) Blair et al. (2013) are real-world tabular datasets where domain shifts exist between the training and testing data. Detailed description of these dataset is provided in Appendix G. All experiments for regressions are repeated with 10 different random seeds; we report the average and worst RMSE in Table 5.

---

Published as a conference paper at ICLR 2025
Table Regression (RMSE) performance under (sub)domain shifts datasets; best method bolded.
Algorithm
SkillCraft
Crime
RCF-MNIST
Avg
Worst
Avg
Worst
Avg
Worst
C-mixup (Yao at al., 2022a)
6.27 50.537
L.000
0.132.20.00 003
00.010
0.180+0.00 0.165=0.001
Elliptic
5.97 #0.283
9.17+1.150
0.132
0.164
0.162=0.002
0.178+0.002
6.4 CLASS IMBALANCE AND ROBUSTNESS TO NOISE
We further evaluate the efficacy of elliptic regularization from two perspectives: imbalance classifica- tion (average and worst class) and robustness to noise. We focus on the real-world medical dataset Med-MNIST (Yang et al., 2023a) where class imbalance exists due to the prevalence of different diseases in the patient populations. We chose sub-datasets of the Med-MNIST dataset to evaluate the effectiveness of the elliptic regularization. Detailed descriptions of these datasets are provided in Appendix G. To showcase the robustness induced by elliptic training, we include 50% label noise to the training data where half of the training data labels are randomly shuffled. Average and worst-class accuracy over 10 seeds are presented in Table 6, where we note the significant improvement of elliptic training and further improvement brought forth by importance weighting. The results in Table 6 indicate that elliptic training performs well the class-imbalance tasks while still maintaining high performance under large label noise.
Table 6: Average and worst class classification accuracy for selected Med-MNIST dataset, best
method is bolded while the second best underlined
Breast
Blood
Path
OrganC
Method
Avg(%)
Worst %6)
Avg(%)
Worst(%)
Avg(%)
Worst(
Avg(%))
Worst(%)
ERM
80.0+3.4
32.4_13.8
81.4
56.3
55.7
10.7
84.1
65.1
mixup
83.1
47.1 10.9
78.8
41.7
54.322.1
1.242.1
85.622.4
56,8211.0
mixupE
76.5
16.7=11.3
74.223.6
31.9216.9
52.5
2.747.6
70.427.6
35.6216.8
Elliptic
87.6
67.6-4.4
85.5
60.2+14.9
62.3+225
26.4=10.2
87.7+0.
65.943.4
IW
87.3
67.6
84.1
68.0+13.7
62.7+1.8
20.9
88.0
66.2
7 DISCUSSION
In this work, we proposed a regularization technique for the loss landscape over defined domain where we borrow ideas from PDEs to impose properties on the landscape. We require that the loss landscape satisfies an elliptic operator and described computational tools for enforcing this. The operator allows us to bound the expected loss on the interior of the domain using the points that we observe while also providing theoretical behavior for the loss in distribution shifts and class imbalance. There exist many avenues for extending the work provided. properties of Radon-Nikodym derivative corresponding to the parameters of the first order derivatives of the PDE should be undertaken. We chose the current parameters due to their ease in computation, but there may be investigation involves studying the coefficients of the operator from a stochastic control perspective understand the behavior under, for example, stochastic gradient descent. One can easily show that the gradient with respect to parameters satisfies another elliptic PDE over the space of inputs using similar tools of analysis, but the implications on the training behavior of the parameters still not clear. It could be helpful to link to other theories, e.g. flatness of minima, to understand the effects on the parameter space; or other types of PDEs, such as the parabolic PDEs, to provide another avenue of study on other learning regimes ( (Yang et al., 2025). Finally, extensions to data supported on manifolds could provide an interesting avenue for regularization for data on complex geometries.
further study on the
ways one can impose coefficients. Related this, important theoretical
There are few limitations with the proposed framework. Without using the Brownian bridge technique, solving the PDE may have infinite time until hitting the first point on the boundary. It is therefore advisable to study what happens when including boundary conditions (e.g. reflecting boundaries) or imposing drift such that the diffusion is guaranteed to hit point on the boundary. Additionally, the computation time is greater for sampling these stochastic processes than with ERM.
Limitations
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
Ali Hasan and Vahid Tarokh were supported in part by the Air Force Office of Scientific Research (AFOSR) under award number FA9550- 20- 1-0397. Haoming Yang was supported in part by Air Force Office of Scientific Research (AFOSR) under award number FA9550-22-1-0315
REFERENCES
Kartik Ahuja, Ethan Caballero, Dinghuai Zhang, Chan-Chrittophe Gagnon-Auder, Yoshua Bengio, Ioannis Mitliagkas, and Irina Rish. Invariance principle meets information bottleneck for out-of- distribution generalization. Advances in Neural Information Processing Systems, $4:3438-3455 2021.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization. arXiv preprint UrXiv:190..02893, 2019.
Randall Balestriero. Leon Bottou, and Yann LeCun. The effects of regularization and data augmenta- tion are class dependent. Advances in Neural Information Processing Systems, 35:37878-37891, 2022.
Peter Bandi, Oscar Geessink, Quirine Manson, Marcory Van Dijk, Maschenka Balkenhol, Meyke Hermsen, Babak Ehteshami Bejnordi, Byungjae Lee, Kyunghyun Paeng, Aoxiao Zhong, al. From detection of individual metastases to classification of lymph node status at the patient level: the camelyon17 challenge. IEEE transactions on medical imaging, 38(2)5550-560, 2018.
Mark Blair, Joe Thompson, Andrew Henrey, and Bill Chen. SkillCraft1 master table dataset. UCI Machine Learning Repository, 2013. DOI: hups://bb.oogg/0..4442/55111
Thomas Brooks, D. Pope, and Michael Marcolini. Airfoil self-noise. UCI Machine Learning Repository, 2014. DOI: http::/loo.or//1.44433///vw
Luigi Carratino, Moustapha Cissé, Rodolphe Jenatton, and Jean-Philippe Vert. On mixup regulariza- tion. The Journal of Machine Learning Research, 23(1):14632-144) 14662, 2022.
Jonas Geiping, Micah Goldblum, Gowthami Somepalli, Ravid Shwartz-Ziv, Tom Goldstein, and Andrew Gordon Wilson. How much data are augmentations worth? an investigation into scaling laws, invariance, and implicit regularization. arXiv preprint arXiv:2210.0644 2022.
Carl Graham and Denis Talay. Stochastic simulation and Monte Carlo methods: mathematical foundations of stochastic simulation, volume 68. Springer Science & Business Media, 2013.
Zongho Han, Zhipeng Liang, Fan Yang, Liu Liu, Lanqing Li, Yatao Bian, Peilin Zhao, Bingzhe Wu, Changqing Zhang, and Jianhua Yao. Umix: Improving importance weighting for subpopulation shift via uncertainty-awar mixup. Advances in Neural Information Processing Systems, 35: 37704-37718, 2022.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference computer vision and pattern recognition, pages 770-778, 2016.
Alex Hernández-Gareía and Peter König. Data augmentation instead of explicit regularization. arXiv preprint arXiv:1806.03855 2018.
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsub- ramani, Weihua Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, Tony Lee, Etienne David, Ian Stavness. Wei Guo, Berton Earnshaw, Imran Haque, Sara M Beery, Jure Leskovec, An- shul Kundaje, Emma Pierson, Sergey Levine, Chelsea Finn, and Percy Liang. Wilds: benchmark of in-the-wild distribution shifts. 139:5637-5664, 18 24 Jul 2021.
Charles Kooperberg. Statlib: an archive for statistical software, datasets, and information. The American Statistician, 51(1)998, 1997.
11

---

Published as a conference paper at ICLR 2025
David Krueger, Ethan Caballero, fbern-Hentie Jacobsen, Amy Zhang, Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Outto--distribution generalization via risk extrap- olation (rex). In International Conference on Machine Learning pages 5815-5826. PMLR, 2021.
Daniel Levy, Yair Carmon, John C Duchi, and Aaron Sidford. Large-scale methods for distributionally robust optimization. Advances in Neural Information Processing Systems, 33:8847-8860, 2020.
Tsung-Yi Yi Lin, Priya Goyal, Ross Girshick, Kaiming He, and Piotr Dollár. Focal loss for dense object detection. In Proceedings of the IEEE international conference on computer vision, pages
2980-2988, 2017.
Evan Z Liu, Behzad Haghgoo, Annie S Chen, Aditi Raghunathan, Pang Wei Koh, Shiori Sagawa, Percy Liang, and Chelsea Finn. Just train twice: Improving group robustness without training group information. In International Conference Machine Learning, pages 6781-6792. PMLR, 2021.
Bernt Øksendal. Stochastic differential equations. Springer, 2003.
E. Pardoux and A. Raşcanu. Stochastic differential equations, backward SDEs, partial differential equations. Stochastic Modelling and Applied Probability. Springer International Publishing, 2014. ISBN 97833
Michael Redmond. Communities and crime. UCI Machine Learning Repository, 2009. DOL: https///ooi.rr//1024433////////
Shiori Sagawa, Pang Wei Koh, Tatsunori Hashimoto, and Percy Liang. Distributionally robust neural networks for group shifts: On the importance of regularization for worst-case generalization. arXiv preprint 1rxiv:I9111 1.08731, 2019.
Yuge Shi, Jeffrey Seely, Philip HS Torr, N Siddharth, Awni Hannun, Nicolas Usunier, and Gabriel Synnaeve. Gradient matching for domain generalization. arXiv preprint arXiv22044009933 2021.
Baochen Sun and Kate Saenko. Deep coral: Correlation aligmment for deep domain adaptation. In Computer Vision-ECCV 2016 Workshops: Amsterdam, The Netherlands, October 8-10 and 15- 16, 2016, Proceedings, Part III 14, pages 443-450. Springer, 2016
Yulin Wang, Gao Huang, Shiji Song, Xuran Pan, Yitong Xia, and Cheng Wu. Regularizing deep networks with semantic data augmentation. IEEE Transactions Pattern Analysis and Machine Intelligence, 44(7):3733-3748, 2021.
Minghao Xu, Jian Zhang, Bingbing Ni. Teng Li, Chengjie Wang, Qi Tian, and Wenjun Zhang Adversarial domain adaptation with domain mixup. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 6502-6509, 2020.
Haoming Yang. Ali Hasan, and Vahid Tarokh. Parabolic continual learning, In The 28th International Conference on Artificial Intelligence and Statistics, 2025.
Jiancheng Yang, Rui Shi, Donglai Wei, Zequan Liu, Lin Zhao, Bilian Ke, Hanspeter Pfister, and Bingbing Ni. Medmnist v2-a large-scale lightweight benchmark for 2d and 3d biomedical image classification. Scientific Data, 10(1):41, 2023a.
Shuo Yang, Yijun Dong, Rachel Ward, Inderjit S Dhillon, Sujay Sanghavi, and Qi Lei. Sample efficiency of data augmentation consistency regularization. In International Conference on Artificial Intelligence and Statistics, pages 3825-3853. PMLR, 2023b.
Huaxiu Yao, Yiping Wang, Linjun Zhang, James Y Zou, and Chelsea Finn, C-mixup: Improving generalization in regression. Advances in Neural Information Processing Systems, 35:3361-3376.
2022a.
Huaxiu Yao, Yu Wang. Sai Li, Linjun Zhang, Weixin Liang, James Zou, and Chelsea Finn. Improving oftttf--isttibutionn robustness via selective augmentation. In International Conference on Machine Learning, pages -5407-25437. PMLR, 2022b.
12

---

Published as a conference paper at ICLR 2025
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. arXiv preprint arXiv: 2016.
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.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. arXiv preprint 17110:1710.09412, 2017.
Linjun Zhang, Zhun Deng, Kenji Kawaguchi, Amirata Ghorbani, and James Zou, How does mixup help with robustness and generalization? In International Conference on Learning Representations, 2020.
Yingtian Zou, Vikas Verma, Sarthak Mittal, Wai Hoh Tang, Hicu Pham, Juho Kannala, Yoshua Bengio, Amo Solin, and Kenji Kawaguchi. Mixupe: Understanding and improving mixup from directional derivative perspective. In Uncertainty in Artificial Intelligence, pages 2597-2607. PMLR, 2023.
13

---

Published as a conference paper at ICLR 2025
ELLIPTIC OPERATORS IN Loss LANDSCAPES (SUPPLEMENTARY MATERIAL)
A ALGORITHM DETAILS
To fully supplement the algorithmic contributions the main paper, we detail the elliptic training procedure that solves the PDE of equation using Brownian bridges in Algorithm 1. There are several hyperaarameters for the training algorithm, such as the number of bridges ni, the number time discretization steps ng, the bridge diffusion constant ob and the strength of importance weighting We will present some ablation studies exploring the effect of hyperparameter & and og in Section F.1. We also note that depending bridges of y. For example, for a classification task, y E [0, I.. so we project the Brownian bridge of onto the simplex by taking the absolute value of the path and normalizing it.
different tasks, there could be constraints imposed on the Brownian
Algorithm 1 Elliptic training algorithm
Input: Data Xan, Van, related Brownian bridge hyperparameters. Initialize: neural network fo-
Compute pairwise 12-distance between all data in Xail (optional).
for X, y in mini-batched Xall, Vall do
Obtain Brownian bridge pairs X', y by sampling inversely proportional distance computed
X'. for arbitrary timestep (see Algorithm 2)
Sample Brownian bridges X Ve ~ BB X,y
Compute loss & as equation 6 with {(f())), Uo) using the Euler's method over E [0, 1].
if use importance weight then
Compute gradient VL, then compute f = f EVE end if
Minimize é using gradient-based optimizer end for
To sample a brownian bridge, we used the Euler-Maruyama Scheme below
Algorithm 2 Sampling Brownian bridge
Input: Initial condition X E Rd, terminal condition X' R diffusion coefficient 0, number of time steps M
Let A, T-te where to = 0, fM 1, and f,41 44 = A, Sample M samples from standard normal N (0, 1), denoted as AW, Integrate samples to form Brownian motion W, = Ei OWW, VA,
Set W, = 0; then compute W.X = W, + X
Output: Brownian bridge BB, W.X t(Wx - X')
A.1 COMPUTATION COMPLEXITY
The most computationally expensive part of our method is computing pairwise distance for every dataset, which is of time-complexity Onn²). However, this can be pre-computed and saved dictionary be called for O(1) speed during training. We also show, through empirical experiment, the computation of distance may not be necessary.
We also tested the computation time between UMIX and the proposed method on the waterbirds dataset. All hyperparameters of UMIX are taken from Han et al. (2022). Three run average training time with the proposed method is 8345 seconds while UMIX requires 11795 seconds (stage 5918, stage 2: 5877). This result suggests that the additional computational cost associated with sampling bridges is marginal and the proposed elliptic regularization remains competitive in performance with lower training time.
14

---

9. [ Published as a conference paper at ICLR 2025
B PROOFS
B.1 PROOF OF PROPOSITION 2
Proof. We will assume T 7r, where Tr is the stopping time for the Brownian motion when starting at the transformed coordinates. This is a natural assumption since the transformed data may be further from other points after transformation and result longer hitting time. From Dynkin's formula, we can write the solution in terms of the stopping time and the boundary condition as:
[R(((X) - 3) X. = T(X), yo y] (J(T(X)) - y)²
E
(ff(x)) - yy) Xo = T(X), yo =
Since (2(()))) - y)² < 2W, Wo) T, we rewrite the solution as
m.. [I(S(X) v) Xo = T(X),M 30 y] (I(T(X)) - y)² 2W, WWWWo( Tr..
Compared with the original expected error given by (f(X) y)² - (X( v) Xu = X, No y| S (f(x) - -2² + 2WW,W,(WWWW))
Taking the difference, we get
error S WWW, TT (f(T(X)) - u)² - (f(x)
s WWWWo(W1WW) (f(((x)) ((X))((((X))) f(X) - 22
s WWWWo((WWW)
+ - -(X))]((((())))+ -
28l
s "/ - br - (A/ - f(X) + + -
x WW] -
-CAAA - -br - x1((f(ArX+br) - f(X) + +f(()) - 2yl)
< (|01:o(WWWW)"7++((C0 -br X((C|ATX-br - -X| (f(X) f(X) 2yl)
s 2W1Wo(W1Wo) (T + bT XI(CC b X| 2z).
B.2 PROOF OF PROPOSITION 3
Proof. We follow a similar proof strategy as in Proposition where we use Dynkin's formula to compare the solution at different points. The first assumption follows from placing uniform distribution over the ReLU activation functions being nonzero. Note that from the intermediate value theorem there exists E (0,) 1) such that q(e*) the expected loss at the two points by
e*. From Dynkin's 's formula, can compare
[ v²(s(X.) - y.) Xo Xv. 20 = yV
u(Xv. = (((x.) - w²
and for the underrepresenteed class,
[ v((((X.) - yy Xo = XA. yo = Un
u(X,, vA) = (f(xx) un) + E
Since 2e* W, Wo(W, S v²(((X,) y.)² < 2W, (W Wo) we let the lower bound
correspond to the integrand of the solution at (X,. yx) and the upper bound correspond to the solution at (Xv, yv). This gives us with probability e-c.
(f(X)) y^)² 2c N,,, W.,, TyP TT, u((x.vn))
(8)
and with probability 1:
u(Xv, vv) S (f(x)) v²)² - + 2W W (WW,W,) W)) Tyy-
(9)
Setting Ty. 2 Tay, the left hand side of equation is greater than the right hand side equation This achieves the desired result.
15

---

Published as a conference paper at ICLR 2025
B.3 PROOF OF APPROXIMATION
We now illustrate how the approximation relates to the expected loss landscape through the following
lemma:
Lemma (Approximation Error). Let (X) represent the loss and (X) represent the expected loss over the domain. Then, under the Brownian bridge loss approximation,
(X) (x) EBBB [[ fo (X.)ds
That is, the difference between expected loss and the real is given by the value of the Brownian bridge objective.
Proof.
Recall that Dynkin's formula states
(X, Xo - x] - + [
(X,j, Xo
(x)
BB
where the expectation is taken over Brownian motion sample paths X,- The left hand side is equal to the expected loss landscape E, the Brownian bridge loss minimizes the BB term.
C LOWER ORDER DERIVATIVES
As mentioned in the main text, we also have control over the parameters in the low order derivatives through an importance sampling technique. Here we will review how this is used modify the Monte Carlo scheme such that certain points are given more weight in the loss. For additional details, please see Øksendal (2003) and Pardoux and Raşcanu (2014). For convenience, will consider the stochastic process (Zt, yt).
C.I EXPONENTIAL MARTINGALE
Consider first the problem of the Laplace equation solved over domain D C [R++ and connection to Brownian motion. Specifically, write the solution of
22 = 0, 2 2 D = l 2 OD
according to the following expectation
u(=) : E [e[ Srop 20 = =]
where
Tap inf > 0 2 E ZD)
and dZ; = dw.
Now suppose we wish write the solution for the case where there exist first order derivatives with coefficients u(z): D + Ba+k corresponding to
u + 4 (2))) = 0, 2 D
(10)
Uy = l, z E 8D
which involves the following expectation
(2) = ( Zo Z
with Tap defined as before and
d p(Z)ddl + dW.
(11)
16

---

Published as a conference paper at ICLR 2025
Instead of sampling paths of equation 11, instead compute the expectation with the original Brownian motions paths weighted by the exponential martingale given by
dP, exx dQ
(f,-0 NO (Z.)d77
1 J
1 (((2)ds))
where Pa denotes the measure associated with the sample paths of equation 11 and Q is the Wiener measure. This then gives the solution to the PDE in equation 10
dpp Zo
Uy == EQ ((ZTOD)
2
dQ
Since want minimize Uy for boundary conditions dependent can apply Jensen's inequality to get rid of the exp term in the exponential martingale which may be subject numerical errors when approximated using an Buler scheme.
learned mapping f.
C.2 CHOOSING THE RIGHT
In our experiments, we choose to be (VL(J(x). v) for some & > 0 to get the following PDE
Vuu + +VV Vu, 0, 2E D.
In practice, one can always make some function (e.g. a neural network) and separately optimize achieve certain properties the loss landscape. We consider the gradient of due the connection between and the uncertainty around a new point. Points with large uncertainty are then given large weight based on the magnitude of the gradient and influence the loss more.
D REVIEW OF THE FEYNMAN-KAC FORMULA
The Feynman- -an formula provides correspondance between expectations of SDEs and PDEs. For full details, we refer to Øksendal (2003, Chapter 9) which describes the correspondence in the case boundary value problems as used this work. For self-containment, we provide a review of the formula here. Often, the Fenyman Kac formula is presented in its form for solving parabolic PDEs. Let X, satisfy the following Itô diffusion: d.2, = p(t, Xt) + XtWW, and consider the following PDE:
du
(1, (u((,)) . ((t, + Tr(aa T (1, (H()s, - r(t, x))(t,, Øt
(12)
with terminal condition u(x,T) g(x). Then the solution equation can be represented by the following expectation:
(tt, x. EE g(x))(x))))))))) exp (-((((x.jas) Xo = x
(13)
In the elliptic case, consider the following PDE solved over domain D:
1 ] Tr(ao (x) -(c)))))r r(z)u(a), x E D.
0 Vu(x) u())
(14)
with boundary condition
u(x) g(x) ÖD.
Then the solution to equation 14 can be represented by the following expectation:
g((x)) exp (-f- r(X.))l))
u(x)
Xo = x
(15)
where inf (8 X, JD).
17

---

Published as a conference paper at ICLR 2025
E VARIATIONS OF THE BRIDGE COMPARED TO THE ORIGINAL PDE
Here we describe a few more interpretations of the Brownian bridge augmentation as compared to the original PDE that we are solving. Recall that solving the full elliptic PDE by computing hitting times of Brownian motion is prohibitively expensive. We therefore need an algorithm that is more scalable for the typical problems in machine learning. We show now how the bridge can be interpreted through two different frameworks. To do this, we will rely on the following lemma which states that the hitting location of Brownian motion is most likely to be the location closest to the starting point for any stopping time.
Lemma 2 (Hitting Location of Brownian Motion). Let ((x)),()) be a set of boundary centroids ordered such that [[(x(0, y( (Xo: yo)l [1( ylitt) (Xo, yo)ll for all i 1...K - 1 and let B(i) {X, y [X X ( yll S e) for some e O. Then,
P(X, E BBO) > P(X, E B(tt))
for all i = 1.... 1.
Proof. Since X₂ ~ (N,, yo, t), the variance of X, increases with time. Then for any 7 > 0,
P(X, E B(0) & exp(- (| Xo + y - yoll/2r). Since J |X Xo + y( +e yol|de de <
50 [|(+1]) No + y(++)) yo + elide, the probabilities of being in B0()) > B(+1).
Lemma 2 allows us to sample Brownian bridges with endpoints mapped to the nearest starting points
rather than sample paths with random stopping times. Next we describe alternative implementations of the bridge algorithm in terms of Dynkin's 's theorem and an elliptic PDE with source term.
Next, we use a fixed for our implementation so we have a fixed time and thus a finite time algorithm for sampling paths for the expectation. Since the variance of the sample paths is given by σ, note that there is an ambiguity between a and т. This is easy to show for two Brownian motions W,, B with W,, r, (0, Ti) and B,, (N(0, then their distributions will be equal if a In We
then fix in our experiments and optimize over in the implementation.
E.1 INTEPRETATION VIA DYNKIN'S FORMULA
Dynkin's formula states the following:
Ele([(=)), Xo = x) (([(())) vo) + +E [ Ae(S(Xa)...dd Xo
which allows to rewrite the expectation of stopping time with respect to the integral of the generator A of X., ys applied to the function É. In our case, related to the PDE equation we are trying to impose. We can consider approximation of this operator through the following
1 AL x (X( -(U())))))))) + e(F(Xo), vo)
B(C(((),) x]
E
lim
Nr
E
(I(((((X--1) - - + e((((X)). vo).
These expectations can be computed over sample paths with fixed r. In general, 7 can be sampled as some function of the distance between the endpoints.
E.2 INTERPRETATION VIA AN ELLIPTIC PDE WITH SOURCE TERM
In the original formulation, all sources of variation arose from the boundary condition which was based the loss value for ball around each data point. We instead think of extending the boundary to be at infinity and only consider instead source terms given by the loss function at different points. This corresponds the following PDE:
( {(oo())))) (X,y) D
(16)
18

---

Published as a conference paper at ICLR 2025
Breast
Blood
Path
Organc
0.6
0.2
0.4
3.2
2014
1281
12/5
180
SOUAE 2112
12:
2141
0.0
0.5
0.
00
Class Index Sorted by Casss Size
Figure 4: Test data uncertainty score (min-max normalized) each class for different datasets. X-axis are class indexes sorted by data size; the errorbars present + standard deviation; the number above each bar the class size.
with u(X, ((fo(()).))) v) for (X. v) OD. This corresponds to the stochastic solution given by
u(x,)) = (((((()))) + [ -e(Jo(X.a).v)dd -
Xo X, yo
Using this representation, we can again sample Brownian bridges between endpoints but instead include an integral over the bridge points instead of computing the expectation over all points along the bridge path. For the new PDE in equation 16, u no longer harmonie; rather subharmonic. Therefore, only the maximum principle is satisfied, not the minimum principle. This is formalized in the following proposition:
Proposition 4. Using the interpretation in equation 16, the expected loss bounded from above by the loss the boundary of the domain.
Proof. In equation 16, u is subharmonic and the maximum principle applies.
E.3 MINIMIZING THE SOURCE TERM
Putting the two together, we can minimize the source term to develop a computationally effective way to satisfy the PDE. We make the substitution in equation 16 such that guarantee the maximum principle holds. We then apply Dynkin's 's formula with A(E(J(X.), v.)) = (()(),), v.) to get
[ E(J(X.) Xo
min BIC(s(x,),, Xo = x = min (((())), vo) + E
Minimizing this as described in the main text has the advantage being computationally efficient well as at the minimum corresponding to the original problem. We also see how the har- monic property is being violated by checking the value of E [S] e(((X,), y. ds)ss| Xo = I] At the minimum, (((())) ya) = and we will recover the original operator given by Al(f(X.). 8) = 0.
F ADDITIONAL EXPERIMENTS
In this section we present additional experiments including ablation studies on some parameters of elliptic training, further empirical results the robustness of the model trained with elliptic regularization, as well as additional benchmarks for the data imbalance/distrribttion shift regime against group-informed algorithms.
F.1 ABLATION STUDY
We provide ablation studies performance against artificial noise induced in test data. For the ablation studies, we mainly focus on the Waterbirds dataset.
hyperparameters used in elliptic training well the robust
The hyperparameters we focus on include & and Øg- Parameter & is related to the strength of the importance weighting, also as the magnitude of the drift derived from Girsanov's Theorem; ds is
19

---

Published as a conference paper at ICLR 2025
Table Ablation Study on Waterbirds
Avg(%)
Worst (%)
Baseline (Ç = 1, 0.1)
92.0+0.25
81.1
No Drift (C = 0, 0- = 0.1)
91.2+0.5
81.7+2.6
Moderate Drift (5 2, 0.1)
92.0+0.05
81.2+2.9
Large Drift (C = 10, 0.1)
91.2+0.6
82.5+0.66
Large Diffusion (5 1,0 1)
89.9
75.640.5
Small Diffusion (5 1, = 0.01)
91.1
81.1+18
Table 8: Noisy evaluation on Waterbirds.
Avg(%)
Worst(%)
Baseline
92.0+0.25
84.1
Noise Std = 0.1
79.843.0 86.0
92. 11.6
Noise Std = 0.5
71.7+11
the diffusion coefficient for data during sampling of Brownian bridges. We show in Table 7 that the importance weighting is necessary to improve model performance in group-imbalanced classification setting, and moderate diffusion coefficient should be selected to achieve the best possible result.
We also show that models trained with elliptic regularization are robust against noise. We first train model using elliptic regularization and clean data and then we add artificial Gaussian noise with different standard deviations to the test data. Table 8 showcases the robust result average and worst-group accuracy.
F.2 CONTROLLED COMPUTATION EXPENSE
In this section we reduce the number of epoch for Elliptic training such that the number function evaluation is similar to the benchmarks. Specifically
the MedMNIST experiments: both ERM and mixup is trained with 500 epochs, while Elliptic training is experimented with 5 bridge timesteps and 100 epochs. We show Table that Elliptic training continues outperform ERM and mixup with similar number of function evaluations.
F.3 COMPUTING PAIRWISE DISTANCE
In the main text, noted that the pairwise distance computation does not affect the model per- formance. Using the medMNIST datasets, we empirically demonstratet this in Table 10 where the model's robust performance is not affected. We suspect that the randomness in the batch leads to general diffusive behavior that enforces the elliptic operator.
Table 9; Average and worst class classification accuracy for selected Med-MNIST dataset with controlled number of function evaluations.
Breast
Blood
Path
OrganC
Algorithm
Avg(%)
Worst(%6
Avg(%)
Worst(%)
Avg(%)
Worst (6)
Avg(%)
Worst(
FRM
80.043.4
32.4413.8
81.422.7
56.3g17.3
55.743.8
10.7
80.7+2.9
52.946.9
mixup
83. 43.1
47.1
78.813.1
41.7+11.3
54.3
1.242.1
81.3
56.8411.0
Elliptic
86.7
67.644.9
81.5
63.4
57.8
12.8
85.840.6
64.64a.4
E IW
86.941.6
67.643.8
81.5
62.949.7
56.3
7.143.5
86.0+1.
60.4
20

---

Published as a conference paper at ICLR 2025
Table 10: Average and worst class classification accuracy for selected Med-MNIST dataset with pairwise distance computation (P-elliptic) and without (elliptic).
Breast
Blood
Path
OrganC
Algorithm
Avg(%)
(orst(%)
Avg(S))
Worst(%)
Avg(%))
Worst(%)
Avg(%)
Worst( %%)
Elliptic
87.641.8
67.6
1..5
14.9
62.3
26.4210.2
20.9
65.9
+IW
87.3
67.613.0
84.1
68.0
62.7
20.9
88.020.7
66.224.8
P-Elliptic
87.341.8
68. 1 +3.9
85.3
61.9
62.7+223
24.749.7
87.9+1.0
63.633.3
P+ IW
87.641.1
68.14a.2
84.3
67. 67.0
62.7+114
23.648.4
88. 88.2
45.5
Table 11: Robust classification accuracy under imbalance subpopulations and domain shift (grouped- informed).
Algorithm
WaterBirds
CelebA
Camelyon17
Avg %
Worst(%)
Avg(%)
Worst(%)
Avg(%)
IRM (Arjovsky et al., 2019)
87.5
75.6
94.010.4
77.813.9
64.2
IB-IRM (Ahuja et al., 2021)
88.5
76.5
93.610.3
85.0
68.916.1
V-Rex (Krueger et al., 2021)
88 88.0+1.0
73.640.2
92.2+0.1
86.7+1.0
71.544..
CORAL (Sun and Saenko, 2016)
90.3+1.1
79,8+1.8
93.840.3
76.943.6
59.5
GroupDRO (Sagawa et al., 2019)
91.840.3
90.6+1.1
92.1
87.2+1.6
68.447.3
Mix (Xu et al., 2020)
76.440.3
53,0+1.3
93.440.1
65.641.7
69.7
Fish (Shi et al., 2021)
85.6 85.640.4
64.0+0.3
93.1
61.2+2.5
74.7+7.1
LISA (Yao al., 2022b)
91.8 1.0.3
89.2
92.4
89.3
77.1 16.5
Elliptic IW
10.25
84.1
91.3
77.4
77.9
F.4 ADDITIONAL COMPARISONS
In this subsection compare Elliptic regularization training with importance weighting with algo- rithms that leverage the grbup-labels during training. This set of algorithms includes IRM (Arjovsky al., 2019), IB-IRM (Ahuja et al., 2021), V-Rex (Krueger al., 2021), CORAL (Sun and Saenko, 2016), GroupDRO (Sagawa et al., 2019), DomainMix (Xu et al., 2020), Fish (Shi et al., 2021), LISA (Yao et al., 2022b). We show in Table 11 that our one-stage, group-oblivious training algo- rithm is comparable to these group-informed algorithms in all three classification datasets. Elliptic regularization achieved the best average and third-best worst group accuracy for waterbirds, results comparable to IRM and CORAL in CelebA, and the best accuracy for Camelyon17.
G DATASET DETAILS
We briefly describe the contents and goals each dataset used in our experiments, provide details about data preprocessing, and reference related publically available codebases.
G.1 BALANCED/IN-DDISRIIUUTOO EXPERIMENTS
CIEARIO, CIFARI00, Timy-Im 200: For these well-known, small-scale image classification datasets, we followed preprocessing pipeline in Zou et al. (2023), which is publically available in the mixupE repository.
AirFoil, NO2: These are publically available tabular datasets with continuous labels for regression tasks. We applied preprocessing to these two datasets following Yao et al. (2022a). The preprocessing publically available in the c-mixup repository
The AirFoil (AirFoil Self-Noise) dataset (Brooks et al., 2014) aims to predict acoustic testing results given the physical features of two and three-dimensional airfoil and the wind tunnel environment.
21

---

Published as a conference paper at ICLR 2025
There are 1503 instances in AirFoil; min-max normalization is applied to the data; there are 1003, 300, and 200 data instances for training, validation, and testing respectively.
The NO2 dataset is included in Statlib (Kooperberg, 1997). The dataset studies air pollution and its relationship to traffic volume and metcorological variables. The dataset is collected by the Norwegian Public Roads Administration response variable column consists of the hourly logged concentration of NO2 particles. There is no preprocessing applied; we split 200, 200, and 100 for training, validation, and testing respectively.
Alnabru in Oslo between October 2001 and August 2003. The
G.2 GROUP IMABLANCE/DISTRIIUTIIO ROBUST EXPERIMENTS
G.2.1 CLASSIFICATION
Waterbirds, CelebA, and Camelyon: For these three datasets, we followed the preprocessing of Han et al. (2022). We applied the same data preprocessing, which is available in the UMIX repository.
The Waterbirds dataset (Koh al., 2021) aims classify whether the bird waterbird or landbird. This dataset has four predefined subpopulations including landbirds on land, landbirds on water, waterbirds on land, and waterbirds on water. Imbalanced group/subpopulation exist in the training set: the largest subpopulation is landbirds land with 3,498 samples, while the smallest subpopulation landbirrls on water with only 56 samples.
CelebA (Sagawa et al., 2019) is well-known large-scale face dataset. We predict the color of the human hair as blond or not blond. There are four imbalanced subpopulations based on gender and hair color including dark hair, female, dark hair, male, blond hair, female and blond hair, male with 71,629, 66,874, 22,880, and 1,387 training samples respectively.
Camelyon17 7 (Bandi et al., 2018) is pathological image dataset with over 450, 000 lymph node scans for predicting the existence of cancer tissue in a patch. The training data consists of scans from three hospitals, while the validation and test data are sampled from other hospitals. Distribution and domain shifts exist in this data as the classification requires generalization across different hospitals and coloring methods. Due the complexity of the data, especially considering that different coloring methods are observed even in samples from the same hospital, there are no reliable subpopulation labels of Camelyon17. We applied the official split scheme of this dataset.
G.2.2 REGRESSION
Crime consists of demographic and economic statistics of different communitics which For SkillCraft, RCF-MNIST is simulated dataset derived from Fashion-MNIST (Yao et al., 2022a).
SkillCraft, Crime: Both of these datasets are obtained from the UCI data repository. The prepro-
cessing applied in experiments is also publically available in the c-mixup repository
SkillCraft (SkillCraftI (aster) et al., 2013) contains video game telemetry data from real-time wal-time strategy (RTS) games to explore the development of expertise. The goal is predict input action latency based on 17 player-related parameters in the game, such as the variables and the Hotkey Usage variables. Missing data are filled by mean padding on each attribute. Levels of competitors, identified through "League Index" variable, are used as the domain information. We split 4,1.3 domains into training, validation, and testing subsets, which contain 1878, 806, 711 data instances, respectively. We train to predict the mean action latency of video game players in perception-actio action cycles; we treat "Leaguelndex" " as domain information and generalize predictions across different leagues.
Crime (Communities And Crimes)(Redmond, 2009) is a tabular dataset that aims to predict violent crimes per capita. The 122 attributes of this dataset combine socio-economic data from the 1990 US Census, law enforcement data from the 1990 US LEMAS survey, and crime data from the 1995 FBI UCR. Following the description of Yao et al. (2022a), we applied min-max to normalize all numeric features into [0.1]. The missing values are filled with the average values of the corresponding attributes. The 46 different State-IDs are used as the domain information, and we split the dataset into training, validation, and test sets into subsets containing 31, 6, and 9 disjoint domains. There are 1, 1,390, 231, and 373 instances in the training, validation, and testing subsets respectively. We train to
22

---

Published as a conference paper at ICLR 2025
predict the total number of violent crimes(per 100K population) of some states and aim to generalize to unseen states.
RCF-MNIST The RCF-MNIST (Rotated-Colore) on Fashion- Yao et al. (2022a). The main goal of this simulation is to understand the model performance under spurious correlation in the training dataset. The simulation for the training set contains two steps: for a specific data sample, 1) choose a normalized angle of rotation E [0. 1], this angle of rotation is used as the label; then 2) color the normalized RGB vector of the original white pixel into [] - 9, 0, .). The color and the angle create a spurious correlation. For testing data, the coloring is reversed; hence the testing data would have the coloring of [g. 0, g] for a chosen angle of rotation g. We used an 80-20 split to split the training and testing data. We train the model to predict the rotation angle of the object in the images.
Aishion-MNISTT is simulated dataset based
G.3 CLASS IMBALANCE MEDMNIST EXPERIMENTS
The Med-MNIST dataset consists of medical image datasets of different scales. The details of these datasets can be found in Yang et al. (2023a). We are using the official training and testing split of these datasets. All of the images for this set of experiments are normalized to mean = 0.5 and std 0.5. Imbalanced class exists in all of the selected datasets in our experiments including BreastMNIST, BloodMNIST, PathMNIST, and OrganCMNNIST.
BreastMNIST aims to use ultrasound images to identify the existence of malignant breast cancer. There are 780 images in total where 70-10-20 data split for training. validation, and testing is applied. Among the training data, there are 147 malignant labels and 399. An official data split of 70-10-20 is applied for each dataset to construct the training, validation, and testing subsets respectively.
BloodMNIST consists of microscopic images of blood cells. The 17092 images are obtained from individuals without infection, hematologic, or oncologic disease and free of any pharmacologic treatment during blood collection. The goal is to identify the 8 different cell types that exist in the dataset. An official data split of 70- 10-20 is applied for each dataset to construct the training, validation, and testing subsets respectively.
PathMNIST is collected from 100,000 non- overlapplgp image patches from hematoxylin & eosin stained histological images. The dataset aims to classify the 9 types of tissues. According to the description of the official data splits, the training and validation data is obtained from one clinical center, while the test data is curated from a different clinical center. Hence, test performance under class imbalance and avoid additional distribution shifts, we only used the official training split as training data and validation data as testing data in our experiments.
OrganCMNISS is derived from 3D computed tomography (CT) images from Liver Tumor Segmenta- tion Benchmark (LiTS). These 3D images are then preprocessed into different views (Yang et al., 2023a). The OrganC images consist of the coronal view of these 3D images. The goal is to classify the 11 body organs from these CT scans. The official data split is applied for this dataset. There are over 10,000 samples in this dataset.
G.4 PREPROCESSING
Our data preprocessing follows the preprocessing pipeline of previous works such as mixupE, UMIX, and c-mixup. For CIFARs, standard normalization applied: for tiny-imagenet, normalization, Ran- domCrop (to height=width=64), and RandomHHorzoottllFiip are applied; for celebA and Camelyon17. normalization and RandomHoorizonallFiip is applied; waterbirds' transformation includes normal- ization,
to crop the resolution to 224*224. The normalizations applied to above image datasets use the recommended means and standard deviation for the respective datasets. Min-max scaling is applied to all features of the regression datasets. We utilize the recommended preprocessing pipeline for MedMNIST datasets, which normalizes the data to 0.5 mean and 0.5 standard deviation.
and a RandomResizedCnnn
23

---

Published as a conference paper at ICLR 2025
Table 12: Architecture and Hyperparameter settings for each dataset/experiments. Optim stands for optimizer; Mom stands for momentum; WD stands for weight decay; and LR stands for learning rate. Note, * means the learning rate annealed by factor of on epoch 100 and 150.
Dataset
Epoch
Batch Size
Optim
Mom
WD
LR
71b
Tit
Tbb
CIFARTO/CIFAR 100
200
100
SGD
0.9
10-1
1-1#
1
0.05
Tiny-ImageNet 200
200
100
SGD
0.9
10-4
1-1#
5
0.05
AirFoil
100
Adam
1-2
20
0.05
NO2
100
32
Adam
1-2
10
0.01
Waterbirds
200
32
Adam
5
0.1
CelebA
50
128
SGD
0.9
0.1
1-4
0.1
Camelyon17
32
SGD
0.9
0.01
1-5
5
0.1
Crime
100
16
Adam
1-3
10
5
0.05
SkillCraft
100
32
Adam
1-2
10
0.05
RCF-MNIST
30
64
Adam
1
10
5
0.05
MedMNIST
500
500
Adam
1 3
10
H HYPERPARAMETERS
Now we describe the hypperaaranttersacchiitttuut information for each set of experiments. For all of our bridge samples, sampled according uniform discretization except for MedMNIST which we used time range [0, 0.01]. We use to denote the number of bridges, and n, to determine the number of discretization. We use db
the time range [0, 1).
All related training hyperparameters for all experiments are listed in Table 12. For all experiments, parameter {
H.1 ARCHITECTURES FOR BALANCED CLASSIFICATION
We followed mixupE's experiment format (Zou et al., 2023) and used PreActResNer50,, PreActRes- Net101, and Wide-ResNet-28 for benchmarking on PreActResNet34, PreActResNet500 for benchmarking on Tiny-imagene2200. These are implementa- tions without pre-trained weights.
We used 8reActResNet11,
H.2 ARCHITECTURES FOR REGRESSION
We applied a 3-layer neural network with hidden layer dimension=125 and the LeakyReLU activation with negative_slope-de which is the same as c-mixup for comparing the regression result (Yao et al., 2022a).
H.3 ARCHITECTURES FOR GROU--MMDLLLNGGDDSSTRRRTTTOOO SHIFT CLASSIFICATION
To be consistent with previous work Han al. (2022), for waterbirds and CelebA, trained Pytorch implementation of ResNet50 that is pre-trained on ImageNet; for Camelyon17, we trained on Pytorch implementation of DenseNet121 without pre-trained weights.
H.4 ARCHITECTURES FOR REGRESSION WITH DISTRIBUTION SHIFT
For Crime and SkillCraft, we used a 3-layer neural network with hidden layer dimension of 128 and the LeakyReLU activation function with negative ssoppesopppppeeee parameter of 0.I. For RCF-MNIST, we used ResNet18 pre-trained on ImageNet, but only up to the second to last layer, as a feature extractor. Then applied linear layer predict the angle.
24

---

Published as a conference paper at ICLR 2025
H.5 ARCHITECTURES FOR MEDMNIST IMBALANCE CLASSIFICATION
For all of the MedMNIST datasets, we used 2-layer neural network with hidden layer dimension of 512. Batch normalization is applied the output of the first layer.
25


Reviews:
[{"review_id": "qE7XQviVz7", "reviewer": "Reviewer_6cLH", "summary": "The paper proposes a new regularisation scheme, elliptic loss regularisation, based on optimising a pseudo loss surface which is the solution a Dirichlet boundary value problem. This regularisation scheme is designed to increase the robustness of models to data shift and imbalance. The authors perform experiments to verify some of their claims which shows that elliptic loss regularisation increase performance over other regularisation schemes like mixup in several settings.", "questions": "****1.**** In the experiments are standard data augmentations used in combination with elliptic loss regularisation and the baselines?\n\n****2.**** Is there a typo on line 166 as $u$ there is defined as a function with the first argument being $f(X)$ not $X$. \n\n****3.**** Eq 6. is confusing to me, specifically the term inside the second expectation seems to be constant over the second expectation, is this a typo or am I missing something?", "limitations": "", "rating": 3, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 2, "strengths": "* I believe the paper proposes quite an original way to perform loss regularisation by defining and basing the regularisation on a pseudo loss surface. I also think the way they define this loss surface as the solution to a boundary value problem is an interesting perspective. \n\n* I think that the paper covers a wide range of experimental settings, from data imbalance to distribution shift. Additionally, across most of these settings elliptic loss regularisation has good performance compared with the given baselines.", "weaknesses": "* I think the greatest weakness of the paper is the lack of clarity and motivation around the actually implemented approach. If this can be fixed I will likely increase my score. I have listed below a few specific points about clarity.\n    * The most significant part of this lack of clarity is how they define the boundary of their pseudo loss surface. I can find nowhere where it is precisely defined and from my understanding for a Dirichlet boundary value problem to be solvable and hence for their theorical motivation to be true the boundary needs to satisfy certain conditions. By reading the paper it seems to me that the authors either assume the boundary to be the set of training data points or the edge of the convex hull of the training points. However, if the boundary is just the training points, then it is discontinuous and to the best of my knowledge not really a boundary in the sense used in defining the Dirichlet boundary value problem. While if it is the edge of the convex hull of the training data then you do not generally have the value of the loss for the whole boundary and hence can’t use it.\n    * In section 5.1 they correctly point out that their approximation needs to be motivated by showing that it maintains the important properties of the true solution. However, I could not find anywhere a justification of why their approximation is accurate. This justification should be clearly presented in the paper to see how near the actual method is to the theory. \n    * There seems to be an overloading of the word ‘loss’, to make the paper clear I think there needs to be a clear distinction between the true loss and the pseudo loss surface created by the method. \n    * Another weakness is that a lot of the theory proved about the robustness of elliptic loss regularisation is only about the pseudo loss surface. However, except at the boundary, no connection is made between this loss surface and the true loss. This currently makes a large part of the theory lack practical use and I think a clear theoretical justification for why the pseudo loss surface will model the true one needs to be given. For example, at the moment, I could propose my own loss surface, perhaps a constant function with its value being the mean of the training loss, and prove similar propositions to propositions 1 to 3 while knowing that my loss surface is not similar to the true one at all. \n\n* While the predictive performance of a model using elliptic loss regularisation is shown to be good, I could not find mention of the computational cost of the method. It seems to me that the method could be quite computationally costly compared to other methods and hence I think a mention of the computational cost of elliptic loss regularisation should be given.\n\n* One final weakness is that the baselines compared to in the paper a relatively old and there are perhaps better ones to use. However, I am not sure what they are and hope that my other reviewers can point out relevant newer baselines, if they exist.", "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": "Ca0x8Z1suN", "reviewer": "Reviewer_tfwK", "summary": "This paper presents a novel regularizer for machine learning problems. The key insight is to construct a training objective (a “loss landscape”) that solves an elliptic PDE with boundary conditions specified by the unregularized loss’ values at the training points. This ensures that the regularized objective cannot be too large away from the training data: In particular, the loss landscape obeys a maximum principle, which implies that the regularized loss is upper-bounded by the maximum of the unregularized loss at the training points.\n\nAs it is impractical to exactly solve a PDE in each training iteration, the authors propose a numerical scheme to approximate the loss landscape. This scheme amounts to sampling pairs of training samples within each minibatch with importance weights inversely proportional to their distance, and then evaluating the training loss along paths connecting these pairs of points. The authors then prove two propositions to illustrate how the properties of the loss landscape can be used to predict the behavior of their regularized loss under affine distribution shift and data imbalance. They finally conduct a range of experiments which show that their regularizer performs comparably to baseline methods on problems without group imbalance or distribution shift, better than one-stage but often worse than two-stage baselines in settings with unbalanced populations and domain shift, and significantly better than baseline methods in class-imbalanced problems.", "questions": "- Why use Brownian bridges to connect pairs of training samples in the numerical procedure detailed in Sec 5.1? Is there any practical advantage to Brownian bridges compared to simply sampling along line segments connecting pairs of training samples (as with mixup)?\n- As I understand, the training algorithm essentially amounts to sampling pairs of nearby training points (distant pairs are presumably sampled with negligible probability), and evaluating the training loss along paths connecting these pairs. If this is correct, is it reasonable to view elliptic loss regularization as a variant of mixup where the interpolated points are more likely to remain on the data manifold? This seems like an intuitive perspective on why one might expect this method to perform well on manifold-supported data embedded in high-dimensional space, which is true of much real-world data under the manifold hypothesis.\n- How costly is elliptic loss regularization relative to simpler baselines like mixup? The proposed method’s performance improvements over baselines are sometimes marginal (see e.g. Table 5), and I wonder if the additional cost of the numerical scheme to approximate the regularized training objective is worth the marginal improvements in such problems.\n- Why are worst-case results not reported for the Camelyon17 dataset in Table 4?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- Elliptic loss regularization is a neat and well-motivated method for obtaining generalization guarantees via a maximum principle.\n- I appreciate the authors including examples of how one might use the properties of their loss landscape function to predict the effects of distribution shifts and class imbalances.\n- The experiments are thorough and adequately study the performance of elliptic loss regularization in a variety of plausible settings.\n- The paper is well-written and generally easy to follow.", "weaknesses": "My main critique of elliptic loss regularization is that it results in sometimes marginal performance improvements (such as in Table 5) over simpler and cheaper baselines. The authors could further strengthen the case for their regularizer by showing that using their numerical scheme, the computational cost of elliptic loss regularization is not much greater than that of simple mixup baselines.", "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": "LHHn7IKX0T", "reviewer": "Reviewer_CpbM", "summary": "The paper proposes a new regularizer based on elliptic loss landscapes. The high-level idea is to imbue the loss landscape of a model with a certain kind of smoothness that allows the loss to be bounded in unseen regions of the landscape. This allows the loss to be bounded under data distribution shifts and data imbalance. In practice, the regularizer works by training on the loss evaluated on paths between data points. Experiments show the experiment is effective on CIFAR, TinyImagenet and distribution shift tests.", "questions": "Can the authors bound the error of the sampling procedure?\nCan the authors explain why the performance of the method is worse in some-cases than other baselines (such as mixupE)?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "**Originality**\n\nThe paper seems quite original- I don't believe previous works have considered using elliptic loss landscapes and the resulting bounds one can derive (in Proposition 1) seem quite powerful. There are previous works exploring other forms of data augmentation (such as mixup) to improve generalization, but I believe few have the strong theoretical guarantees of the proposed method.\n\n**Quality**\n\nOverall, the quality of the paper is strong: the theory looks sound and experiments seem adequate.\n\n**Clarity**\n\nThe main text of the paper is adequately written and figures and tables are well-illustrated.\n\n**Signifiance**\n\nThe paper proposes a new way to bound the loss on unseen points far from any previously observed points, which is generally a very challenging problem. As such, I expect the theory proposed in this paper alone to have great significance to researchers finding ways of improving generalization, and I expect future work would build on the theory in this paper.", "weaknesses": "Regarding the theory, the practical algorithm requires sampling which would likely lead to a gap between the true elliptic loss and the computed loss. Can the authors comment on whether Proposition 1 would theoretically still hold under the sampling technique? Ideally, the error resulting from the sampling process can be bounded.\n\nAlso, unfortunately, it seems that the improvement of the method over previous techniques (particularly mixupE) seem small for CIFAR which raises questions about how effective it is at in-distribution generalization. Similarly, the results for domain shift in Table 4 and Table 5 do not seem conclusive in favoring the proposed regularizer.\n\nClarity-wise, the introduction of elliptic loss landscapes and the proposed algorithm seem a bit confusing due to the complexity of the technique. I would suggest adding an additional figure in the main paper illustrating the nice properties of the loss landscape described by eqns 1 and 2. I would also suggest including an algorithm in the main paper to accompany section 5.1", "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": "KzfklZFBHR", "reviewer": "Reviewer_rpmH", "summary": "This work proposes a regularization technique for a loss landscape where ideas are taken from Partial Differential Equations (PDEs) so that one can impose properties on the landscape. The loss landscape satisfies an elliptic operator allowing thus to bound it on the interior of the domain using the observed points. By doing so one can anticipate the behaviour of the error for points outside the traing set and thus addressing  two important problems : Data shifts and Data imbalance.\nIn this context, a tractable computational method that approximate the elliptic operator is proposed and various experiments are done trying to show its efficiency.", "questions": "1) How the PDE can be justified here ? Why not using OOD technique (for data imbalance) and domain adaptation (for data shift) ?\n2) there is a reference to the Dirichlet problem in the introduction but nothing is said about it in the rest of the paper.\n3) Why you did not compare to a recent Mixupe approach as Yingtian Zou (2023) (see the last reference).\n4) Experiments reflect Mean and Worst, why just these two metrics and not for instance the Best ?", "limitations": "", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "This paper addresses an important problem in machine learning : data shift and data imbalance", "weaknesses": "Although the addressed problem is important the proposed technique to tackle it is not clear nor justified.\nFirst the related work is poor and it's not related to what is proposed in the paper. Thus, after some enumeration of past work on mixup, it's just stated that \"these findings motivate our work where we consider a different type of regularization and its implication on the loss values\". This should be motivated and explained.  \nThen, ERM vs TRM is introduced by saying that the latter may cover more of D depending on how \\phi is chosen, but nothing is said on the price to pay for that.\nThe equation (1) (line 171) is not explained nor justified and as  this equation is the core of the proposed approach, it would be beneficial to explain its derivation\nHow the equation (1) leads to equation (3) ? the derivation is not clear and there is no details in the appendix\nAfter that, the generalization of the operator through equations (4) and (5) is not clear and the Appendix does not help for that", "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)