Second-Order, First-Class：曲率を考慮した学習のための合成可能なスタック
arXiv cs.LG / 2026/3/30
📰 ニュースDeveloper Stack & InfrastructureIdeas & Deep AnalysisModels & Research
要点
- 本論文は、既存の手法が実装が難しく、チューニングが脆く、合成可能なAPIを欠いているため、二次（second-order）の曲率を考慮した学習手法が十分に活用されていないと主張する。
- Optaxネイティブの「composable stack（合成可能なスタック）」であるSomaxを提案し、静的な実行計画に基づいて駆動される1つのJITコンパイル済みステップとして曲率を考慮した学習をパッケージ化する。
- Somaxは、曲率演算子、推定器、線形ソルバ、前処理（preconditioner）、減衰（damping）ポリシーについて、最上位クラスの（first-class）スワップ可能なモジュールを提供しつつ、モメンタム、ウェイトデケイ、学習率スケジュールのような標準的な勾配変換を通じてOptaxとの統合を維持する。
- 計画（planning）と実行（execution）を分離することで、Somaxは中間結果を再利用でき、冗長な再計算を行う「未計画の合成」と比べてステップあたりのオーバーヘッドを削減する。
- 論文中のアブレーション結果では、合成に関する意思決定がスケーリング挙動と精度到達までの時間（time-to-accuracy）に大きく影響し、さらに計画メカニズムが効率を改善することが示されている。
