[論文レビュー] Topology-Aware Network Pruning using Multi-stage Graph Embedding and Reinforcement Learning

Sixing Yu, Arya Mazaheri+1 more|arXiv (Cornell University)|Feb 5, 2021

Advanced Graph Neural Networks被引用数 7

ひとこと要約

この論文では、深層ニューラルネットワーク（DNN）のトポロジーの階層的・移譲可能な埋め込みを学習するためのマルチステージグラフニューラルネットワーク（m-GNN）と、強化学習（PPO）を組み合わせたトポロジーに配慮したネットワーク prune 法 GNN-RL を提案する。FLOPs 制約付きの最適なチャネル prune 政策を自動で発見する。この手法は、ResNets や VGG-16、モバイルフレンドリーなネットワークを含む多様なモデルにおいて、高い圧縮比と最小限のファインチューニングで最先端の精度を達成する。

ABSTRACT

Model compression is an essential technique for deploying deep neural networks (DNNs) on power and memory-constrained resources. However, existing model-compression methods often rely on human expertise and focus on parameters' local importance, ignoring the rich topology information within DNNs. In this paper, we propose a novel multi-stage graph embedding technique based on graph neural networks (GNNs) to identify DNN topologies and use reinforcement learning (RL) to find a suitable compression policy. We performed resource-constrained (i.e., FLOPs) channel pruning and compared our approach with state-of-the-art model compression methods. We evaluated our method on various models from typical to mobile-friendly networks, such as ResNet family, VGG-16, MobileNet-v1/v2, and ShuffleNet. Results show that our method can achieve higher compression ratios with a minimal fine-tuning cost yet yields outstanding and competitive performance.

研究の動機と目的

• 既存のモデル圧縮手法が DNN トポロジーを無視し、手作業で定義されたルールや局所的なパrameter の重要性に依存するという制限に対処すること。
• DNN における構造的モチーフと階層的パターンを捉える包括的かつ移譲可能な prune 政策を開発すること。
• FLOPs 制約を強制する DNN-Graph 環境を導入することで、RL 基盤の prune における報酬と圧縮比のトレードオフを克服すること。
• トポロジーに配慮したグラフ埋め込みに従って駆動される PPO を用いた RL エージェントにより、人為的な干渉を最小限に抑え、自動的かつ適応的な prune を可能にすること。

提案手法

• DNN を階層的な計算グラフとしてモデル化し、構造的モチーフとトポロジーを捉えることで、移譲可能な表現学習を可能にする。
• グラフ層を段階的に処理することで階層的グラフ埋め込みを学習するマルチステージ GNN（m-GNN）を提案する。これにより冗長性が低減され、効率性が向上する。
• 各 prune ステップをグラフ状態に変換する DNN-Graph 環境を設計する。これにより、RL エージェントは圧縮モデルをトポロジー変更プロセスとして操作できる。
• プロキシ・ポリシー最適化（PPO）を RL アルゴリズムとして使用する。エージェントは各層ごとの prune 確率を選択し、FLOPs 制約を満たした際に精度を報酬として受信する。
• グラフ埋め込みと RL を統合することで、エージェントがパrameter の統計のみではなく、構造的重要性に基づいて推論できるようにする。
• FLOPs を主な制約としてチャネル prune に適用し、ターゲットサイズに達するまで自動的に圧縮を継続可能にする。

実験結果

リサーチクエスチョン

• RQ1 グラフニューラルネットワークは、モデル圧縮のための DNN の階層的・再帰的な構造的モチーフを効果的に捉え、表現できるか？
• RQ2 トポロジーに配慮したグラフ埋め込みは、異なる DNN アーキテクチャ間での prune 政策の一般化と移譲性をどのように向上させるか？
• RQ3 手動で定義された報酬形状やヒューリスティックな制約に依存せずに、強化学習エージェントが最適な FLOPs 制約付きの prune 政策を学習できるか？
• RQ4 提案手法は、精度、圧縮比、推論効率の観点で、最先端の prune ベースラインをどの程度上回るか？
• RQ5 学習されたトポロジー埋め込みは、再トレーニングを最小限に抑えながら、異なる DNN を効果的に prune するために転送可能か？

主な発見

• ResNet-20 において、GNN-RL は FLOPs を 49% 減少させたが、精度低下はたった 0.1% にとどまり、VGG-16 では 1.38 倍の推論速度向上を達成した。
• ResNet-110 では、元のモデルと比較して 0.63% の精度向上を達成した。これは、prune が一般化性能を向上させることを示している。
• VGG-16 では、全結合層を prune し、特徴マップのサイズを削減することで、GPU メモリ使用量を 141 MB（26.7%）削減した。
• MobileNet-v2 は元の FLOPs の 79% に圧縮され、精度低下はたった 1.4% にとどまった。これは、モバイルフレンドリーなモデルにおいて優れた性能を示している。
• m-GNN は、標準的な GNN よりも 30–50% 少ないグラフノードとエッジで、より高いノード分類精度を達成した。これは、構造的モチーフを効率的に捉える能力を裏付けている。
• GNN-RL が学習したトポロジーに配慮した prune 政策は、エージェントの再トレーニングなしに多様な DNN を効果的に prune するために転送可能であった。