Label Noise SGD Provably Prefers Flat Global Minimizers

1 Introduction . One of the central questions in modern machine learning theory is the generalization capability of overparametrized models trained by stochastic gradient descent ( SGD ) . Recent work identifies the implicit regularization effect due to the optimization algorithm as one key factor in explaining the generalization of overparameterized models [ 27 , 11 , 19 , 10 ] . This implicit regularization is controlled by many properties of the optimization algorithm including search direction [ 11 ] , learning rate [ 20 ] , batch size [ 26 ] , momentum [ 21 ] and dropout [ 22 ] . The parameter-dependent noise distribution in SGD is a crucial source of regularization [ 16 , 18 ] . Blanc et al . [ 3 ] initiated the study of the regularization effect of label noise SGD with square loss1 by characterizing the local stability of global minimizers of the training loss . By identifying a data-dependent regularizer R ( θ ) , Blanc et al . [ 3 ] proved that label noise SGD locally diverges from the global minimizer θ∗ if and only if θ∗ is not a first-order stationary point of minθ R ( θ ) subject to L ( θ ) = 0 . The analysis is only able to demonstrate that with sufficiently small step size η , label noise SGD initialized at θ∗ locally diverges by a distance of η0.4 and correspondingly decreases the regularizer by η0.4 . This is among the first results that establish that the noise distribution alters the local stability of stochastic gradient descent . However , the parameter movement of η0.4 is required to 1Label noise SGD computes the stochastic gradient by first drawing a sample ( xi , yi ) , perturbing y′i = yi+ with ∼ { −σ , σ } , and computing the gradient with respect to ( xi , y′i ) . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . be inversely polynomially small in dimension and condition number and is thus too small to affect the predictions of the model . HaoChen et al . [ 12 ] , motivated by the local nature of Blanc et al . [ 3 ] , analyzed label noise SGD in the quadratically-parametrized linear regression model [ 29 , 32 , 23 ] . Under a well-specified sparse linear regression model and with isotropic features , HaoChen et al . [ 12 ] proved that label noise SGD recovers the sparse ground-truth despite overparametrization , which demonstrated a global implicit bias towards sparsity in the quadratically-parametrized linear regression model . This work seeks to identify the global implicit regularization effect of label noise SGD . Our primary result , which supports Blanc et al . [ 3 ] , proves that label noise SGD converges to a stationary point of L ( θ ) + λR ( θ ) , where the regularizer R ( θ ) penalizes sharp regions of the loss landscape . The focus of this paper is on label noise SGD due to its strong regularization effects in both real and synthetic experiments [ 25 , 28 , 31 ] . Furthermore , label noise is used in large-batch training as an additional regularizer [ 25 ] when the regularization from standard regularizers ( e.g . mini-batch , batch-norm , and dropout ) is not sufficient . Label noise SGD is also known to be less sensitive to initialization , as shown in HaoChen et al . [ 12 ] . In stark contrast , mini-batch SGD remains stuck when initialized at any poor global minimizer . Our analysis demonstrates a global regularization effect of label noise SGD by proving it converges to a stationary point of a regularized loss L ( θ ) + λR ( θ ) , even when initialized at a zero error global minimum . The learning rate and minibatch size in SGD are known to be important sources of regularization [ 9 ] . Our main theorem highlights the importance of learning rate and batch size as the hyperparameters that control the balance between the loss and the regularizer – larger learning rates and smaller batch sizes lead to stronger regularization . Section 2 reviews the notation and assumptions used throughout the paper . Section 2.4 formally states the main result and Section 3 sketches the proof . Section 4 presents experimental results which support our theory . Finally , Section 6 discusses the implications of this work . 2 Problem Setup and Main Result . Section 2.1 describes our notation and the SGD with label noise algorithm . Section 2.2 introduces the explicit formula for the regularizer R ( θ ) . Sections 2.3 and 2.4 formally state our main result . 2.1 Notation . We focus on the regression setting ( see Appendix E for the extension to the classification setting ) . Let { ( xi , yi ) } i∈ [ n ] be n datapoints with xi ∈ D and yi ∈ R. Let f : D×Rd → R and let fi ( θ ) = f ( xi , θ ) denote the value of f on the datapoint xi . Define ` i ( θ ) = 12 ( fi ( θ ) − yi ) 2 and L ( θ ) = 1n ∑n i=1 ` i ( θ ) . Then we will follow Algorithm 1 which adds fresh additive noise to the labels yi at every step before computing the gradient : Algorithm 1 : SGD with Label Noise Input : θ0 , step size η , noise variance σ2 , batch size B , steps T for k = 0 to T − 1 do Sample batch B ( k ) ⊂ [ n ] B uniformly and label noise ( k ) i ∼ { −σ , σ } for i ∈ B ( k ) . Let ˆ̀ ( k ) i ( θ ) = 1 2 ( fi ( θ ) − yi − ( k ) i ) 2 and L̂ ( k ) = 1B ∑ i∈B ( k ) ˆ̀ ( k ) i . θk+1 ← θk − η∇L̂ ( k ) ( θk ) end Note that σ controls the strength of the label noise and will control the strength of the implicit regularization in Theorem 1 . Throughout the paper we will use ‖ · ‖ = ‖ · ‖2 . We make the following standard assumption on f : Assumption 1 ( Smoothness ) . We assume that each fi is ` f -Lipschitz , ∇fi is ρf -Lipschitz , and∇2fi is κf -Lipschitz with respect to ‖ · ‖2 for i = 1 , . . . , n. We will define ` = ` 2f to be an upper bound on ‖ 1n ∑ i∇fi ( θ ) ∇fi ( θ ) T ‖2 , which is equal to ‖∇2L ( θ ) ‖2 at any global minimizer θ . Our results extend to any learning rate η ∈ ( 0 , 2 ` ) . However , they do not extend to the limit as η → 2 ` . Because we still want to track the dependence on 1η , we do not assume η is a fixed constant and instead assume some constant separation : Assumption 2 ( Learning Rate Separation ) . There exists a constant ν ∈ ( 0 , 1 ) such that η ≤ 2−ν ` . In addition , we make the following local Kurdyka-Łojasiewicz assumption ( KL assumption ) which ensures that there are no regions where the loss is very flat . The KL assumption is very general and holds for some δ > 0 for any analytic function defined on a compact domain ( see Lemma 17 ) . Assumption 3 ( KL ) . Let θ∗ be any global minimizer of L. Then there exist KL > 0 , µ > 0 and 0 < δ ≤ 1/2 such that if L ( θ ) − L ( θ∗ ) ≤ KL , then L ( θ ) − L ( θ∗ ) ≤ µ‖∇L ( θ ) ‖1+δ . We assume L ( θ∗ ) = 0 for any global minimizer θ∗ . Note that if L satisfies Assumption 3 for some δ then it also satisfies Assumption 3 for any δ′ < δ . Assumption 3 with δ = 1 is equivalent to the much stronger Polyak-Łojasiewicz condition which is equivalent to local strong convexity . We will use O , Θ , Ω to hide any polynomial dependence on µ , ` f , ρf , κf , ν , 1/σ , n , d and Õ to hide additional polynomial dependence on log 1/η , logB . 2.2 The Implicit Regularizer R ( θ ) . For L , σ2 , B , η as defined above , we define the implicit regularizer R ( θ ) , the effective regularization parameter λ , and the regularized loss L̃ ( θ ) : R ( θ ) = − 1 2η tr log ( 1− η 2 ∇2L ( θ ) ) , λ = ησ2 B , L̃ ( θ ) = L ( θ ) + λR ( θ ) . ( 1 ) Here log refers to the matrix logarithm . To better understand the regularizer R ( θ ) , let λ1 , . . . , λd be the eigenvalues of ∇2L ( θ ) and let R ( λi ) = − 12η log ( 1− ηλi 2 ) . Then , R ( θ ) = d∑ i=1 R ( λi ) = d∑ i=1 ( λi 4 + ηλ2i 16 + η2λ3i 48 + . . . ) . In the limit as η → 0 , R ( θ ) → 14 tr∇2L ( θ ) , which matches the regularizer in Blanc et al . [ 3 ] for infinitesimal learning rate near a global minimizer . However , in additional to the linear scaling rule , which is implicit in our definition of λ , our analysis uncovers an additional regularization effect of large learning rates that penalizes larger eigenvalues more than smaller ones ( see Figure 1 and Section 6.1 ) . The goal of this paper is to show that Algorithm 1 converges to a stationary point of the regularized loss L̃ = L + λR . In particular , we will show convergence to an ( , γ ) -stationary point , which is defined in the next section . 2.3 ( , γ ) -Stationary Points . We begin with the standard definition of an approximate stationary point : Definition 1 ( -stationary point ) . θ is an -stationary point of f if ‖∇f ( θ ) ‖ ≤ . In stochastic gradient descent it is often necessary to allow λ = ησ 2 B to scale with to reach an -stationary point [ 8 , 15 ] ( e.g. , λ may need to be less than 2 ) . However , for λ = O ( ) , any local minimizer θ∗ is an -stationary point of L̃ = L+ λR . Therefore , reaching a -stationary point of L̃ would be equivalent to finding a local minimizer and would not be evidence for implicit regularization . To address this scaling issue , we consider the rescaled regularized loss : 1 λ L̃ = 1 λ L+R . Reaching an -stationary point of 1λ L̃ requires non-trivially taking the regularizer R into account . However , it is not possible for Algorithm 1 to reach an -stationary point of 1λ L̃ even in the ideal setting when θ is initialized near a global minimizer θ∗ of L̃ . The label noise will cause fluctuations of order √ λ around θ∗ ( see section 3 ) so ‖∇L‖ will remain around √ λ . This causes 1λ∇L to become unbounded for λ ( and therefore ) sufficiently small , and thus Algorithm 1 can not converge to an -stationary point . We therefore prove convergence to an ( , γ ) -stationary point : Definition 2 ( ( , γ ) -stationary point ) . θ is an ( , γ ) -stationary point of f if there exists some θ∗ such that ‖∇f ( θ∗ ) ‖ ≤ and ‖θ − θ∗‖ ≤ γ . Intuitively , Algorithm 1 converges to an ( , γ ) -stationary point when it converges to a neighborhood of some -stationary point θ∗ .

Summary: -The paper studies the behavior of SGD with label noise for over-parametrized models. The label noise setting is linked to an implicit regularization scheme, which is helpful to improve the search for "better" optimal as well as to understand the optimization landscape. The study focuses on the analysis for the "flat" regions. Theoretical analysis has been provided and empirical studies on ResNet18 and CIFAR10 are shown.