xploring GPU-Accelerated Tree Evaluation with JAX (vmap) for Multivariate Symbolic Regression

[cagrisa-creator]

Hello everyone,

I am an independent AI developer working on high-performance architectures for Scientific AI. Inspired by the incredible benchmarks of the symbolic regression community, I have been experimenting with a custom evolutionary engine built entirely from scratch around the JAX ecosystem.

The core architecture utilizes a dynamic tree structure (NoMatematico execution) where mutation, crossover, and node evaluations are fully vectorized across dataset rows using JAX's native vmap on GPU.

The Setup & Benchmark:

Data: Multi-sensor chaotic industrial telemetry dataset (500+ rows, 3 simultaneous sensors influencing a non-linear target variable).

Hardware Parallelization: Heavy usage of vmap to evaluate entire equation populations simultaneously across all data points without standard Python looping bottlenecks.

Results: The engine successfully converged to the underlying physical law, achieving a 99.74% mathematical precision in under 3 minutes on standard GPU instances, maintaining full traceability of the original sensor string names in the final human-readable equation.

Given how critical execution speed is when scaling symbolic regression to high-dimensional datasets, I wanted to open a brief discussion here:

Has anyone here experimented with deep JAX integration or compiler optimizations (jit/vmap) specifically tailored for genetic-tree operations rather than fixed-tensor neural nets?

What are your thoughts on handling highly nested tree logic inside JAX's strict functional paradigm without hitting compilation bottlenecks as the depth grows?

I would highly appreciate any insights, architectural feedback, or thoughts on the potential of GPU-native tree evaluations in this space!

Best regards,

Cesar