Elliptic Loss Regularization

TL;DR

This work introduces a PDE-based elliptic regularization for neural network loss landscapes, enforcing that the loss function $u(X,y)$ satisfies an elliptic PDE with boundary data provided by training losses. The approach leverages a stochastic representation via the Feynman-Kac formula and Brownian bridges to compute an interior loss that generalizes beyond observed data, with optional drift reweighting to address imbalance. Theoretical insights include the maximum principle bounding interior loss by boundary values and diffusion-driven qualitative behavior under distribution shifts and data sparsity; the method also supports data-shift and imbalance-aware extensions through Radon–Nikodym reweighting. Empirically, elliptic regularization achieves competitive performance with Mixup variants on balanced tasks and demonstrates robust improvements under distribution shift and group imbalance across several datasets, indicating practical impact for robust learning in real-world deployments.

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, 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 confirm the utility of the proposed regularization technique.

Figures (4)

• 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 a 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.
• Figure 2: Illustration of the loss values over a domain with $4$ points on the boundary. The expected loss at point $X^\star$ is composed of losses at $\varepsilon-$balls around $X^{(i)}, i=1\ldots4.$ Black paths represent sample paths starting at $X^\star.$
• 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.
• Figure 4: Test data uncertainty score (min-max normalized) of each class for different datasets. X-axis are class indexes sorted by data size; the errorbars present $\pm$ standard deviation; the number above each bar is the class size.

Theorems & Definitions (12)

• Proposition 1: Bounding the Loss for an Interior Point
• proof
• Proposition 2: Expected Error Under Affine Transformations
• Proposition 3: Expected Error in Regions of Low Density
• proof
• proof
• Lemma 1: Approximation Error
• proof
• Lemma 2: Hitting Location of Brownian Motion
• proof