Balancing Learning Rates Across Layers: Exact Two-Step Dynamics and Optimal Scaling in Linear Neural Networks
We study optimal learning-rate selection in two-layer and three-layer linear neural networks trained to learn a single-index target function. In particular, we derive the exact closed-form expressions for the gradients and test loss after one and two steps of gradient descent, enabling a precise cha…