Stiffness: A New Perspective on Generalization in Neural Networks

1 INTRODUCTION . Neural networks are a class of highly expressive function approximators that proved to be successful in approximating solutions to complex tasks across many domains such as vision , natural language understanding , and game-play . They have long been recognized as universal function approximators ( Hornik et al. , 1989 ; Cybenko , 1989 ; Leshno et al. , 1993 ) . The specific details that lead to their expressive power have recently been studied in Montúfar et al . ( 2014 ) ; Raghu et al . ( 2017 ) ; Poole et al . ( 2016 ) . Empirically , neural networks have been extremely successful at generalizing to new data despite their over-parametrization for the task at hand , as well as their proven ability to fit arbitrary random data perfectly Zhang et al . ( 2016 ) ; Arpit et al . ( 2017 ) . The fact that gradient descent is able to find good solutions given the highly over-parametrized family of functions has been studied theoretically in Arora et al . ( 2018 ) and explored empirically in Li et al . ( 2018 ) , where the effective low-dimensional nature of many common learning problems is shown . Fort & Scherlis ( 2019 ) extends the analysis in Li et al . ( 2018 ) to demonstrate the role of initialization on the effective dimensionality , and Fort & Jastrzebski ( 2019 ) use the result to build a phenomenological model of the loss landscape . Du et al . ( 2018a ) and Du et al . ( 2018b ) use a Gram matrix to study convergence in neural network empirical loss . Pennington & Worah ( 2017 ) study the concentration properties of a similar covariance matrix formed from the output of the network . Ghorbani et al . ( 2019 ) investigate the Hessian eigenspectrum and Papyan ( 2019 ) show how it is related to the gradient covariance . Both concepts are closely related to our definition of stiffness . To explain the remarkable generalization properties of neural networks , it has been proposed ( Rahaman et al. , 2018 ) that the function family is biased towards low-frequency functions . The role of similarity between the neural network outputs to similar inputs has been studied in Schoenholz et al . ( 2016 ) for random initializations and explored empirically in Novak et al . ( 2018 ) . 1.1 OUR CONTRIBUTION . In this paper , we study generalization through the lens of stiffness . We measure how stiff a neural network is by analyzing how a small gradient step based on one input example affects the loss on another input example . Mathematically , if the gradient of the loss at point X1 with respect to the network weights is ∇WL ( X1 ) = ~g1 , and the gradient at point X2 is ~g2 , we define stiffness ∝ ~g1 · ~g2 . We specifically focus on the sign of ~g1 · ~g2 as well as the cosine between the two vectors cos ( ~g1 , ~g2 ) = ĝ1 · ĝ2 , where ĝ = ~g/|~g| , which both capture the resistance of the functional approximation learned to deformation by gradient steps . We find the concept of stiffness useful in diagnosing and characterizing generalization . As a corollary , we use stiffness to characterize the regularization power of learning rate , and show that higher learning rates bias the functions learned towards lower stiffness . We show that stiffness is directly related to generalization , and in particular that it starts dropping sharply at the moment training and validation set losses stop evolving together . We explore the concept of stiffness for fully-connected ( FC ) and convolutional neural networks ( CNN ) on 3 classification datasets ( MNIST , FASHION MNIST , CIFAR-10 ) . We focus on how stiffness between data points depends on their 1 ) class membership , 2 ) distance between each other in the space of inputs , 3 ) training epoch , and 4 ) the choice of learning rate . We study stiffness between two images in the training set , one in the training and one in the validation set , and both images in the validation set . We observed the stiffness based on class membership and noticed a clear evolution towards higher stiffness for images between different classes and towards lower stiffness within the same class . We diagnose and characterize the class-dependent stiffness matrix for fully-connected and convolutional neural networks on the datasets mentioned above in different stages of training . We observe the stiffness between inputs to regress to zero with the onset of overfitting , demonstrating the clear connection to generalization . The choice of learning rate effects the stiffness properties of the learned function significantly . High learning rates induce functional approximations that are less stiff over larger distances ( i.e . data points further apart stop responding similarly to gradient updates ) . We define the concept of dynamical critical length to capture the phenomenon . This paper is structured as follows : we introduce the concept of stiffness and the relevant theory in Section 2 . We describe our experimental setup in Section 3 , and discuss their results in Section 4 . We conclude with Section 5 . 2 THEORETICAL BACKGROUND . 2.1 STIFFNESS – DEFINITIONS . Let a functional approximation ( e.g . a neural network ) f be parametrized by tunable parameters W . Let us assume a classification task and let a data point X have the ground truth label y . A loss L ( fW ( X ) , y ) gives us the amount of mismatch between the function ’ s output at input X and the ground truth y . The gradient of the loss with respect to the parameters ~g = ∇WL ( fW ( X ) , y ) ( 1 ) is the direction in which , if we were to change the parameters W , the loss would change the most rapidly ( for infinitesimal step sizes ) . Gradient descent uses this step ( the negative of it ) to update the weights and gradually tune the functional approximation to better correspond to the desired outputs on the training dataset inputs . Let us consider two data points with their ground truth labels ( X1 , y1 ) and ( X2 , y2 ) . We construct a gradient with respect to example 1 as ~g1 = ∇WL ( fW ( X1 ) , y1 ) and ask how do the losses on data points 1 and 2 change as the result of a small change of W in the direction −~g1 , i.e . what is ∆L1 = L ( fW−ε~g1 ( X1 ) , y1 ) − L ( fW ( X1 ) , y1 ) , ( 2 ) which is equivalent to ∆L1 = −ε∇εL ( fW−ε~g1 ( X1 ) , y1 ) = −ε~g1 · ~g1 . ( 3 ) The change in loss on input 2 due to the same gradient step from input 1 becomes equivalently ∆L2 = −ε∇εL ( fW−ε~g1 ( X2 ) , y2 ) = −ε~g1 · ~g2 We are interested in the correlation in loss changes ∆L1 and ∆L2 . We know that ∆L1 < 0 since we constructed the gradient update accordingly . We define positive stiffness to mean ∆L2 < 0 as well , i.e . that losses at both inputs went down . We assign the stiffness of 0 for ∆L2 = 0 . If ∆L2 > 0 , the two inputs would be anti-stiff ( negative stiffness ) . The equations above show that this can equivalently be thought of as the overlap between the two gradients ~g1 · ~g2 being positive for positive stiffness , and negative for negative stiffness . We illustrate this in Figure 1 . The above indicate that what we initially conceived of as a change in loss due to the application of a small gradient update from one input to another is in fact equivalent to analyzing gradient alignment between different datapoints . We will be using 2 different definitions of stiffness : the sign stiffness and the cosine stiffness . We define the sign stiffness to be the expected sign of ~g1 · ~g2 ( or equivalently the expected sign of ∆L1∆L2 ) as Ssign ( ( X1 , y1 ) , ( X2 , y2 ) ; f ) = E [ sign ( ~g1 · ~g2 ) ] , ( 4 ) where stiffness depends on the dataset from which X1 and X2 are drawn . The cosine stiffness is Scos ( ( X1 , y1 ) , ( X2 , y2 ) ; f ) = E [ cos ( ~g1 · ~g2 ) ] , ( 5 ) where cos ( ~g1 · ~g2 ) = ( ~g1/|~g1| ) · ( ~g2/|~g2| ) . We use both versions of stiffness as they are suitable to highlight different phenomena – the sign stiffness shows the stiffness between classes clearer , while the cosines stiffness is more useful for within-class stiffness . 2.2 TRAIN-TRAIN , TRAIN-VAL , AND VAL-VAL . When measuring stiffness between two datapoints , we have 3 options : 1 ) choosing both datapoints from the training set ( we call this train-train ) , 2 ) choosing one from the training set and the other from the validation set ( train-val ) , and 3 ) choosing both from the validation set ( val-val ) . The trainval stiffness is directly related to generalization , as it corresponds to the amount of improvement on the training set transferring to the improvement of the validation set . We empirical observe that all 3 options are empirically remarkably similar , which gives us confidence that they all track generalization . This is further supported by observing their behavior is a function of epoch in Figure 2 . 2.3 STIFFNESS BASED ON CLASS MEMBERSHIP . A natural question to ask is whether a gradient taken with respect to an input X1 in class c1 will also decrease the loss for example X2 with true class c2 . In particular , we define the class stiffness matrix C ( ca , cb ) = E X1∈ca , X2∈cb [ S ( ( X1 , y1 ) , ( X2 , y2 ) ) ] . ( 6 ) The on-diagonal elements of this matrix correspond to the suitability of the current gradient update to the members of a class itself . In particular , they correspond to within class generalizability . The off-diagonal elements , on the other hand , express the amount of improvement transferred from one class to another . They therefore directly diagnose the amount of generality the currently improved features have . A consistent summary of generalization between classes is the off-diagonal sum of the class stiffness matrix Sbetween classes = 1 Nc ( Nc − 1 ) ∑ c1 ∑ c2 6=c1 C ( c1 , c2 ) . ( 7 ) In our experiments , we track this value as a function of learning rate once we reached a fixed loss . The quantity is related to how generally applicable the learned features are , i.e . how well they transfer from one class to another . For example , for CNNs learning good edge detectors in initial layers typically benefits all downstream tasks , regardless of the particular class in question . We do the equivalent for the within-class stiffness ( = on diagonal elements ) . When the within-class stiffness starts going < 1 , the generality of the features improved does not extend to even the class itself . 2.4 STIFFNESS AS A FUNCTION OF DISTANCE . We investigate how stiff two inputs are based on how far away from each other they are . We can think of neural networks as a form of kernel learning and here we are investigating the particular form of the learned kernel . This links our results to the work on spectral bias ( towards slowly-varying , low frequency functions ) in Rahaman et al . ( 2018 ) . We are able to directly measure the characteristic size of the stiff regions in neural networks trained on real tasks which we call dynamical critical length ξ . We work with data normalized to the unit sphere | ~X| = 1 and use their mutual cosine to define their distance as distance ( ~X1 , ~X2 ) = 1− ~X1 · ~X2 | ~X1|| ~X2| . ( 8 ) which has the advantage of being bounded between 0 and 2 . We track this threshold distance ξ as a function of training and learning rate to estimate the characteristic size of the stiff regions of a neural net .

Summary: This paper introduces “stiffness”, a new metric to characterize generalization in neural networks. Stiffness is a pretty simple concept and is relatively straightforward to compute. The authors evaluate this metric on standard datasets using two relatively small neural networks. On the whole, the paper is written clearly and explains its methodology in simple language.