Gradient Accumulation

勾配蓄積は、マイクロバッチと呼ばれる複数の小さなバッチの勾配を加算してから、1回のモデル重み更新を実行するトレーニング手法です。利用可能なGPUメモリにフルバッチを一度に格納できない場合に、より大きなバッチサイズでのトレーニングを近似します。これにより、モデルアーキテクチャを変更したり入力解像度を下げたりせずに、メモリを大量に消費する勾配降下法を実用化できます。

勾配蓄積の仕組み#

通常のトレーニングでは、各バッチは3つの主なステップに従います。損失を計算し、逆伝播を使用して勾配を計算し、オプティマイザーでモデルの重みを更新します。その後、次のバッチの前に勾配をクリアします。

勾配蓄積では、オプティマイザーの更新を遅延させます。各マイクロバッチで順伝播と逆伝播を実行しますが、その勾配はメモリに保持され、後続のマイクロバッチの勾配に加算されます。この動作は、各パラメーターの勾配フィールドに値を蓄積するPyTorchのbackward操作によって自然に実現されます。

設定した蓄積ステップ数に達すると、次の処理が行われます。

1. オプティマイザーがモデルの重みを更新します。
2. Optimizer.zero_gradなどの操作を使用して勾配をリセットします。
3. 新しい蓄積ウィンドウが開始されます。

実効バッチサイズは次のとおりです。

microbatch size x accumulation steps x number of training devices

たとえば、4枚の画像のマイクロバッチを8ステップにわたって蓄積すると、1つのGPU上で32枚の画像のバッチとおおよそ同等に動作します。4つのGPUを使用する場合、グローバル実効バッチサイズは128になります。

メモリ、速度、学習への影響#

勾配蓄積では、一度にネットワークを通過するのが1つのマイクロバッチだけになるため、アクティベーションに必要なピークメモリを削減できます。ただし、モデルパラメーター、勾配、オプティマイザー状態が使用するメモリは、大幅には削減されません。

また、通常はトレーニング速度も向上しません。8つのマイクロバッチを処理するには、依然として8回の順伝播と逆伝播が必要であり、バッチが小さいとGPUの使用効率が低下する場合があります。主な利点は、限られたメモリに目的の実効バッチを収められることです。副次的な利点として、分散トレーニングでは、DistributedDataParallel no_syncなどの機能を使用して、蓄積ウィンドウの終了まで勾配の同期を回避できる実装があります。

手動実装では通常、各マイクロバッチの損失を蓄積ステップ数で割り、得られる勾配が合計ではなく平均を表すようにします。高レベルのトレーナーは正規化を自動的に処理する場合があるため、追加のスケーリングを適用すると更新が不正確になる可能性があります。

蓄積によって真の大規模バッチを近似できますが、結果が常に同一になるとは限りません。バッチ正規化の統計は個々のマイクロバッチから計算されます。また、確率的な操作や浮動小数点演算の順序によっても差異が生じる可能性があります。そのため、学習率は自動的にスケーリングするのではなく、検証が必要になる場合があります。

関連する手法と主な違い#

勾配蓄積は、メモリや安定性に関する他の手法と混同されることがよくあります。

• 勾配チェックポイントは、逆伝播中に選択した順伝播の操作を再計算することで、アクティベーションメモリを削減します。PyTorchのアクティベーションチェックポイントに関するドキュメントでは、この計算とメモリのトレードオフについて説明しています。一方、蓄積では大きなバッチを小さなマイクロバッチに分割します。
• **混合精度**は、選択した操作を低精度のデータ型で保存または計算します。メモリを削減してスループットを向上させられる一方、蓄積はオプティマイザーが更新する頻度を変更します。正しい自動混合精度での蓄積ワークフローに従うことで、これらの手法を組み合わせられます。
• 勾配クリッピングは、異常に大きな更新による不安定性を低減するため、勾配の大きさを制限します。これによって実効バッチが大きくなることはありません。蓄積と組み合わせる場合、通常は完全に勾配を蓄積した後に、clip_grad_normなどの操作を使用してクリッピングを実行します。

実世界での利用例#

高解像度の医用画像セグメンテーションでは、12 GBのGPUに一度に格納できる大規模スキャンは2枚だけの場合があります。8つのマイクロバッチにわたって勾配を蓄積すると、小さな解剖学的構造の識別に必要な解像度を維持しながら、実効バッチを16にできます。トレードオフとして、オプティマイザーの更新間隔が長くなります。

航空画像の物体検出では、エンジニアリングチームが4つのGPUを使用し、デバイスあたり6枚の画像と4つの蓄積ステップでトレーニングする場合があります。グローバル実効バッチは96枚の画像です。これにより、各デバイスをメモリ制限内に保ちながら、小さな車両や建物が多数含まれるシーンの更新を安定させられます。チームは、Ultralytics Platformのクラウドトレーニングを通じて、クラウド実験、データセット、トレーニング実行を管理できます。

Ultralytics YOLOでの勾配蓄積#

Ultralyticsのモデルのトレーニングワークフローでは、物理的なbatch設定と公称nbs設定から蓄積動作を導出します。batch=4とnbs=64を使用すると、トレーナーは複数のマイクロバッチにわたってオプティマイザーの更新を遅延させ、公称バッチサイズ64を目標にします。

from ultralytics import YOLO

# Load a pretrained YOLO26 detection model
model = YOLO("yolo26n.pt")

# Use a small physical batch and a larger nominal batch
results = model.train(
    data="coco8.yaml",
    epochs=10,
    batch=4,
    nbs=64,
)

高レベルのトレーナーは、逆伝播、蓄積、オプティマイザーのステップ、勾配のクリアを管理します。その動作については、Ultralytics BaseTrainerリファレンスに記載されています。実際には、確実に収まる最大のマイクロバッチを選択し、意図する実効バッチを計算したうえで、蓄積ステップ数、学習率、デバイス数を変更する際に検証結果を比較してください。