[論文レビュー] Training Deep Nets with Sublinear Memory Cost

Tianqi Chen, Bing Xu+2 more|arXiv (Cornell University)|Apr 21, 2016

Advanced Neural Network Applications参考文献 19被引用数 538

ひとこと要約

本論文は計算グラフに基づく手法を導入し、サブリニアメモリで深層ネットワークを訓練する方法を提案する。n層ネットワークに対して、追加のフォワード・パスと引換えにメモリを節約し、O(sqrt(n))のメモリを実現する。極端なケースではO(log n)にも達する。非常に深いResNetと長いシーケンスのLSTMで実証。

ABSTRACT

We propose a systematic approach to reduce the memory consumption of deep neural network training. Specifically, we design an algorithm that costs O(sqrt(n)) memory to train a n layer network, with only the computational cost of an extra forward pass per mini-batch. As many of the state-of-the-art models hit the upper bound of the GPU memory, our algorithm allows deeper and more complex models to be explored, and helps advance the innovations in deep learning research. We focus on reducing the memory cost to store the intermediate feature maps and gradients during training. Computation graph analysis is used for automatic in-place operation and memory sharing optimizations. We show that it is possible to trade computation for memory - giving a more memory efficient training algorithm with a little extra computation cost. In the extreme case, our analysis also shows that the memory consumption can be reduced to O(log n) with as little as O(n log n) extra cost for forward computation. Our experiments show that we can reduce the memory cost of a 1,000-layer deep residual network from 48G to 7G with only 30 percent additional running time cost on ImageNet problems. Similarly, significant memory cost reduction is observed in training complex recurrent neural networks on very long sequences.

研究の動機と目的

• 訓練時のメモリ使用量を中間の特徴マップと勾配を少なく格納することで削減する。
• インプレース演算とメモリ共有を用いた計算グラフベースのメモリ最適化を開発する。
• より深いネットワークの訓練を可能にするための、制御された計算とメモリのトレードオフを導入する。
• ディープラーニングフレームワークにメモリ最適化技術を組み込むための実践的ガイドラインを提供する。

提案手法

• 中間結果間でインプレース演算とメモリ共有を可能にするよう、計算グラフを分析する。
• バックプロパゲーション中にドロップした中間結果を再計算するための勾配グラフ構築法（mirror counts）を開発する。
• ネットワークをkセグメントに分割してセグメント出力を保持することでO(sqrt(n))メモリ計画を提案し、追加のフォワード・パスを発生させる。
• このアプローチを任意のグラフへ一般化し、バジェットベースの計画（Alg. 3）を備えたメモリ最適化勾配グラフ（Alg. 2）で実現する。
• 再帰的な細分化（k=1の場合）により、再帰ベースの視点でメモリをO(log n)へ削減できることを示す。
• 既存のDLフレームワークにこれらの手法を実装するためのフレームワーク指針とオープンソース計画を提供する。

実験結果

リサーチクエスチョン

• RQ1 訓練の正確性を維持しつつ、中間の特徴マップと勾配をサブリニアのメモリで格納できるか？
• RQ2 与えられたネットワークの深さとメモリ予算に対して、最適なメモリと計算のトレードオフは何か？
• RQ3 訓練中にインプレース演算とメモリ共有を可能にするよう、計算グラフをどう分析できるか？
• RQ4 このアプローチは、非常に深いアーキテクチャ（例：1000以上の層）や長いシーケンスモデル（例：LSTMs）に対して、過大なオーバーヘッドなしにスケールできるか？
• RQ5 現在のDLフレームワークにメモリ最適化を統合する実践的ガイドラインは何か？

主な発見

• メモリコストはネットワークの深さに対して線形からサブリニアへ削減でき、ミニバッチごとに追加のフォワードパスだけでn層ネットワークに対してO(sqrt(n))のメモリを達成する。
• 極端なケースでは、追加のフォワード計算量がO(n log n)程度で、メモリ使用量をO(log n)に削減できる。
• 提案手法によりImageNetで1000層のResNetのメモリ削減を48Gから7Gへ示す実験結果。
• 長いシーケンス上の複雑なRNN（LSTMs）の訓練で顕著なメモリ削減を観測。
• サブリニアメモリ計画は、線形メモリ割当と比較して約30%の遅延を伴うが、大量のメモリ節約を考慮すれば控えめなオーバーヘッド。
• このアプローチは既存のフレームワークと互換性があり、他のメモリ最適化と組み合わせ可能。