Robust and Fast Training via Per-Sample Clipping

技術要約：パーサンプル・クリッピングによる堅牢かつ高速な学習

問題提起
本論文は、確率的最適化問題 $\min_{x \in \mathbb{R}^d} F(x)$、ここで $F(x) = \mathbb{E}_{\xi \sim \mathcal{D}}[f(x, \xi)]$ を扱う。確率的勾配降下法（SGD）は標準的なソルバーであるが、その収束保証は通常、勾配ノイズが有界な分散（軽裾分布）を持つという仮定に依存している。しかし、近年の経験的および理論的な研究によれば、多くの機械学習アプリケーションにおいて、勾配ノイズは重い裾（ヘビーテイル）を持ち、分散は欠如しているものの $p \in (1, 2]$ の範囲で有限の $p$ 次モーメントを持つ可能性があることが示されている。このような条件下では、標準的なSGDは性能が低下する場合がある。既存の堅牢な代替案である標準的な勾配クリッピング（Clip-SGD）や正規化SGD（Normalized SGD）は、最適なサンプル複雑度の下限を達成できないか、あるいは特定の（例えば初期値ギャップ $\Delta_1$ や失敗確率 $\delta$ の特定の対数因子に依存するような）非現実的なパラメータチューニングを必要とする。

手法
著者らは、バッチ勾配をまとめてからクリッピングを行うのではなく、ミニバッチ内の個々のサンプルに対してクリッピングを適用する堅牢な勾配推定器、**パーサンプル・クリップSGD（PS-Clip-SGD）**を提案する。

サイズ $n$ のバッチにおけるサンプルを $\xi^{(1)}_t, \dots, \xi^{(n)}_t$ とすると、勾配推定器 $G(x_t, \alpha, \beta)$ は次のように定義される：
$$G(x_t, \alpha, \beta) = \frac{1}{n} \sum_{k=1}^n \gamma^{(k)}_t \nabla f(x_t, \xi^{(k)}_t)$$
ここで、パーサンプル・クリッピング係数は $\gamma^{(k)}_t := \min\left\{1, \frac{\alpha k^{1/\beta}}{|\nabla f(x_t, \xi^{(k)}_t)|}\right\}$ である。パラメータ $\alpha$ と $\beta$ は、$\alpha$ はモーメント境界 $\sigma$ に、$\beta$ は裾指数 $p$ に関連している。更新ルールは標準的なSGDに従う：$x_{t+1} = x_t - \eta_t G(x_t, \alpha, \beta)$。

主な貢献

1. 期待値における最適収束： 著者らは、PS-Clip-SGDが重い裾を持つノイズが存在する非凸最適化問題において、最適なサンプル複雑度を達成することを証明した。具体的には、本手法は平均二乗勾配ノルムに対して $O(1/T)$ の収束率を達成し、これはZhangら(2020)によって確立された下限と一致する。これは、ステップサイズに対して初期値ギャップ $\Delta_1$ の知識を必要とせずに、二乗勾配ノルムの和（一部の先行研究で使用されているノルムの和よりも強い基準）に対して境界が成立するという点で、Clip-SGDやNormalized SGDの既存の結果を改善している。
2. 高確率収束： 本論文は、期待値における収束率に対し、失敗確率 $\delta$ およびステップ数 $T$ に関する多項式対数因子を除いて一致する、高確率収束保証を確立している。特筆すべきは、$\log(1/\delta)$ への依存性が最新のClip-SGDの境界と比較して改善されており、かつ、ステップサイズやクリッピング閾値を $\delta$ の関数として明示的にチューニングする必要がない点である。
3. 勾配累積に関する洞察： 著者らは、勾配累積のすべてのステップの後にのみ勾配クリッピングを適用するという一般的な慣習に異議を唱える。彼らは、累積中のミニバッチレベルでのクリッピング（Mini-Batch Clipping）を適用することが、中間的なクリッピングがバイアスを導入するという一般的なコンセンサスに反して、追加の計算コストをほとんど増やすことなく学習性能を向上させ得ることを提案し、経験的に検証した。

結果

• 理論的側面： 理論的分析により、PS-Clip-SGDは$L$-滑らかさと有限の $p$ 次モーメントの仮定の下で強力な収束保証を提供することが示された。導出された境界は、$\delta$ への依存性と収束の強さの尺度（二乗ノルム vs ノルム）の両面において、Normalized SGDおよびClip-SGDよりもタイトであることが示されている。
• 二次最適化： パレート分布のノイズを伴う二次関数の最小化実験において、PS-Clip-SGDは、チューニングされていないハイパーパラメータを使用した場合でも、様々なノイズ領域（$p \in \{1.2, 1.5, 1.8\}$）においてClip-SGDおよびNormalized SGDを大幅に上回った。高確率実験は、$\log(1/\delta)$ に対する劣線形な依存性を裏付けた。
• 画像分類（CIFAR-100上のAlexNet）： モメンタムを用いたAlexNetの学習において、PS-Clip-SGDは、バニラなSGD（モメンタム付き）の検証精度58.90%およびClip-SGD（モメンタム付き）の58.94%に対し、64.84%の高い検証精度を達成した。この向上は、パーサンプル勾配ノルムの計算によって導入される約1.3倍の計算オーバーヘッドを考慮しても維持された。著者らは、これを、標準的なバッチレベルのクリッピングでは検出できない、非常に大きなノルムを持つ少数の外れ値勾配を効果的に制御できたためであるとしている。
• 言語モデリング（GPT-2）： 124MパラメータのGPT-2モデルを訓練する勾配累積の設定において、ミニバッチ・クリッピング（累積ステップごとのクリッピング）は、中間的なクリッピングが有害であるという従来の視点にもかかわらず、複数の学習率において標準的なClip-SGD（全累積後のクリッピング）を上回った。

意義と主張
本論文は、差分プライバシーにおける確立された用途とは異なり、重い裾を持つノイズが存在する場合の学習性能向上のための手法として、パーサンプル・クリッピングを提案した最初の研究であると主張している。その意義は、以下の点にある：

1. 非凸のヘビーテイル問題に対して最適な理論的収束率を達成すること。
2. ハイパーパラメータの選択に対して堅牢であり、精緻化されていない設定でも良好に動作すること。
3. パーサンプル・クリッピングが、計算コストを上回る実用的な性能向上（AlexNetのようなディープラーニングタスクにおいて）をもたらすことを示すこと。
4. 勾配累積戦略におけるパラダイムシフトを示唆し、中間的なクリッピングが有害ではなく有益であることを示すこと。

著者らは限界についても認めており、最適なパラメータを定義するためにノイズのモーメント境界 $\sigma$ と裾指数 $p$ の知識が必要であること、および有限の $p$ 次モーメントの仮定が中心化された $p$ 次モーメントの有界性よりも強い仮定であることなどを挙げている。また、パーサンプル・クリッピングはAlexNetのような小規模なモデルには効果的であるが、非常に大規模なアーキテクチャでは計算オーバーヘッドにより実用的でない可能性があることにも触れており、累積シナリオにおいてはミニバッチ・クリッピングがコスト効率の高い代替案となることを述べている。