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 Augmenting
Topologies (NEAT) is one of the most renowned algorithms
in neuroevolution. Characterized by its openendedness, it
starts with minimal networks and progressively evolves both
the topology and the weights of these networks to optimize
performance. However, the acceleration techniques employed
in prevailing NEAT implementations typically rely on par-
allelism 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