{"id":"cd115f91-2f40-4636-990d-de3188af9efd","arxiv_id":"2502.03047","paper_version":2,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":6.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":3,"one_line_summary":"A JAX-based genetic programming library using matrix-encoded trees gives parallel GPU evaluation and is demonstrated on symbolic regression, control, and loss learning tasks.","lead":"Kozax is a new genetic programming library built on JAX that represents evolving programs as fixed-size matrices, letting a whole population be evaluated in parallel on GPUs. The paper demonstrates it on equation discovery, control policy search, and learned loss functions, arguing it is faster and more general than existing GP tools.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"Core scalability claim is under-supported: Fig. 6 varies dataset size only, not population size, and the O(N*M*D) tensor representation limits simultaneous large N and large D.","rationale":"I read the central claim as: tensorized matrix representation makes Kozax fast and scalable for large populations/datasets, and this enables a general GP library covering multiple tasks. The most load-bearing condition is the runtime evidence for the scaling claim. The paper has real strengths: seven applications, released code, comparison with PySR and DEAP, and an honest discussion of limitations. The reader's JIT-compatibility concern is valid, but all demonstrated tasks are JAX-native; that is a scope caveat rather than a challenge to the mechanism. A stronger problem is that the scaling evidence itself is thin: Fig. 6a varies only data points and omits population size, matrix size, and generations; Fig. 6b confounds population and dataset size. The tensorized evaluator has O(N*M*D) memory, so large populations and large datasets are not independent. The measured sublinear scaling from 10^4 to 10^5 points suggests overhead dominates small runs, so the crossover cannot be robustly inferred without reporting warmup and error bars. These are fixable with additional experiments, so the verdict stays CONDITIONAL/UNCHANGED rather than REJECT.","tokens_in":11691,"tokens_out":8010,"duration_ms":72404,"concrete_test":"Re-run the Fig. 6a symbolic regression benchmark after a warmup call, reporting population size, max tree size, generations, and GPU memory. Sweep population size N in {100, 1000, 10000} at fixed D=100000, and sweep D in {10^4, 10^5, 10^6} at fixed N=1000, with three seeds each and error bars. If wall time per generation grows superlinearly with N/D or allocation fails for N=10000, D=100000, the 'large populations or datasets' claim must be qualified; if it scales linearly, the concern is resolved.","verdict_should_be":"UNCHANGED","load_bearing_attack":"The paper's strongest claim is that the matrix tree representation yields reduced computation times on large populations or datasets. The supporting runtime analysis (Fig. 6a) sweeps only the number of data points, with population size, generations, matrix size, and the exact fitness function not reported in the text. The second runtime comparison (Fig. 6b) changes population size (250 to 1000) together with dataset size and device, so population scaling is never isolated. Moreover, the vectorized interpreter stores the value of every node for every individual and every data point, so memory is O(N * M * D), plus intermediate operator arrays; the paper gives no batching strategy for the data axis. Thus 'large populations or datasets' cannot be read as 'large populations and datasets' without a memory qualification. The wall-clock numbers also appear to include JIT/compilation overhead: Kozax GPU timings grow from 72.3 s at 10^4 points to only 127.0 s at 10^5 points, a 1.76x increase for 10x data, which suggests small-run overhead rather than pure evaluation cost and makes the crossover point at which GPU wins hardware-dependent. None of this proves the mechanism is wrong, but it means the central scalability claim currently rests on a single unreported-configuration sweep with no error bars.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","summary":"The paper introduces Kozax, a genetic programming library implemented in JAX that represents parse trees as fixed-size matrices, enabling vectorized and JIT-compiled evaluation of whole populations on CPU and GPU. The authors claim that this design makes fitness evaluation scalable to large populations or datasets while retaining flexibility through custom operators, custom fitness functions, constant optimization, and simultaneous evolution of multiple trees. Seven experiments are reported: rediscovery of three physical laws, fully and partially observable Lotka-Volterra dynamics, an acrobot swing-up control policy, and evolved objective functions for training small neural networks, with comparisons against PySR and DEAP. Runtime measurements in Figure 6 are used to support claims of GPU scalability.","tokens_in":11973,"tokens_out":3665,"duration_ms":35256,"significance":"If the central claims hold, Kozax would fill a real gap: a general, tree-based GP framework that runs on GPU and is applicable beyond symbolic regression, including control and meta-learning tasks. The paper has notable strengths: the code and reproduction data are publicly available, the seven experiments cover several problem classes, Table 1 gives a useful feature comparison, and the partially observable Lotka-Volterra and objective-function experiments demonstrate functionality that is difficult to realize in PySR and DEAP. The runtime experiments are a reasonable first step, but, as detailed below, the evidence for the central scalability claim is incomplete, and the claimed generality for 'arbitrary problems' is not fully characterized. These issues are fixable and do not undermine the basic soundness of the library design.","major_comments":[{"comment":"The central scalability claim is not isolated by the reported experiments. Figure 6a varies only the number of data points, while population size, number of generations, tree matrix size, operator set, and the exact fitness function are not reported in the text. Figure 6b changes dataset size, population size, and device simultaneously. As a result, the reader cannot determine whether the speedup comes from vectorization over the population or over the data, and the claim that Kozax scales to 'large populations or datasets' is not separately supported for both axes. Please report the full configuration for every runtime run and include separate sweeps over population size and dataset size, with the memory footprint noted at each setting.","section":"§5.2, Fig. 6a"},{"comment":"The matrix-tree representation stores an intermediate value for every node, for every individual, and for every data point, giving O(N * M * D) memory in the population size N, matrix size M, and dataset size D, in addition to intermediate arrays created by operator evaluation. No batching strategy for the data axis is described. The abstract's phrasing 'large populations or datasets' is therefore accurate only as an 'or', while the Discussion's claim that Kozax 'scales efficiently to large populations, datasets and difficult problems' needs a memory qualification. Please state the memory model explicitly and explain how users should batch or limit D when N and M are large.","section":"§4.2 and §6"},{"comment":"The 'arbitrary problems' generality claim is stronger than the evidence supports. Section 4.1 states that custom operators and fitness functions remain compatible with JIT compilation, but all seven experiments use JAX-native operations or a JAX-based environment (Gymnax for acrobot). The paper does not characterize behavior for Python control flow, data-dependent shapes, or non-JAX external simulators, where the speed advantage of vectorized evaluation disappears. Please add an explicit statement of the JIT/vectorization boundary and, ideally, a small experiment or benchmark showing what happens when a fitness function cannot be JIT-compiled.","section":"§4.1 and §5.1"},{"comment":"The cross-library performance comparison is confounded by per-task tuning. The operator sets are 'empirically determined for each experiment' for Kozax, and the number of constant-optimization steps is applied only to Kozax, while the paper does not state whether PySR and DEAP received comparable tuning or their default configurations. This makes the fitness comparisons in Table 3 a comparison of tuned Kozax against default baselines rather than a controlled benchmark. Please report which hyperparameters were tuned for each library and justify that the chosen settings are representative rather than cherry-picked.","section":"Table 2 and §5.1"}],"minor_comments":[{"comment":"The runtime plots in Figure 6 show averages over three runs but no error bars or confidence intervals; given that the reported GPU times jump from 72.3 s to 127.0 s for a 10x increase in data points, run-to-run variability could affect the claimed crossover point. Please add error bars or per-run values.","section":"§5.2"},{"comment":"The paragraph following Figure 6a repeats the same text twice, with the second copy ending in the same sentence as the first. Please remove the duplicated passage.","section":"§5.2"},{"comment":"There is a typo in the Discussion: 'Kozax runs op GPU' should be 'Kozax runs on GPU'. Also, 'A big advantages' should be 'A big advantage'.","section":"§6"},{"comment":"The objective-function experiment compares only against DEAP; adding a standard loss function (e.g., binary cross-entropy) as a baseline would help interpret the reported 98.3% test accuracy and better position the evolved loss.","section":"§5.1.4"},{"comment":"The partially observable Lotka-Volterra row and the 'loss function (big)' row have '--' entries for PySR and DEAP; a brief note in the table caption explaining why these entries are not applicable would improve readability.","section":"Table 3"}],"recommendation":"major_revision","confidential_remarks":"The manuscript is within scope for a GP/evolutionary-computation venue, and the public code and data are commendable. The main risk is that the scalability and generality claims outrun the evidence in the current runtime analysis; the requested additional sweeps and memory analysis are feasible within the manuscript's scope. The per-experiment tuning issue is common in library comparison papers, but it should be made explicit. I would not reject on the current evidence, but the central claims need strengthening before acceptance."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"The paper does something real: it ships a general-purpose, tree-based GP library in JAX, using a fixed-size matrix representation that lets the whole population evaluate in parallel under JIT. The matrix encoding is a legitimate new contribution—I know of no prior tree-based GP library in JAX that does this and still preserves the standard GP hierarchy. The authors also earn credit for the breadth of demos: law discovery, partially observed dynamical systems, control policy search, and learned loss functions, all with code and raw data available. The comparisons to PySR and DEAP are fair, and the discussion openly lists missing features like simplification and strongly typed GP.\n\nThe soft spots are real but not fatal. The runtime analysis in Fig. 6a sweeps only dataset size while population size is fixed, and Fig. 6b changes population size together with dataset size and device. So the claim that Kozax \"scales to large populations\" is never actually demonstrated in isolation—only large datasets. That matters because the vectorized interpreter stores a value for every node of every individual at every data point, i.e. O(N·M·D) memory, plus intermediate operator arrays. No batching strategy is given. The sublinear GPU scaling from 72s at 10^4 points to 127s at 10^5 points also suggests JIT/launch overhead and a hardware-dependent crossover, not a clean evaluation curve. I'd want that clarified.\n\nThe experimental tuning is also a bit loose: operator sets and constant-optimization budgets are chosen per task, and there is no standard-loss baseline for the objective-function experiment. That weakens the generality claim, though it does not invalidate the central mechanism. The paper does not hide this—hyperparameters are in Table 2—but the narrative leans on \"arbitrary problems\" more than the evidence supports.\n\nAll that said, this is a solid systems paper with reproducible artifacts and an honest limitations section. The library is likely to be used. I'd send it to a serious referee, asking specifically for population-only scaling runs, a memory/batching discussion, and runtime error bars. It is not ready as-is, but it deserves engagement.","headline":"Kozax is a genuinely useful JAX GP library with a clean matrix encoding and honest reporting, but the runtime evidence under-supports the headline scalability claim because population scaling is never isolated and memory costs are unanalyzed.","tokens_in":12468,"tokens_out":1460,"would_cite":false,"duration_ms":14967,"reading_group":"maybe","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":["68W50","68W10"],"pacs":[],"model":"deepseek-v4-flash","headline":"Kozax claims that standard tree-based genetic programming becomes GPU-scalable when every candidate expression tree is encoded as a fixed-size matrix, letting one library handle symbolic regression, control policy search, and…","keywords":["genetic programming","symbolic regression","GPU parallelization","matrix representation of trees","constant optimization","control policy evolution","objective function learning","just-in-time compilation"],"falsifier":"Take a problem whose fitness function calls an external simulator with Python-level control flow and data-dependent shapes, keep population size and generations fixed, and compare Kozax's per-generation wall-clock time on GPU against a sequential genetic programming implementation; if the GPU advantage disappears or runtime grows linearly with population size, the claimed general scalability is refuted for that problem class.","tokens_in":11523,"feed_emoji":"🧬","tokens_out":9257,"duration_ms":69877,"temperature":0.7,"pith_summary":"This paper presents Kozax, a genetic programming library that encodes candidate expression trees as fixed-size matrices and evaluates the whole population at once on CPU or GPU. The authors aim to show that a single framework can cover symbolic regression of natural laws and dynamical systems, symbolic control policy search, and evolving loss functions for neural networks, rather than requiring a separate implementation per task. If the central claim holds, users get GPU-scaled evolution for any fitness function that can be expressed in a vectorized numerical-computing stack, with constant optimization and multi-tree evolution included. The paper reports competitive equation discovery, successful acrobot control, and an evolved loss function that trains a network to 98.3 percent test accuracy, with runtime scaling that favors GPU for large datasets.","feed_headline":"One library evolves equations, policies, and loss functions on GPU","feed_subtitle":"Whole populations of expression trees are evaluated at once, so hard problems scale.","key_machinery":"The central object is the matrix representation of a parse tree: a fixed-size matrix whose rows are nodes, with columns for a function index, child node indices, and a stored computed value; constants and variables are represented as special function indices, empty rows are zeros, and the root is the last row. This representation lets every tree in the population share the same shape, so fitness evaluation, initialization, and reproduction can be written as vectorized operations over the whole population and compiled once, removing the structural variability that normally forces sequential evaluation of differently shaped trees.","core_discovery":"The core claim is that the main computational bottleneck of genetic programming, evaluating many candidate programs of different shapes, can be removed without abandoning standard tree-based GP. Kozax represents every parse tree as a matrix of fixed size: each row stores a function index, the indices of its child rows, and a value slot set to zero before execution; the root sits in the last row, and evaluation solves the matrix row by row so children are computed before parents. Because the matrix is the same size for every candidate, initialization, fitness evaluation, crossover, and mutation can all be vectorized and compiled, which lets a population of thousands run in parallel on GPU hardware. The authors demonstrate the resulting system on seven tasks: rediscovering Kepler's, Newton's, and Bode's laws, recovering Lotka-Volterra dynamics in both fully and partially observed settings, solving the acrobot swing-up with a compact symbolic policy, and evolving an objective function that trains a neural network. On the largest objective-function setting, the evolved loss reaches 98.3% test accuracy on the XOR benchmark.","pith_inferences":["Inference: Because trees are padded to a fixed matrix size, large populations with sparse or shallow trees waste compute on empty rows; a benchmark varying tree depth limits would quantify this overhead, which the paper does not report.","Inference: The reported speed advantage depends on the fitness function itself being vectorizable; a natural stress test is evolving policies in a simulator that must be stepped one episode at a time, where the GPU parallelization would only help if multiple episodes are batched.","Inference: The same matrix representation could be extended to directly evolve recurrent or multi-output programs, since the authors note that higher-dimensional inputs and outputs and automatically defined functions are future work.","Inference: The most direct next application suggested by the results is end-to-end meta-learning on GPU, where many neural networks are trained in parallel under evolved losses, an experiment the paper approximates with batched training in the large XOR setting."],"forward_implications":["With the matrix representation, the entire generational loop can be JIT-compiled and run on GPU, so wall-clock time for fixed population and generations scales down as dataset size grows, as shown in the 100,000-point symbolic regression benchmark.","Kozax can evolve several trees simultaneously with different operator and variable sets, enabling tasks such as partially observed dynamical-systems regression where two equations must be integrated together.","Users can define custom operators and fitness functions while retaining compiled execution, which is what lets the same library switch between equation discovery, control policy search, and objective-function learning.","The two included constant-optimization schemes let users tune how many candidates receive gradient-based or genetic constant fitting, trading compute against solution accuracy.","The objective-function experiments show an evolved loss that trains a small neural network to 98.3% test accuracy on XOR, a task outside the reach of standard symbolic regression tools."],"supporting_citations":[{"why":"Supplies the numerical computing library whose vectorization, just-in-time compilation, and automatic differentiation make the matrix-parallel implementation possible.","marker":"[6]"},{"why":"Provides the specialized symbolic regression system that Kozax is benchmarked against on equation discovery and dataset scaling.","marker":"[13]"},{"why":"Provides the general evolutionary computing baseline that Kozax outperforms on the wider control and objective-function tasks.","marker":"[14]"},{"why":"Defines the genetic programming algorithm and parse-tree representation that Kozax keeps while changing the data layout.","marker":"[21]"},{"why":"Shows a prior JAX-based genetic programming implementation restricted to control policies, the gap Kozax aims to close.","marker":"[31]"},{"why":"Introduces gradient-based optimization of constants inside trees, which Kozax implements as one of its two constant-optimization methods.","marker":"[43]"},{"why":"Establishes parallel fitness evaluation as the central bottleneck in genetic programming, the problem this paper targets.","marker":"[16]"},{"why":"Supplies a JAX-compatible acrobot environment used to evaluate the evolved control policy in parallel.","marker":"[28]"}],"fun_headline_variants":["GPU-accelerated genetic programming for equations and policies","Evolve whole expression-tree populations at GPU speed with Kozax","Kozax: tree-based genetic programming without the eval bottleneck","From Kepler to neural loss: one GP framework scales on GPU"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The load-bearing premise is that a user's fitness function and simulation environment can be expressed inside a just-in-time compiled, vectorized computation graph; the demonstrations all use numerical integrators or JAX-compatible environments, so problems with Python control flow, data-dependent shapes, or external non-differentiable simulators may fall back to slow sequential evaluation, and that boundary is not characterized.","fun_headline_variants_meta":{"raw":{"variants":["GPU-accelerated genetic programming for equations and policies","Evolve whole expression-tree populations at GPU speed with Kozax","Kozax: tree-based genetic programming without the eval bottleneck","From Kepler to neural loss: one GP framework scales on GPU"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.000185,"raw_usage":{"total_tokens":1303,"prompt_tokens":908,"completion_tokens":395,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":524,"completion_tokens_details":{"reasoning_tokens":325}},"tokens_in":524,"tokens_out":395,"duration_ms":4478,"temperature":1.0,"reasoning_tokens":325,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-09T10:05:02.782692+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Take a problem whose fitness function calls an external simulator with Python-level control flow and data-dependent shapes, keep population size and generations fixed, and compare Kozax's per-generation wall-clock time on GPU against a sequential genetic programming implementation; if the GPU advantage disappears or runtime grows linearly with population size, the claimed general scalability is refuted for that problem class.","supporting_citations":[{"cited_title":"JAX: composable transformations of Python+NumPy programs","cited_arxiv_id":null,"evidence_quote":"Supplies the numerical computing library whose vectorization, just-in-time compilation, and automatic differentiation make the matrix-parallel implementation possible."},{"cited_title":"DEAP: Evolutionary algorithms made easy","cited_arxiv_id":null,"evidence_quote":"Provides the general evolutionary computing baseline that Kozax outperforms on the wider control and objective-function tasks."},{"cited_title":"Genetic programming as a means for programming computers by natural selection","cited_arxiv_id":null,"evidence_quote":"Defines the genetic programming algorithm and parse-tree representation that Kozax keeps while changing the data layout."},{"cited_title":"Naturally interpretable control policies via graph-based genetic programming","cited_arxiv_id":null,"evidence_quote":"Shows a prior JAX-based genetic programming implementation restricted to control policies, the gap Kozax aims to close."},{"cited_title":"Faster genetic programming based on local gradient search of numeric leaf values","cited_arxiv_id":null,"evidence_quote":"Introduces gradient-based optimization of constants inside trees, which Kozax implements as one of its two constant-optimization methods."},{"cited_title":"Fast genetic programming on GPUs","cited_arxiv_id":null,"evidence_quote":"Establishes parallel fitness evaluation as the central bottleneck in genetic programming, the problem this paper targets."},{"cited_title":"gymnax: A JAX-based Reinforcement Learning Environment Library","cited_arxiv_id":null,"evidence_quote":"Supplies a JAX-compatible acrobot environment used to evaluate the evolved control policy in parallel."}],"review_version":1}