[論文レビュー] Revisiting Locally Supervised Learning: an Alternative to End-to-end Training

Yulin Wang, Zanlin Ni+3 more|arXiv (Cornell University)|Jan 26, 2021

Advanced Neural Network Applications参考文献 65被引用数 24

ひとこと要約

この論文では、局所的教師あり深層学習のための新しい損失関数であるInfoProを提案する。この手法は、情報理論的で非グリーディな目的関数によって、初期層におけるタスク関連情報の崩壊を防ぐ。InfoPro損失の計算可能な上界を再構成項と対照的損失項の組み合わせとして導出し、エンドツーエンド学習と同等の性能を達成しつつ、GPUメモリ使用量を40%未満に抑える。また、非同期かつ並列学習を可能にする。

ABSTRACT

Due to the need to store the intermediate activations for back-propagation, end-to-end (E2E) training of deep networks usually suffers from high GPUs memory footprint. This paper aims to address this problem by revisiting the locally supervised learning, where a network is split into gradient-isolated modules and trained with local supervision. We experimentally show that simply training local modules with E2E loss tends to collapse task-relevant information at early layers, and hence hurts the performance of the full model. To avoid this issue, we propose an information propagation (InfoPro) loss, which encourages local modules to preserve as much useful information as possible, while progressively discard task-irrelevant information. As InfoPro loss is difficult to compute in its original form, we derive a feasible upper bound as a surrogate optimization objective, yielding a simple but effective algorithm. In fact, we show that the proposed method boils down to minimizing the combination of a reconstruction loss and a normal cross-entropy/contrastive term. Extensive empirical results on five datasets (i.e., CIFAR, SVHN, STL-10, ImageNet and Cityscapes) validate that InfoPro is capable of achieving competitive performance with less than 40% memory footprint compared to E2E training, while allowing using training data with higher-resolution or larger batch sizes under the same GPU memory constraint. Our method also enables training local modules asynchronously for potential training acceleration. Code is available at: https://github.com/blackfeather-wang/InfoPro-Pytorch.

研究の動機と目的

• 中間活性化を保存する必要があるため、深層ネットワークにおけるエンドツーエンド（E2E）学習の高いGPUメモリ使用量を緩和すること。
• 初期層で有用な入力情報が崩壊するグリーディで短視眼的な学習によって引き起こされる局所的教師あり学習の性能低下を克服すること。
• 局所モジュール間でタスク関連情報の保持を図りつつ、不要な（不快な）成分を除外する学習目的を開発すること。
• モジュール間の勾配更新を分離することで、モデル並列化と非同期学習を可能にすること。
• モデルの精度を損なわず、E2E学習の実用的でメモリ効率の良い代替手段を提供すること。

提案手法

• 入力と中間特徴間の相互情報量を最大化し、不要な変数への依存度を最小化する情報理論的原則に基づいた情報伝達（InfoPro）損失を提案する。
• 計算可能となるようにInfoPro損失の計算可能な上界を導出し、再構成損失と対照的損失の組み合わせによる近似目的関数を得る。
• 局所モジュールを近似InfoPro損失で学習させ、中間特徴表現が後続層に有用な入力情報を保持するように促進する。
• ネットワーク全体にわたる逆誤差伝搬を必要としない局所的目的関数に最適化を分解し、勾配分離学習を可能にする。
• 後続層からの誤差信号依存性を排除することで、局所モジュールの非同期的かつ並列学習を可能にする。
• PyTorchに実装し、CIFAR、SVHN、STL-10、ImageNet、Cityscapesを含む多様なベンチマークで検証する。

実験結果

リサーチクエスチョン

• RQ1 なぜ深層ネットワークにおけるグリーディな局所的教師あり学習は、エンドツーエンド学習と比較して性能劣化を引き起こすのか？
• RQ2 情報理論的損失関数は、局所的教師あり学習における初期層での情報崩壊を緩和できるか？
• RQ3 実用的な深層学習学習において、非グリーディで情報保持型の目的関数を効果的に近似できるか？
• RQ4 InfoPro損失を用いた局所的教師あり学習では、精度を維持したままどの程度GPUメモリ使用量を削減できるか？
• RQ5 提案手法は、性能を損なわず、ネットワークモジュールの非同期的または並列学習をサポートできるか？

主な発見

• InfoPro学習は、標準的な局所的教師あり学習の主要な失敗モードである初期層におけるタスク関連情報の崩壊を防ぐ。
• CIFAR-10、SVHN、STL-10、ImageNet、Cityscapesにおいて、エンドツーエンド学習と同等の性能を達成し、GPUメモリ使用量は40%未満である。
• 同じメモリ制約下で、InfoProはE2E学習と比較して50%大きなバッチサイズ、または50%高い入力解像度を可能にする。
• 再構成項と対照的損失項の組み合わせによる近似InfoPro損失は、理論的な情報理論的目的関数を効果的に近似する。
• 局所モジュールの非同期学習をサポートし、モデル並列化による潜在的な高速化を実現する。
• 分類およびセマンティックセグメンテーションタスクを含む多様なビジョンベンチマークにおいて、本手法は頑健で効果的である。