勾配蓄積と実効バッチサイズの拡大

なぜ勾配蓄積が要るのか：バッチサイズはメモリで頭打ちになる

大規模モデルの学習では「大きなバッチで学習したい」という要求と「GPU メモリに載らない」という制約が正面衝突します。バッチサイズを上げると勾配推定のノイズが減り、学習率を上げられて収束が安定する場面が多い（この関係は バッチサイズと学習率の相互作用 で詳述）。一方、順伝播・逆伝播では活性（activation）をバッチ内の全サンプル分メモリに保持する必要があり、メモリ消費はバッチサイズにほぼ比例して増えます。

ここで効くのが 勾配蓄積（gradient accumulation） です。発想は単純で、載るだけの小さなマイクロバッチで勾配を計算し、それを複数回ぶんメモリ上で足し込んでから、まとめて1回だけ重みを更新する。こうすると、ピークメモリは1マイクロバッチ分のままで、更新に使う勾配は実質的に「マイクロバッチをすべて連結した大バッチ」のものと一致します。時間をメモリと交換する技法だと言えます。

原理：勾配は和をとれる、だから分割して足せる

勾配蓄積が成立する数学的な根拠は、ミニバッチ勾配がサンプル損失の和（または平均）の勾配であるという線形性です。バッチ全体の損失を平均で書くと、勾配は次のように各サンプルの勾配の平均になります。

L_batch = (1/N) * Σ_i loss(x_i)         # N サンプルの平均損失
∇L_batch = (1/N) * Σ_i ∇loss(x_i)       # 勾配も平均（微分は線形）

つまり N サンプルを K 個のマイクロバッチ（各サイズ N/K）に割っても、各マイクロバッチの勾配を足し合わせれば全体の勾配が復元できます。これがメモリ上で逐次的に実行できるのは、勾配バッファ（param.grad）への足し込みがインプレースで累積されるからです。逆伝播を呼ぶたびに勾配が上書きではなく加算されるフレームワークの挙動を、そのまま利用しています。

optimizer.zero_grad()                    # 勾配バッファを明示的にゼロ化
for k in range(K):                       # K = 蓄積回数（accumulation steps）
    micro = next_microbatch()            # 1マイクロバッチ（メモリに載るサイズ）
    loss = compute_loss(model, micro)
    loss = loss / K                      # ★ ここで蓄積回数で割るのが核心
    loss.backward()                      # param.grad に加算（累積される）
optimizer.step()                         # K回ぶん積んだ勾配で1回だけ更新
optimizer.zero_grad()                    # 次の実効バッチに備えてクリア

実効バッチサイズ（effective batch size）は マイクロバッチサイズ × K × データ並列数 で決まります。K を増やすだけでメモリを増やさずに実効バッチを2倍にも8倍にもでき、ハードウェアを変えずに大バッチ学習を再現できるのが最大の利点です。

等価になる部分・ならない部分

勾配蓄積は「大バッチと完全に同じ」になるわけではありません。何が一致し、何がずれるかを切り分けることが実務では決定的です。

更新に使う勾配そのものは数学的に一致します。一方でピークメモリは大きく下がる代わりに、更新1回に K 回の前後伝播が要るので実時間は約 K 倍かかります。これが「時間とメモリの交換」の意味です。問題になるのは、バッチ全体を見て初めて正しく計算される演算——その代表が BatchNorm です。

BatchNorm との相互作用：ここが最大の落とし穴

BatchNorm はミニバッチ内の平均と分散を使って活性を正規化します（詳細は BatchNorm の理論）。勾配蓄積では正規化は各マイクロバッチ単位で実行されるため、統計は N/K サンプルからしか計算されません。つまり BatchNorm から見える「バッチ」はマイクロバッチであって、実効バッチではない。

この結果、勾配蓄積は BatchNorm 入りモデルでは大バッチと等価になりません。マイクロバッチが小さいほど統計推定のノイズが大きくなり、正規化の挙動が変わります。さらに実行時の移動平均（running mean/var）も K 倍の頻度で更新されるため、推論時に使う統計の性質もずれます。

分散学習の勾配同期との関係：通信を「蓄積中は止める」

データ並列の分散学習（分散学習の並列化戦略）では、各 GPU が自分の勾配を計算し、backward() の直後に All-Reduce で全 GPU の勾配を平均同期します。ここに勾配蓄積を素朴に組み合わせると、致命的な無駄が生じます。

naive: 蓄積K回 × backward → 毎回 All-Reduce が走る（K回の通信）
       しかし step は1回。途中K-1回の同期は捨てられる

蓄積中の中間勾配はまだ更新に使わないのに、毎 backward() で重い通信が走ってしまう。正しくは、蓄積している K-1 回は同期を抑制し、最後の1回でだけ All-Reduce する。PyTorch DDP の no_sync() コンテキストはまさにこのための仕組みで、蓄積中は勾配を各 GPU でローカルに積むだけにし、最終マイクロバッチで一度だけ同期します。

for k in range(K):
    micro = next_microbatch()
    if k < K - 1:
        with model.no_sync():            # 同期を抑制してローカル蓄積のみ
            (compute_loss(model, micro) / K).backward()
    else:
        (compute_loss(model, micro) / K).backward()  # 最後だけAll-Reduce
optimizer.step()
optimizer.zero_grad()

これで通信回数が K 回から1回に減り、蓄積による計算時間増にさらに通信オーバーヘッドが乗るのを防げます。なお最終的な実効バッチは マイクロバッチ × K × GPU数 であり、勾配蓄積とデータ並列は乗算的に実効バッチを拡大します。

まとめ：時間でメモリを買い、等価性の境界を見極める

勾配蓄積は、勾配が和（平均）に分解できるという線形性を使い、マイクロバッチの勾配を K 回足し込んでから1回更新することで、ピークメモリを据え置いたまま実効バッチを自在に拡大する技法でした。骨格は「各損失を K で割って逆伝播し、K 回累積してから step」。更新に使う勾配は大バッチと数学的に一致しますが、BatchNorm のようにバッチ全体の統計に依存する演算は等価にならない——ここが等価性の境界です。さらに分散環境では no_sync で蓄積中の通信を止めるのが定石でした。実効バッチを大きくしたら学習率もそれに合わせて調整する必要があり、その指針は バッチサイズと学習率 と、メモリ削減の別アプローチである 混合精度学習 と合わせて読むと、大規模学習のメモリ・速度・安定性のトレードオフが立体的に見えてきます。