分析双层ReLU网络训练动态,证明能快速收敛到全局最优解。
When Both Layers Learn: Training Dynamics of Representing Linear Models via ReLU Networks

- 同时训练隐藏层和输出层,通过梯度下降优化。
- 在高斯输入下,以线性速率收敛至全局最小值。
- 适合研究深度学习优化机制的科研人员参考。
本文研究在可实现设定下,从随机初始化出发联合训练单隐层ReLU网络两层参数以拟合线性目标函数的梯度下降动态。输入独立同分布于高斯分布,标签服从预设线性模型。该框架捕捉了反问题和某些自编码器模型端到端训练的关键特征。尽管形式简单,但其损失曲面存在多个非严格鞍点,导致梯度下降能否有效避开不良驻点尚不明确。本文详细刻画了优化景观,证明从适度小的随机初始化出发,同时训练两层参数时,梯度下降以线性速率收敛至全局最小值,且样本复杂度阶数最优。分析揭示了三个阶段:对齐阶段,隐藏权重逐步对齐于预设方向,输出权重保持正确符号模式;增长阶段,两层范数增大但仍维持对齐;局部精修阶段,对齐神经元迅速逼近预设方向,实现快速局部收敛。为严格证明避免非严格鞍点,提出了轨迹级控制方法;并建立了沿整个训练轨迹成立的新型一致集中不等式,对获得阶数最优样本复杂度至关重要。实验在多种配置下验证了理论结果。
原文摘要 · Abstract (English)
In this paper, we study the gradient descent dynamics for jointly training both layers of a one-hidden-layer ReLU network to fit a linear target function. Concretely, we consider a realizable setting where inputs are drawn i.i.d. from a Gaussian distribution and labels follow a planted linear model. This stylized framework captures salient features of end-to-end training in inverse problems and certain auto-encoder models. Despite its apparent simplicity, the dynamics remain poorly understood, in part because the loss landscape contains multiple non-strict saddle points, making it unclear why gradient descent from random initialization reliably escapes bad stationary regions. We provide a detailed characterization of the optimization landscape and prove that gradient descent from a moderately small random initialization-simultaneously training both layers-converges to a global minimizer at a linear rate with order-wise optimal sample complexity. Our analysis tracks the trajectory through three phases: an alignment phase in which hidden weights progressively align with the planted direction while the output weights maintain the correct sign pattern; a growth phase in which the norms of both layers increase while preserving alignment; and a local refinement phase in which the aligned neurons rapidly converge to the planted direction, yielding fast local convergence. To rigorously show that GD avoids non-strict saddles, we develop trajectory-level control arguments for the end-to-end dynamics. In addition, we establish novel uniform concentration results that hold along the entire trajectory, and are essential for obtaining order-wise optimal sample complexity. We corroborate our theory with extensive experiments across a range of configurations.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。