Tensorized NeuroEvolution of Augmenting Topologies for GPU Acceleration

TL;DR

This work tackles the scalability bottleneck of NeuroEvolution of Augmenting Topologies (NEAT) by introducing a tensorization method that encodes variable-topology networks as uniformly shaped tensors, enabling parallel evaluation across an entire population. Built atop JAX, TensorNEAT vectorizes NEAT operations and leverages hardware acceleration (GPUs/TPUs) to deliver large speedups, reportedly up to 500x over NEAT-Python in robotics benchmarks. The library supports NEAT variants such as CPPN and HyperNEAT and integrates with environments like Gym, Brax, and gymnax, while offering flexible interfaces and problem templates. The results demonstrate substantial practical impact for scalable neuroevolution in complex control tasks, with future work aimed at distributed computation and broader variant support, including DeepNEAT and CoDeepNEAT.

Abstract

The NeuroEvolution of Augmenting Topologies (NEAT) algorithm has received considerable recognition in the field of neuroevolution. Its effectiveness is derived from initiating with simple networks and incrementally evolving both their topologies and weights. Although its capability across various challenges is evident, the algorithm's computational efficiency remains an impediment, limiting its scalability potential. In response, this paper introduces a tensorization method for the NEAT algorithm, enabling the transformation of its diverse network topologies and associated operations into uniformly shaped tensors for computation. This advancement facilitates the execution of the NEAT algorithm in a parallelized manner across the entire population. Furthermore, we develop TensorNEAT, a library that implements the tensorized NEAT algorithm and its variants, such as CPPN and HyperNEAT. Building upon JAX, TensorNEAT promotes efficient parallel computations via automated function vectorization and hardware acceleration. Moreover, the TensorNEAT library supports various benchmark environments including Gym, Brax, and gymnax. Through evaluations across a spectrum of robotics control environments in Brax, TensorNEAT achieves up to 500x speedups compared to the existing implementations such as NEAT-Python. Source codes are available at: https://github.com/EMI-Group/tensorneat.

Figures (7)

• Figure 1: Illustration of the fundamental acceleration principle underlying our tensorization method. In traditional network computations, each network individually processes its inputs. By contrast, with our tensorization method, a single computation suffices to derive the batched output for all networks.
• Figure 2: Illustration of the network encoding process. In our method, networks with varying topological structures are transformed into uniformly shaped tensors, enabling the representation of the entire network population as batched tensors. The orange cells symbolize attributes specific to the NEAT algorithm, such as historical markers and enabled flags. The yellow cells denote the attributes of the network's nodes and connections, including biases, weights, and activation functions. The gray cells represent sections filled with NaN to ensure consistent tensor shapes.
• Figure 3: An illustration of tensorized network operations. It demonstrates how traditional network operations are converted into equivalent tensor operations during tensorization. Changes from the original format are highlighted in red, underscoring the modifications made within the tensor.
• Figure 4: Illustration of tensorized network inference process. It transforms a feedforward network's topology into tensors and the subsequent calculation of node values for network inference. The network is first encoded into node and connection tensors, which are then ordered and expanded for processing. Finally, values are calculated through connection and node functions to produce the output.
• Figure 5: Average fitness and wall-clock time against generation.