Provably Learning Diverse Features in Multi-View Data with Midpoint Mixup
Authors: Muthu Chidambaram
Xiang Wang
Chenwei Wu
Rong Ge
Venue: Submitted to ICLR 2023 (2023)
Keywords: mixup
data augmentation
theory
optimization
multi-view
feature learning
generalization
deep learning

Abstract: Mixup is a data augmentation technique that relies on training using random convex combinations of data points and their labels. In recent years, Mixup has become a standard primitive used in the training of state-of-the-art image classification models due to its demonstrated benefits over empirical risk minimization with regards to generalization and robustness. In this work, we try to explain some of this success from a feature learning perspective. We focus our attention on classification problems in which each class may have multiple associated features (or views) that can be used to predict the class correctly. Our main theoretical results demonstrate that, for a non-trivial class of data distributions with two features per class, training a 2-layer convolutional network using empirical risk minimization can lead to learning only one feature for almost all classes while training with a specific instantiation of Mixup succeeds in learning both features for every class. We also show empirically that these theoretical insights extend to the practical settings of image benchmarks modified to have additional synthetic features.

Under review as a conference paper at ICLR 2023
PROVABLY LEARNING DIVERSE FEATURES IN MULTI- VIEW DATA WITH MIDPOINT MIXUP
Anonymous authors Paper under double-blind review
ABSTRACT
Mixup is data augmentation technique that relies training using random convex combinations of data points and their labels. In recent years, Mixup has become standard primitive used in the training of state-of-the-art image classification models due to its demonstrated benefits over empirical risk minimization with regards to generalization and robustness. In this work, we try to explain some of this success from feature learning perspective. We focus our attention on classification problems which each class may have multiple associated features (or views) that can be used to predict the class correctly. Our main theoretical results demonstrate that, for non-trivial class of data distributions with two features per class, training 2-layer convolutional network using empirical risk minimization can lead to learning only one feature for almost all classes while training with a specific instantiation of Mixup succeeds in learning both features for every class. We also show empirically that these theoretical insights extend to the practical settings of image benchmarks modified to have additional synthetic features.
INTRODUCTION
Data augmentation techniques have been a mainstay in the training of state-of-the-art models for wide array of tasks particularly in the field of computer vision due their ability to artificially inflate dataset size and encourage model robustness to various transformations of the data, One such technique that has achieved widespread use is Mixup (Zhang al., 2018), which constructs new data points as convex combinations of pairs of data points and their labels from the original dataset. Mixup has been shown to empirically improve generalization and robustness when compared to standard training over different model architectures, tasks, and domains (Liang al., 2018; He et al., 2019; Thulasidasan et al., 2019; Lamb et al., 2019; Arazo et al., 2019; Guo, 2020; Verma et al., 2021b; Wang al., 2021). It has also found applications to distributed private learning (Huang et al., 2021), learning fair models (Chuang and Mroueh, 2021), semi-supervised learning (Berthelot et al., 2019b; Sohn et al., 2020; Berthelot et al., 2019a), self-supervised (specifically contrastive) learning (Verma et al., 2021a; Lee et al., 2020; Kalantidis et al., 2020), and multi-modal learning (So et al., 2022).
The success of Mixup has instigated several works attempting to theoretically characterize its potential benefits and drawbacks (Guo et al.. 2019; Carratino et al., 2020; Zhang et al., 2020; 2021; high-level, the
Chidambaram et al., 2021). These works have focused mainly on analyzing, beneficial (or detrimental) behaviors encouraged by the Mixup-version of the original empirical loss for a given task.
As such, none of these previous works (to the best of our knowledge) have provided an algorithmic analysis of Mixup training in the context of non-linear models (i.e. neural networks), which is the main use case of Mixup. In this paper, we begin this line of work by theoretically separating the full training dynamics of Mixup (with a specific set of hyperparameters) from empirical risk minimization (ERM) for a 2-layer convolutional network (CNN) architecture class of data distributions exhibiting multi-view nature. This multi-view property essentially requires (assuming classification data) that each class in the data is well-correlated with multiple features present in the data.
Our analysis is heavily motivated by the recent work of Allen-Zha and Li (2021), which showed that this kind of multi-view data can provide a fruitful setting for theoretically understanding the benefits of ensembles and knowledge distillation in the training of deep learning models. We show that Mixup

---

Under review as a conference paper at ICLR 2023
can, perhaps surprisingly, capture some of the key benefits of ensembles explained by Allen- and Li (2021) despite only being used to train a single model.
Main Contributions and Outline. Our main contributions three-fold. In Sections 2 and 3, we introduce the main ideas behind Mixup and analyze a simple, linearly separable multi-view data distribution which we to lay the groundwork for our main results. In analyzing this distribution, we motivate the use of a particular setting of Mixup- which we refer to as Midpoint Mixup - in which training is done on the midpoints of data points and their labels.
Section 4 contains our main results; we prove that, for a highly noisy class of data distributions with Two features per class, minimizing the empirical cross-entropy using gradient descent can lead to learning only one of the features in the data while minimizing the Midpoint Mixup cross-entropy succeeds in learning both features. While our theory focuses on the case of two features/views per class to be consistent with Allen-Zhu and Li (2021), our techniques can readily be extended to more general multi-view data distributions.
Lastly, in Section 5, we conduct experiments illustrating that our theoretical insights in Sections 3 and 4 can apply to the training of realistic models on image classification benchmarks. We show for each benchmark that, after modifying the training data to include additional spurious features correlated with the true labels, both Mixup (with standard settings) and Midpoint Mixup outperform ERM on the original test data, with Midpoint Mixup closely approximating the performance of regular Mixup.
Related Work. The idea of training on midpoints (or approximate midpoints) is not new; both Guo (2021) and Chidambaram et al. (2021) empirically study settings resembling what we consider in this paper, but they do not develop theory for this kind of training (beyond an information theoretic result in the latter case). As mentioned earlier. there are also several theoretical works analyzing the Mixup formulation and it variants (Carratino et al., 2020; Zhang et al., 2020; 2021; Chidambaram et al., 2021; Park et al., 2022), but none of these works contain optimization results (which are the focus of this work). Additionally, we note that there are many Mixup-like data augmentation techniques and training formulations that are not (immediately) within the scope of the theory developed in this paper. For example, Cut Mix (Yun et al., 2019), Manifold Mixup (Verma et al., 2019), Puzzle Mix (Kim et al., 2020), Co-Mixup (Kim et al., 2021), and Noisy Feature Mixup (Lim et al., 2021) are all such variations.
Our work also influenced by the existing large body of work theoretically analyzing the benefits of data augmentation (Bishop. 1995; Dao et al., 2019; Wu et al., 2020; Hanin and Sun, 2021: Rajput et al., 2019; Yang et al., 2022; Wang et al., 2022; Chen et al., 2020; Mei et al., 2021). The most relevant such work to ours is the recent work of Shen et al. (2022), which also studies the impact of data augmentation on the learning dynamics of a 2-layer network in a setting motivated by that of Allen-Zhu and Li (2021). However, Midpoint Mixup differs significantly from the data augmentation scheme considered in Shen et al. (2022), and consequently our results and setting are also of a different nature (we stick much more closely to the setting of Allen-Zhu and Li (2021)). As such, our work can be viewed as a parallel thread to that of Shen et al. (2022).
2 PRELIMINARIES AND MOTIVATION FOR MIDPOINT MIXUP
We will introduce Mixup in the context of k-class classification, although the definitions below easily extend to regression. As a notational convenience, we will use [k] to indicate {1, (1,2,....k, k...
Recall that, given finite dataset X C Rdd x [k] with |X| cross-entropy loss [[g X)) of a model g R2d -> T as:
N, we can define the empirical
J(g, X) =
1
log dith (g(x:)) where SE[N]
Σ
(x)( (9(x()) E.Ej (x)(
(2.1)
N
With e being the standard softmax function and the notation 9", 6Z indicating the y-th coordinate functions of g and & respectively. Now let us fix distribution D2 whose support is contained in [o, 1] and introduce the notation 2A,(() = Az; + (1 - A)ι, (using zi,j when A is clear from context) where (I,, vi), (Iz). Mi) E X. Then we may define the Mixup cross-entropy JM (g. X. Dx) as:
JM X, DA) = N2
1
Σ Σ
Ex-DA [ log (g(2,))) + (1 A) log の. (g(2+,)) (2.2)
JE[M
2

---

Under review as a conference paper at ICLR 2023
We mention minor differences between Equation 2.2 and the original formulation of Zhang et al. (2018). Zhang et al. (2018) consider the expectation term in Equation 2.2 over N randomly sampled pairs of points from the original dataset X, whereas explicitly consider mixing all N2 possible pairs of points. This is, however, just to make various parts of our analysis easier to follow one could also sample N mixed points uniformly, and the analysis would still carry through with an additional high probability qualifier (the important aspect is the proportions with which different mixed points show up: i.e. mixing across classes versus mixing within a class).
MOTIVATING MIDPOINT MIXUP: THE LINEAR REGIME
As can be seen from Equation 2.2, the Mixup cross-entropy Jas (g. X, Dx) depends heavily on the
choice of mixing distribution x-- Zhang et al. (2018) took DA to be Bcta(a, a) with a being a hyperparameter. In this work, we will specifically be interested in the case of - oc, for which the distribution DA takes the value 1/2 with probability 1. We refer to this special case as Midpoint Mixup, and note that it can also be viewed Guo (2021). We will write the Midpoint Mixup loss as JMM (g. X)) (here 2 = (xt + 22)/2 and there is no D> dependence as it is deterministic):
case of the Pairwise Label Smoothing strategy introduced by
JMM(g, = 2N
Σ Σ [E[N]
(log 6th (g(2a,)) + log の (g(2+,)))
(3.1)
We focus on this version of Mixup for a few key reasons. Firstly, we will show that Jaras (g, X) exhibits the nice property that its global minimizer corresponds to a model in which all of the features in the data are learned equally (in a sense to be made precise below). We will also show that this is not the case for JM (g, X, Dx) when Dx is any other non-trivial distribution. Additionally, from technical perspective, the Midpoint Mixup loss lends itself to a much cleaner optimization analysis due to the fact that the structure of its gradients is not changing with each optimization iteration (i.c. we do not need to sample new mixing proportions at each step). This allows us to more casily show how the gradient descent dynamics encourage learning all of the features in the data. That being said, we are not trying to claim that Midpoint Mixup is a superior practical alternative to standard Mixup our goal is simply show that better accentuates the theoretical benefits of Mixup, and is empirically comparable to standard Mixup settings. Full proof' for all of the results presented in the next subsection can be found in Section of the Appendix.
3.1 MIDPOINT MIXUP WITH LINEAR MODELS ON LINEARLY SEPARABLE DATA
To make clear what we mean by feature learning, we first turn our attention to the simple setting of learning linear models (x) (w 2² (i.e. one weight vector associated per class) on linearly separable data, as this setting will serve as a foundation for our main results. Namely, we consider k-class classification with a dataset of N labeled data points generated according to the following data distribution (with N sufficiently large):
Definition 3.1. [Simple Multi-View Setting] For each class y E [kt], let Ug Uy.2 E Rdd be orthonor- mal unit vectors also satisfying Uy.E 1 Va,e' when y j 8 for any e, l' E [2]. Each point (a, v) ~ D is
then generated by sampling y E [k] uniformly and constructing x as:
I By + (1-B,)),,2
By Uni (0.1,
(3.2)
Definition 3.1 is multi-view in the following sense: for any class y. it suffices (from an accuracy perspective) to learn a model g that has a significant correlation with either the feature vector 1y or 1y In this context, one can think of feature learning as corresponding to how positively correlated the weight wy is with each of the same class fcature vectors Uy and ug.1 (we provide a more rigorous definition in our main results).
If one now considers the empirical cross- cross-entropy loss (g, X). it is straightforward to see that is possible to achieve the global minimum of J(g X²) by just considering models g in which we take (Wy, 2g + oo for every class y. This means we can minimize the usual cross-entropy loss without learning both features in the dataset X.
However, this is not the case for Midpoint Mixup. Indeed, we show below that a necessary (with extremely high probability) and sufficient condition for linear model g to minimize Jar (when
3

---

Under review as a conference paper at ICLR 2023
taking its scaling to oc) is that it has equal correlation with both features for every class (sufficieney relies also on having weaker correlations with other class features). In what follows, we use inf JM ar (h, X) to indicate the global minimum of JMM over all functions h : R - R (i.e. this is the smallest achievable loss).
Lemma 3.2. [Midpoint Mixup Optimal Direction] A linear model satisfies the following:
lim JMM(19, X)= = inf JM 7-9000
(3.3)
If has the property that for every class y we have (Wy, Vy,&,) = (w,, v,,(a) >> 0 and (wy, ve,e,) S 0 for every s * y and l1, lg E [2]. Furthermore, with probability - exp(- randomness of 2), the condition (wy, (2,01) (Wa, (s,02) is necessary for g satisfy Equation 3.3.
(over the
Proof Sketch, The idea is that if has equal correlation with both features for every class, its predictions will be constant on the original data points due to the fact that the coefficients for each feature in each data point are mirrored per Equation 3.2. With the condition (wy, Ca,e) S 0 (this can be weakened significantly), this implies the softmax output of on the Midpoint Mixup points will be exactly 1/2 for each of the classes being mixed (and 0 for all other classes), which optimal.
As mentioned earlier, we can also show that if we consider JAr (g, X, Dx) for any other non-point- mass distribution, we can prove that the analogue of Lemma 3.2 does not hold true.
Proposition 3.3. For any distribution DA that not point mass on 0, 1, or 1/2, and any linear model satisfying the conditions of Lemma 3.2, we have that with probability 1 - (NN)) (over the randomness of X') there exists an 60 > 0 depending only on DA such that:
Jas (g. X, Dx) 2 inf JM(h Dx) 6
(3.4)
Proof Sketch. In the case of general mixing distributions, we cannot achieve the Mixup optimal behavior of 6hth (g(2z,) (A))) = X for every X the outputs g" are constant on the original data points.
Lemma 3.2 outlines the key theoretical benefit of Midpoint Mixup namely that its global optimizers exist within the class of models that we consider, and such optimizers learn all features in the data equally. And although Lemma 3.2 is stated in the context of linear models, the result naturally carries through to when we consider two-layer neural networks the type we define in the next section. That being said, the interpretation of Proposition 3.3 is not intended to disqualify the possibility that the minimizer of (M. X, Dx) when restricted to a specific model class is model in which all features are learned near-equally (we expect this to be the case in fact for any reasonable Dz). Proposition 3.3 is moreso intended to motivate the study of Midpoint Mixup as a particularly interesting choice of the mixing distribution D)-
We now proceed one step further from the above results and show that the feature learning benefit of Midpoint Mixup manifests itself even in the optimization process (when using gradient-based methods). We show that, if significant separation between feature correlations exists, the Midpoint Mixup gradients correct the separation. For simplicity, we suppose WLOG that (wy, Ug,1) > (Wyg 2y.2)) Now letting A, = (wy, Up Uy.2) and using the notation Vay for Bwy," we can prove:
Proposition 3.4. [Mixup Gradient Lower Bound] Let y be any class such that A, 2 log k, and suppose that both (wy, 1y > 2 0 and the cross-class orthogonality condition (Wa, Vate) = 0 hold for all # u and & [2]. Then we have with high probability that:
1
w.y JM.M(, 22
θ
(3.5)
A-2
Proof Sketch. The key idea is to amalyze the gradient correlation with the direction Uy.1 1y2 via concentration of measure argument.
Proposition 3.4 shows that, assuming nonnegativity of within-class correlations and an orthogonality condition across classes (which we will show to be approximately true in our main results), the feature correlation that is lagging behind for any class y will receive significant gradient when optimizing the Midpoint Mixup loss. On the other hand, we can also prove that this need not be the case for empirical risk minimization:
4

---

Under review as a conference paper at ICLR 2023
Proposition 3.5. [ERM Gradient Upper Bound] For every y [k]. assuming the same conditions as in Proposition 3.4, if Ay 2 C log k for any C >> 0 then with high probability we have that:
V wyy Jg. x), (222) s
(3.6)
R..IC-1
Proof Sketch. This follows directly from the form of the gradient for [gg 2).
While Proposition 3.5 demonstrates that training using ERM can possibly fail to learn both features associated with class due increasingly small gradients, one can verify that this does not naturally occur in the optimization dynamics of linear models on linearly separable data of the type in Definition 3.1 (see for example, the related result in Chidambaram et al. (2021)). On the other hand, if we move away from linearly separable data and linear models to more realistic settings, the situation described above does indeed show up, which motivates our main results.
ANALYZING MIDPOINT MIXUP TRAINING DYNAMICS ON GENERAL MULTI-VIEW' DATA
For our main results, we now consider a data distribution and class of models that are meant to more closely mimic practical situations.
4.1 GENERAL MULTI- VIEW DATA SETUP
We adopt slightly simplified version of the setting of Allen-Zhr and Li (2021). We still consider the problem of k-class classification on a dataset X' of labeled data points, but our data points are now represented as ordered tuples x = (x(1), (P)) of P input patches (1)) with each x(() E Rd (so X C R Pd x (..).
As was the case in Definition 3.1 and in Allen Zhu and Li (2021), we assume that the data is multi- view in that each class y associated with 2 orthonormal feature vectors yy and 2y and we once again consider N and l to be sufficiently large. As mentioned Allen-Zhu and Li (2021), )) we could alternatively consider the number of classes of associated features to be large, and our theory would still translate. We now precisely define the data generating distribution D that we will focus on for the remainder of the paper.
be fixed (i.e. binary classification) and the number
Definition 4.1. [General Multi-View Data Distribution] Identically to Definition 3.1, each class y is associated with two orthonormal feature vectors, after which each point (x, ~ 'D is generated as:.
1. Sample a label y uniformly from [k]].
2. Designate via any method two disjoint subsets (x. Py.2(z) [P] with (
(z))(z) = Cp for a universal constant Cr. and additionally choose via any method a bijection y Py.1(1) (1:) > (a))(z). We then generate the signal patches of corresponding pairs (x()) = By, and ((((p)) = (82 - By.p)iy. = By.ptp) for every p E Py, (1) with the Sy.z chosen according
symmetric distribution (allowed to vary per class y) supported on [81,82 - 81] satisfying the anti-concentration property that 8y.p takes values in a subset of its support whose Lebesgue measure is ((1 log k) with probability o(1).
3. Fix, via any method, Q distinct classes $1, $2, ..., sQ E [k] \ with Q = +(1). The remaining [P] \ (2r) U (x))(xr) patches not considered above the feature noise patches of x. and are defined be (P)) Ejeeo EtEp) U0, where the Yj [d3, 84] can be arbitrary.
Note that there are parts of the data- data-generating process that we leave underspecified, our results will work for any choice. Henceforth, we use X' to refer to a dataset consisting of N i.i.d draws from the distribution D. Our data distribution represents very low signal-to-noiss (SNR) setting in which the true signal for a class exists only in a constant (2Cp) number of patches while the rest of the patches contain low magnitude noise in the form of other class features.
This assumption is true for any distribution with reasonable variance; for example, the uniform distribution.
5

---

Under review as a conference paper at ICLR 2023
We focus on the case of learning the data distribution D with the same two-layer CNN-like architecture used in Allen-Zhu and Li (2021). We recall that this architecture relies on the following polynomially- smoothed ReLU activation, which we refer to ReLU:
if 0
if [0, pl
ReLU(x)
Orpra
I
if 2
The polynomial part of this activation function will be very useful for us in suppressing the feature noise in DD. Our full network architecture, which consists of m hidden neurons, can then be specified as follows.
Definition 4.2. [2-Layer Network] We denote our network by g R Pa + R For each E [kt], we define g" follows.
((War,
(0(x)= Σ E ReL.U
(4.1)
rE(m PE(P)
We will use wy.. to refer to the weights of the network g at initialization (and wy. after steps of gradient descent), and similarly 9r refer to the model after t iterations of gradient descent. We consider the standard choice of Xavier initialization, which, in our setting, corresponds to (0) (0, ala).
For model training, we focus on full batch gradient descent with a fixed learning rate of n applied to [g, X²) and M 2'). Once again using the notation V for 9 the updates to the weights of the network y are thus of the form:
awy.
w((+1) = ((()
wy. (M
(4.2)
In defining our data distribution and model above, we have introduced several hyperparameters. Throughout results, make the following assumptions about these hyperparameters. Assumption 4.3. [Choice of Hyperarrameters] We assume that:
d = ())22))
P ( 83, 84 (k-5
Cr (1))
(k:)
81, 62 (())
p = ((1kk)
a 8
Discussion of Hyperparameter Choices. We make concrete choices of hyperparameters above for the sake of calculations (and we stress that these are not close to the tightest possible choices), but only the relationships between them are important. Namely, we need d to be significantly larger polynomial of & than P, we need d3, 64 = o(1) but large enough so that Påg >> 82 (to avoid learnability by linear models, shown below), suppress the low-mage magnitude feature noise, and we need 61, 62 = (1) so that the signal feature coefficients significantly outweigh the noise feature coefficients.
need sufficiently large so that the network can
To convince the reader that our choice of model is not needlessly complicated given the setting, we prove the following result showing that there exist realizations of the distribution D on which linear classifiers cannot achieve perfect accuracy.
Proposition 4.4. There exists D satisfying all of the conditions of Definition 4.1 and Assumption 4.3 such that with probability at least - k exp(e(-N//E))) the form (x:) Epe[r] x()) and any X' consisting of N i.i.d. draws from D, there exists point (z. v) X and a class # y such that hx)) 2 >²()))
for any classifier h Pd - R of
Proof Sketch. The idea, as was originally pointed out by Allen-Zhu and Li (2021), is that there (A2)) feature noise patches with coefficients of order (kk-1.5). Thus, because the features orthogonal, these noise patches can influence the classification by order @( term away from the direction of the true signal.
The full proof can be found in Section of the Appendix.

---

Under review as a conference paper at ICLR 2023
4.2 MAIN RESULTS
Having established the setting for our main results, we now concretely define the notion of feature learning in our context.
Definition 4.5. [Feature Learningl Let (z, y) ~ D. We say that feature Vyye is learned by if argmax, 9 (x') = y where s' is x with all instances of feature vy,3-r replaced by the all-zero vector.
Our definition of feature learning corresponds to whether the model is able to correctly classify data points the presence of only a single signal feature instead both (generalizing the notion of weight- feature correlation to nonlinear models). By analyzing the gradient descent dynamics of g for the empirical cross-entropy J, we can then show the following.
Theorem 4.6. For k and N sufficiently large and the settings stated in Assumption 4.3, we have that the following hold with probability at least 1 1 ((1/k) after running gradient descent with step size 77
poly(k)) (k)) for ((poly(k) (k)/n) iterations on J[g, g, X) (for sufficiently large polynomials in k):
1. (Training accuracy is perfect): For all (rx, E X, we have argmax, 98 (2:.) y/-
2. (Only one feature is learned); For (1 o(1)) classes, there exists exactly one feature that is learned the sense Definition 4.5 by the model g:-
Furthermore, the above remains true for all (ppol((k)) for any polynomial in &.
Proof Sketch. The proof is in spirit very similar to Theorem in Allen-Zhu and Li (2021), and relies on many of the tools therein. The main idea that, with high probability, there exists separation between the class y weight correlations with the features Uy and 2/y22 at initialization.
This separation is then amplified throughout training due to the polynomial part of ReLU, Once one feature correlation becomes large enough, the gradient updates to the class y weights rapidly decrease, leading to the remaining feature not being learned.
Theorem 4.6 shows that only one feature is learned (in our sense) for the vast majority of classes. As mentioned, our proof is quite similar to Allen-Zhu and Li (2021), but due to simplifications in our setting (no added Gaussian noise for example) and some different ideas the proof is much shorter we hope this makes some of the machinery from Allen-Zhu and Li (2021) accessible to wider audience.
The reason we prove Theorem 4.6 is in fact to highlight the contrast provided by the analogous result for Midpoint Mixup.
Theorem 4.7. For k and N sufficiently large and the settings stated in Assumptions 4.3, have that the following hold with probability 1 ((1k)) after running gradient descent with a step size n O(1/polykk)) for O(poly(k)/n) iterations on Jua (g, X) (for sufficiently large polynomials in k):
1. (Training accuracy is perfect): For all (24, y:) X, we have argmax, (x:) y/-
2. (Both features are learned): For each class y E [k:. both U9 and U9.2 are learned in the sense of Definition 4.5 by the model g.
Furthermore, the above remains true for all = (poly(k)) for any polynomial in
Proof Sketch. The core idea of the proof relies similar techniques to that of Proposition 3.4, but the nonlinear part of the ReLU activation introduces a few additional difficulties due to the fact that the gradients in the nonlinear par are much smaller than those in the linear part of ReLU. Nevertheless, we show that even these smaller gradients are sufficient for the feature correlation that lagging behind catch up polynomial time.
The full proofs of Theorems 4.6 and 4.7 can be found in Section B of the Appendix.
Remark 4.8. Theorems 4.6 and 4.7 show a separation between ERM and Midpoint Mixup with respect feature learning, trained models on the distribution D; even learning only a single feature per class sufficient for perfect test accuracy D. The significance (and our desired interpretation) of these results is that,
have defined. They
results regarding the test accuracy of the
7

---

Under review as a conference paper at ICLR 2023
when the training distribution D has some additional spurious features when compared to the testing distribution, ERM can potentially fail to learn the true signal features whereas Midpoint Mixup will likely learn all features (including the true signal). One may also interpret the results generalization that is robust to distributional shift; the test distribution this case has dropped some features present in the training distribution.
5 EXPERIMENTS
The goal of the results of Sections 3 and 4 was to provide theory (from feature learning and optimization perspective) for why Mixup has enjoyed success over ERM in many practical settings. The intuition is that, for image classification tasks, could reasonably expect images from the same class to be generated from a shared set of latent features (much like our data distribution in Definition 4.1), in which case it may be possible achieve perfect training accuracy by learning a strict subset of these features when doing empirical risk minimization. On the other hand, based on our ideas, we would expect Mixup to learn all such latent features associated with each class (assuming some dependency between them), and thus potentially generalize better.
direct empirical verification of this phenomenon on image datasets is tricky (and a possible avenue for future work) due to the fact that one would need to clearly define notion of latent features with respect to the images being considered, which outside the scope of this work. Instead, we take for granted that such features exist, and attempt verify whether Mixup is able to learn the "true" features associated with each class better than ERM when spurious features are added.
For our experimental setup, we consider training ResNet- et al., 2015) on versions of Fashion MNIST (FMNIST) (Xiao et al., 2017), CIFAR-10, and CIFAR- 100 (Krizhevsky, 2009) in which every training data point is transformed such that randomly sampled training point from different (but randomly fixed) class is concatenated (along the channels dimension) to the original point. Additionally, introduce dependency structure akin to what we have in Definitions 3.1 and 4.1 1,1, we sample y ~ Uni([O, 1J) and scale the first part of the training point (the true image) by while scaling the concatenated part by - during training.
y- 1-
Figure 1: Visualization of data modification in CIFAR-10.
If we work under the intuition that images from each class generated by relatively different latent features, then this modification process corresponds to adding patches of (fixed) spurious features to each class that have a dependency (from the scaling factor 7) the original features of the data. We leave the test data for each dataset unmodified, except for the concatenation of an all-zeros vector of the same shape each point so that the shape of the test data matches that of the training data (in effect, this penalizes models that learned only the spurious features we concatenated in the training data). This zeroing out of the additional channels also intended to replicate Definition 4.5 in our experimental setup.
While we consider the above setup to be intuitive and resemble our theoretical setting, it fair to ask why we chose this setup compared to the many possible alternatives. Firstly, we found that using synthetic spurious features (i.e. random orthogonal vectors scaled have the same norm as the images) opposed to images from different classes was far too noisy (training error went to 0 immediately); the test errors on each dataset degraded to near-random levels, so it was difficult to make comparisons. Additionally, we found the same be true if considered adding spurious features as opposed to concatenating them.
For each of our image classification tasks, train models using Mixup with Dx Beta(1, 1) (the choice used in Zhang et al. (2018) for CIFAR, which we refer to as Uniform Mixup), Midpoint Mixup, and ERM. Our implementation is in PyTorch (Paszke et al., 2019) and uses the ResNet
8

---

Under review as a conference paper at ICLR 2023
Rr1 CLENES, PURAT OF Nesturs, MAY,
Potee Curant, CHARRS FROELeNS AdWS, RRan
Foar DIVES CRAATIO FARUERS ASNYL S.RAPO
67
64
035
(a) FMNIST
(b) CIFAR-10
(c) CIFAR-100
Figure 2: Test error comparison between Uniform Mixup (green), Midpoint Mixup (orange), and ERM (blue). Each curve represents the average of 5 model runs (over the randomness of the data augmentations and model initializations), while the surrounding area represents standard deviation.
implementation of Kuang Liu, released under an MIT license. All models were trained for 100 epochs with batch size of 750, which was the largest feasible size on our compute setup of a single P100 GPU (we use a large batch size to approximate the full batch gradient descent aspect of our theory). For optimization, we use Adam (Kingma and Ba, 2015) with the default hyperparameters of B1 = 0.9, 32 = 0.999 and a learning rate 0.001. We did a modest amount of hyperparameter tuning in preliminary experiments, where compared Adam and SGD with different log-spaced learning rates in the range [0.001, 0.1], and found that Adam with the default hyperparameters almost always worked the best. We report our results for each dataset in Table 1, and accompanying test error plots are shown in Figure 2.
Model Uniform Mixup
FMNIST
CIFAR-10
CIFAR- 100
9.66
18.52 +1
53.42 +1
Midpoint Mixup
14.84 +1
22.29 #2
53.61 #2
ERM
16.55 #2
27.77 12
69.28 #2
Table 1: Final test errors on unmodified test data (mean over 5 runs) along with 1 standard deviation range for Uniform Mixup, Midpoint Mixup, and ERM.
From Table we see that Uniform Mixup performs the best in all cases, and that Midpoint Mixup tracks the performance of Uniform Mixup reasonably closely. We stress that the ordering of model performance is unsurprising; a truly fair comparison with Midpoint Mixup would require training on all N2 possible mixed points, which is infeasible in our compute setup (we opt to randomly mix points per batch, as is standard). Our experiments are intended to show that Midpoint Mixup still non-trivially captures the benefits of Mixup in empirical setting that is far from the asymptotic regime of our theory, while Mixup using standard hyperparameter settings significantly outperforms ERM in the presence of spurious features. A final observation worth making is that find Midpoint Mixup performs significantly better than ERM when moving from the 10-class settings of FMNIST and CIFAR-10 to the 100-class setting of CIFAR-100, and this is in line with what our theory predicts (a larger number of classes more closely approximates our setting).
6 CONCLUSION
To summarize, the main contributions of this work have been theoretical motivation for an extreme case of Mixup training (Midpoint Mixup), as well as an optimization analysis separating the learning dynamics of a 2-layer convolutional network trained using Midpoint Mixup and empirical risk minimization.
Our results show that, for a class of data distributions satisfying the property that there multiple, dependent features correlated with each class in the data, Midpoint Mixup can outperform ERM (both theoretically and empirically) in learning these features. We hope that the ideas introduced in the theory can be a useful building block for future theoretical investigations into Mixup and related methods in the context of training neural networks.

---

Under review as a conference paper at ICLR 2023
7 ETHICS STATEMENT
As this work is almost entirely theoretical in nature, we do not anticipate any (direct) negative broader impacts potential for misuse.
8 REPRODUCIBILITY STATEMENT
All of the results discussed in this paper have accompanying complete proofs in Sections B and C of the Appendix. While the proofs are quite technical in nature, we have tried our best to provide intuitive explanations at each step, and have also included derivations of various calculations and well-known concentration of measure results in Section A of the Appendix.
We have also included in the supplementary material the code necessary to run the experiments in Section 5, along with detailed instructions explaining how to recreate each our experiments.
REFERENCES
Zeyuan Allen-Zhu and Yuanzhi Li. Towards understanding ensemble, knowledge distillation and self-distillation in deep learning, 2021.
Eric Arazo, Diego Ortego, Paul Albert, Noel E O' Connor, and Kevin McGuinness. Unsupervised label noise modeling and loss correction. arXiv preprint arXiv:1904.1123: 2019.
David Berthelot, Nicholas Carlini, Ekin D Cubuk, Alex Kurakin, Kihyuk Sohn, Han Zhang, and Colin Raffel. Remixmatch: Semi-supervised learning with distribution alignment and augmentation anchoring. arXiv preprint arXiv:1911.09785 2019a.
David Berthelot, Nicholas Carlini, Ian Goodfellow, Nicolas Papernot, Avital Oliver, and Colin A Raffel. Mixmatch: A holistic approach to semi-supervised learning. In H. Wallach, H. Larochelle. Beygelzimer, F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019b. URL htwpt:tt///teett httptppss////occeeiiii ngs. neurips ccc/paaper/2219//11ee 1cd1 10100 99a 68 f.b44bbb72 (oboo772e0e0eomm 4bbp:ee0000bcc2d---aaa
Chris M Bishop. Training with noise is equivalent to tikhonov regularization. Neural computation, 7 (1):108-116, 1995.
Luigi Carratino, Moustapha Cissé, Rodolphe Jenatton, and Jean-Philippe Vert. On mixup regulariza- tion, 2020.
Shuxiao Chen, Edgar Dobriban, and Jane Lee. A group- theorete framework for data augmentation. Advances in Neural Information Processing Systems, 33:21321-21333, 2020.
Victor Chernozhukov, Denis Chetverikov, and Kengo Kato. Comparison and anti-concentration bounds for maxima of gaussian random vectors, 2014.
Muthu Chidambaram, Xiang Wang, Yuzheng Hu, Chenwei Wu, and Rong Ge. Towards understanding the data dependency of mixup-style training. CoRR, abs/21 10.07647, 2021. URL https:
// //arxiv. org.vvorrgaoo http://bs/211100/.//
Ching-Yar Yao Chuang and Youssef Mroueh. Fair mixup: Fairness via interpolation. abs/2103.06503, 2021. URL https: https:///rxivv http:////axiv.ooo//////// wtt.:o////.......///533
CoRR,
Tri Dao, Albert Gu, Alexander Ratner, Virginia Smith, Chris De Sa, and Christopher Ré. A kernel theory of modern data augmentation. In International Conference Machine Learning, pages 1538. PMLR, 2019.
Hongyu Guo. Nonlinear mixup: Out-of-manifold data augmentation for text classification. In The Thirty-Fourth AAAI Conference Artificial Intelligence, AAAI 2020, The Thirty- -ieity-Second Innovative Applications of Artificial Intelligence Conference. IAAI 2020, The Tenth AAAI Symposium on Educational Advances in Artificial Intelligence, EAAI 2020. New York, NY, USA, February 7-12, 2020, pages 4044 4051, AAAI Press, 2020. URL https: / /ta.aaaa..rggoaaaaaa.../.a.....a orr///rg///// jjs// /inddxnx.x./////// /ndex.php/ AAAT/ A.AT/artticlee/veew////
10

---

Under review as a conference paper at ICLR 2023
Hongyu Guo. Midpoint regularization: from high uncertainty training to conservative classification, 2021.
Hongyu Guo, Yongyi Mao, and Richong Zhang. Mixup as locally linear out-of-manifold regu- larization. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 3714-3722, 2019.
Boris Hanin and Sun. How data augmentation affects optimization for linear regression. Advances in Neural Information Processing Systems, 34, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition, 2015.
Tong He, Zhi Zhang, Hang Zhang, Zhongyue Zhang, Junyuan Xie, and Mu Li. Bag of tricks for image classification with convolutional neural networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), June 2019.
Yangsibo Huang. Zhao Song, Kai Li, and Sanjeev Arora. Instahide: Instance-hiding schemes for private distributed learning, 2021.
Yannis Kalantidis, Mert Bulent Sariyildiz, Noe Pion, Philippe Weinzaepfel, and Diane Larlus. Hard negative mixing for contrastive learning. Advances in Neural Information Processing Systems, 33: 21798-21809, 2020.
Jang-Hyun Kim, Wonho Choo, and Hyun Oh Song. Puzzle mix:
Exploiting saliency and local statistics for optimal mixup. In International Conference Machine Learning, pages 5275-5285.
PMLR, 2020.
JangHyun Kim, Wonho Choo, Hosan Jeong, and Hyun Oh Song. Co-mixup: Saliency guided joint
mixup with supermodular diversity. In International Conference Learning Representations, 2021. URL https: http::::///ppnrevieww net/ forum? wtw.:mmdidgxxxxxxkkkk.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http:// /atp///aavvvvv/ee///e/ .rwg abs 1411441244 6980.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, 2009.
Alex Lamb, Vikas Verma, Juho Kannala, and Yoshua Bengio. Interpolated adversarial training: Achieving robust neural networks without sacrificing too much accuracy. In Proceedings of the 12th ACM Workshop Artificial Intelligence and Security, pages 95-103, 2019.
Kibok Lee, Yian Zhu, Kihyuk Sohn, Chun-Liang Li, Jinwoo Shin, and Honglak Lee. i-mix: A strategy for regularizing contrastive representation learning. CoRR, abs/2010.08887, 2020. URL
https: /ttpaaaaivvv//.////vv//eaaeaaaaa orgg/aaggao alss/ 2010 08887.
Daojun Liang, Feng Yang, Tian Zhang, and Peter Yang. Understanding mixup training methods. IEEE Access, 6:58774-58783, 2018.
Soon Hoe Lim, N. Benjamin Erichson, Francisco Utrera, Winnie Xu, and Michael W. Mahoney. Noisy feature mixup, 2021. URL https: / http///axvvv.//////////// wwwpaa/swwr...///..e/ao....b..s/ / 2110. C2180.
pettps://arxxr
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Learning with invariances in random features and kernel models. In Conference on Learning Theory, pages 3351-3418. PML....... 2021.
Chanwoo Park, Sangdoo Yun, and Sanghyuk Chun. A unified analysis of mixed sample data augmentation: A loss function perspective. arXiv preprint arXiv:2208. 09913, 2022.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary Vito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems, volume 32, pages 8026-8037. Curran Associates, Inc., 2019.
11

---

Under review as a conference paper at ICLR 2023
Shashank Rajput, Zhili Feng, Zachary Charles, Po-Ling Loh, and Dimitris Papailiopoulos. Does data augmentation lead to positive margin? In International Conference on Machine Learning, pages 5321-5330. PMLR. 2019.
Ruoqi Shen, Sébastien Bubeck, and Suriya Gunasekar. Data augmentation as feature manipulation: a story of desert cows and grass cows, 2022. URL https: http:///wwwwi...//a/aa/////a//// 221022223..........2222 01572.
Junhyuk So, Changdae Oh, Minchul Shin, and Kyungwoo Song. Multi-modal mixup for robust fine-tuning, arXiv preprint Stt.::220000389977 2022.
Kihyuk Sohn, David Berthelot, Nicholas Carlini, Zizhao Zhang, Han Zhang, Colin A Raffel, Ekin Do- gus Cubuk, Alexey Kurakin, and Chun-Liang Li. Fixmatch: Simplifying cemi-supervised learning with consistency and confidence. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 596- 608. Curran Associates, Inc., 2020. URL https: newrips ec/ heppaprr/ 2020 / http://666660.ecooooooooooeo66..
Sunil Thulasidasan, Gopinath Chennupati, Jeff A Bilmes, Tanmoy Bhattacharya, and Sarah Michalak.
On mixup training: Improved calibration and predictive uncertainty for deep neural networks. Advances in Neural Information Processing Systems, 32: 13888 11899, 2019.
Vikas Verma, Alex Lamb, Christopher Beckham, Amir Najafi. Ioannis Mitliagkas, David Lopez-Paz. and Yoshua Bengio. Manifold mixup: Better representations by interpolating hidden states. In International Conference on Machine Learning, pages 344876447. PMLR, 2019.
Vikas Verma, Minh-Thang Luong, Kenji Kawaguchi, Hieu Pham, and Quoc V. Le. Towards domain- agnostic contrastive learning, 2021a.
Vikas Verma, Meng Qu, Kenji Kawaguchi. Alex Lamb, Yoshua Bengio, Juho Kannala, and Jian Tang. Graphmix: Improved training of gnns for semi-supervised learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 10024-10032, 2021b.
R. Vershynin. frigh-Dimensiona Probability: An Introduction with Applications in Data Sci- ence. Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press, 2018. ISBN 9781108415194, URL /tttps:/bbokke google com/books?id NDdeDWAAQRAJ.
Rui Wang, Robin Walters, and Rose Yu. Data augmentation vs. equivariant networks: theory of generalization on dynamics forecasting. arXiv preprint arXiv:2206.094: 2022.
Yiwei Wang, Wei Wang, Yuxuan Liang, Yujun Cai, and Bryan Hooi. Mixup for node and graph classification. In Proceedings of the Web Conference 2021, pages 3663-3674, 2021.
Sen Wu, Hongyang Zhang, Gregory Valiant, and Christopher Ré. On the generalization effects of linear transformations data augmentation. In International Conference on Machine Learning, pages 10410-10420. PMLR, 2020.
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms, 2017. URL sittps://artt org abs/ 07747.
Shuo Yang, Yijun Dong, Rachel Ward, Inderjit S Dhillon, Sujay Sanghavi, and Qi Lei. Sample efficiency of data augmentation consistency regularization. arXiv preprint arXiv:2202. 12230, 2022.
Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy to train strong classifiers with localizable features. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 6023-6032, 2019.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. In International Conference on Learning Representations, 2018.
Linjun Zhang, Zhun Deng, Kenji Kawaguchi, Amirata Ghorbani, and James Zou. How does mixup help with robustness and generalization?, 2020.
Linjun Zhang, Zhun Deng, Kenji Kawaguchi, and James Zou. When and how mixup improves calibration, 2021.
12

---

Under review as a conference paper at ICLR 2023
A SUPPORTING LEMMAS AND CALCULATIONS
In this section we collect several technical lemmas and computations that will be necessary for the proofs of our main results.
A.1 GAUSSIAN CONCENTRATION AND ANTI-CONCENTRATION RESULTS
The following are well-known concentration results for Gaussian random variables; we include proofs for the convenience of the reader.
Proposition A.I. Let X; NNO,, of) with [m] and let max; di Then,
E max X₂) S a / log m
Proof. Let Z max; X,. Then by Jensen's inequality and the MGF of NV of), have:
exp (tE < Exxp(tZ) E/exp maax maxxx.)
E
exp((X,))
exp((oo712)
m mexp(ff2212)
log to²
B| <
Minimizing the RHS yields = logom m/o, from which the result follows.
Proposition A.2. Let X; be as Proposition A.1. Then,
P(max X; 2 L ov o m) S {exp(-12/(2o))) i
Proof. We simply union bound and use the fact that P(X, 21) S for zero mean Gaussians) to get:
(Chernoff bound
P(max X, 2 t t 22gg m)
S E P(X; 2 t /2gg
I mexp
-(1 a / 22g m) /
s exp(-tt/11 (f(-t(/(207))
Proposition A.3. Let X1, X2,.., Xm be i.i.d. Gaussian variables with mean 0 and variance o².
Then we have that:
P maxxx: X, @ (ovoo////@///////////// ((//log(m/log((/81)) -1-
Proof. We recall that:
P(G; > 2) = 4 (=e--~11200)
proof of this fact can be found in Vershynin (2018). We additionally have that:
P(max X; > x) = (1 11-PPo -PX, > x²)²
So from the previous asymptotic characterization of P(X, > x) have that choosing @ ( (ovlog(m//og((1) gives P(X; > r)
from which the result follows.
We will also have need for a recent antiiconcentration result due to Chemoozhukvv et al. (2014), which restate below.
13

---

Under review as a conference paper at ICLR 2023
Proposition A.4 (Theorem 3 (i) Chernozhukov al. (2014)). Let X; ~V(00, for [m] with r² > 0. Defining Am = max X,/o], we then have for every > 0:
sup P( sER
max X; =x{ S e) s 4e(1 + am)/o
Corollary A.5. Applying Proposition A.1, we have suPNER P([max; X; - x| S c) 1 4c(1 +
V2logm)) m)/o.
A.2 GRADIENT CALCULATIONS
Here we collect the gradient calculations used the proofs the main results. We recall that we use and Fi.j (z; + x2)/2. Additionally, we will omit parentheses after ReLU
V to to indicate Su
when function application is clear.
Calculation A.6. For any (Ii Mi) E X:
Dt Di (I,) = Σ
ReLU (unir 2 (11) x
(t) Wy
PE[P]
Proof.
Wygar 96 (z.)
Σ
Σ
ReLU whh offe)
Σ ReLU ( (u ( 22/01 x!!0)
ouig (1)
we(m) PE(P)
PE[P)
Calculation A.7. For any (1t, vh) X, if max ( (whhrr Un.e)
(1)
< p/(62 - 51) and # yes then:
Sa (wbirr U Uyril)
Dt
St
(Ii), uyr.e)
pa-1
Pf l e(ms)
Po maxuga y Wyyer Buie
Hi (2i), Vole 5
pp-1
Proof. When max ( (orers vu,e) p/(822 -51), are in the polynomial part of ReL.U for (t)
every patch in Fg since maxpePy.eeees ( (Wyrrr xx! p since St.p 82 81. The first line then follows from Calculation A.6 and the fact that all of the feature vectors are orthonormal (so only those patches that have the features uy are relevant). The second line follows from the fact that there are at most P 2C'p feature noise patches containing the vector Us,s, and in each of these patches there are only a constant number of feature vectors (which we do not constrain).
Calculation A.8. For any (Ii, yb) X, (ugh)):)r.)) Vos.)) 2 0/81, then:
Mi (Ix)) vycl))
Σ
Bip
un
((P
Proof. When (ogers ky..)
(t)
p/8, we necessarily have min (()))))))))) ( (w x
2 p since
Si.p 2 81 and then the result again follows from Calculation A.6 and the fact that ReLU = in the linear regime.
Calculation A.9 (ERM Gradient).
V Wy..
J[et X))
1
Σ (1 (1 -
-
10 (xx)) (x:)
N
with
SE[N]
14

---

Under review as a conference paper at ICLR 2023
Proof. First let observe that:
log 6hh (gr (3e() - (x:) - log Σ exp(gi(x:))
log @ (ge(x:)) w V (xx) 6" (g(xxt)
1o 1o (xx)
Bwyr Summing (and negating) the above over all points I; gives the result.
Calculation A.10 (Midpoint Mixup Gradient).
(t) (9r,
1 Σ Σ
(Iu:zy + Lur=y 20" 20" (9((2(()))))
(t (2rg)
2N
JE[N]
Proof. Follows from applying Calculation A.9 each part of the summation in JMM X).
B PROOFS OF MAIN RESULTS
This section contains the proofs of the main results in this paper. We have opted to present the proofs in linear fashion inlining several claims and their proofs along the way we find this to be more readable than the alternative. The proofs inlined claims are ended with the symbol, while the proofs of the overarching results are ended with the D symbol.
For convenience, we recall the assumptions (as they were stated in the main body) that are used in these results:
Assumption 4.3. [Choice of Hyperparameters] We assume that:
d = (1)220)
P = (
Cp = (11)
m ((k)
81, 62 @(1)
83, 64 ((k-1.)
p ((1kk)
8
B.1 PROOF OF THEOREM 4.6
Theorem 4.6. For l and N sufficiently large and the settings stated in Assumption 4.3, we have that the following hold with probability least 0 after running gradient descent with a step size 77 (11ppoly (tt//oll(k))) (k)) for (pooy (k)/n) iterations on J [g, X) (for sufficiently large polynomials in k);
1. (Training accuracy is perfect): For all (z,, W:) E X, we have argmax, g (zi) = V/-
2. (Only one feature is learned): For (1 - o(1)) classes, there exists exactly one feature that is learned the sense Definition 4.5 by the model gs.
Furthermore, the above remains true for all (poly(kk) for any polynomial in
Proof. We break the proof into two parts. In part one, we prove that (with high probability) each class output of becomes large (but not too large) on data points belonging to class and stays small on other data points, which consequently allows us obtain perfect training accuracy the end (thereby proving the first half of the theorem). In part two, we show that (again with high probability) max correlations with features 1y and 6,.2 for class y have a separation at initialization that gets amplified over the course of training, and due to this separation one of the feature correlations becomes essentially irrelevant, which will be used to prove the second half of the theorem.
Part I.
In this part, we show that the network output gf (zx) reaches and remains Olog k) while of (x:)) = o(1) for all t (ppol((k) and + yi- These two facts together allow to control the 1 6th (gx()) terms that show up throughout our analysis (see Calculation A.9), while also being sufficient for showing that get perfect training accuracy. The intuition behind these results is that, when 9f" (x,) > clog k. we have that exp(ge" (x,)) > kf so the 6t (g. (x;)) terms the gradient updates quickly become small and 9e stops growing. Throughout this part of the proof and the next,
15

---

Under review as a conference paper at ICLR 2023
will use the following notation (some of which has been introduced previously) to simplify the presentation.
Ny = (i: : E [N] and yi = =ソ ( = pp= [P) and (z,²) m220) > 0}
B0 2,6 = fr:re: e [m]] and (w m/.e) 2 p/ó1}
(B.1)
Here, Ny represents the indices corresponding to class y points, (v.e( (as used in Definition 4.1) represents the patch support of the feature Vy, in æ; (recall the features are orthonormal), and B(t) represents the set of class weights that have achieved big enough correlation with the feature dg.6 necessarily be in the linear regime of ReLU on all class y points iteration t.
9.6
Prior to beginning our analysis of the network outputs g². we first prove a claim that will serve as the setting for the rest of the proof.
Claim B.1. With probability 011/k), all of the following are (simultancously) true for every class E [
N, = (NN/k)
max,
Warrs vy Uye) Ollgg vd)
V E [2], max+e(m) (wyrr v/.e) ((1/VQ)
Proof of Claim B.I. We prove each part of the claim in order, starting with showing that |NV
O(N/ fk) with the desired probability for each To see this, we note that the joint distribution of the Ng multinomial with uniform probability 1/k. Now by a Chernoff bound, we have that (NN/k) with probability at least exp( (ap((((-)) (O(-N/))) Conditioning on |N1 = (N/kk) we have that the joint distribution of [N.l,..., [N&] is multinomial with uniform probability 1/(k - 1). so we obtain an identical Chernoff bound for [N2]. Repeating this argument and taking union bound gives that Ng = (N/k) for all y E [k] with probability at least - k cxxp(O(-N/k)..
[N1
The fact that for every y we have maxx Ws.r: Uy.s) = (log k/ Vd) with probability
(0)
1 - O(1/k) follows from Proposition A.2. Namely, using Proposition A.2 with t = 20 log (here 1/vd by our choice of initialization) yields that max, (... 2/.e) 2 3 2 Te/ Vd with
2lgm
probability bounded above by 1/k3 for any s, y. Taking a union bound over s, y then gives the result. The final fact follows by near identical logic but using Proposition A.3 (note that the correlations
(0) Ug.r) are i.i.d. (0, 1/d) due to the fact that the features are orthonormal and the weights themselves i.i.d.).
In everything that follows, we will always assume the conditions of Claim B.1 unless otherwise stated. We begin by proving result concerning the size of softmax outputs d" (z)()) that we will repeatedly use throughout the rest of the proof.
Claim B.2. Consider E Ny and suppose that both max,e(k), re[m), (c[2] ws.r. vs.e = O(lg k)
(wel), Wairr uy.e) (llog())/Vi) hold true. we have (xx)) 2 a log le for
and max rE[m), #E[2]
some a E oo), then:
(1 (0())))))))(()))))))))))))))))))
ifa>
1 1- (c))) ( (1)) otherwise
Proof of Claim B.2. By assumption, all of the weitht-feature correlations are Ollog k) at Further- more, for + y, all of the off-diagonal correlations
( (ws,7 uy.r) are Ollog(k)/V)) This implies
16

---

Under review as a conference paper at ICLR 2023
that (using 84 -.6 @(1/k) P= ((k2), and 8):
mPP maxuzy (wbit Vest
pa-1
k2+a log(k)))
1og(())
0
1l.5a
A-d
[kg(()))
exp(s & 11+0
(B.2)
1:3
Where above we disregarded the constant number (2Cp) of very low order correlations
(Wsrr Vy.e)
and used the inequality that exp(x) < 1 + x + x² for x I 1. Now by the assumption that
ge (zx) 2 log k we have exp(gi (z:)) 2 ka, s0:
ka
2 (xx(x))) 1
ka + (k - 1) + o(1)
k 1 o1))
(B.3)
kat (k: - 1) + o(1)
From which the result follows.
Corollary B.3. Under the same conditions as Claim B.2, for + y, we have:
1
6" (oe())) = {max((.1)
Proof of Corollary B.3. Follows from Equations B.2 and B.3.
With these softmax bounds hand, we now show that the "diagonal" correlations (wyer vy.e) grow much more quickly than the the "off-diagonal" correlations (wg.rr vo.t) (where & * y). This will allow satisfy the conditions of Claim B.2 throughout training.
Claim B.4. Consider an arbitrary E [k). Let A S p/ - 61) and let Ta denote the first iteration (Ta)
at which max/e[m]]eE(2)
Wyer Vg 2 A. Then we must have both that TA = O(poly(k)) and
that max
Wy
Ve.e = O (log(+)/Vdd)
Proof of Claim B.4. Firstly, all weight-feature correlations are o(p) at initialization (see Claim B.1). (0) (wgr: Ve.e) > O, we have for every which Wy.r. Vo,e < p/(62 81) that
Now for # y and (using Calculations A.7 and A.9):
(t)
(t) Wgrt, Vo.6
a-1
Pip
wi 02.2)
"
Σ
のd (9((x;))
N
1a-1
i.
pEPr.(Zd))
P8 maxu4y, (E(2)
(witt Vnee
a-1
Σ
1 -0-0" (9((x)))
+ N
pa-1
iE
(t) P8g min m6(2) Uueee pa-1
a-
Σ
()((zz))
iEN,UN.
P81 max MEl) pa-1
Wyrre vall
(c)
1
Σ (1 (1-
N
SEN
(B.4)
17

---

Under review as a conference paper at ICLR 2023
Similarly, for (prrr vv.e)
(0)
> O. we have for every at which (wger Byi)) < p/ - -ô1) that:
(
a-1
ga up (grrs U/ee
V 2.r
J(g 1),
2 N
Σ (1 iE
2d (c((2s)))
pr-
(EP))((Z))
a-1
Po maxu 6222
Wyere Vull
Σ
の (a([x)))
N N
pa-1
(B.5)
From Equation B.5, Claim B.2, and Corollary B.3 we get that for TA:
17 ( (t) Wgrrr ey.e)
(e) w...
X), Uy.l
2 +
kp
(B.6)
Where above also used the fact that [NV, (NV/k). On the other hand, also using Claim B.2 and Corollary B.3, we have that for all for which ( (wg Vo.e) < p/(62 -81):
(t)
n PP max (E(2) (witt: Uu,e"
[g, X), U0
(B.7)
pP-1
Now suppose that (
(Wg.r Ve,e is the maximum off-diagonal correlation at initialization. Then using
(0)
(t) Equation B.7, we can lower bound the number of iterations T it takes for (wgeer 03.2) grow by fixed constant factor from initialization:
11 PP80Ca- Wy.rr Us.E
(C- 1C- (1(0) U'..)
Te
pa-1
pa-1
720
a-2
nP68
Wy Usle
(0)
11.5a-2 pa-1d /2
=
(B.8)
As there exists at least one w (g,c) = 2(1/Vd)) it immediately follows from comparing to Equation B.6 and recalling that Assumption 4.3 that T >> TA and that Ta O(poly(k)) so the claim is proved.
Having established strong control over the off-diagoona correlations, we are now ready to prove the first half of the main result of this part of the proof that g (x;) reaches N(log k) for all E N, in (ppoyy(k) (k)) iterations. In proving this, it will help us to have some control over the network outputs oe across different points 2; and zy at the later stages training. which we take care of below.
Claim B.5. For every E
[k] and all t such that maxieN, g% 2 log k o (log(k))VV)). we have maxieN, (,r,) =
(1) and =F (y.r,, vo,e) A(minee,, de (x:)).
Proof of Claim B.5. Let argmax,e,, of (z.). Since of (zz) 2 log k. necessarily have that By is non-empty for at least one l [2] (since mp O(1)). Only those weights
(0) with By
18

---

Under review as a conference paper at ICLR 2023
for some & [2] are asymptotically relevant (as any weights not considered can only contribute (1) term), and we can write:
(01e)) = Σ
Pop
Σ
wilt) Birr U2 + oloo k)
EE[2]
pE (x.)).(
TEB sie
For any other j E Ny. we have that $j.p 2 61 8/31.p/ (02 81), from which the result follows.
Now we may show:
Claim B.6. For each y [k], let Ty denote the first iteration such that maxieN, S.t, (x:) log k.
Then T3 = =((polykk)) and sm). minseN, de (2x) (log k) for all / 2 Ty
Wayr Vs.e O (log(k)/VV) Furthermore,
(Ty)
Proof of Claim B.6. Applying Claim B.4 to an arbitrary y E [k]] yields the existence of a corre- 2 p/ 51). Reusing the logic of Claim B.4, but this time replacing (w(t) y.rr; : in Equation B.6 with p/(62 - -81), shows that in ((poly(k)) additional iterations we have "yrr Vyye") 2 p/ō, (implying E B while the off-diagonal correlations still remain within a constant factor of initialization.
lation
(w (wyrrr U2,
( (w v2,e+) as (using Calculation A.8):
(1)
Now we may lower bound the update
wo X(, Uy.e-
7
(1 6y (g((z)))
Σ Bip
N
NN
P8 max CE|2 ( (0)) Vuye'
Σ N
2y (x((z))
N
pa-1
(B.9)
log k (which necessarily still the case at this point, as again mp = (1)), we have by the logic of Claim B.2 that we can simplify Equation B.9 to:
So long as ma o (xx4 <
[(g X(, X), X), uy. V wi 2:0
> O(n/k)
(B.10)
Where again we used the fact that |N, O(N/k). Now we can upper bound T, by the number of iterations it takes ( (wyr Vy,e-) to grow log(k)/8, From Equation B. clearly have that T, (polyy(k) (k:) for some polynomial in k. Furthermore, comparing to Equation B.7, we necessarily still have sax (y.. V., O (log(k)/Vd) Finally, as the update in Equation B.9 is positive at Tr (and the absolute value of a gradient update is o(1)), it follows that mingeN, de (x:) 2log k) for all / Ty by Claim B.5.
(T₂)
The final remaining task to show that gi (z,) = Ollog k) and (x,) = o(1) for = (ppoly(k)) and E N2 for every y E [k].
Claim B.7. For all t () for any universal constant C, and for every y E [k] and + y. (8) have that 9f (x.) Ollog k) and (a ( (Wyrr, vs.t) O (log(k)//V) for all E Ny.
Proof of Claim B.7. Let us again consider any class y E [k]] and t Tg. The idea is show that maXieN, 1 (g (x:)) is decreasing rapidly as minjeN, g² (x;) grows to successive levels of log l: for 1.
19

---

Under review as a conference paper at ICLR 2023
Firstly, following Equation B.9, we can form the following upper bound for the gradient updates to B 3.8
r E
w (e) yint 2, Vy,e*
Σ (1 (1- 8² (1((zz))) iENy
Σ
Big
pEPoela.)
$ (11-min
min o² (ge(x))) (((())))))((/1))) ((//k) iEN,
(B.11)
From Equation B.11 it follows that it takes at least O log(k)/ (mn) iterations (since the correlations must grow at least log(k)/m) from Ty for g (æ;) to reach 2 log k. Now let T, denote the number of iterations takes for minkeN, ge [x;) to cross a log after crossing (a - 1) log & for the first time. For 2 3, we necessarily have that To (((TT--1) by Claim B.2 and Equation B.11.
Let us now further define T to be the first iteration which max,eN, 9 (ax) 2 f(k) log for some f(k) = (1). By Claim B.5, at this point minje 9t, (a;) = [)(((k) log k). However, we have from the above discussion that:
[1k)
h2 log k:
Ty z N(poly(k)) +
Σ Ω
17
a=0
log k (k=((k)-2 1)
1i(k 1)
2 w(poly(k))
(B.12)
So maxieN, of (xi) O(logh k) for all / (poly(k))) An identical analysis also works for the off-diagonal correlations ( (u//er vs,e) but forming an upper bound using Equation B.4, so we are done.
We get the following two corollaries straightforward consequences of Claim B.7.
Corollary B.8 (Perfect Training Accuracy). We have that there exists universal constant C' such that argmax, of (I:) = for every (x,, 3:) E X' for all / 2 F.c but with f. (p))y(k)).
Corollary B.9 (Softmax Control). We have that for all E [k]] and any t = O(poly(k)) for any polynomial in & that max,eN, EnSn (xp(g/(xx)) = k + o11).
Corollary B.8 finishes this part of the proof.
Part II.
For the next part of the proof, we characterize the separation between max,e]m)
(w.rr 0g and
(0)
max+e[m] Wyyer 0y.2) and show that this separation (when it is significant enough) gets amplified the course of training. To show this, we will rely largely the techniques found in Allen-Zhu and Li (2021), and finish in a near-identical manner to the proof of Claim B.7.
As with Part I, we first introduce some notation that we will throughout this part of the proof.
Sy,
Σ E
3o
A 1
max [m]]
(Wyrr Ug.e
(t)
N
SEN, PEP
Here, Sy.t represents the data-dependent quantities that show up in the gradient updates to the correlations during the phase of training which the correlations are in the polynomial part of ReLU, while qe represents the max class correlation with feature Uy,e at time 1.
Now we can prove essentially the same result as Proposition B.2 in Allen-Zhu and Li (2021), which quantifies the separation between A (0) and A (0) after taking into account Sy, and Sy
20

---

Lemma B.11 (Lemma C.19 from Allen-Zhi and Li (2021)). Let 2 constant and mo, yo o(1). Let {zt, Ve he be two positive sequences updated as Under review as a conference paper at ICLR 2023
Claim B.10 (Feature Separation at Initialization). For each class y, we have that either:
(0)
2
Sg.2
--z
1 1
llog k) )
1
A A or
yy.1
A Sy,1
-
1 Q
log² k ))
1
A (0)
Sg.2
with probability 0
log k
Proof of Claim B0 Suppose WLOG that Sy,1 2 Sy.2. If neither of the inequalities in the claim hold, then we have that:
y,22
1
1
))
A (0)
Sp,1
log² k
Which follows from the fact that, for constant A, we have:
A
log" k
log
Now recall that A A and A A are both maximums over i.i.d. /N 1/d) variables (again, since the feature vectors are orthonormal), so we can apply Corollary A.5 (Gaussian anti-concentration)
A [0) while taking e (Sy,2/S,,))-- + (1/ log* k)) A and (Sy,2/Sy.1) A It is crucial note that we can only do this because A (0 independent of A ( and both take values over all of R. From this we get that:
4e(1 + /2 logm)
P 1 E (S,,2/S,,1)
A A +e) s
a V m
v log m
+
log 2
1
1 m
0
with probability
log
Where used the fact that (k) and Proposition A.2 characterize Agg (also noting that Sy.2/S, Sy, is ((1)). Thus, neither of the inequalities hold with probability O(1/ log k), we have the desired result.
We the separation from Claim B. 10 to show that, in the initial stages of training, the max correlated weight/feature pair grows out of the polynomial region of ReLU and becomes large much faster than the correlations with the other feature for the same class. For y E [k)] let t' be
y. such that A (0) is the left-hand side of the satisfied inequality from Claim B.10. Additionally, let
r" argmax, (w..r.yy.e*)
(0)
Ug.e.) i.e. the strongest weight/feature correlation pair initialization. We
will show that when (w.(1) [ 29,e- becomes (p)). the other correlations remain small. In order to do so, need useful lemma from Allen-Zhu and Li (2021) that restate below.
T+41 34 C for some C. @(1), and
Me+1 S + SCttt 11 for some constant $ = 0(1)
Where 17 (k1pollykk)
for sufficiently large polynomial in k. Suppose Zo > For every A = O(1), letting T, be the first iteration such that
S 1 θ polybog(k)
It 2 A, we must have that
yr. (xoooolloog(k)
21

---

Under review as a conference paper at ICLR 2023
To apply Lemma B.11 in setting, first prove the following claim.
Claim B.12. For a class E [k]. we define the following two sequences:
Sy.e. pa-1
2-2
(wit) vy,3-0-
dy.t
W yiroo Vyye-)
and by.t
(pp-)
pP-1
Where the r in the definition of by. is arbitrary. Then with probability 1 1- O logg
there exist
C₁, S = @ such that for all for which (w vv.e-) p/(02 - -d)):
d av n0agg a-1
0y.e-1 & by, + 1
Additionally (with the same probability), we have that ay.0 2 S (1 11+ (potvlotts) )) 1y.0-
Proof of Claim B.12. The update to (w vy.e0) (w.er-,,,.)) in this regime can be bounded as follows (using Corollary B.9 and recalling Equation B.5):
(1) Wairr Uyye
(t) J [gg x), Vy,e- v.r-
Σ iEN,
2y (a((x)))
Σ
pa-1
PE (F)
P8T max ('|2]
Wyirr Vuyer
1
Σ
2 (9x(s:))
N
pa-1
i
(0)
(1 1 1 (1) S.ee
,
(B.13)
pa-1
Similarly, have (noting also that (Wyerr v < 0/(02 -1))::
(1)
(1)
Wiiry Vy,3-1
(2) J[g X), 2 21.3-1)) S 91.3 N.
(B.14)
pr
Multiplying the above inequalities by (S / p -- we see that aye and by satisfy the inequal- ities in the claim with C. = - O (+) and S = (Sy,3- 1 (1 θ (+)). Now by Claim B.10 have:
Sy
log k ) by. (polylo(k))
0g.0 2
1+0
Sy,
2 SoLa
1+0
by.0
So are done.
Now by the fact that |, = ((N/k), we have Sy,1, Sy.2 = θ (1/k) = O(p). which implies
that (Sy,e- 1p-) = (1/p). From this we get that while ay.t < C/((22 - 61) for some
appropriately chosen constant C. we have ( (ug.r, uy,e-) p/ 81).
Since Claim 1.12 holds in this regime, we can apply Lemma B. with A = C/(82 - 61), which gives us that when ay,t C02 - 81), we have by (loyooooly)og(k) From this we obtain that
when ( (wg).. vy,00) 2 p0 61) we have that (wyrr v is still within polylog(k) factor
of (wgrr U/,3-8)) for any
(0)
Now from the same logic as the proof of Claim B.7, we can show that this separation remains throughout training.
22

---

Under review as a conference paper at ICLR 2023
Claim B.13. For any class y [k], with probability - O(1/ log k), have that Whir vg.3-e-) for all = ((poly(k) for ma Wyrer 0g.3-e+) O(polylog(k) maxxe(m] any polynomial k.
Proof of Claim 13. It follows from the same logic in the proof of Claim B.6 that at the first iteration for which we have minjeN, 9 (m:) 2 log ki we still have ( (w (t)
(Wyrr Vy,3-e+) is within some (ig Uy,3-r-) can be viewed the same as an off-diagonal correlation from the proof of Claim B.6). The rest of the proof then follows from identical logic to that of Claim B.7; namely, we can show that for (wger vb,3-e* e. to grow by more than polylog( (k) factor we need (poly(kk) training iterations.
(klylog(k) factor of initialization (here the correlation (
From Claim B.13 along with Claim B.7, it follows that with probability 1 0(1/ log k), for any class (after polynomially many training iterations) we have:
(x(( = =0 (mpoyylo(())
kpolylog(k)
Vd
vd for Ps, #
(x.) = Ω (p log(k)
kpolylog(k) Vd
(B.15)
1,2.0
Where æ² is any point 2; with i N2 modified that all instances of feature vy replaced by O, and the second line above follows from the fact that by Claim B.7 must have Wa,rr Vo,e)
(6)
(k)log())/))))))))) for at least some r, for every (and d (l)2))) This proves that feature y. is not learned in the sense of Definition 4.5.
Using Claim B.13 for each class, we have by Chernoff bound that with probability at least 1 - (1/k) that for (1 (k1) classes only a single feature is learned, which proves the theorem.
B.2 PROOF OF THEOREM 4.7
Theorem 4.7. For i and sufficiently large and the settings stated in Assumptions 4.3, we have that the following hold with probability ((1kk) after running gradient descent with step size r) = O(1/poly(kk) for (ppoly(k)/m) iterations on JMM (g, X) (for sufficiently large polynomials in k))
1. (Training accuracy is perfect): For all (æ,, y:) X, we have argmax, gzxz)) = ys-
2. (Both features are learned): For each class y E [k]. both 1y and u2 learned in the sense of Definition 4.5 by the model g.
Furthermore, the above remains true for all = (ppoy((k)) for any polynomial in
Proof. As the proof of Theorem 4.6, we break the proof into two parts. The first part mirrors most of the structure of Part of the proof of Theorem 4.6, in that we analyze the off-diagonal correlations and also show that the network outputs of can grow to (and remain) 2ioog k) as training progresses. However, we do not show that the outputs stay Ollog k) in Part I (as we did in the ERM case). as there are additional subtleties in the Midpoint Mixup analysis that require different techniques which we find are more easily introduced separately.
The second part of the proof differs significantly from Part II of the proof of Theorem 4.6, as our goal is to now show that any large separation between weight-feature correlations for each class are corrected over the course of training. At a high level, we show this by proving a gradient correlation lower bound that depends only on the magnitude of the separation between correlations and the variance of the feature coefficients in the data distribution, after which we can conclude that any feature lagging behind will catch up in polynomially many training iterations. We then use the techniques from the gradient lower bound analysis to prove that the network outputs g stay O k) throughout training, which wraps up the proof.
23

---

Under review as a conference paper at ICLR 2023
Part I. We first recall that 2., (x) + 22)/2, and we refer to such 2., as "mixed points". In this part of the proof, we show that g." (2t,) crosses log ls on at least one mixed point 21, in polynomially many iterations (after which the network outputs remain (loo k)). As before, this requires getting (t)
handle on the off-diagonal correlations (Wy.r. Vs.t) (with + y).
Throughout the proof, we will continue to rely on the notation introduced in Equation B.1 in the proof of Theorem 4.6. However, we make one slight modification to the definition of ((t) for the v.e
Mixup case (so as to be able to handle mixed points), which is as follows:
f u.6 = (r r E [m]]] and (uff), y.r Ug,e 2 2p/ó13
(
(B.16)
We again start by proving a claim that will constitute our setting for the rest of this proof.
Claim B.14. With probability - ((1/k), all of the following are (simultaneously) true for every class y E [k]:
I, = (N//k)
max (warr vyy O(log Vd)
AA E [2), maxre[m] (yyrr (g.e) 2(1/vo)
(0)
For (2k) tuples (s, E) E [K] [2] have
(0) w v3,e) > O.
Proof of Claim B.14. The first three items in the claim are exactly the same as in Claim B.1, and the last item is true because the correlations (Wy.r Ve,t/ are mean zero Gaussians.
Once again, in everything that follows, we will always assume the conditions Claim B. unless otherwise stated. We now translate Claim B.2 to the Midpoint Mixup setting.
Claim B.15. Consider i e Ny. j E N, for s * y and suppose that max,ę (v.s) 9 (21.g) = (llog())) holds true. If we have g (zu.) = a log and gi (z+z) = blog k for a, b = 0(), then:
(1( if > 1, a- -- 211 a > 1, a-aale b- bb (1) otherwise
1 1 -0
(gr (zx(z))))
+0(1)
Where in the second item above the sign of 240 (gx (zi.)) depends on the sign of b.
Proof of Claim 1.5 In comparison to Claim B.2, the Midpoint Mixup case is slightly more involved in that gf (zx,)) can be quite large due to the x; part zi.j As a result, we directly assume some control over the different class outputs on the mixed points (which we will prove to hold throughout training later). By assumption, we have for + y, 8:
9i (21,,)) (llog(A)/H)
exp(o"(21,) < ((lgg())/))
(B.17)
Where above we used the inequality exp(x)) < 1 + + x² for E [0, 1J. Now by the assumptions that 9% (za.5) = log and g; (21,)) = b blog k. we have:
2k
1 - 20" (g (x:)) $
k.t 1h (k: 2) + okk)
bb ka + (k - 2) o(k)
(B.18)
k.a bb + (k - 2) + o(k)
From which the result follows.
Corollary B.16. Under the same conditions as Claim B.15, for u + y. s we have:
1 - 20 (00"(99(21,))
=1)
24

---

Under review as a conference paper at ICLR 2023
Proof of Corollary B.16. Follows from Equations B.17 and B.18.
We observe that Claim B.15 and Corollary B. are less precise than Claim B.2, largely because there is now a dependence the gap between the class y and class network outputs as opposed to just the class y network output. We are now again ready to compare the growth of the diagonal correlations (Wigrr Mg,e) with the off-diagonal correlations (wg)r Va,e) However, this is not as straightforward as it was in the ERM setting. The issue is that the off-diagonal correlations can actually grow significantly, due to the fact that the features Uy, can show up when mixing points in class with class 8.
Claim B.17. Fix arbitrary class y E [k]. Let A [9(p), p/ - 81)1 and let TA be the first iteration at which maxxe[m))xE(e) that, for every 8 + y and & E [2]:
(Ta) Wigh uy.t) 2 A; must have both that Ta Oppoly(k)) and
(wlTa)) yor
Usse) =0
max [EEE[2]
(Taa yor
vyye) /k)
Additionally, for all s, with ( (wbir vo,e) vo
(Ta) have that War U2.e polylog(k)
(Ta)
Proof of Claim 1.17. By our setting, we must have that there exists diagonal correlation (w Vy,e (L/V/d), which we will focus our attention on. Using Calculation A.10 and
the ideas from Calculation A.7, we can lower bound the update to (w vyyt) from initialization up Ta
(1(()) Vyge- pa-l
a-1
(g, vyge-
Σ Σ (1 2000 2400 ((o) ((((=))))))
N2
NEN JēN,
a-1
Wyre Vy pa-1
"
Σ Σ
(1 2tl (Ca((cs))))
N2
Ne JEN,
( (w Unig
a-1
Po maxuelMl.qq(e)
" Σ
Σ
8² (9((2))))
N2
pa-1
JKN,
(B.19)
Above made use of the fact that, for N2 and j & Ny. we have ( (woft) vyye-) /2 for at least O(IN,IN) mixed points since the correlation y Vo.g is pos- itive for l(k) tuples (u, q) E [k] [2] (under Setting B. 14). We similarly upper bound the update
(w (0)
( (p)
25

---

Under review as a conference paper at ICLR 2023
to { (grrr velk for an arbitrary r E [m]] as:
(9
Whjre Vo.e
-n w0 Jor vale)
N Σ
Σ
2 (a((z)))
pa-
iq Ng JEN,
(w Wgitr Un.o
Pog minuelMM
N2 Σ E
("(((2()))
pa-1
iENyUN, JKN,UN.
a-1
maxqe12] Wyins ustt + uy. pa-1
Σ
Σ
(1 -1-200 20t 20 (9((2
+
N
Ny JEN,
a-1
61 max pa-1
wpier vu.g
7 E
Σ (1-203 200"(((
N2
EENy N,UNN
a-1
Po2 max
Whire Vu.g
M E (1- (((((r)))) N2 KENg JEN,
+
pa-1
(B.20)
As the above may be rather difficult to parse on first glance, let us take a moment to unpack the individual terms on the RHS. The first two terms are precise splitting of the -2)"(ge) term from Calculation 10; namely, the case where we mix with the class allows for constant size coefficients on the feature us while the other cases only allow for U,, to show up in the feature noise patches. The next three terms consider all cases of mixing with the class The first of these terms considers the case of mixing class y with class 8, in which case it is possible have patches in 2i,j that have both e and 2y with constant coefficients. The next term considers mixing class with class that is neither y nor 8, in which case the feature u, can only show up when mixing with a feature noise patch, so we suffer factor of at least 54 = O @(1 /k1-5) (note we not suffer of factor as "gge- can still be in ei.)) from the (Zi.j, U,,6) part of the gradient. Finally, the last term considers mixing within class y.
The first of the three positive terms in the RHS of Equation B.20 presents the main problem the fact that the diagonal correlations can show up non-trivially the off-diagonal correlation gradient means the gradients can be much larger than in the ERM case. However, the key is that there are only ((N/k2) mixed points between classes y and class s. Thus, once more using the fact that ((NN.) ((N/k) for every [k), the other conditions in setting, and Claim B.15, we obtain that for all < TA:
" (w Uir Dyge
a=1
V m0 JMMs
2 0
(B.21)
kpa-1
(1)
a-1
n) m
mighr Vang
V with JMas (g. X),
5
(B.22)
122 papa-1
Crucially we have that Equation B.22 is @(1/k) factor smaller than Equation B.21. Recalling that all correlations are (log(k)) /d) at initialization, we see that the difference in the updates Equations B.21 and B.22 at least of the same order as Equation B.21. Thus, in (poly(k)) (wgr Vs.e) (this necessarily occurs for definition of A and comparison to the bounds above), after which it follows from Equations B.21 and (Tx)
iterations, follows that < (w Vy,e-)
TA by
B.22 that "yyh" va.e) 0 ( g.r" uy,e)) /k) (and clearly TA = O(poly(k))) This proves the first part of the claim.
(Ta)
26

---

Under review as a conference paper at ICLR 2023
It remains to show that the off off-diagonal correlations also do not decrease by much, if they were to become negative that would potentially cause problems in Equation B.19 due to ReLU becoming
0. Using Equation B.20, we can form the following lower bound to ( (wy os,e)
1/ (wwiir Va,e
a-1
V (9. vele) 2 - why
(B.23)
k2pt 12pp--1
Now let denote the number of iterations starting from initialization that takes (wf.t. Ua,e
(0) to decrease to Wigrr Va.e) /polylog(k) for some fixed polylog(k) factor. Then it follows from Equation B.21 that in T iterations (t) Ey.c) has increased by at least a k/polylog(k) factor. As ryye.)
( (1)
result, we have that at Ta the correlation factor for some universal constant C, proving the claim.
(wy.rs Va.1) has decreased by at most a polylog(k))
Corollary B.18. For any class y lk], and any TA (for any TA satisfying the definition in Claim B.17), we have for any # and & [2]:
(wlf
uirs Uale) Ua,e)
mmxx
[E[2
Additionally, for all 8, with (wggrt Ws,t) we have that ( (wbbt Vs,e >0.
Proof of Corollary B. 18. The O(1/k) factor separation between the updates diagonal and off- diagonal correlations shown in Equations B.21 and B.22 continue to hold once we pass into the linear regime of ReLU. Furthermore, the logic used to prove the lower bound for positive correlations in Claim B.17 easily extends to showing that the correlations remain positive throughout training.
As noted above, the bound on the off off-diagonal correlations obtained in Claim B.17 and Corollary 1. much weaker than what was in Claim B.4, which is why we weakened the assumptions in Claim B. 15. We now prove the Midpoint Mixup analogues to Claims B.5, B.6, and B.7.
Claim B.19. Consider y E [R]] and such that max,eN, JE(N 9 (zz,5) @llog Then max,eN, JE[N] ge (z1,)) O(min JEIN o (z1,)))
Proof of Claim B.19. For any t satisfying the conditions of the claim, necessarily have that Corollary B.18 holds. As result, we have:
ΣΣC PE(M FE
wo(e) v,,0) (log k)
Σ
(with vo,s) (log(k))kk)
rE[m) RE(2]
Thus, we may disregard the off-diagonal correlations in considering the class y output on žij (i.e. we do not need to worry about the Iú part of zı.j), and the rest is identical to Claim B.5.
Claim B.20. For each y E [k]. let Ty denote the first iteration at which max4eN, JēN, 97, (ziz) Z log(k- 1). Then we have that T, (poly(k)) (for sufficiently large polynomial in k) and that minjeN,, j 9 (Za,) = (log k) for all 2 Ty,
Proof of Claim B.20. As in the proof of Claim B.6, applying Claim B.17 any class y yields (t)
the existence of correlation (w.rr, Uy, and = (ppoly(k)) such that w ggrr- Vyye" >
p/(62 5 61). And again, reusing the logic of Claim B.17 but replacing ( (w "y. in Equation B.21
with p/ 81) yields that in additional O(poly (k)) iterations have (n ( Uyge+ Uyge+
2p/ 51
(implying that w has reached the linear regime of ReLU on effectively all mixed points) while the off-diagonal correlations continue to lag behind by a (1/kk) factor.
(1)
27

---

Under review as a conference paper at ICLR 2023
At this point we may lower bound the update to ( grr Eyytt) as:
V in
(2, vyye-
z
Σ
0 - 200 20 (((((zz))))
N
NENy JEN,
"
Σ
Σ
e Sith ((gr ((.(3)))))))
N
JEN
( (w Cuig
Po max
7
Σ
Σ
/ ((((()))
N
pa-1
CENUJEN
(B.24)
Using Claim B.15, have that long maxieN,, JaN 26 (26.) < log(k - 1), get (using [N.]] ((N/k) for all E [k:)):
(9, Ey,
(n/kk)
(B.25)
This also implies, by the logic of Claim B.17, that the off-diagonal correlations (wgers Us.0) have updates that can be upper bounded as:
w (g, 3 ((1/k²)
(B.26)
Comparing Equations B.25 and B.26, we have that g (Zi,g) 2 log(k- - 1) (and, consequently,
gl (axz) 2 log(k - 1)) for least one mixed point ži, with i N, in O(poly(k)) iterations while the off-diagonal correlations are ((log(k)/k)) This also implies that minseN, JEIN] o (Zx,) = Alog k) by Claim 1.19. Finally, since Equation B.25 is positive, the class y network outputs remain (log k) for tort2 2 T3 (as again we cannot decrease below log(k: 1) by more than o(1) since the gradients o(1)).
Part II. Having analyzed the growth of diagonal and off-diagonal correlations in the initial stages of training. we now shift gears to focusing on the gaps between the correlations for each class. The key idea is that JMM will push the correlations for the features Vy and Uy closer together throughout training (so as long they are sufficiently separated), for every class y.
In order to prove this, we will rely on analyzing an expectation form of the gradient for JMM- As the expressions involved in this analysis can become cumbersome quite quickly, we will first introduce slew new notation to make the presentation of the results bit easier
Firstly, in everything that follows, we assume y/y, to be the better correlated feature at time t for every class y E [k] in the following sense:
Σ
(wlt) Mirs 1g.1
(((()) girs Uy.
(B.27)
TEB
rEB 6.2
Where the sets (tt) u.l are as defined in Equation B.16. Furthermore, as we will refer to the quantities
in Equation B.27 many times, we use c(0) and (1())
9.2 to denote the LHS and RHS of Equation B.27,
and (10) = 2.1 ((0 ( 5.2-
Now for the aforementioned expectation analysis we introduce several relevant random variables. We use 3y,p (for every y E [k)) to denote a random variable following the distribution of the signal coefficients for class y from Definition 4.1 and we further use By to denote a random variable representing the sum of Cp i.i.d. 8y.p- Similarly, we use zn,s to denote the average of two random variables following the distributions of class y and class points respectively. Finally, define A (Be, 3y) and ( By) as:
(33,,80) 2 260
(B.28)
A3 Bry) 6z ((9((zy,))
(B.29)
28

---

Under review as a conference paper at ICLR 2023
In context, this notation will imply that A₁ (9) =1- 2Φ" ((9o(Zy,))) (i.e. swapping the order of arguments changes which coordinate of the softmax is being considered).
Now we will first prove upper bound on the difference of gradient correlations in the linear regime, and then use these ideas to prove that correlations in the poly part of ReLU will still get significant gradient. After we have done that, will revisit this next claim show that the separation between feature correlations continues to decrease even after they reach the linear regime.
Claim B.21. Suppose that max (t)
Ollog k). Then for any class y [k: and any ri B
and r2
( we let:
(2), 2 (-V whire
JMM 2),
X0,
(B.30)
Wairg
After which we have that:
(2), S
Σ BR [A] (3) -Cob///)
Cp62/2)]
1:2
+ +
Es, [A2 By) -Cpd2/2)
122
+ 0 ( P6 (log)))
(B.31)
Proof of Claim B.21. Using the logic from Equation B.24 as well as the fact that in E B(t) and 12 E B (t) 9.2 (i.e. we are considering weights in the linear regime of ReL.U for each feature), we get:
((r1,2)) S N
Σ Σ (1 ((e((()))))
Σ
Bip- - Cpd2/2
iEN, JEN,
(PPr))(e))
E (29(((()))
Σ
Σ
3G - Cpò2/2
IEN, JEN,
O (P57 (log (logk)a k)a-1)
(B.32)
Now since we took N sufficiently large in Assumption 4.3, by concentration for bounded random variables we can replace the expressions on the RHS above with their expected values, as the deviation will be within O (P87 (log k)=-1) (with probability setting). However, the expectations will be over all of the random variables 8, for [kt, not just
O(1kk), consistent with our
the classes s and being mixed (or in the case of the second term above, just the class y).
Fortunately, observe that for the mixed point random variable 2y.s: the 3., for * y. s can only show up in the feature noise patches of zy.s. Thus, by an identical calculation to the one controlling the feature noise contribution to the gradient above (once again, refer to Equation B.24), that we may consider the expectation over just By and 3, while marginalizing out the other random variables and staying within the error term above, thereby obtaining Equation B.31.
We will now show that Es.,s, [Ar (3, Cp02/2)] is significantly negative long as the separation between feature correlations A/" is sufficiently large. Once again, to simplify notation even further, we will use B, B, - Cp02/2 and use P((,) to refer to its associated probability measure. Furthermore, we will use:
D (C c( (C.9) (.9)
In other words, DSf represents the difference in the linear outputs of classes y and s. With this in mind, we can prove the aforementioned result.
Claim B.22. Suppose that max c(u.] (log k). Let y be any class such that A (0) log k
(1), and suppose that there exists at least one class E [k] such that there is set U C 00, Cr(62/2-
29

---

Under review as a conference paper at ICLR 2023
61]] [-Cp(82/2 51), Cp(62/2 - 61)l with (P(3,) (() 2 0.01 (i.e. its measure is at least 0.01) and for all (a, E U we have:
(b - Cpd2
g. -a- Cp22pp)p2/
(B.33)
Then we have:
Eg [A1(3e, By) Cpóz/2)]=-
(B.34)
Proof of Claim B.22. We begin by first showing that the expectation on the LHS of Equation B.34 is negative. Indeed, this is almost immediate from the fact that By is a symmetric, mean zero random variable we need only show that A is monotonically decreasing in By.
From the definition of A1, we observe that it suffices to show that g (z is monotonically increasing in 3y. However, this straightforward see from the assumption that A (4) 2 .log k - this
implies that an e increase in By leads to ( log k) - O()) increase in 9t since the feature noise and weights that are in the polynomial part of ReLU can contribute at most O(1) by the logic of Claim B.19.
Now we need only show that the expectation is sufficiently negative. To do this, we will rely on the following facts, which will allow us to write things purely in terms of C and c(t) (i.e. disregarding the weights that are not in the linear regime):
of (2 E
c (+) B. 3.1 + (Cpd2 - P,) 9.2 s,C, + (Cpd2 By) + o()] + (Cro2 B.)C + + 0())
(B.35)
(Fy.o) E
[,,0%] + (Cpd2 3.) B, 8,C
(B.36)
(Zy.a) = 0 log k
for + 8.8
(B.37)
k
Which follow from Claim 1. and Corollary B.18 (alongside the assumption that maX C Olog k)) respectively. Now we perform the substitution 9i + De Cpo2(C g.1 c/ for all [k]. as this can be done without changing the value of o6" (g) Under this transformation we have that (using Equation B.35):
o (2y.s) E (3. Cr - Cr52D 0. /2, (3. Crō2/2)A - Cr 0(1) L (2y,)) E
(B, - Cp62/2) y (1)
(8, - Cpbp/2)4 + +(1]
(B.38)
(B.39)
Which isolates the correlation gap term ( Prior to proceeding further we will let A ه Enzu exp(oe (2y.))), so as prevent the equations follow from becoming too unwieldy. Now have:
Es,,, (B,, - Cpd2/2)]
- (v((6-66)
( explgi' (2g,0)) (B, - Cr02/2) dT (TP(B))
Scrb
A + exp(gi ( (z,.x)
A - exp
3 +0(1)) +
(Cr161/2-5)) (0)
By ddP(8))dP(B))
(Ce(h-69/7)
CP(61 ( + exp
(3 ( +0(1)
Cp(62/2-81)
Cp(óg/2-61) A! exp
(B,A(²)
+
By dP (3
(B.40)
Cp(61-8//2)
(1) + exp
(,,4
30

---

Under review as a conference paper at ICLR 2023
Where above we used Equation B.38 to get the upper bound in the last step. We next focus bounding the inner integral in Equation B.40. Using the symmetry of Sy. we have that:
Eg, (A...,3,))0 - Cpd2/2)]
Cp(ba/2-61)
A exp
-, + +0(1)
-eep A. - -ep (B
By dP(3,)
A! exp
-,,A, +0(1)
AS + cxp (8,4
2A (4)
(exp (8,A,0
exp ( + +0(1))
Cp(62/2-61)
B, d(((g))
(A 2 (
exp (3 eepp (-8, +0(1)) + exp(((11)
(B.41)
One can sanity check that Equation B.41 is bounded below by -1, we would expect. The only tricky aspect of Equation B.41 is the O(1) term the exponential, which can lead to a positive contribution (via a negative integrand) when Sy is close to 0. However, we can safely restrict the bounds of integration in Equation B.41 region [o. (6 61) )l for o = ((1/ log k) (with an appropriately chosen constant), as in such region the integrand is guaranteed to be positive since A 2 log o(1) Furthermore, this restriction does not us anything (like additional positive term), as concern from Equation B.41 is purely a consequence of how we bounded 9 and gi Indeed, by our earlier monotonicity argument it is clear that we can cut out the region corresponding to [-0,2 from the first line of Equation B.40 without decreasing the RHS.
Furthermore, we also have from A (t) N log k o(1) and Equation B.37 that Emgyy exp(g"' (zy,=))
O(1/kCréz-1) (after making the adjustment g +- 9 Cr 9.1 + c(t2) //2 that we did above).
Now using our assumption in the statement of the claim that Equation B.33 holds for some set U, we obtain:
Eg,,8, (A,(8,,,)(3, - Cp62/2)]
By exp (3 Crdd)) + exp (
(o(6)/2-61)
(c(66/2-61)
θ
( (d.)
Cp(01-82/2 -2/2)
-1(
(B.42)
Where the last line follows after restricting the bounds of integration of the two integrals to their intersections with V (this allows to disregard of in the asymptotic expression above via Equation B.39). This proves the claim.
We proved Claim B.22 in terms of By (the sum of the individual Sy.p) keep notation manageable (avoids Cp iterated integrals) and to more closely mirror the proof of Proposition 3.4. However, what will really use for our remaining analysis is the following corollary, which gives the same result as Claim B.22 but for each of the individual terms Sy.p- Below we use ECp By,s make explicit the dependence between the sum and each individual random variable Sy.p (so as to not mislead one to think of them as independent random variables).
Corollary B.23. Under the same conditions as Claim B.22, we have for every p [Cp]:
Cp Ai
Eg By..
Ay.
Σ (y,) (oo.p-88//)) (1)
(B.43)
Proof of Claim B.23. The proof follows identically to that of Claim B.23 (effectively the only change in the computations that By becomes By.p). as the functions A₁ and Ay satisfy the same mono- tonically increasing property in each of the i.i.d. By.p (and there are only Cp O(1) many). One could also have simply seen this from the symmetry of the By,p in Equation B.43; indeed, we expect Equation B.43 to only differ by factor 1 1/Cp from Equation B.34.
31

---

Under review as a conference paper at ICLR 2023
Now show using Corollary B.23 that there is significant gradient component towards correcting the separation between the feature correlations even when the second feature correlation is in the polynomial part of ReLU (which is where got stuck for significant number of classes in the ERM proof).
Claim B.24. Suppose that maX Cl.) Olog k). Let y be any class such that A log k:
(11). Then for any 12 & B,y satisfying Wy.r2 vy.2 2 0, long as there exists an 71 B 8,1 satisfying:
(0)
to murr
())
2
(B.44)
We have: Proof of Clain B.24. By near-identical logic to the steps leading to Equation B.31 and using Equation B.21, we obtain (following the same notation as before):
V 1iig
(9,
T 0
(B.45)
p
122
Cp Σ
Cp
20), 2 2ir2
Σ 8F4
Bp Bo Be. Px.Op
AI
By.p)
Σ
(62 By )a
pA12
p=1
p=1
(B.46)
Where above have absorbed a factor of 1/2 resulting from mixing classes and the feature noise component into the asymptotic term in front of the summation.
Now we break the of the proof into cases: whether Equation B.33 holds or not. In the former case, using Corollary B.23 and our assumption in the statement of the claim that Equation B.44 holds, we get (from linearity of expectation):
Cp Σ p=l
2(1 k22
θ
ER 1. By
A1
2
(B.47)
122
*
Now observing that Cov(A2(2) Epnr
(82 By.p)
>> 0 for every p. obtain:
(t) JMM (g) Xy,
warr
Cp
Cp
Ba Sp.cr
A1
Σ By. 80,D) E 1(02 - By,p) p=1
BFy pm
(B.48)
And the result then follows for this case from Equation B.47 and the fact that 11.. ((62 -- is dana-distrbbutiin-deeeeden constant.
For the second case when Equation B.33 does not hold, we have that every class + y there exists set U' c [0, Cp2 81)] (-Cp(62/22 -1)), Cp(02/2 di)] such that (IP(5)) P(.))) > 0.49 (note the total measure the set which U' subset of 0.5, by symmetry) and for all
(a, b) E U we have:
(b (b-Cr CpCos/2)) - Cr (() /2 (a- Crb,/2)
(B.49)
By anti-concentration assumption the Sy.p it immediately follows that D (loo k). from which we obtain that the the expectation terms in Equation B.46 are all @(1), so we are done.
Having proved Claim B.24, it remains to prove that both Equation B.44 and max y Ollog k) hold throughout training, as after doing so we can conclude that the second feature correla- tion will escape the polynomial part of ReLU and become sufficiently large in polynomially many training steps.
32

---

Under review as a conference paper at ICLR 2023
Claim B.25. For any y [k], [2], and [ml]. have that:
(++ (o 2). pue)) 0.99 wigh JMM X), vyy))))
So long as Lozy e 2 Eesy (2).g for all mixed points 2,,j with i E Ng.
Proof of Claim B.25. We proceed by brute force; namely, as long as is sufficiently small, we can prove that the gradient for JMM does not decrease too much between successive iterations. As notation is going to become cumbersome quite quickly, we will use the following place-holders for the gradient correlations time 1 and 1:
- G2
(o Jas 0g.e) Myir
We will now prove the result assuming r B(( (we will have the upper bound shown below with additional o(1) factors). We have that (compare to Equation B.24):
n.8. as the the case where r & B is strictly better u.
G S
1
Σ
Σ
8 ((Ee)(((())))) 6² ((9r(()))))
N2
JEN,
+ Σ
Σ
((9++1 Sz ((se(z())))
iEN, JENy
(w Vu.g
Po ma
3
d (((et1(()))))) 2²" (c
N
pp-1
IEN, JEN,
(B.50)
Let us now focus on the 6² (g(+1(21.))) -3 ((g((())) terms present in Equation B.50 above. We will just consider the first case above (mixing between class and non-y class), as the other analyses follow similarly. Furthermore, we will omit the 2i,j in what follows (in the interest brevity) and simply write o Additionally, similar Claim B.22, will use the notation A, E. exp(g.)
Now by the assumption in the statement of the claim, have that A++1 2 As, and since O(k) (number of weights per class), we have that 9'41 < g + ((knGG) (all of the updates for weights in the linear regime are identical and strictly larger than updates for those in the poly regime). Thus,
exple + (kk)G))
exp(9")
(g++1) (9)( S
Aot exp(o" + (kr)G)) A cxp(o"))
A exp(o)) (exp(((()))))(( -
A + A, exxpap + + exp(2ge + (kn)).)
< (kmG, +
O(k)),)
(B.51)
Where in the last line we again used the inequality exp(x) S 1 + x + x² for E [0, 1]. Plugging Equation B.51 into Equation B.50 (after similar calculations for the other two pieces of Equation B.50) yields the result (since again, O(1/poly(k)) is suitably small).
Corollary B.26. For any y [K],E E [2], r E [ml, and t, we have that:
(o: JMM (', My.e) 2 - my
JMM X)) 20
myr
Proof of Corollary B.26. For every class y. we have E exp 1ap(gf+122..0 (21,5) E.ty exp(ge for mixed points Zij with i Nr for t = 0 (see the proof Claim B.17), and the corollary then follows from induction argument and Claim B.25.
And with Corollary B.26 may prove that maxye(K) c( O(log k) holds throughout training.
33

---

Under review as a conference paper at ICLR 2023
Claim B.27. For all O(poly(k)). for any polynomial in k, have that maxge[M] Cy, Ollog k).
Proof of Claim B.27. The idea is to consider the sum of gradient correlations across classes, and show that the cross-class mixing term in this sum becomes smaller (as this would be our only concern already know the same-class mixing term will become smaller by the logic of Claim B.6).
As in the previous claims in this section, we will proceed with an expectation analysis. We will focus on the weights Wyrr that are in the linear regime for the feature 2/, for each class y, these are the only relevant weights for ((f)- Additionally, instead of considering the sum of gradient correlations (()) over all Wyr with r E 1.1. it will suffice for our purposes to just consider the sum of gradient correlations over classes while using an arbitrary weight Wyrr in the linear regime. Thus, will abuse notation slightly and use Wy.r to indicate such a weight for each class for the remainder of the proof of this claim (note that we do mean to imply by this that weight is in the linear regime for every class simultaneously, but rather that there exists some for every class that is in the linear regime).
Now in the same vein as Equation B.31 (referring again to Equation B.24), we have that:
h
Σ yе[k]
50 JA
2/.1
E Eg..8, [Ar(Bs, + Es,,9, (3,.
1:2
y=
++
k
+ +
E3
12
N=1
(B.52)
And we recall that is sufficiently large so that the deviations from the expectations above are negligible compared to the subtracted term. We have carefully paired the expectations in the leading term Equation B.52 so as to make use of the following fact:
A1 A1 By) = - (1 (,) + Ene(W](ure) exp(g) (2y.s))
(B.53)
Euel exp(ol' (zy,s))
The second term on the RHS of Equation B.53 is of course of P(P o(Pog/poly({) so long as gi (Zy.e) and/or gy (Zy.s) are greater than C Cloog k for large enough constant C, so we obtain:
Σ VE(K]
X) 5
E s, (A. 3a]]
wy.
122
y=1
1+1
+ +
Ep,
0 /poly(k))
1
(B.54)
Now again by the logic of Claim B.22 we have that Cov(AA(B,) By), -Ba) 0, it follows that:
k
(..
5
Es,
0 OPP
K2
(B.55)
And if of (2y.s) 2 C Clog l for a sufficiently large constant C, we have that the RHS above would be negative, which contradicts Corollary B.26, proving the claim.
We have now wrapped up all of the pieces necessary to prove Theorem 4.7. Indeed, we can now show that for every class the correlation with both features becomes large over the course of training.
Claim B.28. For every class y E [k)), in (p)oy(kk) iterations (for a sufficiently large polynomial in k)) we have that both c(I) 2(oog k) and c(() 2(log)).
34

---

Under review as a conference paper at ICLR 2023
Proof of Claim B.28. Claim B.20 guarantees C 2log k) in polynomially many iterations. If at this point C Qlog k), we done. If this is not the case, but we have B * 0, then we done by Claim B.2 Corollary B.26, and Claim B.27. On the other hand, if g y.2 0, then we have by Claim B.24 that in polynomially many iterations (as we can take T = (1/ d) by our setting) 1 * 0, after which we have reverted back to the previous case and we are still done.
Now we may conclude the overall proof by observing that Claims B.20 and B.27 in tandem imply that we achieve (and maintain) perfect training accuracy in polynomially many iterations, while Claim B.28 implies that both features are learned in the sense of Definition 4.5.
C PROOFS OF AUXILIARY RESULTS
C.I PROOFS OF LEMMA 3.2 AND PROPOSITION 3.3
Lemma 3.2. [Midpoint Mixup Optimal Direction] A linear model g satisfies the following:
lim (1M X² = inf JMM JM, 22) 7-220
(3.3)
If g has the property that for every class y we have (wy, Fyy,,)) = (w., U2 > 0 and (Wy, vs,ea) S 0 for every s * y and l E [2]. Furthermore, with probability 1- exp(-O(N)) (over the randomness of 2), the condition (Wy, "y.e,) = (Wa, Us.ea) is necessary for to satisfy Equation 3.3.
Proof. We first prove sufficiency. g satisfies the conditions in the Lemma, then we have for any data point (zi.yi that 9dd (a:) = (Wyg l/ + 22.,2) > O. We also have that g" (x;) = O for any + y/s (by the cross-class orthogonality condition). Letting C = (Wy. Vy,1 + fy,2) (note that this correlation is the same independent of y due to the conditions of the lemma), we then get:
exp(>C/2)
6" (Tg((Zz)) = (11g' (=1S)) O(k 2) 2exxp(r//2)
(C.I)
For any mixed point žij with ye + yj. Equation C.I tends to 1/2 7 + co, and one can easily check that this is the global optimal prediction for the classes y/ and yj on the Mixup point zi,j (for any such mixed point). Similarly, if 21,y is mixed point with y = yj. then Equation C.1 becomes the ERM case, we obtain the optimal prediction of for the correct class in the limit.
On the other hand, if there exists a pair of classes (y, s) with s + y and é,, lg [2] such that (wyr )p.e.) # (w., U,, then with probability 1 - eex((-((N) there exists a mixed point 2i in X (where yi- y. MJ = and j s) such that g" (=x.)) & g (zx,), and hence lim d² (rgg m())) of 1/2, so we cannot achieve the infimum of the Midpoint Mixup loss.
Proposition 3.3. For any distribution DA that is not point mass on 0, 1, or 1/2, and any linear model y satisfying the conditions of Lemma 3.2, we have that with probability 1 - exp(-O(N))
(over the randomness of X)) there exists an eo > 0 depending only on Dx such that:
Ja (g. X, Dx) Z inf JM (h, X, Dx) + co
(3.4)
Proof. Firstly, we observe (just from properties of cross-entropy):
inf (h, = -EA-PA log + (1 - A) log(1 - A)]
(C.2)
= Now suppose a model g satisfies the conditions of Lemma 3.2. Then we have that g" (xx)
nDو (xx,) C' > 0 for some constant C' and every pair (24, y:) and (z),
As before, with at least probability - exp(-O(A)) we have that there exist a pair of points (z,, yi) and (z,, 2;) in X with yı + yı- The Mixup loss restricted to this pair (for which we use the notation Ja (g. Zij, Dx)) is then:
Jg Jij, D2) -EANDA [Alog 6th (o(2),)) + (1 A) 6t: (g(2
(C.3)
35

---

Under review as a conference paper at ICLR 2023
Furthermore, we have that:
exp(AC)
8th (79(22) (A))) = O(k - 2) exp(XC) exp((1 XJC)
(C.44
From Equations C.3 and C.4 we can see that, since D2 is supported on more than just 0. 1, and 1/2, JM Zi.j Px) - oe as C + oo (Equation C.4 implies that in the limit 6"i (1g(2z,) (A)) can only take the values, O, 1, or 1/22.. It is also easy to see that the same behavior occurs if one considers C --00 0. Thus suffices constrain our attention C E [M1, M2] for some M1, M2 > depending only on Dx
However, Ju(y, Zi.j DX) - inf Ja (h, 2z.j, DA) > (note that inf JM (h, Zi,j+ Dx) inf Jur (h, X, Dx)) for all C E [M, Ma]. Since this is a continuous function of C over com- pact set, it must obtain a minimum greater than 0, and we may choose Eo to this minimum (rescaled by factor of ((1/N2),, thereby finishing the proof.
C.2 PROOFS OF PROPOSITIONS 3.4 AND 3.5
Proposition 3,4. [Mixup Gradient Lower Bound] Let y be any class such that 4g 2 log k. and suppose that both (wy: 0y.1) 2 and the cross-class orthogonality condition (ws, va,e) hold for all * u and l E [2). Then we have with high probability that:
1 R2
JM (222)
+
(3.5)
Proof. The idea of proof will be to analyze the gradient correlation with 2ly, 12y,2 and either show that this is significantly negative or, in the case where it not, the gradient correlation with 2,, is still significant. Firstly, using the cross-class orthogonality assumption and Calculation A.10, we can compute:
(1-20" 1 20th ((o( ((o(z.,)))) ( 8 _
(-V my JM, u222)
N2 Σ
Σ
2
iEN, J
Σ Σ
(1-0z
1 -
(C.5)
N2
KEN, N
Where above we used Ny to indicate the indices corresponding class y data points (as do in the proofs of the main results). Now using concentration of measure for bounded random variables and the fact that N is sufficiently large, we have from Equation C.5 that with high probability (and with poly(k) representing a very large polynomial in k):
wy JoJMM(y uy.1 -,y22)
5
EB BA,,, (B, 1/22]]
1
+
Eg, [A2 - + O(L/poly(k)) (C.6)
k2
Where we define the functions A and A2 as:
(1 By) A 20 ((g (2y,s)) A2(3,, 3,y)
(C.7)
2d ((9e(2g.))))
(C.8)
With Zy.s being a random variable denoting the sum of class y point and a class s point (distributed according to Definition 4.1). Note that Equations C.7 and C.8 are not abuses of notation the functions Aj and A2 depend only on the random variables 3, and By, since we can ignore the cross-class correlations due to orthogonality.
Let us immediately observe that the first two terms (the expectation terms) in Equation C.6 are bounded above by 0. This is due the fact that 8g 18--1/1 is symmetric, centered random variable and the functions A, and Ag are monotonically decreasing in B,, We will focus on showing that the first term is significantly negative, that will be sufficient for our purposes.
36

---

Under review as a conference paper at ICLR 2023
Now let E, (wy: Uy,1) (Wy, ly,2) and perform the transformation 9i < gi -Ey/2 2 for all s E [k] (note this doesn't change the value the softmax outputs). Under this transformation we have that of (zy,s) (8,- 1y which isolates the gap term Ay.
For further convenience, let us use A, - Enply exp(g" (Z),))()z)))) (Zy,a)), and observe that A, depends only on 33 due to orthogonality. Using the change of variables By By 1/2 we can then compute the expectation in the first term of Equation C.6 as:
1
0.9
0.9
A, - explai d (2n,0)) (B - 1/2) d,, ds
Bg.,8 (A,(8.,,,)66, - 1/2)]
0.64
A, + exp(9) (2
0.1
0.1
0.1 A. exp (,
0.4
By dBg d3,
(C.9)
0.64
00.1
A, + exp 8 8,
We will focus on the inner integral in Equation C.9. Using the symmetry of Bgy we have that:
0.1 A, - exp By
1
Eg, (8,,
1/2)]
d,3y
0.8
A, + exp
(ByD,)
0.4
Ag exp
- -,
Ag - exp (3 (xp((()))))
By dg
0.8 0
Ag exp
(-8,A,)
A. + exp (3
0.4
A. exp
(ByAw)
exp
(By
By d3
0.8
Jo
A3
+ Ag
eep
(B
+ exp
(-8g4n) +
(C.10)
From orthogonality assumption and the facts that Ag log and (wy, 0y,2) 2 0, we have that Eujy, exp(9" (2y.))) = O()). Additionally, if we let:
yy. (Cy.1 C2.2) (C.,1 C2,2)
Then get from Equation C.10 that:
B3.,s (A. (By 1/2)]
By exp B,D. - Dy,o/) exp (B,A,)
0.1
0.4
+
dī, dE,
0.4
exp 23,A, Dy.e) + texp (B.A. - D2,o/2) exp (3
(C.11)
Now we consider two cases. First, suppose there exists set U c [-0.4,0.] x [0, 0.4] with probability measure at least 0.01 such that for all (5,, B,)) U:
B.D. Dg.e/ S B,A,
Then we immediately have that Equation C. ((1). On the other hand, if this is not the case, one can see that 18,,.,9, [A, (Bx, (,),,,,)() (1) - 3,)] ((1), so we are done.
Proposition 3.5. [ERM Gradient Upper Bound] For every y E [k], assuming the same conditions as in Proposition 3.4, if Ay 2 C' log & for any C > 0 then with high probability we have that:
V wy [g. u. S
h0.C-1
(3.6)
Proof. From the facts that (wy, 0g,2) 0 and A, 2 C' Cog every i E Ny (where By represents the coefficient in front of u in x,). Since By,s E [0.1, 0.9], we immediately have the result from the logic of Claim B.2 and Calculation A.9.
have that (x;) 2 By log k for
37

---

Under review as a conference paper at ICLR 2023
C.3 PROOF OF PROPOSITION 4.4
Proposition 4.4. There exists D satisfying all of the conditions of Definition 4.1 and Assumption 4.3 such that with probability least 1 k2 exp(O(-N/<*)), for any classifier h R R of the form h (xx) = [p((P) (W,,x(2) and any X consisting of i.dd. draws from D, there exists point (z, y) E X' and a class # such that h²(x)) 2 hr))..
Proof. For hyperparameters, we can choose 6 = 82 81 = 1, consistent with Assumption 4.3. For the distribution D, for each point (2 24), sample class 8 E [K] \ Ys uniformly and choose it to be the single class used in the feature noise patches for ry. This clearly falls within the scope of Definition 4.1.
k-1 while being
Now in ii.d samples from D as specified above, we have with probability at least 1 222 exp ( (-NN1))) (Chernoff bound, in Claim B.1) that sample with each possible pair y/ E [k] of signal and noise classes exists. Suppose now that there exist classes y/, E [K] such that u & argmax,e[4 (We, v, + U9.2). This necessarily implies that u # max g" (x) for all points with label u but having feature noise class y (since the order of the sum of the feature noise is e( @(v)), which gives the result.
On the other hand, if there does not exist such a class pair y, u, then we are also done as that implies all of the feature correlations are the same.
38


Reviews:
[{"review_id": "jK95Ubo3e-6", "reviewer": "Reviewer_Pxqu", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper studies a specific instance of Mixup data augmentation: Midpoint Mixup. It provides a theoretical characterization of the learning dynamics of neural networks with mixup-augmented data, and proved how it can improve feature learning by helping the neural network pick up diverse features, under the multi-view data framework proposed by Allen-Zhu and Li.\n\n\n[1] Zeyuan Allen-Zhu, Yuanzhi Li, Towards Understanding Ensemble, Knowledge Distillation and Self-Distillation in Deep Learning", "summary_of_the_review": "This paper studies an interesting instance of Mixup augmentation and obtained some decent theory results with the help of an established theoretical framework, but the authors are not good at justifying their choice of study and their theoretical claims, and they fail at explaining the key technical ingredient of the paper, as well as how well their assumptions can connect to practical scenarios. ", "strength_and_weaknesses": "Strength: Most of the paper is clearly written and it provides a rigorous analysis of the subject it studies. Mixup is an interesting data augmentation technique and tries to explain why it can help in improving the feature learning of neural nets. The theory result goes beyond the previous works that are based on linear models (as they are based on nonlinear neural networks) and is reasonable due to the feature learning technique it uses. It might be an important paper if the authors can fix the issues I find below.\n\nWeaknesses: \n1. The authors did not provide sufficient motivation why they specifically study Midpoint Mixup rather than the general mixup augmentations. Moreover, they did not provide an empirical comparison between the original mixup and the midpoint mixup, making it hard to justify their choice of study. A potential improvement is to explain the technical difficulty of analyzing the original mixup, and justifies that the midpoint mixup is competitive (or particularly interesting).\n2. Assumption 3.10, which seems to be the key assumption to migrate the proof technique in [2] to the setting of this paper, is very technical, without intuitive explanation, and without support from empirical evidence. It is hard to make sense of this assumption as assembling any real-world structure. This weakens the validity of the proposed explanation as the theory result could be just an artificial construction.\n3. The experimental results in this paper are not informative, as many of their metrics are too low. I do not understand why this paper constructs such special data by mixing images using Dirichlet distribution (see equation 4.1 and Fig 1) to compare ERM with Midpoint-Mixup. Moreover, the test errors often exceed 60%-90% for experiments on CIFAR-10/100 when L>1 (that is over the special dataset rather than the original Cifar10/100), meaning that the model struggles a lot to predict the mixup label. I don't think these experiments are sufficient to support the authors' choice of study or their theoretical claims.\n\n\n[1] Ruoqi Shen, Sebastien Bubeck, Suriya Gunasekar, Data Augmentation as Feature Manipulation\n[2] Zeyuan Allen-Zhu, Yuanzhi Li, Towards Understanding Ensemble, Knowledge Distillation and Self-Distillation in Deep Learning", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "a02INeviT1", "reviewer": "Reviewer_9dJZ", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This work adopts the feature learning framework by Allen-Zhu and Li (2021), and further introduces latent orthogonal feature views. The authors prove that under certain assumptions on features, a two-layer smooth-ReLU network can learn all features with midpoint mixup, whereas ERM can only learn one feature.", "summary_of_the_review": "In general, I believe this work contributes significantly to the feature learning topic.\nI hope the authors could fix my concerns above and add more clarifications.", "strength_and_weaknesses": "Strength:\n1. On top of Allen-Zhu and Li (2021), the authors further decompose the input into multiple orthonormal feature views. This allows the authors to analyze the contribution of each individual feature and their gradient correlations.\n2. Although the change of data settings (section 3.1) and assumptions (3.10) is simple, the authors provide significant efforts in proving the feature learning under the mixup loss.\n\nWeakness:\n* Theory:\n1) In Eq. B.34, do you need a lower bound or dependence on the training iterations, such that the feature correlations will not be over-corrected?\n2) Could you add more intuitive explanations of the meaning of function f in Assumption 3.10? Is it possible to analytically prove the monotonicity of f?\n\n* Experiments:\n1) Thm 3.9 and 3.11 did not directly imply better test accuracy (generalization). Instead, is it possible to empirically verify the one or all feature(s) (v) learned (defined by 3.8) by the Midpoint Mixup?\n2) It would be better to add variances in Table 1.\n3) Many test errors in Table 1 are far worse than random guess (90% error rate). From Eq. 4.1, I did not see any noise is introduced. What is the reason for this large test error?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "Dd6yeNe0tBn", "reviewer": "Reviewer_uhyy", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper focuses on mid-point mixup (a variant of mixup) and provides theoretical analysis on the effect of using mid-point mixup on multi-view data. The theoretical result is specific to the setting of the 2-layer convolutional networks and the datasets where samples at each class have 2 features. Under this setting, this paper shows that empricial risk minimization learn only single feature at each class, while mid-point mixup learns both features at each class. ", "summary_of_the_review": "This paper focuses on an interesting problem, but has much room for improvement in terms of organizing the results, validating the assumption, and explaining how the experimental results are helpful in practical settings. ", "strength_and_weaknesses": "- Strength\n    - The direction of analyzing mid-point mixup on mulit-view data is new and interesting.\n\n- Weakness\n\n    - Theory\n        - Why is the assumption of Lemma 2.1 (z_{i,j} is unique for all i,j pair) true in practice? Simple 2-dimensional dataset where each class follows Gaussian distribution cannot follow this assumption. (how about high-dimensional case?)\n        - It seems like Lemma 2.1 is important to show that Midpoint Mixup is not poor, but the proof of Lemma 2.1 is not fully understood on my side. (I put the details in the questions section)\n\n    - Experiments\n        - This paper shows that when L is large, Mid-point Mixup is having less error than ERM. Of course if L is large, it will go to the ideal multi-view dataset this paper is assuming, but then why don’t we even start with real data? Instead, making a synthetic multi-view data to show the effectiveness of the theory seems more reasonable. I personally feel the experiments in this paper is similar to just “simulation” instead of real-world experiments.\n        - In Table 1, the authors are comparing “test error 96%” versus “test error 97%”, when we have 100 classes. What does it mean to increase 1% test error from “nearly random guessing”? Does these empirical results have any positive impact on the practical deployment of mid-point mixup?\n\n    - Writing\n        - The assumptions and results are less organized (thus I have many questions below). The definition 3.4 has no intuitive explanation or visualization.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject