アクティベーション再計算（アクティベーションサイケイサン）
ニューラルネットワークの逆伝播時に、メモリから読み出す代わりに中間活性値を順伝播から再計算する手法。勾配チェックポインティングの中核メカニズムであり、メモリ使用量と計算コストのトレードオフを実現する。
アクティベーション再計算とは
アクティベーション再計算（Activation Recomputation）は、ディープラーニングの訓練における逆伝播フェーズで、保存されていない中間活性値（Activation）をその場で再計算する処理を指す。勾配チェックポインティングの中核をなすメカニズムであり、「Rematerialization」とも呼ばれる。
大規模モデルの訓練では、各層の活性値がGPUメモリの大半を占有する。例えばLlama 3 8Bモデルのシーケンス長2,048・バッチサイズ4での訓練では、活性値だけで約20GBのGPUメモリを消費する。パラメータ（16GB）と勾配（16GB）、オプティマイザ状態（32GB）を含めると84GB以上が必要で、A100 80GBでも収まらない。アクティベーション再計算により活性値メモリを60〜80%削減することで、限られたGPUリソースでの訓練が可能になる。
再計算の粒度
Full Recomputation
全ての中間活性値を破棄し、逆伝播時に全てを再計算する。メモリ削減効果は最大だが、順伝播が完全に2回実行されるため計算コストが2倍になる。実用的ではないため、通常は使用されない。
Selective Recomputation
計算コストが低い操作（ReLU、LayerNorm、Dropout等）の活性値のみを破棄し、コストが高い操作（行列乗算、Attention計算）の活性値は保持する。Megatron-LMで採用されている方式で、メモリ削減と計算コスト増のバランスが最も優れている。
再計算コストの目安：
| 操作 | 計算コスト | メモリ占有 | 再計算推奨 |
|---|---|---|---|
| ReLU/GELU/SiLU | 極小 | 中 | 推奨（コスト微小） |
| LayerNorm/RMSNorm | 小 | 中 | 推奨 |
| Dropout | 極小 | 中 | 推奨（マスク再生成にシード保持必要） |
| 線形変換（MatMul） | 大 | 大 | 非推奨（計算コスト高） |
| Attention Score | 大 | 大 | 非推奨（FlashAttentionで代替） |
Block-level Recomputation
Transformerブロック単位で活性値を保持/破棄する。最もよく使われる粒度で、PyTorchの torch.utils.checkpoint はこの方式を採用している。各ブロックの入力のみを保持し、ブロック内部の全中間値は逆伝播時に再計算する。
FlashAttentionとの関係
FlashAttention（Tri Dao, 2022）は、Attention計算のメモリ使用量をO(N²)からO(N)に削減するアルゴリズムである。従来のAttention計算ではN×Nのスコア行列を中間活性値として保持する必要があったが、FlashAttentionはタイリングとオンライン正規化によりスコア行列を保持せずにAttention出力を計算する。
これは実質的にAttention部分の「アクティベーション再計算」を効率化した手法と言える。FlashAttentionとブロックレベルのチェックポインティングを組み合わせることで、Attention以外の活性値も削減でき、総合的なメモリ効率が最大化される。
フレームワーク別の実装
PyTorch
torch.utils.checkpoint.checkpoint と checkpoint_sequential が基本API。カスタムの再計算ポリシーは CheckpointPolicy で指定可能。PyTorch 2.0以降では torch.compile との互換性も改善されている。
Megatron-LM
NVIDIAのMegatron-LMでは、Selective Recomputationがデフォルトで有効。--recompute-activations フラグで全再計算、--recompute-granularity selective で選択的再計算を指定する。テンソル並列・パイプライン並列との組み合わせで最大効率を発揮。
DeepSpeed
Activation Checkpointing APIとして実装。deepspeed.checkpointing.checkpoint 関数でブロック単位の再計算を指定。CPU Offloading（活性値をCPUメモリに退避）との組み合わせも可能。
FAQ
Q1: Dropoutの再計算時にランダム性は保たれますか？
PyTorchの実装では、チェックポイント境界でRNGの状態を保存し、再計算時に復元する。これによりDropoutマスクが順伝播と逆伝播で同一になることが保証される。ユーザーが明示的にRNG管理する必要はない。
Q2: アクティベーション再計算はGPU利用率に影響しますか？
再計算処理はGPUの計算ユニットを使用するため、逆伝播中のGPU利用率は向上する。一方でメモリ帯域のボトルネックが緩和されるため、全体の計算効率（TFLOPS/GPU）は維持または改善される場合もある。実測では5〜15%の訓練時間増加が典型的。
Q3: 推論時にアクティベーション再計算は適用されますか？
推論は順伝播のみなので逆伝播用の活性値保持が不要であり、再計算も発生しない。推論時のメモリ最適化は量子化（INT8/INT4）やKVキャッシュ管理が主な手法となる。