机器之心报道
此前,由谷歌大脑团队科学家 Aleksandra Faust 和研究员 Srivatsan Krishnan 发布的深度强化学习模型在解决导航、核物理、机器人和游戏等现实世界的顺序决策问题方面取得巨大进步。它很有应用前景,但缺点之一是训练时间过长。
虽然可以使用分布式计算加快复杂困难任务的强化学习的训练,但是需要数百甚至数千个计算节点,且要使用大量硬件资源,这使得强化学习训练成本变得极其高昂,同时还要考虑对环境的影响。最近的研究表明,对现有硬件进行性能优化可以减少模型训练的碳足迹(即温室气体排放总量)。
借助系统优化,可以缩短强化学习训练时间、提高硬件利用率、减少二氧化碳(CO2)排放。其中一种技术是量化,将全精度浮点(FP32)数转换为低精度(int8)数,然后使用低精度数字进行计算。量化可以节省内存成本和带宽,实现更快、更节能的计算。量化已成功应用于监督学习,以实现机器学习(ML)模型的边缘部署并实现更快的训练。同样也可以将量化应用于强化学习训练。
近日,谷歌的研究者在《Transactions of Machine Learning Research》期刊上发表了《QuaRL:快速和环境可持续强化学习的量化》,介绍了一种称为「ActorQ」的新范式。该范式使用了量化,在保持性能的同时,将强化学习训练速度提高 1.5-5.4 倍。作者证明,与全精度训练相比,碳足迹也减少了 1.9-3.8 倍。
量化应用于强化学习训练
在传统的强化学习训练中,learner 策略会应用于 actor,actor 使用该策略探索环境并收集数据样本,actor 收集的样本随后被 learner 用于不断完善初始策略。定期地,针对 learner 的训练策略被用来更新 actor 的策略。为了将量化应用于强化学习训练,作者开创了 ActorQ 范式。ActorQ 执行上面描述的相同序列,其中关键区别是,从 learner 到 actor 的策略更新是量化的,actor 使用 int8 量化策略探索环境以收集样本。
以这种方式将量化应用于强化学习训练有两个关键好处。首先,它减少了策略的内存占用。对于相同的峰值带宽,learner 和 actor 之间传输的数据较少,这降低了 actor 与 learner 之间的策略更新通信成本。其次,actor 对量化策略进行推理,以生成给定环境状态的操作。与完全精确地执行推理相比,量化推理过程要快得多。
传统 RL 训练概述(左)和 ActorQ RL 训练(右)。
在 ActorQ 中,作者使用了 ACME 分布式强化学习框架。量化模块执行统一量化,将 FP32 转换为 int8。actor 使用优化的 int8 计算进行推理。虽然在设计量化模块块时使用的是均匀量化,但其它量化技术可以取代均匀量化并产生类似的结果。actor 收集的样本供 learner 用于训练神经网络策略。学习的策略被量化模块定期量化并广播给 actor。
以量化提升强化学习训练效率
作者在一系列环境中评估 ActorQ,包括 Deepmind Control Suite 和 OpenAI Gym。作者演示了 D4PG 和 DQN 的加速和性能改进。选择 D4PG 是因为它是用于 Deepmind Control Suite 任务的 ACME 中最好的学习算法,而 DQN 是一种广泛使用的标准强化学习算法。
作者在实验中观察到训练强化学习策略的速度显著加快(1.5 倍至 5.41 倍之间)。更重要的是,即使 actor 进行了基于 int8 的量化推理,也可以保持性能。下图显示了用于 Deepmind Control Suite 和 OpenAI Gym 任务的 D4PG 和 DQN 智能体的这一点。
对 D4PG 智能体在各种 Deepmind Control Suite 任务中使用 FP32 策略(q=32)和量化 int8 策略(q=8)的强化学习训练进行比较。量化实现了 1.5 倍到 3.06 倍的速度提升。
OpenAI Gym 环境中 DQN 智能体使用 FP32 策略(q=32)和 int8 量化策略(q=8)进行强化学习训练的比较。量化实现了 2.2 倍到 5.41 倍的加速。
量化减少碳排放
使用 ActorQ 在强化学习中应用量化可以在不影响性能的情况下缩短训练时间。更有效地使用硬件,并减少了碳足迹。作者通过在训练时使用 FP32 的碳排放量与在训练时采用 int8 的排放量之比来衡量碳足迹的改善。
为了测量强化学习训练实验的碳排放量,作者使用了之前工作中提出的 experiment-impact-tracker (https://github.com/Breakend/experiment-impact-tracker)。ActorQ 系统配备碳监测器 API,以测量每个训练实验的碳排放量。
与全精度运行时的碳排放量(FP32)相比,量化可将碳排放量从 1.9 倍降至 3.76 倍,具体取决于任务。由于强化学习系统可以在数千个分布式硬件核心和加速器上运行,绝对碳减排量(以千克二氧化碳计)可能会非常显著。
比较使用 FP32 和 int8 的训练之间的碳排放。X 轴标度标准化为 FP32 的碳排放量,用红色条表示。ActorQ 减少了碳排放。
结论和未来方向
作者引入了 ActorQ,这是一种新的范式,将量化应用于强化学习训练,并在保持性能的同时实现了 1.5-5.4 倍的加速改进。与未应用量化的全精度训练相比,ActorQ 可以将强化学习训练的碳足迹减少 1.9-3.8 倍。
ActorQ 证明量化可以有效地应用于强化学习的许多方面,从获得高质量和高效的量化策略到减少训练时间和碳排放。随着强化学习在解决现实问题方面继续取得长足进步,我们有理由相信,使强化学习训练实现可持续发展将是关键。当将强化学习扩展到数千个 CPU 和 GPU 时,即使 50% 的改进也会在成本、能源和碳排放方面显著降低。作者的工作是将量化应用于强化学习训练以实现高效和环境可持续训练的第一步。
作者在 ActorQ 中的量化设计基于简单的均匀量化,但实际可以应用其他形式的量化、压缩和稀疏性(如蒸馏、稀疏化等)。未来的工作将考虑应用更积极的量化和压缩方法,这可能会为强化学习在训练性能和精度的权衡上带来更多的好处。
原文链接:
https://ai.googleblog.com/2022/09/quantization-for-fast-and.html