Towards understanding how momentum improves generalization in deep learning

Introduction

It is commonly accepted that adding momentum to an optimization algorithm is required to optimally train a large-scale deep network. Most of the modern architectures maintain during the training process a heavy momentum close to 1 (Krizhevsky et al., 2012; Simonyan & Zisserman, 2014; He et al., 2016; Zagoruyko & Komodakis, 2016). Indeed, it has been empirically observed that architectures trained with momentum outperform those which are trained without (Sutskever et al., 2013). Several papers have attempted to explain this phenomenon. From the optimization perspective, (Defazio, 2020) assert that momentum yields faster convergence of the training loss since, at the early stages, it cancels out the noise from the stochastic gradients. On the other hand, (Leclerc & Madry, 2020) empirically observes that momentum yields faster training convergence only when the learning rate is small. While these works shed light on how momentum acts on neural network training, they fail to capture the generalization improvement induced by momentum (Sutskever et al., 2013). Besides, the noise reduction property of momentum advocated by (Defazio, 2020) contradicts the observation that, in deep learning, having a large noise in the training improves generalization (Li et al., 2019; HaoChen et al., 2020). To the best of our knowledge, there is no existing work which theoretically explains how momentum improves generalization in deep learning. Therefore, this paper aims to close this gap and addresses the following question:

Why does momentum improve generalization? What is the underlying mechanism of momentum improving generalization in deep learning?

In computer vision, practitioners usually train their architectures with stochastic gradient descent with momentum (SGD+M). It is therefore natural to investigate whether the generalization improvement induced by momentum is tied to the stochasticity of the gradient. We train a VGG-19 (Simonyan & Zisserman, 2014) using SGD, SGD+M, gradient descent (GD) and GD with momentum (GD+M) on the CIFAR-10 image classification task. To further isolate the regularization effect of momentum, we turn off data augmentation and batch normalization. Figure 1 displays the training loss and test accuracy of the four models. Not only momentum improves generalization in the full batch setting but the generalization improvement increases as the batch size is larger. Motivated by this empirical observation, we focus on the contribution of momentum in gradient descent. We emphasize that this setting allows to isolate the contribution of momentum on generalization since the stochastic gradient noise influences generalization (Li et al., 2019; HaoChen et al., 2020).

Given the success of momentum in different deep learning tasks such as image classification (Simonyan & Zisserman, 2014; He et al., 2016) or language modelling (Vaswani et al., 2017; Devlin et al., 2018), we start our investigation by raising the following question:

Does momentum unconditionally improve generalization in deep learning?

We respond in the negative to this question through the following synthetic binary classification example. We consider a Gaussian dataset where each data-point is sampled from a standard normal distribution. We generate the labels using multiple teacher networks. Starting from the same initialization, we train several student networks on this dataset using GD and GD+M and compare their test accuracies in Table 1. Whether the target function is simple (linear) or complex (neural network), momentum does not improve generalization for any of the student networks. The same observation holds for SGD/SGD+M as shown in Appendix A. Therefore, momentum does not always lead to a higher generalization in deep learning. Instead, such benefit seems to heavily depend on both the structure of the data and the learning problem.

Motivated by the aforementioned observations, this paper aims to determine the underlying mechanism produced by momentum to improve generalization. Our work is a first step to formally understand the role of momentum in deep learning. Our contributions are divided as follows:

In Section 2, we empirically confirm that momentum consistently improves generalization when using different architectures on a wide range of batch sizes and datasets. We also observe that as the batch size increases, momentum contributes more significantly to generalization.

In Section 3, we introduce our synthetic data structure and learning problem to theoretically study the contribution of momentum to generalization.

In Section 4, we present our main theorems along with the intermediate lemmas. We theoretically show that a 1-hidden layer neural network trained with GD+M on our synthetic dataset is able to generalize better than the same model trained with GD. Above all, we rigorously characterize the mechanism by which momentum improves generalization. A sketch of the proof is presented in Section 5 and Section 6.

The previous experiments suggest that momentum improves generalization in CIFAR-10 while it does not for Gaussian datasets. This means that this generalization improvement must be specific to the data structure and the learning problem. In Section 3, we devise a binary classification problem where the data are linearly separated by a hyperplane directed by the vector $\bm{w^{*}}$ as depicted in Figure 2. We refer to this vector as the feature and the goal is to learn it. Each data-point is a vector constituted of a single signal patch equal to $\theta\bm{w^{*}}$ and of multiple noise patches. For $\mu\ll 1$, we assume that with probability $1-\mu$, the sampled data-point has large margin i.e. $\theta=\alpha\gg 1$ while it has small margin i.e. $\theta=\beta\ll 1$ with probability $\mu$. The noise patches are Gaussian random vectors with small variance. We underline that all the examples share the same feature but differ in their margins. Our dataset can be viewed as an extreme simplification of real-world object-recognition datasets with data of different level of difficulty. Indeed, images are divided into signal patches that are helpful for the classification such as the nose of a dog and noise patches e.g. the background of an image that are uninformative. Besides, the signal patch may be strong i.e. the feature is clearly visible or weak when the feature is indistinguishable e.g. in a car image, the wheel feature is more or less visible depending on the orientation of the car.

This paper proposes a theory to explain why momentum improves generalization. The following informal theorems characterize the generalization of a 1-hidden layer convolutional neural network trained with GD and GD+M on the aforedescribed dataset. They dramatically simplify Theorem 4.1 and Theorem 4.2 but highlight the intuitions.

There exists a dataset of size $N$ such that a 1-hidden layer (over-parameterized) convolutional network trained with GD:

initially only learns the $(1-\mu)N$ large margin data.

has small gradient after learning these data.

memorizes the remaining small margin data from the $\mu N$ examples.

The model thus reaches zero training loss and well-classifies the large margin data at test. However, it fails to classify the small margin data because of the memorization step during training.

There exists a dataset of size $N$ such that a one-hidden layer (over-parameterized) convolutional network trained with GD+M:

initially only learns the $(1-\mu)N$ large margin data.

has large historical gradients that contain the feature $\bm{w^{*}}$ present in small margin data.

keeps learning the feature in the small margin data using its momentum historical gradients.

The model thus reaches zero training error and perfectly classify large and small margin data at test.

Theorem 1.1 and Theorem 1.2 indicate that since the large margin data are dominant, the two models learn in priority these examples to decrease their training losses. Since the training loss is the logistic one, this implies that the gradient terms stemming from the large margin data thus become negligible. Consequently, the current gradient becomes a sum of the small margin data gradients. Thus, it is in the direction of $\beta\bm{w^{*}}$ (signal patch) and Gaussian vectors $\mathbf{g}$ (noise patches). Since $\|\beta\bm{w^{*}}\|_{2}\ll\|\mathbf{g}\|_{2}$, the current gradient is noisy. Therefore, the GD model keeps decreasing its training loss and memorizes the small margin data. On the other hand, contrary to GD, GD+M updates its weights using a weighted average of the historical gradients. In particular, it has large past gradients (stemming from large margin data) that are in the direction $\alpha\bm{w^{*}}$. Therefore, even though the current gradient is noisy, the GD+M uses its historical gradients to learn the small margin data since all the examples share the same feature. We name this process historical feature amplification and believe that it is key to understand why momentum improves generalization.

Our theory relies on the ability of momentum to well-classify small margin data. We first perform experiments in our theoretical setting described in Section 3. We set the dimension to $d=30$, the number of training examples to $N=20000$, the test examples to $2000$. Regarding the architecture, we set the number of neurons to $m=5$ and the number of patches to $P=5$. The parameters $\alpha,\beta,\mu$ are set as in Section 3. We refer to stochastic gradient descent optimizer with full batch size as GD/GD+M. Note that for each optimizer, we grid-search over stepsizes to find the best one in terms of test accuracy. We trained the models for 50 epochs. We set the momentum parameter to 0.9. We apply a linear decay learning rate scheduling during training. Figure 3 shows that the models trained with GD and GD+M get zero training loss and well-classify large-margin data at test time. Contrary to GD, GD+M well-classifies small margin data.

To further validate our theory, we artificially generate small-margin data in CIFAR-10. We first randomly sample 10% of the training and test images. As displayed in 4(c), for each image, we randomly shuffle the RGB channels. We train a VGG-19 without data augmentation nor batch normalization. While the GD and GD+M models reach 100% training accuracy, Figure 4 shows that GD+M gets higher test accuracy than GD. Above all, GD+M generalizes better than GD on small-margin data as the accuracy drop factor for GD+M is $79.47/53.30=1.49$ while for GD, this drop factor is $68.33/34.80=1.96$.

Related Work

A long line of work consists in understanding the convergence speed of momentum methods when optimizing non-convex functions. (Mai & Johansson, 2020; Liu et al., 2020; Cutkosky & Mehta, 2020; Defazio, 2020) show that SGD+M reaches a stationary point as fast as SGD under diverse assumptions. Besides, (Leclerc & Madry, 2020) empirically shows that momentum accelerates neural network training for small learning rates and slows it down otherwise. Our paper differs from these works as we work in the batch setting and theoretically investigate the generalization benefits brought by momentum (and not the training ones).

Momentum-based methods such as SGD+M, RMSProp (Tieleman & Hinton, 2012) and Adam (Kingma & Ba, 2014) are standard in deep learning training since the seminal work of (Sutskever et al., 2013). Although it is known that momentum improve generalization in deep learning, only a few works formally investigate the role of momentum in generalization. (Leclerc & Madry, 2020) empirically report that momentum yields higher generalization when using a large learning rate. However, they assert that this benefit can be obtained by applying an even larger learning rate on vanilla SGD. We suspect that this is due to data augmentation and batch normalization (Ioffe & Szegedy, 2015) which are known to bias the algorithm’s generalization (Bjorck et al., 2018). To our knowledge, our work is the first that theoretically investigates the generalization of momentum in deep learning.

Numerical performance of momentum

To evaluate the contribution of momentum to generalization, we conducted extensive experiments on CIFAR-10 and CIFAR-100 (Krizhevsky et al., 2009). We used VGG-19 (Simonyan & Zisserman, 2014) and Resnet-18 (He et al., 2016) as architectures. In this section, we only present the plots obtained with VGG-19 and invite the reader to look at Appendix A for the Resnet-18 experiments.

In all of our experiments, we refer to the stochastic gradient descent optimizer with batch size 128 as SGD/SGD+M and the optimizer with full batch size as GD/GD+M. We turn off data augmentation and batch normalization to isolate the contribution of momentum to the optimization. Note that for each algorithm, we grid-search over stepsizes and momentum parameter to find the best one in terms of test accuracy. We train the models for 300 epochs. The stepsize is decayed by a factor 10 at epochs 190 and 265 during training. All the results are averaged over 5 seeds.

Momentum improves generalization. Figure 5 shows the performance of GD, GD+M, SGD and SGD+M when training a VGG-19 on CIFAR-100. We observe that GD+M/SGD+M consistently outperform GD/SGD. Besides, we highlight that the generalization improvement induced is more significant for GD than for SGD. Similar observations hold for Resnet-18 (see Appendix A).

Influence of batch size. 6(a) shows the test accuracy of a VGG-19 trained on CIFAR-10 with the stochastic gradient descent optimizer on a wide range of batch sizes. We compare the generalization obtained with momentum and without. We remark that momentum does not improve generalization when the batch size is tiny. However, as the batch size increases, the gap between the momentum curve and the no momentum one widens.

Batch normalization and data augmentation. Practitioners usually add batch normalization and data augmentation when training their architectures. 6(b) displays the test accuracy obtained when training a VGG-19 with these two regularizers. We remark that they inhibit the generalization improvement of momentum for small and middle range batch sizes. For large batch sizes, momentum slightly improves generalization. Additional experiments on the influence of batch normalization and data augmentation are in Appendix A.

Setting and algorithms

In this section, we introduce our theoretical setting to analyze the implicit bias of momentum. We first formally define the data distribution sketched in the introduction and the neural network model we use to learn it. We finally present the GD and GD+M algorithms.

We define a data distribution $\mathcal{D}$ where each sample consists in an input $\bm{X}$ and a label $y$ such that:

Uniformly sample the label $y$ from $\{-1,1\}.$

Signal patch: one patch $P(X)\in[P]$ satisfies

$c$ is distributed as $c=\alpha y$ with probability $1-\mu$ and $c=\beta y$ otherwise.

Noisy patches: $\bm{X}[j]\sim\mathcal{N}(0,(\mathbf{I}_{d}-\bm{w^{*}}\bm{w^{*\top}})\sigma^{2}),$ for $j\in[P]\backslash\{P(\bm{X})\}$.

When $\lambda>0$, if the loss $\frac{1}{N}\sum_{i=1}^{N}\log\left(1+\exp\left(-y_{i}f_{\bm{W}}(\bm{X}_{i})\right)\right)$ is convex, then there is a unique global optimal solution, so the choice of optimization algorithm does not matter. In our case, due to the non-convexity of the training objective, GD+M converges to a different (approximate) global optimal compared to GD, with better generalization properties.

We assess the quality of a predictor $\bm{\widehat{W}}$ using the classical 0-1 loss used in binary classification. Given a sample $(\bm{X},y),$ the individual test (classification) error is defined as $\mathscr{L}(\bm{X},y)=\mathbf{1}\{f_{\bm{\widehat{W}}}(\bm{X})y<0\}.$ While $\mathscr{L}$ measures the error of $f_{\bm{\widehat{W}}}$ on an individual data-point, we are interested in the test error that measures the average loss over data points generated from $\mathcal{D}$ and defined as

We solve the training problem (P) using GD and GD+M. GD is defined for $t\geq 0$ by

where $\eta>0$ is the learning rate. On the other hand, GD+M is defined by the update rule

where $\bm{g}^{(0)}=\bm{0}_{m\times d}$ and $\gamma\in(0,1)$ is the momentum factor. We now detail how to set parameters in GD and GD+M.

Our 3.1 matches with the parameters used in practice as the weights are generally initialized from Gaussian with small variance and momentum is set close to 1 (Sutskever et al., 2013).

Main results

We now formally state our main theorems regarding the generalization of models trained using (GD) and (GD+M) in the setting described in Section 3. We first introduce some notations.

Let $r\in[m]$, $i\in[N]$, $j\in P\backslash\{P(\bm{X}_{i})\}$ and $t\geq 0.$ Our analysis tracks $\bm{w}_{r}^{(t)}$ the $r$-th weight of the network, $\nabla_{\bm{w}_{r}}\widehat{L}(\bm{W}^{(t)})$ the gradient of $\widehat{L}$ with respect to $\bm{w}_{r}$, $\bm{g}_{r}^{(t)}$ the momentum gradient defined by $\bm{g}_{r}^{(t+1)}=\gamma\bm{g}_{r}^{(t)}+(1-\gamma)\nabla_{\bm{w}_{r}}\widehat{L}(\bm{W}^{(t)})$. We introduce the projection of these objects on the feature $\bm{w^{*}}$ and noise patches $\bm{X}_{i}[j]$:

– Projection on $\bm{w^{*}}$: $c_{r}^{(t)}=\langle\bm{w}_{r}^{(t)},\bm{w^{*}}\rangle$.

– Projection on $\bm{X}_{i}[j]:$ $\Xi_{i,j,r}^{(t)}=\langle\bm{w}_{r}^{(t)},\bm{X}_{i}[j]\rangle$.

– Total noise: $\Xi_{i}^{(t)}=\sum_{r=1}^{m}\sum_{j\in[P]\backslash\{P(\bm{X}_{i})\}}y_{i}(\Xi_{i,j,r}^{(t)})^{3}.$

– Maximum signal: $c^{(t)}=\max_{r\in[m]}c_{r_{\max}}^{(t)}$.

Lastly, we define the negative sigmoid $\mathfrak{S}(x)=1/(1+e^{x}).$

We now provide our first result which states that the learner model trained with GD does not generalize well on $\mathcal{D}$.

Assume that we run GD on P for $T$ iterations with parameters set as in 3.1. With high probability, the weights learned by GD

Intuitively, the training process of the GD model is described as follows. Given $|\mathcal{Z}_{1}|\gg|\mathcal{Z}_{2}|$ and our choice of parameters for $\alpha,\beta,\sigma$, the gradient points mainly in the direction of $\bm{w^{*}}$. Therefore, GD eventually learns the feature in $\mathcal{Z}_{1}$ (Lemma 5.1) and the gradients from $\mathcal{Z}_{1}$ quickly become small. Afterwards, the gradient is dominated by the gradients from $\mathcal{Z}_{2}$ (Lemma 5.2). Because $\mathcal{Z}_{2}$ has small margin, the full gradient is now directed by the noisy patches. It implies that GD memorizes noise in $\mathcal{Z}_{2}$ (Lemma 5.4). Since these gradients also control the amount of remaining feature to be learned (Lemma 5.3), we conclude that the GD model partially learns the feature and introduces a huge noise component in the learned weights. We provide a proof sketch of Theorem 4.1 in Section 5. On the other hand, the model trained with GD+M generalizes well on $\mathcal{D}$.

Assume that we run GD+M on (P) for $T$ iterations with parameters set as in 3.1. With high probability, the weights learned by GD+M

Intuitively, the GD+M model follows this training process. Similarly to GD, it first learns the feature in $\mathcal{Z}_{1}$ (Lemma 6.1). Contrary to GD, the momentum gradient is still highly correlated with $\bm{w^{*}}$ after this step (Lemma 6.2). Indeed, the key difference is that momentum accumulates historical gradients. Since these gradients were accumulated when learning $\mathcal{Z}_{1}$, the direction of momentum gradient is highly biased towards $\bm{w^{*}}$. Therefore, the GD+M model amplifies the feature from these historical gradients to learn the feature in small margin data (Lemma 6.3). Subsequently, the gradient becomes small (Lemma 6.4) and the GD+M model manages to ignore the noisy patches (Lemma 6.5) and learns the feature from both $\mathcal{Z}_{1}$ and $\mathcal{Z}_{2}.$ We provide a proof sketch of Theorem 4.2 in Section 6.

Our analysis is built upon a decomposition of the updates (GD) and (GD+M) on $\bm{w^{*}}$ and $\bm{X}_{i}[j]$. The projection of the vanilla and momentum gradients along these directions are

$\mathscr{G}_{r}^{(t)}=\langle\nabla_{\bm{w}_{r}}\widehat{L}(\bm{W}^{(t)}),\bm{w^{*}}\rangle$ and $\mathcal{G}_{r}^{(t)}=\langle\bm{g}_{r}^{(t)},\bm{w^{*}}\rangle.$

$\texttt{G}_{i,j,r}^{(t)}=\langle\nabla_{\bm{w}_{r}}\widehat{L}(\bm{W}^{(t)}),\bm{X}_{i}[j]\rangle$ and $G_{i,j,r}^{(t)}=\langle\bm{g}_{r}^{(t)},\bm{X}_{i}[j]\rangle.$

We now define the projected updates as follows:

Analysis of GD

In this section, we provide a proof sketch for Theorem 4.1 that reflects the behavior of GD with $\lambda=0$. A more detailed proof extending to $\lambda>0$ can be found in the Appendix.

At the beginning of the learning process, the gradient is mostly dominated by the gradients coming from the $\mathcal{Z}_{1}$ samples. Since these data have large margin, the gradient is thus highly correlated with $\bm{w^{*}}$ and $c_{r}^{(t)}$ increases as shown in the following Lemma.

For all $r\in[m]$ and $t\geq 0$, (1) is simplified as:

Lemma 5.3 implies that quantifying the decrease rate of $\nu_{2}^{(t)}$ provides an estimate on the quantity of feature learnt by the model. We remark that $\nu_{2}^{(t)}=\mathfrak{S}(\beta^{3}\sum_{s=1}^{m}(c_{s}^{(t)})^{3}+\Xi_{i}^{(t)})$ for some $i\in\mathcal{Z}_{2}$. We thus need to determine whether the feature or the noise terms dominates in the sigmoid.

We now show that the total correlation between the weights and the noise in $\mathcal{Z}_{2}$ data increases until being large.

Let $i\in\mathcal{Z}_{2}$, $j\in[P]\backslash\{P(\bm{X}_{i})\}$ and $r\in[m]$. For $t\geq 0$, (2) is simplified as:

By Lemma 5.4, the noise $\Xi_{i}^{(t)}$ dominates in $\nu_{2}^{(t)}$. Consequently, the algorithm memorizes the $\mathcal{Z}_{2}$ data which implies a fast decay of $\nu_{2}^{(t)}$.

Combining Lemma 5.5 and Lemma 5.3, we prove that GD partially learns the feature.

Lemma 5.4 and Lemma 5.6 respectively yield the first two items in Theorem 4.1. Bounds on the training loss and test errors are obtained by plugging these results in (P) and (TE).

Analysis of GD+M

In this section, we provide a proof sketch for Theorem 4.2 that reflects the behavior of GD+M with $\lambda=0$. A proof extending to $\lambda>0$ can be found in the Appendix.

Similarly to GD, by our initialization choice, the early gradients and so, momentum gradients are large. They are in the span of $\bm{w^{*}}$ and therefore, the GD+M model also increases its correlation with $\bm{w^{*}}$.

Contrary to GD, GD+M has a large momentum that contains $\bm{w^{*}}$ after Step 1.

Lemma 6.2 hints an important distinction between GD and GD+M: while the current gradient along $\bm{w^{*}}$ is small at time $\mathcal{T}_{0},$ the momentum gradient stores historical gradients that are spanned by $\bm{w^{*}}$. It amplifies the feature present in previous gradients to learn the feature in $\mathcal{Z}_{2}$.

With this fast convergence, Lemma 6.4 implies that the correlation of the weights with the noisy patches does not have enough time to increase and thus, remains small.

Lemma 6.3 and Lemma 6.5 respectively yield the two first items in Theorem 4.2.

Discussion

Our work is a first step towards understanding the algorithmic regularization of momentum and leaves room for improvements. We constructed a data distribution where historical feature amplification may explain the generalization improvement of momentum. However, it would be interesting to understand whether this phenomenon is the only reason or whether there are other mechanisms explaining momentum’s benefits. An interesting setting for this question is NLP where momentum is used to train large models as BERT (Devlin et al., 2018). Lastly, our analysis is in the batch setting to isolate the generalization induced by momentum. It would be interesting to understand how the stochastic noise and the momentum together contribute to the generalization of a neural network.

References

Appendix A Additional Experiments

In this section, we present additional experiments to further strengthen our empirical results. We verify on Resnet-18 that momentum induces generalization improvement when trained without batch normalization and data augmentation. We then check that when these two regualizers are used, momentum does not improve generalization. We then confirm on Resnet-18 that the generalization improvement gets larger as the batch size increases. Then, we provide the performance of SGD and SGD+M on the Gaussian experiment introduced in the introduction. Lastly, we give additional plots showing that momentum allows to well-classify small margin data as mentioned at the end of the introduction.

Figure 7 displays the training loss and test accuracy obtained by training a Resnet-18 on CIFAR-10 and CIFAR-100. Similarly to the case where we trained a VGG-19, momentum significantly improves generalization whether in the stochastic case (SGD) or in the full batch setting (GD).

A.2 Influence of the batch size

Figure 8 shows the test accuracy obtained with a Resnet-18 using the stochastic gradient descent optimizer on CIFAR-10. Similarly to the VGG-19 experiment in Section 2, the generalization improvement induced by momentum gets larger as the batch size increases.

A.3 Influence of batch normalization and data augmentation

As mentioned in Section 2, batch normalization and data augmentation significantly reduce the generalization improvement induced by momentum. We further confirm this observation in Figure 9 and Figure 10.

A.4 Synthetic Gaussian data experiments

We provide a complete table with mean and standard deviations obtained by using different student networks to learn the Gaussian synthetic experiment mentioned in the introduction.

A.5 Additional justification for the theory

In this section, we present further experiments to consolidate the experiment on the artificially decimated CIFAR-10 dataset described in the introduction.

In 11(a), we observe that using a Resnet-18, momentum still improves generalization on the small margin images.In 11(d) and 12(b), we see that using stochastic updates lead SGD to classify small margin images as well as SGD+M. Lastly, Figure 13 and Figure 14 show that batch normalization and data augmentation also reduce the generalization improvement of momentum: GD/SGD perform similarly as well as GD+M/SGD+M on the small margin data.

Appendix B Additional related work

GD+M (a.k.a. heavy ball or Polyak momentum) consists in using an exponentially weighted average of the past gradients to update the weights. For convex functions near a strict twice-differentiable minimum, GD+M is optimal regarding local convergence rate (Polyak, 1963, 1964; Nemirovskij & Yudin, 1983; Nesterov, 2003). However, it may fail to converge globally for general strongly convex twice-differentiable functions (Lessard et al., 2015) and is no longer optimal for the class of smooth convex functions. In the stochastic setting, GD+M is more sensitive to noise in the gradients; that is, to preserve their improved convergence rates, significantly less noise is required (d’Aspremont, 2008; Schmidt et al., 2011; Devolder et al., 2014; Kidambi et al., 2018). Finally, other momentum methods are extensively used for convex functions such as Nesterov’s accelerated gradient (Nesterov, 1983). Our paper focuses on the use of GD+M and contrary to the aforementioned papers, our setting is non-convex. Besides, we mainly focus on the generalization of the model learned by GD and GD+M when both methods converge to global optimal. Contrary to the non-convex case, generalization is disentangled from optimization for (strictly) convex functions.

The question we address concerns algorithmic regularization which characterizes the generalization of an optimization algorithm when multiple global solutions exist in over-parametrized models (Soudry et al., 2018; Lyu & Li, 2019; Ji & Telgarsky, 2019; Chizat & Bach, 2020; Gunasekar et al., 2018; Arora et al., 2019). This regularization arises in deep learning mainly due to the non-convexity of the objective function. Indeed, this latter potentially creates multiple global minima scattered in the space that vastly differ in terms of generalization. Algorithmic regularization is induced by and depends on many factors such as learning rate and batch size (Goyal et al., 2017; Hoffer et al., 2017; Keskar et al., 2016; Smith et al., 2018), initialization (Allen-Zhu & Li, 2020), adaptive step-size (Kingma & Ba, 2014; Neyshabur et al., 2015; Wilson et al., 2017), batch normalization (Arora et al., 2018; Hoffer et al., 2019; Ioffe & Szegedy, 2015) and dropout (Srivastava et al., 2014; Wei et al., 2020). However, none of these works theoretically analyzes the regularization induced by momentum.

Appendix C Notations

In this section, we introduce the different notations used in the proofs. We start by defining the notations that appear for GD and GD+M. We first consider the case when $\lambda=0$, we will extend the proof to $\lambda>0$ in section H

Our paper rely on the notions of signal and noise components of the iterates.

– Signal intensity: $\theta=\alpha$ if $i\in\mathcal{Z}_{1}$ and $\beta$ otherwise.

– Signal: $c_{r}^{(t)}=\langle\bm{w^{*}},\bm{w}_{r}^{(t)}\rangle$ for $r\in[m].$

– Noise: $\Xi_{i,j,r}^{(t)}=\langle\bm{w}_{r}^{(t)},\bm{X}_{i}[j]\rangle$ for $i\in[N]$ and $j\in[P]\backslash\{P(\bm{X}_{i})\}.$

– Max noise: ${\Xi}_{\max}^{(t)}=\max_{i\in[N],j\not=P(\bm{X}_{i}),r\in[m]}|\Xi_{i,j,r}^{(t)}|^{2}.$

– Total noise: $\Xi_{i}^{(t)}=\sum_{r\in[m],j\in[P],j\not=P(\bm{X}_{i})}y_{i}\left(\Xi_{i,j,r}^{(t)}\right)^{3}.$

We also use the following notations when dealing with the loss function and its gradient.

– Noise loss: $\widehat{\mathcal{L}}^{(t)}(\Xi_{i}^{(t)})=\log\left(1+\exp\left(-\Xi_{i}^{(t)}\right)\right)$.

– Full derivative: $\nu^{(t)}=\nu_{1}^{(t)}+\nu_{2}^{(t)}.$

– Gradient on signal: $\mathscr{G}_{r}^{(t)}=\langle\nabla_{\bm{w}_{r}}\widehat{L}(\bm{W}^{(t)}),\bm{w^{*}}\rangle$ for $r\in[m].$

– Gradient on noise: $\texttt{G}_{i,j,r}^{(t)}=\langle\nabla_{\bm{w}_{r}}\widehat{L}(\bm{W}^{(t)}),\bm{X}_{i}[j]\rangle$ for $i\in[N]$, $j\in[P]\backslash\{P(\bm{X}_{i})\}$ and $r\in[m].$

C.2 Notations specific to GD+M

We now introduce the notations that only appear in the proofs involving GD+M.

– Momentum gradient oracle: $\bm{g}_{r}^{(t)}=\gamma\bm{g}_{r}^{(t-1)}+(1-\gamma)\nabla_{\bm{w}_{r}}\widehat{L}(\bm{W}^{(t)})$ for $r\in[m].$

– Signal momentum: $\mathcal{G}_{r}^{(t)}:=\langle\bm{g}_{r}^{(t)},\bm{w^{*}}\rangle$ for $r\in[m].$

– Noise momentum: $G_{i,j,r}^{(t)}=\langle\bm{g}_{r}^{(t)},\bm{X}_{i}[j]\rangle$ for $i\in[N]$, $j\in[P]\backslash\{P(\bm{X}_{i})\}$ and $r\in[m].$

Appendix D Induction hypotheses

We prove our main result using an induction. More specifically, we make the following assumptions for every time $t\leq T.$

Throughout the training process using GD for $t\leq T$, we maintain that:

(Large signal data have small noise component). For every $i\in\mathcal{Z}_{1}$, for every $j\in[P]\backslash\{P(\bm{X}_{i})\}$ and $r\in[m],$ we maintain:

(Small signal data have large noise component). For every $i\in\mathcal{Z}_{2}$, for every $j\in[P]\backslash\{P(\bm{X}_{i})\}$ and $r\in[m],$ we have:

Throughout the training process using GD for $t\leq T$, the signal component is bounded for every $r\in[m]$ as

Throughout the training process using GD for $t\leq T$, we maintain:

Throughout the training process using GD+M for $t\leq T$, for every $i\in[N]$, for every $j\in[P]\backslash\{P(\bm{X}_{i})\}$, we have that:

Throughout the training process using GD+M for $t\leq T$, for $r\in[m]$, we have that:

In what follows, we assume these induction hypotheses for $t<T$ to prove our generalization results. We then prove these hypotheses for $t+1.$

Appendix E Gradients and updates

In this section, we first derive the gradient of the loss $\widehat{L}$. We then provide its projection on $\bm{\bm{w^{*}}}$ (signal gradient) and on $\bm{X}_{i}[j]$ (noise gradient). We first derive the gradient of the loss $\widehat{L}.$

For $t\geq 0$ and $r\in[m]$, the gradient of the loss $\widehat{L}$ with respect to $\bm{w}_{r}$ is:

. We derive $\widehat{L}$ with respect to $\bm{w}_{r}$ and obtain:

By rewriting (9), we obtain the desired result. ∎

To track the signal learnt by our models, we compute the signal gradient which is the projection of the gradient on $\bm{w^{*}}.$

For all $t\geq 0$ and $r\in[m]$, the signal gradient is:

We obtain the desired result by projecting the gradient from Lemma E.1 on $\bm{w^{*}}$ and using $\bm{X}_{i}[j]\perp\bm{w^{*}}.$ ∎

E.2 Noise gradient

To prove the memorization of GD and the non-memorization of GD+M, we also need to compute the noise gradient which is the projection of the gradient $\nabla_{\bm{w}_{r}}\widehat{L}$ on $\bm{X}_{i}[j].$

For all $t\geq 0$, $i\in[N]$ and $j\in[P]\backslash\{P(X_{i})\}$ and $r\in[m]$, the noise gradient is:

Similarly to Lemma E.2, we obtain the desired result by projecting the gradient from Lemma E.1 on $\bm{X}_{i}[j]$ and using $\bm{X}_{i}[j]\perp\bm{w^{*}}.$ ∎

Intuitively, (10) means that the sum of the sigmoid terms for all time steps is bounded (up to a logarithmic dependence).

Appendix F Learning with GD

In this section, we detail the proofs of the lemmas in Section 5 and Theorem 4.1. We first characterize the dynamics of the signal $c_{r}^{(t)}$ in subsection F.1. We then analyze the dynamics of the noise $\Xi_{i,j,r}^{(t)}$ in subsection F.2 and show the memorization of the GD model. We finally prove Theorem 4.1 in subsection F.3 and the induction hypotheses in subsection F.4.

To track the amount of signal learnt by GD, we make use of the following update.

For all $t\geq 0$ and $r\in[m]$, the signal update (1) is equal:

The signal update is obtained by using (1) and the signal gradient (Lemma E.2). This yields

To obtain the desired upper bound, we apply the same reasoning as above to bound the $\mathcal{Z}_{1}$ term. ∎

For $t\in[0,T_{0}]$, we know that for all $s\in[m]$, we have $c_{s}^{(t)}\leq\frac{\kappa}{m^{1/3}\alpha}$. Therefore, we have

Plugging (15) in the left-hand side of (11) yields the desired lower bound.

We now prove Lemma 5.1 that quantifies the amount of signal learnt by GD when the derivative is large.

Let $r\in[m]$. From Lemma F.2, the signal update for $t\in[0,T_{0}]$ is

where $A$ and $B$ are respectively defined as:

We now prove Lemma 5.2. It states that since the signal $c^{(t)}$ has significantly increased, the $\mathcal{Z}_{1}$ derivative $\nu_{1}^{(t)}$ is now small. Before proving this result, we introduce an auxiliary Lemma.

From Lemma F.3, we deduce an upper bound on $\nu_{1}^{(t)}$:

On the other hand, using Lemma E.2, the signal difference is bounded as:

We now bound (20) by a loss term by applying Lemma K.20. Using Lemma 5.1 and D.2, we have:

From Lemma I.7, we have the convergence rate of $\widehat{\mathcal{L}}^{(t)}(\alpha)$. We use it to bound $\nu_{1}^{(t)}.$

The bound on $\nu^{(t)}$ is obtained by using its definition $\nu^{(t)}=\nu_{1}^{(t)}+\nu_{2}^{(t)}$. ∎

We earlier proved that after $T_{0}$ iterations, the signal $c^{(t)}$ learnt by the GD model significantly increases until making $\nu_{1}^{(t)}$ small. We therefore need to rewrite the signal update in this case.

For $t\in[T]$, the maximal signal $c^{(t)}$ updates as:

From the signal update given by Lemma F.1, we know that:

To obtain the desired result, we need to prove for $i\in\mathcal{Z}_{1}$:

By using D.1 and D.2, (27) is bounded as:

Using Remark 1, the sigmoid term in (F.1.2) becomes small when $\alpha c^{(t)}\geq\kappa^{1/3}$. To summarize, we have:

Besides, we use D.2 to bound $(c^{(t)})^{2}$ in the right-hand side of (25). ∎

We now show that once $\nu_{1}^{(t)}$ is small, the amount of learnt signal is controlled by $\nu_{2}^{(t)}$ .

Let $\tau\in[T_{0},T)$. From Lemma F.4, we know that:

Let $t\in[\tau,T)$. We now sum up (30) for $\tau=T_{0},\dots,t$ and obtain:

We now plug the bound on $\nu_{1}^{(t)}$ from Lemma 5.2 in (31). This implies:

F.2 Memorization process of GD

Lemma 5.2 shows that after $T_{0}$ iterations, the gradient is controlled by $\nu_{2}^{(t)}$. In this section, we show that this yields the GD model to memorize.

Using Lemma F.1, we simplify the noise update.

Let all $t\geq 0$, $i\in[N]$, $j\in[P]\backslash\{P(\bm{X}_{i})\}$ and $r\in[m]$. Then, with probability at least $1-o(1)$, the noise update (2) is bounded as

Let $i\in[N]$, $j\in[P]\backslash\{P(\bm{X}_{i})\}$ and $r\in[m]$. From Lemma E.3, we know that the noise update satisfies:

We now apply Lemma K.5 and Lemma K.7 to respectively bound $\|\bm{X}_{i}[j]\|_{2}^{2}$ and $\langle\bm{X}_{a}[k],\bm{X}_{i}[j]\rangle$ in (34) and obtain the desired result. ∎

In the next lemma, we further simplify the noise update from Lemma F.5.

Let $i\in\mathcal{Z}_{2}$, $j\in[P]\backslash\{P(\bm{X}_{i})\}$ and $r\in[m]$. Our starting point is Lemma I.4 which states that:

Lemma F.7 indicates that $y_{i}\Xi_{i,j,r}^{(t)}$ is not non-decreasing but overall, this quantity gets large over time. We now want to determine the time $T_{1}$ where one of the $y_{i}\Xi_{i,j,r}^{(t)}$ becomes large.

Lemma F.7 indicates that the noise iterate satisfies for $t\in[0,T_{1}]$:

We proved in the previous section that after $T_{1}$ iterations, the amount of noise memorized by the GD model significantly increases. We want to show that after this phase, $\nu_{2}^{(t)}$ is well-controlled.

On the other hand, from Lemma I.4, we know that for all $r\in[m]$:

Moreover, by applying Lemma K.20 to (49), we have:

We now apply Lemma I.8 to bound the loss in (50).

Using Lemma F.9, we can obtain a bound on the sum over time of $\mathcal{Z}_{2}$ derivatives .

Combining the bound on $\sum_{j=T_{1}}^{T}\nu_{2}^{(j)}$ from Lemma F.9 and (52) yields:

We have thus a control on the sum over time of $\nu_{2}^{(t)}$. We can make use of Lemma 5.3 to get the final control on the signal iterate $c^{(t)}.$

Let $t\in[T].$ From Lemma 5.3, we know that the signal is bounded as

We plug the bound from Lemma 5.5 to bound the last term in the right-hand side of (54).

F.3 Proof of Theorem 4.1

We proved that the weights learnt by GD satisfy for $r\in[m]$

We now bound the training and test error achieved by GD at time $T.$

Train error. Lemma I.8 provides a convergence bound on the training loss.

Test error. Let $(X,y)$ be a datapoint. We remind that $\bm{X}=(\bm{X},\dots,\bm{X}[P])$ where $\bm{X}[P(\bm{X})]=\theta y\bm{w^{*}}$ and $\bm{X}[j]\sim\mathcal{N}(0,\sigma^{2}\mathbf{I}_{d})$ for $j\in[P]\backslash\{P(\bm{X})\}.$ We bound the test error as follows:

We now want to compute the probability terms in (58) and (59). We remind that $(\bm{X},y)\sim\mathcal{Z}_{1},$ $yf_{\bm{W}^{(T)}}(\bm{X})$ is given by

We now apply Lemma 5.6 in (60) and obtain:

Let $(\bm{X},y)\sim\mathcal{Z}_{2}$. Similarly, by applying Lemma 5.6, $yf_{\bm{W}^{(T)}}(\bm{X})$ is bounded as:

Therefore, using (122), we upper bound the test error (120) as:

Since $y$ is taken uniformly from $\{-1,1\},$ we further simplify (63) as:

F.4 Proof of the GD induction hypotheses

To prove Theorem 4.1, we used the induction hypotheses stated in Appendix D. The goal of this section is to prove them for $t+1.$

We prove here the main hypotheses we made on the noise when using GD.

Let’s start with the upper bound $y_{i}\Xi_{i,j,r}^{(t)}$ for $i\in\mathcal{Z}_{2}$. Using Lemma I.6, Lemma 5.5 and D.1, we deduce from (66) that:

which proves the induction hypothesis for $t+1.$ Regarding the lower bound, using D.1 and Lemma 5.5, we deduce from (66) that:

which proves the induction hypothesis for $t+1.$

Using D.1, we bound $y_{i}\Xi_{i,j,r}^{(0)}$ and $(\Xi_{a,k,r}^{(\tau)})^{2}$ in (69). We obtain:

We now apply Lemma I.5 and Lemma 5.5 to bound the derivative terms in (71).

Using the same type of reasoning as for the upper bound, one can show that (72) yields:

(73) shows the induction hypothesis for $t+1.$

We prove the induction hypotheses for the signal $c_{r}^{(t)}.$

Appendix G Learning with GD+M

In this section, we prove the Lemmas in Section 6 and Theorem 4.2.

To track the amount of signal learnt by GD, we make use of the following update.

For all $t\geq 0$ and $r\in[m]$, the signal momentum in (3) is equal to:

By definition of the momentum update, we have: $\bm{g}_{r}^{(t+1)}=\gamma\bm{g}_{r}^{(t)}+(1-\gamma)\nabla_{\bm{w}_{r}}\widehat{L}(\bm{W}^{(t)}).$ We project this update onto $\bm{w^{*}}$ and use Lemma E.2 to get:

From Lemma G.1, we can simplify the momentum update as:

For $t\in[0,\mathcal{T}_{0}]$, we know that for all $s\in[m]$, we have $c_{s}^{(t)}\leq\frac{\kappa}{m^{1/3}\alpha}$ . Thus, we have:

Plugging (78) and (80) in (76) yields the desired result. ∎

We now prove Lemma 6.1 that quantifies the signal learnt by GD when $\nu_{1}^{(t)}$ is non-zero.

By Lemma G.2, the signal update for $t\in[0,\mathcal{T}_{0}]$ satisfies:

We now show that contrary to GD, GD+M still has a large momentum in the $\bm{w^{*}}$ direction. In other words, we want to show that $-\mathcal{G}^{(t)}$ is still large after $\mathcal{T}_{0}$ iterations. Given that the small margin and large margin data share the same feature $\bm{w^{*}}$, this large momentum helps to learn $\mathcal{Z}_{2}$.

Before proving such result, we need some intermediate lemmas.

Using the momentum update rule, we know that:

Let’s define $t^{\prime}:=\mathcal{T}_{0}-\frac{1}{\sqrt{1-\gamma}}$. We start by summing the GD+M update (3) for $\tau=t^{\prime},\dots,\mathcal{T}_{0}$ to get

Applying Lemma G.3 to bound the momentum gradient, we further bound (83) to get:

We now use the fact that $\mathcal{T}_{0}-t^{\prime}=\frac{1}{\sqrt{1-\gamma}}$ in (84) to get:

Since $\gamma=1-\varepsilon$ with $\varepsilon\ll 1$, we linearize the right-hand side in (85) to obtain:

Using Lemma G.4, we can therefore show that once we learn $\mathcal{Z}_{1},$ $\mathcal{G}^{(t)}$ still stays large.

Using Lemma G.4, we bound $(c_{r}^{(\tau)})$ in (87) and get:

Since the signal momentum is large (Lemma 6.2), we want to argue that GD+M keeps learning the feature to eventually have a large signal.

Let $\mathcal{T}_{1}\in[T]$ such that $\mathcal{T}_{0}<\mathcal{T}_{1}.$ From the signal momentum update, we deduce:

We now apply Lemma 6.2 to bound $-\mathcal{G}^{(\mathcal{T}_{1})}$ in (89) and get:

We would like to find the time $\mathcal{T}_{1}$ such that $\gamma^{\mathcal{T}_{1}-\mathcal{T}_{0}}$ is a constant factor $a\leq 1$ i.e. such that

Let $t\in(\mathcal{T}_{1},T]$. Using (3) update rule, we have

where we used the fact that $-\mathcal{G}^{(\tau)}\geq 0$ in (94). Plugging (93) in (94) yields the desired bound.

G.2 GD+M does not memorize

Lemma 6.3 implies that after $\mathcal{T}_{1}$ iterations, the learnt signal is very large. We would like to show that this implies that the full derivative quickly decreases (Lemma 6.4) which implies that the GD+M cannot memorize (Lemma 6.5). Before proving Lemma 6.4, we need an auxiliary lemma that connects the signal momentum and the full derivative $\nu^{(t)}$.

For $t\in[\mathcal{T}_{1},T]$, the signal momentum is bounded as

From Lemma G.1 we know that the signal momentum is equal to

We finally apply Lemma 6.3 to bound $c^{(t)}$ in (96) to obtain the desired result. ∎

Lemma G.5 provides an upper bound on $\nu^{(t)}$ since:

We now would like to give a convergence rate on the iterates $\mathcal{G}^{(t+1)}-\gamma\mathcal{G}^{(t)}.$ Since Lemma J.9 gives a rate on the loss function, we connect the momentum increment with a loss term. Applying Lemma G.1, we have:

We now show that for $t\in[\mathcal{T}_{1},T]$, we have:

Indeed, by using Lemma 6.3 and D.5, we have:

Given our choice of $\alpha$, $\beta$ and $\hat{\mu}$, we finally bound (102) as:

We now apply Lemma K.20 to link (104) with a loss term. By Lemma 6.3 and D.5, we have:

Therefore, applying Lemma K.20 in (104) gives:

We finally apply Lemma J.9 to bound the loss term in (107) and get the desired result. ∎

After $\mathcal{T}_{1}$ iterations, the gradient is now very small and the noise component learnt by GD+M stays very small.

We sum up (108) for $s=\mathcal{T}_{1},\dots,t$ and obtain:

We apply the triangle inequality in (109) and obtain:

We now use D.4 to bound $|\Xi_{i,j,r}^{(\mathcal{T}_{1})}|$ in (110):

We now plug the bound on $\sum_{s=\mathcal{T}_{1}}^{t}|G_{i,j,r}^{(s+1)}|$ given by Lemma J.6 and obtain:

Given the values of $\mathcal{T}_{1},$ $\eta$, $\gamma$ and $\beta$, we can deduce that

Plugging (113) in (112) proves the induction hypothesis for $t+1.$

G.3 Proof of Theorem 4.2

We proved that the weights learnt by GD+M satisfy for $r\in[m]$

We now bound the training and test error achieved by GD+M at time $T.$

Train error. Lemma J.9 provides a convergence bound on the fake loss. Indeed, we know that:

Using Lemma K.24 along with D.4, we lower bound the loss term in (115) by the true loss.

Combining (115) and (116), we obtain a bound on the training loss.

Test error. Let $(\bm{X},y)$ be a datapoint. We remind that $\bm{X}=(\bm{X},\dots,\bm{X}[P])$ where $\bm{X}[P(\bm{X})]=\theta y\bm{w^{*}}$ and $\bm{X}[j]\sim\mathcal{N}(0,\sigma^{2}\mathbf{I}_{d})$ for $j\in[P]\backslash\{P(\bm{X})\}.$ We bound the test error as follows:

We now want to compute the probability terms in (119) and (120). We remind that $yf_{\bm{W}^{T}}(\bm{X})$ is given by

We now apply Lemma 6.3, (121) is finally bounded as:

Therefore, using (122), we upper bound the test error (120) as:

Since $y$ is uniformly sampled from $\{-1,1\},$ we further simplify (123) as:

We know that $\langle\bm{v}_{s}^{(T)},\bm{X}[j]\rangle\sim\mathcal{N}(0,\|\bm{v}_{s}^{(T)}\|_{2}^{2}\sigma^{2})$. Therefore, $\langle\bm{v}_{s}^{(T)},\bm{X}[j]\rangle^{3}$ is the cube of a centered Gaussian.This random variable is symmetric. By Lemma K.1, we know that $\sum_{s=1}^{m}\sum_{j\neq P(\bm{X})}\langle\bm{v}_{s}^{(T)},\bm{X}[j]\rangle^{3}$ is also symmetric. Therefore, we simplify (124) as:

From Lemma K.14, we know that $\sum_{s=1}^{m}\sum_{j\neq p}\langle\bm{v}_{s}^{(T)},\bm{X}[j]\rangle^{3}$ is $\sigma^{3}\sqrt{P-1}\sqrt{\sum_{s=1}^{m}\|\bm{v}_{s}^{(T)}\|_{2}^{6}}$-subGaussian. Therefore, by applying Lemma K.3, (125) is further bounded by:

Using the fact that $\|\bm{v}_{s}^{(T)}\|_{2}\leq 1$ in (126) finally yields:

G.4 Proof of the GD+M induction hypotheses

We prove the induction hypotheses for the signal $c_{r}^{(t)}.$

We sum up (128) for $\tau=\mathcal{T}_{1},\dots,t$ and obtain:

We apply the triangle inequality in (129) and obtain:

We now use D.5 to bound $|c_{r}^{(\mathcal{T}_{1})}|$ in (130):

We now plug the bound on $\sum_{\tau=\mathcal{T}_{1}}^{t}|\mathcal{G}_{r}^{(\tau+1)}|$ given by Lemma J.3. We have:

Appendix H Extension to λ>0𝜆0\lambda>0

After iteration $T$, by Lemma I.8 and Lemma J.9, we know that for GD:

On the other hand, for $\Xi_{i,j,r}^{(t)}$ we know that:

Appendix I Technical lemmas for GD

This section presents the technical lemmas needed in Appendix F. These lemmas mainly consists in different rewritings of GD.

I.2 Signal lemmas

In this section, we present a lemma that bounds the sum over time of the GD increment.

Let $t,\mathscr{T}\in[T]$ such that $\mathscr{T}<t.$ Then, the $\mathcal{Z}_{1}$ derivative is bounded as:

Let $\mathscr{T},t\in[T]$ such that $\mathscr{T}<t.$ We now sum up (135) for $\tau=\mathscr{T},\dots,t$ and get:

Subcase 1: $\mathscr{T}<T_{0}.$ From Lemma 5.3, we know that:

Subcase 2: $\mathscr{T}>T_{0}.$ From Lemma 5.3, we know that:

We therefore managed to prove that in all the cases, (140) holds.

I.3 Noise lemmas

In this section, we present the technical lemmas needed in subsection F.2. The following lemma bounds the projection of the GD increment on the noise.

Let $i\in[N]$, $j\in[P]\backslash\{P(\bm{X}_{i})\}$ and $r\in[m]$. Let $\mathscr{T},t\in[T]$ such that $\mathscr{T}<t.$ Then, the noise update (2) satisfies

Let $i\in[N]$, $j\in[P]\backslash\{P(\bm{X}_{i})\}$ and $r\in[m]$. We set up the following induction hypothesis:

Let’s first show this hypothesis for $t=\mathscr{T}.$ From Lemma F.5, we have:

Now, we apply D.3 to bound $(\Xi_{a,k,r}^{(\mathscr{T})})^{2}$ in (143) and obtain:

Therefore, the induction hypothesis is verified for $t=\mathscr{T}.$ Now, assume (LABEL:eq:noiseindhypoth) for $t.$ Let’s prove the result for $t+1.$ We start by summing up the noise update from Lemma F.5 for $\tau=\mathscr{T},\dots,t$ which yields:

We apply D.3 to bound $(\Xi_{a,k,r}^{(t)})^{2}$ in (LABEL:eq:noiseupd20) and obtain:

To bound the first term in the right-hand side of (LABEL:eq:noiseupd202), we use the induction hypothesis (LABEL:eq:noiseindhypoth). Plugging this inequality in (LABEL:eq:noiseupd202) yields:

By rearranging the terms, we finally have:

which proves the induction hypothesis for $t+1.$

Now, let’s simplify the sum terms in (LABEL:eq:noiseindhypoth). Since $P\ll\sqrt{d}$, by definition of a geometric sequence, we have:

Plugging (151) in (LABEL:eq:noiseindhypoth) yields

Now, let’s simplify the second sum term in (152). Indeed, we have:

where we used (151) in the last inequality. Plugging (153) in (152) gives the final result. ∎

Let $i\in[N]$, $j\in[P]\backslash\{P(\bm{X}_{i})\}$ and $r\in[m]$. Let $\mathscr{T},t\in[T]$ such that $\mathscr{T}<t.$ Then, the noise update (2) satisfies

On the other hand we know from Lemma I.5 that:

By applying D.1, (157) is eventually bounded as:

By combining (51) and (158) we deduce that for all $j\in[P]\backslash\{P(\bm{X}_{i})\}$ and $r\in[m]$:

I.4 Convergence rate of the training loss using GD

In this section, we prove that when using GD, the training loss converges sublinearly in our setting.

Let $t\in[T_{0},T]$. Run GD with learning rate $\eta$ for $t$ iterations. Then, the $\mathcal{Z}_{1}$ loss sublinearly converges to zero as:

Let $t\in[T_{0},T].$ From Lemma F.1, we know that the signal update is lower bounded as:

Let’s now assume by contradiction that for $t\in[T_{0},T]$, we have:

From the (3) update, we know that $c_{r}^{(\tau)}$ is a non-decreasing sequence which implies that $\sum_{r=1}^{m}(\alpha c_{r}^{(\tau)})^{3}$ is also non-decreasing. Since $x\mapsto\log(1+\exp(-x))$ is non-increasing, this implies that for $s\leq t$, we have:

Plugging (164) in the update (162) yields for $s\in[T_{0},t]$:

Let $t\in[T_{0},T]$. We now sum (165) for $s=T_{0},\dots,t$ and obtain:

Given the values of $T,\eta,\alpha,\hat{\mu}$, we finally have:

Let $t\in[T_{1},T]$. Run GD with learning rate $\eta\in(0,1/L)$ for $t$ iterations. Then, the loss sublinearly converges to zero as:

We first apply the classical descent lemma for smooth functions (Lemma K.18). Since $\widehat{L}(W)$ is smooth, we have:

Lemma I.9 provides a lower bound on the gradient. We plug it in (170) and get:

Applying Lemma K.19 to (171) yields the aimed result. ∎

I.4.3 Auxiliary lemmas for the proof of Lemma I.8

To obtain the convergence rate in Lemma I.8, we used the following auxiliary lemma.

Let $t\in[T_{1},T]$. Run GD for $t$ iterations. Then, the norm of gradient is lower bounded as follows:

Let $t\in[T_{1},T]$. To obtain the lower bound, we project the gradient on the the signal and on the noise.

Since $\|w^{*}\|_{2}=1$, we lower bound $\|\nabla_{\bm{w}_{r}}\widehat{L}(\bm{W}^{(t)})\|_{2}^{2}$ as

By successively applying Lemma E.2 and Lemma I.1, $(\mathscr{G}_{r}^{(t)})^{2}$ is lower bounded as

For a fixed $i\in\mathcal{Z}_{2}$ and $j\in[P]\backslash\{P(\bm{X}_{i})\}$, we know that $\|\nabla_{\bm{w}_{r}}\widehat{L}(\bm{W}^{(t)})\|_{2}^{2}$ is lower bounded as

Combining (172), (175), (173) and (176) and using $2a^{2}+2b^{2}\geq(a+b)^{2},$ we thus bound $\|\nabla_{\bm{w}_{r}}\widehat{L}(\bm{W}^{(t)})\|_{2}^{2}$ as:

We now sum up (177) for $r=1,\dots,m$ and apply Cauchy-Schwarz inequality to get:

We apply Lemma I.1 to further lower bound (178) and get:

Using Lemma I.10, Lemma I.11 and Lemma I.12 we have:

Plugging (180), (181) and (182) in (179) yields:

Finally, we use Lemma I.13 and lower bound (183) by $\widehat{L}(\bm{W}^{(t)})^{2}$. This gives the aimed result.

We now present auxiliary lemmas that link the gradient terms with their corresponding loss.

Let $t\in[T_{1},T].$ Run GD for $t$ iterations. Then, we have:

Therefore, we can apply Lemma K.20 and get the lower bound:

Let $t\in[T_{1},T].$ Run GD for $t$ iterations. Then, we have:

We again verify that the conditions of Lemma K.20 are met. By using D.1, D.2 and Lemma 5.1, we have:

Lastly, we want to link the loss term in (187) with $\widehat{\mathcal{L}}^{(t)}(\alpha)$. By applying D.1 and Lemma K.24 in (187), we finally get:

Combining (187) and (188) yields the aimed result. ∎

Let $t\in[T_{1},T].$ Run GD for $t$ iterations. Then, we have:

We again verify that the conditions of Lemma K.20 are met. Using D.1, D.2 and Lemma 5.4, we have:

Lastly, we want to link the loss term in (LABEL:eq:neknfce) with $\widehat{\mathcal{L}}^{(t)}(\Xi_{i}^{(t)})$. By applying D.1 and Lemma K.24 in (LABEL:eq:neknfce), we finally get:

Combining (LABEL:eq:neknfce) and (191) yields the aimed result. ∎

Let $t\in[0,T]$ Run GD for for $t$ iterations. Then, we have:

we need to lower bound $\widehat{\mathcal{L}}^{(t)}(\alpha)$. By successively applying Lemma K.24 and D.1, we obtain:

By successively applying Lemma K.24 and D.1, we obtain:

Combining (193) and (194) yields the aimed result.

Lastly, to obtain Lemma I.8, we need to bound $G_{r}^{(t)}$ which is given by the next lemma.

For $r\in[m]$, the gradient of the loss $\widehat{L}(\bm{W}^{(t)})$ projected on the normalized noise $\textstyle\chi$ satisfies with probability $1-o(1)$ for $r\in[m]$:

Projecting the gradient (given by Lemma E.1) on $\textstyle\chi$ yields:

Since $\frac{\frac{1}{N}\sum_{i\in\mathcal{Z}_{2}}\sum_{j\neq P(\bm{X}_{i})}\bm{X}_{i}[j]}{\|\frac{1}{N}\sum_{b\in\mathcal{Z}_{2}}\sum_{l\neq P(\bm{X}_{i})}\bm{X}_{b}[l]\|_{2}}$ is a unit Gaussian vector, using Lemma K.8, we bound the right-hand side of (LABEL:eq:Grbd2) with probability $1-o(1)$, as:

Now, using Lemma Lemma K.10 , we can further lower bound the left-hand side of (LABEL:eq:Grbd3) as:

Remark that $\frac{1}{N}\sum_{b\in\mathcal{Z}_{2}}\sum_{l\neq P(\bm{X}_{i})}\bm{X}_{b}[l]\sim\mathcal{N}(0,\frac{\hat{\mu}P}{N}\sigma^{2})$. By applying Lemma K.9, we have:

Appendix J Auxiliary lemmas for GD+M

This section presents the auxiliary lemmas needed in Appendix G.

J.2 Signal lemmas

In this section, we present the auxiliary lemmas needed to prove D.5. We first rewrite the (3) update to take into account the case where the signal $c^{(\tau)}$ becomes large.

For $t\in[T]$, the maximal signal momentum $\mathcal{G}^{(t)}$ is bounded as:

Let $t\in[T]$. Using the signal momentum given by Lemma G.1, we know that:

To obtain the desired result, we need to prove for $i\in\mathcal{Z}_{1}$:

By using D.4 and D.5, (203) is bounded as:

Using Remark 1, the sigmoid term in (J.2) becomes small when $\alpha c^{(\tau)}\geq\kappa^{1/3}$. To summarize, we have:

A similar reasoning implies for $i\in\mathcal{Z}_{2}$:

Plugging (202) and (206) in (201) yields the aimed result. ∎

For $t\in[\mathcal{T}_{1},T)$, the sum of maximal signal momentum is bounded as:

Let $s\in[\mathcal{T}_{1},T]$. From Lemma J.2, the signal momentum is bounded as:

For $\tau\in[\mathcal{T}_{0}-1]$, we have $\gamma^{s-\tau}\leq\gamma^{s-\mathcal{T}_{0}+1}$ and for $\tau\in[\mathcal{T}_{1}-1]$, $\gamma^{s-\tau}\leq\gamma^{s-\mathcal{T}_{1}+1}$. From Lemma J.8 and Lemma 6.4, we can bound $\nu_{1}^{(\tau)}$ and $\nu_{2}^{(\tau)}$. Therefore, (209) is further bounded as:

We now use Lemma K.25 to bound the sum terms in (210). We have:

We now sum (211) for $s=\mathcal{T}_{1},\dots,t$. Using the geometric sum inequality $\sum_{s}\gamma^{s}\leq 1/(1-\gamma)$ and obtain:

We plug $\sum_{s}\sqrt{\gamma}^{s}\leq 1/(1-\sqrt{\gamma})$ and $\sum_{s=1}^{t-\mathcal{T}_{1}+1}1/s\leq\log(t)+1$ in (212). This yields the desired result. ∎

J.3 Noise lemmas

In this section, we present the technical lemmas to prove Lemma 6.5.

Run GD+M on the loss function $\widehat{L}(\bm{W}).$ Let $i\in[N]$, $j\in[P]\backslash\{P(\bm{X}_{i})\}$. At a time $t$, the noise momentum is bounded with probability $1-o(1)$ as:

Let $i\in[N]$ and $j\in[P]\backslash\{P(\bm{X}_{i})\}$. Combining the (4) update rule and Lemma E.3 to get the noise gradient $\texttt{G}_{i,j,r}^{(t)}$, we obtain

Using Lemma K.5 and Lemma K.7, (LABEL:eq:diffmoms1) becomes with probability $1-o(1),$

We upper bound the second term in (LABEL:eq:diffmoms2) by again using D.4:

Let $t\in[T]$. The noise momentum is bounded as

Let $\tau\in[T].$ From Lemma J.4, we know that:

We unravel the recursion (217) rule for $\tau=0,\dots,t$ and obtain:

For $t\in[\mathcal{T}_{1},T)$, the sum of noise momentum is bounded as:

Let $s\in[\mathcal{T}_{1},T)$. We first apply Lemma J.5 and obtain:

Using the bound from Lemma 6.4, (218) becomes

For $\tau\in[0,\mathcal{T}_{1}-1]$, we have $\gamma^{s-1-\tau}\leq\gamma^{s-\mathcal{T}_{1}+1}$. Plugging these two bounds in (219) implies:

We now use Lemma K.25 to bound the sum terms in (220). We have:

We now sum (221) for $s=\mathcal{T}_{1},\dots,t$. Using the geometric sum inequality $\sum_{s}\gamma^{s}\leq 1/(1-\gamma),$ we obtain:

We finally use the harmonic series inequality $\sum_{s=1}^{t-\mathcal{T}_{1}}1/s\leq 1+\log(t)$ in (222) to obtain the desired result. ∎

J.4 Convergence rate of the training loss using GD+M

In this section, we prove that when using GD+M, the training loss converges sublinearly in our setting.

For $t\in[\mathcal{T}_{0},T]$ Using GD+M with learning rate $\eta$, the loss sublinearly converges to zero as

Let $t\in[\mathcal{T}_{0},T].$ Using Lemma J.11, we bound the signal momentum as:

We now plug (225) in the signal update (3).

We now apply Lemma K.22 to lower bound (226) by loss terms. We have:

Let’s now assume by contradiction that for $t\in[\mathcal{T}_{0},T]$, we have:

From the (3) update, we know that $c_{r}^{(\tau)}$ is a non-decreasing sequence which implies that $\sum_{r=1}^{m}(\alpha c_{r}^{(\tau)})^{3}$ is also non-decreasing for $\tau\in[T]$. Since $x\mapsto\log(1+\exp(-x))$ is non-increasing, this implies that for $s\leq t$, we have:

Plugging (229) in the update (228) yields for $s\in[\mathcal{T}_{0},t]$:

We now sum (230) for $s=\mathcal{T}_{0},\dots,t$ and obtain:

Given the values of $\alpha,\eta,T$, we finally have:

We now link the bound on the loss to the derivative $\nu_{1}^{(t)}.$

The proof is similar to the one of Lemma 6.4.

For $t\in[\mathcal{T}_{1},T]$ Using GD+M with learning rate $\eta>0$, the loss sublinearly converges to zero as

Let $t\in[\mathcal{T}_{1},T].$ From Lemma J.10, we know that the signal gradient is bounded as $-\mathscr{G}^{(t)}\geq-\mathscr{G}^{(s)}$ for $s\in[\mathcal{T}_{1},t].$

By combining (236) and (238), we finally obtain:

We now plug (239) in the signal update (3).

We now apply Lemma K.22 to lower bound (240) by loss terms. We have:

Let’s now assume by contradiction that for $t\in[\mathcal{T}_{1},T]$, we have:

From the (3) update, we know that $c_{r}^{(\tau)}$ is a non-decreasing sequence which implies that $\sum_{r=1}^{m}(\theta c_{r}^{(\tau)})^{3}$ is also non-decreasing for $\tau\in[T]$. Since $x\mapsto\log(1+\exp(-x))$ is non-increasing, this implies that for $s\leq t$, we have:

Plugging (243) in the update (241) yields for $s\in[\mathcal{T}_{1},t]$:

We now sum (244) for $s=\mathcal{T}_{1},\dots,t$ and obtain:

Given the values of $\alpha,\beta,\eta,T,\hat{\mu}$, we finally have:

J.4.3 Auxiliary lemmas

We now provide an auxiliary lemma needed to obtain (J.9).

Let $t\in[\mathcal{T}_{1},T].$ Then, the signal gradient decreases i.e. $-\mathscr{G}^{(s)}\geq-\mathscr{G}^{(t)}$ for $s\in[\mathcal{T}_{1},t].$

The proof is similar to the one of Lemma J.10.

Appendix K Useful lemmas

In this section, we provide the probabilistic and optimization lemmas and the main inequalities used above.

In this section, we introduce the probabilistic lemmas used in the proof.

The sum of of symmetric random variables is symmetric.

Let $\sigma_{1},\sigma_{2}>0.$ Let $X$ and $Y$ respetively be $\sigma_{1}$- and $\sigma_{2}$-subGaussian random variables. Then, $X+Y$ is $\sqrt{\sigma_{1}+\sigma_{2}}$-subGaussian random variable.

Let $t>0.$ Let $X$ be a $\sigma$-subGaussian random variable. Then, we have:

We know that the $\|\cdot\|_{2}$ is $1$-Lipschitz and by applying Theorem K.1, we therefore have::

By rewriting (253) and using Lemma K.4, we have with probability $1-\delta,$

By squaring (254) and using $(a+b)^{2}\leq a^{2}+b^{2}$, we obtain the aimed result. ∎

We know that the $\|\cdot\|_{2}$ is $1$-Lipschitz and by applying Theorem K.1, we therefore have:

We use Lemma K.4 and set $\epsilon=\frac{\sigma\sqrt{d}}{2}$ in (255) to finally get:

Let’s define $Z:=\langle\bm{X},\bm{Y}\rangle.$ We first remark that $Z$ is a sub-exponential random variable. Indeed, the generating moment function is:

where we used $\log(1-x)\geq-x$ for $x<1$ in the last inequality. Therefore, by definition of a sub-exponential variable, we have:

Since $\sigma^{2}d\leq 1$ and $\epsilon\in,$ (256) is bounded as:

Let $\bm{U}:=\bm{X}/\|\bm{X}\|_{2}$ and $Z:=\langle\bm{U},\bm{Y}\rangle.$ We know that the pdf of $\bm{U}$ in polar coordinates is $f_{\bm{U}}(\theta)=\frac{\Gamma(d/2)}{2\pi^{d/2}}.$ Therefore, the generating moment function of $Z$ is: