The Feature Speed Formula: a Flexible Approach to Scale Hyper-parameters of Deep Neural Networks

Deep learning succeeds by doing hierarchical feature learning, yet tuning hyper-parameters (HP) such as initialization scales, learning rates etc., only give indirect control over this behavior. In this paper, we introduce a key notion to predict and control feature learning: the angle theta(l) between the feature updates and the backward pass (at layer index l). We show that the magnitude of feature updates after one GD step, at any training time, can be expressed via a simple and general feature speed formula in terms of this angle theta(l), the loss decay, and the magnitude of the backward pass. This angle theta(l) is controlled by the conditioning of the layer-to-layer Jacobians and at random initialization, it is determined by the spectrum of a certain kernel, which coincides with the Neural Tangent Kernel when l = depth. Given theta(l), the feature speed formula provides us with rules to adjust HPs (scales and learning rates) so as to satisfy certain dynamical properties, such as feature learning and loss decay. We investigate the implications of our approach for ReLU MLPs and ResNets in the large width-then-depth limit. Relying on prior work, we show that in ReLU MLPs with iid initialization, the angle degenerates with depth as cos(theta(l)) = Theta(1/root l). In contrast, ResNets with branch scale O(1/root depth) maintain a non-degenerate angle cos(theta(l)) = Theta(1). We use these insights to recover key properties of known HP scalings (such as mu P), and also introduce a new HP scaling for large depth ReLU MLPs with favorable theoretical properties.