Improving Neural Symbolic Regression with a Diverse Prior and Post-training Alignment

David Otte; Jörg K. H. Franke; Frank Hutter

cs.LGOct 13, 2026

Improving Neural Symbolic Regression with a Diverse Prior and Post-training Alignment

Authors: David Otte, Jörg K. H. Franke, Frank Hutter

Organizations: LMU Munich · Munich Center for Machine Learning (MCML) · University of Freiburg · LAION · ELLIS Institute Tübingen · Prior Labs

Abstract

Symbolic regression is the problem of finding an expression that fits a dataset. Solving it requires searching a combinatorial space of candidate expressions. Neural symbolic regression can amortize that search by training a model once on synthetic data and then predicting expressions for a new dataset. But existing models still underperform genetic programming on real, noisy data. We argue that the deficit lies in the training rather than in the paradigm, and present RASR (Reinforced Amortized Symbolic Regression), a method with two key changes to the training recipe. First, we design a synthetic prior that narrows the gap with real-world data: inputs come from the prior of a tabular foundation model; constants are sampled bottom-up, conditioned on the data; targets carry heteroscedastic and structured noise; and everything is standardized. Second, we split training into two stages. Supervised pre-training maximizes the token-level likelihood of a reference expression. This is followed by reinforcement-learning post-training that aligns the model with its inference-time use by rewarding held-out fit and penalizing expression size. On the black-box track of SRBench 2.0, RASR achieves the best mean test $R^2$ rank of all evaluated methods. Three RASR variants, trained with different size penalties, make up most of the accuracy-complexity Pareto front. On Strogatz, our method recovers more ground-truth expressions than any baseline, and on Feynman it is competitive with the best-performing methods.

Figures & tables

Figure 1: Accuracy-complexity trade-off on the black-box track of SRBench 2.0. Mean expression size (log scale) against mean improvability, the relative RMSE gap to the best method on each dataset (%, lower is better). The dashed line marks the Pareto front, and triangles on the border mark methods outside the plotted range. Three of the five methods on the front are our RASR variants, differing only in the expression-size penalty used during post-training, and the default RASR-94M-L attains the best fit of any evaluated method. Figure J.1 shows every method.

Figure 2: Overview of the synthetic training data generation. Each training pair is produced by sampling an expression skeleton, drawing the inputs together with an auxiliary noise node from the TabICLv2 regression prior, filling the constant placeholders bottom-up conditioned on those inputs, and optionally injecting noise before standardizing and tokenizing. One skeleton is followed throughout: its constants, the target panel of stage 4 and the token row all come from a single draw through the pipeline.

Panel of Figure 3 — Figure 3: The two-stage training process of RASR. Pre-training maximizes the likelihood of the ground-truth prefix sequence of each synthetic pair under teacher forcing. Post-training discards the ground-truth expression entirely: the policy samples a group of candidates per dataset, each candidate’s constants are refined on a training split and scored on a held-out split, and the resulting fit and size form a reward turned into a risk-sensitive group-relative advantage.

Figure 4: Black-box track overview. Methods are ordered by mean test

R^{2}

rank across the twelve datasets (left, lower is better), while the middle and right panels show mean expression size and fitting time on log scales. Markers are means across datasets, bars bootstrapped

95\%

confidence intervals. Amortized methods are flagged with a blue square next to their name. For readability only selected baselines are shown, with ranks computed against all baselines. Figure J.2 shows every method.

Figure 5: Ground-truth recovery on the Strogatz (left) and Feynman (right) tracks. Panels give the accuracy solution rate, the fraction of runs reaching test

R^{2}>0.999

(higher is better), and the mean expression size on a log scale (lower is better). Markers show target-noise levels with

95\%

confidence intervals over problems. Methods are ordered by solution rate averaged over noise levels. Amortized methods are flagged with a blue square next to their name.

Appendix figures & tables23 assets

Supplementary material from the paper’s appendix.

Appendix

Parameter	Value
Distinct variables $n_{v}$	uniform $\{1,\dots,11\}$
Binary ops $n_{b}$	uniform $\{0,\dots,8+n_{v}\}$
Unary ops $n_{u}$	uniform $\{0,\dots,5\}$
Max expression length	$512$ tokens
Constant mantissa	$3$ significant figures
Constant magnitude bound	$[10^{-4},10^{4}]$

Appendix

Table B.1: Skeleton-sampling hyperparameters. The upper bound on the binary-operator count grows with the number of distinct variables

n_{v}

, so that expressions over more variables may be structurally larger.

Component	Setting
Function-type mixture	$4\times$ {mlp, lin, quad, gp, prod}, $1\times$ {tree, disc}
Removed activations	sign, Heaviside, unit-interval, argmax indicators
Min. random-tree depth	$1\rightarrow 3$
Min. discretization centers	$2\rightarrow 5$
Per-column diversity filter	added ( $\geq 8$ clusters; no cluster $>1/3$ of rows)

Appendix

Table B.2: Changes to the published TabICLv2 regression prior. Arrows denote a changed default value.

Operator(s)	$c_{0}$ $(\mu,\sigma)$	$c_{1}$ params	Shift type
$\sin,\cos$	$(0.45,1.0)$	$(0,2.0)$	generic
$\exp$	$(-0.69,0.8)$	$(0,1.0)$	generic
$\tanh$	$(0.0,0.7)$	$(0,1.5)$	generic
$x^{2},x^{3},x^{4},x^{5},x^{1/3},x^{1/5}$	$(0.45,1.0)$	$(0,4.0)$	generic
$\log$	$(0.4,0.8)$	$(1.5,1.0)$	positive
$x^{1/2},x^{1/4}$	$(1.0,0.8)$	$(1.5,1.0)$	positive

Appendix

Table B.3: Per-operator affine-constant distributions.

c_{0}

is log-normal with the listed

(\mu,\sigma)

on the log scale (then normalized by the child’s standard deviation);

c_{1}

follows the listed shift type with parameters

(m,\sigma_{1})

. Powers, roots, and the two trigonometric functions share parameters and are grouped. The signed shift of

\mathrm{inv}

is parameterized in the text.

Parameter	Value
Observations $n$	uniform integer $[8,1024]$
Distractor columns	half-normal, scale $1.5$ , capped by pool size
Clean-dataset fraction	$0.4$
Noise strength $s$	$\mathrm{Exponential}(\text{mean }0.08)$
Gaussian / structured form	$0.5$ / $0.5$
Blend weight $w$ (Gaussian form)	$\mathcal{U}[0,1]$

Appendix

Table B.4: Dataset- and noise-level hyperparameters.

Figure B.1: Statistics of the pre-training prior over

20{,}000

sampled pairs. Top row, left to right: the number of distinct variables

n_{v}

in the expression, the number of distractor features appended to the dataset, and the noise strength

s

(in units of the clean target’s standard deviation

\sigma_{y}

). Bottom row, left to right: the number of constants per expression, the magnitude of the sampled constants (log scale, bounded by the sampler’s own rejection threshold), and the expression size in nodes.

Panel of Figure 11 — Figure C.1: Overview of the model architecture. The table, whose shaded first column holds the target, is encoded cell by cell and refined by a stack of encoder blocks that alternate attention over its two axes. The grids beside them mark the cells attended together, along a row for feature-axis attention and down a column for sample-axis attention. The encoder returns one latent per cell, and the decoder cross-attends to those of the target column alone, shaded in the grid, while generating the expression as a prefix token sequence. Every sublayer is pre-normalized and wrapped in a residual connection.

Model	Params	Model dim.	Enc. layers	Dec. layers	Heads
$24$ M	${\sim}24.37$ M	$384$	$5$	$5$	$6$
$94$ M	${\sim}94.13$ M	$512$	$11$	$11$	$8$

Appendix

Table C.1: Configuration of the two model sizes.

Model	Batch Size	Steps	Datasets	Peak LR	GPUs	Duration
$24$ M	$64$	$250$ k	${\approx}16$ M	$4.6{\times}10^{-4}$	$4\times$ L40S	${\approx}31$ h
$94$ M	$256$	$242$ k	${\approx}62$ M	$2.2{\times}10^{-4}$	$8\times$ H200	${\approx}71$ h

Appendix

Table D.1: Pre-training budget per model size.

Black-box	Rows	Cols	PMLB identifier
SWD	1000	11	1028_SWD
USCrime	47	14	1089_USCrime
lowbwt	31104	10	1193_BNG_lowbwt
echoMonths	17496	10	1199_BNG_echoMonths
vineyard	52	3	192_vineyard
cloud	108	6	210_cloud

Appendix

Table G.1: The two tracks of SRBench 2.0. Adapted from the benchmark ( Aldeia et al., 2025 ) . Datasets are listed under the short names used throughout the paper, with the corresponding PMLB identifier given for the black-box track and the data source, with the additive noise level for the generated datasets, for the first-principles track. The column count includes the target, so a dataset with

c

columns has

c-1

features. The first-principles identifiers are the short names prefixed with first_principles_ .

Method	SWD	USCrime	lowbwt	echoMonths	vineyard	cloud	pm10	apnea1	fri_c0_250_5	fri_c2_1000_10	fri_c0_500_50	environmental
RASR-94M-L	0	0	0	0	0	0	0	0	0	0	0	0
RASR-94M-M	0	0	0	0	0	0	0	0	0	0	0	0
RASR-94M-S	0	0	0	0	0	0	0	0	0	0	0	0
AFP	0	0	0	0	0	0	0	0	0	0	0	0
AFP_EHC	0	0	0	0	0	0	0	0	0	0	0	0
AFP_FE	0	0	0	0	0	0	0	0	0	0	0	0

Appendix

Table I.1: Failed runs per black-box dataset. Number of failed runs out of the

30

runs of each method. These are the runs imputed as described above. Zeros are shown in gray, and the color of the other counts shifts from orange to dark red with the number of failures.

Figure J.1: Accuracy-complexity trade-off with all baselines. As Figure 1 , with every evaluated baseline shown individually instead of only the methods named on the plot.

Figure J.2: Black-box track overview with all baselines and RASR variants. As Figure 4 over all twelve datasets, but with every baseline shown and with RASR-24M-L, RASR-94M-base, RASR-94M-M, and RASR-94M-S alongside RASR-94M-L. Each RASR variant is ranked against the baselines alone, so the baseline ranks match Figure 4 .

Figure J.3: Black-box track overview on unreduced datasets. As Figure J.2 , restricted to the eight datasets RASR neither subsamples nor feature-reduces.

Figure J.4: Accuracy-complexity trade-off on unreduced datasets. As Figure 1 , restricted to the eight datasets RASR neither subsamples nor feature-reduces.

Figure K.1: Test

R^{2}

on the ground-truth tracks. As Figure 5 for the Strogatz (left) and Feynman (right) tracks, but scoring mean test

R^{2}

(higher is better) rather than the accuracy solution rate. The

R^{2}

axis is clipped at

0.5

, so a method with a lower mean at some noise level falls outside it. Amortized methods are flagged with a blue square next to their name.

Figure L.1: First-principles track. Per-problem accuracy-complexity Pareto fronts of median expression size against median test

R^{2}

. Only non-dominated methods appear. The black star is the accepted expression.

Figure M.1: Post-training dynamics of RASR-94M-L. Mean

R^{2}

of the sampled candidates, mean best-of-

256

R^{2}

and mean expression size over post-training steps, on a held-out validation set drawn from the post-training prior.

Variant	Base	Post-trained	Base	Post-trained
	Expression size		Fitting time (s)
RASR-94M-L	37.4	27.5	762	133
RASR-94M-M	37.4	18.3	762	69
RASR-94M-S	37.4	10.3	762	26
RASR-24M-L	39.3	28.5	646	131

Appendix

Table M.1: Effect of post-training on expression size and fitting time on the black-box track. Mean expression size (nodes) and fitting time (seconds), both lower is better, for the pre-trained-only base model and the post-trained model at

K=1024

sampled candidates.

Figure M.2: Accuracy versus test-time compute on the black-box track. Mean improvability (%, lower is better) against the number of sampled candidates

K

for the base model, a model post-trained with GRPO, and the main model post-trained with RS-GRPO, all at the

94

M scale.

Prior	Base	Post-trained
Simple prior	26.56	18.31
Our prior	9.36	9.51

Appendix

Table N.1: Prior ablation on the black-box track (

24

M scale). The simple prior draws inputs from independent uniform distributions and constants from fixed ranges, and applies no standardization. Mean improvability (%, lower is better).

Variant	Improvability
RASR-24M-L	9.51
Skeleton-only	12.09
No init noise	10.57
No refinement	28.33

Appendix

Table N.2: Constant-strategy ablation on the black-box track (

24

M scale). Mean improvability (%, lower is better) for our constants-as-initialization approach against a skeleton-only variant, refinement without initialization noise, and no least-squares refinement. The last two rows are RASR-24M-L with the respective inference step disabled. The skeleton-only variant is pre-trained separately and post-trained with the same size penalty.

Figure O.1: Performance profile on the black-box track. The fraction of problems on which a method reaches at least a given test

R^{2}

, computed per seed and aggregated by the median of the

30

per-seed curves at every threshold, using the SRBench 2.0 post-processing. The legend gives the area under each curve. Higher is better throughout.

Figure O.2: Performance profile with the maximum over seeds. As Figure O.1 , but taking the maximum instead of the median of the per-seed curves, the aggregation the benchmark uses. Methods are split into four panels by their area under the curve so that the curves stay legible, with the methods outside a panel’s band drawn in gray for reference.

Explore similar work

Neural Symbollic Regression Using Deep Learning and Sparse Modelling
Sep 1, 2026Ravi Kumar U, Sumitra SScientific ML Neurosymbolic Learning
Department of Mathematics Indian Institute of Space Science and Technology Thiruvananthapuram, India
Bayesian Symbolic Regression with Entropic Reinforcement Learning
Aug 10, 2026Oussama Boussif, Mohammed Mahfoud, Younesse Kaddar +6Bayesian Inference Symbolic Regression
Mila – Québec AI Institute · Université de Montréal · Independent +5
Symbolic Regression via Latent Iterative Refinement
May 26, 2026Xieting Chu, Sriram Vishwanath, Vijay GaneshLatent Variable Models Amortized Inference
Georgia Institute of Technology

Sep 1, 2026J/K move · Enter open · S save

Neural Symbollic Regression Using Deep Learning and Sparse Modelling

Ravi Kumar U, Sumitra S

Department of Mathematics Indian Institute of Space Science and Technology Thiruvananthapuram, India

Scientific ML Neurosymbolic Learning Symbolic Regression

Open paper AskPDF arXiv

Symbolic Regression (SR) seeks to find succinct mathematical expressions that represent the fundamental relationships within data, providing interpretability and scientific understanding that exceeds that of black-box models. Nevertheless, traditional methods like Genetic Programming face challenges with scalability and are highly sensitive to noise, while sparse regression techniques such as SINDy rely significantly on predetermined feature libraries. In this work, we present a Neural Symbolic Regression (NSR) framework that treats neural networks as functional preconditioners for symbolic discovery. Our approach uses a decoupled pipeline: a neural network first learns a smooth, noise-robust approximation of the target function in an interaction- aware nonlinear feature space. LASSO is then applied to extract sparse, interpretable closed-form expressions. To improve predictive accuracy and symbolic fidelity by integrating distributed hyperparameter optimization with Ray Tune and ASHA scheduling. Experiments on the Nguyen benchmark suite show that our approach consistently outperforms SINDy and non-tuned neural baselines in RMSE, noise robustness, and out-of-distribution generalization. Ablation studies confirm the significance of feature interactions, neural depth, and tuning strategies. In general, this study presents a scalable and understandable neural-symbolic framework, creating a solid link between neural approximation and the discovery of sparse equations for scientific machine learning.

Abstract

Explore similar work

Neural Symbollic Regression Using Deep Learning and Sparse Modelling

Bayesian Symbolic Regression with Entropic Reinforcement Learning

Symbolic Regression via Latent Iterative Refinement