RobertTLange/evosax

何を解決するか

evosax は、ニューロエボリューションを実装する際に複雑な非同期処理を扱う必要を排除することを目的としています。研究者は、分散実装に通常伴うオーバーヘッドなしに、GPU や TPU などの高スループットハードウェアアクセラレータを活用して、Evolution Strategies (ES) に対して大規模なベクトル化を実現できます。

仕組み

このライブラリは JAX に基づいて構築されており、XLA コンパイルと JAX の変換プリミティブ（jit、vmap、lax.scan など）を活用して高いパフォーマンスを実現しています。標準化された ask-eval-tell API ループを採用しています：

1. Ask：戦略が評価対象の候補解（集団）を生成します。
2. Eval：ユーザーがその集団の適応度を評価します。
3. Tell：ユーザーが適応度の結果を戦略に返し、内部状態を更新して次世代の集団を進化させます。

対象ユーザー

進化的アルゴリズムやニューロエボリューションに取り組む研究者や実務家で、高性能な JAX ネイティブライブラリを使って現代のハードウェアアクセラレータに実験をスケーリングしたい方を対象としています。

特徴

• 広範なアルゴリズムライブラリ：CMA-ES や Differential Evolution といった古典的手法から、OpenAI-ES や Diffusion Evolution といった最新のアプローチまで、30 以上を実装しています。
• JAX ネイティブ：JAX の変換を完全にサポートし、効率的なスケーリングと XLA コンパイルを可能にします。
• 標準化された API：SimpleGA などの集団ベースのアルゴリズムを含め、異なる戦略間で一貫した ask-eval-tell ループを採用しています。
• 高スループット：従来の分散システムのオーバーヘッドを回避するように、ハードウェアアクセラレータに最適化されています。