[论文解读] Single-Shot Pruning for Offline Reinforcement Learning

Samin Yeasar Arnob, Riyasat Ohib+2 more|arXiv (Cornell University)|Dec 31, 2021

Machine Learning and Data Classification被引用 4

一句话总结

本文提出一种在离线强化学习（RL）训练前应用的单次剪枝方法，利用固定数据集识别稀疏但高性能的子网络。通过在初始化时应用SNIP和GraSP剪枝，该方法在95%剪枝率下仍能达到与密集网络相当的性能，内存使用减少4倍且无需压缩，同时在资源受限设备上实现更快推理。

ABSTRACT

Deep Reinforcement Learning (RL) is a powerful framework for solving complex real-world problems. Large neural networks employed in the framework are traditionally associated with better generalization capabilities, but their increased size entails the drawbacks of extensive training duration, substantial hardware resources, and longer inference times. One way to tackle this problem is to prune neural networks leaving only the necessary parameters. State-of-the-art concurrent pruning techniques for imposing sparsity perform demonstrably well in applications where data distributions are fixed. However, they have not yet been substantially explored in the context of RL. We close the gap between RL and single-shot pruning techniques and present a general pruning approach to the Offline RL. We leverage a fixed dataset to prune neural networks before the start of RL training. We then run experiments varying the network sparsity level and evaluating the validity of pruning at initialization techniques in continuous control tasks. Our results show that with 95% of the network weights pruned, Offline-RL algorithms can still retain performance in the majority of our experiments. To the best of our knowledge, no prior work utilizing pruning in RL retained performance at such high levels of sparsity. Moreover, pruning at initialization techniques can be easily integrated into any existing Offline-RL algorithms without changing the learning objective.

研究动机与目标

• 通过在训练前实现高效的模型压缩，解决大型神经网络在离线RL中计算和内存成本过高的问题。
• 探索原本为监督学习设计的单次剪枝技术是否可有效应用于具有固定数据集的离线RL。
• 在极端剪枝率（高达95%）下评估剪枝网络在连续控制任务中的性能。
• 证明在初始化时进行剪枝可无缝集成到现有离线RL算法中，而无需修改学习目标。
• 在机器人和边缘AI等实时、低资源应用场景中，实现更快的训练和推理。

提出的方法

• 在离线RL训练前，对随机初始化的神经网络应用单次剪枝方法——SNIP和GraSP。
• 使用D4RL基准中的固定数据集，基于梯度敏感性（GraSP）或连接重要性（SNIP）计算剪枝标准。
• 在初始化时一次性剪枝，无需任何微调或迭代优化，以生成稀疏子网络。
• 在连续控制环境（HalfCheetah、Hopper、Walker2D）中，使用标准离线RL算法（BCQ、BC）评估剪枝网络的性能。
• 通过测量性能、内存使用量和逐层剪枝率，评估剪枝对泛化能力和效率的影响。
• 使用PyTorch的稀疏张量功能存储并比较密集模型与剪枝模型的内存需求。

实验结果

研究问题

• RQ1 在训练前应用初始化时的单次剪枝，是否能在离线RL中保持性能？
• RQ2 在离线RL中，可实现的最高剪枝率（如95%）是多少，且性能不会下降？
• RQ3 不同的剪枝标准（SNIP与GraSP）如何影响离线RL中剪枝网络的结构与性能？
• RQ4 剪枝在多大程度上减少了内存使用并提升了离线RL模型的推理效率？
• RQ5 在不修改学习目标的前提下，初始化时的剪枝能否无缝集成到现有离线RL算法中？

主要发现

• 即使在95%剪枝率下，所提方法的性能仍与密集网络相当，多个环境中回报无显著下降。
• 使用GraSP剪枝在最终层保留了更多权重，与其保持梯度流动的目标一致。
• 对于95%剪枝率的模型，内存使用量相比密集模型减少了4倍，即使未使用压缩技术。
• 该方法即插即用：可应用于任何离线RL算法，无需修改学习目标或增加额外训练。
• 实验表明，改变批量大小或剪枝迭代次数均无法在单次剪枝循环之外提升性能。
• 结果表明，固定数据集的离线RL使有效的单次剪枝成为可能，而在线RL由于数据分布变化无法实现此方法。