[論文レビュー] GST: Group-Sparse Training for Accelerating Deep Reinforcement Learning

Juhyoung Lee, Sangyeob Kim+3 more|arXiv (Cornell University)|Jan 24, 2021

Reinforcement Learning in Robotics参考文献 31被引用数 9

ひとこと要約

本稿では、ブロック巡回行列圧縮と報酬に配慮したプルーニングを組み合わせることで、深層強化学習（DRL）の学習を高速化する新規な重み圧縮手法、Group-Sparse Training（GST）を提案する。GSTは初期学習段階から高い圧縮比を維持し、報酬フィードバックに基づいてスパarsityを動的に調整することで安定した学習を実現する。Mujoco環境でTD3を用いた実験では、反復的プルーニングと比較して平均圧縮比が25–41.5ポイント高い結果を得た。

ABSTRACT

Deep reinforcement learning (DRL) has shown remarkable success in sequential decision-making problems but suffers from a long training time to obtain such good performance. Many parallel and distributed DRL training approaches have been proposed to solve this problem, but it is difficult to utilize them on resource-limited devices. In order to accelerate DRL in real-world edge devices, memory bandwidth bottlenecks due to large weight transactions have to be resolved. However, previous iterative pruning not only shows a low compression ratio at the beginning of training but also makes DRL training unstable. To overcome these shortcomings, we propose a novel weight compression method for DRL training acceleration, named group-sparse training (GST). GST selectively utilizes block-circulant compression to maintain a high weight compression ratio during all iterations of DRL training and dynamically adapt target sparsity through reward-aware pruning for stable training. Thanks to the features, GST achieves a 25 \%p $\sim$ 41.5 \%p higher average compression ratio than the iterative pruning method without reward drop in Mujoco Halfcheetah-v2 and Mujoco humanoid-v2 environment with TD3 training.

研究の動機と目的

• リソース制限のあるエッジデバイスにおけるDRL学習の長時間化という課題に対処すること。
• 反復的プルーニング手法における初期学習段階での低圧縮比という課題を克服すること。
• 固定スパarsityスケジュールに代えて報酬に配慮した動的スパarsity調整を導入することで、DRL学習の安定性を向上させること。
• 大規模なモデルパラメータによるメモリバンド幅のボトル neck を軽減することで、モバイルおよびエッジデバイス上での効率的なDRL学習を可能にすること。
• さまざまなDRLベンチマーク（Atari、Google Research Football、ImageNet分類タスクなど）にわたる一般化の有効性を示すこと。

提案手法

• 初期学習段階でブロック巡回圧縮を選択的に適用し、初期段階からの高い圧縮比を維持する。
• 報酬に配慮したプルーニングを用いて、学習報酬の傾向に基づき、目標スパarsityを動的に調整し、固定スケジュールによる不安定化を回避する。
• スパarsity増加のタイミングを制御するためのフェーズシフトパラメータ $ S_{\text{shift}} $ を導入し、ブロック巡回からプルーニングされた重みへの滑らかな移行を実現する。
• 圧縮と性能のバランスを取るために、3種類のブロックサイズ変換手法——投影、block4フレンドリーな block2、block4フレンドリーな block4——を実装する。
• 最初および最後の層を除き、ネットワークパラメータの94–99%にGSTを選択的に適用することで、重要な特徴表現を保持する。
• ブロック変換時のパラメータ値の変化を最小限に抑えるフレンドリーな行列変換を導入し、学習の安定性を確保する。

実験結果

リサーチクエスチョン

• RQ1 ブロック巡回圧縮を反復的プルーニングと効果的に組み合わせることで、DRL学習全体にわたり高い圧縮比を維持できるか？
• RQ2 固定スパarsityスケジュールと比較して、報酬に配慮した動的スパarsityスケジューリングはDRLにおける学習安定性を向上させるか？
• RQ3 GSTは最終的な性能を損なわずに、反復的プルーニングと比較してより高い平均圧縮比を達成できるか？
• RQ4 GSTは、連続的制御（Mujoco）、離散的制御（Atari）、および複雑な環境（Google Research Football）を含む多様なDRLベンチマークで効果を発揮するか？
• RQ5 GSTはDRL以外のタスク、例えば画像分類タスクに対しても一般化可能であり、より広範な適用可能性を示せるか？

主な発見

• TD3を用いたMujoco Halfcheetah-v2およびHumanoid-v2環境において、GSTは反復的プルーニングと比較して平均圧縮比が25–41.5ポイント高く、報酬低下なしに達成した。
• Mujoco Humanoid-v2では、block4フレンドリーな block2手法と $ S_{\text{shift}} = 0.25 $ を用いることで、最大66.5%の圧縮比を達成し、報酬に影響を与えないことを確認した。
• A2Cを用いたAtari Breakoutでは、$ B=2 $ および $ S_{\text{shift}} = 1.0 $ の条件下で71.9%の平均圧縮比を達成し、報酬低下なしに実現した。
• PPOを用いたGoogle Research Footballでは、$ B=4 $ および $ S_{\text{shift}} = 1.0 $ の条件下で73.6%の平均圧縮比を達成し、報酬損失なしに実現した。
• ResNet-32を用いたCIFAR-10では、68.2%の圧縮率と91.4%の精度（ベースライン比0.8%低下）を達成し、分類タスクへの一般化を示した。
• AlexNetを用いたImageNetでは、62.9%の圧縮率と55.8%の精度（ベースライン比0.4%低下）を達成し、モデルアーキテクチャの多様性に対しても堅牢であることを確認した。