SearcharxivSearch

arXiv · 2502.03047

Kozax: Flexible and Scalable Genetic Programming in JAX

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.

Explore related subjects

Keep this discovery

Explore connections, maps & timelines

BibTeXRIS

Sigur de Vries, Sander W. Keemink, Marcel A. J. van Gerven. 2025-02-05. Kozax: Flexible and Scalable Genetic Programming in JAX. https://arxiv.org/abs/2502.03047

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

Breaking the Central Bias: Spatially Partitioned Experts for Coordinate-Based Neuroevolution

Evolvable-Substrate HyperNEAT (ES-HyperNEAT), a bio-inspired indirect encoding that determines neuron placement and connection weights from spatial coordinates, exhibits a failure mode on MNIST as a diagnostic benchmark. Because input pixels map to a coordinate space centered at the origin, evolved networks converge on a small central cluster of input pixels, a spatial-concentration bias; prior work observed only 21% mean accuracy in this regime. Is this bias an optimization artifact or an architectural ceiling? Inspired by Mixture-of-Experts (MoE) principles, we partition the input into non-overlapping spatial segments, each assigned to a separately evolved specialist network. With 13 such experts, this design reaches 43% mean accuracy, a 106% relative improvement over the baseline. The architectural gain does not depend on data-driven aggregation: equal-weighted averaging, which uses no validation data, already yields a 70% improvement; the gain comes from partitioning, not the weighting. Receptive-field analysis shows the mechanism: partitioning forces evolution to discover features across the entire image, expanding active pixel coverage from 4% to 79%. Absolute accuracy stays below gradient-trained baselines, but the relative gain points to central bias, not the evolutionary search. Two tools are designed to generalize beyond MNIST: a receptive-field diagnostic for silent input-coverage collapse, and a spatial-partitioning remedy that restores coverage.

cs.NE

A Bio-Plausible Visual Neural Network for Locust-Inspired Collision Perception

Locust visual systems have long served as an important biological paradigm for studying looming perception and collision avoidance. Numerous computational models have successfully reproduced the selective responses of Lobula Giant Movement Detector (LGMD) neurons to approaching objects, thereby emulating the fundamental functionality of the biological system. However, existing models remain limited in biological plausibility and robustness when operating in complex and dynamic visual environments. To address these limitations, we propose a biologically plausible neural network for locust-inspired looming detection. The proposed framework incorporates a spatially isotropic sampling strategy that mimics the ommatidial organization of the locust compound eye, a population-voting mechanism inspired by population coding in biological neural systems, and leaky integrate-and-fire neuronal dynamics to replace conventional sigmoid-based membrane activation. Systematic experiments on synthetic stimuli, laboratory sequences, and real-world driving scenarios demonstrate that the proposed model improves robustness under challenging visual conditions while preserving computational efficiency and enhancing biological fidelity. These results highlight the potential of biologically grounded neural computation for robust and efficient collision perception.

cs.NE

Structural Fusion of Bayesian Networks with Limited Treewidth Using Genetic Algorithms

This paper introduces an evolutionary computation approach for consensus in structural Bayesian Network (BN) fusion under the constraint of limited treewidth. The consensus BN aims to reconcile multiple input BNs into a single one that retains key structural features present in the original networks. Treewidth, a graph-based parameter associated with computationally tractable inference, is utilized to restrict the complexity of the resulting network. A genetic algorithm is proposed to look for a BN that codifies as much information about the unrestricted fusion as possible while ensuring the treewidth restriction. Experimental evaluation demonstrates the genetic algorithm's ability to obtain consensus BNs with limited treewidth, providing a valuable tool for aggregating information from diverse sources while returning a computationally actionable model.

cs.NE