Elliptic Loss Regularization

Thank you for your constructive feedback and efforts in writing the review. We hope that we adequately address all the concerns in the review, please let us know if there is anything else that requires discussion.

Related Work

The closest augmentation based method to the proposed method is mixup, especially its implication of solving a regularized problem through augmentation as pointed out in the related work. Using TRM to solve a regularization problem motivates us to develop our method of using diffusion to solve an optimization problem regularized by a PDE. We have updated the related work accordingly to better demonstrate our motivation stemming from TRM and mixup.

Theoretical Development

The motivation for constraining the regularization with equation (1) is its theoretical properties, as stated in the introduction (additionally see our response below). The theoretical development from equation 1 to equation 3 follows the application of Feynman-Kac theorem, then we applied the Girsanov theorem to obtain equations 4 and 5.

PDE Justification

The PDE is justified for two reasons: 1) the theoretical properties described throughout the paper that can be used for the problems we're interested in tackling and 2) the computationally scalable approach that we use to implement the method. As such, we provide theoretical evidence for how the PDE perspective provides a unified approach to addressing the problems of distribution shift and imbalance. We further show that the PDE approach provides some empirical improvement when compared to other imbalance/shift approaches as the reviewer mentions.

Dirichlet Problem

We included a reference to a Dirichlet problem at the beginning to bridge the gap to traditional PDE disciplines. In the revision, we removed the reference to a Dirichlet problem to avoid confusion.

Comparison to MixupE

We did include a comparison to the MixupE approach (please see Tables 1, 2, and 6).

Best Group

For data imbalance experiments, we are interested in quantifying worst group performance which is usually the one that has the fewest samples as average performance across all groups. We want to balance the average performance to the worst case performance. Best group performance may not be as meaningful since, for example, a method can overfit to a particular group. Finally, these are the metrics commonly reported within the literature for assessing performance for group imbalance (see the references for comparable tables).

Thank you for all your time and effort in reviewing and providing commentary on the manuscript. Below we respond to all of the concerns, please let us know if there's anything else that we should address.

Proposition 1 with Approximation

This is a good point, the reviewer is correct there is some error associated with sampling with a discrete scheme such as Euler-Maruyama. The error with Euler-Maruyama can be reduced by decreasing the step size, we included a reference to Chapter 7.2 of [1] in the revision. This is going to be the main source of error introduced by the numerical scheme.

The other consideration is the difference between Brownian bridges which we sample versus unconstrained sample paths. In the case of the Laplace equation, sampling bridges with endpoints according to a distribution based on minimum distance is exact (which we discuss in Lemma 1 in appendix E). While sampling should be according to the pairwise distance, sampling without the pairwise distance would induce a first order term but would not remove any of the theoretical properties of the proposition.

Finally, in terms of training behavior, the minimum bound in proposition 1 could be violated if the loss is very large for points along the bridge. This would mean that $\ell(f(x),y)$ is not harmonic but is subharmonic, meaning the minimum principle would not hold, only the maximum principle (i.e. the lower bound does not exist but the upper bound does). At the minimum when $\ell(f(x),y) =0$, though, approximation guarantees that proposition 1 holds.

Adding Figure and Algorithm

We thank the reviewer for these suggestions. Due to page limitations, the proposed algorithm outlined in section 5.1 is presented in Appendix A. We updated Figure 1 to illustrate the properties of a loss landscape that satisfies eq (1).

Error of Sampling Procedure

The error of the sampling procedure is based on the discretization of the sample paths. We updated the manuscript to include a discussion on this as well as the reference to [1, Theorem 7.15] which describes the optimal convergence rate for these types of equations.

Performance of Different Experiments

The reviewer brings up an interesting question. We first note that our method is universal but is comparable to the methods built specifically for particular settings. We hypothesize that the experiments with a more complex dataset (i.e a larger number of classes, a larger number of subpopulations and groups) are more effective with the regularization since datasets with more classes/groups have a more complex domain that can be better estimated with elliptic regularizers as the data-augmentation spend more of its time in the interior of the domain.

[1] Graham, Carl, and Denis Talay. Stochastic simulation and Monte Carlo methods: mathematical foundations of stochastic simulation. Vol. 68. Springer Science & Business Media, 2013.

Thank you for your time and efforts in providing comments and feedback on the manuscript. We believe we address all the concerns mentioned in the review below and updated the manuscript accordingly, please let us know if there's anything else that requires clarification.

Boundary Definition

We consider the boundary to be the union of the boundary of the convex hull of the training data (to ensure a compact set) and the balls surrounding each data point as described in section 4.1. The boundary can be thought of as a union of disjoint boundaries with each boundary being the boundary of the ball centered at each data point, which is simply defining multiple boundary conditions. This is done in a variety of settings, one of the most common applications is in inverse scattering [1]. For points on the boundary of the convex hull, a value of the $\ell(f_\theta(X_\tau), y_t)$ does exist and can be computed, though these points do not lie within the training data. Note that even if we do not include the boundary of the convex-hull, Dirichlet problems can be defined and are studied in the literature as exterior Dirichlet problems [2]. We rewrote section 4.1 to make this clearer.

[1] Hariharan, S. I. (1982). Inverse scattering for an exterior Dirichlet problem. Quarterly of Applied Mathematics, 40(3), 273-286.

[2] Meyers, N., & Serrin, J. (1960). The exterior Dirichlet problem for second order elliptic partial differential equations. Journal of Mathematics and Mechanics, 513-538.

Properties of the Approximation

In Lemma 1 in Appendix E we describe the relationship between the Brownian bridges and the probable paths for computing first hitting times of the boundary. We then use these paths, since they can be easily sampled within a finite amount of time, to apply Dynkin's formula to satisfy the PDE. Recalling Dynkin's formula $\mathbb{E}[\ell(f(x_\tau), y_\tau) \mid x_0,y_0 = x,y] = \ell(f(x_0),y_0) + \mathbb{E} \left [\int_0^\tau \mathcal{A} \ell(f(x_s),y_s) d s \mid x_0,y_0 = x,y \right]$ we compute the expectation over these bridge paths instead of the unconstrained sample paths which have unknown sampling time where $\mathcal{A}$ is the generator of the stochastic process $x_t,y_t$. Recall that the operator we're imposing is $\mathcal{A} \ell(f(x),y) = 0$, but we want to avoid computing $\mathcal{A}$ directly each iteration since it is costly. Instead, we solve the problem $\mathcal{A} \ell(f(x),y) = \ell(f(x),y)$ which requires integrating $\ell(f(x),y)$ over the sample paths. Then, since we're minimizing $\ell(f(x),y) \geq 0$, at the minimum of $\ell(f(x),y) = 0$ we obtain the desired relationship $\mathcal{A} \ell(f(x),y) = 0$. Therefore at minimum, the properties all hold and the expected loss landscape is equal to the true loss. Note that if the value is greater than zero then $\ell(\cdot)$ is subharmonic meaning the maximum principle will still hold (but the minimum principle will not). However, at minimum it is then equal to the harmonic solution.

Overloading the Word "Loss"

Apologies for this, we opted to use the term loss landscape for the object that we perform our analysis on and describe it as an expected error. We are open to any suggestions for improving the terminology.

Connection Between Expected Loss and True Loss

This follows the response given above. The two are equal when $\mathcal{A} \ell(f(x),y) =0$. Using Dynkin's formula again $E[\ell(f(X_\tau), y_\tau) \mid X_0 = x, y_0= y] = \ell(f(x),y) + E[\int_0^\tau \mathcal{A} \ell(f(X_s), y_s) \mathrm{d}s \mid X_0 = x, y_0= y]$ where $\mathcal{A}$ is the generator of the stochastic process. We're interested in comparing $\ell(f(x),y)$ to $E[\ell(f(X_\tau), y_\tau) \mid X_0 = x, y_0= y]$. The difference is given by $E[\int_0^\tau \mathcal{A} \ell(f(X_s), y_s) \mathrm{d}s \mid X_0 = x, y_0= y]$, which, when $\ell(f(X_s), y_s)$ satisfies the PDE, is zero. Then the difference between the loss landscape and the true loss is then given by how well the function $\ell(f(X_s),y_s)$ satisfies PDE. By the substitution described in the previous point, the better trained the model is, the closer its adherence to satisfying the equation. We included this note within the revised manuscript.

Thank you for your time and efforts in reading the manuscript and providing helpful feedback. We appreciate the encouragement, and we believe we address all the related concerns below. Please let us know if there is anything that requires additional clarification.

Why Brownian Bridges

The reviewer brings up a good point. The construction of the regularization scheme is derived from the PDE using the connection between the PDE and stochastic processes via the Feynman-Kac theorem, which requires sampling of diffusion paths to the boundary condition. Sampling only on the linear interpolation between two data is not sufficient to compute the expectation in Eq (3) since it restricts the samples to only lie on the line between any two samples rather than being absolutely continuous with respect to the Lebesgue measure. This would look like convective behavior rather than diffusive, which would have different properties, but is an interesting question nevertheless.

Connection to mixup

The reviewer is correct that the proposed method is related to the mixup algorithm in the sense that both our method and mixup apply data transformation to solve a regularized problem. Elliptic regularization, however, solves a different regularized problem compared to the mixup algorithm motivated by the viewpoint of the PDE.

The reviewer also raises an interesting point regarding the sampling scheme. The sampling scheme was set up to approximate the PDE in expectation. While the distance based pairing may induce some of the properties that the reviewer mentions regarding the manifold supported data, it is likely the case that the overall diffusive dynamics are contributing to the method's performance since including pairwise calculations has a small effect on the performance. However, trying this on explicitly manifold supported data (e.g. graphs/solving on irregular domains) is a great avenue for future work to test this idea. We updated the manuscript to include this.

Cost of Imposing Regularization

The most computationally expensive part of our method is computing pairwise distance for every dataset, which is of time-complexity $\mathcal{O}(n^2)$. However, this can be pre-computed and saved as a dictionary to be called for $\mathcal{O}(1)$ speed during training.

As stated on line 370, none of our results apply this sampling technique because the improvements are negligible (see Appendix F). If we do not compute distances, the bridge paths correspond to drifted Brownian motion which induces a first-order term on the operator. This has the same effect as the change of measure described in the importance weighting section.

Therefore, whether computing the pairwise distance beforehand or not computing the pairwise distance at all, will result in an algorithm that has marginal computational overhead compared to the other benchmarks and the additional cost comes from the additional function evaluation controlled by the number of bridges and the number of time samples for each Brownian bridge. We experiment with controlled function evaluation where we train fewer epochs for elliptic regularization compared to other benchmarks. These results are in Appendix F. We also note that our method trains faster than the SOTA baseline UMIX, and in general has the same number of function evaluations as UMIX and JTT.

Finally, we compared epoch-wise training speed between elliptic regularization and mixupE and the results are hardware dependent. On an NVIDIA RTX 6000, mixupE trains 3 minutes and 12 seconds per epoch while elliptic regularization trains 3 minutes and 23 seconds per epoch. On an NVIDIA RTX 3090, mixupE trains 4 minutes and 51 seconds per epoch while elliptic regularization trains 4 minutes and 20 seconds per epoch. These experiments are conducted by training CIFAR10 with PreActResNet101 architecture; the reported training time is averaged across 10 epochs. Every hyperparameter is the same as those reported in the appendix.

Camelyon Table

The worst group is not shown for the Camelyon dataset because the groups are undefined in this dataset due to datashifts occuring within each group but no labels are associated with these shifts. We include details regarding the unknown subpopulation are provided in Appendix G.2.1.