毎秒240万ゲームフレームを処理し、AIトレーニングコストを80％削減、GoogleがRL並列コンピューティングフレームワークをオープンソース化

この記事はAI新メディアQuantum Bit（公開アカウントID：QbitAI）より許可を得て転載しています。転載の際は出典元にご連絡ください。

金持ちにとって最も辛いのは、ハードウェアをたくさん持っていても、1+1=2 の効果を達成できないことです。

これは、AI トレーニングにおける並列コンピューティングの場合です。1,000 個の GPU があっても、単一マシン トレーニングの 1,000 倍の効果を達成することはできません。

最近、資金に困っていない Google は、数千台のマシンで AI トレーニングを簡単に実行できる SEED RL フレームワークをオープンソース化しました。その効果は従来の方法よりも最大 4 倍優れています。

十分な資金があれば、クラウドで大規模な並列コンピューティングを実行することで、トレーニング コストを 80% 節約できます。大規模な AI モデルのトレーニング コストが簡単に数百万に達する可能性があることを考慮すると、これは本当に相当な額です。

サッカーゲームをプレイするための AI のトレーニングに関しては、SEED RL は 1 秒あたり 240 万フレームの速度で処理できます。 60fps で計算すると、1 秒あたり 11 時間分のゲーム映像を処理することに相当します。

SEED RLアーキテクチャ

前世代の分散強化学習エージェント IMPALA には、アクターと学習者の 2 つの部分を含むアーキテクチャがあります。

アクターは通常、CPU 上で実行され、環境内で実行されるステップと、モデルで推論を実行して次のアクションを予測するステップを反復します。

アクターは推論モデルのパラメータを頻繁に更新し、十分な数の観測を収集した後、観測とアクションの軌跡を学習者に送信して学習者を最適化します。

このアーキテクチャでは、学習者は数百台のマシンからの分散推論入力を使用して GPU 上でモデルをトレーニングします。

しかし、IMPALA には多くの欠点があります。

1. ニューラル ネットワークの推論に CPU を使用するのは非効率的です。そして、モデルが大きくなり、計算量が増加するにつれて、問題はますます深刻になります。

2. アクターと学習者間のモデルパラメータの帯域幅がパフォーマンスのボトルネックになります。

3. リソース利用効率が低い。アクターは環境タスクと推論タスクを交互に実行しますが、これら 2 つのタスクの計算要件が異なるため、同じマシン上のリソースを最大限に活用することが困難です。

SEED RL アーキテクチャはこれらの欠点を解決します。アクターは、GPU や TPU などの AI ハードウェア アクセラレータで推論を実行できるため、モデルのパラメータと状態がローカルに保持され、推論が高速化され、データ転送のボトルネックを回避できます。

IMPALA アーキテクチャとは対照的に、SEED RL のアクターは環境内でのみアクションを実行します。学習者は、複数のアクターからのデータのバッチを使用して、ハードウェア アクセラレータ上で集中的に推論を実行します。

SEED RL は、gPRC フレームワークのネットワーク ライブラリを使用して、各環境ステップで学習者に観測値を送信する際のレイテンシを低く抑えます。これにより、SEED RL は単一のマシンで 1 秒あたり最大 100 万クエリを達成できるようになります。

学習者は数千のコアに拡張でき、アクターの数も数千のマシンに拡張できるため、1 秒あたり数百万フレームのトレーニング速度を実現できます。

SEED RL は、V-trace と R2D2 という 2 つの最先端のアルゴリズムを使用します。

V-trace は、サンプリングされたアクションからアクションの分布を予測する役割を担い、R2D2 は、アクションの予測された将来の値に基づいてアクションを選択する役割を担います。

V-trace はポリシー勾配ベースの方法であり、IMPALA によって最初に採用されました。 Actor と Learner は非同期で実行されるため、V-trace は非同期アーキテクチャで適切に機能します。

2 番目のアルゴリズムは R2D2 です。これは、DeepMind が Atari ゲームでの強化学習エージェントを 4 倍向上させ、52 のゲームで人間のパフォーマンスを上回るために使用した Q 学習手法です。

このアプローチにより、RNN を使用しながら、Q 学習アルゴリズムを大規模なハードウェア上で実行できるようになります。

実験結果

Google は、DeepMind が最近開発したオープンソースのサッカー ゲーム プロジェクトである Google Research Football のベンチマーク テストを実施しました。

64 個の Cloud TPU コアを使用することで、1 秒あたり 240 万フレームのデータ転送速度が達成され、従来の最先端の分散型 IMPALA と比べて 80 倍の向上が実現しました。

IMPALA では同じ速度を達成するには 14,000 個の CPU が必要でしたが、SEED RL では 4,160 個の CPU しか使用されませんでした。同じ速度を得るには、IMPALA は SEED RL の 3 ～ 4 倍の CPU を必要とします。

並列コンピューティング用にハードウェア アクセラレータを最適化することで、モデルのサイズを安全かつ大胆に増やすことができます。

上記のフットボール ゲーム タスクでは、モデルのサイズと入力解像度を増やすことで、これまで解決できなかったいくつかの問題を解決でき、モデルのトレーニングの効率が大幅に向上します。

ポータル

論文の宛先:

https://arxiv.org/abs/1910.06591

GitHub アドレス:

https://github.com/google-research/seed_rl