[論文レビュー] Single-Shot Pruning for Offline Reinforcement Learning

Samin Yeasar Arnob, Riyasat Ohib+2 more|arXiv (Cornell University)|Dec 31, 2021

Machine Learning and Data Classification被引用数 4

ひとこと要約

本論文は、オフライン強化学習（RL）における学習前に行われる1ショット prune 法を提案する。固定データセットを活用して、スパースで高性能なサブネットワークを特定する。初期化時に SNIP および GraSP prune を適用することで、95% のスパarsity であっても、密なネットワークと同等の性能を達成し、圧縮なしでメモリ使用量を4倍に削減し、リソース制約のあるデバイスでの高速推論を可能にする。

ABSTRACT

Deep Reinforcement Learning (RL) is a powerful framework for solving complex real-world problems. Large neural networks employed in the framework are traditionally associated with better generalization capabilities, but their increased size entails the drawbacks of extensive training duration, substantial hardware resources, and longer inference times. One way to tackle this problem is to prune neural networks leaving only the necessary parameters. State-of-the-art concurrent pruning techniques for imposing sparsity perform demonstrably well in applications where data distributions are fixed. However, they have not yet been substantially explored in the context of RL. We close the gap between RL and single-shot pruning techniques and present a general pruning approach to the Offline RL. We leverage a fixed dataset to prune neural networks before the start of RL training. We then run experiments varying the network sparsity level and evaluating the validity of pruning at initialization techniques in continuous control tasks. Our results show that with 95% of the network weights pruned, Offline-RL algorithms can still retain performance in the majority of our experiments. To the best of our knowledge, no prior work utilizing pruning in RL retained performance at such high levels of sparsity. Moreover, pruning at initialization techniques can be easily integrated into any existing Offline-RL algorithms without changing the learning objective.

研究の動機と目的

• オフラインRLにおける大規模ニューラルネットワークの高い計算コストとメモリ使用量を低減するため、学習前における効率的なモデル圧縮を可能にする。
• 教師あり学習向けに開発された1ショット prune 法が、固定データセットを用いたオフラインRLに効果的に適用可能かどうかを検討する。
• 連続的制御タスクにおける極度のスパarsity（最大95%）下での prune ネットワークの性能を評価する。
• 初期化時の prune が、学習目的を変更せずに、既存のオフラインRLアルゴリズムにスムーズに統合可能であることを示す。
• ロボティクスやエッジAIのようなリアルタイムでリソースが限られたアプリケーションにおける、高速な学習と推論を実現する。

提案手法

• オフラインRL学習の前に、ランダムに初期化されたニューラルネットワークに1ショット prune 法（SNIP および GraSP）を適用する。
• D4RLベンチマークからの固定データセットを用いて、勾配感受性（GraSP）または接続の重要度（SNIP）に基づく prune 準拠を計算する。
• 再訓練や反復的ファインチューニングを一切行わず、初期化時に一度だけネットワークを prune し、スパースサブネットワークを生成する。
• 標準的なオフラインRLアルゴリズム（BCQ、BC）を用いて、連続的制御環境（HalfCheetah、Hopper、Walker2D）で prune ネットワークの性能を評価する。
• 性能、メモリ使用量、レイヤーごとのスパarsity を測定し、prune が一般化性能および効率性に与える影響を評価する。
• PyTorch のスパーステンソル機能を用いて、密なモデルと prune モデルのメモリ要件を格納および比較する。

実験結果

リサーチクエスチョン

• RQ1 初期化時の1ショット prune が、学習前に適用された場合、オフラインRLにおいて性能を維持できるか？
• RQ2 95% などの高いスパarsity であっても、オフラインRLで性能が劣化しない最大のスパarsity 水準は何か？
• RQ3 SNIP と GraSP のような異なる prune 準拠が、オフラインRLにおける prune ネットワークの構造と性能にどのように影響を与えるか？
• RQ4 prune が、オフラインRLモデルにおけるメモリ使用量の削減と推論効率の向上にどの程度寄与するか？
• RQ5 初期化時の prune が、学習目的を変更せずに、既存のオフラインRLアルゴリズムに統合可能か？

主な発見

• 提案手法は、95% のスパarsity であっても、複数の環境で密なネットワークと同等の性能を維持し、帰還報酬に顕著な低下がない。
• GraSP を用いた prune は、SNIP よりも最終層により多くの重みを保持しており、勾配フローを維持するという目的と整合する。
• 圧縮なしで、95% スパースモデルのメモリ使用量は、密なモデルと比較して4倍に削減される。
• 本手法は即挿し可能である：学習目的を変更せず、任意のオフラインRLアルゴリズムに適用可能で、追加の訓練を要しない。
• 実験の結果、バッチサイズや prune の反復回数を変更しても、1回のショット prune ループに比べて性能が向上しない。
• 結果から、固定データセットを用いたオフラインRLが、有効な1ショット prune を可能にすることが示され、オンラインRLではデータ分布の変化のため、これは実現不可能である。