In-Context Symbolic Regression: Leveraging Large Language Models for Function Discovery (2024)

Abstract: State of the art Symbolic Regression (SR) methods currently build specialized models, while the application of Large Language Models (LLMs) remains largely unexplored. In this work, we introduce the first comprehensive framework that utilizes LLMs for the task of SR.We propose In-Context Symbolic Regression (ICSR), an SR method which iteratively refines a functional form with an LLM and determines its coefficients with an external optimizer. ICSR leverages LLMs’ strong mathematical prior both to propose an initial set of possible functions given the observations and to refine them based on their errors.Our findings reveal that LLMs are able to successfully find symbolic equations that fit the given data, matching or outperforming the overall performance of the best SR baselines on four popular benchmarks, while yielding simpler equations with better out of distribution generalization.

## introduction classical machine learning regression methods can be divided into two broad categories: statistical methods, which learn an implicit statistical (black-box) model of the relationship between the observations, and rule-based methods, which instead attempt to extract an explainable set of rules that explicitly model the transformation between the inputs and outputs @xcite . symbolic regression (sr) is a particular subset of the latter category, which searches the set of all possible explicit mathematical expressions to find the equation that best fits the given set of observations. this has the clear advantage of explainability, as well as a potential for better generalization, if the trend holds outside of the observed data. the traditional approach for sr algorithms is genetic programming @xcite (gp), which combines fundamental blocks for mathematical expressions (e.g., basic operators, trigonometric functions, etc.) into more complex formulas using strategies borrowed from evolutionary biology, such as mutations and fitness. the recent success of transformer models, first introduced by @xcite , has revolutionized various fields of artificial intelligence, notably natural language processing @xcite @xcite and computer vision @xcite . transformer-based methods have also been proposed for sr @xcite , typically by employing a model pre-trained on a large amount of synthetic sr datasets. large language models (llms), also based on the transformer, have proven to possess unprecedented reasoning and generalization abilities, based on their capability for in-context learning (icl) @xcite . this refers to the ability to perform tasks based on the context provided in the input text without any additional fine-tuning. with the help of icl, these models can be leveraged for a wide range of different tasks, suggesting a potential use case for symbolic regression. in this paper, we examine the integration of llms into the sr pipeline, with the aim of using them to search for new equations that could fit the data. inspired by the optimization by prompting (opro) approach presented by @xcite , we propose in-context symbolic regression (icsr) foot_0 . this approach leverages pre-trained language models by providing a number of previously tested equations and their fitness scores in the prompt, tasking them to generate a new candidate that could be a better fit. the method is repeated until convergence is reached or the computational budget is exhausted. to the best of our knowledge, only a contemporary work by @xcite has ever explored the use of llms for sr. however, they focus on working with equations from a scientific domain where natural language knowledge can be directly incorporated, while this work aims to generally explore the capabilities of llms for sr without any additional information, in order to lay a foundation that can be expanded later. we discuss in depth the differences between the two works in section 2. our approach presents several advantages compared to models specifically trained for sr: as the llm is not fine-tuned for this task, improvements in the underlying base model can improve icsr without any changes to the method itself. further, llms provide a natural language interface that can be leveraged to include additional information about the problem, like the domain of the equation and the interpretation of the observation values. the models could also be asked to explain the reasoning behind the proposed functions, potentially leading to a more interpretable process. in summary, we make the following contributions: 1) we propose icsr, the first general framework to leverage llms for the sr task. 2) we compare the method with a range of competitive sr baselines, matching or outperforming state of the art results on four popular sr benchmarks: nguyen @xcite , constant @xcite , r (krawiec and pawlak, 2013) and keijzer @xcite . 3) we show that the equations generated with our method tend to exhibit lower complexity, which correlates with stronger out of distribution performance. more recently, deep learning methods have been applied to enhance the available toolkit for sr. udrescu and tegmark (2020) proposed an iterative simplification of the problem relying on insights from physics and outsourcing the function approximation part to a neural network. petersen et al. (2021) used a recurrent neural network (rnn) with a risk-seeking policy to perform a hierarchical search over the space of user-defined operators and mathematical functions. the main drawback of these methods, including gp, is the fact that the algorithms start from scratch for every new expression, with very limited abilities of knowledge preservation between tasks. to address this limitation, numerous transformer-based methods inspired by language modelling have been developed. symbolicgpt by @xcite , nesymres by @xcite and cl-sr by @xcite proposed different generative transformer models specifically trained for sr. these models generate a functional form ('skeleton') of the equation with a special token for coefficients which are fitted via an external numerical optimizer. subsequently, @xcite presented e2e, a transformer model able to produce the full expression including the coefficient values. while retaining knowledge between tasks, transformer-based methods are quite limited in refining their solutions for the given set of points. to this end, @xcite presented a method integrating a pre-trained transformer with monte carlo tree search to guide the equation generation merging the strength of the search and model pre-training. the proposed framework can be also viewed as a combination of a pre-trained model and an iterative refinement process. however, none of the prior methods employ a foundation model @xcite , such as an llm, in order to leverage mathematical knowledge, but either pre-train an sr model @xcite , or learn from scratch for every new function @xcite . mathematical reasoning with llms. as llms form the backbone of the method presented in this work, we rely entirely on their mathematical reasoning capabilities, such as icl @xcite , to explore the solution space. @xcite show that llms are able to recognize patterns from in-context examples and can extrapolate them to complete related tasks in the input. similarly, @xcite find that llms can extrapolate zero-shot the pattern from a timeseries (although they do not extract any functional representation). furthermore, @xcite present a study in which they find that transformer models can learn higher-order optimization methods (similar to newton's method). contemporary to our work, @xcite also propose to perform sr with an llm aided by an external coefficient optimizer. however, they focus exclusively on the case where llms can leverage scientific knowledge for sr, by including a description of the input and output variables in the llm prompt. in contrast, we focus on the general case where no extra knowledge is given and test on standard benchmarks within the sr community and include a wider range of established baselines, with the aim to directly evaluate the capability of llms on the task of sr. furthermore, we propose advancements in the structure of the prompt, including the coordinates of the points to be regressed, the score of previous attempts, and more in-context examples. finally, rather than asking the llm to optimize a mean squared error (mse) objective, we employ a more advanced loss function, presented in section 4, which jointly optimizes for the accuracy and complexity of the function for improved generalization properties. ## background the optimization by prompting (opro) framework was introduced by @xcite for prompt optimization, i.e., for increasing the performance of models (such as llms) that receive a textual prompt in the input and have to perform a specific task (such as mathematical reasoning). closer to our interest, the authors also present experiments on classical optimization problems (linear regression and travelling salesman problem), suggesting that opro can solve such tasks. the key idea of the method is the use of a socalled meta-prompt, a higher level prompt which contains a description of the task to be optimized and previous attempts (examples) in solving it with their corresponding scores. an example of such task can be querying the model to find a linear function that fits a set of points. in this case, the prompt is augmented by the functions that have been tried out and the mean squared error on the data, obtained with an external evaluation procedure. the assumption behind it is that llms have the ability to extrapolate the pattern formed by the examples, thanks to icl, and propose a better alternative. the meta-prompt is given as input to the llm and the model's output is then evaluated and added back to the meta-prompt if the score is good enough. this approach can then be iterated until a satisfying result is achieved or a certain computational budget is exhausted. ## method we consider a regression dataset d with n observations and target variables {x i , y i } n @xmath0, also denoted with (x, y ) more compactly, to be used for producing a function f that well approximates the data. to leverage the opro approach for sr, we need to design a meta-prompt suitable for the task and fill it with the available observations (x, y ), an initial set of k functions @xmath1, f (2) 0 , . . . , f (k) 0 } (either hand-written or model-generated) and a measure of their fitness (score) on d. for our purposes, we frame the refinement process as a minimization problem over an objective function, also called the error function, such that generated equations with the lowest error have the highest fitness. the goal is then to iteratively refine the set of functions fi , @xmath2, . . . , n] for each iteration i, until a sufficiently low error is obtained by one of them or a maximum number of iterations is reached; we denote this process as the optimization loop. due to the finite size of the llm context window, we only keep the k best performing previous attempts in the set fi , where k is a design choice (@xmath3. the meta-prompt used in the experiment can be found in appendix c. ## experiments we empirically evaluate icsr and compare it against a set of competitive baselines, checking both in-domain and out of distribution performance of the proposed approach. ## discussion optimizing for out of distribution. a general framework for optimizing the out of distribution performance of a predictive model (such as a symbolic equation) is to regularise its complexity, following the occam's razor principle that simpler explanations are preferable to more complex ones, all other things being equal. in our work we use the working definition of complexity as the number of nodes in the expression tree of an equation. however, more optimal choices could be available: for instance, equations containing expressions not defined on all the real domain (such as logarithms and square roots) could be penalised more, as they could be undefined when extrapolating to larger domains. knowing in advance the full domain in which an equation is supposed to hold could also greatly improve out of distribution performance by filtering out invalid candidate functions. in the case of icsr, it could also be leveraged as extra information by the llm. furthermore, we observe that numerous equations that we derive with icsr have extra terms with very small coefficients (e.g. o(10 -3 )) that do not contribute significantly to the shape of the equation and could be safely suppressed, resulting in expressions with a lower complexity. this could be done by modifying the optimization procedure of the coefficients, to eliminate coefficients under a certain threshold, which would be a hyperparameter of the method. vocabulary. in general, most sr methods are limited to a predefined vocabulary of operators and tokens, while llms can virtually explore any possible function and combination. an example of this is with the x x 2 1 function in the nguyen benchmark: in @xcite , the authors mention that it is not included in the set of equations that their model can fit, while our approach can recover the exact expression. we also observe a similar trend with the other baselines for this specific expression. in our prompts (see appendix c) we include a vocabulary for the llm, but this is meant more to guide the llm into the correct search space and is by no means a hard restriction: for example, we observe that icsr can produce the erf function even if it wasn't reported in this list. furthermore, any function that can be found in the model's pre-training corpus (fundamentally the in-ternet) can be potentially added to the prompt at any time if desired, which is impossible for other fixed-vocabulary methods. ## conclusion we show that llms paired with the icsr approach are able to perform symbolic regression tasks on classical sr benchmarks. the proposed method matches or outperforms a variety of established sr baselines, while producing simpler expressions that more closely resemble the complexity of the ground truth equations and result in better out of distribution performance. this work exposes yet another task that llms can be leveraged for, thanks to specialized techniques such as icsr, and shows promise for integrating these models with mathematical reasoning methods.