Breaking the Simplification Bottleneck in Amortized Neural Symbolic Regression

AI summary

Overview

Research area: Machine learning / symbolic regression / amortized neural program synthesis, with connections to term-rewriting systems and computer algebra.

Technical level: Intermediate to Advanced. The paper assumes familiarity with transformer encoder-decoder architectures, set transformers, and the basics of symbolic regression; the simplification engine itself is explained accessibly, but the architectural and evaluation details require some background.

Scope in one sentence: The paper introduces SimpliPy, a hash-based rule-driven expression simplifier that is roughly 100x faster than SymPy, and uses it to build Flash-ANSR, a transformer framework for amortized symbolic regression that trains on 512M on-the-fly generated expressions and matches genetic-programming baselines such as PySR.

What This Paper Is About

Symbolic regression tries to recover interpretable closed-form equations from observed data. Amortized neural approaches learn a distribution over expressions from millions of synthetic examples, which is far faster at inference time than traditional genetic programming, but they require the training data to be simplified — otherwise redundant forms like x + x and 2x dilute the learning signal. Existing methods either skip simplification entirely (hurting inference efficiency) or rely on general-purpose computer algebra systems like SymPy, whose cost makes data generation far slower than model training. This paper removes that bottleneck so that large-scale, high-quality, on-the-fly expression generation becomes practical.

Key Contributions

1. SimpliPy, a rule-based simplification engine. SimpliPy applies a pre-computed term-rewriting system (TRS) plus a multiplicity-based cancellation procedure to reduce arbitrary expressions to a normalized, length-minimal form. It operates directly on tokenized prefix sequences and achieves up to a ~100x speedup over SymPy at comparable simplification quality.

2. Flash-ANSR, a scalable amortized SR framework. Built on an improved encoder-decoder transformer (pre-RMSNorm decoder, masked RMSSetNorm, RoPE positional encoding, 32-bit IEEE-754 input encoding, softmax sampling, Levenberg-Marquardt constant optimization), it trains on 512M on-the-fly generated and simplified data-expression pairs, with support for higher-dimensional inputs and a broader operator set than prior work.

3. State-of-the-art results under a strict evaluation protocol. Flash-ANSR dominates the inference-time recovery-rate Pareto frontier against amortized baselines (NeSymReS, E2E) on the FastSRB benchmark and matches PySR while producing more concise expressions as the inference budget grows.

4. Rigorous decontamination and evaluation standards. The authors apply symbolic and numeric decontamination of training data against test expressions, demand machine-precision recovery (FVU < 1e-7) rather than lenient R² thresholds, and report test-time compute trade-offs explicitly — practices they find largely absent from prior literature.

Main Findings

• The simplification bottleneck is the dominant obstacle to scaling amortized SR. Because CAS-based simplification is orders of magnitude slower than gradient updates, prior work has had to choose between data quality, data diversity, and dimensionality (typically restricted to D ≤ 3 with a limited operator set).

• Simplification can be amortized offline. SimpliPy enumerates and validates rewriting rules once for expressions up to 7 symbols and 4 variables using numeric equivalence testing in increasing-length order, at a cost of about 100 hours on 32 threads. Runtime simplification then reduces to hash-table lookups and subtree matching, with worst-case online complexity of O(n² · (K·R·P + log n)).

• Normalization improves both learning and inference. Training only on normalized forms gives a one-to-one mapping from function to target sequence, so model capacity and search budget are spent on distinct functional hypotheses rather than on rediscovering algebraic identities or re-optimizing constants for syntactic variants.

• Flash-ANSR outperforms amortized baselines and matches genetic programming. It is substantially more accurate than NeSymReS and E2E on FastSRB, and performs on par with PySR — with the added property that increasing the inference budget yields more concise recovered expressions rather than more complex ones.

• Evaluation in the field is systematically lax. The authors report that most prior work omits confidence intervals, uses lenient success thresholds (e.g., R² > 0.9) that conceal failures, compares methods with vastly different compute budgets, and — with the partial exception of Biggio et al. (2021) — performs no data decontamination, risking overestimated performance.

• Architectural modernization matters at scale. Switching from post-norm to pre-norm, adopting a masked RMSSetNorm for variable-size input sets, and using RoPE instead of sinusoidal positional encodings are presented as necessary to exploit the larger, higher-quality training stream.

Methodology in Plain English

The authors treat simplification as a discrete rewriting problem rather than solving algebra from first principles each time. In an offline phase, they enumerate candidate expression patterns and replacement patterns by increasing length, test numerically whether each candidate pair is genuinely equivalent, and accept a rule only if it strictly shortens the expression, introduces no new variables, and does not duplicate a metavariable more often than it appears on the left-hand side. These three conditions guarantee that the rule set terminates and always reduces size. Rules are organized into operator-specific buckets, with ground rules hashed for constant-time lookup and pattern rules stored as trees for subtree matching.

At runtime, SimpliPy takes the tokenized prefix sequence the generator already produces, alternates rule application with a cancellation pass that merges associative-commutative clusters and removes inverse pairs, iterates up to five times or until nothing changes, sorts commutative operands into a canonical order, and returns the shortest result — falling back to the original if simplification somehow made things longer.

For training, Flash-ANSR generates data in four stages: sample an operator count from a length-exponential prior and build a prefix skeleton with constant placeholders; simplify the skeleton with SimpliPy and reject non-finite results; check the expression against the test set for symbolic and numeric contamination; then render the dataset by sampling inputs and evaluating the expression. This stream feeds a transformer encoder-decoder that predicts the expression sequence given the dataset, with the encoder using a Set Transformer variant to handle variable numbers of data points.

Why This Matters

Impact on research. The paper reframes a data-engineering problem as the central scaling obstacle in amortized symbolic regression, and it supplies both the tool (SimpliPy) and a strict evaluation protocol to the community. It also surfaces a methodological problem — widespread absence of decontamination and inference-time accounting — that likely affects reported numbers in prior work.

Real-world applications:

• Scientific discovery: recovering governing equations from experimental data in physics, chemistry, biology, and climate science, where interpretable formulas matter more than black-box accuracy.

• Automated data science and feature engineering: generating compact, human-readable features or formulas as part of AutoML pipelines.

• Engineering and control: deriving compact system models for controller design, where complex or redundant expressions are operationally unusable.

• Regulated domains: medicine and finance, where models must be auditable and explainable, and where concise symbolic forms are preferable to opaque networks.

Industry relevance. Any organization doing large-scale automated model discovery benefits directly from a 100x faster simplification kernel plus an on-the-fly generation pipeline, since the cost of training data generation is often the binding constraint. The released rule set and code make the approach directly reusable.

Future Directions

• Extending rule coverage. The offline discovery phase is bounded at 7 symbols and 4 variables; scaling it, or porting the engine to broader operator sets and higher dimensions, would widen applicability.

• Closing the remaining gap to genetic programming. Flash-ANSR matches PySR but does not clearly beat it; hybrid schemes that use the neural prior to seed a direct search, or that combine both at inference time, are natural next steps.

• Better constant handling. Constant optimization (Levenberg-Marquardt here) remains a separate, expensive step; tighter integration between structural search and constant fitting could reduce inference cost further.

• Uncertainty and robustness at scale. The paper emphasizes strict evaluation but does not focus on calibrated uncertainty over recovered expressions, nor on behavior under heavy noise or out-of-distribution data — both important for scientific deployment.

Target Audience

Researchers and practitioners in machine learning and scientific computing who work on symbolic regression, equation discovery, or simulation-based inference; engineers building AutoML or interpretable modeling systems; and methodologists interested in evaluation rigor, data decontamination, and test-time compute trade-offs. Readers primarily interested in pure genetic-programming or classical computer algebra will find the paper tangential, though the SimpliPy engine may be independently useful as a fast preprocessor.

Authors’ 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.