SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Parshin Shojaee\*1, Kazem Meidani\* 2 Amir Barati Farimani 2,3, Chandan K. Reddy
$^{1}$  Department of Computer Science, Virginia Tech  
 $^{2}$  Department of Mechanical Engineering, Carnegie Mellon University  
 $^{3}$  Machine Learning Department, Carnegie Mellon University
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 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 objectives borrowed from text generation and overlook equation discovery goals like accuracy and complexity. To address this, we propose TPSR, a Transformer-based Planning strategy for Symbolic Regression that incorporates Monte Carlo Tree Search planning algorithm into the transformer decoding process. Unlike conventional decoding strategies, TPSR enables the integration of non-differentiable equation verification feedback, such as fitting accuracy and complexity, as external sources of knowledge into the transformer 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<sup>1</sup>.
Symbolic regression (SR) is a powerful method to discover mathematical expressions for governing equations of complex systems and to describe data patterns in an interpretable symbolic form. It finds extensive applications in science and engineering, enabling the modeling of physical phenomena in various domains such as molecular dynamics, fluid dynamics, and cosmology [1-6]. Symbolic representations provide valuable insights into complex systems, facilitating a better understanding, prediction, and control of these systems through the design of accurate, generalizable, and efficient models [7-9]. SR models establish the functional relationship between independent and target variables by mapping them to mathematical equations. The input data can be obtained from simulations, experimental measurements, or real-world observations. Symbolic regression, however, poses several challenges, including the combinatorial nature of the large optimization search space, vulnerability to the quality of input data, and the difficulty of striking a balance between model fitting, complexity, and generalization performance [10, 11].
Symbolic regression encompasses a wide range of methods, spanning different categories. Traditional approaches, such as Genetic Programming (GP), use a heuristic population-based search strategy where each individual represents a potential solution to the problem [12, 13]. Though GP algorithms are capable of finding solutions for nonlinear and complex problems, they are typically slow to converge due to the vast functional search space. Also, as they need to start the search from scratch for each dataset, they tend to be computationally expensive, prone to overfitting, and sensitive to the
*Equal contribution. Contact email: parshinshojaae@vt.edu
<sup>1</sup>The codes are available at: https://github.com/deep-symbolic-mathematics/TPSR
(a) Black-box dataset
(b) Feynman dataset
Figure 1: Pareto plot comparing the rankings of all methods in terms of the  $R^2$  performance and identified equation complexity for (a) SRBench Black-box datasets and (b) Feynman datasets. Our results with Transformer-based Planning (TPSR) applied on top of E2E transformer SR model improves its average accuracy on both data groups while maintaining a similar range of equation complexity. TPSR can successfully reach the first Pareto-front which is better than E2E baseline on both data groups. Connecting lines denote Pareto dominance rankings, colors denote the families of models, and "*" indicates SR versus ML methods in Black-box datasets.
choice of parameters [14]. Recent works in SR have shown promising results by using pre-trained transformers [15] for generating equations as sequences of tokens. These models leverage the prior knowledge learned through large-scale pre-training and can generate equations with a single forward pass, leading to considerably faster inference time compared to the GP-based methods [16-19]. However, one of the limitations of these models is that they focus on the supervised language model pre-training objective borrowed from text generation, i.e., they are trained solely with the token-level cross-entropy loss, which can result in equations that may exhibit high token-level similarities but are suboptimal with respect to equation discovery objectives such as fitting accuracy and complexity which are critical in this task. To mitigate this issue, beam search [20, 21] or sampling [22] approaches have been employed as decoding strategies to propose multiple candidate equations for a given dataset, and then select the optimal candidate equation based on the fitting accuracy after optimizing for constants. Nonetheless, both beam search and sampling decoding strategies primarily rely on the pretrained transformer's logits and next token probabilities, and therefore do not receive any performance feedback during the generation of equation candidates.
To consider the equation discovery objectives in the transformer generation process and still benefit from the pre-trained model logits, we propose TPSR, a Transformer-based Planning strategy for Symbolic Regression. TPSR leverages a lookahead planning algorithm, using Monte Carlo Tree Search (MCTS) as a decoding strategy on top of pre-trained transformer SR models to guide equation sequence generation. TPSR significantly improves performance of the discovered equations by considering verification feedback during the generation process and still remains considerably faster than GP-based models which do not leverage the pre-training priors and learn expressions for each dataset from scratch. Notably, our approach is model-agnostic and can be applied to any pre-trained SR model, enabling optimization of generated equation sequences for non-differentiable objectives that may encompass combinations of fitting accuracy, complexity, and equation forms. Additionally, we incorporate different caching mechanisms to reduce the overall inference time. Our experimental results demonstrate that applying TPSR on top of the pre-trained E2E SR model [18] significantly enhances its performance across various benchmark datasets. As depicted in Fig. 1, TPSR achieves a strong balance between fitting accuracy and model complexity compared to other leading baselines. It also effectively drives the E2E model towards the optimal trade-off, represented by the first Pareto front. The major contributions of this work are summarized below:
- Proposing TPSR, a new method that combines pre-trained transformer SR models with Monte Carlo Tree Search (MCTS) lookahead planning to optimize the generation of equation sequences while considering non-differentiable performance feedback.
- Developing a new reward function that balances equation fitting accuracy and complexity to optimize the generated equations for an effective trade-off.
- Demonstrating that TPSR consistently outperforms state-of-the-art baselines across various SR benchmark datasets, generating equations with higher fitting accuracy while maintaining lower complexity to avoid non-parsimonious solutions. TPSR still achieves considerably faster inference time than GP-based models which do not use pre-trained priors.
- Showcasing the extrapolation and noise robustness of TPSR compared to the baseline and conducting an ablation study to investigate the impact of various model components.
Symbolic Regression without Learned Priors. Genetic Programming (GP) algorithms are typically employed for single-instance SR, aiming to find the best-fit equation for a "single" dataset at hand [12]. Recently, alternative neural network-based search algorithms have been explored, including deep reinforcement learning (RL) [14, 23, 24], combinations of GP and RL [25], and Monte Carlo Tree Search (MCTS) as a standalone framework [26]. Despite their successes, all these methods lack the benefits of prior knowledge learned from large-scale pre-training. Consequently, they are slow during inference as they need to restart the search from scratch for new datasets.
Pre-trained Transformers for Symbolic Regression. In recent years, pre-trained transformers have shown remarkable performance in natural language and programming language tasks [27-29]. This success has inspired researchers to develop pre-trained transformer models for SR [16-19, 30]. For example, Biggio et al. [16] introduced a Neural Symbolic Regression model that scales (NeSymReS) with the amount of synthetic training data and generates equation skeletons where all the numerical constants are represented by a single token “ $C$ ”. Kamienny et al. [18] proposed an end-to-end transformer SR framework that predicts the complete equation form along with its constants. More recent works [30, 31] introduced unified frameworks that include a transformer-based pre-training stage as the prior for subsequent RL or GP optimization steps. While GP and RL methods have to start anew for each problem, the transformer approaches rely on synthetic data and the power of large-scale pre-trained priors to generate equations in a single forward pass. However, these models are pre-trained on token-level language modeling loss function and thus can perform suboptimal for other equation discovery objectives critical in SR such as fitting accuracy to the observed data as well as equation's complexity. Our model, TPSR, utilizes lookahead planning to guide the generation of equations towards better performance by employing these objectives as feedback during the transformer generation process.
Planning in Sequence Generation. Recently, planning algorithms have been utilized in NLP tasks to optimize text output for specific objectives, such as controlling generated text to meet certain constraints like non-toxicity or conveying certain emotions [32-34]. Recent advances in programming language models developed in code generation have also yielded promising techniques that could be adapted for SR, as they share several vital similarities with each other. Both involve generating sequences of symbols for a given input and typically require optimizing the generated sequences for specific criteria which is different from the pre-trianing objective. For code generation, this may involve optimizing objectives like code compiling ability, readability, or passing test cases [35-37]. Similarly, in SR, the focus may be on equation-specific sequence-level objectives such as fitting accuracy or minimizing complexity. Motivated by these successes, we develop an approach that combines MCTS planning with pre-trained transformer SR models for improved equation discovery.
In SR, the main goal is to find a symbolic expression for the unknown function  $f(\cdot)$  mapping the  $d$ -dimensional input  $\pmb{x} \in \mathbb{R}^d$  to the target variable  $y = f(\pmb{x}) \in \mathbb{R}$ . Given a dataset of  $n$  observations  $\mathcal{D} = (x_i, y_i)_{i=1}^n$ , SR methods try to generate an equation  $\tilde{f}(\cdot)$  such that  $y_i \approx \tilde{f}(x_i)$  for all  $i \in \mathbb{N}_n$ . Also, the proposed equation is desired to generalize well and to effectively balance the fitting accuracy and complexity. The transformer SR models are trained on a large-scale dataset comprising equation instances paired with their corresponding observations,  $\{(\mathcal{D}_1, f_1(\cdot)) \dots (\mathcal{D}_M, f_M(\cdot))\}$ , where  $M$  is the dataset size (number of paired samples). During inference, the trained model directly generates the equation  $\tilde{f}(\cdot)$  as a sequence of tokens in an autoregressive manner. An effective way to represent the expression tree of equations in a sequence is to use prefix notation as in [38]. For embedding the
(b) Inference with Beam Search/ Sampling
(c) Inference with MCTS
Figure 2: An overview of our proposed method with MCTS-guided decoding at inference compared to the concurrent works with beam search/sampling decoding strategy.
observations, we adopt the pre-trained SR model backbone from [18]. Notably, given the potential for large input sequences with tokenized numeric data and the quadratic complexity of transformers, the method introduced in [18] deploys a linear embedder module to map tokenized inputs to a singular embedding space before introducing them to the transformer encoder and decoder. Subsequent to embedding, these models encode the input observations and then pass the encoded representation along with the masked tokens to decode the equation sequence. To train the model, token-level cross-entropy loss with teacher forcing is employed to learn the distribution of next token prediction conditioned on the encoded dataset and the current state of sequence (Fig. 2(a)).
Achieving a good fitting performance from the model's predicted sequence demands generating accurate constants in the equation. To address this, the generated skeleton or equation can undergo a round of optimization to estimate their constants using nonlinear methods, such as Broyden-Fletcher-Goldfarb-Shanno algorithm (BFGS) [39]. Previous works [18, 16] employ beam search and sampling strategies for transformer decoding in combination with constant optimization to propose several candidate equations. Subsequently, they use fitting metrics such as  $R^2$  to order these candidates and output the final equation with the best performance (Fig. 2(b)). Transformer models utilizing beam search or sampling decoding strategies can generate multiple high-likelihood equation sequences, but their generation process is based on logits obtained from model parameters pre-trained with token-matching loss relative to the reference equation. As a result, such models lack the capability to receive verification feedback during generation and optimize sequence for equation discovery objectives such as fitting or complexity of equations.
3.2 MCTS-Guided Equation Generation
To generate equations that are both better-fitting and less-complex, it is crucial to incorporate feedback into the equation generation process. To achieve this, we utilize Monte Carlo Tree Search (MCTS) lookahead planning during inference, guiding the decoder towards optimal solutions for fitting and complexity objectives (as shown in Fig. 2(c)). The MCTS-guided transformer decoding explores different possibilities, identifying the most promising paths based on the objectives.
We frame the SR equation generation task as a Markov Decision Process (MDP) where state  $s$  represents the current sequence at generation step (token)  $t$ . If  $s$  has not reached the terminal state (i.e., the <EOS> token), we select the next token from the vocabulary as action  $a$ , updating state  $s'$  by concatenating  $s$  and  $a$ . Upon reaching the terminal state, the reward  $r$  is computed and used to update the decoding model. MCTS represents states as nodes and actions as edges within a tree structure, navigating state-space from the root node (i.e., initial state) to reach terminal states with maximum rewards. MCTS balances exploration and exploitation, considering nodes that lead to higher quality equations (i.e., higher Q-values) and under-explored nodes (i.e., those with fewer visits). During the generation process of the transformer, we utilize the MCTS algorithm iteratively to conduct lookahead planning and determine the next token. However, the large search-space requires more than the sole application of MCTS to discover high-quality equations. We need to effectively share information between the pre-trained transformer model and MCTS for better generations. To achieve this, we incorporate the probabilities of the next-token that are acquired through the pre-trained
Figure 3: Overview of TPSR's key steps: Selection, Expansion, Evaluation, and Backpropagation. MCTS-guided lookahead planning in decoding interacts with the pre-trained transformer SR model in the expansion and evaluation steps employing the transformer top-k sampling and beam search, respectively. The designed reward is used to guide the backpropagation.
transformer SR models into the MCTS planning process. This incorporation helps to enhance the search process, leading to more efficient and effective results. The key steps of MCTS for transformer decoding in SR models, as depicted in Fig. 3, are as follows:
Selection. The Upper Confidence Bound for Trees [40] criterion is employed to select actions (i.e., next tokens) for fully extended nodes in the search tree, balancing exploration and exploitation. We use the P-UCB heuristic [41] as
$$
\mathrm {P} - \mathrm {U C B} (s, a) = Q (s, a) + \beta \cdot P _ {\theta} (a | s) \cdot \sqrt {\frac {\ln (N (s))}{1 + N \left(s ^ {\prime}\right)}}, \tag {1}
$$
where  $Q(s, a)$  is the maximum return for action  $a$  in state  $s$  across all simulations, promoting the exploitation of the optimal child node. The second term encourages exploration of less-visited children, with  $N(s)$  as state  $s$ 's visit count and  $s'$  as the subsequent state.  $P_{\theta}(a|s)$  is the probability of the next token  $a$  given the partial sequence state  $s$  from pre-trained transformer model parameterized by  $\theta$ . The exploration-exploitation trade-off is adjusted by hyperparameter  $\beta$ . Lastly, the next token action maximizes the P-UCB:  $\operatorname{Select}(s) = \arg \max_{a} \mathrm{P-UCB}(s, a)$ .
Expansion. In the expansion stage, after selecting a node that is not fully expanded, a new child (next token) for the current state is explored. Random expansion of the node from the vocabulary, however, might result in an invalid equation (that does not comply with the prefix notation) and makes the search process very time-consuming. Therefore, given partial equations, only top- $k$  most likely choices of the next token are considered as the possible children of the node for expansion. In other words, we are restricting the actions to be only from the top- $k$  high-likelihood options which are retrieved from the pre-trained transformer SR model's logits. These options are then ordered to determine the sequence in which the children will be expanded.
Evaluation. To evaluate the newly expanded nodes, we perform simulations to complete the equation sequence. This is necessary because the new state may still be a partial equation and performance feedback can only be obtained at the end of the sequence when the equation generation is completed. In MCTS, it is common to employ random actions during the simulation stage. Nevertheless, random action selection for equation generation, much like during expansion, suffers from certain drawbacks in terms of time and the possibility of generating invalid equations. Consequently, the pre-trained transformer SR model is invoked again, this time utilizing beam search with a beam size of  $b$ , to generate complete equation candidates based on the current state. The beam size  $b$  determines the number of complete equations to be generated from the current partial equation. Following these simulations, the highest reward among all the complete equation candidates is assigned to the new node value.
Backpropagation. After generating a complete equation  $\tilde{f} (\cdot)$ , the corresponding reward  $r(\tilde{f} (\cdot))$  can be computed. The highest reward among all simulations is then assigned to the new node, which recursively backpropagates its estimated value to its parents until it reaches the root of the tree. This
update process involves updating the  $Q$  values of all state-action pairs, denoted as  $s'$  and  $a'$ , along the trajectory in the tree to reach the root. Specifically, for each state-action pair, the  $Q$  value is updated by taking the maximum of the current  $Q$  value and the new value  $r$ :  $Q(s', a') \gets \max(Q(s', a'), r)$ . More details on TPSR, including its steps and implementation can be found in Appendix C.
3.3 Reward Definition
We define a numerical reward  $r \in \mathbb{R}$  to evaluate complete equation candidate  $\tilde{f}(\cdot)$ , promoting fitting accuracy and regulating complexity. After optimizing constants in the complete sequence, we compute the reward. We first calculate the normalized mean squared error (NMSE) between ground-truth target variable  $y$  and predicted target variable  $\tilde{y} = \tilde{f}(\boldsymbol{x})$ , and formulate the reward as:
$$
r (\tilde {f} (\cdot) | \boldsymbol {x}, y) = \frac {1}{1 + \operatorname {N M S E} (y , \tilde {f} (\boldsymbol {x}))} + \lambda \exp \left(- \frac {l (\tilde {f} (\cdot))}{L}\right), \tag {2}
$$
where  $l$  represents equation complexity as the sequence length in prefix notation [18, 42, 16];  $L$  denotes the model's maximum sequence length; and  $\lambda$  is a hyperparameter balancing fitting and complexity reward terms. Higher  $\lambda$  values favor less complex equations, encouraging best-fitting and penalizing non-parsimonious solutions. NMSE is calculated as  $(\frac{1}{n}\| y - \tilde{f} (\pmb {x})\|_{2}^{2}) / (\frac{1}{n}\| y\|_{2}^{2} + \epsilon)$ , where  $\epsilon$  is a small constant to prevent numerical instability.
3.4 Efficient Implementation with Caching
During MCTS evaluation, the transformer model generates complete equation sequences from a given state, constructing implicit tree structures for beam search and computing top-k next tokens for visited states. These computations are required in future MCTS iterations, so we employ two caching mechanisms, top-k caching and sequence caching, to reduce redundancy and improve efficiency. Top-k caching stores computed top-k values for given states. For example, in Fig. 4, when evaluating state  $s = [+, \sin]$  in MCTS iteration  $t$ , top-k tokens are computed for  $s$  and subsequent visited states, such as  $[+, \sin, x_2]$ . State and top-k value pairs are cached for future use, avoiding re
Figure 4: An illustration of caching mechanisms in TPSR.
dundant token retrieval. Sequence caching caches complete equations generated with the provided beam size. If a state matches a stored equation partially, the cached equation can be used directly in future iterations, bypassing iterative sequence generation. Both caching strategies are designed to enhance efficiency without compromising performance. More details are provided in Appendix C.
In this section, we present our experimental results that evaluate the effectiveness and efficiency of TPSR. While the proposed decoding strategy is generally model-agnostic, here we showcase the results of using TPSR for the end-to-end (E2E) pre-trained SR transformer backbone [18], as E2E is the state-of-the-art open-source pre-trained SR model with publicly accessible model weights. Additional results of using TPSR with the NeSymReS pre-trained SR backbone [16] can be found in Appendix D.4. We evaluate our framework by answering the following research questions (RQs):
RQ1. Does TPSR perform better than other decoding strategies (beam search/sampling) and competing baseline methods over standard SR benchmark datasets?
RQ2. Does TPSR provide better extrapolation and robustness to noise?
RQ3. Are TPSR's caching mechanisms effective in reducing computation time?
RQ4. What is the role of individual components in TPSR's overall performance gain?