[論文レビュー] POET: Training Neural Networks on Tiny Devices with Integrated Rematerialization and Paging

Shishir G. Patil, Paras Jain+3 more|arXiv (Cornell University)|Jul 15, 2022

Advanced Neural Network Applications被引用数 6

ひとこと要約

POETは、メモリ制限のあるエッジデバイスにおける大規模ニューラルネットワークのエネルギー最適トレーニングを可能にする、混合整数線形計画法（MILP）に基づくフレームワークを導入する。リマテリアライゼーションとページングを統合的に最適化することで、単独手法に比べ最大40％のエネルギー消費削減を達成し、Cortex-MクラスのデバイスでResNet-18やBERTのトレーニングを32 KBのRAMのみで実行可能にし、エネルギー効率とスケーラビリティにおいて既存のエッジトレーニング手法を上回る性能を発揮する。

ABSTRACT

Fine-tuning models on edge devices like mobile phones would enable privacy-preserving personalization over sensitive data. However, edge training has historically been limited to relatively small models with simple architectures because training is both memory and energy intensive. We present POET, an algorithm to enable training large neural networks on memory-scarce battery-operated edge devices. POET jointly optimizes the integrated search search spaces of rematerialization and paging, two algorithms to reduce the memory consumption of backpropagation. Given a memory budget and a run-time constraint, we formulate a mixed-integer linear program (MILP) for energy-optimal training. Our approach enables training significantly larger models on embedded devices while reducing energy consumption while not modifying mathematical correctness of backpropagation. We demonstrate that it is possible to fine-tune both ResNet-18 and BERT within the memory constraints of a Cortex-M class embedded device while outperforming current edge training methods in energy efficiency. POET is an open-source project available at https://github.com/ShishirPatil/poet

研究の動機と目的

• モバイル端末やマイコンチップを含む、メモリが限られたバッテリー駆動のエッジデバイス上で、大規模で現代的なニューラルネットワークのトレーニングを可能にすること。
• エッジトレーニングにおけるバックプロパゲーションの高いメモリおよびエネルギー要件に対処すること。現状では、ロジスティック回帰のような小規模なアーキテクチャに限定されている。
• リマテリアライゼーションとページングの2つの活性化メモリ削減技術を、統一的なエネルギー最小化フレームワーク内で同時に最適化すること。
• 実世界のエッジ展開に一般的なハード制約（例：アイドル時や充電時におけるトレーニング）を満たすこと。
• 既存のディープラーニングモデルにアーキテクチャの変更を加えずに、数学的に正しく、値を保持するソリューションを提供すること。

提案手法

• エッジトレーニングを、メモリ制約と実行時間制約の下でリマテリアライゼーションとページングの意思決定を統合的に最適化する混合整数線形計画問題（MILP）として定式化する。
• ターゲットハードウェア上でモデルオペレータをプロファイリングし、エネルギー、メモリ、実行時間の細粒度コストモデルを構築する。
• MILPソルバーを用いて、メモリ制約と時間予算を満たしつつ、合計エネルギー消費量を最小化するエネルギー最適スケジュールを特定する。
• ページング操作のパイプライン処理をリマテリアライゼーションと統合し、レイテンシを隠蔽し、厳しい実行時間制約を満たす。
• オペレータレベルのコストプロファイルに一般化することで、線形モデルおよび非線形モデル（ResNet-18 や BERT を含む）をサポートする。
• 安価な演算（例：ReLU）のリマテリアライゼーションと高コスト演算（例：畳み込み）のページングを統合した一貫した探索空間を採用し、エネルギーオーバーヘッドを低減する。

実験結果

リサーチクエスチョン

• RQ1 リマテリアライゼーションとページングを統合的に最適化することで、単独で使用する場合よりもエッジデバイスにおけるエネルギー消費を低減できるか？
• RQ2 32 KBのRAMでの制限下でも、BERT や ResNet-18 といった複雑なモデルを、厳密な実行時間制約を満たしながらエッジデバイスでトレーニング可能か？
• RQ3 リマテリアライゼーションとページングを統合することで、ヒューリスティック的または逐次的アプローチと比較して、エネルギー効率とメモリ使用量にどのような影響を与えるか？
• RQ4 MILPに基づく定式化が、現実的なハードウェア制約下で、証明可能な最適トレーニングスケジュールを特定できるか？
• RQ5 メモリおよび時間予算を変化させた場合、エネルギー効率と実行時間性能のトレードオフはどのように変化するか？

主な発見

• A72プラットフォームで、ResNet-18をトレーニングする際、タイトなメモリ予算下でPOETは、完全メモリ使用のベースラインに比べ最大141％のエネルギーオーバーヘッド低減を達成した。
• POETは、完全メモリトレーニングと比較して1％未満のエネルギーオーバーヘッドに抑えられたが、ヒューリスティックベースの手法であるCapuchinは、同様の条件下で73％～141％のエネルギー消費が高くなった。
• Jetson TX2では、線形モデルに限定された先行手法POFOと比較し、POETはピークメモリ消費量を8.3％削減し、スループットを13％向上させた。
• POETの統合的ソリューションは、リマテリアライゼーションまたはページングを単独で使用する場合に比べ、最大40％のエネルギー消費削減を達成した。特に、0.5 ms/エポックという厳しい実行時間制約下で顕著であった。
• MILPソルバーは、Bert や ResNet-18 といった複雑なモデルに対しても、商用ハードウェア上で10分未満で最適解を発見できた。
• POETは、32 KBのRAMでの制限下でも、Cortex-Mクラスのデバイス上でResNet-18およびBERTのトレーニングに成功し、プライバシー保護型のオンデバイスパーソナライゼーションの実現可能性を示した。