Discovering ordinary differential equations that govern time-series

Sampling symbolic expressions.  To exploit large-scale supervised pretraining we generate a dataset of $\sim$63M ODEs in symbolic form along with numerical solutions for randomly sampled initial values. Since we assume ODEs to be in canonical form $\dot{y}=f(y)$, generating an ODE is equivalent to generating a symbolic expression $f(y)$. We follow Lample and Charton [12], who sample such an expression $f(y)$ as a unary-binary tree, where each internal node corresponds to an operator and each leaf node corresponds to a constant or variable. The algorithm consists of two phases: (1) A unary-binary tree is sampled uniformly from the distribution of unary-binary trees with up to $k\in\mathbb{N}$ internal nodes, which crucially does not overrepresent small trees corresponding to short expressions. Here the maximum number of internal nodes $K$ is a hyperparameter of the algorithm. (2) The sampled tree is “decorated”, that is, each binary internal node is assigned a binary operator, each unary internal node is assigned a unary operator, and each leaf is assigned a variable or constant. Hence, we have to specify a distribution over the $N_{\mathrm{bin}}$ binary operators, one over the $N_{\mathrm{una}}$ unary operators, a probability $p_{\mathrm{sym}}\in(0,1)$ to decide between symbols and constants, as well as a distribution $p_{\mathrm{c}}$ over constants. For constants we distinguish explicitly between sampling an integer or a real value. Together with $K$, these choices uniquely determine a distribution over equations $f$ and are described in detail in Appendix B. Figure 1 depicts an overview of the data generation procedure.

The pre-order traversal of a sampled tree results in the symbolic expression for $f$ in prefix notation. After conversion to infix notation, we simplify each expression using the computer algebra system SymPy [16], and filter out constant equations $f(y)=c$ as well as expressions that contain operators or symbols that were not part of the original distribution. We call the structure modulo the value of the constants of such an expression a skeleton. Any skeleton containing at least one binary operator or constant can be represented by different unary-binary trees. Vice versa many of the generated trees will be simplified to the same skeleton. To ensure diversity and to mitigate potential dataset bias towards particular expressions, we discard duplicates on the skeleton level. To further cheaply increase the variability of ODEs we sample $N_{\mathrm{const}}$ unique sets of constants per skeleton. When sampling constants we take care not to modify the canonical expression by adhering to the rules listed in Section B.1. We provide summary statistics on operator frequencies and expression complexities for the generated dataset in Appendix C. Here, complexity refers to overall count of symbols (e.g., $y$, or constants) as well as operators in an expression, a simple yet common measure in the symbolic regression literature.

Computing numerical solutions.  We obtain numerical solutions for all ODEs via SciPy’s interface [25] to the LSODA software package [7] with both relative and absolute tolerances set to $10^{-9}$. We solve each equation on a fixed time interval $t\in[0,T]$ and store solutions on a regular grid of $N_{\mathrm{grid}}$ points. For each ODE, we sample up to $N_{\mathrm{iv}}$ initial values $y(0)=y_{0}$ uniformly from $(y_{0}^{\min},y_{0}^{\max})$.¹¹1Due to a timeout per ODE, fewer solutions may remain in cases when the numerical solver fails repeatedly. While LSODA attempts to select an appropriate solver, numerical solutions still cannot be trusted in all cases. Therefore, we check the validity of solutions via the following quality control check: we use 9th order central finite differences to approximate the temporal derivative of the solution trajectory (on the same temporal grid as the proposed solution), denoted by $\dot{y}_{\mathrm{fd}}$, and filter out any solution for which $\|\dot{y}_{\mathrm{fd}}-\dot{y}\|_{\infty}>\epsilon$, where we use $\epsilon=1$.

NSODE consists of an encoder-decoder transformer with architecture choices listed in Section B.3. We provide a visual overview in Figure 1.

Representing input trajectories.  A key difficulty in feeding numerical solutions $\{y_{i}\}_{i=1}^{n}$ as input is that their range may differ greatly both within a single solution as well as across ODEs. For example, the linear ODE $\dot{y}=c\cdot y$ for a constant $c$ is solved by an exponential $y(t)=y_{0}\exp(ct)$ for initial value $y(0)=y_{0}$, which may span many orders of magnitude on a fixed time interval. To prevent numerical errors and vanishing or exploding gradients caused by the large range of values, we assume each representable 64-bit float value is a token and use its IEEE-754 encoding as the token representation [2]. We thus convert all pairs $(t_{i},y_{i})$ to their IEEE-754 64 bit representations, channel them through a linear embedding layer before feeding them to the encoder.

Representing symbolic expressions.  The target sequence (i.e., the string for the symbolic expression of $f$) is tokenized on the symbol-level. We distinguish two cases: (1) Operators and variables: for each operator and variable we include a unique token in the vocabulary. (2) Numerical constants: constants may come from both discrete (integers) as well as continuous distributions, as for example in y**2+1.64*cos(y). Hence, it is unfeasible to include individual tokens “for each constant”. Naively tokenizing on the digit level, i.e., representing real values literally as the sequence of characters (e.g., "1.64"), not only significantly expands the length of target sequences and thus the computational cost, but also requires a variable number of prediction steps for every single constant.

Instead, we take inspiration from Schrittwieser et al. [19] and encode constants in a two-hot fashion. We fix a finite homogeneous grid on the real numbers $x_{1}<x_{2}<\ldots<x_{m}$ for some $m\in\mathbb{N}$, which we add as tokens to the vocabulary. The grid range $(x_{1},x_{m})$ and number of grid points $m$ are hyperparameters that can be set in accordance to the problems of interest. Our choices are described in Section B.3.

For any constant $c$ in the target sequence we then find $i\in\{1,\ldots,m-1\}$ such that $x_{i}\leq c\leq x_{i+1}$ and encode $c$ as a distribution supported on $x_{i},x_{i+1}$ with weights $\alpha,\beta$ such that $\alpha x_{i}+\beta x_{i+1}=c$. That is, the target in the cross-entropy loss for a constant token is not a strict one-hot encoding, but a distribution supported on two (neighboring) vocabulary tokens resulting in a lossless encoding of continuous values in $[x_{1},x_{m}]$. This two-hot representation can be used directly in the cross-entropy loss function.

Decoding constants.  When decoding a predicted sequence, we check at each prediction step whether the argmax of the logits corresponds to one of the $m$ constant tokens $\{x_{1},\ldots,x_{m}\}$. If not, we proceed by conventional one-hot decoding to obtain predicted operators and variables. If instead the argmax corresponds to, for example, $x_{i}$, we also pick its largest-logit neighbor ($x_{i-1}$ or $x_{i+1}$; suppose $x_{i+1}$), renormalize their probabilities by applying a softmax to all logits and use the resulting two probability estimates as weights $\alpha,\beta$. Constants are then ultimately decoded as $\alpha x_{i}+\beta x_{i+1}$.

Sampling solutions.  To infer a symbolic expression for the governing ODE of a new observed solution trajectory $\{(t_{i},y_{i})\}_{i=1}^{n}$, all the typical policies such as greedy, sampling, or beam search are available. In our evaluation, we use beam search with 1536 beams and report top-$k$ results with $k$ ranging from 1 to 1536.

Metrics.  We evaluate model performance both numerically and symbolically. For numerical evaluation we follow Biggio et al. [2]: suppose the ground truth ODE is given by $\dot{y}=f(y)$ with (numerical) solution $y(t)$ and the predicted ODE is given by $\hat{\dot{y}}=\hat{f}(y)$. To compute numerical accuracy we first evaluate $f$ and $\hat{f}$ on $N_{\mathrm{eval}}$ points in the interval $[\min(y(t)),\max(y(t))]$ (i.e., the interval traced out by the observed solution), which yields function evaluations $\texttt{gt}=\{\dot{y}_{i}\}_{i=1}^{N_{\mathrm{eval}}}$ and $\texttt{pred}=\{\hat{\dot{y}}_{i}\}_{i=1}^{N_{\mathrm{eval}}}$. We then assess whether numpy.allclose²²2numpy.allclose returns True if abs(a - b) <= (atol + rtol * abs(b)) holds element-wise for elements $a$ and $b$ from the two input arrays. We use atol=1e-10 and rtol=0.05; $a$ corresponds to predictions, $b$ corresponds to ground truth. returns True as well as whether the coefficient of determination $\mathrm{R}^{2}\geq 0.999$.³³3For observations $y_{i}$ and predictions $\hat{y}_{i}$ we have $\mathrm{R}^{2}=1-(\sum_{i}(y_{i}-\hat{y}_{i})^{2})/(\sum_{i}(y_{i}-\overline{y})^{2})$. Numerical evaluations capture how closely the predicted function approximates the ground truth function within the interval $[\min(y(t)),\max(y(t))]$.

However, a key motivation for symbolic regression is to uncover a symbolic mathematical expression that governs the observations. Testing for symbolic equivalence between ground truth expression $f(y)$ and a predicted expression $\hat{f}(y)$ is unsuitable in the presence of real-valued constants as even minor deviations between true and predicted constants render the equivalence false. Instead, we regard the predicted expression $\hat{f}(y)$ to be symbolically correct if $f(y)$ and $\hat{f}(y)$ can be made equivalent by modifying only the values of constants in the predicted expression $\hat{f}(y)$. This is implemented using SymPy’s match function. In order not to alter the structure of the predicted expression, we constrain modifications of constants such that all constants remain non-zero and retain their original sign. This definition is thus primarily concerned with the structure of an expression, rather than precise numerical agreement. Once the structure is known, the inference problem becomes conventional parameter estimation. We report percentages of samples in a given test set that satisfies any individual metric (numerical and symbolic), as well as percentages satisfying symbolic and numerical metrics simultaneously.

We construct several test sets to evaluate model performance and generalization in different settings.

• •

  testset-iv: Our first test set assesses generalization within initial values not seen during training. It consists of 5793 ODEs picked uniformly at random from our generated dataset but re-sampled initial values. We also employ the following constraints via rejection sampling: (a) All skeletons in testset-iv are unique. (b) As the number of unique skeletons increases with the number of operators, we allow at most 2000 examples per number of operators (with substantially fewer unique skeletons existing for few operators).

• •

  testset-constants: Our second test set assesses generalization within unseen initial values and constants. It consists of 2720 ODEs picked uniformly at random from our dataset (ensuring unique skeletons and at most 1000 examples per number of operators as above), but re-sampled intial values and constants.

• •

  testset-skeletons: In principle, we can train NSODE on all possible expressions (using only the specified operators and number ranges) up to a specified number of operators. However, even with the millions examples in our dataset, we have by far not exhausted the huge space of possible skeletons (especially for larger numbers of operators). Hence, our third test set contains 100 novel random ODEs with skeletons that were never seen during training.

• •

  testset-iv-163: This is a subset of testset-iv motivated by the fact that most symbolic regression models we want to compare to require a separate optimization for every individual example, which was computationally infeasible for our testset-iv. For a fair comparison, we therefore subsampled up to 10 ODEs per complexity uniformly at random, yielding 163 examples.

• •

  testset-textbook: To assess how NSODE performs on “real problems”, we manually curated 12 scalar, non-linear ODEs from Wikipedia pages, physics textbooks, and lecture note from university courses on ODEs. These equations are listed in Table 7 in Appendix D. We note that they are all extremely simple compared to the expressions in our generated dataset in that they are ultimately mostly low order polynomials, some of which with one fractional exponent.

• •

  testset-classic: To validate our approach on existing datasets we turn to benchmarks in the classic symbolic regression literature (inferring just the functional relationship between input-ouput pairs) and simply interpret functions as ODEs. In particular we include all scalar function listed in the overview in [15] which includes equations from many different benchmarks [9, 10, 11, 22, 26]. For example, we interpret the function $f(y)=y^{3}+y^{2}+y$ from Uy et al. [22] as an autonomous ODE $\dot{y}(t)=f(y(t))=y(t)^{3}+y(t)^{2}+y$, which we solve numerically for a randomly sampled initial value as described before.

We compare our method to recent popular baselines from the literature (see Section 2). We briefly describe them including some limitations here and defer all details to Appendix E. First, no baseline is suited directly to infer dynamical laws, but only to infer functional relationships. Therefore, all baselines fit a separate regression function mapping $y(t)\mapsto\hat{\dot{y}}(t)$ per individual ODE, using the coefficient of determination $\mathrm{R}^{2}$ as optimization objective. Since derivatives $\hat{\dot{y}}(t)$ are typically not observed, we approximate them via finite differences using PySindy [6]. Hence, all these methods crucially rely on regularly sampled and noise-free observations, whereas our approach can easily be extended to take those into account (see Appendix A).

• •

  Sindy [3]: Sindy builds a (sparse) linear combination of a fixed set of (non-linear) basis functions. The resulting Lasso regression is efficient, but suffers from limited expressiveness. In particular, Sindy cannot easly represent nested functions or non-integer powers as all non-linear expressions have to be added explicitly to the set of basis functions. We cross-validate Sindy over a fairly extensive hyperparameter grid of 800 different combinations for each individual trajectory.

• •

  GPL⁴⁴4gplearn.readthedocs.io/ (genetic programming): GPL(earn) maintains a population of programs each representing a mathematical expression. The programs are mutated for several generations to heuristically optimize a user defined fitness function. While not originally developed for ODEs, we can apply GPLearn on our datasets by leveraging the finite difference approximation. We use a population size of 1000 and report the best performance across all final programs. Compared to sindy, GPLearn is more expressive yet substantially slower to fit.

• •

  AIFeynman [20, 21]: AIFeynman is a physics-inspired approach to symbolic regression that exploits the insight that many famous equations in natural sciences exhibit well-understood functional properties such as symmetries, compositionality, or smoothness. AIFeynman implements a neural network based heuristic search that tests for such properties in order to identify a symbolic expression that fits the data. For every test sample AIFeynman computes a pareto front of solutions that trade off complexity versus accuracy. We report the best performance across all functions on the pareto front. Notably, AIFeynman performs quite an exhaustive search procedure such that running it even on a single equation took on the order of tens of minutes.

Model Performance.  Figure 2 shows NSODE’s performance on our testset-iv, testset-constants, and testset-skeletons according to our numerical and symbolic metrics as well as combined skeleton recovery and allclose as we vary $k$ in the top-k considered candidates of the beam search. Investing sufficient test-time-compute (i.e., considering many candidates) continuously improves performance. While we capped $k$ at 1536 due to memory limitations, we did not observe a stagnation of the roughly logarithmic scaling of all performance metrics with $k$. This cannot be attributed to “exhaustion effects”, where one may assume that all possible ODEs will eventually be among the candidates, because (a) the space of possible skeletons grows much faster than exponentially, and (b) the numerical metrics are extremely sensitive also to the predicted constant values in continuous domains.

As one may expect, performance decreases as we move from only new initial conditions, to also sampling new constants, to finally sampling entirely unseen skeletons. On testset-iv with $k=1536$ we achieve about 50% skeleton recovery and still successfully recover more than a third skeletons of testset-skeletons with similar numbers for allclose. The fact that the combined metric (symbolic + numerical) is only about half of that indicates that numerical and symbolic fit are indeed two separate measures, none of which need to imply the other. Hence, a thorough evaluation of both is crucial to understand model performance in symbolic regression tasks.

Comparison to baselines.  In Table 1 we compare NSODE to all baselines using $k=1536$ in our beam search; full results on all datasets can be found in Appendix F. We also include the average wallclock runtime per expression for each of the datasets.

First, we note that on our subsampled testset-iv-163, NSODE outperforms competing approaches in terms of skeleton recovery by a wide margin and also performs best in terms of joint skeleton recovery and numerical measures, which is a strong indication of actually having recovered the governing ODE accurately. By spending over 200x more time on its exhaustive heuristic search, AIFeynman manages to outperform NSODE in terms of numerical accuracy ($\mathrm{R}^{2}$ and allclose). Figure 3 shows the number of skeletons recovered by each method given the complexity of equations, results for other datasets can be found in Appendix F. ⁵⁵5Due to simplification, complexity is not upper bounded by the number of nodes in a unary-binary tree. While AIFeynman and Sindy recover some of the low complexity expressions, NSODE is the only method to also recover some of the more complex skeletons.

On testset-textbook, AIFeynman outperforms all other methods on numerical and symbolic metrics. This can be understood with regard to the dataset where $8/12$ expressions are polynomials with the remaining 4/12 expressions having a polynomial skeleton with fractional or negative exponents. These expressions are particularly favorable for the heuristics implemented by AIFeynman which explicitly attempt to fit a polynomial to the data. However, even on these simple examples AIFeynman takes over 200x longer than our method, which in turn clearly outperforms Sindy and GPLearn in terms of skeleton recovery.