サーベイ: Training Deep Nets with Sublinear Memory Cost

Chen, Tianqi, et al. "Training deep nets with sublinear memory cost." arXiv preprint arXiv:1604.06174 (2016).

@article{chen2016training,
  title={Training deep nets with sublinear memory cost},
  author={Chen, Tianqi and Xu, Bing and Zhang, Chiyuan and Guestrin, Carlos},
  journal={arXiv preprint arXiv:1604.06174},
  year={2016}
}

paper: https://arxiv.org/pdf/1604.06174

slide: https://www.cs.cmu.edu/~zhihaoj2/15-849/slides/12-tensor-rematerialization.pdf:image=https://www.cs.cmu.edu/~zhihaoj2/15-849/slides/12-tensor-rematerialization.pdf

github: https://github.com/dmlc/mxnet-memonger

背景

深層学習モデルの学習にはパラメタ保存の何倍ものメモリが必要

どんなもの?

訓練時（forward/backward）に必要なメモリ用を減らす方法の提案

先行研究と比べてどこがすごい?

学習中における計算とメモリのトレードオフを考慮できる
（少なくともCNN, RNNを含む）一般的なモデルに適用できる

技術や手法のキモはどこ?

sec3: backward時に、ノード（演算）ごとに、メモリを確保するか既に不要なメモリを再利用（上書き）するかをシンプルなラベル付グラフアルゴリズムにより求める

（赤いところはメモリの共有ができる）

sec4: モデルを分割し、分割されたブロックの最初の層の活性状態量（隠れ層の出力）を保存する
- 保存するデータ量が減少
- backward時には再びforwardを行って活性化状態を復元する（追加の計算コストあり）

（シンプルでメモリ削減効果が高い）

どうやって有効だと検証した?

全ての活性状態量、特徴マップを保存する場合
sec3の方法を用いてメモリ使用量削減を試みた場合
sec3+4の方法（checkpint）で計算コストと引き換えにさらにメモリ削減を試みた場合で学習させてメモリ使用量を比較層数1000のLSTMだとcheckpointで約256GB → 2GB以下に削減できた、計算時間は30%増