QuaRL:強化学習を量子化して高速化と環境負荷を低減(1/2)

１．QuaRL:強化学習を量子化して高速化と環境負荷を低減(1/2)まとめ

・深層強化学習は逐次意思決定問題大きな進展を続けているが学習時間が長い
・計算負荷が高い事により温室効果ガスの総排出量も多くなってしまう
・量子化を適用する事で性能を維持したまま強化学習の高速化に成功した

２．QuaRLとは？

以下、ai.googleblog.comより「Quantization for Fast and Environmentally Sustainable Reinforcement Learning」の意訳です。元記事の投稿は2022年9月27日、Srivatsan KrishnanさんとAleksandra Faustさんによる投稿です。

量子化と聞くと、量子コンピュータをイメージしてしまうかもしれませんが、全く関係なくて、数字を「たくさんの桁を扱える浮動小数点」ではなく「整数」として扱う事で計算量を減らす技法の事です。

機械学習の世界では「80%の確率で猫である事」と「80.0000001%の確率で猫である事」の違いを追究したくなる事はほぼないので、計算効率を上げるために有効な手法です。

ただし、大規模モデルの場合は精度の違いが学習に影響を及ぼす事もあるようなので、今回の研究は量子化する対象の設計が良かったという事なのかな、と思います。

アイキャッチ画像はstable diffusionの生成でQuaRLとはクラゲやメデューサの意味があるそうなので、トトロにクラゲと遊んでもらいました

深層強化学習(RL:Reinforcement Learning)は、気球の航行、原子物理学、ロボット工学、ゲームなどの実世界の逐次意思決定問題(sequential decision-making problems)の解決において、大きな進展を続けています。

その有望さにもかかわらず、その限界要因の1つは長い学習時間です。複雑で難しいタスクのRL学習を高速化するための現在のアプローチは、数百から数千のコンピューター間で分散学習を行うように規模を拡大していますが、それでもかなりのハードウェア資源を必要とし、RL学習を高価にするとともに、環境への影響を増大させることになります。

しかし、最近の研究では、既存のハードウェアの性能を最適化することで、学習と推論に伴うカーボンフットプリント(carbon footprint：温室効果ガスの総排出量)を削減できることが示されています。

RLも同様のシステム最適化技術により、学習時間の短縮、ハードウェア利用率の向上、二酸化炭素(CO2)排出量の削減などの恩恵を受けることができます。

量子化(quantization)とは、全精度の浮動小数点数(FP32)を低精度の整数(int8)に変換し、低精度数を用いて計算を実行するプロセスです。

量子化により、メモリストレージのコストとバンド幅を節約し、より高速でエネルギー効率の高い計算を実現することができます。量子化は、教師あり学習において、機械学習(ML:machine learning)モデルを計算能力の低い機器内で利用可能にする事や学習の高速化に成功しました。しかし、RLのトレーニングにも量子化を適用する機会が残されています。

そのため、私達は「QuaRL:Quantization for Fast and Environmentally Sustainable Reinforcement Learning」を発表します。この論文では、量子化を適用することで、性能を維持したままRL学習を1.5～5.4倍に高速化するActorQという新しいパラダイムを紹介しています。さらに、全精度で学習する場合と比較して、カーボンフットプリントも1.9～3.8倍と大幅に削減されることを実証しています。

強化学習への量子化の適用

従来のRLでは、学習者ポリシー(learner policy)がアクター(actor、俳優)に適用され、アクターはそのポリシーを使用して環境を探索し、データサンプルを収集します。アクターが収集したサンプルは、学習者ポリシーにより、最初の学習者ポリシーを継続的に改良するために用いられます。

定期的に学習者側で学習されたポリシーがアクターのポリシーを更新するために用いられます。RLのトレーニングに量子化を適用するために、私達はActorQパラダイムを開発しました。

ActorQは上述と同じ順に実行されますが、1つの重要な違いは、学習者からアクターへのポリシーの更新が量子化され、アクターはint8で量子化されたポリシーを用いて環境を探索し、サンプルを収集する点です。

このようにRLの学習に量子化を適用することで、2つの重要な利点があります。

まず、量子化により，学習ポリシーのメモリ使用量を削減することができます。同じピーク帯域幅であれば、学習者とアクター間で転送されるデータは少なくなり、学習者からアクターへのポリシー更新のための通信コストが削減されます。

第二に、アクターは量子化されたポリシーに対して推論を行い、与えられた環境状態に対するアクションを生成します。量子化された推論処理は、完全な精度で推論を行う場合と比較して、はるかに高速に実行されます。

従来の強化学習(左)とActorQ強化学習(右)の概要

ActorQでは、ACME分散RLフレームワークを用いています。

量子化ブロックは、FP32ポリシーをint8に変換する均一な量子化(uniform quantization)を行います。アクターは最適化されたint8計算を用いて推論を行います。

量子化ブロックを設計する際に均一な量子化を用いましたが、他の量子化技術でも均一な量子化の代わりに同様の結果が得られると私達は考えています。アクターによって収集されたサンプルは学習者によってニューラルネットワークのポリシーを学習するために使用されます。学習されたポリシーは定期的に量子化ブロックによって量子化され、アクターに配信されます。