[论文解读] Predict Globally, Correct Locally: Parallel-in-Time Optimal Control of Neural Networks

Panos Parpas, Corey Muir|arXiv (Cornell University)|Feb 7, 2019

Stochastic Gradient Optimization Techniques参考文献 22被引用 11

一句话总结

该论文提出了一种并行时间最优控制方法，用于通过将残差神经网络的层建模为动力系统中的离散时间点来训练网络。通过多级时间离散化和全局预测阶段的协态近似，实现了层间并行，相较于数据并行的SGD方法，实现了更高的加速比和效率，且在弱于先前方法的条件下仍保持相近的精度和收敛性保证。

ABSTRACT

The links between optimal control of dynamical systems and neural networks have proved beneficial both from a theoretical and from a practical point of view. Several researchers have exploited these links to investigate the stability of different neural network architectures and develop memory efficient training algorithms. We also adopt the dynamical systems view of neural networks, but our aim is different from earlier works. We exploit the links between dynamical systems, optimal control, and neural networks to develop a novel distributed optimization algorithm. The proposed algorithm addresses the most significant obstacle for distributed algorithms for neural network optimization: the network weights cannot be updated until the forward propagation of the data, and backward propagation of the gradients are complete. Using the dynamical systems point of view, we interpret the layers of a (residual) neural network as the discretized dynamics of a dynamical system and exploit the relationship between the co-states (adjoints) of the optimal control problem and backpropagation. We then develop a parallel-in-time method that updates the parameters of the network without waiting for the forward or back propagation algorithms to complete in full. We establish the convergence of the proposed algorithm. Preliminary numerical results suggest that the algorithm is competitive and more efficient than the state-of-the-art.

研究动机与目标

• 为克服分布式深度学习中前向传播和反向传播的串行瓶颈。
• 通过利用残差网络的动力系统解释，实现在神经网络训练中的真正层间并行。
• 开发一种分布式优化算法，避免等待完整的前向和反向传播。
• 在弱于现有合成梯度方法的假设下确保收敛性。
• 在深层网络上实现高并行效率和加速比，尤其适用于大规模模型。

提出的方法

• 将残差神经网络建模为离散时间最优控制问题，其中各层对应于动力系统中的时间步。
• 在全局预测阶段使用粗粒度时间离散化，以近似网络中所有位置的最优状态和协态（伴随变量）。
• 通过局部校正阶段，利用精细离散化对预测的状态和协态进行修正，从而实现并行更新。
• 采用多级时间离散化（如欧拉法和Verlet格式）以加速协态估计，并减少对完整反向传播的依赖。
• 将网络沿时间轴（层）分解，并使用基于伴随的梯度近似并行优化子系统。
• 在反向传播中引入回归步骤以校正协态估计，提高准确性和收敛性。

实验结果

研究问题

• RQ1 能否通过动力系统视角克服深度神经网络中前向和反向传播的串行性？
• RQ2 最优控制理论能否实现残差网络的并行时间训练，且精度损失最小？
• RQ3 与合成梯度方法相比，采用全局预测与局部校正的多级时间离散化方案是否能提升收敛性和效率？
• RQ4 协态近似质量对分布式设置下训练稳定性与加速比的影响如何？
• RQ5 所提方法在不同网络深度和数据集上的加速比与并行效率如何扩展？

主要发现

• 所提方法在加速比方面与数据并行SGD相当或更优，尤其在深层网络（如>32层）中表现突出，在MNIST和Swiss-Roll数据集上并行效率超过75%。
• 在Swiss-Roll数据集上，该方法在32–64层网络中优于数据并行SGD，且随着网络深度增加，加速比进一步提升。
• 反向传播中回归步骤的均方误差在使用全局预测阶段后降低一个数量级，趋近于零，验证了方法的稳定性和收敛性假设。
• 多级Verlet离散化方案在MNIST和合成数据集上的结果表明，其精度与随机梯度下降相当。
• 该方法的并行效率（75%）显著高于以往并行时间方法（如Günther等人，2018年报告的3–4%）和Huo等人（2018年）报告的50%。
• 数值结果表明，该方法与最先进方法具有竞争力，且随着模型深度增加，其性能因粗粒度模型求解的相对成本降低而进一步提升。