[論文レビュー] evosax: JAX-based Evolution Strategies

Robert Tjarko Lange|arXiv (Cornell University)|Dec 8, 2022

Evolutionary Algorithms and Applications被引用数 5

ひとこと要約

evosax は JAX をベースにしたライブラリであり、JAX の JIT コンパイル、自動ベクトル化、マルチデバイス並列処理を活用することで、ハイパフォーマンスでハードウェアアクセラレーションを実現する進化戦略（ES）を可能にします。30 種類以上の ES アルゴリズム—有限差分法、自然進化戦略、遺伝的アルゴリズムを含む—をサポートしており、GPU や TPU でスケーラブルなブラックボックス最適化を実現し、1 行の並列化コマンドで実行できます。

ABSTRACT

The deep learning revolution has greatly been accelerated by the 'hardware lottery': Recent advances in modern hardware accelerators and compilers paved the way for large-scale batch gradient optimization. Evolutionary optimization, on the other hand, has mainly relied on CPU-parallelism, e.g. using Dask scheduling and distributed multi-host infrastructure. Here we argue that also modern evolutionary computation can significantly benefit from the massive computational throughput provided by GPUs and TPUs. In order to better harness these resources and to enable the next generation of black-box optimization algorithms, we release evosax: A JAX-based library of evolution strategies which allows researchers to leverage powerful function transformations such as just-in-time compilation, automatic vectorization and hardware parallelization. evosax implements 30 evolutionary optimization algorithms including finite-difference-based, estimation-of-distribution evolution strategies and various genetic algorithms. Every single algorithm can directly be executed on hardware accelerators and automatically vectorized or parallelized across devices using a single line of code. It is designed in a modular fashion and allows for flexible usage via a simple ask-evaluate-tell API. We thereby hope to facilitate a new wave of scalable evolutionary optimization algorithms.

研究の動機と目的

• 進化計算分野において、長年にわたり CPU による並列処理に依存してきた現代のハードウェアアクセラレータ（GPU/TPU）の活用が不十分であるという問題に対処すること。
• JAX のハイパフォーマンス計算スタックに進化戦略を移植することで、効率的でスケーラブルなブラックボックス最適化（BBO）を実現すること。
• ハードウェア並列実行、自動ベクトル化、モジュラー設計を通じて、次世代の進化アルゴリズムの実現を促進すること。
• 拡張可能なユーティリティとエンコーディングを提供することで、ニューロエボリューション、メタラーニング、アーキテクチャ探索などの多様な最適化ワークロードをサポートすること。
• 勾配ベースのディープラーニングと勾配フリー最適化の間のギャップを埋め、ES を現代のアクセラレータハードウェアにネイティブに適合させること。

提案手法

• JAX を用いて進化戦略を実装し、JIT コンパイル、自動微分、デバイス間のハードウェア並列実行を活用すること。
• 任意の目的関数（確率的ロールアウトやニューラルネットワーク評価を含む）とシームレスに統合できるモジュラーな ask-evaluate-tell API を提供すること。
• 有限差分勾配ベース手法（例：OpenAI-ES、ARS）、自然進化戦略（例：SNES、xNES）、CMA-ES 変種を含む、30 種類の異なる ES アルゴリズムをサポートすること。
• JAX の `pmap` と `vmap` を用いて、1 行のコードで複数のランダムシードやハイパーパramータ設定に対して、バッチ化された ES 実行をハードウェア並列処理で実現すること。
• パラメータのリシェーピング（例：フラットベクトルからニューラルネットワーク重みへの変換）、フィットネスの形状調整（z スコア、ランクベース）、間接エンコーディング（例：ハイパーネットワーク、ランダム行列射影）のためのユーティリティを提供すること。
• JAX の集約演算（例：`pmean`）を活用したメモリ最適化とデバイス間のサブポピュレーション配布により、大規模な問題における効率的な実行を実現すること。

Figure 1. Comparison of different evolution strategies (OpenAI-ES, PGPE, SNES and Sep-CMA-ES) implemented by evosax on four continuous control tasks using the Brax library (Freeman et al . , 2021 ) . The simulations use a population size of 256 and the results are averaged over 3 independent runs an

実験結果

リサーチクエスチョン

• RQ1 JAX の機能的プログラミングスタックを用いることで、現代のハードウェアアクセラレータ（GPU/TPU）上で進化戦略を効率的に高速化できるか？
• RQ2 JAX の自動ベクトル化とハードウェア並列処理は、ブラックボックス最適化のスループットとスケーラビリティをどの程度向上できるか？
• RQ3 JAX コンパイル済み進化戦略は、多様な最適化タスクにおいて、従来の CPU による実装と比較して、性能と安定性に優れているか？
• RQ4 evosax のモジュラーで合成可能なユーティリティは、メタラーニングされた進化戦略や間接エンコーディング分野における新たな研究方向性を可能にするか？
• RQ5 ハードウェアアクセラレーションは、高次元かつ微分不能な最適化問題における ES の収束性とサンプル効率にどのような影響を与えるか？

主な発見

• evosax は JAX の JIT コンパイルとマルチデバイス並列処理を活用することで、進化戦略の完全なハードウェアアクセラレーションを実現し、大規模な ES 実行におけるウォールクロック時間の短縮を実現しました。
• 30 種類以上の多様な進化戦略（有限差分法、自然進化、CMA-ES 変種を含む）をサポートしており、Brax 環境を用いた連続制御タスクにおいて一貫したパフォーマンスを発揮しました。
• 4 つのニューロエボリューションタスクにおける実験では、JAX コンパイル済み ES（例：OpenAI-ES、SNES、Sep-CMA-ES）が、最小限のコード変更で安定的かつ再現性のあるパフォーマンスを達成しました。
• `pmap` を用いた 1 行の並列化により、アーキテクチャの変更なしに、異なるランダムシードやハイパーパramータ設定に対して複数の ES インスタンスを実行できるようになりました。
• フィットネスの形状調整、パラメータのリシェーピング、間接エンコーディングユーティリティは、ES をディープラーニングパイプラインや複雑なアーキテクチャに統合するプロセスを大幅に簡素化しました。
• ライブラリのモジュラー設計により、学習された進化戦略、サブポピュレーション管理、戦略アンサンブルといった高度な研究分野が、従来の計算ボトル neck に制限されずに実現可能になりました。

Figure 2. Comparison of different ES (OpenAI-ES, PGPE, SNES and Sep-CMA-ES) across population sizes on four continuous control tasks. The results are averaged over 3 independent runs.