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: The goal of this paper is to extend existing results on stable minima found by SGD to the non-differentiable setting of ReLU networks.   In particular, in its main contribution (Theorem 1) the paper shows that the solution found by SGD for 1-dimensional ReLU networks has bounded second derivative, under a weighted total variation norm. Moreover, this weighted norm is controlled by the step-size of SGD.  This implies that also for ReLU networks, SGD with large learning rate is biased towards smooth solutions.