[论文解读] Learning to Prune Deep Neural Networks via Reinforcement Learning

Manas Gupta, Siddharth Aravindan+3 more|arXiv (Cornell University)|Jul 9, 2020

Advanced Neural Network Applications参考文献 14被引用 7

一句话总结

该论文提出PuRL，一种基于强化学习的深度神经网络剪枝方法，通过在每个剪枝步骤使用密集奖励来加速训练并提高样本效率。PuRL在ResNet-50上实现了80.27%的稀疏度，ImageNet上的Top-1准确率为75.37%，将强化学习训练episode数从AMC的400次减少至仅55次，同时达到最先进性能。

ABSTRACT

This paper proposes PuRL - a deep reinforcement learning (RL) based algorithm for pruning neural networks. Unlike current RL based model compression approaches where feedback is given only at the end of each episode to the agent, PuRL provides rewards at every pruning step. This enables PuRL to achieve sparsity and accuracy comparable to current state-of-the-art methods, while having a much shorter training cycle. PuRL achieves more than 80% sparsity on the ResNet-50 model while retaining a Top-1 accuracy of 75.37% on the ImageNet dataset. Through our experiments we show that PuRL is also able to sparsify already efficient architectures like MobileNet-V2. In addition to performance characterisation experiments, we also provide a discussion and analysis of the various RL design choices that went into the tuning of the Markov Decision Process underlying PuRL. Lastly, we point out that PuRL is simple to use and can be easily adapted for various architectures.

研究动机与目标

• 通过强化学习自动化神经网络剪枝，以提升效率和可扩展性。
• 解决现有基于强化学习的剪枝方法因仅在episode结束时给予稀疏奖励而导致的收敛缓慢问题。
• 设计一种具有密集奖励和高效动作空间的马尔可夫决策过程（MDP），以实现更快、更有效的剪枝策略学习。
• 在多种架构（包括ResNet-50、MobileNet-V2和EfficientNet-B2）上评估PuRL，证明其泛化能力和易于适配性。
• 通过消融研究分析MDP各组件（状态、动作、奖励）的影响，为未来基于强化学习的剪枝设计提供指导。

提出的方法

• 将网络剪枝形式化为包含状态、动作、奖励、转移和折扣因子的马尔可夫决策过程（MDP）。
• 采用密集奖励函数，在每个剪枝步骤后提供反馈，相比仅在episode结束时给予奖励的方法，可实现更快的策略学习。
• 采用基于权重幅度的剪枝准则，使用由权重标准差导出的阈值：若|w| < ασ(w)，则剪除该权重。
• 引入对α值（例如0.0, 0.1, ..., 2.2）的离散动作空间，并通过将步长增加至0.2来减少搜索复杂度。
• 采用迭代剪枝策略，逐步提升稀疏度目标，并结合微调以在压缩过程中保持准确率。
• 采用两种状态表示：低维元组（层索引、准确率、已剪枝比例）和高维的每层准确率与稀疏度向量。

实验结果

研究问题

• RQ1 与仅在episode结束时给予稀疏奖励相比，在每个剪枝步骤提供密集奖励是否能显著提升学习效率和收敛速度？
• RQ2 不同的状态表示如何影响智能体学习有效剪枝策略的能力？
• RQ3 减少动作空间中的动作数量是否可通过改善探索和策略学习，带来更好的性能？
• RQ4 所提出方法在不同神经网络架构（包括MobileNet-V2和EfficientNet-B2等高效模型）上的泛化能力如何？
• RQ5 奖励设计、动作空间设计和状态表示对最终稀疏度和准确率的相对影响是什么？

主要发现

• PuRL将ResNet-50上所需的强化学习episode数从AMC的400次减少至55次，实现85%的减少，同时保持了具有竞争力的准确率。
• 与稀疏奖励相比，密集奖励使最终准确率提升24个百分点，稀疏度提升4个百分点，证明其学习效率更优。
• 使用更大的动作步长（0.2而非0.1）可获得帕累托占优解，在稀疏度和准确率上均优于基线。
• PuRL在ResNet-50上实现80.27%的稀疏度，ImageNet上的Top-1准确率为75.37%，达到最先进性能。
• 该方法在高效架构上泛化良好：在不进行超参数调优的情况下，PuRL在MobileNet-V2和EfficientNet-B2上的稀疏度超过AMC的1.5倍。
• 消融研究证实，密集奖励和减少的动作空间对快速收敛和高性能至关重要，而状态表示的选择影响最小。