Small Batch Size Training for Language Models: When Vanilla SGD Works, and Why Gradient Accumulation Is Wasteful

基本情報

• arXiv ID: 2507.07101v1 (https://arxiv.org/abs/2507.07101)
• 著者: Martin Marek, Sanae Lotfi, Aditya Somasundaram, Andrew Gordon Wilson, Micah Goldblum
• 所属: New York University, Columbia University
• 投稿日: 2025年07月11日
• カテゴリ: cs.LG, cs.CL

簡単に説明すると

言語モデルの訓練では「大きなバッチサイズが安定性のために必要」という常識がありますが、本研究はこの常識を覆します。
バッチサイズ1という極端に小さな値でも、Adamのハイパーパラメータを適切にスケーリングすれば、安定した訓練が可能であることを実証しました。
重要な発見は、β₂パラメータを固定せずに「半減期」（トークン数で測定）を固定することで、小さなバッチサイズが大きなバッチサイズと同等またはそれ以上の性能を達成できるというものです。
さらに、小さなバッチサイズでは、モーメンタムを使わないバニラSGDでさえも競争力のある性能を示すことが分かりました。
コード: https://github.com/martin-marek/batch-size

1. 研究概要

1.1 背景と動機

大きなバッチサイズは言語モデル訓練の安定性を向上させると広く信じられており、その結果、大きなバッチサイズで良好に動作する洗練されたオプティマイザが標準的に使用されています。
小さなバッチサイズでの事前訓練実験では、損失スパイクや重い不安定性を観察することがあります。
学習率は小さなバッチサイズに対して減少させることが一般的ですが、Adamのβパラメータなどの他のハイパーパラメータは、バッチサイズ間で固定されることが多いです。
本研究は、β₂を固定する代わりに、β₂の「半減期」（トークン数で測定）を固定することで、バッチサイズ1に至るまで安定した訓練が可能であることを示します。

1.2 主要な貢献

本研究は、言語モデル訓練における小さなバッチサイズの優位性を実証し、実践的な推奨事項を提供します。
主な貢献として、次のような点があります。

• 小さなバッチサイズ（バッチサイズ1を含む）が安定して訓練でき、大きなバッチサイズと同等以上の性能を達成することの実証
• Adamハイパーパラメータの新しいスケーリング規則（β₂* = β₂^(B*/B)）の提案
• 小さなバッチサイズがハイパーパラメータ選択に対してより堅牢であることの発見
• バニラSGD（モーメンタムなし）が小さなバッチサイズで競争力のある性能を示すことの実証
• グラデーション累積を避け、スループットを最大化する最小バッチサイズを使用するという実践的推奨

2. 提案手法

2.1 手法の概要

本研究の核心は、Adamオプティマイザのβ₂パラメータをバッチサイズに応じて適切にスケーリングする新しい方法です。
従来の手法では、β₂値を固定していましたが、本研究では「半減期」の概念を導入します。
半減期とは、ある勾配の寄与が半分に減衰するまでに処理されるトークン数を指します。
この半減期をバッチサイズ間で固定することで、小さなバッチサイズでも安定した訓練が可能になります。

2.2 技術的詳細

Adamオプティマイザの第一および第二モーメントは、過去のミニバッチ勾配の指数移動平均です。
各ステップで、過去のミニバッチ勾配の寄与はβ係数によって減衰します。
n回のステップ後、任意のミニバッチ勾配の寄与は半分になります：β^n = 1/2。
トークン数で表される半減期t₁/₂は、β^(t₁/₂/(B×T)) = 1/2として定義されます（Bはバッチサイズ、Tはシーケンス長）。

提案されるスケーリング規則は以下の通りです：
β₂* = β₂^(B*/B)

ここで、B*は新しいバッチサイズ、Bは元のバッチサイズです。

2.3 新規性

既存手法との主な違いは、以下の点です。
第一に、β₂パラメータを固定せずに、トークン数で測定される半減期を固定するという新しいアプローチを採用しています。
第二に、この単純な変更により、小さなバッチサイズが大きなバッチサイズと同等以上の性能を達成できることを示しています。
第三に、小さなバッチサイズでは、複雑なオプティマイザが不要になり、バニラSGDでさえも競争力のある性能を示すことを発見しています。

3. 実験結果

3.1 実験設定

30Mパラメータモデルから1.3Bパラメータ（GPT-3スケール）までの様々なサイズのモデルで実験を実施しました。
FineWeb-EduおよびFineWebデータセットを使用し、バッチサイズ{1, 4, 16, 64, 256, 1024, 4096}で訓練しました。
SGD、Adam、Adafactor、Muonオプティマイザを比較し、各設定で慎重にハイパーパラメータを調整しました。

3.2 主要な結果

すべてのオプティマイザが最小のバッチサイズで優れた性能を達成しました。
バッチサイズが増加するにつれて、性能が低下するだけでなく、オプティマイザ間の差も拡大しました。
小さなバッチサイズでは、バニラSGD（モーメンタムなし）でさえも競争力のある性能を示しました。
GPT-3（1.3B）スケールでは、バッチサイズ1のバニラSGDがAdamWベースラインと同等の性能を達成しました。

3.3 既存手法との比較

提案されたβ₂スケーリング規則を使用して、Xiao et al. (2024)の実験を再現しました。
追加の調整なしで、小さなバッチサイズ（16、32）が大きなバッチサイズ（128、256）と同等の性能を達成しました。
ハイパーパラメータの堅牢性の実験では、バッチサイズ1がテストされたハイパーパラメータの全範囲でほぼ最適な損失を達成しました。
Gemma 3（4B）モデルのファインチューニングでは、小さなバッチサイズのAdafactorがLoRAを上回る性能を示しました。

4. 実用性評価

4.1 実装の容易性

提案された方法は非常にシンプルで、既存のコードベースに容易に統合できます。
必要な変更は、β₂パラメータをβ₂* = β₂^(B*/B)に従ってスケーリングすることだけです。
小さなバッチサイズはハイパーパラメータに対してより堅牢であるため、調整の労力が約90%削減されます。

4.2 計算効率

小さなバッチサイズは、FLOPあたりの性能で大きなバッチサイズと同等またはそれ以上の結果を達成します。
バニラSGDを使用する場合、オプティマイザ状態を保存する必要がないため、メモリ使用量が最大50%削減されます。
グラデーション累積が不要になるため、実装が簡素化され、メモリ使用量が削減されます。

4.3 応用可能性

本手法は、事前訓練とファインチューニングの両方で有効であることが実証されています。
メモリ制約のある環境では、SGDやAdafactorと小さなバッチサイズの組み合わせが特に有用です。
GPT-3（13B）規模では、Adamとグラデーション累積の使用が消費者向けGPUでは実行不可能ですが、SGDやAdafactorなら実行可能です。

5. まとめと所感

5.1 論文の意義

本研究は、言語モデル訓練における長年の常識を覆す重要な発見を提示しています。
小さなバッチサイズの優位性を実証することで、より効率的でアクセシブルな訓練方法への道を開いています。
特に、メモリ制約のある環境でのLLM訓練をより多くの研究者が利用できるようにします。
また、複雑なオプティマイザが必ずしも必要ではないという発見は、訓練パイプラインの簡素化につながります。

5.2 今後の展望

著者らは、今後の研究課題として以下を挙げています。
バッチサイズスケジュールとの相互作用の理解。
小さなバッチサイズ領域向けのより効果的な小状態オプティマイザの設計。
第二モーメントの半減期を固定することがなぜ汎化するのかの理論的理解。
これらの研究により、言語モデル訓練のさらなる効率化と民主化が期待されます。