Gradient Accumulation

勾配蓄積は、単一のモデルウェイトの更新を実行する前に、マイクロバッチと呼ばれる複数の小さなバッチからの勾配を足し合わせるトレーニング技術です。利用可能なGPUメモリで一度にフルバッチを処理できない場合に、より大きなバッチサイズでのトレーニングを近似します。これにより、モデルのアーキテクチャを変更したり入力解像度を下げたりすることなく、メモリを大量に消費する勾配降下法を実用的なものにします。

勾配蓄積の仕組み#

通常のトレーニング中、各バッチは3つの主要なステップに従います。損失を計算し、バックプロパゲーションを使用して勾配を計算し、オプティマイザにモデルウェイトを更新させます。その後、次のバッチの前に勾配がクリアされます。

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

設定された蓄積ステップ数に達した後、以下の処理が行われます。

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の効率的な利用が低下する可能性があります。主な利点は、希望する有効バッチを制限されたメモリに収めることです。2つ目の利点は分散トレーニングで発生する可能性があり、DistributedDataParallel no_syncなどの機能を使用して蓄積ウィンドウの終了まで勾配の同期を回避できる場合があります。

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

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

関連するテクニックと主な違い#

勾配蓄積は、他のメモリまたは安定性のテクニックと混同されることがよくあります。

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

実社会での応用#

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

航空画像検出の場合、エンジニアリングチームは、デバイスごとに6つの画像と4つの蓄積ステップを使用して、4つのGPUでトレーニングを行う場合があります。有効なグローバルバッチは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リファレンスに dokument化されています。実際には、確実に収まる最大のマイクロバッチを選択し、意図した有効バッチを計算し、蓄積ステップ、学習率、またはデバイス数を変更するときに検証結果を比較してください。