Discovering Mathematical Formulas from Data via GPT-guided Monte Carlo Tree Search

TL;DR

The paper tackles symbolic regression, the NP-hard task of discovering concise and interpretable formulas from data, by introducing SR-GPT, a framework that couples a Generative Pre-Trained Transformer (GPT) with Monte Carlo Tree Search (MCTS) to guide and refine expression search. The GPT provides a symbol-policy and state-value, which steer MCTS during expansion and simulation, while self-search data are used to iteratively train the GPT, forming a mutual feedback loop. A novel entropy-aware loss on the GPT and an $S_{NRMSE}$-based reward are introduced to sharpen guidance and address multivariate variable omissions. Across 222 expressions from 10+ datasets, SR-GPT achieves state-of-the-art recovery, shows strong noise robustness, and demonstrates high generalization on the AI Feynman suite, indicating practical potential for interpretable formula discovery.

Abstract

Finding a concise and interpretable mathematical formula that accurately describes the relationship between each variable and the predicted value in the data is a crucial task in scientific research, as well as a significant challenge in artificial intelligence. This problem is referred to as symbolic regression, which is an NP-hard problem. In the previous year, a novel symbolic regression methodology utilizing Monte Carlo Tree Search (MCTS) was advanced, achieving state-of-the-art results on a diverse range of datasets. although this algorithm has shown considerable improvement in recovering target expressions compared to previous methods, the lack of guidance during the MCTS process severely hampers its search efficiency. Recently, some algorithms have added a pre-trained policy network to guide the search of MCTS, but the pre-trained policy network generalizes poorly. To optimize the trade-off between efficiency and versatility, we introduce SR-GPT, a novel algorithm for symbolic regression that integrates Monte Carlo Tree Search (MCTS) with a Generative Pre-Trained Transformer (GPT). By using GPT to guide the MCTS, the search efficiency of MCTS is significantly improved. Next, we utilize the MCTS results to further refine the GPT, enhancing its capabilities and providing more accurate guidance for the MCTS. MCTS and GPT are coupled together and optimize each other until the target expression is successfully determined. We conducted extensive evaluations of SR-GPT using 222 expressions sourced from over 10 different symbolic regression datasets. The experimental results demonstrate that SR-GPT outperforms existing state-of-the-art algorithms in accurately recovering symbolic expressions both with and without added noise.

Figures (5)

• Figure 1: a. self-search. During the self-search phase, the operators with the highest probabilities(max($\pi$)) are selected sequentially for each state $s_1, s_2, ..., s_t$. For each state $s_t$, multiple simulations are carried out, and during each simulation, an MCTS $\alpha(\theta)$ guided by the latest neural network $\mathcal{N}_{\theta}$ is executed (as depicted in Fig \ref{['fig2']}). Finally, by normalizing the number of times each child state of the current state has been visited, we can determine the selection probability $a_{t}\sim\pi_{t})$ for each state. Finally, when the expression is complete at time $s_r$, we calculate the reward value $z$ and perform backpropagation. b. GPT training. In SR-GPT, the GPT is designed to take the state $s_t$ as input, which is then passed through an GPT with parameters $\theta$. The output comprises a vector $p_t$ and a scalar value $v_t$, where $p_t$ represents a probability distribution over candidate symbols and $v_t$ represents the possible reward value after generating a complete expression starting from the current state $s_t$. During training, the GPT parameters, $\theta$, are updated to maximize the similarity between the policy vector, $p_t$, and the search probabilities, $\pi_t$, while minimizing the difference in the predicted reward $v_t$ and the actual reward $z$. The new parameters are used in the next iteration of self-search as in Fig (a).
• Figure 2: MCTS in SR-GPT. Select, Starting from the root node, the child node with the largest UCT value is selected recursively, where $\mathcal{UCT}=\mathcal{Q}(s, a)+\mathcal{U}(s, a)$, and finally a leaf node is reached. Expand, If the current leaf is not a terminal node, then create multiple (number of symbols) child nodes, and select the one with the greatest probability $p(s, a)$ to expand, where $p(s, \bullet) = \mathcal{N}_{\theta}(s)$, and the $p(s, a)$ values are stored in the corresponding new child nodes. evaluate, Use GPT to calculate the $V$ of the new extension node. backpropagation, After each simulation, backpropagation is performed, where $\mathcal{N}_{visit}$ of all visited nodes during the simulation is incremented by one. Additionally, Action value $Q$ is updated to track the mean of all evaluations V in the subtree below that node\ref{['e2']}. Search, The probability $\pi$ of selecting each symbol can be determined during the self-search phase after the completion of the simulation phase. Herein, $\pi$ is directly proportional to $(\mathcal{N}_{visit})^{\frac{1}{\tau}}$, where $\mathcal{N}_{visit}$ is the visit count of each move from the root state and $\tau$ is a parameter controlling temperature.
• Figure 3: Fig (a) illustrates the recovery for various ablations of SR-GPT on all Nguyen benchmarks, with error bars indicating the standard errors. Fig (b) describes the recovery rate of SR-GPT and four other excellent algorithms on all benchmarks under different levels of noise. Fig (c) It shows the general ability of five algorithms, from which we can see that SR-GPT is better than other algorithms in general ability.
• Figure B.1: The series of figures above presents line graphs depicting the reward values of SR-GPT on the Nguyen dataset over time. As observed from the figures, throughout the search process, the reward values for all expressions demonstrate an oscillatory ascent with the increase in training iterations. Notably, Expression 8 is an exception due to its comparatively simple structure, achieving its best result within just one epoch.
• Figure C.1: Fig (a) illustrates the recovery for various ablations of SR-GPT on all Nguyen benchmarks, with error bars indicating the standard errors. Fig (b) describes the recovery rate of SR-GPT and four other excellent algorithms on all benchmarks under different levels of noise. Fig (c) It shows the general ability of five algorithms, from which we can see that SR-GPT is better than other algorithms in general ability.