evosax: JAX-based Evolution Strategies
Robert Tjarko Lange

TL;DR
evosax is a JAX-based library that enables scalable, hardware-accelerated evolutionary strategies, facilitating efficient black-box optimization on GPUs and TPUs with a flexible, modular API.
Contribution
The paper introduces evosax, a novel JAX-based library that implements 30 evolutionary algorithms optimized for modern hardware accelerators, enhancing scalability and performance.
Findings
Enables direct execution of algorithms on GPUs and TPUs.
Supports automatic vectorization and parallelization across devices.
Provides a modular, easy-to-use API for evolutionary optimization.
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…
Peer Reviews
No public reviews on file for this paper yet. If you reviewed it on a platform where reviews are public (OpenReview, ICLR, NeurIPS, ICML), you can paste yours below so the community can read it here.
Code & Models
Videos
No videos yet. Explain this paper in a talk, walkthrough, or lecture? Add one.
Taxonomy
TopicsEvolutionary Algorithms and Applications · Metaheuristic Optimization Algorithms Research · Advanced Multi-Objective Optimization Algorithms
MethodsLib
