TensorNEAT: A GPU-accelerated Library for NeuroEvolution of Augmenting Topologies

TL;DR

TensorNEAT addresses NEAT’s scalability bottlenecks by introducing a tensorization strategy that encodes heterogeneous networks as fixed-shape tensors and leverages JAX-based GPU/TPU acceleration to parallelize evolutionary operations across populations. The approach enables full GPU-accelerated NEAT (including CPPN and HyperNEAT) with seamless integration to Gym, Brax, and gymnax, yielding up to 500x speedups over NEAT-Python. Key contributions include the tensorized encoding/operations, two-stage tensorized inference, multi-GPU support, visualization tools, and a comprehensive benchmark suite. This work significantly improves the practicality of large-scale neuroevolution for robotics and control tasks, and outlines a path toward distributed deployment and incorporation of additional NEAT variants.

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. To address these limitations, this paper introduces TensorNEAT, a GPU-accelerated library that applies tensorization to the NEAT algorithm. Tensorization reformulates NEAT's diverse network topologies and operations into uniformly shaped tensors, enabling efficient parallel execution across entire populations. TensorNEAT is built upon JAX, leveraging automatic function vectorization and hardware acceleration to significantly enhance computational efficiency. In addition to NEAT, the library supports variants such as CPPN and HyperNEAT, and integrates with benchmark environments like Gym, Brax, and gymnax. Experimental evaluations across various robotic control environments in Brax demonstrate that TensorNEAT delivers up to 500x speedups compared to existing implementations, such as NEAT-Python. The source code for TensorNEAT is publicly available at: https://github.com/EMI-Group/tensorneat.

Figures (11)

• 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 grids symbolize attributes specific to the NEAT algorithm, such as historical markers and enabled flags. The yellow grids denote the attributes of the network's nodes and connections, including biases, weights, and activation functions. The gray grids 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: Visualization methods for networks in TensorNEAT. The left side shows the topology diagram, and the right side represents the network formulations. This network has three input nodes, one output node, and three hidden nodes. In the topology diagram, yellow nodes represent input nodes, white nodes represent hidden nodes, and the blue node represents the output node. In the formula, the values of the input nodes are denoted by $i_0$, $i_1$, and $i_2$; the hidden nodes by $h_0$, $h_1$, and $h_2$; and the output node by $o_0$.