LLM-SR: Scientific Equation Discovery via Programming with Large Language Models

In the task of data-driven equation discovery, also known as symbolic regression (SR), the goal is to find a concise and interpretable symbolic mathematical expression $\tilde{f}$ that closely approximates an unknown underlying equation $f:\mathbb{R}^{d}\rightarrow\mathbb{R}$, which maps a $d$-dimensional input vector $\mathbf{x}$ to output $y$. Given a dataset $\mathcal{D}=\{(\mathbf{x}_{i},y_{i})\}_{i=1}^{n}$ consisting of input-output pairs, SR methods seek to uncover the hidden mathematical relationship, such that $\tilde{f}(\mathbf{x}_{i})\approx y_{i}$ for all $i$. The discovered equation should not only accurately fit the observed data points but also exhibit strong generalization capabilities to unseen data while maintaining interpretability.

Current SR methods typically represent equations using techniques such as expression trees (Cranmer, 2023), prefix sequences (Petersen et al., 2021; Biggio et al., 2021), or context-free grammars (Brence et al., 2021), constructing a search space $\mathcal{F}$. These representations provide limited and structured search spaces, enabling evolutionary search approaches like genetic programming (GP) to explore and find candidate expressions. In contrast, our work employs program functions to directly map inputs $\mathbf{x}$ to output targets $y$: def f(x): ... return y. This offers a flexible and expressive way to represent equations, allowing for a diverse set of mathematical operations and step-by-step instructions to capture complex mathematical relations. However, the space of possible programs can be vast, with extremely sparse solutions, which can hinder the effectiveness of traditional search methods like GP in finding the desired equation programs (Devlin et al., 2017; Parisotto et al., 2017). To tackle this challenge, we leverage the scientific knowledge and coding capabilities of LLMs to effectively navigate the program space for equation discovery. LLMs have the ability to generate syntactically correct and knowledge-guided code snippets, guiding the efficient exploration of the large equation program optimization landscape.

Let $\pi_{\theta}$ denote a pre-trained LLM with parameters $\theta$. In this task, we iteratively sample a set of equation program skeletons from LLM, $\mathcal{F}=\{f:f\sim\pi_{\theta}\}$. The goal is to find the skeleton that maximizes the reward (evaluation score) $\text{Score}_{\mathcal{T}}(f,\mathcal{D})$ for a given scientific problem $\mathcal{T}$ and data observations $\mathcal{D}$. The optimization problem can be expressed as: $f^{*}=\arg\max_{f}\mathbb{E}_{f\in\mathcal{F}}\left[\text{Score}_{\mathcal{T}}(f,\mathcal{D})\right]$, where $f^{*}$ denotes the optimal equation program discovered at the end of the LLM-guided search. LLM-SR prompts LLM to propose hypotheses given the problem specification, and experience demonstrations containing previously discovered promising equations. In LLM-SR, LLM only generates hypotheses as equation program skeletons where the coefficients or constants of each equation are considered as a placeholder parameter vector that can be optimized. The parameters are optimized using robust off-the-shelf optimizers and evaluated for their fit. Promising hypotheses are then added to a dynamic experience buffer, which guides the equation refinement process, while non-promising hypotheses are discarded. Below we explain the key components of this framework, shown in Fig. 1, in more detail.

As shown in Fig. 1 (a), the hypothesis generation step relies on a pre-trained LLM to propose diverse and promising equation program skeletons. By leveraging the LLM’s extensive knowledge and generative capabilities, we aim to efficiently explore the vast space of potential equations while maintaining syntactic correctness and scientific plausibility.

Following the hypothesis generation step, we evaluate and score the generated equation skeleton hypotheses using observed data. In this step, demonstrated in Fig. 1 (b), we first optimize the parameters of each hypothesis with regard to the dataset, and then use the optimized parameters to evaluate and score each hypothesis.

To better navigate the search landscape and avoid local minima, LLM-SR contains an experience management step, shown in Fig. 1 (c), which maintains a diverse population of high-quality equation programs in the experience buffer and effectively samples from this population to construct informative prompts for the LLM in next iteration. This step consists of two key components: a) experience manager for maintaining the experience buffer, and b) sampling from this experience buffer.

Current symbolic regression benchmarks (La Cava et al., 2021; Strogatz, 2015; Udrescu & Tegmark, 2020), mostly based on Feynman physics equations (Udrescu et al., 2020), may not effectively capture the true discovery process and reasoning capabilities of LLMs. Our findings show that LLM-SR rapidly achieves low Normalized MSE scores within very few iterations on Feynman problems, suggesting that LLMs have likely memorized these commonly known equations due to their prevalence in training data (check App. C for full results). Feynman equations are often commonly known, and easily memorized by LLMs, making them unsuitable evaluation benchmarks for this study. To overcome this limitation, we introduce novel benchmark problems across three diverse scientific domains that are designed to challenge the model’s ability to uncover complex mathematical relations while leveraging its scientific prior knowledge. These new benchmark problems are designed to simulate the conditions for scientific discovery, ensuring that the equations cannot be trivially obtained through memorization of the language model. By focusing on diverse scientific domains and introducing custom modifications to known equations, we encourage LLM-SR to creatively explore the equation search space and identify mathematical relations. We next discuss these new benchmark problems.

Nonlinear Oscillators. Nonlinear damped oscillators are ubiquitous in physics and engineering. The dynamics of oscillators are governed by differential equations that describe the relationship between the oscillator’s position, velocity, and the forces acting upon it. Specifically, the oscillator’s motion is given by a second-order differential equation: $\ddot{x}+f(t,x,\dot{x})=0$, involving time ($t$), position ($x$), and nonlinear function $f(t,x,\dot{x})$ accounts for forces. In this study, we explore two custom nonlinear forms: Oscillation 1: $\dot{v}=F\sin(\omega x)-\alpha v^{3}-\beta x^{3}-\gamma x\cdot v-x\cdot\cos(x)$; and Oscillation 2: $\dot{v}=F\sin(\omega t)-\alpha v^{3}-\beta x\cdot v-\delta x\cdot\exp(\gamma x)$, with $v=\dot{x}$ as velocity and ${\omega,\alpha,\beta,\delta,\gamma}$ as constants. These forms are designed to challenge the model’s ability to uncover complex equations, avoiding simple memorization of commonly known oscillators. More details on the design and data generation of these oscillators are provided in App. D.1.

Bacterial Growth. The growth of Escherichia coli (E. coli) bacteria has been widely studied in microbiology due to its importance in various applications, such as biotechnology, and food safety (Monod, 1949; Rosso et al., 1995). Discovering equations governing E. coli growth rate under different conditions is crucial for predicting and optimizing bacterial growth (Tuttle et al., 2021). The bacterial population growth rate can be modeled using a differential equation with the effects of population density ($B$), substrate concentration ($S$), temperature ($T$), and $pH$ level, which is commonly formulated as $\frac{dB}{dt}=f(B,S,T,pH)=f_{B}(B)\cdot f_{S}(S)\cdot f_{T}(T)\cdot f_{pH}(pH)$. To invoke the LLM’s prior knowledge about this problem, yet avoid memorization, we consider nonlinear arbitrary designs for $f_{T}(T)$ and $f_{pH}(pH)$ that share similar characteristics with the common models but are not trivial to recover from memory. More details on the specific differential equation employed and data for this problem are provided in App. D.2.

Material Stress Behavior. Understanding the stress-strain behavior of materials under various conditions, particularly the relationship between stress, strain, and temperature, is crucial for designing and analyzing structures in various engineering fields. In this problem, we use a real-world dataset from (Aakash et al., 2019), containing tensile tests on aluminum at six temperatures ranging from 20°C to 300°C. The data covers the elastic, plastic, and failure regions of the stress-strain curves, providing valuable insights into the material’s behavior. More details on the data and visualization of these stress-strain curves can be found in App. D.3.

We compare LLM-SR against several state-of-the-art symbolic regression (SR) baselines, including GPlearn¹¹1https://gplearn.readthedocs.io/en/stable/, Deep Symbolic Regression (DSR) (Petersen et al., 2021), Unified DSR (uDSR) (Landajuela et al., 2022), and PySR (Cranmer, 2023). GPlearn is a pioneering Genetic Programming (GP) based SR approach with the open-source gplearn package. DSR is a deep learning-based SR approach that employs reinforcement learning over symbolic expression sequence generations. uDSR extends DSR by incorporating large-scale pretraining with Transformers and neural-guided GP search at the decoding stage. PySR²²2https://github.com/MilesCranmer/PySR is a recent SR method that uses multi-island asynchronous GP-based evolution with the aim of scientific equation discovery. For a fair comparison with LLM-SR, we allow all baselines to run for over $2$M iterations until convergence to their best performance. More details on the implementation and parameter settings of each baseline are provided in App. A.

We experiment with GPT-3.5-turbo and Mixtral-8x7B-Instruct as 2 different LLM backbones for LLM-SR. At each iteration, we sample $b=4$ equation skeletons per prompt from the LLM using a temperature $\tau=0.8$, optimize equation skeleton parameters via BFGS ($30$s timeout per program), and sample $k=2$ in-context examples from the experience buffer for the refinement. We run LLM-SR variants for $2.5$K iterations in all experiments. For more implementation details (incl. the experience buffer structure, prompt refinement strategy, and parallel evaluation), please refer to App. A.

Table 1 shows the performance comparison of LLM-SR (using GPT-3.5 and Mixtral backbones) against state-of-the-art symbolic regression methods on various scientific benchmark problems. In these experiments, symbolic regression methods were allowed to run for considerably longer iterations (over $2$M) while LLM-SR variants ran for around $2$K iterations. The table reports the Normalized Mean Squared Error (NMSE), where lower values indicate better performance. Under the in-domain (ID) test setting, where the test set is sampled from the same domain as the data used for equation discovery, LLM-SR variants significantly outperform leading symbolic regression baselines across all benchmarks, despite running for far fewer iterations. Among the symbolic regression baselines, PySR and uDSR exhibit superior ID performance compared to DSR and GPlearn.

The ability to generate equations that generalize well to domains beyond the training data, i.e., out-of-domain (OOD) data, is a crucial characteristic of effective equation discovery methods. To assess the generalization capability of the predicted equations, we evaluate LLM-SR and the symbolic regression baselines in the OOD test setting, focusing on their ability to predict systems outside the region encountered during training. Table 1 reveals that the performance gap between LLM-SR and the baselines is more pronounced in the OOD test setting. This observation suggests that equations discovered by LLM-SR exhibit superior generalization capabilities, likely attributable to the scientific prior knowledge embedded by LLMs in the equation discovery process. For instance, on the E. coli growth problem, LLM-SR achieves an OOD NMSE of around 0.0037, considerably better than symbolic regression methods which all obtain OOD NMSE $>1$. Fig. 3 also compares the ground truth distribution of E. coli growth problem with the predicted distributions obtained from LLM-SR, PySR, and uDSR. The shaded region and black points indicate in-domain (ID) data, while the rest represent out-of-domain (OOD). Results show that the distributions obtained from LLM-SR align well with the ground truth, not only for ID data but also for OOD regions. This alignment demonstrates the better generalizability of equations discovered by LLM-SR to unseen data, likely due to the integration of scientific prior knowledge in the equation discovery process. In contrast, leading symbolic regression methods like PySR and uDSR tend to overfit the observed data, with their predicted distributions deviating significantly from the true distributions in OOD regions. This overfitting behavior highlights their limited ability to generalize beyond the training data and capture the true physical underlying patterns. For more detailed results and analyses for other benchmark problems, check App. E.

Fig. 4 shows the performance trajectories of LLM-SR variants and symbolic regression baselines across different scientific benchmark problems, depicting the maximum fitting scores achieved over search iterations. By leveraging scientific prior knowledge, LLM-SR needs to explore a considerably lower number of equation candidates in the vast optimization space compared to symbolic regression baselines that lack this knowledge. This is evident from the sharp drops in the error curves for LLM-SR variants, indicating they efficiently navigate the search space by exploiting domain knowledge to identify promising candidates more quickly. In contrast, the symbolic regression baselines show much more gradual improvements and fail to match LLM-SR’s performance even after running for over $2$M iterations, as shown by their final results in Fig. 4. The performance gap between LLM-SR and symbolic regression baselines also widens as iterations increase, particularly in the Oscillation 2 and E. coli growth problems, highlighting the effectiveness of LLM-SR’s iterative refinement process.

Fig. 5 presents an ablation study on the Oscillation 2  problem with GPT-3.5 backbone to investigate the impact of different components in the LLM-SR performance.

Problem Specification. We create a “w/o Problem Specification” variant by removing the natural language description of the problem and variables of the system from the LLM’s input prompt. In this setup, the model operates solely as an adaptive evolutionary search agent without domain-specific information about the variables and the problem being studied. We find that this variant performs worse, with increases in the Normalized MSE from 2.12e-7 to 4.65e-5 for in-domain data and from 3.81e-5 to 7.10e-3 for OOD data. This highlights the significance of incorporating prior domain knowledge into the equation discovery process.

Iterative Refinement. The “w/o Iterative Refinement” variant is equivalent to the LLM sampling baseline, which generates more than $2$K samples for the initial LLM-SR’s prompt without computing feedback and updating in-context examples over the sampling span. The absence of iterative refinement leads to a substantial performance drop, with the Normalized MSE increasing to 1.01e-1 for in-domain data and 1.81e-1 for OOD data. This emphasizes the importance of the evolutionary search and iterative refinement process in LLM-SR’s success.

Coefficient Optimization. We also create a “w/o Coeff Optimization” variant that requires the LLM to generate complete equations, including numerical coefficients and constants, in a single end-to-end step. This helps us study the role of parameter optimization using external optimizers in the equation discovery process. We find that this variant shows significantly worse results, with the Normalized MSE increasing to 3.75e-1 for in-domain data and 3.78e-1 for OOD data. These results indicate that the two-stage approach of generating equation skeletons followed by data-driven skeleton parameter optimization is essential for effective performance, as it helps navigate the intricate combinatorial optimization landscape involving both continuous parameter values and discrete equation structures.

Optimization method. LLM-SR relies on direct and differentiable parameter optimization, a capability not present in current symbolic regression methods. We compare two different optimization frameworks: numpy+BFGS and torch+Adam. The numpy+BFGS variant yielded better-performing equations on both ID (2.12e-7) and OOD (3.81e-5) evaluations compared to torch+Adam (ID: 1.60e-6, OOD: 2.4e-4). However, our initial analysis indicated that this discrepancy could be primarily due to the LLM’s higher proficiency in generating numpy code compared to pytorch, rather than the superiority of one optimization method over the other. Combining LLM-SR with LLM backbones that are better in generating pytorch code could potentially leverage differentiable parameter optimization to achieve better performance in future works.

Fig. 6 presents the final discovered equations for both Oscillation problems using LLM-SR and other symbolic regression baselines. A notable observation is that equations discovered by LLM-SR have better recovered the symbolic terms of the true equations compared to baseline models. Moreover, LLM-SR provides explanations and reasoning steps based on the scientific knowledge about the problem, leading to more interpretable terms combined as the final equation. For example, in both problems, LLM-SR identifies the equation structure as a combination of driving force, damping force, and restoring force terms, relating them to the problem’s physical characteristics. In contrast, baselines (DSR, uDSR, PySR) generate equations lacking interpretability and understanding of the physical meanings or relations behind the terms. The equations appear as combination of mathematical operations and variables without clear connection to the problem’s underlying principles. More detailed results and qualitative analysis on the discovered equations in other benchmark problems and across performance progression curves can also be found in App. E.