SoftServe:深層学習向けに設計された正定値を保つスケーラブルな準ニュートン法
この論文は「SoftServe」と呼ぶ新しい準ニュートン法の一族を紹介します。準ニュートン法は二次情報(曲率)を使って最適化を速める古典的手法です。しかし深層学習では目的関数が非凸であることと、モデルのパラメータ数が非常に大きいことが障害でした。本研究はこれらの問題を避けつつ、バッチ全体の損失を使う線検索や場当たり的な曲率修正を要さない手法を提案します。SoftServeは「Soft Secant Equation Regularized with Variational Entropy」の略で、変分的な目的から曲率の正定値な推定を導きます。負の曲率が存在する場合でも、推定は常に正定値になります。
研究者たちはまず、従来の準ニュートン法で使われる「セカント方程式」を緩める考えを採用しました。セカント方程式はパラメータの差と勾配の差から曲率を推定する式です。しかし非凸な場合にはこの式を厳密に満たすと逆に推定行列が非正定(降下方向にならない)になり得ます。SoftServeはこの等式を最低二乗のペナルティに変えた変分目的を最小化します。ペナルティ強さを表すハイパーパラメータでトレードオフを調整でき、有限の値では常に正定値解が得られると説明されています。
スケール面では、密行列を扱うと計算量やメモリが膨大になります。そこで著者らは対角近似(-Diag)とクロネッカー積構造の近似(-Kron)という構造制約を導入しました。対角版は各座標ごとの正のスケールを保持するだけで、計算・記憶どちらもパラメータ数に線形に比例します。クロネッカー版はパラメータをブロックに分けて、各ブロックごとに行列構造を保存し、ShampooやKFACと同程度のコストに抑えます。また必要な行列演算には、固有値分解のような高コストの分解を避け、GPUで高速に動く行列乗算を多用する安定したNewton–Schulz反復を用いています。
著者らはひどく条件の悪い問題での有用性を示しています。具体的には再帰型ネットワーク、深いオートエンコーダ、物理情報を取り入れたニューラルネットワーク(physics-informed neural networks)、そして1億を超えるパラメータ(136M)の物理情報拡散モデルなどで評価し、AdamやMuon、SOAPといった確立された最適化手法と比べてしばしば低い損失を達成したと報告しています。SoftServeは構成上常に正定値を保つため、線検索や全バッチ評価を必要としない点も利点として挙げられています。実装は公開リポジトリで入手可能とされています。
重要な注意点もあります。提案法はペナルティの重みや構造化の選択といったハイパーパラメータに依存します。クロネッカー版は既存の構造化二次法と同様に計算負荷や実装の複雑さが残る場合があります。また、論文の抜粋は結果の詳細や一般化性能の限界を完全には示していません。著者らの報告では多くのタスクで良好な結果が出ていますが、すべての問題で常に優れるとは明言されておらず、深層学習全般への適用可能性は今後の検証が必要です。