[论文解读] GST: Group-Sparse Training for Accelerating Deep Reinforcement Learning

Juhyoung Lee, Sangyeob Kim+3 more|arXiv (Cornell University)|Jan 24, 2021

Reinforcement Learning in Robotics参考文献 31被引用 9

一句话总结

本文提出了一种名为分组稀疏训练（Group-Sparse Training, GST）的新颖权重压缩方法，结合块循环矩阵压缩与奖励感知剪枝，以加速深度强化学习（DRL）训练。GST在训练初期即保持高稀疏率，并通过基于奖励反馈动态调整稀疏度来实现稳定训练，在使用TD3算法的MuJoCo环境中，平均压缩率比迭代剪枝方法高出25–41.5个百分点，且无性能下降。

ABSTRACT

Deep reinforcement learning (DRL) has shown remarkable success in sequential decision-making problems but suffers from a long training time to obtain such good performance. Many parallel and distributed DRL training approaches have been proposed to solve this problem, but it is difficult to utilize them on resource-limited devices. In order to accelerate DRL in real-world edge devices, memory bandwidth bottlenecks due to large weight transactions have to be resolved. However, previous iterative pruning not only shows a low compression ratio at the beginning of training but also makes DRL training unstable. To overcome these shortcomings, we propose a novel weight compression method for DRL training acceleration, named group-sparse training (GST). GST selectively utilizes block-circulant compression to maintain a high weight compression ratio during all iterations of DRL training and dynamically adapt target sparsity through reward-aware pruning for stable training. Thanks to the features, GST achieves a 25 \%p $\sim$ 41.5 \%p higher average compression ratio than the iterative pruning method without reward drop in Mujoco Halfcheetah-v2 and Mujoco humanoid-v2 environment with TD3 training.

研究动机与目标

• 解决资源受限的边缘设备上DRL训练时间过长的挑战。
• 克服迭代剪枝方法在训练初期压缩率低的问题。
• 通过用奖励感知的动态稀疏度调整替代固定稀疏度调度，稳定DRL训练。
• 通过减少大模型参数带来的内存带宽瓶颈，实现在移动设备和边缘设备上的高效DRL训练。
• 在包括Atari、Google Research Football和ImageNet分类任务在内的多样化DRL基准上展示该方法的泛化能力。

提出的方法

• 在训练初期选择性地应用块循环压缩，以维持初始阶段的高压缩率。
• 使用奖励感知剪枝，根据训练奖励趋势动态调整目标稀疏度，避免固定调度带来的不稳定性。
• 引入相位偏移参数 $ S_{\text{shift}} $ 以控制稀疏度提升的时间点，实现从块循环权重到剪枝权重的平稳过渡。
• 实现三种块大小转换方法——投影法、适合block4的block2，以及适合block4的block4，以在压缩率与性能之间取得平衡。
• 选择性地对94–99%的网络参数应用GST，排除首层和末层，以保留关键的特征表示。
• 使用友好的矩阵变换，最小化块转换过程中的参数值偏移，确保训练稳定性。

实验结果

研究问题

• RQ1 块循环压缩能否与迭代剪枝有效结合，以在整个DRL训练过程中维持高压缩率？
• RQ2 与固定稀疏度调度相比，奖励感知的动态稀疏度调度是否能提升DRL训练的稳定性？
• RQ3 GST能否在不牺牲最终性能的前提下，实现比迭代剪枝更高的平均压缩率？
• RQ4 GST在多样化DRL基准上的表现如何，包括连续控制（MuJoCo）、离散控制（Atari）和复杂环境（Google Research Football）？
• RQ5 GST在非DRL任务（如图像分类）上的泛化程度如何，是否展现出更广泛的应用潜力？

主要发现

• 在使用TD3算法的MuJoCo Halfcheetah-v2和Humanoid-v2环境中，GST的平均压缩率比迭代剪枝高出25–41.5个百分点，且无奖励下降。
• 在MuJoCo Humanoid-v2环境中，使用block4-friendly block2方法和 $ S_{\text{shift}} = 0.25 $ 时，GST实现了66.5%的最大压缩率，且无奖励损失。
• 在Atari Breakout环境中使用A2C算法，当 $ B=2 $ 且 $ S_{\text{shift}} = 1.0 $ 时，GST实现了71.9%的平均压缩率，且无奖励下降。
• 在Google Research Football环境中使用PPO算法，当 $ B=4 $ 且 $ S_{\text{shift}} = 1.0 $ 时，GST实现了73.6%的平均压缩率，且无奖励损失。
• 在CIFAR-10数据集上使用ResNet-32，GST实现了68.2%的压缩率，准确率为91.4%（比基线低0.8%），证明了其在分类任务上的泛化能力。
• 在ImageNet数据集上使用AlexNet，GST实现了62.9%的压缩率，准确率为55.8%（比基线低0.4%），证实了其在不同模型架构下的鲁棒性。