Stanford 大学が 2023 年に提案した二次最適化アルゴリズム。Second-order Clipped Stochastic Optimization の略称で、Hessian(ヘッセ行列)の対角近似を用いてパラメータごとの曲率情報を効率的に活用する。AdamW と同等の最終性能に約半分のステップ数で到達できるとされ、LLM 訓練の計算コスト削減が期待されている。
関連する技術記事・ガイドを検索
Sophia(Second-order Clipped Stochastic Optimization)は、2023 年に Stanford 大学の Hong Liu らが論文「Sophia: A Scalable Stochastic Second-order Optimizer for Language Model Pre-training」で提案した最適化アルゴリズムである。名称はギリシャ語の「知恵」に由来する。
従来の Adam 系オプティマイザが勾配の 1 次情報(勾配値と勾配の 2 乗の移動平均)のみを使用するのに対し、Sophia は損失関数の 2 次情報(曲率、すなわち Hessian 行列の対角成分の近似)を活用する。曲率が大きい方向ではステップサイズを小さくし、曲率が小さい方向では大きくすることで、損失ランドスケープの形状に適応した効率的な更新を実現する。
ニュートン法では、パラメータ更新は以下の形をとる:
θ_{t+1} = θ_t − η · H^{−1} · g_t
ここで H は Hessian 行列(損失関数の 2 階偏微分行列)、g_t は勾配ベクトルである。Hessian の逆行列を使うことで、損失ランドスケープの曲率を考慮した最適なステップサイズが各方向に対して自動的に決まる。
しかし、LLM のような巨大モデルで Hessian 行列を完全に計算・逆行列化することは不可能である(70B パラメータでは 70B × 70B のフル Hessian が必要)。Sophia はこの問題を 2 つのアプローチで解決する:
ステップ 4 のクリッピングが Sophia の核心である。曲率が非常に小さい方向では m_t / h_t が爆発的に大きくなる可能性があるが、上限値 ρ でクリッピングすることで更新量を制限し安定性を確保する。
Sophia は 2 種類の Hessian 推定手法を提供する:
| 手法 | Sophia-G | Sophia-H |
|---|---|---|
| 推定対象 | Gauss-Newton 行列の対角 | Hessian の対角 |
| 計算方法 | ミニバッチの勾配外積 | Hutchinson 推定量 |
| 計算コスト | 勾配計算の約 1.0 倍 | 勾配計算の約 0.5 倍追加 |
| 推定頻度 | 毎ステップ | k ステップごと(k=10 推奨) |
| 精度 | 高い | やや低いがコスト効率的 |
| 推奨用途 | 小〜中規模モデル | 大規模 LLM |
Sophia-H は k ステップごとに Hessian を推定するため、追加の計算コストは実質的に 1/k に希釈される。k=10 の場合、1 ステップあたりの追加コストは約 5% に抑えられる。
論文で報告された GPT-2 規模のモデルでの比較結果:
| 指標 | AdamW | Sophia-H | 改善率 |
|---|---|---|---|
| 同じ perplexity 到達ステップ | 100K | 50K | 2x 高速 |
| 同じ perplexity 到達時間 | 100% | 約 50% | 2x 高速 |
| 1 ステップの計算コスト | 1.0x | 約 1.05x | 5% 増 |
| メモリ使用量 | m + v(2 状態) | m + h(2 状態) | 同等 |
| 最終 perplexity(同ステップ数) | ベースライン | −0.5〜1.0 改善 | — |
Sophia が同じ最終性能に約半分のステップ数で到達できるという結果は、計算コスト(GPU 時間 × 電力)の大幅な削減を意味する。
Sophia を実装する際の重要なハイパーパラメータ:
| パラメータ | 推奨値 | 説明 |
|---|---|---|
| η(学習率) | 3e-4〜1e-3 | AdamW より高めに設定可能 |
| β1 | 0.96 | AdamW (0.9) より大きめ |
| β2 | 0.99 | Hessian 推定の指数移動平均 |
| ρ(クリッピング閾値) | 0.01〜0.04 | 大きすぎると不安定、小さすぎると収束遅延 |
| k(Hessian 更新間隔) | 10 | 計算コストと精度のトレードオフ |
| ε | 1e-12 | ゼロ除算防止 |
クリッピング閾値 ρ の設定が最も重要で、モデルサイズやタスクに応じた調整が必要である。ρ が大きすぎると勾配爆発と同様の不安定さが生じ、小さすぎると AdamW と同程度の収束速度に退化する。
Sophia は理論的に優れた特性を持つが、実用化にはいくつかの課題がある:
これらの課題が解決されれば、LLM 訓練の計算コストを大幅に削減できる可能性があり、活発な研究が続いている。
ステップ数ベースでは約 2 倍だが、実時間では 1.5〜1.8 倍程度になる。理由は、Hessian 推定の追加計算(約 5%)と、Hessian 推定ステップでの通信オーバーヘッド(分散学習時)が存在するため。ただし、ステップ数の削減は GPU 電力コストの直接的な削減に繋がるため、大規模訓練では総コストの改善効果が大きい。
完全なドロップイン置き換えではない。学習率を高めに設定し、β1 を調整し、ρ(クリッピング閾値)を新たに設定する必要がある。また、Hessian 推定のための追加の backward pass が必要であり、訓練ループの実装にも変更が必要。PyTorch 用の参考実装は公開されているが、Hugging Face Trainer 等への統合はコミュニティ主導で進行中である。
従来の二次最適化手法(L-BFGS, K-FAC 等)は Hessian 行列の計算・保存コストが膨大で、LLM スケールでは実用的でなかった。K-FAC はブロック対角近似を用いるが、Transformer のアテンション層への適用が困難という問題があった。Sophia は対角近似 + 確率的推定 + クリッピングの組み合わせにより、これらの問題を実用的なレベルまで解決した点が画期的である。