[論文レビュー] QuaRL: Quantization for Sustainable Reinforcement Learning.

Srivatsan Krishnan, Maximilian Lam+5 more|arXiv (Cornell University)|Nov 29, 2021

Reinforcement Learning in Robotics参考文献 35被引用数 4

ひとこと要約

QuaRLは、報酬の低下を伴わずに強化学習における8ビット推論を可能にする新しい量子化手法ActorQを提案する。この手法により、DeepMind Control Suiteにおいてエンドツーエンドの学習速度が1.5倍〜2.5倍向上し、完全精度学習と量子化されたエージェント推論を組み合わせることで実現された。本研究は、適切に適用された場合、量子化の計算上の利点が精度の損失を上回ることを示している。

ABSTRACT

Deep reinforcement learning has achieved significant milestones, however, the computational demands of reinforcement learning training and inference remain substantial. Quantization is an effective method to reduce the computational overheads of neural networks, though in the context of reinforcement learning, it is unknown whether quantization's computational benefits outweigh the accuracy costs introduced by the corresponding quantization error. To quantify this tradeoff we perform a broad study applying quantization to reinforcement learning. We apply standard quantization techniques such as post-training quantization (PTQ) and quantization aware training (QAT) to a comprehensive set of reinforcement learning tasks (Atari, Gym), algorithms (A2C, DDPG, DQN, D4PG, PPO), and models (MLPs, CNNs) and show that policies may be quantized to 8-bits without degrading reward, enabling significant inference speedups on resource-constrained edge devices. Motivated by the effectiveness of standard quantization techniques on reinforcement learning policies, we introduce a novel quantization algorithm, extit{ActorQ}, for quantized actor-learner distributed reinforcement learning training. By leveraging full precision optimization on the learner and quantized execution on the actors, extit{ActorQ} enables 8-bit inference while maintaining convergence. We develop a system for quantized reinforcement learning training around extit{ActorQ} and demonstrate end to end speedups of $>$ 1.5 $ imes$ - 2.5 $ imes$ over full precision training on a range of tasks (Deepmind Control Suite). Finally, we break down the various runtime costs of distributed reinforcement learning training (such as communication time, inference time, model load time, etc) and evaluate the effects of quantization on these system attributes.

研究の動機と目的

• 量子化の計算上の利点が深層強化学習における精度の損失を上回るかどうかを調査すること。
• 多様なRLタスク、アルゴリズム、アーキテクチャにおいて、標準的な量子化技術（PTQおよびQAT）の効果を評価すること。
• 分散型エージェント・リーナーRLを対象とした、新しい量子化に配慮した学習手法ActorQを設計し、8ビット推論を維持しながら収束を保証すること。
• 量子化が与える影響を踏まえ、分散RL学習におけるシステムレベルのランタイムコスト（通信、推論、モデルロード）を分析・最適化すること。
• 量子化を用いることで、ポリシー性能を維持したままエンドツーエンドの学習速度向上を実現すること。

提案手法

• Atマーイ、Gymなどの多様なRL環境、A2C、DDPG、DQN、D4PG、PPOなどのアルゴリズム、MLP、CNNなどのモデルに対して、トレーニング後量子化（PTQ）および量子化に配慮した学習（QAT）を適用する。
• 学習者では完全精度の最適化を、エージェントでは量子化された推論を用いる、新しいトレーニングフレームワークActorQを導入し、8ビットポリシー実行を可能にする。
• 最小限の性能低下でエンドツーエンドの量子化RL学習を実現するためのActorQを基盤とするシステムを設計する。
• 分散RL学習コスト（通信、推論、モデルロード）を分解し、量子化が各要素に与える影響を測定する。
• 標準的な量子化技術を、適切なキャリブレーションを伴って用いることで、学習および推論中のポリシー性能の維持を図る。
• 複数のベンチマークおよびモデル設定において、計算効率とポリシー精度のトレードオフを評価する。

実験結果

リサーチクエスチョン

• RQ1 PTQやQATのような標準的な量子化技術を、顕著な報酬低下を伴わずに深層強化学習ポリシーに効果的に適用できるか？
• RQ2 8ビット量子化ポリシーは、多様なRLタスクおよびアーキテクチャにおいてどの程度の性能を維持できるか？
• RQ3 ActorQのような新しいトレーニングアルゴリズムは、8ビット推論を用いながらも分散RL学習における安定した収束を実現できるか？
• RQ4 量子化は、分散RL学習ワークロードの各要素（通信、推論、モデルロードなど）にどのように影響を与えるか？
• RQ5 ポリシー性能を犠牲にせずに、量子化を用いることでRL学習でどの程度のエンドツーエンドの速度向上が達成できるか？

主な発見

• AtariおよびGym環境において、ポリシーを8ビットに量子化しても累積報酬に劣化が生じず、エッジデバイスでの推論速度向上が顕著に実現された。
• ActorQは、分散エージェント・リーナーRLにおいて安定した収束を実現しながら、8ビット推論を維持し、ポリシー性能を保持した。
• 提案された量子化学習システムを用いることで、DeepMind Control Suiteでエンドツーエンドの学習速度が1.5倍〜2.5倍向上した。
• 量子化により推論およびモデルロード時間が顕著に短縮されたが、分散学習における通信オーバーヘッドへの影響は最小限に抑えられた。
• QATやActorQのような適切な技術を用いることで、量子化の計算上の利点が精度の損失を上回ることが本研究で確認された。
• システムレベルのプロファイリングにより、推論およびモデルロードが学習遅延の主な要因であることが判明し、量子化によってこれらが効果的に短縮された。