Abstract
Neuroevolution is a subfield of artificial intelligence that
leverages evolutionary algorithms to generate and optimize
artificial neural networks. This technique has proven to be
successful in solving a wide range of complex problems
across various domains. The NeuroEvolution of Augment-
ing Topologies (NEAT) is one of the most renowned al-
gorithms in neuroevolution. Characterized by its openness,
it starts with minimal networks and progressively evolves
both the topology and the weights of these networks to opti-
mize performance. However, the acceleration techniques em-
ployed in prevailing NEAT implementations typically rely on
parallelism on CPUs, failing to harness the rapidly expand-
ing computational resources of today. To bridge this gap,
we present NEATAX, an innovative framework that adapts
NEAT for execution on hardware accelerators. Built on top
of the JAX, NEATAX represents networks with varying topo-
logical structures as tensors with the common shape, facili-
tating efficient parallel computation using function vectoriza-
tion. Upon rigorous testing across various tasks, we found
that NEATAX has the capacity to shrink the computa-
tion time from hours or even days down to a matter of
minutes. These results demonstrate the potential of NEATAX
as a scalable and efficient solution for neuroevolution tasks,
paving the way for the future application of NEAT in more
complex and demanding scenarios. NEATAX is available at
https://github.com/WLS2002/neatax.