SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 1: Performance of TPSR compared with beam search and sampling decoding strategies on the SRBench [42] and In-domain Synthetic [18] datasets.
Data Group	Model	Feynman	Strogatz	Black-box
↑R2>0.99	↓Complexity	↑R2>0.99	↓Complexity	↑R2	↓Complexity
SRBench	E2E+Beam	0.815	54.19	0.357	53.21	0.847	83.61
E2E+Sampling	0.848	50.73	0.357	50.14	0.864	82.78
TPSR (λ=0)	0.952	84.42	0.928	82.78	0.938	129.85
TPSR (λ=0.1)	0.949	57.22	0.785	56.14	0.945	95.71
TPSR (λ=0.5)	0.924	50.01	0.714	47.02	0.931	82.58
TPSR (λ=1)	0.916	47.24	0.571	43.42	0.924	79.43
Data Group	Model	↑R2>0.99	↑R2	↑Acc0.1	↑Acc0.01	↑Acc0.001	↓Complexity
In-domain	E2E+Beam	0.657	0.782	0.461	0.298	0.2	38.37
E2E+Sampling	0.640	0.794	0.472	0.332	0.208	39.82
TPSR (λ=0)	0.702	0.828	0.550	0.416	0.333	67.11
TPSR (λ=0.1)	0.708	0.833	0.514	0.326	0.213	40.31
TPSR (λ=0.5)	0.697	0.830	0.459	0.274	0.184	36.55
TPSR (λ=1)	0.691	0.827	0.439	0.271	0.176	35.67
We evaluate TPSR and various baseline methods on standard SR benchmark datasets from Penn Machine Learning Benchmark (PMLB) [43] studied in SRBench [42], as well as In-domain Synthetic Data generated based on [38, 18]. The benchmark datasets include 119 equations from Feynman Lectures on Physics database series $^{2}$  [44], 14 symbolic regression problems from the ODE-Strogatz database $^{3}$  [45], and 57 Black-box $^{4}$  regression problems without known underlying equations. We limit the datasets to those with continuous features and input dimension  $d \leq 10$ , as the transformer SR model [18] is pre-trained with  $d_{max} = 10$ . The In-domain Synthetic Data consists of 400 validation functions with different levels of difficulty and number of input points. This data is referred to as "in-domain" because the validation functions and their corresponding input points are generated using the same approach as the data on which the backbone transformer model [18] is pre-trained. More details on each of these datasets are provided in Appendix A.
We evaluate our model using the following three metrics:  $R^2$  score [42], accuracy to tolerance  $\omega$  [16, 46], and equation complexity [18, 42].
$$
R ^ {2} = 1 - \frac {\sum_ {i} ^ {N _ {t e s t}} \left(y _ {i} - \tilde {y} _ {i}\right) ^ {2}}{\sum_ {i} ^ {N _ {t e s t}} \left(y _ {i} - \bar {y}\right) ^ {2}}, \quad A c c _ {\omega} = \mathbb {1} \left(\max  _ {1 \leq i \leq N _ {t e s t}} \left| \frac {\tilde {y} _ {i} - y _ {i}}{y _ {i}} \right| \leq \omega\right), \quad C o m p l e x i t y = \left| \mathcal {T} (\tilde {f} (\cdot)) \right|,
$$
where  $R^2$  measures fitting performance with  $\bar{y}$  as the mean of  $y$  in test set,  $Acc_{\omega}$  evaluates equation precision based on tolerance threshold  $\omega$ , and equation complexity is determined by the number of nodes in the expression tree  $\mathcal{T}$  of the generated equation  $\tilde{f}(\cdot)$ . Following [18, 42], we set  $R^2 = 0$  for rare pathological examples and discard the worst 5% predictions for  $Acc_{\omega}$  to reduce outlier sensitivity.
4.3 (RQ1) Effectiveness of TPSR
Table 1 presents the performance comparison results of TPSR with the baseline decoding strategies on the SRBench benchmark and the In-domain synthetic dataset. For the E2E baseline, we use the settings reported in [18], including beam/sample size of  $C = 10$  candidates, and the refinement of all the candidates  $K = 10$ . For our model, we use the width of tree search as  $k_{max} = 3$ , number of rollouts  $r = 3$ , and simulation beam size  $b = 1$  as the default setting. For PMLB datasets that contain more than 200 points, we follow [18] and use  $B$  bags of data, each containing  $N = 200$  points, due to the limitation that the baseline method is pre-trained with  $N \leq 200$  data points. In the baseline method [18], a total of  $BC$  candidates are generated ( $C$  candidates for  $B$  bags), which are then sorted and refined to generate the best equation. However, for TPSR, since we need to train an MCTS for each bag, we use an iterative decoding approach, starting with the first bag and continuing
$^{2}$ https://space.mit.edu/home/tegmark/aifeynman.html
<sup>3</sup>https://github.com/lacava/ode-strogatz
<sup>4</sup>https://github.com/EpistasisLab/pmlb/tree/master/datasets
Figure 5: Performance comparison of TPSR and SRBench algorithms in terms of Accuracy-Complexity-Time on Feynman (top) and Black-box (bottom) datasets. For Feynman dataset, algorithms are sorted based on mean accuracy defined as the ratio of solutions with  $R^2 > 0.99$  on test set under various noise levels, and for Black-box dataset, the algorithms are sorted based on the median  $R^2$  score on test set. TPSR demonstrates a strong balance of performance with relatively low model complexity and lower inference time compared to GP-based algorithms. The error bars represent the  $95\%$  confidence interval and " *" refers to SR methods for Black-box dataset.
with subsequent bags until a criterion  $(R^2 > 0.99)$  is met or we use a maximum of  $B = 10$  bags. To ensure a fair comparison, we use  $B = 10$  for the E2E baseline method as well. In this table, we demonstrate the results of our proposed framework, TPSR, with varying values of the  $\lambda$  parameter that controls the trade-off between fitting performance and complexity in the hybrid reward function defined in Eq. (2). For a detailed comparison of the experimental settings across different approaches, refer to Table 2 in Appendix B.
As shown in Table 1, when  $\lambda = 0$ , the framework generates complex equations that overoptimize for fitting performance. However, as we increase  $\lambda$ , the framework generates less complex equations with a slight reduction in fitting performance. Notably, even for large values of  $\lambda$ , such as  $\lambda = 1$ , the fitting performance of TPSR significantly outperforms that of the baseline methods. Based on the results, we recommend a default setting of  $\lambda = 0.1$  as it offers a balanced trade-off between complexity and accuracy, while also mitigating potential overfitting (as detailed in Appendix D.1). These findings demonstrate the superiority of TPSR over the baseline methods in terms of fitting performance across all datasets, while generating equations with comparable or reduced complexity than those generated by the baseline methods. Table 1 shows that TPSR exhibits a more significant gap in fitting performance when compared to E2E baselines on SRBench datasets, while this gap is smaller for In-domain datasets (even performing slightly worse on  $Acc_{\omega}$  for larger  $\lambda = 0.5, 1$ ). This is due to the In-domain dataset being generated using the same approach as the E2E pre-training data, resulting in the E2E model's superior performance on this synthetic dataset. Furthermore, qualitative comparisons of TPSR with baseline symbolic and black-box regression models demonstrate the
Figure 6: TPSR with  $\lambda \in \{0,0.1,0.5,1\}$  compared to E2E for (a) Extrapolation performance where in-domain accuracy is shown for different input sampling variance  $(\sigma)$ , and (b) Robustness to noise, where mean accuracy  $(R^2 > 0.99)$  is shown for various target noise levels  $(\gamma)$ .
superior performance of TPSR in learning the underlying equation and out-of-domain extrapolation (see Appendix D.3).
Fig. 5 presents a detailed comparison of our proposed TPSR with the baseline E2E transformer model and all the SRBench baselines on the PMLB Feynman and Black-box datasets. This figure illustrates the relative position of each algorithm with respect to (1) fitting performance, (2) model complexity, and (3) inference time. The results indicate that transformer-based planning in the TPSR significantly enhances the performance of E2E and outperforms even the state-of-the-art GP baselines, achieving the highest fitting performance on the black-box datasets. This is achieved while the complexity of the generated equations in TPSR is not greater than that of E2E, and shows a great fitting-complexity-time balance compared to other SR algorithms. The pareto plots provided in Fig. 1 and Appendix D.2 also demonstrate the effectiveness of TPSR in balancing fitting-complexity as well as fitting-time compared to all other SRBench baselines. Our TPSR effectively pushes this balanced performance to the first pareto front for both the Feynman and Black-box datasets. Moreover, it is important to note that, while the inference time of TPSR is longer than the baseline E2E transformer model, it still has significantly lower inference time than RL or GP-based SRBench baselines. Further results on the SRBench and In-domain datasets are provided in Appendix D.
4.4 (RQ2) Extrapolation and Robustness
The ability to extrapolate well is inherently linked to the quality of the equations discovered through symbolic regression. To investigate the extrapolation performance of TPSR to out-of-training regions, we normalize the input test data points to different scales  $(\sigma)$  instead of unit variance (used for training points) as per [18]. Fig. 6(a) depicts the average performance of TPSR compared to E2E with sampling decoding on the training data as well as testing data in scales of  $\sigma \in \{1,2,4,8,16\}$  for the In-domain Synthetic dataset. Also, we investigate the effect of different complexity controlling levels  $(\lambda \in \{0,0.1,0.5,1.0\})$  on the extrapolation performance. It can be observed that, while  $\lambda = 0$  (i.e., no complexity regularization) achieves the best fitting accuracy on the training data, it has a sub-par performance for  $\sigma > 8$ . This can be due to the overfitting issue when the symbolic model is much more complex than the true function, similar to the common overfitting issue in ML models. Results highlight the importance of controlling complexity in the extrapolation of identified equations. For values of  $\lambda > 0$ , the overfitting issue is mitigated as the generated equations become less complex. However, very high values of  $\lambda$  (e.g.,  $\lambda = 1$ ) mostly result in poor accuracy performance. The flexibility of TPSR for allowing different values of  $\lambda$  to balance fitting and complexity for a given task is crucial for optimal performance. Fig. 6(b) also presents the robustness of TPSR with different  $\lambda$  levels compared to the E2E transformer baseline on the Feynman dataset. The results indicate that MCTS-guided lookahead planning can offer robust performance with a smaller drop in accuracy compared to the baseline in the presence of noise.
In this section, we investigate the effect of different MCTS parameters and caching mechanisms on the performance of TPSR by conducting ablation experiments on the Feynman datasets.
Figure 7: Ablation study on the modules and parameters of TPSR. (a) Effect of caching mechanisms: Sequence caching and top-k caching improve the inference time of TPSR ( $\lambda = 0.1$ ). (b) Efficiency and parameters of TPSR: Average accuracy of TPSR (varying model parameters), and baseline E2E (varying sampling size) across number of generated candidates.
(RQ3) Caching Mechanisms. In Fig. 7(a), we illustrate the effectiveness of the sequence and top-k caching mechanisms in reducing the total inference time of TPSR ( $\lambda = 0.1$ ). Our experiments show that sequence caching has more effect in dropping the inference time as it replaces the time-consuming sequence generation process. Overall, these two mechanisms can reduce the total inference time by around  $28\%$ .
(RQ4) Search Parameters. Fig. 7(b) shows the accuracy performance against the number of generated equations during the decoding process for both TPSR ( $\lambda = 0.1$ ) and the baseline E2E with sampling decoding. In this figure, the 'number of generated equation candidates' represents the total number of complete equation sequences generated by each method. Specifically, this refers to the sample size in the E2E with sampling decoding, and the function calls of the beam search sub-routine multiplied by beam size  $b$  in TPSR. The results show that under the same number of generated equation candidates, TPSR significantly outperforms the E2E baseline. This is primarily attributed to the fact that the E2E baseline is deprived of any performance feedback during the equation generation process. We report the results for variants of TPSR with different MCTS parameters. We assess the performance with varying number of rollouts,  $r = \{1,3,6,9\}$ , number of beams in simulations,  $b = \{1,3\}$ , and the maximum number of possible expansions at each state,  $k_{max} = \{2,3,4\}$ . The default setting of TPSR parameters are  $b = 1$ ,  $k_{max} = 3$ , and  $r = 3$ . Results indicate that increasing  $r$ ,  $k_{max}$ , and  $b$  all contribute to the better performance of TPSR, with the most significant improvement observed when increasing  $r$ . This is because more rollouts provide model with more opportunities to learn from trials and learn better values.
In this work, we propose TPSR, a model-agnostic decoding strategy for symbolic regression that leverages the power of pre-trained SR transformer models combined with MCTS lookahead planning, and outperforms the existing methods in generating equations with superior fitting-complexity-time trade-off. We demonstrate the flexibility of TPSR in controlling discovered equation complexity without fine-tuning the pre-trained model. We also show that TPSR performs  $100\mathrm{x}$  faster than state-of-the-art genetic algorithms by leveraging the pre-trained priors. Additional results show that better expressions obtained with lookahead planning can further improve model performance in terms of noise robustness and extrapolation to unseen data. Future research could focus on enhancing the adaptability of feedback-based expression generation mechanisms, potentially by modulating the flexibility of MCTS or transformer model weights, and the integration of MCTS with the training or fine-tuning of transformer SR models. Furthermore, employing parallelization and distributed computing could potentially improve planning efficiency.