# Kozax: Flexible and Scalable Genetic Programming in JAX

arXiv: https://arxiv.org/abs/2502.03047 · HTML: https://mlsanity.com/paper/2502.03047/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 13 of its 48 references that are in the map.
- First posted on arXiv 5 Feb 2025.

## Builds on (most cited first, 13 of 13)

- [Contemporary Symbolic Regression Methods and their Relative Performance](https://mlsanity.com/paper/2107.14351/) (2021)
- [Explainable Reinforcement Learning: A Survey](https://mlsanity.com/paper/2005.06247/) (2020)
- [Meta-Learning via Learned Loss](https://mlsanity.com/paper/1906.05374/) (2019)
- [Interpretable Machine Learning for Science with PySR and SymbolicRegression.jl](https://mlsanity.com/paper/2305.01582/) (2023)
- [Interpretable Policies for Reinforcement Learning by Genetic Programming](https://mlsanity.com/paper/1712.04170/) (2017)
- [evosax: JAX-based Evolution Strategies](https://mlsanity.com/paper/2212.04180/) (2022)
- [EvoJAX: Hardware-Accelerated Neuroevolution](https://mlsanity.com/paper/2202.05008/) (2022)
- [EvoX: A Distributed GPU-accelerated Framework for Scalable Evolutionary Computation](https://mlsanity.com/paper/2301.12457/) (2023)
- [Enabling Population-Level Parallelism in Tree-Based Genetic Programming for GPU Acceleration](https://mlsanity.com/paper/2501.17168/) (2025)
- [Accelerating Genetic Programming using GPUs](https://mlsanity.com/paper/2110.11226/) (2021)
- [Faster GPU Based Genetic Programming Using A Two Dimensional Stack](https://mlsanity.com/paper/1601.00221/) (2016)
- [GSGP-CUDA -- a CUDA framework for Geometric Semantic Genetic Programming](https://mlsanity.com/paper/2106.04034/) (2021)
- [Fast and Efficient Local Search for Genetic Programming Based Loss Function Learning](https://mlsanity.com/paper/2403.00865/) (2024)

## Built on by (most cited first, 4 of 4)

- [Neuromorphic Intelligence](https://mlsanity.com/paper/2509.11940/) (2025)
- [Efficient Constant Optimization for Symbolic Regression with GPU-Accelerated Tree-Based Genetic Programming](https://mlsanity.com/paper/2609.03352/) (2026)
- [Evolving Layer-Specific Scalar Functions for Hardware-Aware Transformer Adaptation](https://mlsanity.com/paper/2605.14047/) (2026)
- [MorphingNet: Symbolic Regression Where the Network Is the Formula](https://mlsanity.com/paper/2610.19942/) (2026)

## Abstract (arXiv)

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.