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 submission introduces a metric, termed stiffness, to evaluate the generalization capability of neural networks. The metric is novel and straightforward, it measures how stiff a network is by looking at how a small gradient step on one example affects the loss on another example. The authors study several configurations  on three small datasets. They demonstrate that stiffness is a useful concept for diagnosing and characterizing generalization.