EvoJAX: Hardware-Accelerated Neuroevolution (2202.05008v2)

Abstract: Evolutionary computation has been shown to be a highly effective method for training neural networks, particularly when employed at scale on CPU clusters. Recent work have also showcased their effectiveness on hardware accelerators, such as GPUs, but so far such demonstrations are tailored for very specific tasks, limiting applicability to other domains. We present EvoJAX, a scalable, general purpose, hardware-accelerated neuroevolution toolkit. Building on top of the JAX library, our toolkit enables neuroevolution algorithms to work with neural networks running in parallel across multiple TPU/GPUs. EvoJAX achieves very high performance by implementing the evolution algorithm, neural network and task all in NumPy, which is compiled just-in-time to run on accelerators. We provide extensible examples of EvoJAX for a wide range of tasks, including supervised learning, reinforcement learning and generative art. Since EvoJAX can find solutions to most of these tasks within minutes on a single accelerator, compared to hours or days when using CPUs, our toolkit can significantly shorten the iteration cycle of evolutionary computation experiments. EvoJAX is available at https://github.com/google/evojax

• The paper introduces EvoJAX, which accelerates the entire neuroevolution pipeline by leveraging JAX's just-in-time compilation and hardware parallelism.
• It implements a unified evolutionary framework with vectorized task evaluations and modern ML optimizers for efficient neural network training.
• Experimental results on tasks like MNIST, Cart-Pole, and generative art show speed-ups of up to 20 times over CPU-based methods, underlining its practical impact.

EvoJAX: Hardware-Accelerated Neuroevolution

The paper presents EvoJAX, a toolkit for conducting neuroevolution experiments efficiently using state-of-the-art hardware accelerators. EvoJAX is built on the JAX library, allowing for scalable implementations of neuroevolution algorithms across multiple TPU/GPUs. This paper provides an in-depth look at the design and application of EvoJAX, emphasizing its capability to significantly reduce the computational time involved in training neural networks through evolutionary methods.

Key Contributions

EvoJAX addresses the need for hardware-accelerated neuroevolution by implementing the entire evolutionary pipeline using JAX, enabling faster execution by leveraging just-in-time compilation and other advanced features. The toolkit includes support for modern ML optimizers, offering more effective optimization strategies through gradient estimation-based approaches. EvoJAX adopts a "Single-Program, Multiple-Data" (SPMD) model, which reduces resource duplication and enhances computational throughput.

The authors highlight several innovative components:

• Global Policy: A single computational graph is shared for policy evaluation, optimizing resource use and aligning with modern deep learning frameworks.
• Vectorized Tasks: Task evaluations are grouped in a vectorized format, allowing for effective parallel computing.
• Device Parallelism: JAX's support for device parallelism is utilized, allowing EvoJAX to scale nearly linearly with hardware.

Experimental Validation

The paper provides comprehensive examples to showcase EvoJAX's capabilities across diverse tasks:

1. Supervised Learning: Demonstrated with MNIST classification and a seq-to-seq learning task, EvoJAX exhibits significant speed-ups, achieving high test accuracies rapidly.
2. Control Tasks: Includes tasks like Cart-Pole and robotic control, where EvoJAX efficiently handles undetermined task steps and demonstrates automatic training complexity.
3. Novel Tasks: Addresses more creative and complex scenarios such as multi-agent WaterWorld and generative art, illustrating EvoJAX's capacity to foster innovative solutions.

The presented experimental results show that EvoJAX outperforms traditional CPU-based neuroevolution methods, often achieving speed-ups in the range of 10 to 20 times on modest hardware accelerators.

Implications and Future Directions

EvoJAX represents a considerable advancement in hardware-accelerated evolutionary computation, providing researchers with tools to iterate swiftly on new ideas. The toolkit's design facilitates the exploration of novel policy architectures, neuroevolution algorithms, and complex, real-world tasks. Its modular design ensures easy integration and extendibility, further catalyzing research in neuroevolution.

The practicality of EvoJAX culminates in its potential to significantly enhance experimental efficiency, creating new opportunities for innovation in AI and machine learning research. Future work may focus on expanding the range of algorithms and examples included with EvoJAX, thereby promoting broader adoption and fostering further advancements in the field.

In conclusion, EvoJAX stands out as a powerful resource for researchers seeking to leverage the computational potential of hardware accelerators for neuroevolutionary experiments, thus providing a valuable contribution to the ongoing development of evolutionary computation methodologies.