提出线性网络中分层学习率的最优平衡方法，揭示训练初期不等速率更优。
Balancing Learning Rates Across Layers: Exact Two-Step Dynamics and Optimal Scaling in Linear Neural Networks
- 推导两步梯度下降的闭式解，精确刻画早期训练动态
- 发现初始步骤用不等学习率可最小化测试损失，后续则等率最优
- 理论与实验结合，适合研究优化机制的学者参考
我们研究在两层和三层线性神经网络中，为学习线性目标函数而进行的最优学习率选择问题。具体地，我们推导了梯度和测试损失在梯度下降前一步和两步后的精确闭式表达式，从而能够对早期训练动态进行精准刻画。我们分析了在梯度近似下学习率在前两步应如何缩放，并证明使用该近似进行更新可得到一个可处理的代理损失，其近似误差小且紧致。这一形式化使分层学习率的理论分析成为可能，揭示出一个独特的早期训练阶段：测试损失可通过初始步骤采用不等学习率来最小化，而在后续步骤中等学习率变为最优。数值实验验证了理论结果，并展示了早期平衡分层学习率的重要性。代码已公开于：https://github.com/TDCSZ327/Layer-Balancing。
原文摘要 · Abstract (English)
We study optimal learning-rate selection in two-layer and three-layer linear neural networks trained to learn linear target functions. 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 characterization of early training dynamics. We characterize how learning rates should scale under the gradient approximation in the first two steps, and prove that performing updates with this approximation yields a tractable surrogate loss with a tight, small approximation error. This formulation enables the theoretical analysis of layer-wise learning rates and reveals a distinct early-training regime: test loss can be minimized by unequal learning rates at the initial step, while equal learning rates become optimal in subsequent steps. Our numerical experiments validate the theory and demonstrate the importance of balancing layer-wise learning rates early during training. The code is available at: https://github.com/TDCSZ327/Layer-Balancing.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。