[论文解读] 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

一句话总结

该论文提出 GNN-RL，一种拓扑感知的网络剪枝方法，结合多阶段图神经网络（m-GNN）学习深度神经网络（DNN）拓扑的分层可迁移表征，并利用强化学习（PPO）自动发现最优的 FLOPs 约束下的通道剪枝策略。该方法在多种模型（包括 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 拓扑结构、依赖手工规则或局部参数重要性的问题。
• 开发一种统一的、可迁移的剪枝策略，以捕捉 DNN 中的结构模式与分层特征。
• 通过引入一种 DNN-Graph 环境以缓解强化学习剪枝中的奖励-剪枝率权衡问题，该环境强制执行 FLOPs 约束。
• 通过基于 PPO 的强化学习智能体，结合拓扑感知图嵌入，实现自动、自适应的剪枝，最大限度减少人工干预。

提出的方法

• 将 DNN 建模为分层计算图，以捕捉结构模式与拓扑信息，支持可迁移的表征学习。
• 提出多阶段图神经网络（m-GNN），通过顺序处理图层来学习分层图嵌入，降低冗余并提升效率。
• 设计 DNN-Graph 环境，将每个剪枝步骤转化为图状态，使强化学习智能体能够将压缩过程视为拓扑结构的修改过程。
• 采用近端策略优化（PPO）作为强化学习算法，智能体按层选择剪枝比例，并在满足 FLOPs 约束后接收精度作为稀疏奖励。
• 将图嵌入与强化学习结合，使智能体能够基于结构重要性进行推理，而不仅依赖参数统计信息。
• 将方法应用于以 FLOPs 为主要约束的通道剪枝，实现自动压缩直至达到目标模型大小。

实验结果

研究问题

• RQ1 图神经网络能否有效捕捉并表征 DNN 中的分层与重复性结构模式，以支持模型压缩？
• RQ2 拓扑感知的图嵌入在多大程度上提升了剪枝策略在不同 DNN 架构间的泛化性与可迁移性？
• RQ3 强化学习智能体能否在不依赖人工设计奖励函数或启发式约束的情况下，学习到最优的 FLOPs 约束剪枝策略？
• RQ4 与当前最先进的剪枝基线相比，该方法在精度、剪枝率和推理效率方面表现如何？
• RQ5 所学习的拓扑嵌入在无需重新训练智能体的前提下，能否有效用于剪枝不同 DNN，实现最小化微调？

主要发现

• 在 ResNet-20 上，GNN-RL 实现了 49% 的 FLOPs 减少，精度仅下降 0.1%，同时在 VGG-16 上实现 1.38 倍的推理加速。
• 在 ResNet-110 上，剪枝后的模型相比原始模型精度提升了 0.63%，表明剪枝可改善泛化能力。
• 通过剪枝全连接层并减小特征图尺寸，VGG-16 的 GPU 显存使用量减少了 141 MB（26.7%）。
• MobileNet-v2 的 FLOPs 被压缩至原始值的 79%，精度仅下降 1.4%，显示出在轻量级模型上的优异性能。
• 与标准 GNN 相比，m-GNN 在减少 30–50% 图节点与边数的前提下，实现了更高的节点分类准确率，证明其在捕捉结构模式方面的高效性。
• GNN-RL 学习到的拓扑感知剪枝策略具备可迁移性，无需重新训练智能体即可有效剪枝多种 DNN。