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 θℓ between the feature updates and the backward pass (at layer index ℓ). 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 θℓ, the loss decay, and the magnitude of the backward pass. This angle θℓ 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 ℓ=depth. Given θℓ, 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(1/√ℓ). In contrast, ResNets with branch scale O(1/√depth) maintain a non-degenerate angle cos(θℓ)=Theta(1). We use these insights to recover key properties of known HP scalings and also to introduce a new HP scaling for large depth ReLU MLPs with favorable theoretical properties.