[论文解读] 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 上实现可扩展的黑箱优化。

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.

研究动机与目标

• 解决进化计算中现代硬件加速器（GPU/TPU）长期未被充分利用的问题，该领域传统上依赖于基于 CPU 的并行计算。
• 通过将进化策略迁移到 JAX 的高性能计算栈，实现高效、可扩展的黑箱优化（BBO）。
• 通过硬件并行执行、自动向量化和模块化设计，推动下一代进化算法的发展。
• 通过提供可扩展的工具和编码方式，支持多样化的优化工作负载，包括神经进化、元学习和架构搜索。
• 通过使 ES 天然兼容现代加速器硬件，弥合基于梯度的深度学习与无导数优化之间的差距。

提出的方法

• 在 JAX 中实现进化策略，以利用其 JIT 编译、自动微分以及跨设备的硬件并行执行能力。
• 提供模块化的 ask-evaluate-tell API，支持与任意目标函数（包括随机 rollout 和神经网络评估）无缝集成。
• 支持 30 种不同的 ES 算法，包括基于有限差分的梯度方法（如 OpenAI-ES、ARS）、自然进化策略（如 SNES、xNES）以及 CMA-ES 变体。
• 通过 JAX 的 `pmap` 和 `vmap` 实现硬件并行执行，仅用一行代码即可在多个随机种子或超参数设置下批量运行 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 环境下的连续控制任务中表现出一致的性能。
• 在四个神经进化任务上的实验表明，JAX 编译的 ES（如 OpenAI-ES、SNES、Sep-CMA-ES）在仅需极少代码修改的情况下，实现了稳定且可复现的性能。
• 通过 `pmap` 实现的一行并行化，用户可在不改变架构的前提下，跨不同随机种子或超参数并行运行多个 ES 实例。
• 适应度塑造、参数重塑和间接编码工具显著简化了 ES 与深度学习流水线及复杂架构的集成。
• 该库的模块化设计使高级研究方向（如学习型进化策略、子种群管理与策略集成）成为可能，这些方向此前受限于计算瓶颈。

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.