A Non-Monotone Preconditioned Trust-Region Method for Neural Network Training

技术摘要：用于神经网络训练的非单调预条件信赖域方法

问题陈述
在大规模尺度上训练深度神经网络（NN）面临着超参数调优、多设备流水线利用以及处理非凸、噪声目标函数等挑战。虽然一阶方法如随机梯度下降（SGD）和 Adam 是标准方法，但它们严重依赖手动步长和动量调度，且无法自然地利用模型并行性（即网络被分割为跨设备的顺序块）。相反，域分解（DD）方法提供了一种将参数划分为子域以进行并行更新的途径，但它们需要稳健的全局化策略以确保在非凸设置下的收敛性。先前的方法，如加性预条件信赖域策略（APTS），利用信赖域（TR）机制，但往往遭受粗空间步长的高拒绝率，导致搜索方向需要昂贵的重新计算。

方法论：NAPTS
本文提出了非单调加性预条件信赖域策略（NAPTS），这是 APTS 框架的扩展，集成了非单调信赖域（NTR）原理。该方法在每个迭代中通过三个不同的阶段运行：

1. 并行子域更新：NN 参数空间被划分为非重叠子域 $\{C_d\}_{d=1}^N$。在每个迭代 $k$，全局前向和反向传播缓存激活值和下游梯度。每个子域随后并行执行 $\ell$ 次局部约束 Adam（CAdam）步。这些局部步利用缓存的边界信息计算，以近似块梯度，确保子域在独立运行的同时尊重全局信息。局部步被提升并加性组合，形成全局试验步 $s_k$。
2. 非单调全局化：与要求目标函数严格单调递减的标准 TR 方法不同，NAPTS 采用最近 $\nu$ 次成功迭代的滑动窗口。它定义了一个参考索引 $r(k)$，对应于该窗口内的最高目标值。步长接受度由组合一致比率 $\rho_k = \max\{\rho^c_k, \rho^h_k\}$ 决定，其中：
   • $\rho^c_k$ 衡量当前步的实际减少量与预测减少量之间的一致性。
   • $\rho^h_k$ 衡量相对于历史参考值 $f(\theta_{r(k)})$ 和累积历史项 $\sigma^h_k$ 的一致性。
     该机制允许接受暂时增加目标函数（$f(\theta_k + s_k) > f(\theta_k)$）的步长，只要它们与长期下降趋势一致，从而减少对有效粗空间方向的拒绝。
3. 校正与平滑：如果试验步被拒绝，则计算校正步 $c_k$，它是梯度下降步与原始提议的凸组合，避免了已计算梯度信息的浪费。最终的 NTR 更新细化迭代点并调整信赖域半径 $\Delta_k$。

主要贡献

• 非单调 APTS（NAPTS）：主要贡献是将 NTR 原理集成到 APTS 框架中。这具体解决了在噪声和非凸 NN 训练景观中粗空间步长被激进拒绝的问题。
• 受控的非单调性：本文引入了一种窗口化接受准则，允许受控的目标函数增加。这防止了对步长的“不加区分”的接受，同时仍允许算法坚持那些可能暂时未改善但全局有益的粗空间方向。
• 高效并行性：该方法保持了加性 Schwarz 结构，允许在局部迭代期间进行并行子域训练而无需子域间通信，仅需在梯度缓存和步长聚合时进行全局同步。

结果
实验在 CIFAR-10 数据集上进行，使用卷积神经网络（120 万参数）在四块 Nvidia A100 GPU 上训练。

• 性能：NAPTS 实现了每轮115.88 秒的训练时间，与标准 APTS（169.14 秒/轮）相比，CPU 时间减少了 30%，并且比激进接受变体 APTS$_A$（152.39 秒/轮）有进一步改进。
• 拒绝率：该方法显著减少了被拒绝的步长数量。NAPTS 记录了约1,921 次总拒绝迭代，而 APTS 为 13,972 次，APTS$_A$ 为 6,648 次。这表明非单调机制成功保留了原本会被丢弃的有效粗空间提议。
• 准确率：虽然基线 TR 和 NTR 方法难以达到较高的验证准确率，但预条件变体（APTS 和 NAPTS）实现了显著更高的验证准确率和更低的训练损失，表明在不妨碍泛化能力的情况下优化得到了改善。

意义
本文声称，NAPTS 提供了一种稳健的全局化策略，特别适用于大规模神经网络训练的非凸和噪声特性。通过放宽严格的单调性要求，该方法避免了因反复拒绝和重新计算搜索方向而产生的计算成本。作者认为，这种方法有效地利用域分解进行模型并行，同时通过信赖域提供自动步长调整，减少了与一阶方法相关的手动超参数调优需求。结果表明，NAPTS 在保持预条件方法准确率的同时，显著提高了计算效率。