[論文レビュー] N2N Learning: Network to Network Compression via Policy Gradient Reinforcement Learning

Anubhav Ashok, Nicholas Rhinehart+2 more|arXiv (Cornell University)|Sep 18, 2017

Advanced Neural Network Applications参考文献 35被引用数 116

ひとこと要約

この論文は、教師ネットワークを高精度の軽量な学生ネットワークへ自動圧縮するための、2段階の強化学習手法（層の削除→層の縮小）を提案し、ポリシー勾配と知識蒸留を利用します。

ABSTRACT

While bigger and deeper neural network architectures continue to advance the state-of-the-art for many computer vision tasks, real-world adoption of these networks is impeded by hardware and speed constraints. Conventional model compression methods attempt to address this problem by modifying the architecture manually or using pre-defined heuristics. Since the space of all reduced architectures is very large, modifying the architecture of a deep neural network in this way is a difficult task. In this paper, we tackle this issue by introducing a principled method for learning reduced network architectures in a data-driven way using reinforcement learning. Our approach takes a larger `teacher' network as input and outputs a compressed `student' network derived from the `teacher' network. In the first stage of our method, a recurrent policy network aggressively removes layers from the large `teacher' model. In the second stage, another recurrent policy network carefully reduces the size of each remaining layer. The resulting network is then evaluated to obtain a reward -- a score based on the accuracy and compression of the network. Our approach uses this reward signal with policy gradients to train the policies to find a locally optimal student network. Our experiments show that we can achieve compression rates of more than 10x for models such as ResNet-34 while maintaining similar performance to the input `teacher' network. We also present a valuable transfer learning result which shows that policies which are pre-trained on smaller `teacher' networks can be used to rapidly speed up training on larger `teacher' networks.

研究の動機と目的

• ハードウェア制約を満たすための自動的でデータ駆動のネットワーク圧縮を動機付ける。
• 教師ネットワークから派生したコンパクトなアーキテクチャを探索する principledなRLフレームワークを開発する。
• アーキテクチャ空間を効率的に探索するための二段階のアクションスキーム（層削除に続く層縮小）を提示する。
• 圧縮された学生モデルの訓練に知識蒸留を組み込む。
• 複数データセットで圧縮の有効性と学習した方策の転移性を示す。

提案手法

• 教師-学生圧縮をネットワークアーキテクチャ上のマルコフ決定過程として定式化する。
• 二段階のポリシーネットワークを使用：層削除ポリシー（バイナリで保持/削除）と層縮小ポリシー（層パラメータの連続的減衰）。
• 報酬R = Rc × Ra で圧縮と精度を組み合わせたREINFORCEポリシー勾配でポリシーを最適化する。
• Rcはパラメータ数に基づく非線形圧縮報酬、Raは学生の検証精度と教師の比率。
• Ax ≤ bの緩和を用いてハードウェア制約を報酬条件に組み込み、段階的なペナルティを適用。
• 教師のロジットを用いた知識蒸留（学生の出力と教師ロジットのL2損失）により学習を導く。

実験結果

リサーチクエスチョン

• RQ1 自動化された強化学習は、より大きな教師ネットワークからの圧縮時に精度を保つコンパクトな学生アーキテクチャを自動的に発見できるか？
• RQ2 二段階のアクション戦略（層削除→層縮小）は現代のアーキテクチャやデータセットにスケールするか？
• RQ3 学習した圧縮方策は類似のアーキテクチャ間または大きな教師間で転移できるか？
• RQ4 ハードウェア制約を報酬に効果的に統合して実用的なモデルを生み出せるか？
• RQ5 教師からの蒸留は圧縮された学生ネットワークの性能を高めるか？

主な発見

• ResNet-34 などのモデルで高い圧縮を達成しつつ教師に近い精度を維持できた。
• 二段階の方策学習により macro（層削除）と micro（層縮小）決定を分離して探索を加速。
• 小さな教師で学んだ方策が大きな教師へ転移し、新しい設定での訓練を加速。
• 複数データセット（MNIST, CIFAR-10/100, SVHN, Caltech-256）で剪定や手作り知識蒸留ベースラインを上回った。
• ハードウェア制約を持つ報酬により、サイズ制約下で実用的なモデルが得られ、実用性が示された。