The Implicit Bias of Minima Stability: A View from Function Space

1 Introduction . Understanding the overwhelming success of deep learning requires unveiling the mechanisms that allow over-parametrized models to generalize well , a phenomenon that is in sharp contrast to classical wisdom . It has been suggested that one of the sources for this behavior is the implicit bias that training algorithms have towards certain solutions . Implicit biases of generic optimization methods had been known for decades [ 49 ] and have also been discussed in the context of neural network training already in [ 30 ] . However , their dramatic effect on modern deep learning has only recently started to unveil [ 36 , 57 ] . For example , in binary classification tasks with linearly separable data , it has been shown that among all global minima , gradient descent ( GD ) converges to the maximum margin separator [ 48 ] . Similarly , in over-parametrized linear regression , GD converges to the minimum norm solution when initialized at zero [ 57 ] . Implicit biases were recently studied in many different settings , including in linear convolutional networks [ 16 ] , matrix and tensor factorization [ 14 , 41 ] , with weight normalization [ 54 ] and with various loss functions [ 15 ] . Some works have drawn analogies between these biases and traditional regularization schemes . For example , it has been conjectured that the implicit bias of GD can be expressed as some regularization term that is added to the training loss [ 14 ] . This turned out to be true for several settings , such as over-parametrized linear regression with the quadratic loss [ 57 ] , matrix factorization under certain assumptions [ 26 ] , and linear classification on separable data using losses with an exponential tail [ 48 ] . But recent work has pointed out that this is not true in general , 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . by illustrating that the implicit bias of GD is often not equivalent to a loss term of the form of any function of the network ’ s weights [ 40 , 44 ] . The current paradigm for analyzing implicit bias seeks to ensure convergence to global minima as a first step . This requires making sure that GD is well behaved , i.e . does not escape the minimum . A common way to guarantee this is by considering a sufficiently small learning rate , or even the limit of infinitesimal step size , known as gradient flow ( GF ) . In this case , the minima to which GD converges tend to be close to the initialization [ 39 ] . Therefore , existing results are initialization-dependent and relevant only to small step-sizes . Unfortunately , these type of analyses fail to capture phenomena that are known to be directly related to the step size in practical training , e.g . [ 8 , 17 , 19 , 24 , 29 , 45 , 46 , 53 ] . For instance , Fig . 1 demonstrates that when initializing GD near a minimum , different step sizes lead to convergence to completely different solutions . This behavior is not reflected by current studies . Following this understanding , here we seek an initialization-independent analysis that explicitly takes into account the learning rate and reveals its effect on the learned model . To this end , instead of analyzing the trajectory of the parameters from initialization to convergence , we study the properties of the minima to which stochastic GD ( SGD ) can converge . Specifically , it is well known that SGD can not stably converge to minima that are too sharp relative to its step size [ 8 , 19 , 45 , 53 ] . This property has been studied in [ 53 ] for twice differential minima , but its implication on the learned function has not been discussed . Here we extend this analysis to non-differentiable minima , which are common in ReLU networks , and use this property to characterize the end-to-end functions that a neural network can implement upon convergence . As in [ 23 , 42 , 51 ] , we study an over-parameterized single-hidden-layer univariate ReLU network , and focus on the quadratic loss . Under this setting , the network can implement infinitely many piecewise linear functions f that globally minimize the loss ( i.e . interpolate the data ) . Our key result ( see Theorem 1 ) is that when using a step size of η , SGD can only converge to solutions satisfying∫ R |f ′′ ( x ) |g ( x ) dx ≤ 1 η − 1 2 , ( 1 ) with some particular weight function g ( see Fig . 1 ) . In other words , the larger the step size , the smoother the solutions that a network learns . The weight g is larger around the center of the support of the training distribution , and smaller towards its boundaries . This implies that the solutions found by SGD tend to be smoother around the center of the distribution . 2 Minima stability . SGD is routinely used to minimize objective functions that have multiple global minima . However , it is well known that not all minima are accessible to SGD [ 53 ] . Understanding to which minima SGD can converge , requires analyzing its dynamics in the vicinity of minima . If once SGD arrives near a minimum , it converges to it , then we say that this is a stable minimum . If SGD repels from the minimum , then we say that this is an unstable minimum . Here we briefly survey and extend the existing knowledge about stable minima for SGD . Let ` j : Rd 7→ R be differentiable almost everywhere for all j ∈ [ n ] . Consider a loss function L and its stochastic counterpart L̂ , given by L ( θ ) = 1 n n∑ j=1 ` j ( θ ) and L̂t ( θ ) = 1 B ∑ j∈Bt ` j ( θ ) , ( 2 ) where Bt is a batch of size B sampled at iteration t. We assume that the batches { Bt } are drawn uniformly from the training set , independently across iterations . SGD ’ s update rule is given by θt+1 = θt − η∇L̂t ( θt ) , ( 3 ) where η is the step size . In the following we define the notion of stability for minima , and provide conditions for a minimum to be stable . 2.1 Twice differentiable minima . We start by examining a twice differentiable minimum θ∗ . Using a Taylor expansion about θ∗ , we have L̂t ( θ ) ≈ L̂t ( θ∗ ) + ( θ − θ∗ ) T∇L̂t ( θ∗ ) + 1 2 ( θ − θ∗ ) T∇2L̂t ( θ∗ ) ( θ − θ∗ ) , ( 4 ) where∇L̂t ( θ∗ ) and∇2L̂t ( θ∗ ) are the gradient and Hessian of L̂t at θ∗ . Therefore , in the vicinity of the minimum , ( 3 ) is approximately given by θt+1 ≈ θt − η ( ∇L̂t ( θ∗ ) +∇2L̂t ( θ∗ ) ( θt − θ∗ ) ) . ( 5 ) This approximation gets more accurate as θt gets closer to θ∗ . We can thus use this linearized dynamics to learn about the stability of minima [ 53 ] . Definition 1 ( Linear stability ) . Let θ∗ be a twice differentiable minimum ofL . Consider the linearized stochastic dynamical system θt+1 = θt − η ( ∇L̂t ( θ∗ ) +∇2L̂t ( θ∗ ) ( θt − θ∗ ) ) . ( 6 ) Then θ∗ is said to be ε linearly stable if for any θ0 in the ε-ball Bε ( θ∗ ) , we have lim sup t→∞ E [ ‖θt − θ∗‖ ] ≤ ε . In other words , θ∗ is ε linearly stable if once we have arrived at a distance of ε from it ( at t = 0 without loss of generality ) , we end up at a distance no greater than ε around it in expectation . Some of our results do not depend on ε , in which case we simply refer to “ linear stability ” ( without the ε ) . Wu et al . [ 53 ] provided a sufficient condition for linear stability under the assumption that ∇L̂t ( θ∗ ) = 0 for all t ≥ 1 . For our analysis , we need a necessary condition . This is simple to obtain from the well known stability criterion for GD and the fact that GD ’ s trajectory corresponds to the expectation of SGD ’ s steps . Specifically , we have the following condition ( which only requires that E [ ∇L̂t ( θ∗ ) ] = 0 ; see formal proof in Appendix II ) . Lemma 1 ( Necessary condition for stability ) . Consider SGD with step size η , where batches are drawn uniformly from the training set , independently across iterations . If θ∗ is an ε linearly stable minimum of L , then λmax ( ∇2L ( θ∗ ) ) ≤ 2 η . ( 7 ) Note that this condition is also sufficient when considering GD ( i.e . full batch SGD ) . 2.2 Non-differentiable minima . We now generalize the definition of linear stability to non-differentiable minima . In deep learning , non-differentiability is typically caused by a mode switch in the system , e.g . switching of ReLU activations or max-pooling layers . However , if we fix the mode ( e.g . constant activation or pooling selection patterns per sample ) , then the loss becomes infinitely differentiable everywhere . In these cases , we can therefore model SGD ’ s dynamics as a switching dynamical system . Let { Sm } be a partition of Rd that represents the regions of the different modes , i.e . ∀i 6= j Si ∩ Sj = ∅ , and ⋃ m Sm = Rd . ( 8 ) Additionally , let ψm : Rd 7→ R be an analytic function representing the loss for the mth mode . We assume the overall loss ( and its stochastic version ) can be written as1 L ( θ ) = ψm ( θ ) , L̂t ( θ ) = ψ̂ ( t ) m ( θ ) if θ ∈ Sm , ( 9 ) where ψ̂ ( t ) m is the stochastic counterpart of ψm at time t. Therefore , near a minimum θ ∗ we can approximate the loss as L̂t ( θ ) ≈ L̂t ( θ∗ ) + ( θ − θ∗ ) T ĝ ( t ) θ + 1 2 ( θ − θ∗ ) T Ĥ ( t ) θ ( θ − θ ∗ ) , ( 10 ) where ( defining Int ( A ) as the interior of A ) ∀θ ∈ Int ( Sm ) ĝ ( t ) θ , ∇ψ̂ ( t ) m ( θ ∗ ) , and Ĥ ( t ) θ , ∇2ψ̂ ( t ) m ( θ ∗ ) . ( 11 ) The update rule of SGD ( 3 ) can thus be approximated as θt+1 ≈ θt − η ( ĝ ( t ) θt + Ĥ ( t ) θt ( θt − θ ∗ ) ) . ( 12 ) Note that this approximation is relevant only when θt is close to θ∗ . Let2 I = { m : θ∗ ∈ S̄m } be the indices of the sets around θ∗ , and let A = ⋃ m∈I Sm be their union . We assume there is a finite number of modes inA ( this is always the case for ReLU networks with a finite number of parameters ) . Consider an ε-neighborhood around the minimum such that Bε ( θ∗ ) ⊆ A . We define linear stability in this neighborhood as follows . Definition 2 ( Generalized linear stability ) . Let θ∗ be a minimum point of L. Consider the switching stochastic dynamical system θt+1 = θt − η ( ĝ ( t ) θt + Ĥ ( t ) θt ( θt − θ ∗ ) ) ( 13 ) and assume that θt ∈ A with probability one for all t > 0 . • θ∗ is said to be ε linearly stable if lim sup t→∞ E [ ‖θt − θ∗‖ ] ≤ ε for any θ0 ∈ Bε ( θ∗ ) . • θ∗ is said to be ε linearly strongly stable if sup t E [ ‖θt − θ∗‖ ] ≤ ε for any θ0 ∈ Bε ( θ∗ ) . In other words , we consider a situation where at some point in time ( which we call t = 0 without loss of generality ) , the parameter vector is ε-close to θ∗ . If this guarantees that from that moment on , θt is always ε-close to θ∗ in expectation , then θ∗ is called strongly stable . If this only guarantees that from some later point in time , θt is always ε-close to θ∗ in expectation , then θ∗ is called stable . Our results are stated in terms ofHm = ∇2ψm ( θ ∗ ) and ĝ ( t ) m = ∇ψ̂ ( t ) m ( θ ∗ ) . The next lemma gives a necessary condition for linear stability ( see proof in Appendix III ) . Lemma 2 ( Necessary condition for stability ) . Assume that SGD with step size η draws batches uniformly from the training set , independently across iterations . Let θ∗ be a minimum point of L , for which there is a finite number of modes in I . Suppose there exist q ∈ Sd−1 and { λm } such that ‖Hmq − λmq‖ ≤ δ for all m ∈ I and denote λlower = min m∈I { λm } , ( 14 ) and3 γ = maxm∈I E [ |qT ĝ ( t ) m | ] . If λlower > 2 η + δ + γ ε , ( 15 ) then θ∗ is not an ε strongly stable minimum . Furthermore , if δ = 0 ( i.e . q is a common eigenvector ) then θ∗ is not an ε stable minimum . 1This is true since the modes corresponding to a batch are a subset of those of the full data . 2S̄m denotes the closure of Sm . 3Note that { ĝ ( t ) m } ∞t=1 are i.i.d . and therefore E [ |qT ĝ ( t ) m | ] is time invariant . Note that the lemma considers the case where the Hessians of all subsystems have an approximate common eigenvector q . In this setting , a necessary condition for stability is that the smallest ( approximate ) eigenvalue associated with q is not too large w.r.t . 2/η . Interestingly , in the setting of ReLU networks , this ( approximate ) common eigenvector assumption holds true ( see Lemma 5 ) . We thus make use of this property in our analysis of the solutions to which SGD can converge . For completeness , we also present here a sufficient condition for stability in the special case of GD where ∇ψm ( θ ∗ ) = 0 for all m ∈ I ( see proof in Appendix IV ) . This setting is of interest in our case , since the global minima of overparameterized ReLU networks under the quadratic loss always satisfy this condition ( see proof in Appendix VII for our setting ) . Lemma 3 ( Sufficient conditions for stability ) . Consider full batch SGD ( i.e . GD ) with step size η . Let θ∗ be a minimum point of L , for which∇ψm ( θ ∗ ) = 0 for all m ∈ I. Denote λuppermax = max m∈I λmax ( Hm ) . ( 16 ) If λuppermax ≤ 2 η ( 17 ) then θ∗ is linearly strongly stable . Here we see that if all the subsystems are stable , then the overall switching system is strongly stable .

Summary: This paper studies the smoothness of the functions represented by single variate two-layer neural networks at minima with different flatness, and connect this smoothness with the implicit bias of (stochastic) gradient descent to flat minima. Firstly, for a twice differentiable minimum found by SGD, an upper bound for the leading eigenvalue of the loss function's Hessian matrix (with respect to parameters) is derived depending on the learning rate. Then, a weighted L1 norm of the second derivative of the model function (with respect to input data) is upper bounded by the Hessian's leading eigenvalue. Combining these two bounds, the weighted L1 norm of the model function's second derivative is bounded by a quantity depending on the learning rate. This result implies that SGD with bigger learning rate finds flatter minima, and thus finds smoother functions. The result only considers the minima found by optimization algorithms, hence it is independent with initialization.   For neural networks with ReLU activation, minima are sometimes not differentiable. The analysis above is extended to this case and similar bounds are also provided. The theorem works in the mildly over-parameterized case in which the number of neurons is slightly bigger than the number of training data.   Numerical experiments show that GD with bigger learning rate can indeed find smoother functions, especially in the region where training data are densely distributed. This observation is consistent with the theoretical results.