Kozax: Flexible and Scalable Genetic Programming in JAX

TL;DR

This work addresses the need for a general, scalable genetic programming framework capable of evolving symbolic expressions across diverse scientific tasks. It introduces Kozax, a JAX-based GP library that vectorizes population evaluation, runs on CPU or GPU, and supports custom fitness functions, operators, and multi-tree evolution. A matrix-based tree representation enables parallel execution of GP operations and scalable constant optimization, showing competitive performance with PySR on symbolic regression and advantages over DEAP in flexibility and speed, especially with GPU acceleration. The results across law discovery, dynamical systems, control policies, and objective-function learning demonstrate Kozax's potential for scalable, interpretable white-box modeling in scientific computing, with room for further enhancements such as simplification during evolution and higher-dimensional or multi-output trees.

Abstract

Genetic programming is an optimization algorithm inspired by evolution which automatically evolves the structure of interpretable computer programs. The fitness evaluation in genetic programming suffers from high computational requirements, limiting the performance on difficult problems. Consequently, there is no efficient genetic programming framework that is usable for a wide range of tasks. To this end, we developed Kozax, a genetic programming framework that evolves symbolic expressions for arbitrary problems. We implemented Kozax using JAX, a framework for high-performance and scalable machine learning, which allows the fitness evaluation to scale efficiently to large populations or datasets on GPU. Furthermore, Kozax offers constant optimization, custom operator definition and simultaneous evolution of multiple trees. We demonstrate successful applications of Kozax to discover equations of natural laws, recover equations of hidden dynamic variables, evolve a control policy and optimize an objective function. Overall, Kozax provides a general, fast, and scalable library to optimize white-box solutions in the realm of scientific computing.

Figures (7)

• Figure 1: We introduce Kozax as a general framework for genetic programming, utilizing JAX for fast and parallelizable computation. It allows for highly flexible problem and fitness function specifications, making Kozax applicable to diverse and complex tasks, including symbolic regression of dynamical systems, control policy optimization and objective function evolution for training neural networks.
• Figure 1: Genetic programming algorithm
• Figure 2: Overview of trees and reproduction in genetic programming. (a) An example of the parse tree representation used in genetic programming. (b) An example of crossover on a pair of trees, where the blue and orange subtrees are swapped. (c) An example of mutation on a tree, where the operator in blue is replaced with the new operator in orange.
• Figure 3: Step-by-step mapping of a tree to a matrix. The node added in the tree and the corresponding row in the matrix are marked blue at every step. The references to child nodes are added to the relevant rows once the child node itself has been added to the tree. The matrix is inverted, as the execution starts with the leaf nodes.
• Figure 4: Processing of a matrix in Kozax. The matrix is iteratively solved with the input values: $y_1 = \frac{\pi}{2}$ and $y_2 = 1$. A blue row in the matrix corresponds to a blue node in the tree, where the computed value of the node is stored in the last column. The final value of the tree is obtained by taking the stored value in the last row.