fcmaes - a Python 3 gradient-free optimization library

Hardware-Accelerated Neuroevolution

This tutorial

Shows how to apply CR-FM-NES to train deep neural networks (DNNs) for real-world applications. See also the EvoJax intro.
Shows how to wrap the C++ implementation of CR-FM-NES for use in EvoJax.
Shows how to wrap the C++ implementation of PGPE for use in EvoJax.
Shows how to adapt the Python implementation of CR-FM-NES for use in EvoJax.
Compares the performance (wall time per iteration and convergence) of different evolutionary EvoJax optimization algorithms.

Introduction

The EvoJax environment differs significantly from other applications of fcmaes:

The objective function is the loss-function of a DNN for specific configured weights - the decision variables. CPUs are not well suited for this task. You usually need one or more GPUs or TPUs (Tensor Processing Units).
The number of decision variables is extremely high, up to many thousands. Our CMA-ES implementations become very slow with > 1000 variables. EvoJax addresses this by re-implementing CMA-ES using jax, see cma_jax.py. Running CMA-ES’s heavy matrix operations on a GPU/TPU can speed things up by a factor > 100, which makes CMA-ES a valid choice again.
Parallelization on the level of optimization runs (parallel retry) or parallel objective function evaluation as provided by fcmaes does not help much here, because GPU/TPU resources are the bottleneck. What matters instead is parallelizing the optimization algorithm itself. In this setting, the fcmaes Python implementation of CMA-ES is faster than its C++ counterpart, because numpy parallelizes its matrix operations. And cma_jax.py is faster again than any fcmaes CMA-ES implementation.
CMA-ES has not only performance issues, it may converge slower in case of dependencies between decision variables.

fcmaes provides two algorithms specialized for high-dimensional problems: CR-FM-NES and PGPE. We will compare them with other EvoJax algorithms.

All CPU-based algorithms have one clear disadvantage. Argument and result vectors must be transferred between CPU and GPU/TPU, which adds overhead. We need experiments to see how much this matters in practice.

We will show that both fcmaes C++ implementations of CR-FM-NES and PGPE outperform most EvoJax algorithms, at least on our hardware: an AMD 5950x CPU and an NVIDIA 3090 GPU. This is a typical configuration, and the NVIDIA 3090 is well suited for machine learning because of its 24GB RAM.

Both "jaxified" Python implementations of CR-FM-NES and PGPE, implemented only in EvoJax, perform similarly to the C++ implementations. The main differences are:

The C++ implementations use fewer GPU/TPU resources.
The C++ implementations perform better if the EvoJax optimization task runs only on CPUs.
The C++ implementations have higher accuracy (64 bit). For most benchmarks we did not observe a significant difference in the optimization results. In most cases, 32 bits are sufficient.
The "jaxified" Python implementations may perform better on fast (multi-) GPUs/TPUs if the number of decision variables is very high.

Our C++ implementation of PGPE is derived from its JAX-based EvoJax implementation.

Implementation

The wrapper of C++-CR-FM-NES used as an EvoJax algorithm in our tests is here: fcrfmc.py. The implementation was straightforward. It is derived from existing wrappers such as cma_wrapper.py
The wrapper of C++-PGPE used as an EvoJax algorithm in our tests is here: fpgpec.py. It was derived from existing wrappers such as the one for C++-CR-FM-NES.
The "jaxified" Python implementation of CR-FM-NES is here: crfmnes.py. Since there are no for-loops, there are not many useful applications of jax.jit. We mainly converted np.arrays into jnp.arrays deployed on GPUs/TPUs.

Our benchmark configurations are here:

We used higher iteration counts in our own benchmarks to identify late bloomers.

Remarks

Note that

Our results are partly inconsistent with the ones reported in benchmarks even if you compare only the score after n iterations, which should be independent of the hardware used. Since our results are better, there may have been improvements in the problem formulation itself, or our hyper-parameter settings may have been better.
The benchmarked JAX implementation of CMA-ES is based on _cma.py, see 1604.00772. It is from the same author as the original CR-FM-NES: Masahiro Nomura.
Our C++ implementation of CR-FM-NES is very close to the original Python version, apart from the added ask/tell interface.
We performed experiments with both fcmaes CMA-ES implementations. For high dimensions, they converge similarly to CMA_ES_JAX. Since CMA_ES_JAX is a simpler version of CMA-ES, it is easier to "jaxify" and slightly faster.

Benchmark Results

We compared several algorithms using these benchmarks with the following modifications:

We monitor the progress of the optimization over time, number of fitness evaluations and number of iterations. We use a higher iteration limit to identify late bloomers that improve their rank near the end.
We show both the iteration and evaluation numbers, because EvoJax sometimes profits from higher population size due to parallelization. This effect may increase with multiple / faster GPUs/TPUs.

All tests used the same hardware: an AMD 5950x CPU with 128GB RAM and an NVIDIA 3090 GPU with 24GB RAM.

The tasks used for the comparison are described here: tasks.

Most algorithms create a population for each generation from a random normal distribution vector. Between generations, both the mean and standard deviation vectors are updated using algorithm-specific methods. Some algorithms shape the distribution using a covariance matrix. A full covariance update is expensive, so CRFMNES, FCRFMC and Sep-CMA use a diagonal covariance matrix representation.

Table 1. Algorithm properties
	ARS	PGPE	FPGPEC	CMA-ES	CRFMNES	FCRFMC
population based on random normal distribution vector	+	+	+	+	+	+
symmetric sampling	+	+	+	-	+	+
ADAM for mean update	+	+	+	-	-	-
distribution shaped by covariance matrix	-	-	-	+	+	+
covariance matrix based on diagonal (fast)	-	-	-	-	+	+
rank one update only on ridge structures	-	-	-	-	+	+

Since ARS (Augmented Random Search) is the simplest of all algorithms we additionally show the relative score to ARS_native, its JAX-based implementation. PGPE/FPGPEC is very similar to ARS, and almost as simple, yet it performs much better here. That makes it a serious competitor to the covariance-based algorithms, not only in wall time.

Cartpole Easy

This benchmark is too easy to derive meaningful conclusions. PGPE/FPGPEC and FCRFMC/CRFMNES (our CR-FM-NES wrapper) lead the pack, CMA_ES_JAX is very slow on our hardware. The "iterations" diagram shows that convergence, independent of the hardware used, is also slightly worse than for the other algorithms.

cartpole easy EvoJax Score relative to ARS native

Cartpole Hard

Almost the same result as for Cartpole Easy. PGPE/FPGPEC and FCRFMC/CRFMNES are in the lead, with CMA_ES_JAX lagging behind. All algorithms are > 800 after 1000 iterations, inconsistent with the results reported here.

cartpole hard EvoJax Score relative to ARS native

Brax

Again PGPE/FPGPEC and FCRFMC/CRFMNES lead the pack. Since this task is harder, the pattern is clearer. PGPE is better at the beginning, but FCRFMC/CRFMNES improves faster and takes the lead in the end. OpenES also performs very well here, while CMA_ES_JAX loses again.

brax EvoJax Score relative to ARS native

MNIST

PGPE/FPGPEC is in the lead at first, but both FCRFMC/CRFMNES and OpenES catch up in the end. CMA_ES_JAX improves more slowly in later stages, even in the hardware-independent "iterations" diagram.

mnist EvoJax Score relative to ARS native

Waterworld

Our Waterworld results after 1000 iterations are generally much higher than those reported in benchmarks, but the distance between algorithms is quite consistent. Again PGPE leads at first, then FCRFMC/CRFMNES catches up in the end. OpenES performs strongly, while CMA_ES_JAX lags behind. Note that OpenES shows no further improvement up to 3000 iterations, while the scores of PGPE/FPGPEC and FCRFMC/CRFMNES are still growing.

waterworld EvoJax Score relative to ARS native

Waterworld MA

This benchmark has a very small fixed population size (16). Only PGPE/FPGPEC and FCRFMC/CRFMNES are successful. This is the only benchmark where FCRFMC is faster than CRFMNES, perhaps because of its 64-bit accuracy.

waterworld ma EvoJax Score relative to ARS native

Slimevolley

This final benchmark is clearly dominated by CR-FM-NES. Even OpenES and CMA-ES can surpass PGPE/FPGPEC here. Slimevolley has only 323 decision variables, a small number compared to the other tasks. That seems to move the problem out of typical PGPE territory, although PGPE’s Waterworld MA results contradict that somewhat. We expect CR-FM-NES to be very competitive for "low"-dimensional machine learning tasks in general. Even CMA-ES is back in the game, since its wall-time disadvantage shrinks significantly, especially for the JAX-based implementation. It still trails behind CR-FM-NES. The highest dimensionality where we observed top performance using CMA-ES is multi-UAV task assignment with 104 parameters.

slimevolley EvoJax Score relative to ARS native

Increasing the Optimization Budget

We further increase the optimization budget for the two Waterworld tasks to investigate:

How far did the results listed in benchmarks miss the potential maximal score?
Does the relative maximal scores of the algorithms change if we increase the budget?

Waterworld

We increased the population size to 384 and applied the best algorithms from the first test. The best score is now 16.0625, a lot better than the 11.64 reported at benchmarks.

Table 2. Score/Time for Waterworld task
algorithm	score	time in sec
CRFMNES	16.0625	6948
FCRFMC	16.0625	7781
PGPE	16.0625	7983

All algorithms reach the same score, but the Python/JAX implementation of CR-FM-NES is the fastest, beating both its C++ variant and PGPE. The pattern is now very clear: PGPE is faster at the beginning, while CR-FM-NES is better at the end. This strongly suggests an idea that has worked well in other domains with fcmaes: use different algorithms within the same optimization process. One option would be to use PGPE in the first phase and then pass the resulting PGPE population (or just the best result so far) to CR-FM-NES. That would require extending the solver API so a solver can be initialized with a whole population.

Waterworld MA

Population size cannot be changed here. It remains 16. Increasing the budget reveals: - A score of 3.03125 can be reached, far better than what is reported in benchmarks. - No algorithm can compete with PGPE, which was not visible before.

Summary

All measured tasks show consistent results:

PGPE is slightly superior for lower optimization budgets and shows, together with CRFMNES, that JAX-based optimization algorithms are very competitive in the machine learning domain.
FCRFMC/CRFMNES shows the highest improvement rate as the optimization budget increases and may be an alternative for even more complex tasks. Note that FCRFMC, despite being single CPU-threaded and using very low CPU/GPU/TPU resources, is quite competitive. The overhead of transferring data between CPU and GPU/TPU does not seem to be a decisive disadvantage. Use FCRFMC/FPGPEC (C++) on CPUs and on older operating systems or with a lower number of decision variables.
CRFMNES, the 'jaxified' Python implementation of CR-FM-NES has no significant wall time disadvantage compared to the C++-version FCRFMC, sometimes it is even faster. The reduced accuracy doesn’t harm the convergence. Same holds for PGPE compared to FPGPEC.
OpenES is a valid alternative only slightly behind.
CMA_ES_JAX: Although JAX brings CMA-ES the biggest performance boost for all algorithms, CMA_ES_JAX is still lagging behind. The low convergence of CMA-ES for high dimensional problems makes it the worst choice in the machine learning domain. Note that as the name of my library (fcmaes) indicates, I am a big fan of this algorithm for lower dimensions.
Wrapping a C++ algorithm based on Eigen can perform and converge as fast as the best JAX-based implementations, even single threaded, while saving CPU/GPU/TPU resources, as long as no computationally heavy matrix operations are involved, such as maintaining a full covariance matrix.
Testing revealed that the random generator plays an important role in both wall time and convergence. This led to replacement of the random generator used for all fcmaes algorithms by EigenRand.