Shape Matters: Understanding the Implicit Bias of the Noise Covariance

Introduction

One central mystery of deep artificial neural networks is their capability to generalize when having far more learnable parameters than training examples Zhang et al. (2016). To add to the mystery, deep nets can also obtain reasonable performance in the absence of any explicit regularization. This has motivated recent work to study the regularization effect due to the optimization (rather than objective function), also known as implicit bias or implicit regularization Gunasekar et al. (2017, 2018a, 2018b), Soudry et al. (2018), Arora et al. (2019a). The implicit bias is induced by and depends on many factors, such as learning rate and batch size (Smith et al., 2017, Goyal et al., 2017, Keskar et al., 2016, Li et al., 2019b, Hoffer et al., 2017), initialization and momentum (Sutskever et al., 2013), adaptive stepsize (Kingma and Ba, 2014, Neyshabur et al., 2015, Wilson et al., 2017), batch normalization Ioffe and Szegedy (2015), Hoffer et al. (2018), Arora et al. (2018) and dropout Srivastava et al. (2014), Wei et al. (2020).

Among these sources of implicit regularization, the SGD noise is believed to be a vital one (LeCun et al., 2012, Keskar et al., 2016). Previous theoretical works (e.g., Li et al. (2019b)) have studied the implicit regularization effect from the scale of the noise, which is directly influenced by learning rate and batch size. However, people have empirically observed that the shape of the noise also has a strong (if not stronger) implicit bias. For example, prior works show that mini-batch noise or label noise (label smoothing) – noise in the parameter updates from the perturbation of labels in training – is far more effective than adding spherical Gaussian noise (e.g., see (Shallue et al., 2018, Section 4.6) and Szegedy et al. (2016), Wen et al. (2019)). We also confirm this phenomenon in Figure 1 (left). Thus, understanding the implicit bias of the noise shape is crucial. Such an understanding may also be applicable to distributed training because synthetically adding noise may help generalization if parallelism reduces the amount of mini-batch noise (Shallue et al., 2018).

In this paper, we theoretically study the effect of the shape of the noise, demonstrating that it can provably determine generalization performance at convergence. Our analysis is based on a nonlinear quadratically-parameterized model introduced by (Woodworth et al., 2020, Vaskevicius et al., 2019), which is rich enough to exhibit similar empirical phenomena as deep networks. Indeed, Figure 1 (right) empirically shows that SGD with mini-batch noise or label noise can generalize with arbitrary initialization without explicit regularization, whereas GD or SGD with Gaussian noise cannot. We aim to analyze the implicit bias of label noise and Gaussian noise in the quadratically-parametrized model and explain these empirical observations.

We choose to study label noise because it can replicate the regularization effects of minibatch noise in both real and synthetic data (Figure 1), and has been used to regularize large-batch parallel training (Shallue et al., 2018). Moreover, label noise is less sensitive to the initialization and the optimization history than mini-batch noise, which makes it more amenable to theoretical analysis. For example, in an extreme case, if we happen to reach or initialize at a solution that overfits the data exactly, then mini-batch SGD will stay there forever because both the gradient and the noise vanish (Vaswani et al., 2019). In contrast, label noise will not accidentally vanish, so analysis is more tractable. Understanding label noise may lead to understanding mini-batch noise or replacing it with other more robust choices.

In our setting, we prove that with a proper learning rate schedule, SGD with label noise recovers a sparse ground-truth classifier and generalizes well, whereas SGD with spherical Gaussian noise generalizes poorly. Concretely, SGD with label noise biases the parameter towards sparse solutions and exactly recovers the sparse ground-truth, even when the initialization is arbitrarily large (Theorem 2.1). In this same regime, noise-free gradient descent quickly overfits because it trains in the NTK regime (Jacot et al., 2018, Chizat and Bach, 2018, Oymak and Soltanolkotabi, 2020, Du et al., 2018b, Arora et al., 2019b). Adding Gaussian noise is insufficient to fix this, as this algorithm would end up sampling from a Gibbs distribution with infinite partition function and fail to converge to the ground-truth (Theorem 2.2). In summary, with not too small learning rate or noise level, label noise suffices to bias the parameter towards sparse solutions without relying on a small initialization, whereas Gaussian noise cannot.

Our analysis suggests that the fundamental difference between label or mini-batch noise and Gaussian noise is that the former is parameter-dependent, and therefore introduces stronger biases than the latter. The conceptual message highlighted by our analysis is that there are two possible implicit biases induced by the noise: 1. prior work (Keskar et al., 2016) shows that by escaping sharp local minima, noisy gradient descent biases the parameter towards solutions which are more robust (i.e, solutions with low curvature, or “flat” minima), and 2. when the noise covariance varies across the parameter space, there is another (potentially stronger) implicit bias effect toward parameters where the noise covariance is smaller. Label or mini-batch noise benefit from both biases, whereas Gaussian noise is independent of the parameter, so it benefits from the first bias but not the second. For the quadratically-parameterized model, this first bias is not sufficient for finding solutions with good generalization because there is a large set of overfitting global minima of the training loss with reasonable curvature. In contrast, the covariance of label noise is proportional to the scale of the parameter, inducing a much stronger bias towards low norm solutions which generalize well.

There has been a line of work empirically studying how noise influences generalization. Keskar et al. (2016) argued that large batch training will converge to “sharp” local minima which do not generalize well. Hoffer et al. (2017) argued that large batch size doesn’t hurt generalization much if training goes on long enough and additional noise is added with a larger learning rate. Goyal et al. (2017) and Shallue et al. (2018) showed large batch training with proper learning rate and additional label noise can achieve similar generalization as small batch. Agarwal et al. (2020) disentangled the effects of update direction and scale for a variety of optimizers. Wei and Schwab (2019), Chaudhari and Soatto (2018), Yaida (2018) (heuristically) suggested that SGD may encourage solutions with smaller noise covariance. Martin and Mahoney (2018) used random matrix theory to analyze implicit regularization effects of noises. The noise induced by dropout has been shown to change the expected training objective, hence provides a regularization effect (Mianjy et al., 2018, Mianjy and Arora, 2019, Wei et al., 2020, Arora et al., 2020). Wei et al. (2020) showed that there also exisits an implicit bias induced by dropout noise.

Blanc et al. (2019) and Zhu et al. (2019) also studied implicit regularization effects which arise due to shape, rather than scale, of the noise, but only considered the local effect of the noise near some local minimum of the loss. In contrast, our work analyzes the global effect of noise. For a more detailed comparison with (Blanc et al., 2019), see Section 2.2.

Langevin dynamics or the closely-related stochastic gradient descent with Gaussian noise, has been studied in previous works Welling and Teh (2011), Teh et al. (2016), Raginsky et al. (2017), Zhang et al. (2017), Mou et al. (2017), Roberts et al. (1996), Ge et al. (2015), Negrea et al. (2019), Neelakantan et al. (2015). In particular, Raginsky et al. (2017) and Li et al. (2019a) provided generalization bounds for SGLD using algorithmic stability.

A number of works have theoretically analyzed implicit regularization in simplified settings (Soudry et al., 2018, Gunasekar et al., 2018b, Ji and Telgarsky, 2018a). Gunasekar et al. (2017) and Li et al. (2017) showed that gradient descent finds low rank solutions in matrix completion. Gradient descent also been shown to maximize the margin in linear and homogeneous models (Soudry et al., 2018, Ji and Telgarsky, 2018b, Nacson et al., 2018, Lyu and Li, 2019, Gunasekar et al., 2018a, Nacson et al., 2019, Poggio et al., 2017). Du et al. (2018a) showed that gradient descent implicitly balances the layers of deep homogeneous models. Other works showed that it may not be always possible to characterize implicit biases in terms of some norm (Arora et al., 2019a, Razin and Cohen, 2020). Gissin et al. (2019) showed that gradient descent dynamics exhibit different implicit biases based on depth. Hardt et al. (2015) derived stability-based generalization bounds for SGD based on training speed.

Woodworth et al. (2020), Vaskevicius et al. (2019) analyze the effect of initialization for the same model that we study, showing that a large initialization trains in the NTK regime (shown to generalize poorly (Wei et al., 2019, Ghorbani et al., 2019)) whereas small initialization does not. We show that when the initialization is large, adding noise helps avoid the NTK regime (Li and Liang, 2018, Jacot et al., 2018, Du et al., 2018b, Woodworth et al., 2020) without explicit regularization.

Recent works also suggest that explicit regularization may mitigate the lack of implicit regularization, especially in noisy or imbalanced settings. For example, Wei and Ma (2019) show that Lipschitz-ness regularization improves the performance in clean or noisy label setting when the learning rate is sub-optimal. Cao et al. (2019) show that additional regularization improves the generalization performance of rare classes. Nakkiran et al. (2020) show that explicit regularization can mitigate the double descent phenomenon in linear regression, which is caused by the fact that the implicit regularization of gradient descent with zero initialization is insufficient for the regime when the number of parameters is close to the number of datapoints.

Setup and Main Results

We remark that we can recover $v^{\star}$ by re-parameterizing $u=v^{\odot 2}$ and applying LASSO (Tibshirani, 1996) in the $u$-space when $n\geq\widetilde{O}(r)$, which is minimax optimal (Raskutti et al., 2012). However, the main goal of the paper, similar to several prior works (Woodworth et al., 2020, Vaskevicius et al., 2019, Li et al., 2017), is to prove that the implicit biases of non-convex optimization can recover the ground truth without explicit regularization in the over-parameterized regime when $n=\textup{poly}(r)\ll d$.We also remark that it’s common to obtain only sub-optimal sample complexity guarantees in the sparsity parameters with non-convex optimization methods (Li et al., 2017, Ge et al., 2016, Vaskevicius et al., 2019, Chi et al., 2019) due to technical limitations. We also assume throughout the paper that $n,d$ are larger than some sufficiently large universal constant.

Initialization. We use a large initialization of the form $v^{[{0}]}=\tau\cdot\mathds{1}$ where $\mathds{1}$ denotes the all 1’s vector, where we allow $\tau$ to be arbitrarily large (but polynomial in $d$).

SGD with label noise. We study SGD with label noise as shown in Algorithm 1. We sample an example, add label noise sampled from $\{\pm\delta\}$ to the label, and apply the gradient update. Computing the gradient, we obtain the update rule written explicitly as:

Langevin dynamics/diffusion. We compare SGD with label noise to Langevin dynamics, which adds spherical Gaussian noise to gradient descent (Neal et al., 2011):

where the noise $\xi\sim\mathcal{N}(0,\mathcal{I}_{d\times d})$ and $\lambda>0$ controls the scale of noise. Langevin dynamics (LD) or its more computationally-efficient variant, stochastic gradient Langevin dynamics (SGLD), is known to converge to the Gibbs distribution $\mu(v)\propto e^{-\lambda\mathcal{L}({v})}$ under various settings with sufficiently small learning rate (Roberts et al., 1996, Dalalyan, 2017, Bubeck et al., 2018, Raginsky et al., 2017). In our negative result about Langevin dynamics/diffusion, we directly analyze the Gibbs distribution in order to disentangle the convergence and the generalization.

This paper equates discrete time Langevin dynamics (equation (2)) with gradient descent with Gaussian noise, because LD with learning rate $\eta$ and temperature parameter $\lambda$ is exactly equivalent to gradient descent with learning rate $\eta$ and spherical Gaussian noise with standard deviation $\sigma=\sqrt{2/(\lambda\eta)}$. Thus technically the negative result for the Gibbs distribution (Theorem 2.2) applies to gradient descent with $\sigma$-Gaussian noise when keeping $\eta\sigma^{2}$ fixed (to be any number) and letting $\eta$ be sufficiently small.We also note that when $\eta\sigma^{2}$ also tends to zero, the effect of the noise will vanish and very likely gradient descent with Gaussian noise perform similarly to gradient descent.

Notations. Unless otherwise specified, we use $O(\cdot),\Omega(\cdot),\Theta(\cdot)$ to hide absolute multiplicative factors and $\widetilde{O}(\cdot),\widetilde{\Theta}(\cdot),\widetilde{\Omega}(\cdot)$ to hide poly-logarithmic factors in problem parameters such as $d$ and $\tau$. For example, every occurrence of $\widetilde{O}(x)$ is a placeholder for a quantity $f(x)$ that satisfies that for some absolute constants $c_{1},c_{2}>0$, $\forall x$, $|f(x)|\leq c_{1}|x|\cdot\log^{c_{2}}(d\tau)$.

2 Main Results

Our main result can be summarized by the following theorem, which suggests that stochastic gradient descent with label noise can converge to the ground truth despite a potentially large initialization.

In the setting of Section 2.1, given a target error $\epsilon>0$. Suppose we have $n\geq\widetilde{\Theta}(r^{2})$ samples. For any label noise level $\delta\geq\widetilde{\Theta}(\tau^{2}d^{2})$, we run SGD with label noise (Algorithm 1) with the following learning rate schedule:

learning rate $\eta_{0}=\widetilde{\Theta}(1/\delta)$ for $T_{0}=\widetilde{\Theta}(1)$ iterations,

learning rate $\eta_{1}=\widetilde{\Theta}(1/\delta^{2})$ for $T_{1}=\widetilde{\Theta}(1/\eta_{1})$ iterations,

learning rate $\eta_{2}=\widetilde{\Theta}(\epsilon^{2}/\delta^{2})$ for $T_{2}=\widetilde{\Theta}(1/\eta_{2})$ iterations.

Then, with probability at least $0.9$, the final iterate $v^{[{T}]}$ at time $T=T_{0}+T_{1}+T_{2}$ satisfies

Here $\widetilde{\Theta}(\cdot)$ omits poly-logarithmic dependencies on $1/\epsilon$, $d$ and $\tau$.

In other words, with arbitrarily large initialization scale $\tau$, we can choose large label noise level and the learning rate schedule so that SGD with label noise succeeds in recovering the ground truth. In contrast, when $\tau$ is large, gradient flow without noise trains in the “kernel” regime as shown by (Woodworth et al., 2020, Chizat and Bach, 2018). The solution in this kernel regime minimizes the RKHS distance to initialization, and in our setting equates to finding a zero-error solution with minimum $\|v^{\odot 2}-v^{\odot 2}\|_{2}$. Such a solution could be arbitrarily far away when initialization scale $\tau$ is large and therefore have poor generalization. Figure 1 (right) confirms GD performs poorly with large initialization whereas SGD with minibatch or label noise works. We outline the analysis of Theorem 2.1 in Section 3.

On the other hand, the following negative result for Langevin dynamics demonstrates that adding Gaussian noise fails to recover the ground truth even when $v^{\star}=0$. This suggests that spherical Gaussian noise does not induce a strong enough implicit bias towards low-norm solutions.

Assume in addition to the setting in Section 2.1 that the ground truth $v^{\star}=0$. When $n\leq d/3$, with probability at least $0.9$ over the randomness of the data, for any $\lambda>0$, the Gibbs distribution is not well-defined because the partition function explodes:

As a consequence, Langevin diffusion does not converge to a proper stationary distribution.

Theorem 2.2 helps explain the behavior in Figure 1, where adding Gaussian noise generalizes poorly for both synthetic and real data. In particular, in Figure 1 (right) adding Gaussian noise causes the parameter to diverge for synthetic data, and Theorem 2.2 explains this observation. A priori, the intuition regarding Langevin dynamics is as follows: as $\lambda\rightarrow+\infty$, the Gibbs distribution (if it exists) should concentrate on the manifold of global minima with zero loss. The measure on the manifold of global minima should be decided by the geometry of $\mathcal{L}({\cdot})$, and in particular, the curvature around the global minimum. As $\lambda\rightarrow+\infty$, the mass should likely concentrate at the flattest global minimum (according to some measure of flatness), which intuitively is $v^{\star}=0$ in this case.

However, our main intuition is that when $n<d$, even though the global minimum at $v^{\star}$ is the flattest, there are also many bad global minima with only slightly sharper curvatures. The vast volume of bad global minima dominate the flatness of the global minimum at $v^{\star}=0$ for any $\lambda$,In fact, one can show that if this phenomenon happens for some $\lambda>0$, then it happens for all other $\lambda$. and hence the partition function blows up and the Gibbs distribution doesn’t exist. More details in Section 4.

Analysis Overview of SGD with Label Noise (Theorem 2.1)

Towards building intuition and tools for analyzing the parameter-dependent noise, in this subsection we start by studying an extremely simplified random walk in one dimensional space. The random walk is purely driven by mean-zero noisy updates and does not involve any gradient updates:

Indeed, attentive readers can verify that when dimension $d=1$, sample size $n=1$, and $v^{\star}=0$, equation (2) degenerates to the above random walk if we omit the gradient update term (second to last term in equation (2)). We compare it with the standard Brownian motion (which is the analog of gradient descent with spherical Gaussian noise under this extreme simplification)

However, the parameter-dependent random walk (5) has dramatically different behavior when $\eta<1$: the random variable $v$ will eventually converge to $v=0$ with high probability (though the variance grows and the mean remains at 1.). This is because the variance of the noise depends on the scale of $v$. The smaller $v$ is, the smaller the noise variance is, and so the random walk tends to get “trapped” around . In fact, this claim has the following informal but simple proof that does not strongly rely on the exact form of the noise and can be extended to more general high-dimensional cases.

From the 1-D case to the high-dimensional case. In one dimension, it may appear that the varying scale of noise or norm of the covariance introduces the bias. However, in the high dimensional case, the shape of the covariance also matters. For example, if we generalize the random walk (5) to high-dimensions by running $d$ of the random walks in parallel, then we will observe the same phenomenon, but the noise variances in different dimensions are not identical — they depend on the current scales of the coordinates. (Precisely, the noise variance for dimension $k$ is $\eta^{2}v_{k}^{2}$.) However, suppose we instead add noise of the same variance to all dimensions. Even if this variance depends on the norm of $v$ (say, $\eta^{2}\|v\|_{2}^{2}$), the implicit bias will be diminished, as the smaller coordinates will have relatively outsized noise and the larger coordinates will have relatively insufficient noise.

Outline of the rest of the subsections. We will give a proof sketch of Theorem 2.1 that consists of three stages. We first show in the initial stage of the training that label noise effectively decreases the parameter on all dimensions, bringing the training from large initialization to a small initialization regime, where better generalization is possible (Section 3.2). Then, we show in Section 3.3 that when the parameter is decently small, with label noise and a decayed learning rate, the algorithm will increase the magnitude of those dimensions in support set of $v^{\star}$, while keep decreasing the norm of the rest of dimensions. Finally, with one more decay, the algorithm can recover the ground truth.

2 Stage 0: Label Noise with Large Learning Rate Reduces the Parameter Norm

We first analyze the initial phase where we use a relatively large learning rate. When the initialization is of a decent size, GD quickly overfits to a bad global minimum nearest to the initialization. In contrast, we prove that SGD with label noise biases towards the small norm region, for a similar reason as the random walk example with parameter-dependent noise in Section 3.1.

In the setting of Theorem 2.1, recall that we initialize with $v^{[{0}]}=\tau\cdot\mathds{1}$. Assume $n\geq\Theta(\log d)$. Suppose we run SGD with label noise with noise level $\delta\geq\widetilde{\Theta}(\tau^{2}d^{2})$ and learning rate $\eta_{0}\in[\widetilde{\Theta}(\tau^{2}d^{2}/\delta^{2}),\widetilde{\Theta}(1/\delta)]$ for $T_{0}=\widetilde{\Theta}(1/(\eta^{2}\delta^{2}))$ iterations. Then, with probability at least $0.99$ over the randomness of the algorithm,

Moreover, the minimum entry of $v^{[{T}]}$ is bounded below by $\exp(-\widetilde{O}((\eta\delta)^{-1}))$.

We remark that our requirement of $\eta$ being large is consistent with the empirical observation that large initial learning rate helps generalization Goyal et al. (2017), Li et al. (2019b). We provide intuitions and a proof sketch of the theorem in the rest of the subsection and defer the full proof to Section B . Our proof is based on the construction of a concave potential function $\Phi$ similar to Section 3.1. We will show that, at every step, the noise has a second order effect on the potential function and decrease the potential function by a quantity on the order of $\eta^{2}\delta^{2}$ (omitting the $d$ dependency).In general, any mean-zero noise has a second order effect on any potential function. Therefore, when the noise level is fixed, as $\eta\rightarrow 0$, the effect of the noise diminishes. This is why a lower bound on the learning rate is necessary for the noise to play a role. On the other hand, the gradient step may increase the potential by a quantity at most on the order of $\eta$ (omitting $d$ dependency again). Therefore, when $\eta^{2}\delta^{2}\gtrsim\eta$, we expect the algorithm to decrease the potential and the parameter norm.

In particular, we define $\Phi(v)\triangleq\sum_{k=1}^{d}\phi(v_{k})=\sum_{k=1}^{d}\sqrt{v_{k}}$.In the formal proof we will use a slightly different version of potential function (see Definition B.1). By the update rule (1), the update for a coordinate $k\in[d]$ can be written as

where $s_{t}$ is sampled from $\{-\delta,\delta\}$ and $i_{t}$ is sampled from $[n]$. Let $g_{k}^{(i_{t})}\triangleq(({v^{[{t}]}}^{\odot 2}-{v^{\star}}^{\odot 2})^{\top}x^{({i_{t}})})x^{({i_{t}})}_{k}$ be the component coming from the stochastic gradient. Using the fact that $\phi(ab)=\phi(a)\phi(b)$ for any $a,b>0$, we can evaluate the potential function at time $t+1$,

Here the expectation is over $s_{t}$ and $i_{t}$. We perform Taylor-expansion on the term $\phi(1-\eta s_{t}x^{({i_{t}})}_{k}-\eta g_{k}^{(i_{t})})$ to deal with the non-linearity and use the fact that $\eta s_{t}x^{({i_{t}})}_{k}$ is mean-zero:

In the setting of Theorem 3.1, for some failure probability $\rho>0$, let $b_{0}\triangleq 6\tau d/\rho$. Then, with probability at least $1-\rho/3$, we have that $\|v^{[{t}]}\|_{2}\leq b_{0}$ for any $t\leq T_{0}$.

By Lemma 3.2 and the bound on $|g^{i_{t}}_{k}|$ in terms of $\|v^{[{t}]}\|_{2}$, we have $|g^{i_{t}}_{k}|\leq(b_{0}^{2}+r)\|x^{({i_{t}})}\|_{\infty}^{2}\leq\widetilde{O}(b_{0}^{2}+r)$ with $b_{0}$ defined in Lemma 3.2 (up to logarithmic factors). Here we use again that each entry of the data is from $\mathcal{N}(0,1)$. Plugging these into equation (10) we obtain

In the setting of Section 2.1, given a target error bound $\epsilon_{1}>0$, we assume that $n\geq\widetilde{\Theta}(r^{2}\log^{2}(1/\epsilon_{1}))$. We run SGD with label noise (Algorithm 1) with an initialization $v^{[{0}]}$ whose entries are all in $[\epsilon_{\textup{min}},1/d]$, where $\epsilon_{\textup{min}}\geq\exp(-\widetilde{O}(1))$. Let noise level $\delta\geq\widetilde{\Theta}(\log(1/\epsilon_{1}))$ and learning rate $\eta=\widetilde{\Theta}({1}/{\delta^{2}})$, and number of iterations $T=\widetilde{\Theta}(\log(1/\epsilon_{1})/\eta)$. Then, with probability at least $0.99$, after $T$ iterations, we have

The proof of this theorem balances the contribution of the gradient against that of the noise on $S$ and $\bar{S}$. On $S$, the gradient provides a stronger signal than label noise, whereas on $\bar{S}$, the implicit bias of the noise, similarly to the effect in Section 3.2, outweighs the gradient and reduces the entries to zero. The analysis is more involved than that of Theorem 3.1, and we defer the full proof to Section C.

The conclusion of Theorem 3.3 still allows constant error in the support, namely, $\left\lVert v_{S}-v^{\star}_{S}\right\rVert_{\infty}\leq 0.1$. The following theorem shows that further annealing the learning rate will let the algorithm fully converge to $v^{\star}$ with any target error $\epsilon$.

[informal version of Theorem D.1] Assume initialization $v^{[{0}]}$ satisfies $\|v^{[{0}]}_{S}-v^{\star}_{S}\|_{\infty}\leq 0.1$. Suppose we run SGD with label noise with any noise level $\delta\geq 0$ and small enough learning rate $\eta$ for $T=\Theta(1/\eta)$ iterations. Then, with high probability over the randomness of the algorithm and data, there is $\|v^{[{T}]}_{S}-v^{\star}_{S}\|_{\infty}\leq\|v^{[{0}]}_{S}-v^{\star}_{S}\|_{\infty}/10.$

The formal version of Theorem 3.4 and its proof can be found in Section D.

Proof of Theorem 2.1. In Section E of Appendix, we combine Theorem 3.1, Theorem 3.3, and Theorem D.1 to prove our main Theorem 2.1.

Analysis Overview of Langevin Dynamics (Theorem 2.2)

Conclusion

In this work, we study the implicit bias effect induced by noise. For a quadratically-parameterized model, we theoretically show that the parameter-dependent noise has a strong implicit bias, which can help recover the sparse ground-truth from limited data. In comparison, our negative result shows that such a bias cannot be induced by spherical Gaussian noise. Our result provides an explanation for the empirical observation that replacing mini-batch noise or label noise with Gaussian noise usually leads to degradation in the generalization performance of deep models.

Acknowledgements

JZH acknowledges support from the Enlight Foundation Graduate Fellowship. CW acknowledges support from an NSF Graduate Research Fellowship. JDL acknowledges support of the ARO under MURI Award W911NF-11-1-0303, the Sloan Research Fellowship, and NSF CCF 2002272. TM acknowledges support of Google Faculty Award. The work is also partially supported by SDSI and SAIL at Stanford.

References

Appendix A Experimental Details

A.2 Experimental Details for Deep Neural Networks on CIFAR100

We train a VGG19 model (Simonyan and Zisserman, 2014) on CIFAR100, using a small and large batch baseline. We also experiment with adding Gaussian noise to the parameters after every gradient update as well as adding label noise in the following manner: with some probability that depends on the current iteration count, we replace the original label with a randomly chosen one.

To add spherical Gaussian noise to the parameter every update, we simply set $W\leftarrow W+\sigma z$ after every gradient update, where $z$ is a mean-zero Gaussian whose coordinates are drawn independently from $\mathcal{N}(0,1)$. We tune this $\sigma$ over the values shown in Figure 1.

We turn off weight decay and BatchNorm to isolate the regularization effects of just the noise alone. Standard data augmentation is still present in our runs. Our small batch baseline uses a batch size of 26, and our large batch baseline uses a batch size of 256. In runs where we add noise, the batch size is always 256. For all runs, we use an initial learning rate of 0.004. We train for 410550 iterations (i.e., minibatches), annealing the learning rate by a factor of 0.1 at the 175950-th and 293250-th iteration. Our models take around 20 hours to train on a single NVIDIA TitanXp GPU when the batch size is 256. The final performance gap between label noise or small minibatch training v.s. large batch or Gaussian noise is around 13% accuracy.

Appendix B Proof of Stage 0 (Theorem 3.1)

In this section, we will first prove several lemmas on which the proof of Theorem 3.1 is built upon. Then we will provide a proof of Theorem 3.1.

Since the gradient descent with label noise algorithm will blow up with some very small chance, we first define a coupled version of each optimization trajectory such that it is bounded and behaves similarly to the original trajectory.

Recall the update at $t$-th iteration is:

where the first inequality is Markov Inequality, the second is by the previous equation, and the third is by assumption of $\eta$ and the definition of $b_{0}$. ∎

We also define the following potential function which is similar to the one introduced in Section 3.2 but is only non-zero in a bounded area.

($b$-bounded potential function) For a vector $v$ that is positive on each dimension, we define the $b$-bounded potential function $\Phi(v)$ as follows: if $\left\lVert v\right\rVert_{1}\leq b$, we let $\Phi(v)\triangleq\sum_{k=1}^{d}\sqrt{v_{k}}$; otherwise $\Phi(v)\triangleq 0$.

Next, we prove that this potential function decreases to less than $\sqrt{\epsilon_{0}}$ with high probability after some number of iterations.

where $M$ is upper bound on $|g^{\prime\prime\prime}(1+x^{\prime})|$ for $x^{\prime}$ in to $x$, which is less than $3$ if $|x|\leq\frac{1}{2}$. So in our theorem if $\Delta\triangleq\eta s_{t}x^{({i_{t}})}_{k}+\eta(b_{0}^{2}+r)b_{x}^{2}\in[-\frac{1}{2},\frac{1}{2}]$, we have

where the first inequality if by Markov Inequaltiy, the second inequality is by the previous inequality, and the last inequality is because $T=\lceil\frac{32}{\eta^{2}\delta^{2}}\log(\frac{3d\sqrt{\tau}}{\rho\sqrt{\epsilon_{0}}})\rceil$. ∎

Now we are ready to prove Theorem 3.1 by combining the lemmas above.

Assume $\delta\geq 6\times 32^{2}b_{x}^{3}\frac{(b_{0}^{2}+r)}{\rho}\log(\frac{3d\sqrt{\tau}}{\rho\sqrt{\epsilon_{0}}})$, then we only need $\eta\in\left[\frac{6\times 32b_{x}^{2}}{\rho\delta^{2}}(b_{0}^{2}+r)\log(\frac{3d\sqrt{\tau}}{\rho\sqrt{\epsilon_{0}}}),\frac{1}{32\delta b_{x}}\right]$, and then all the above assumptions are satisfied.

where the first inequality is by update rule and the second is because $\delta>(b_{0}^{2}+r)$. Putting in the value of $T$, we have

Appendix C Proof of Stage 1 (Theorem 3.3)

In this section, we will first prove several lemmas on which the proof of Theorem 3.3 is built upon. Then we will provide a proof of Theorem 3.3.

Similar to Section B, we first define a coupled version of each optimization trajectory such that it is bounded and behaves similarly to the original trajectory. The difference here is that since those dimensions in $S$ are expected to grow to larger than those dimensions not in $S$, we use different boundaries for these two type of dimensions.

We first show that dimensions in $S$ don’t become much larger than the ground truth (which is $1$ for these dimensions).

where the first inequality is by Lemma G.3, the second inequality is by taking the sum of denominator.

where the last inequality is by assumption. ∎

Then, we prove that those dimensions not in $S$ don’t become much larger than ground truth (which is for these dimensions).

Next, we prove that suppose all the dimensions (in $S$ or not) are never much larger than the ground truth, for each dimension in $S$, there is some time such that this dimension is very close to the ground truth.

Next, we prove that for each dimension in $S$, whenever it gets close to ground truth, it never becomes much smaller than the ground truth.

where the first inequality is by Lemma G.3, the second inequality is by taking the sum of denominator.

Similar to Section B, we define a potential function in a bounded area. Since now we only want those dimensions not in $S$ to decrease, the potential function is only defined on these dimensions.

Then we prove that the this potential function decreases to less than $\sqrt{\epsilon_{1}}$ after proper number of iterations.

Now we are ready to prove Theorem 3.3 by combining the lemmas above.

We show the assumptions in the previous lemmas are all satisfied. The assumption $\frac{c_{1}^{2}}{8\eta\delta^{2}b_{x}^{2}}\geq\log\frac{6rT^{2}}{\rho}$ in Lemma C.2 is satisfied by

where the first is by $T\leq\frac{1}{\eta^{2}}$, the second is by $\log\frac{6r}{\rho}+4\log\frac{1}{\eta}\leq 4\log\frac{6r}{\rho}\log\frac{1}{\eta}$, and the last line is true because

The assumption $\frac{T}{\delta^{2}}\geq\frac{2^{9}}{c_{1}^{2}}\log\frac{6r}{\rho}$ in Lemma C.4 is satisfied by

Appendix D Proof of Stage 2 (Theorem 3.4)

The conclusion of Theorem 3.3 still allows constant error in the support, namely, $\left\lVert v_{S}-v^{\star}_{S}\right\rVert_{\infty}\leq 0.1$. To prove that further annealing the learning rate will let the algorithm fully converge to $v^{\star}$, we leverage a “bootstrapping” type of proof, where we first prove that whenever the support dimensions of the iterate is already somewhat close to $v^{\star}$, it can always become even closer (by a factor of $10$) to ground truth, while at the same time the other dimensions don’t increase by too much. By repeatedly using this analysis, we can prove that eventually the iterates will be arbitrarily close to the ground truth. Formally, we index the number of rounds that we use this analysis to be $s=2,3,\cdots$, and assume initially the iterate’s distance to $v^{\star}$ at the end of Theorem 3.3 is $\|v_{S}-v^{\star}_{S}\|_{\infty}\leq c_{1}\triangleq 0.1$ and $\|v_{\bar{S}}-v^{\star}_{\bar{S}}\|_{1}\leq\epsilon_{1}$. We prove the following theorem:

Let $s\geq 2$ be the index of the current round of bootstrapping. Let constant $c_{0}=1/10$. In the setting of Section 2.1, assume $v^{[{0}]}$ is an initial parameter satisfying $\|v^{[{0}]}_{S}-v^{\star}_{S}\|_{\infty}\leq c_{s-1}$ and $\|v^{[{0}]}_{\bar{S}}-v^{\star}_{\bar{S}}\|_{1}\leq\epsilon_{s-1}$, where $0<\epsilon_{s-1}\leq c_{s-1}\leq c_{0}$. Given a failure rate $\rho>0$. Assume $n\geq\widetilde{\Theta}(r^{2})$. Suppose we run SGD with label noise with noise level $\delta\geq 0$ and learning rate $\eta\leq\widetilde{\Theta}({c_{s}^{2}}/{(\delta^{2}+r^{2})})$ for $T=\log(4/c_{0})/\eta$ iterations. Then, with probability at least $1-\rho$ over the randomness of the algorithm and data, there is $\|v^{[{T}]}_{S}-v^{\star}_{S}\|_{\infty}\leq c_{s}\triangleq c_{s-1}c_{0}$ and $\|v^{[{T}]}_{\bar{S}}-v^{\star}_{\bar{S}}\|_{1}\leq\epsilon_{s}\triangleq(4/c_{0})^{2c_{s-1}}\epsilon_{s-1}.$ Here $\widetilde{\Theta}(\cdot)$ omits poly logarithmic dependency on $\rho$.

In the rest of this section, we will first prove several lemmas on which the proof of Theorem D.1 is built upon. Then we will provide a proof of Theorem D.1.

To begin with, we define the following coupled version of trajectories that are bounded to a region close to the ground truth.

First, we show that with high probability, those dimensions in $S$ don’t get too far away from ground truth (which is $1$ for these dimensions).

where the first inequality is by Lemma G.3, the second inequality is by taking the sum of denominator.

where the first inequality is by Lemma G.3, the second inequality is by taking the sum of denominator.

where the first inequality is by union bound, the second inequality is by previous results, the third inequality is by assumption of this lemma. ∎

The next step is to show that those dimensions not in $S$ remain close to ground truth .

here we are using $\epsilon_{s}>(1+\eta c_{s-1})^{T}\epsilon_{s-1}$ by assumption and the last step is by assumption. ∎

Then we prove that when every dimension (in $S$ or not) remains close to the ground truth, each dimension in $S$ will become even closer ($c_{s}/2$-close) to ground truth at some time.

Here the first inequality is by assumption, the second inequality is because $(\epsilon_{s}^{2}+4c_{s-1}r)C_{x}b_{x}^{2}\leq\frac{1}{10}c_{s}$ and $c_{s-1}\leq\frac{1}{10}$.

where the second inequality is because of assumption.

Next we show that once one dimension in $S$ become $c_{s}/2$-close to the ground truth (which is $1$), this dimension’s distance to ground truth will never become larger than $c_{s}$ any more.

where the first inequality is by Lemma G.3, the second inequality is by taking the sum of denominator.

Finally, we finish the proof with a union bound:

where the first inequality is by union bound, the second inequality is by previous results, the third inequality is by assumption of this lemma. ∎

Now we are ready to combine these lemmas to prove Theorem D.1.

Set $\eta$ small enough such that $\frac{c_{s}^{2}}{8\eta b_{x}^{2}(\delta^{2}+b_{x}^{2}(\epsilon_{s}^{2}+r)^{2})}\geq\log\frac{10rT^{2}}{\rho}$. Obviously we only need $\eta\leq\widetilde{\Theta}(\frac{c_{s}^{2}}{\delta^{2}+r^{2}})$, where $\widetilde{\Theta}(\cdot)$ omits poly logarithmic dependency on $d$ and $\rho$. Assume $(\epsilon_{s}^{2}+4c_{s-1}r)C_{x}b_{x}^{2}\leq\frac{c_{s}}{10}$, which can be represented as $C_{x}\leq\widetilde{\Theta}(\frac{1}{r})$. Recall $T=\frac{1}{\eta}\log\frac{4}{c_{0}}$, $\epsilon_{s}=e^{2c_{s-1}\log\frac{4}{c_{0}}}\epsilon_{s-1}$.

We first show that the additional assumptions in the previous lemmas are satisfied. There is

The assumption $\epsilon_{s}>(1+\eta c_{s-1})^{T}\epsilon_{s-1}$ in Lemma D.4 is therefore satisfied by definition of $\epsilon_{s}$. The assumption $\frac{((1+\eta c_{s-1})^{-T}\epsilon_{s}-\epsilon_{s-1})^{2}}{2T\eta^{2}b_{x}^{2}(\delta^{2}+b_{x}^{2}(\epsilon_{s}^{2}+r)^{2})\epsilon_{s}^{2}}\geq\log\frac{5}{\rho}$ in Lemma D.4 is satisfied because:

which is larger than $\log\frac{5}{\rho}$ by the definition of $\eta$. The assumption $(1-\eta)^{T}2c_{s-1}\leq\frac{c_{s}}{2}$ in Lemma D.5 is satisfied because $(1-\eta)^{T}2c_{s-1}\leq(\frac{1}{e})^{\log\frac{4}{c_{0}}}2c_{s-1}=\frac{c_{s}}{2}$. All the other assumptions in Lemma D.3, Lemma D.4, Lemma D.5and Lemma D.6 naturally follows from the definition of $\eta$ and the requirement of $C_{x}$.

Appendix E Proof of Theorem 2.1

Starting from initialization $\tau\cdot\mathds{1}$, by Theorem 3.1, running SGD with label noise with noise level $\delta>\widetilde{O}(\frac{\tau^{2}d^{2}}{\rho^{3}})$ and $\eta_{0}=\widetilde{\Theta}(\frac{1}{\delta})$ for $T_{0}=\widetilde{\Theta}(1)$ iterations gives us that with probability at least $0.99$, $\epsilon_{min}\leq v^{[{T_{0}}]}_{k}\leq\frac{1}{d}$ where $\epsilon_{min}=exp(-\widetilde{\Theta}(1))$. Now $v^{[{T_{0}}]}$ satisfies the initial condition of Theorem 3.3.

Recall the final target precision is $\epsilon$, set $\epsilon_{1}=\frac{1}{40^{3}}\epsilon$. By Theorem 3.3, with $n\geq\widetilde{\Theta}(r^{2}\log^{2}(1/\epsilon))$ data, after running SGD with label noise with learning rate $\eta_{1}=\widetilde{\Theta}(\frac{1}{\delta^{2}})$ for $T_{1}=\widetilde{\Theta}(\frac{\log(1/\epsilon)}{\eta_{1}})$ iterations, with probability at least $0.99$, there is,

So we have $v^{[{T_{0}+T_{1}}]}$ satisfies the initial condition of Theorem D.1.

Finally, set $\rho=0.01/\lceil\log_{10}(1/\epsilon)\rceil$, and apply Theorem D.1 for $n_{s}=\lceil\log_{10}(1/\epsilon)\rceil=\widetilde{\Theta}(1)$ rounds. Since $c_{s}$ gets smaller by $1/10$ for each round, the final $c_{n_{s}}$ satisfies $\frac{1}{10}\epsilon\leq c_{n_{s}}\leq\epsilon$. Since the requirement of $\eta$ for round $s$ is $\eta\leq\widetilde{\Theta}(\frac{c_{s}^{2}}{\delta^{2}+r^{2}})$, we can set $\eta_{2}\leq\widetilde{\Theta}(\frac{\epsilon^{2}}{\delta^{2}})$ to satisfy all the rounds at the same time. Set $T_{2}$ be the total number of iterations in all of these rounds, obviously $T_{2}=\widetilde{\Theta}(\frac{1}{\eta_{2}})$. Notice that $\epsilon_{s}\leq e^{\sum_{s=2}^{\infty}2c_{s-1}\log\frac{4}{c_{0}}}\epsilon_{1}\leq 40^{3}\epsilon_{1}=\epsilon$, we have with probability at least $0.99$,

The total failure rate of above three stages is $0.03$, so with probability at least $0.97$, there is $\left\lVert v^{[{T_{0}+T_{1}+T_{2}}]}-v^{\star}\right\rVert_{\infty}\leq\epsilon$, which finishes the proof. ∎

Appendix F Proof of Theorem 2.2

We first prove that with high probability, there is always some element-wise positive vector that is orthogonal to the subspace spanned by data.

By Theorem 1 of Amelunxen et al. (2014), we only need to prove

where $\delta(\cdot)$ is the statistical dimension of a set. By equation (2.1) of Amelunxen et al. (2014), there is

To calculate $\delta(C)$, we use Proposition 2.4 from Amelunxen et al. (2014),

where $g$ is a standard random vector, $\Pi_{C}$ is projection of $g$ to $C$, the expectation is over $g$. Since $C$ is the set of all points with element-wise positive coordinate, $\Pi_{C}(g)$ is simply setting all the negative dimension of $g$ to and keep the positive ones. Therefore,

Now we use this lemma to prove Theorem 2.2.

Let $X^{\perp}$ be the subspace that is orthogonal to the subspace $X$ spanned by data. Since data is random, with probability $1$ the random subspace $X$ is of $n$ dimension. Therefore, according to the previous lemma, with probability at least $0.999$, there is $X^{\perp}\cap C\neq\{0\}$, where $C$ is the coordinate-wise positive cone. Let $\mu\in X^{\perp}$ be such a vector such that $\mu_{i}>0$ for $\forall i\in[d]$, and we scale it such that $\left\lVert\mu\right\rVert_{2}=1$. We can construct the following orthonormal matrix

such that $span\{a^{1},\cdots a^{n}\}=X$ and $a^{n+1}=\mu$. Consider the following transformation

We can lower bound the partition function with

Here the inequality is because of the definition of $c$.

Appendix G Helper Lemmas

Suppose $n$ random data points are sampled i.i.d: $\forall i\in[n],x^{({i})}\sim\mathcal{N}(0,\mathcal{I}_{d\times d})$ Then, with probability at least $1-\rho$, for every $i\in[n]$ there is

By Gaussian tail bound, there is $\Pr\left(|x^{({i})}_{k}|>b_{x}\right)\leq 2e^{-\frac{b_{x}^{2}}{2}}$. So by union bound we have $\Pr\left(\max_{i,k}|x^{({i})}_{k}|>b_{x}\right)\leq 2nde^{-\frac{b_{x}^{2}}{2}}$. Let $b_{x}=\sqrt{2\log\frac{2nd}{\rho}}$ we complete the proof. ∎

Suppose $n$ random data points are sampled i.i.d: $\forall i\in[n],x^{({i})}\sim\mathcal{N}(0,\mathcal{I}_{d\times d})$. Then, when $n>24\log\frac{d}{\rho}$, with probability at least $1-\rho$, for every $k\in[d]$ there is

Therefore, when $n\geq 24\log\frac{d}{\rho}$, by union bound we finish complete the proof. ∎

where $A^{[{0:T}]}\in\left[\frac{c}{3},c\right]$ means for any $0\leq t<T$, $A^{[{t}]}\in\left[\frac{c}{3},c\right]$.

We only need to consider when $A^{[{0}]}\leq\frac{c}{2}$. Let $\hat{A}^{[{t}]}$ be the following coupling of $A^{[{t}]}$: starting from $\hat{A}^{[{0}]}=A^{[{0}]}$, for each time $t<T$, if $A^{[{t}]}=A^{[{t+1}]}=\cdots=A^{[{T}]}$ or there is $t^{\prime}\leq t$ such that $A^{[{t^{\prime}}]}\notin[\frac{c}{3},c]$, we let $\hat{A}^{[{t+1}]}\triangleq(1-\gamma)\hat{A}^{[{t+1}]}$; otherwise $\hat{A}^{[{t+1}]}=A^{[{t+1}]}$. Intuitively, whenever $A^{[{t}]}$ stops updating or exceeds proper range, we only times $\hat{A}^{[{t}]}$ by $1-\gamma$ afterwards, otherwise we let it be the same as $A^{[{t}]}$. Notice that if the event in Equation 179 happens, there has to be $\hat{A}^{[{T}]}=A^{[{T}]}$ (otherwise $A^{[{t}]}$ stops updating or exceeds range at some time, contradicting the event). So we only need to bound $\Pr\left(\hat{A}^{[{T}]}>c\right)$.

where $\neg E_{[0:T]}$ means for any $0\leq t<T$, $E_{t}$ doesn’t happen.

Let $\hat{A}^{[{t}]}$ be the following coupling of $A^{[{t}]}$: starting from $\hat{A}^{[{0}]}=A^{[{0}]}$, for each time $t<T$, if exists $t^{\prime}\leq t$such that $E_{t^{\prime}}$ happens or $A^{[{t^{\prime}}]}\notin[c_{1},c_{2}]$, we let $\hat{A}^{[{t+1}]}\triangleq(1-\gamma)\hat{A}^{[{t+1}]}$; otherwise $\hat{A}^{[{t+1}]}=A^{[{t+1}]}$. Intuitively, whenever $A^{[{t}]}$ exceeds proper range, we only times $\hat{A}^{[{t}]}$ by $1-\gamma$ afterwards, otherwise we let it be the same as $A^{[{t}]}$. Notice that if the event in Equation 185 happens, there has to be $\hat{A}^{[{T}]}=A^{[{T}]}$ (otherwise $E_{t}$ happens sometimes or $A^{[{t}]}\notin[c_{1},c_{2}]$, contradicting the event). So we only need to bound $\Pr\left(\hat{A}^{[{T}]}>c_{1}\right)$.