Layer rotation: a surprisingly powerful indicator of generalization in deep networks?

TL;DR

This paper demonstrates that the evolution of layer rotation, measured by cosine distance from initialization, is a strong indicator of generalization in deep networks, with optimal performance when layers reach a cosine distance of 1.

Contribution

It introduces layer rotation as a novel, consistent indicator of generalization and shows how it can be monitored and controlled to improve training outcomes.

Findings

Larger layer rotation correlates with better generalization.

Optimal generalization occurs when layers reach a cosine distance of 1 from initialization.

Layer rotation can be used to guide hyperparameter tuning.

Abstract

Our work presents extensive empirical evidence that layer rotation, i.e. the evolution across training of the cosine distance between each layer's weight vector and its initialization, constitutes an impressively consistent indicator of generalization performance. In particular, larger cosine distances between final and initial weights of each layer consistently translate into better generalization performance of the final model. Interestingly, this relation admits a network independent optimum: training procedures during which all layers' weights reach a cosine distance of 1 from their initialization consistently outperform other configurations -by up to 30% test accuracy. Moreover, we show that layer rotations are easily monitored and controlled (helpful for hyperparameter tuning) and potentially provide a unified framework to explain the impact of learning rate tuning, weight decay, learning rate warmups and adaptive gradient methods on generalization and training speed. In an attempt to explain the surprising properties of layer rotation, we show on a 1-layer MLP trained on MNIST that layer rotation correlates with the degree to which features of intermediate layers have been trained.