[Paper Review] Neural Symbolic Regression that Scales

Luca Biggio, Tommaso Bendinelli+3 more|arXiv (Cornell University)|Jun 11, 2021

Model Reduction and Neural NetworksPhysics and Astronomy31 references27 citations

TL;DR

A Transformer-based symbolic regressor pre-trained on millions of generated equations, which then predicts equation skeletons from input-output pairs and fits constants, enabling scalable, data-driven symbolic regression that improves with more data and compute.

ABSTRACT

Symbolic equations are at the core of scientific discovery. The task of discovering the underlying equation from a set of input-output pairs is called symbolic regression. Traditionally, symbolic regression methods use hand-designed strategies that do not improve with experience. In this paper, we introduce the first symbolic regression method that leverages large scale pre-training. We procedurally generate an unbounded set of equations, and simultaneously pre-train a Transformer to predict the symbolic equation from a corresponding set of input-output-pairs. At test time, we query the model on a new set of points and use its output to guide the search for the equation. We show empirically that this approach can re-discover a set of well-known physical equations, and that it improves over time with more data and compute.

Motivation & Objective

• Motivate symbolic regression as a data-driven, scalable alternative to hand-designed methods.
• Introduce a large-scale pre-training framework for symbolic regression using generated equations.
• Learn a parametric symbolic regressor that predicts equation skeletons from input-output data.
• Evaluate how pre-training size, test-time compute, and number of data points affect performance.
• Demonstrate scalability and robustness across multiple symbolic regression benchmarks.

Proposed method

• Pre-train a Transformer to map a set of input-output pairs to an equation skeleton (with placeholders for constants).
• Represent equations in prefix notation and use a placeholder symbol for constants to be fitted later.
• Train on hundreds of millions of procedurally generated equations to learn a prior over symbolic expressions.
• At test time, encode the given data, sample skeleton candidates via beam-search, and fit constants with a nonlinear optimizer (e.g., BFGS).
• Use Set Transformer encoder for permutation-invariant handling of variable-sized input sets; decode with a standard Transformer decoder.
• Evaluate by comparing predicted equations against ground-truth on various datasets (AIF, SOOSE, Nguyen) with in-distribution and out-of-distribution metrics.

Experimental results

Research questions

• RQ1 Can a symbolic regressor learn from large-scale pre-training to outperform hand-designed symbolic regression methods?
• RQ2 How does pre-training data size affect symbolic regression performance at test time?
• RQ3 Does the method scale with the number of input-output pairs and input variables?
• RQ4 How does the approach perform on in-distribution versus out-of-distribution equation sets?
• RQ5 What are the trade-offs between test-time compute and accuracy?

Key findings

• NeSymReS improves as pre-training data increases, outperforming state-of-the-art baselines across datasets under similar compute budgets.
• At test time, NeSymReS achieves higher accuracy with less compute compared to baselines, particularly on AI-Feynman (AIF).
• Performance improves with more test-time input-output pairs and remains robust when the number of data points varies widely.
• The method scales linearly with the number of input-output points and input dimensions due to Set Transformer design.
• Pre-training on large equation distributions enables the model to generalize to unseen equations (SOOSE) and longer expressions.
• NeSymReS can be substantially faster (orders of magnitude) than baselines on CPU for comparable accuracy.