Resource-Aware Neural Network Pruning Using Graph-based Reinforcement Learning
作者: Dieter Balemans, Thomas Huybrechts, Jan Steckel, Siegfried Mercelis
分类: cs.LG, cs.AI
发布日期: 2025-09-04
💡 一句话要点
提出基于图强化学习的资源感知神经网络剪枝方法，提升剪枝效率和性能。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 神经网络剪枝 图神经网络 强化学习 AutoML 模型压缩
📋 核心要点
- 现有剪枝方法依赖手工启发式和局部优化，导致性能次优和效率低下。
- 提出基于图的强化学习剪枝框架，利用图注意力网络编码网络拓扑结构，学习最优通道重要性。
- 实验表明，该方法在CIFAR和ImageNet等数据集上优于传统剪枝方法，实现更优的性能。
📝 摘要（中文）
本文提出了一种新的神经网络剪枝方法，该方法将基于图的观察空间集成到AutoML框架中，以解决现有方法的局限性。传统剪枝方法依赖于手工设计的启发式方法和局部优化，导致次优性能和低效的剪枝策略。本文通过引入神经网络的图表示来转换剪枝过程，该图表示捕获层和通道之间的完整拓扑关系，用网络结构的全局视图取代了有限的逐层观察空间。核心创新包括一个图注意力网络（GAT）编码器，用于处理网络的图表示并生成丰富的嵌入。此外，对于动作空间，我们从连续剪枝率过渡到细粒度的二元动作空间，使agent能够直接从数据中学习最佳通道重要性标准，从而摆脱预定义的评分函数。这些贡献在约束马尔可夫决策过程（CMDP）框架中建模，允许agent在满足资源约束（如目标压缩率）的同时做出明智的剪枝决策。为此，我们设计了一个自我竞争奖励系统，鼓励agent在满足约束的同时超越其之前的最佳性能。通过在CIFAR-10、CIFAR-100和ImageNet等基准数据集上的大量实验，证明了该方法的有效性。实验表明，该方法始终优于传统的剪枝技术，在学习特定于任务的剪枝策略的同时，展示了最先进的结果，这些策略识别了功能上冗余的连接，而不仅仅是简单的权重幅度考虑。
🔬 方法详解
问题定义：现有神经网络剪枝方法通常依赖于手工设计的启发式规则和局部优化策略，无法充分利用网络结构的全局信息，导致剪枝后的网络性能下降，且难以适应不同的任务和资源约束。这些方法通常基于权重幅度等简单指标进行剪枝，忽略了网络中复杂的连接关系和功能冗余。
核心思路：本文的核心思路是将神经网络的剪枝过程建模为一个强化学习问题，并利用图神经网络来学习网络结构的全局表示。通过图表示，agent可以更好地理解网络中不同层和通道之间的依赖关系，从而做出更明智的剪枝决策。此外，采用细粒度的二元动作空间，允许agent直接从数据中学习通道的重要性，避免了手工设计评分函数的局限性。
技术框架：该方法采用Constrained Markov Decision Process (CMDP)框架。整体流程如下：首先，将目标神经网络表示为一个图结构，节点代表层或通道，边代表连接关系。然后，使用图注意力网络（GAT）对图进行编码，生成节点的嵌入表示。agent根据嵌入表示选择要剪枝的通道，并执行剪枝操作。剪枝后，对网络进行微调，并根据性能和资源约束计算奖励信号。agent通过与环境的交互不断学习，优化剪枝策略。
关键创新：该方法最重要的创新点在于将图神经网络引入到神经网络剪枝中，利用图结构来表示网络的拓扑关系，从而使agent能够更好地理解网络的全局结构和功能。此外，采用细粒度的二元动作空间，允许agent直接从数据中学习通道的重要性，避免了手工设计评分函数的局限性。自我竞争奖励系统鼓励agent在满足约束的同时超越其之前的最佳性能。
关键设计：图注意力网络（GAT）用于编码网络图结构，学习节点的嵌入表示。奖励函数设计为同时考虑剪枝后的性能和资源约束，采用自我竞争机制，鼓励agent不断提升性能。动作空间为细粒度的二元动作空间，每个动作对应一个通道是否被剪枝。使用Adam优化器训练强化学习agent。
🖼️ 关键图片
📊 实验亮点
实验结果表明，该方法在CIFAR-10、CIFAR-100和ImageNet等基准数据集上取得了优于传统剪枝方法的结果。例如，在ImageNet上，该方法在保持相同精度的情况下，可以实现更高的压缩率。与基于权重幅度的剪枝方法相比，该方法能够识别功能上冗余的连接，从而实现更有效的剪枝。
🎯 应用场景
该研究成果可应用于资源受限的设备（如移动设备、嵌入式系统）上的深度学习模型部署。通过智能剪枝，可以在保证模型性能的前提下，显著降低模型的计算复杂度和存储空间，从而实现更高效的模型推理和部署。此外，该方法还可以用于自动化模型压缩和加速，降低人工调参的成本。
📄 摘要（原文）
This paper presents a novel approach to neural network pruning by integrating a graph-based observation space into an AutoML framework to address the limitations of existing methods. Traditional pruning approaches often depend on hand-crafted heuristics and local optimization perspectives, which can lead to suboptimal performance and inefficient pruning strategies. Our framework transforms the pruning process by introducing a graph representation of the target neural network that captures complete topological relationships between layers and channels, replacing the limited layer-wise observation space with a global view of network structure. The core innovations include a Graph Attention Network (GAT) encoder that processes the network's graph representation and generates a rich embedding. Additionally, for the action space we transition from continuous pruning ratios to fine-grained binary action spaces which enables the agent to learn optimal channel importance criteria directly from data, moving away from predefined scoring functions. These contributions are modelled within a Constrained Markov Decision Process (CMDP) framework, allowing the agent to make informed pruning decisions while adhering to resource constraints such as target compression rates. For this, we design a self-competition reward system that encourages the agent to outperform its previous best performance while satisfying the defined constraints. We demonstrate the effectiveness of our approach through extensive experiments on benchmark datasets including CIFAR-10, CIFAR-100, and ImageNet. The experiments show that our method consistently outperforms traditional pruning techniques, showing state-of-the-art results while learning task-specific pruning strategies that identify functionally redundant connections beyond simple weight magnitude considerations.