勾配蓄積は、ニューラルネットワークやその他の機械学習モデル、特に深層学習の訓練において、メモリ制約によりすべてのサンプルを一度に処理できない場合に、大きなバッチサイズの効果を近似するために用いられる手法である。大きなバッチ全体に対する勾配を単一の順伝播および逆伝播で計算する代わりに、オプティマイザは複数の小さなミニバッチにわたって勾配を蓄積し、蓄積された勾配が所望の実効バッチサイズに達した後にのみ重み更新を実行する。このアプローチにより、実践者は、GPUやTPUなどの単一デバイスのメモリ容量を超えるようなバッチサイズでもモデルを訓練できるようになる。
この手法は、1950年代にRobbins–Monroアルゴリズムを通じて導入された反復最適化アルゴリズムである確率的勾配降下法（SGD）と密接に関連している。SGDでは、データセット全体に対して計算される目的関数の真の勾配を、ランダムに選択されたデータの部分集合に対する勾配で近似する。一般的な妥協点はミニバッチアプローチであり、各ステップで小さなサンプル集合に対して勾配が計算される。勾配蓄積は、モデルパラメータを更新する前に複数のミニバッチにわたって勾配を平均化することでこの考え方を拡張し、バッチサイズをメモリフットプリントから実質的に切り離す。
動機とメモリ制約
大規模言語モデルやトランスフォーマーなどの大規模モデルの訓練には、アクティベーション、勾配、オプティマイザ状態を格納するための相当なメモリが必要である。バッチサイズはメモリ消費に直接影響する。バッチサイズが大きいほど、中間計算に多くのメモリが必要となる。コンシューマ向けGPUやエッジデバイスなどメモリが限られたハードウェアでは、実現可能な最大バッチサイズは小さくなる可能性がある。しかし、バッチサイズが小さいと、勾配推定がノイズを含み、訓練ダイナミクスが不安定になることがある。勾配蓄積は、順伝播と逆伝播に小さなバッチを使用しながら、複数のステップにわたって勾配を蓄積してより大きな実効バッチサイズを達成することで解決策を提供し、ピークメモリ使用量を増やさずに勾配の安定性を向上させる。
勾配蓄積の仕組み
標準的なミニバッチ訓練では、モデルはサイズ\(b\)のミニバッチを処理し、損失を計算し、逆伝播を行って勾配を取得し、SGDやAdamなどのオプティマイザを使用して直ちに重みを更新する。勾配蓄積では、このプロセスが変更される。モデルは\(k\)個のミニバッチを順次処理し、各逆伝播後に勾配を蓄積する（通常は加算または平均化）。\(k\)ステップ後、蓄積された勾配を使用してモデル重みを更新し、勾配アキュムレータはゼロにリセットされる。実効バッチサイズは\(k \times b\)である。このアプローチは、損失関数がサンプルに対する和または平均である場合、バッチ正規化層が注意深く処理されることを条件に、バッチサイズ\(k \times b\)での訓練と数学的に等価である。
利点と使用例
勾配蓄積は、大きなバッチサイズが有益であるがメモリ制限のために実現不可能なシナリオで広く使用されている。例えば、生成モデルの訓練や大規模言語モデルのファインチューニングでは、訓練を安定させ収束を改善するために、数百または数千のバッチサイズが必要となることが多い。勾配を蓄積することで、研究者は単一のGPUやCPUベースのシステムなどメモリが限られたハードウェア上でこれらの大きなバッチをシミュレートできる。さらに、勾配蓄積により、毎ステップで勾配を同期することなく複数のデバイスにわたって訓練が可能になり、分散訓練における通信オーバーヘッドを削減できる。これは、異なるメモリ容量のGPUインスタンスを提供するAmazon Web Services、Microsoft Azure、Google Cloudなどのクラウドベースのプラットフォームに特に関連する。
バッチサイズと学習率との関係
勾配蓄積を使用する場合、実効バッチサイズが増加するため、学習率の調整が必要になることが多い。実際には、大規模バッチ訓練に関する先行研究で示唆されているように、学習率をバッチサイズに比例して線形にスケーリングするのが一般的なヒューリスティックである。ただし、最適なスケーリングはモデルアーキテクチャ、オプティマイザ、データセットに依存する。勾配蓄積は、重み更新あたりに参照されるサンプル数を変更しない。勾配の集約方法のみを変更する。したがって、実効バッチサイズが増加すると総訓練ステップ数は減少し、収束速度と最終的な性能に影響を与える可能性がある。実践者は、真の大規模バッチ訓練と同等の結果を達成するために、学習率と蓄積ステップ数を調整する必要がある。
実装上の考慮事項
勾配蓄積の実装には、訓練ループの変更が含まれる。各順伝播と逆伝播の後、勾配はモデルの勾配バッファに蓄積される（例えば、PyTorchのaccumulate_gradやTensorFlowのGradientTapeを永続変数とともに使用）。オプティマイザステップは、所望の蓄積ステップ数に達した後にのみ実行される。損失を適切にスケーリングすることが重要である。損失がミニバッチに対して平均化されている場合、蓄積された勾配は蓄積ステップ数で除算して正しい勾配の大きさを維持する必要がある。あるいは、損失が合計されている場合、スケーリングは不要である。さらに、バッチ正規化層は現在のミニバッチに対して統計量を計算するため、特別な注意が必要である。勾配蓄積では、正規化のための実効バッチサイズは依然としてミニバッチサイズであり、真の大規模バッチ訓練とは異なる挙動につながる可能性がある。一部のフレームワークでは、デバイス間の同期バッチ正規化による自動処理を提供している。
他の手法との比較
勾配蓄積は、逆伝播中にアクティベーションを再計算することでメモリを削減する勾配チェックポイントや、低精度演算を使用してメモリを削減する混合精度訓練としばしば比較される。これらの手法は組み合わせることができる。例えば、勾配蓄積と混合精度を併用することで、さらに大きな実効バッチサイズでの訓練が可能になる。複数のGPUにわたる分散訓練とは異なり、勾配蓄積はデバイス間の通信を必要としないため、実装が簡単で同期オーバーヘッドが少ない。ただし、順次ステップ数が増加するため、ミニバッチサイズが非常に小さく、複数の逆伝播のオーバーヘッドが大きい場合、訓練が遅くなる可能性がある。
制限と代替案
勾配蓄積は、真の大規模バッチ訓練の完全な代替ではない。勾配推定は異なるミニバッチから計算されるため、データシャッフルやバッチ正規化統計量によるわずかな違いが生じる可能性がある。場合によっては、これが異なる収束挙動につながることがある。代替案には、CerebrasシステムやGroqアクセラレータなど、高いメモリ帯域幅と容量を提供するより大きなメモリデバイスの使用や、モデル並列性とパイプライン並列性を使用してバッチをデバイス間で分散することが含まれる。さらに、LAMBやLARSなどの一部のオプティマイザは、大きなバッチサイズを直接処理するように設計されており、蓄積の必要性を減らす可能性がある。
歴史的背景と採用
複数のミニバッチにわたって勾配を蓄積する概念は、現代の深層学習より前から存在し、1980年代の「バンチモード逆伝播アルゴリズム」にルーツがある。これは、2010年代に深層学習モデルが成長し、メモリ制約がボトルネックになるにつれて重要性を増した。今日、勾配蓄積はPyTorch、TensorFlow、JAXなどの主要な深層学習フレームワークの標準機能であり、OpenAI、Anthropic、Google DeepMindなどの組織での最先端モデルの訓練に広く使用されている。また、研究および本番環境、特に限られたハードウェアでの大規模モデルのファインチューニングにおいて一般的な手法である。
結論
勾配蓄積は、メモリ制約下で大規模モデルを訓練するための実用的で広く採用されている手法である。複数のミニバッチにわたって勾配を蓄積することで、比例したメモリ増加なしに大きなバッチサイズのシミュレーションを可能にする。いくつかのオーバーヘッドを導入し、学習率と正規化層の注意深い調整を必要とするが、特に大規模な訓練とファインチューニングタスクにおいて、深層学習実践者のツールキットの不可欠なツールであり続けている。