PaperlayerView exact revision on arXiv

Source-linked AI summary

Transformer-based Planning for Symbolic Regression

Parshin Shojaee, Kazem Meidani, Amir Barati Farimani, Chandan K. Reddy

arXiv:2303.06833v5cs.LGcs.AI

TL;DR

Symbolic regression needs to balance equation fit and complexity, while pretrained transformers can optimize token-level objectives instead of discovery goals. TPSR adds MCTS planning and performance feedback to transformer decoding; the paper reports stronger fitting-complexity trade-offs, extrapolation, and noise robustness than existing methods.

  • Problem

    Pretrained SR transformers trained on token-level sequence losses can generate equations suboptimal for fitting accuracy and complexity.

    Source: §1 · p. 2

  • Method

    TPSR combines pretrained SR transformers with MCTS lookahead planning and a reward balancing fitting accuracy and equation complexity.

    Source: §1 · p. 2

  • Results

    TPSR outperforms state-of-the-art SR baselines across benchmark datasets, with higher fitting accuracy while maintaining lower equation complexity.

    Source: §1 · p. 3

  • Takeaways & Limitations

    TPSR controls equation complexity without fine-tuning and improves noise robustness and extrapolation to unseen data.

    Source: §5 · p. 10

    TPSR takes longer to infer than beam search and sampling because it searches and incorporates performance feedback during decoding.

    Source: E Discussion and Future Work · p. 23

Abstract

from arXiv · show

Symbolic regression (SR) is a challenging task in machine learning that involves finding a mathematical expression for a function based on its values. Recent advancements in SR have demonstrated the effectiveness of pre-trained transformer-based models in generating equations as sequences, leveraging large-scale pre-training on synthetic datasets and offering notable advantages in terms of inference time over classical Genetic Programming (GP) methods. However, these models primarily rely on supervised pre-training goals borrowed from text generation and overlook equation discovery objectives like accuracy and complexity. To address this, we propose TPSR, a Transformer-based Planning strategy for Symbolic Regression that incorporates Monte Carlo Tree Search into the transformer decoding process. Unlike conventional decoding strategies, TPSR enables the integration of non-differentiable feedback, such as fitting accuracy and complexity, as external sources of knowledge into the transformer-based equation generation process. Extensive experiments on various datasets show that our approach outperforms state-of-the-art methods, enhancing the model's fitting-complexity trade-off, extrapolation abilities, and robustness to noise.

1 Introduction

Symbolic regression searches a combinatorial space of equations while balancing fitting, complexity, and generalization; pretrained transformers accelerate generation but can prioritize token-level similarity over discovery objectives. TPSR adds lookahead planning and performance feedback to transformer decoding, improving the fitting-complexity trade-off against baselines.

1 Introduction

  • Symbolic regression searches for interpretable equations but faces a combinatorial optimization space and must balance fitting, complexity, and generalization.Input-data quality is another challenge.Source: §1 · p. 1
  • Genetic Programming uses population-based heuristic search but can converge slowly and incur computational expense by restarting from scratch for each dataset.Source: §1 · p. 1
  • Pretrained transformers generate equations in a single forward pass, but token-level cross-entropy training can yield equations suboptimal for fitting accuracy and complexity.Beam search and sampling use transformer logits but receive no performance feedback during candidate generation.Source: §1 · p. 2
  • TPSR applies Monte Carlo Tree Search lookahead to pretrained transformer decoding, using feedback to guide equation-sequence generation while retaining pretrained-model logits.The approach is described as model-agnostic and compatible with non-differentiable objectives.Source: §1 · p. 2
  • A new reward balances equation fitting accuracy and complexity to promote an effective trade-off in generated equations.Source: §1 · p. 2
  • TPSR consistently outperforms state-of-the-art SR baselines, achieving higher fitting accuracy while maintaining lower equation complexity across benchmark datasets.The conclusion also reports improved fitting-complexity trade-offs, extrapolation, and noise robustness.Source: §1 · p. 3
  • The work also highlights TPSR’s extrapolation and noise robustness relative to the baseline and studies the impact of model components.Source: §1 · p. 3

2 Related Work

Prior symbolic-regression approaches include search without learned priors, pretrained transformer generation, and planning for sequence objectives. TPSR combines pretrained transformer priors with MCTS to guide equation generation using fitting and complexity feedback.

Symbolic Regression without Learned Priors

  • GP, reinforcement-learning, hybrid GP-RL, and standalone MCTS approaches lack large-scale pretraining priors and restart search for new datasets.This makes inference slow compared with methods that use learned priors.Source: §2 · p. 3

Pre-trained Transformers for Symbolic Regression

  • Pretrained SR transformers generate equations from learned priors, but token-level sequence losses can leave fitting accuracy and complexity suboptimal.TPSR uses lookahead planning with fitting and complexity feedback during generation.Source: §2 · p. 3

Planning in Sequence Generation

  • Planning in NLP and code generation targets sequence-level objectives beyond pretraining, motivating MCTS-guided transformer generation for equation discovery.SR objectives include fitting accuracy and minimizing complexity.Source: §2 · p. 3

3 Methodology

TPSR uses MCTS to guide transformer-based equation generation with feedback on fitting accuracy and complexity. Its reward balances those objectives, and caching reduces redundant computations.

3.1 Preliminaries

  • Transformer SR models generate equations autoregressively from encoded observations, but beam search and sampling rank candidates without feedback on equation-discovery objectives.Prior models use token-level cross-entropy training and may optimize constants and rank completed candidates by fitting metrics such as R2.Source: §3.1 · p. 4

3.2 MCTS-Guided Equation Generation

  • MCTS treats partial token sequences as states and next tokens as actions, using transformer token probabilities and P-UCB to balance promising and under-explored paths.P-UCB combines action return, transformer next-token probability, and state visit counts; β adjusts the exploration-exploitation trade-off.Source: §3.2 · pp. 4–5
  • TPSR expands only with the transformer’s top-k next tokens, completes partial equations using beam search, then backpropagates the highest simulated reward toward the root.Limiting expansion avoids arbitrary vocabulary actions, while beam-search simulations produce completed candidates for reward evaluation.Source: §3.2 · pp. 5–6

3.3 Reward Definition

  • The reward scores completed equations using fitting accuracy and sequence-length complexity, with λ controlling the balance between fit and simpler solutions.Constants are optimized before reward calculation; sequence length in prefix notation represents complexity, and higher λ favors less complex equations.Source: §3.3 · p. 6

3.4 Efficient Implementation with Caching

  • Top-k caching reuses next-token values for states, while sequence caching reuses completed equations matching a partial state to reduce redundant generation.The paper states that both caching strategies improve efficiency without compromising performance.Source: §3.4 · p. 6

4 Experiments

Experiments evaluate TPSR with an E2E transformer backbone on SRBench and in-domain data, testing benchmark performance, extrapolation, noise robustness, caching, and search parameters. TPSR improves fitting while balancing equation complexity, and its caching mechanisms reduce inference time.

4.1 Datasets

  • Evaluation covers Feynman, Strogatz, and Black-box benchmarks plus 400 in-domain synthetic functions, with continuous features and d ≤ 10.The in-domain functions and input points are generated using the same approach as the E2E backbone’s pre-training data.Source: §4.1 · p. 7

4.2 Evaluation Metrics

  • The evaluation uses R2, tolerance accuracy Accω, and equation complexity, with complexity measured by expression-tree node count.For Accω, the worst 5% of predictions are discarded to reduce outlier sensitivity.Source: §4.2 · p. 7

4.3 (RQ1) Effectiveness of TPSR

  • TPSR outperforms baseline methods in fitting across datasets while producing equations of comparable or lower complexity.The reported fitting-performance gap is larger on SRBench than on in-domain data; on the latter, E2E has an advantage because its training data use the same generation approach.Source: Table 1 · p. 7, §4.3 · p. 8
  • λ controls the fitting-complexity trade-off: λ = 0 favors complex equations, while λ = 0.1 is recommended for balance and overfitting mitigation.Higher λ reduces complexity with a slight fitting-performance reduction; λ = 1 still significantly outperforms baseline methods in fitting.Source: §4.3 · p. 8
  • On Feynman and Black-box datasets, TPSR improves on E2E and balances fitting, complexity, and inference time against SRBench methods.TPSR has the highest fitting performance on Black-box data, complexity no greater than E2E, and lower inference time than GP-based methods.Source: Figure 5 · p. 8, §4.3 · p. 9

4.4 (RQ2) Extrapolation and Robustness

  • Complexity control mitigates overfitting and supports extrapolation, while λ = 1 can impair fitting; MCTS-guided decoding also drops less under noise than E2E.The extrapolation test varies input scale σ; the noise test compares mean accuracy at different target noise levels.Source: §4.4 · p. 9, Figure 6 · p. 9

4.5 Ablation Study

  • Sequence and top-k caching reduce total inference time by around 28%, with sequence caching contributing more.Sequence caching replaces time-consuming sequence generation; top-k caching reuses computed values for visited states.Source: §4.5 · p. 10, Figure 7 · p. 10
  • At equal numbers of generated candidates, TPSR significantly outperforms E2E; increasing rollouts, simulation beams, and expansions improves TPSR, especially rollouts.The default search settings are r = 3, b = 1, and kmax = 3.Source: §4.5 · p. 10, Figure 7 · p. 10

5 Conclusion

TPSR combines pretrained transformer SR models with MCTS lookahead planning and improves the fitting-complexity trade-off. The conclusion also reports controllable equation complexity without fine-tuning and improved extrapolation and noise robustness.

5 Conclusion

  • TPSR is a model-agnostic decoding strategy combining pretrained SR transformers with MCTS lookahead planning, and it outperforms existing methods on fitting-complexity trade-off.The conclusion also identifies integration of MCTS with transformer training or fine-tuning as a direction for future research.Source: §5 · p. 10
  • TPSR controls equation complexity without fine-tuning the pretrained model.The authors suggest λ=0.1 as a starting point elsewhere, while noting that the appropriate balance depends on the use case.Source: §5 · p. 10
  • Lookahead-planned expressions further improve noise robustness and extrapolation to unseen data.The conclusion identifies MCTS adaptability and parallel or distributed computing as possible future research directions.Source: §5 · p. 10

A Dataset Details

The evaluation uses three SRBench datasets and an in-domain synthetic validation set, covering known-equation problems, black-box datasets, and generated equations.

A Dataset Details

  • Evaluation covers four datasets: Feynman, Black-box, and Strogatz from SRBench, plus in-domain synthetic data.The synthetic validation set contains 400 equation examples generated using the same approach as the pretrained model's data.Source: A Dataset Details · p. 14
  • Feynman contains 119 equations and Strogatz 14 problems, with known underlying functions; Feynman inputs are split into bags of 200 when N exceeds 200.Feynman input dimension is d ≤10; Strogatz problems have d = 2.Source: A Dataset Details · p. 14
  • The Black-box evaluation uses 57 continuous-feature datasets with d ≤10, including real-world and synthetic data with varying noise.The study seeks interpretable expressions that fit these datasets effectively.Source: A Dataset Details · p. 14
  • Synthetic validation varies input dimension, unary and binary operator counts, and input-point counts of 50, 100, 150, or 200, then tests on 200 points.Test inputs are sampled from a multimodal distribution.Source: A Dataset Details · p. 14

B Implementation Details

The appendices detail TPSR's implementation, MCTS decoding and caching, and comparisons across complexity control and qualitative extrapolation. They report that TPSR can balance equation complexity and fitting performance and outperform E2E on an unseen evaluation region.

B Implementation Details

  • Experiments use E2E as the pretrained transformer backbone, with L=200, kmax=3, beam size 1, and 3 rollouts by default.SRBench observations are split 75%/25% for training and testing; synthetic evaluations vary sampling variance from 1 to 16.Source: B Implementation Details · p. 14

C.1 MCTS-Guided Decoding Details

  • MCTS selects states, expands using the transformer's top-k tokens, evaluates completed equations with beam search and reward, then backpropagates values.Simulation-generated sequences complete partial equations for feedback; they are not counted as tree visits.Source: §C.1 · p. 15

C.2 Distinguishing TPSR from other MCTS Approaches in SR

  • Unlike [48]'s mutation-based search over full equations, TPSR searches partial token sequences and uses pretrained transformer distributions without parameter updates.TPSR also uses the pretrained model to guide candidate generation, addressing the extensive action space and variable estimates of standalone MCTS.Source: §C.2 · pp. 16–17

C.3 Caching Details

  • Top-k and sequence caches reuse token predictions and completed equations to reduce redundant computation without compromising performance.Sequence caching applies when beam size is one and a later state matches part of a stored equation.Source: §C.3 · pp. 17–18

D.1 Controlling the Fitting-Complexity Trade-off

  • At λ=1, TPSR scores 0.916 versus E2E+Sampling's 0.848 in fitting performance, with complexity 47.24 versus 50.73.At λ=0.1, complexity falls below 60 with minimal fitting-performance effect; the authors recommend it as a starting setting, while noting use case affects the choice.Source: Figure 10 · p. 18, §D.1 · p. 18

D.2 Qualitative Study

  • On the x2sin(x) example, TPSR fits the unseen (−5, 5) region better than E2E, while symbolic regression methods outperform XGBoost there.Training uses 200 points sampled in (−2, 2).Source: §D.2 · pp. 18–19, Figure 11 · p. 19

D.3 Evaluating the Model-Agnostic Capability

On 52 Feynman equations with D ≤3, TPSR improves NeSymReS fitting accuracy without changing average equation complexity at λ=0.1; at λ=0, complexity rises slightly.

D.3 Evaluating the Model-Agnostic Capability

  • NeSymReS predicts equation skeletons only, requiring a more complex constant-optimization process for integration.It is limited to datasets with at most three dimensions.Source: §D.3 · p. 19
  • The NeSymReS integration evaluation uses 52 Feynman equations and respects its D ≤3 dimensionality constraint.The dataset matches the setup in [16].Source: §D.3 · p. 19
  • TPSR significantly improves NeSymReS fitting accuracy without changing average equation complexity at λ=0.1, with a slight complexity increase at λ=0.The reported comparison covers 52 Feynman datasets with D ≤3.Source: Table 3 · p. 19, §D.3 · p. 19

D.4 Additional SRBench Results

TPSR ranks first on Friedman and other PMLB black-box datasets, and performs strongly on Strogatz despite its time-ordered data differing from E2E pre-training. Across black-box input dimensions, TPSR generally outperforms competitors while its inference time remains relatively stable.

D.4 Additional SRBench Results

  • Despite E2E’s weaker performance on Strogatz’s time-ordered data, TPSR ranks among the top three for fitting while matching or slightly improving complexity and inference time.The Strogatz dataset contains 14 equations from a two-state first-order ODE system, unlike E2E’s pre-training data.Source: §D.4 · pp. 19–20, Figure 12 · p. 20
  • TPSR outperforms competitors across most black-box input dimensions, with inference time roughly constant and peaking at 10^3 seconds for d = 9.At d = 9, the passage compares this with 3 hours for Operon and 30 hours for SBP-GP; complexity comparisons vary by dimension and λ.Source: §D.4 · pp. 20–21
  • TPSR maintains top-1 fitting performance across challenging Friedman and remaining PMLB black-box datasets, while several baselines degrade on Friedman datasets.Figure 13 separates Friedman datasets, other PMLB datasets, and their combined results; error bars show 95% confidence intervals.Source: §D.4 · p. 20, Figure 13 · p. 21

D.5 Additional In-Domain Results

On in-domain synthetic data, increasing formula complexity degrades performance across models, but TPSR generally has smaller drops than E2E sampling. Its relative performance also varies with data quantity, diversity, and target noise.

D.5 Additional In-Domain Results

  • For N < 100, TPSR with λ = 0 or 0.1 performs considerably better than E2E, while the evaluated maximum is Nmax = 200.The passage attributes this cap to E2E pre-training with N ≤ 200 and the encoding stage’s substantial training compute and GPU requirements above 200 points.Source: §D.5 · p. 21
  • TPSR with λ = 0 or 0.1 consistently outperforms E2E as the number of input data centroids varies.Performance improves for all models as centroids increase, representing greater diversity across distribution clusters.Source: §D.5 · p. 21
  • As target-noise variance increases, all models’ performance deteriorates, but TPSR with λ > 0 performs slightly better than E2E, particularly at larger variance.The tested noise is multiplicative: y → y(1 + σ), with σ ∼ N(0, γ).Source: §D.5 · p. 21
  • As formula complexity increases, performance degrades across models, while TPSR with λ = 0 or 0.1 always has lower drops than E2E and λ = 0.5 or 1 mostly does.Figure 15 varies binary and unary operator counts, input dimension, input points, centroids, and noise variance.Source: §D.5 · p. 21, Figure 15 · p. 22

D.6 Additional Ablation Studies

The ablation finds an exploration–exploitation balance for β, with best performance over a reported range. Example expressions also show TPSR outputs more closely aligned with true functions and with superior fitting performance than E2E outputs.

D.6 Additional Ablation Studies

  • Increasing β increases the number of equation-sequence candidates, but beyond β > 0.1 the increase is marginal.The passage suggests repetitive sequence generation may activate caching mechanisms.Source: §D.6 · p. 23
  • β values from 0.1 to 10 achieve optimal performance across 119 Feynman datasets, while β = 0 and β = 100 are suboptimal.The passage attributes the extremes to constrained exploration at β = 0 and over-emphasized exploration at β = 100.Source: Figure 16 · p. 23, §D.6 · p. 23

D.7 Examples of Generated Symbolic Expressions

  • For 200 observation points, TPSR expressions with λ = 1 align more closely with true functions and have notably superior R2 fitting performance than E2E sampling outputs.The comparison uses the same model weights; generated constants are represented by the token “C”.Source: §D.7 · p. 23

E Discussion and Future Work

TPSR’s search and performance-feedback decoding increases inference time relative to beam search and sampling, though it remains faster than most GP-based SRBench baselines. Its dependence on a pretrained SR model also constrains the approach to that model’s learned priors and structural limits.

E Discussion and Future Work

  • TPSR takes longer to infer than simpler decoding methods such as beam search and sampling because it searches and incorporates performance feedback during generation.The passage says inference time remains lower than that of most GP-based SRBench baselines.Source: E Discussion and Future Work · p. 23
  • TPSR depends on the pretrained SR model’s learned priors and inherits its limits on input dimensionality, expression length, and vocabulary.Source: E Discussion and Future Work · p. 23
  • The discussion places TPSR’s inference-time trade-off between simpler decoding and most GP-based SRBench methods.The supplied table caption identifies example E2E and TPSR expressions with fitting performance, but does not state specific comparison values.Source: E Discussion and Future Work · p. 23, Table 4 · p. 24

C sin (CX1 + CX2) + CX2

TPSR’s scope depends partly on the pretrained model: with the E2E backbone, input dimension is limited to d ≤10, while the framework could integrate more advanced models. The discussion also outlines possible applications and ethical risks of symbolic regression.

C sin (CX1 + CX2) + CX2

  • TPSR with the E2E backbone is limited to input dimensions d ≤10, although the framework could integrate more advanced pretrained SR models.The bound follows from E2E’s maximum pretraining input dimension, dmax = 10.Source: C sin (CX1 + CX2) + CX2 · p. 24
  • TPSR considerably elevates pretrained-model performance on out-of-distribution datasets, especially in SRBench comparisons, but remains constrained by the model’s priors.The paper describes decoding-stage search and planning as a partial solution and suggests fine-tuning with non-differentiable rewards as future work.Source: C sin (CX1 + CX2) + CX2 · p. 24

F Broader Impacts

  • TPSR is presented as balancing fitting accuracy and complexity, with potential applications where accurate and interpretable mathematical models support understanding and prediction.The discussion identifies science and engineering as relevant domains.Source: F Broader Impacts · p. 24
  • The authors argue TPSR may improve the efficiency and scalability of symbolic regression and inspire other work combining MCTS and non-differentiable feedback with transformer models.They identify resource-constrained settings, scientific research, engineering design, and optimization as potential applications.Source: F Broader Impacts · p. 25
  • Symbolic regression may expose sensitive data patterns, so the authors call for separate modules to screen inputs and deny requests that could enable harmful pattern extraction.They cite potential privacy breaches as a concern, including in settings such as healthcare.Source: F Broader Impacts · p. 25

Related papers

Loading 2303.06833v5…