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

技術概要：ニューラルネットワーク学習のための非単調前処理トラスト領域法

問題定義
大規模な深層ニューラルネットワーク（NN）の学習には、ハイパーパラメータの調整、マルチデバイスパイプラインの活用、および非凸でノイズの多い目的関数の処理に関する課題が存在する。第一階次法である確率的勾配降下法（SGD）や Adam は標準的であるが、これらは手動によるステップサイズとモーメンタムのスケジューリングに大きく依存しており、ネットワークをデバイス間で逐次ブロックに分割するモデル並列性を自然に活用することはできない。一方、ドメイン分解（DD）法はパラメータを部分領域に分割して並列更新を行う手段を提供するが、非凸設定での収束を確保するためには、堅牢な大域化戦略が必要となる。加法的に前処理されたトラスト領域戦略（APTS）などの従来手法はトラスト領域（TR）メカニズムを利用するが、粗空間ステップの拒絶率が高く、探索方向の再計算に莫大なコストを要する傾向がある。

手法：NAPTS
本論文は、非単調トラスト領域（NTR）の原理を統合した APTS フレームワークの拡張である「非単調加法的に前処理されたトラスト領域戦略（NAPTS）」を提案する。本手法は、各反復において以下の 3 つの明確なフェーズで動作する。

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 学習環境における、過激な粗空間ステップの拒絶という課題に対処する。
• 制御された非単調性: 本論文は、制御された目的関数の増加を許容するウィンドウ付き受諾基準を導入する。これにより、ステップの「無差別な」受容を防ぎつつ、一時的には改善をもたらさないが全体的には有益な粗空間方向に対してアルゴリズムが継続することを可能にする。
• 効率的な並列性: 本手法は加法的シュワルツ構造を維持しており、局所反復中の部分間通信なしに並列部分領域学習を可能にする。グローバルな同期は、勾配のキャッシュとステップの集約時のみで必要となる。

結果
実験は、4 台の Nvidia A100 GPU で学習された畳み込み NN（120 万パラメータ）を用いて CIFAR-10 データセットで行われた。

• 性能: 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 が前処理法の精度を維持しながら、計算効率を大幅に向上させることを示唆している。