AI・機械学習
上級

クリティカルバッチサイズ(クリティカルバッチサイズ)

バッチサイズを拡大しても学習速度(収束に必要な計算量)の改善が頭打ちになる閾値。McCandlish et al.(2018)が理論化し、勾配のシグナル対ノイズ比(SNR)から導出される。クリティカルバッチサイズ以下では計算量がほぼ線形にスケールし、それを超えると追加のバッチ拡大は計算資源の浪費となるため、LLM の分散学習設計における最重要指標の一つである。

0 回閲覧
0 いいね

クリティカルバッチサイズとは

クリティカルバッチサイズ(Critical Batch Size, B_crit)は、確率的勾配降下法(SGD)系のオプティマイザにおいて、バッチサイズの拡大による学習効率の向上が飽和する転換点を定量化した概念である。McCandlish et al. が 2018 年の論文「An Empirical Model of Large-Batch Training」で提唱し、OpenAI のスケーリング則(Scaling Laws)研究の基盤となった。

直感的に説明すると、小さなバッチサイズでは勾配のノイズが大きく、各パラメータ更新の方向が不安定である。バッチサイズを増やすとノイズが減り、更新の質が向上する。しかし、ある大きさ(B_crit)を超えると、すでに勾配の方向が十分に正確であり、さらにサンプルを追加しても情報利得がほぼゼロとなる。この閾値がクリティカルバッチサイズである。

理論的導出

クリティカルバッチサイズは勾配のノイズスケール(B_noise)から導出される。損失関数 L に対する勾配のノイズスケールは以下で定義される:

B_noise = tr(Σ) / |G|²

ここで Σ は個々のサンプル勾配の共分散行列、G は真の勾配(全データの勾配)、tr は行列のトレース(対角和)である。

B_noise の物理的意味は「勾配の真のシグナルに対するノイズの相対的な大きさ」を表す。バッチサイズ B でのパラメータ更新を考えると:

  • B ≪ B_noise の場合:勾配がノイズ支配的。バッチサイズを 2 倍にすると、同じ目標損失に到達するのに必要なステップ数がほぼ半分になり、計算量はほぼ一定。「ステップ効率的」な領域
  • B ≫ B_noise の場合:勾配がシグナル支配的。バッチサイズを 2 倍にしてもステップ数はほとんど減らず、計算量が 2 倍になる。「計算浪費」の領域
  • B ≈ B_noise(= B_crit)の場合:ステップ効率と計算効率のバランスが最適

クリティカルバッチサイズは学習の進行に伴い変化する。一般に学習初期は損失関数の勾配が大きく方向が明確なため B_crit が大きく、学習が進み損失が減少するにつれて B_crit も小さくなる傾向がある。

LLM におけるクリティカルバッチサイズの測定

LLM のクリティカルバッチサイズは以下の手順で測定できる:

  1. 異なるバッチサイズ(例: 2^10, 2^12, 2^14, 2^16, 2^18 トークン)で同一モデルを学習
  2. 各バッチサイズで目標損失に到達するまでのステップ数 S(B) と処理トークン数 D(B) = B × S(B) を記録
  3. S(B) と D(B) を B の関数としてプロットし、以下のモデルにフィッティング:

S(B) = S_min × (1 + B_crit / B) D(B) = D_min × (1 + B / B_crit)

ここで S_min は理想的な(ノイズゼロの)ステップ数、D_min は最小限必要なトークン数。

主要 LLM のクリティカルバッチサイズ推定

モデルパラメータ数推定 B_crit(トークン)実運用バッチサイズB/B_crit 比
GPT-3 Small125M約 60K500K〜8×
GPT-3175B約 500K〜2M3.2M〜2〜6×
Chinchilla70B約 1M1.5M〜1.5×
Llama 270B約 1〜2M4M〜2〜4×
Llama 3405B約 2〜4M16M(最終)〜4〜8×

この表から、フロンティアモデルの多くが B_crit の 2〜8 倍のバッチサイズで運用されていることがわかる。これは計算効率の観点からは次善(suboptimal)だが、壁時計時間(学習完了までの実時間)を優先した設計である。B_crit ちょうどで運用すると計算効率は最適だが、GPU クラスタの稼働時間が伸び、H100 のレンタルコスト($3〜4/GPU-hour)が嵩むため、総コスト最適化は B_crit より大きなバッチサイズになることが多い。

クリティカルバッチサイズとスケーリング則

OpenAI の Kaplan et al.(2020)および Hoffmann et al.(2022, Chinchilla)のスケーリング則研究では、クリティカルバッチサイズが損失 L のべき乗則に従うことが示された:

B_crit ∝ L^(-α)

つまり、損失が小さくなる(モデルが賢くなる)ほどクリティカルバッチサイズが大きくなる。これは直感的に理解できる:モデルが賢くなるとデータから抽出すべきシグナルが微細になり、より多くのサンプルを見ないとシグナルを検出できないためである。

この関係は LLM の学習スケジュール設計に直接的な示唆を与える:

  • 学習初期:B_crit が小さいため、小バッチで多くのステップを踏むほうが効率的
  • 学習後期:B_crit が大きくなるため、大バッチで計算効率を活かせる
  • 結論:バッチサイズウォームアップ(学習中にバッチサイズを段階的に増加)は理論的に最適な戦略

コスト最適化への応用

B_crit を用いた学習コストの最適化手順:

  1. 小規模プロキシ実験:1B〜7B パラメータのプロキシモデルで複数バッチサイズの学習を実施し、B_crit を測定
  2. スケーリング外挿:パラメータ数と B_crit の関係をべき乗則でフィッティングし、目標モデルサイズの B_crit を推定
  3. 壁時計 vs 計算のトレードオフ曲線作成:B_crit の 1〜8 倍の範囲でバッチサイズを変えた場合の学習時間と GPU コストを試算
  4. 総コスト最小化:GPU レンタル単価 × 学習時間の総コストが最小となるバッチサイズを選択

Meta が Llama 3 405B の学習で実施したとされるプロキシ実験でも、このアプローチが採用されている。

クリティカルバッチサイズの動的変化

B_crit は学習中に変動するため、固定バッチサイズは常に次善策である。理想的には B_crit をリアルタイムで推定し、バッチサイズを動的に調整する「適応的バッチサイズ」が最適だが、実装の複雑さから研究段階にとどまっている。

近似的なアプローチとして、勾配のノイズスケールをオンラインで推定する手法がある:

  • Smith et al.(2018)の Simple Noise Scale:異なるバッチサイズで計算した勾配の分散比から B_noise を推定
  • Grad-Norm Ratio:勾配のノルムの分散と平均の比から近似
  • AdaScale SGD:勾配の統計量を元にステップサイズとバッチサイズを同時に適応調整

FAQ

Q: クリティカルバッチサイズは学習率の選択にも影響しますか?

A: はい、密接に関連する。B < B_crit の領域では学習率のリニアスケーリング(バッチサイズ k 倍 → 学習率 k 倍)が有効に機能する。B > B_crit の領域ではリニアスケーリングの仮定が崩れ、学習率を上げすぎると学習が発散する。平方根スケーリング(学習率 √k 倍)や LAMB のような適応型オプティマイザが必要になるのは、B_crit を超える領域で運用しているからである。B_crit を知ることで、リニアスケーリングが安全な範囲を見積もれる。

Q: モデルサイズが大きくなるとクリティカルバッチサイズはどう変わりますか?

A: 一般にモデルサイズが大きくなるほど B_crit も大きくなる傾向がある。これはパラメータ数が多いモデルほど勾配空間の次元が高く、個々のサンプルからの勾配方向のばらつきが大きくなるためである。経験的には B_crit ∝ N^(0.5〜0.7)(N はパラメータ数)程度のスケーリングが観測されており、175B モデルの B_crit は 1B モデルの 10〜30 倍程度と推定される。

Q: ファインチューニングにもクリティカルバッチサイズの概念は適用できますか?

A: 適用可能であり、かつ重要性が高い。ファインチューニングのデータセットは事前学習より桁違いに小さいため、バッチサイズがデータセットの有意な割合を占めることがある。たとえば 10,000 サンプルの SFT データセットでバッチサイズ 512 なら、1 エポック = 約 20 ステップしかない。このような状況では B_crit が比較的小さく(数百〜数千サンプル)、それを超えるバッチサイズは汎化性能を悪化させる傾向がある。SFT では小バッチ(4〜32)が推奨されることが多いのは、このためである。