訓練速度提高最多5.4倍，谷歌提出RL訓練新范式ActorQ

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

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

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

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

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

量化應用于強化學(xué)習訓練

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

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

傳統 RL 訓練概述（左）和 ActorQ RL 訓練（右）。

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

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

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

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

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

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

量化減少碳排放

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

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

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

比較使用 FP32 和 int8 的訓練之間的碳排放。X 軸標度標準化為 FP32 的碳排放量，用紅色條表示。ActorQ 減少了碳排放。

結論和未來(lái)方向

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

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

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

原文鏈接：

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