[Paper Review] Stiffness: A New Perspective on Generalization in Neural Networks

Stanislav Fort, Paweł Krzysztof Nowak+2 more|arXiv (Cornell University)|Jan 28, 2019

Neural Networks and Applications28 references60 citations

TL;DR

The paper introduces stiffness as a measure of how gradient updates from one example affect losses on others, linking gradient alignment to generalization and analyzing it across datasets, architectures, and learning rates.

ABSTRACT

In this paper we develop a new perspective on generalization of neural networks by proposing and investigating the concept of a neural network stiffness. We measure how stiff a network is by looking at how a small gradient step in the network's parameters on one example affects the loss on another example. Higher stiffness suggests that a network is learning features that generalize. In particular, we study how stiffness depends on 1) class membership, 2) distance between data points in the input space, 3) training iteration, and 4) learning rate. We present experiments on MNIST, FASHION MNIST, and CIFAR-10/100 using fully-connected and convolutional neural networks, as well as on a transformer-based NLP model. We demonstrate the connection between stiffness and generalization, and observe its dependence on learning rate. When training on CIFAR-100, the stiffness matrix exhibits a coarse-grained behavior indicative of the model's awareness of super-class membership. In addition, we measure how stiffness between two data points depends on their mutual input-space distance, and establish the concept of a dynamical critical length -- a distance below which a parameter update based on a data point influences its neighbors.

Motivation & Objective

• Motivate and formalize the stiffness concept as a probe of generalization in neural networks.
• Investigate how stiffness depends on class membership, data-point distance in input space, training epoch, and learning rate.
• Demonstrate stiffness behavior across vision (MNIST, FASHION-MNIST, CIFAR-10/100) and a transformer-based NLP model.
• Examine the dynamical critical length and semantic group structure (super-classes) revealed by stiffness.

Proposed method

• Define stiffness via two gradent-based measures: sign stiffness (sign of g1·g2) and cosine stiffness (cosine similarity of g1 and g2).
• Compute how a small update from one input X1 with gradient g1 changes the loss on another input X2.
• Construct class stiffness matrices C(ca, cb) and analyze between-class versus within-class stiffness.
• Assess stiffness in train-train, train-val, and val-val settings to relate to generalization.
• Measure stiffness as a function of input-space distance using a dynamical critical length xi.
• Evaluate stiffness across learning rates and epochs to observe how higher rates bias toward lower, more local stiffness.

Experimental results

Research questions

• RQ1 How is neural network stiffness defined and what does it reveal about generalization?
• RQ2 How does stiffness vary with class membership and semantic groupings (including super-classes) across datasets?
• RQ3 How does stiffness depend on the input-space distance between datapoints?
• RQ4 What is the impact of training epoch and learning rate on stiffness and the dynamical critical length xi?
• RQ5 Does stiffness behavior generalize across vision and language models (including CNNs, ResNets, and BERT)?

Key findings

• Stiffness correlates with generalization: higher stiffness within and between classes is observed during learning, but declines with overfitting.
• Within-class stiffness remains high early and during learning, while between-classes stiffness increases as the model learns; both fade as overfitting begins.
• Stiffness reveals semantically meaningful group structure: higher stiffness within super-classes and even super-super-classes than random baselines in CIFAR-100.
• A dynamical critical length xi exists: stiffness decays to zero as input-space distance increases; xi decreases with training and with higher learning rates.
• Higher learning rates produce functions with smaller xi, i.e., more local and more easily bendable updates, indicating a regularizing effect on the learned function.
• Stiffness concepts extend to NLP (BERT fine-tuned on MNLI), exhibiting similar within-class and between-class dynamics as vision models.