訓(xùn)練速度提高最多5.4倍，谷歌提出RL訓(xùn)練新范式ActorQ

近日，谷歌的研究者提出了一種稱為「ActorQ」的新范式，使用量化將強化學(xué)習(xí)訓(xùn)練速度提高 1.5-5.4 倍。與全精度訓(xùn)練相比，碳足跡也減少了 1.9-3.8 倍。

此前，由谷歌大腦團(tuán)隊科學(xué)家 Aleksandra Faust 和研究員 Srivatsan Krishnan 發(fā)布的深度強化學(xué)習(xí)模型在解決導(dǎo)航、核物理、機器人和游戲等現(xiàn)實世界的順序決策問題方面取得巨大進(jìn)步。它很有應(yīng)用前景，但缺點之一是訓(xùn)練時間過長。

雖然可以使用分布式計算加快復(fù)雜困難任務(wù)的強化學(xué)習(xí)的訓(xùn)練，但是需要數(shù)百甚至數(shù)千個計算節(jié)點，且要使用大量硬件資源，這使得強化學(xué)習(xí)訓(xùn)練成本變得極其高昂，同時還要考慮對環(huán)境的影響。最近的研究表明，對現(xiàn)有硬件進(jìn)行性能優(yōu)化可以減少模型訓(xùn)練的碳足跡（即溫室氣體排放總量）。

借助系統(tǒng)優(yōu)化，可以縮短強化學(xué)習(xí)訓(xùn)練時間、提高硬件利用率、減少二氧化碳（CO2）排放。其中一種技術(shù)是量化，將全精度浮點（FP32）數(shù)轉(zhuǎn)換為低精度（int8）數(shù)，然后使用低精度數(shù)字進(jìn)行計算。量化可以節(jié)省內(nèi)存成本和帶寬，實現(xiàn)更快、更節(jié)能的計算。量化已成功應(yīng)用于監(jiān)督學(xué)習(xí)，以實現(xiàn)機器學(xué)習(xí)（ML）模型的邊緣部署并實現(xiàn)更快的訓(xùn)練。同樣也可以將量化應(yīng)用于強化學(xué)習(xí)訓(xùn)練。

近日，谷歌的研究者在《Transactions of Machine Learning Research》期刊上發(fā)表了《QuaRL：快速和環(huán)境可持續(xù)強化學(xué)習(xí)的量化》，介紹了一種稱為「ActorQ」的新范式。該范式使用了量化，在保持性能的同時，將強化學(xué)習(xí)訓(xùn)練速度提高 1.5-5.4 倍。作者證明，與全精度訓(xùn)練相比，碳足跡也減少了 1.9-3.8 倍。

量化應(yīng)用于強化學(xué)習(xí)訓(xùn)練

在傳統(tǒng)的強化學(xué)習(xí)訓(xùn)練中，learner 策略會應(yīng)用于 actor，actor 使用該策略探索環(huán)境并收集數(shù)據(jù)樣本，actor 收集的樣本隨后被 learner 用于不斷完善初始策略。定期地，針對 learner 的訓(xùn)練策略被用來更新 actor 的策略。為了將量化應(yīng)用于強化學(xué)習(xí)訓(xùn)練，作者開創(chuàng)了 ActorQ 范式。ActorQ 執(zhí)行上面描述的相同序列，其中關(guān)鍵區(qū)別是，從 learner 到 actor 的策略更新是量化的，actor 使用 int8 量化策略探索環(huán)境以收集樣本。

以這種方式將量化應(yīng)用于強化學(xué)習(xí)訓(xùn)練有兩個關(guān)鍵好處。首先，它減少了策略的內(nèi)存占用。對于相同的峰值帶寬，learner 和 actor 之間傳輸?shù)臄?shù)據(jù)較少，這降低了 actor 與 learner 之間的策略更新通信成本。其次，actor 對量化策略進(jìn)行推理，以生成給定環(huán)境狀態(tài)的操作。與完全精確地執(zhí)行推理相比，量化推理過程要快得多。

傳統(tǒng) RL 訓(xùn)練概述（左）和 ActorQ RL 訓(xùn)練（右）。

在 ActorQ 中，作者使用了 ACME 分布式強化學(xué)習(xí)框架。量化模塊執(zhí)行統(tǒng)一量化，將 FP32 轉(zhuǎn)換為 int8。actor 使用優(yōu)化的 int8 計算進(jìn)行推理。雖然在設(shè)計量化模塊塊時使用的是均勻量化，但其它量化技術(shù)可以取代均勻量化并產(chǎn)生類似的結(jié)果。actor 收集的樣本供 learner 用于訓(xùn)練神經(jīng)網(wǎng)絡(luò)策略。學(xué)習(xí)的策略被量化模塊定期量化并廣播給 actor。

以量化提升強化學(xué)習(xí)訓(xùn)練效率

作者在一系列環(huán)境中評估 ActorQ，包括 Deepmind Control Suite 和 OpenAI Gym。作者演示了 D4PG 和 DQN 的加速和性能改進(jìn)。選擇 D4PG 是因為它是用于 Deepmind Control Suite 任務(wù)的 ACME 中最好的學(xué)習(xí)算法，而 DQN 是一種廣泛使用的標(biāo)準(zhǔn)強化學(xué)習(xí)算法。

作者在實驗中觀察到訓(xùn)練強化學(xué)習(xí)策略的速度顯著加快（1.5 倍至 5.41 倍之間）。更重要的是，即使 actor 進(jìn)行了基于 int8 的量化推理，也可以保持性能。下圖顯示了用于 Deepmind Control Suite 和 OpenAI Gym 任務(wù)的 D4PG 和 DQN 智能體的這一點。

對 D4PG 智能體在各種 Deepmind Control Suite 任務(wù)中使用 FP32 策略（q=32）和量化 int8 策略（q=8）的強化學(xué)習(xí)訓(xùn)練進(jìn)行比較。量化實現(xiàn)了 1.5 倍到 3.06 倍的速度提升。

OpenAI Gym 環(huán)境中 DQN 智能體使用 FP32 策略（q=32）和 int8 量化策略（q=8）進(jìn)行強化學(xué)習(xí)訓(xùn)練的比較。量化實現(xiàn)了 2.2 倍到 5.41 倍的加速。

量化減少碳排放

使用 ActorQ 在強化學(xué)習(xí)中應(yīng)用量化可以在不影響性能的情況下縮短訓(xùn)練時間。更有效地使用硬件，并減少了碳足跡。作者通過在訓(xùn)練時使用 FP32 的碳排放量與在訓(xùn)練時采用 int8 的排放量之比來衡量碳足跡的改善。

為了測量強化學(xué)習(xí)訓(xùn)練實驗的碳排放量，作者使用了之前工作中提出的 experiment-impact-tracker (https://github.com/Breakend/experiment-impact-tracker)。ActorQ 系統(tǒng)配備碳監(jiān)測器 API，以測量每個訓(xùn)練實驗的碳排放量。

與全精度運行時的碳排放量（FP32）相比，量化可將碳排放量從 1.9 倍降至 3.76 倍，具體取決于任務(wù)。由于強化學(xué)習(xí)系統(tǒng)可以在數(shù)千個分布式硬件核心和加速器上運行，絕對碳減排量（以千克二氧化碳計）可能會非常顯著。

比較使用 FP32 和 int8 的訓(xùn)練之間的碳排放。X 軸標(biāo)度標(biāo)準(zhǔn)化為 FP32 的碳排放量，用紅色條表示。ActorQ 減少了碳排放。

結(jié)論和未來方向

作者引入了 ActorQ，這是一種新的范式，將量化應(yīng)用于強化學(xué)習(xí)訓(xùn)練，并在保持性能的同時實現(xiàn)了 1.5-5.4 倍的加速改進(jìn)。與未應(yīng)用量化的全精度訓(xùn)練相比，ActorQ 可以將強化學(xué)習(xí)訓(xùn)練的碳足跡減少 1.9-3.8 倍。

ActorQ 證明量化可以有效地應(yīng)用于強化學(xué)習(xí)的許多方面，從獲得高質(zhì)量和高效的量化策略到減少訓(xùn)練時間和碳排放。隨著強化學(xué)習(xí)在解決現(xiàn)實問題方面繼續(xù)取得長足進(jìn)步，我們有理由相信，使強化學(xué)習(xí)訓(xùn)練實現(xiàn)可持續(xù)發(fā)展將是關(guān)鍵。當(dāng)將強化學(xué)習(xí)擴展到數(shù)千個 CPU 和 GPU 時，即使 50% 的改進(jìn)也會在成本、能源和碳排放方面顯著降低。作者的工作是將量化應(yīng)用于強化學(xué)習(xí)訓(xùn)練以實現(xiàn)高效和環(huán)境可持續(xù)訓(xùn)練的第一步。

作者在 ActorQ 中的量化設(shè)計基于簡單的均勻量化，但實際可以應(yīng)用其他形式的量化、壓縮和稀疏性（如蒸餾、稀疏化等）。未來的工作將考慮應(yīng)用更積極的量化和壓縮方法，這可能會為強化學(xué)習(xí)在訓(xùn)練性能和精度的權(quán)衡上帶來更多的好處。

原文鏈接：

https://ai.googleblog.com/2022/09/quantization-for-fast-and.html