深度强化学习中的剪枝网络：提升性能的新途径

人工智能咨询培训老师叶梓 转载标明出处

尽管深度强化学习在现实世界问题中的应用取得了成功，但训练这些网络时出现的挑战和病态行为也日益明显。研究表明，DRL代理在训练过程中存在隐式欠参数化问题，大量神经元在训练过程中变得不活跃。 Google DeepMind、魁北克人工智能研究所 Mila 和蒙特利尔大学的研究团队通过引入稀疏训练技术的优势，展示了渐进式幅度剪枝技术如何使代理最大化参数有效性，从而在性能上取得显著提升。

剪枝提升深度强化学习性能

对于使用Impala-based ResNet架构的DQN和Rainbow代理，网络宽度的扩展与性能的关系

作者选择了Jax实现的Dopamine库作为实验的基础框架，因为它提供了高效的实现和灵活性。Dopamine库中的代理默认使用了Impala架构，这是一个基于ResNet的15层网络结构，它在先前的研究中已被证明能够带来性能上的提升。

实验中，作者使用了JaxPruner库来实现渐进式幅度剪枝，这个库与Dopamine库紧密集成，方便了剪枝过程的实现。剪枝计划遵循了Graesser等人（2022年）的实验设置，即从训练的20%开始剪枝，直到80%结束，之后保持网络的稀疏状态不变。这种剪枝计划通过逐渐减少网络中的参数数量，旨在提高剩余参数的效率和网络的整体性能。

作者在Arcade Learning Environment（ALE）上进行了实验，选择了15种不同的游戏来测试代理的性能，这些游戏被选中是因为它们在难度和类型上具有多样性。实验在40M帧的范围内进行，尽管这个帧数少于通常的200M帧，但作者发现在40M帧时算法之间的定性差异与在100M帧时基本一致。

实验遵循了Agarwal等人（2021年）的评估准则，每个实验都使用了5个独立的随机种子进行，以确保结果的可靠性和普遍性。所有的实验结果都报告了人类标准化的四分位均值（IQM），并给出了95%的分层自举置信区间。这些实验都在NVIDIA Tesla P100 GPU上运行，每个实验大约需要2天的时间来完成。通过这些详细的实现和评估方法，作者确保了实验结果的准确性和复现性。

实验中使用的逐步幅度剪枝计划，目标是达到95%的稀疏度

在线强化学习部分作者探讨了在网络宽度扩大的情况下，剪枝网络能否维持高水平的稀疏性而不降低性能。实验结果表明，当对DQN使用Impala架构并将卷积层的宽度扩大3倍时，90%和95%的稀疏性可以带来33%的性能提升，而99%的稀疏性能够保持性能。作者主要关注95%的稀疏性水平，因为它在初步探索中始终提供了最佳性能。

对比了使用Mnih et al. (2015)的原始CNN架构的DQN和Rainbow代理，在网络宽度扩展时的性能

在数据受限的环境中，作者评估了剪枝代理的性能。他们发现，在只有100k环境交互的Atari 100k基准测试中，剪枝代理并没有表现出增益。然而，当训练环境交互达到40M时，剪枝代理表现出显著的性能提升。特别是，在DrQ(𝜋?)训练中，剪枝代理避免了基线在长时间训练时的性能下降。

DrQ(𝜋?)和DER代理在训练40M帧时的性能，两者都使用宽度为3的ResNet架构

在离线强化学习部分，作者使用了CQL和CQL+C51这两种最新的先进方法，并发现使用剪枝网络时，这两种方法都观察到了显著的性能提升，特别是对于更宽的网络。有趣的是，在离线环境中，剪枝还有助于避免使用浅层网络（宽度比例为1）时的性能崩溃，或者在CQL+C51的情况下提高最终性能。

离线代理CQL和CQL+C51使用ResNet架构时，网络宽度扩展的性能

作者将注意力转向了价值函数方法之外的策略梯度方法，特别是Soft Actor-Critic (SAC)算法，以评估剪枝技术是否同样适用于这类算法。SAC是一种流行的策略梯度算法，它在连续控制任务中表现出色。实验在MuJoCo环境中进行，涉及多个具有挑战性的环境。结果显示，在某些环境如Walker2d-v2和Ant-v2中，剪枝技术能够带来性能上的提升。这些结果表明，剪枝不仅在价值函数方法中有效，也能够增强策略梯度方法的性能。

在MuJoCo环境中，改变稀疏度参数对SAC性能的影响

作者还探讨了剪枝网络的稳定性。他们遵循了之前研究中提出的剪枝计划，并质疑了仅在训练后期保持稀疏网络是否足够稳定。为了验证这一点，他们对DQN进行了长达100M帧的训练，并应用了两种不同的剪枝计划。结果令人鼓舞，即使在训练周期压缩的情况下，剪枝网络仍然保持了其强大的性能，并且在某些情况下，剪枝网络比使用常规计划的网络更快地达到了高性能。这些发现支持了剪枝技术可以作为一种提高代理性能的稳定手段的观点。

对于DQN使用Impala-based ResNet（宽度=3）时，不同剪枝计划的影响

最后作者探讨了学习率和批量大小对大型神经网络性能的影响。对于深度Q网络（DQN），默认学习率为6.25×10^-5。作者进行了实验，将学习率除以宽度缩放因子（例如，宽度因子为3时为2.08×10^-5，宽度因子为5时为1.25×10^-5）。虽然这些调整后的学习率改善了基线模型的性能，但剪枝技术仍然展现出更强的性能优势（见图28）。此外，当评估不同的批量大小值时，也观察到了类似的趋势。所有基于价值的代理默认批量大小为32，实验中尝试了16和64的批量大小。在所有情况下，剪枝都保持了其显著的优势（见图29）。这些结果与剪枝可以作为提高代理性能的一种即插即用机制的观点相一致。

对于使用ResNet架构的DQN代理，在不同宽度因子（分别是3和5）下的学习能力评估

对于使用ResNet架构的DQN代理，当网络宽度乘数为3时，批量大小对性能的影响

为什么剪枝如此有效？

作者探讨了剪枝技术与其他几种方法相比较的效果。这些方法包括传统的权重衰减（Weight Decay）和最近提出的网络重置（Network Reset）技术。权重衰减是一种通过在损失函数中加入权重的平方和来防止权重过大的正则化手段，而网络重置则是为了解决训练过程中网络可塑性下降的问题，通过周期性地重置网络参数来维持其灵活性。

实验结果揭示了剪枝技术在不同网络宽度下的性能表现。作者发现，即使在调整了学习率以适应不同宽度的网络之后，剪枝仍然能够带来性能上的显著提升。具体来说，当学习率按照宽度缩放因子进行调整时，虽然基线模型的性能有所改善，但剪枝后的网络性能仍然超越了这些调整。作者还测试了不同的批量大小，包括16和64，以评估批量大小对性能的影响。在所有测试的批量大小下，剪枝后的网络都保持了其性能上的优势。

这些发现表明，剪枝技术的效果并不仅仅依赖于正则化或可塑性增强，而是提供了一种更为全面和根本的改进。剪枝通过减少网络中的冗余参数，不仅提高了参数的效率，还可能通过降低梯度的共线性来增强网络的学习能力。这些结果进一步支持了剪枝可以作为一种通用的策略，用于提升各种深度强化学习代理的性能。

在四个代表性游戏中应用剪枝时的实证分析，包括训练回报、平均𝐴?目标方差、网络参数平均范数、𝐴?估计平均范数、有效秩和休眠神经元

作者进一步分析了剪枝对网络可塑性的影响。可塑性是指神经网络在面对不断变化的数据分布时快速调整自身参数的能力，这对于动态和非静态的强化学习环境至关重要。研究指出，随着训练的进行，强化学习网络往往会失去这种可塑性，导致性能下降。作者通过评估梯度的协方差结构来量化网络的可塑性，发现剪枝网络展现出了较低的梯度相关性，这表明剪枝有助于减少梯度干扰，从而维持网络的灵活性和适应性。剪枝后的网络在训练过程中能够更快地达到更高的性能水平，并且在更长的训练周期中保持稳定，这进一步证实了剪枝技术通过提升网络的可塑性，有助于提高其在复杂环境中的长期学习效率和性能。

Breakout和VideoPinball游戏中密集网络和剪枝网络的梯度协方差矩阵

研究表明，通过剪枝，网络能够以更少的参数实现更高的效率和稳定性，这对于构建更复杂任务中的强化学习代理具有重要意义。剪枝技术的普适性和对不同网络架构的适用性，为未来在多任务泛化、样本效率和泛化能力方面的研究提供了新的思路。作者期望这项工作能激发对非标准网络架构和拓扑的探索，以最大化强化学习代理的性能，并为硬件加速器和边缘设备应用带来新的可能性。

论文链接：https://arxiv.org/abs/2402.12479