勾配チェックポイントは、Deep learningで使用される技術であり、ニューラルネットワークのトレーニングにおけるメモリ使用量を削減することを目的としています。標準的なバックプロパゲーションでは、ネットワークは順伝播で計算されたすべての中間活性化を保存し、逆伝播で勾配を計算する必要があります。非常に深いモデル、例えば大規模言語モデルやトランスフォーマーでは、この保存量が利用可能なハードウェアのメモリ容量を超えることがあります。勾配チェックポイントは、すべての活性化を保存するのではなく、一部のみを保持し、逆伝播中に破棄されたものを必要に応じて再計算することで、この問題に対処します。これにより、計算コストの増加と引き換えにメモリ使用量を大幅に削減し、同じハードウェア上でより大きなモデルのトレーニングや、より大きなバッチサイズの使用を可能にします。
この技術は、2016年にCarnegie Mellon UniversityとOpenAIの研究者によって、「Training Deep Nets with Sublinear Memory Cost」という論文で紹介されました。著者であるTianqi Chen、Bing Xu、Chiyuan Zhang、Carlos Guestrinらは、特定のチェックポイント（例えば、数層ごと）にのみ活性化を保存し、残りを再計算することで、n層のネットワークのトレーニングにおけるメモリコストをO(n)からO(sqrt(n))に削減できることを実証しました。その代償として、おおよそ1回の追加の順伝播が必要です。この基礎的な研究は、モデルサイズが劇的に増大する中で、機械学習コミュニティにおける標準的なツールとなっています。
標準的なバックプロパゲーションのメモリ使用
従来のトレーニングループでは、順伝播がネットワークの各層の活性化を計算します。これらの活性化は、逆伝播が連鎖律を介して勾配を計算するために必要であるため、メモリに保存されます。L層のネットワークでは、Lセットの活性化を保存する必要があり、それぞれが大きくなる可能性があります。例えば、Residual Network (ResNet)で数百層、または数十のアテンションブロックを持つトランスフォーマーでは、単一のトレーニング例に対して数ギガバイトの活性化データが蓄積されることがあります。大きなバッチサイズでトレーニングする場合、メモリ要件はバッチサイズに比例して増加し、しばしば主要なボトルネックとなります。
チェックポイント戦略
勾配チェックポイントは、ネットワークをセグメントに分割し、各セグメントの境界にチェックポイントを設定します。順伝播中、チェックポイント境界の活性化のみがメモリに保存されます。セグメント内の他のすべての中間活性化は破棄されます。逆伝播がセグメントに到達すると、保存されたチェックポイント活性化を使用してそのセグメントの順伝播を再計算し、勾配計算に必要な中間活性化を再生成します。この再計算は、トレーニングステップごとにおおよそ1回の追加の順伝播に相当する計算オーバーヘッドを追加しますが、ピークメモリ使用量を劇的に削減します。
チェックポイントの配置の選択はトレードオフです。チェックポイントが多いほど再計算は少なくなりますが、メモリ使用量は増加します。チェックポイントが少ないほどメモリは低くなりますが、計算は増加します。n層のネットワークに対する最適なチェックポイント数はおおよそsqrt(n)であり、メモリと計算のバランスを取ります。実際には、PyTorchやTensorFlowなどのフレームワークは、ユーザーがチェックポイント間隔を指定したり、自動ヒューリスティックを使用したりすることを可能にします。
バリアントと改善
元の技術にはいくつかの改良が開発されています。一般的なバリアントの1つは選択的チェックポイントであり、特定の層タイプ（アテンションブロックや畳み込み層など）のみをチェックポイントし、他は通常どおり保存します。別のアプローチであるメモリ効率的な勾配チェックポイントは、より洗練されたスケジュールを使用して、複数の粒度レベルで活性化を保存し、追加の再計算と引き換えにメモリをさらに削減します。一部のフレームワークは「オフロード」も実装しており、チェックポイントをCPUメモリやディスクに移動しますが、これにはデータ転送のオーバーヘッドが伴います。
Transformer (architecture)モデルの文脈では、勾配チェックポイントはMixed Precision TrainingやGradient Clippingなどの他のメモリ節約技術と組み合わせられることがよくあります。例えば、1750億のパラメータを持つGPT-3のようなモデルのトレーニングは、そのような最適化なしでは不可能です。この技術は、大規模モデルのファインチューニングでも使用され、メモリ節約により、クラスターではなく単一のGPUで実行できるようになります。
実践的な実装
現代の深層学習フレームワークでは、勾配チェックポイントは通常、シンプルなAPIとして公開されています。例えば、PyTorchでは、torch.utils.checkpointモジュールが、モジュールや一連の操作をラップするcheckpoint関数を提供します。ラップされたモジュールが実行されると、その活性化は保存されず、逆伝播中に再計算されます。TensorFlowはtf.recompute_gradを通じて同様の機能を提供します。これらの実装は簿記を自動的に処理するため、研究者はモデルアーキテクチャを変更せずに技術を簡単に採用できます。
勾配チェックポイントの計算オーバーヘッドは無視できません。sqrt(n)のチェックポイントを持つネットワークでは、トレーニング中の総順伝播計算は標準的なトレーニングと比較して約30〜40%増加します。しかし、このコストは、バッチサイズやモデルサイズを減らすことが収束やモデル品質に悪影響を与える可能性があるため、しばしば許容されます。多くの場合、より大きなバッチサイズを使用することによる高速化は、再計算オーバーヘッドを上回ります。
大規模モデルトレーニングへの影響
勾配チェックポイントは、非常に大規模なモデルのトレーニングの基盤となっています。OpenAI、Anthropic、Google DeepMindなどの企業は、数千億のパラメータを持つモデルのトレーニングに依存しています。例えば、8つのGPUを持つ単一ノードで700億パラメータのモデルをトレーニングする場合、チェックポイントなしでは、それらのGPUの合計メモリを超える活性化を保存する必要があります。勾配チェックポイントを使用することで、これらの組織はトレーニングジョブを利用可能なハードウェアに適合させることができますが、トレーニング時間は長くなります。
この技術は、ドキュメント要約やコード生成など、長いシーケンスを伴うGenerative AIアプリケーションにも不可欠です。これらの場合、活性化メモリはシーケンス長に応じて増加し、チェックポイントにより、メモリ制限を超えずに長いコンテキストを可能にします。これにより、10万トークン以上のコンテキストウィンドウを持つモデルの開発が直接可能になりました。
他のメモリ最適化との関係
勾配チェックポイントは、他の技術と一緒に使用されることがよくあります。Batch NormalizationやLayer Normalizationは直接メモリを削減しませんが、トレーニングの安定性を向上させることができ、チェックポイントを補完します。Model Pruningはパラメータ数を削減しますが、活性化は依然としてボトルネックであるため、チェックポイントは依然として必要です。Data Augmentationは実効データセットサイズを増加させますが、活性化メモリには影響しません。分散トレーニングでは、勾配チェックポイントをパイプライン並列処理と組み合わせることができ、異なる層を異なるデバイスに割り当てることで、デバイスごとのメモリ圧力をさらに削減できます。
注目すべき代替案の1つはGradient Accumulationであり、複数の小さなバッチで勾配を蓄積することで、より大きなバッチサイズをシミュレートします。これによりオプティマイザ状態のメモリは削減されますが、活性化メモリは削減されないため、チェックポイントの代替にはなりません。もう1つの関連するアイデアは、一部のResidual Network (ResNet)バリアントで使用される可逆層であり、出力から活性化を再構築できますが、これにはアーキテクチャの変更が必要であり、チェックポイントほど一般的ではありません。
制限とトレードオフ
勾配チェックポイントの主な制限は、トレーニングステップごとのウォールクロック時間の増加です。すでに計算に制約されているモデルでは、追加の順伝播によりトレーニングが20〜40%遅くなる可能性があります。さらに、この技術はモデルパラメータやオプティマイザ状態のメモリを削減しません。これらも大規模モデルではかなりの量になる可能性があります。非常に大規模なモデルでは、パラメータのオフロードとチェックポイントを組み合わせるか、AWS TrainiumやGroqなどのオンチップメモリが大きい専用ハードウェアを使用する必要があるかもしれません。
もう1つの微妙な問題は、再計算が数値的な違いを導入する可能性があることですが、実際には通常無視できる程度です。この技術は、同じ活性化を複数回再計算しないように注意深く実装する必要があり、そうしないとオーバーヘッドが増加します。これらの課題にもかかわらず、勾配チェックポイントは広く使用され、信頼性の高い方法であり続けています。
将来の方向性
モデルが成長し続けるにつれて、研究者はより効率的なチェックポイント戦略を模索しています。最近の研究では、モデルの構造とハードウェアのメモリプロファイルに基づいて、どの活性化を保存するかを決定する学習済みヒューリスティックを使用しています。他の研究では、層ではなく個々の操作レベルでのチェックポイントを調査しており、より細かい制御を可能にしています。また、チェックポイントを低精度トレーニングと組み合わせてメモリをさらに削減することへの関心もありますが、これには精度のトレードオフが伴います。
この技術は、メモリが深層学習の制約である限り、関連性を維持する可能性があります。大規模言語モデルやマルチヘッドアテンションアーキテクチャの台頭により、メモリ効率的なトレーニング手法への需要はさらに高まるでしょう。勾配チェックポイントは、他の最適化とともに、Artificial intelligenceの進歩の重要な推進力であり続けるでしょう。
結論
勾配チェックポイントは、シンプルでありながら強力なアイデアです。活性化を選択的に保存し、バックプロパゲーション中に再計算することで、深いネットワークをはるかに少ないメモリでトレーニングできるようにします。2016年に導入されて以来、この技術は分野の標準的な実践となり、利用可能なハードウェアでは不可能だったモデルのトレーニングを可能にしました。計算オーバーヘッドを追加しますが、特に大規模モデルでは、そのトレードオフはしばしば価値があります。分野がさらに大きなアーキテクチャに向かうにつれて、勾配チェックポイントは機械学習ツールボックスにおける不可欠なツールであり続けるでしょう。