Breaking the Simplification Bottleneck in Amortized Neural Symbolic Regression
Abstract
Symbolic regression (SR) aims to discover interpretable analytical expressions that accurately describe observed data. Amortized SR promises to be much more efficient than the predominant genetic programming SR methods, but currently struggles to scale to realistic scientific complexity. We find that a key obstacle is the lack of a fast reduction of equivalent expressions to a concise normalized form. Amortized SR has addressed this with general-purpose Computer Algebra Systems (CAS) like SymPy, but the high computational cost severely limits training and inference speed. We propose SimpliPy, a rule-based simplification engine achieving a 100-fold speed-up over SymPy at comparable quality. This enables substantial improvements in amortized SR, including scalability to much larger training sets, more efficient use of the per-expression token budget, and systematic training set decontamination with respect to equivalent test expressions. We demonstrate these advantages in our Flash-ANSR framework, which achieves much better accuracy than amortized baselines (NeSymReS, E2E) on the FastSRB benchmark. Moreover, it performs on par with state-of-the-art direct optimization (PySR) while recovering more concise rather than more complex expressions with increasing inference budget.
Lay Summary
Scientists often want to discover an unknown mathematical formula that explains a set of measurements: which equation, when evaluated, produces the numbers we observed? This task, called symbolic regression, is hard because the space of possible formulas is enormous. Recent methods train large neural networks to read in measurements and predict a candidate formula, learning from hundreds of millions of synthetic examples. A practical bottleneck is that each training example needs to be cleaned up into a canonical form (for instance, x + x should become 2 * x), and the standard tool for this job (a computer algebra system) is so slow that it dominates training time. We introduce SimpliPy, a lightweight engine that performs this clean-up step roughly a hundred times faster by matching patterns rather than doing full symbolic reasoning. Built on top of SimpliPy, our model Flash-ANSR trains on much larger and more varied formula collections than prior work. Compared to existing methods, it recovers correct formulas more often within the same compute budget, returns shorter and more interpretable expressions, and scales to higher-dimensional inputs and a broader set of mathematical operations.