勾配チェックポインティングとは

Gradient Checkpointing / Activation Checkpointing

順伝播時の中間活性化の一部を保存せず、逆伝播時に再計算することでメモリ使用量を削減する学習手法

概要

勾配チェックポインティングは、ニューラルネットワークの学習時に必要なメモリ使用量を削減する手法。通常の逆伝播では、順伝播時に計算した各層の中間活性化をすべて保持しておき、逆伝播時の勾配計算に使う。層数が多い、あるいは1層あたりの活性化が大きいモデルでは、この保持だけで大量のメモリを消費する。 勾配チェックポインティングでは、一部の層(チェックポイント)の活性化のみを保持し、それ以外は逆伝播で必要になった時点で順伝播を部分的にやり直して再計算する。これによりメモリ使用量を削減できる一方、再計算の分だけ計算量と学習時間は増加する。 2016年にChenらが発表した論文「Training Deep Nets with Sublinear Memory Cost」で提案された手法。PyTorchやTensorFlowなど主要フレームワークで標準的にサポートされる。

利点

• 学習時のメモリ使用量を削減でき、より大きなバッチサイズやモデルを扱える
• モデル構造を変更せずに既存フレームワークの機能として適用できる

欠点

• 再計算が発生する分、学習にかかる計算時間が増加する(一般に20〜30%程度のオーバーヘッド)

関連用語

参考文献

• Research PaperTraining Deep Nets with Sublinear Memory Cost