Differentiable composition for model discovery

Gilles Louppe
Montefiore Institute
University of Liègeg.louppe@uliege.be

Omer Rochman SharabiMontefiore Institute
University of Liègeo.rochman@uliege.be

Abstract
We propose DiffComp, a symbolic regressor that can learn arbitrary function compositions, including derivatives of various orders. We use DiffComp in conjunction
with a Physics Informed Neural Network (PINN) to discover differential equations
from data. DiffComp has a layered structure where a set of user-defined basic
functions are composed up to a specified depth. As it is differentiable, it can betrained using gradient descent. We test the architecture using simulated data from
common PDEs and compare to existing model discovery frameworks, including
PySINDy and DeePyMoD. We then test on oceanographic data.

1

Introduction

Plenty of processes, many of them physical systems, can be modeled using differential equations;wave propagation, heat dissipation, the motion of celestial bodies, or car traffic are some examples.
Usually, a scientist, using hard-to-come-by domain knowledge, designs a model and then runs
simulations (ranging from pen and paper calculations to programs on supercomputers) which arelater compared to real data. While this process works and has given us most of the models we use
today, it has limitations. As we move on to more complex systems, it may be beneficial to be able tolearn the dynamics that govern a given system from data, i.e. observations of the system.
Previous works approached the problem of model discovery from multiple angles [3, 16]. One
approach is black box regression, where a black box model (usually some neural network) canbe used to learn the dynamics of a system [9, 15, 17]. This approach, owing to the power of its
black box approximators, allows us to learn complex dynamics. On the other hand, it suffers fromcertain drawbacks; most notably failure to predict new out-of-training-distribution data, and lack ofhuman interpretability. Several works propose to mitigate those problems by learning a symbolic
representation of the model dynamics. A black box model is used in [11] to learn the dynamics
of bodies in the Solar System, then a tree-based symbolic regressor, PySR [7], is used to find asymbolic expression for the black box model, and finally that symbolic expression is fine tuned.
Biggio et al. [2] propose a transformer-based symbolic regressor. Both and Kusters [5] introduce afully differentiable model discovery algorithm based on PINNs [13] and Sparse Bayesian Learning(SBL) [6]. In the latter framework, one defines a library of possible terms that they think the model
might include and then performs sparse linear regression on those terms.
Model discovery is usually framed as follows: we observe data generated by an unknown process,generally an ODE or a PDE, and want to discover the underlying dynamics of the process. Formally,
we have ODEs: t ∈ R, u(t) ≡ u, 0 = F (u, t); and PDEs: t ∈ R, x ∈ D ⊂ Rd , u(t, x) ≡ u ∈ R,
0 = F (u, t, x). The unknown operator F may include partial derivatives of u with respect to x and t.
It is the goal of model discovery to learn this F . In the following we will denote both u(t, x) and∂uu(t) simply as u, write dudt ≡ u̇ and ∂t ≡ ut , and focus on the subproblem ut = F (u, t). Boldsymbols are vectors, and in particular u ∈ RK is the vector u({ti , xi }K1 ), where K is the number ofobservations or, during training, the batch size. Some known ODE and PDE examples from physics
Machine Learning and the Physical Sciences workshop, NeurIPS 2022.

(r)are Newton’s law r̈ = Fm
, the heat equation ut = ∆u and the wave equation utt = c2 ∆u, where
the laplacian ∆u = ∇ · (∇u). In this work, we build upon [5] and extend their model by substituting
the SBL component with DiffComp, a fully differentiable architecture that can learn composite
functions and is therefore not limited in its search by the user-defined library.

2

DiffComp

We combine DiffComp, a differentiable symbolic regressor, with PINNs, following [5] but substituting
the SBL component. PINNs is a method to solve differential equations in which a NN ansatz is placedon the solution, making solving the equation equivalent to minimizing the PINN loss [13], which
consists of a residual term, a boundary term, and an initial term. Another option, more suited to ourcontext, is to constrain the equation using measured data. Having access to a dataset {ui , xi , ti }N
1 ofmeasurements, the new loss becomes
1 X
1 X
L=
∥∂t ûi − F (ûi , ti )∥2 +∥ûi − ui ∥2 ,
(1)
N i
N i
where û is our approximation of u, which will usually be a neural network. Both and Kusters [5] domodel discovery by adding an extra term to the loss, given by ∥∂t û − Θξ∥2 , where Θ is a libraryof manually defined candidate functions and ξ a vector of coefficients. Using SBL, ξ can be madesparse. We substitute SBL with DiffComp.
DiffComp starts with a user-defined library of basic functions, and attempts to learn a compositionof those functions to approximate some objective. The library is arranged in a layered graph. A
layer consists of nodes that are fully connected to the previous layer via weighted edges. Each nodecontains one function from the user-defined library. During the forward pass, the outputs of theprevious layers are weighted by the edges of a node and fed into the node as inputs to its function. By
stacking layers we can compose complex functions even if the library contains only simple ones.
Mathematically, we are trying to learn F in F (u, t) = ut . Let z 0 := [û, t, x, ...] ∈ Rn be theinputs, where [û, t, x, ...] ∈ RM +M +M ×d+...=n is a concatenation of û, t, x and possibly other
inputs, such as terms known through domain expertise. Further, let z k := fk (ξ k z k−1 ) ∈ Rmk be
the output of the k-th layer and M the total number of layers. The layers are defined componentwise as z ik = fki (ξ ik z k−1 ) ∈ R, where the f i s are user-defined functions such as derivatives,
arithmetic operations or trigonometric functions, ξ ik ∈ Rmk−1 are sparse learnable weights, andz k = (z 1 , ..., z mk ) ∈ Rmk . Each node selects, using ξ ik , which components of the previous layer touse. In summary, the differentiable composition is given byû, t, x → f1i (ξ i1 [û, t, x]) = z1i → z 1 → f2i (ξ i2 z 1 ) = z2i → z 2 → ... → z M .

(2)

We remark that, in contrast to other frameworks, by virtue of the layered composition and PyTorch’s[12] automatic differentiation, we can take derivatives not just of u but of any composition ofi∂f i (ξ z)functions i.e. k ∂kx k−1 (the x dependence is implicit in z k−1 ).
This architecture works best when learning a model where the dimensionality of the input x is highand the terms that appear in the real model are shallow compositions of library functions. The
reason is that building a library containing all possible (shallow) terms combining the inputs becomesprohibitive for other methods, while ours can find those compositions. However, we expect thearchitecture to perform poorly when the terms that appear in the equation are a very deep compositionof library functions, making domain knowledge necessary both for the method presented here andfor other library-based methods. In that case, prior knowledge can be incorporated into the model.
The model can also be used with a human in the loop that runs the model on pure inputs, selects
relevant terms from the outputs of the model, and then uses those terms as inputs again to discover
more complex terms. We have used this last training methodology with oceanographic data.
During training, we assume we observe points x ∈ R

3Experiments

To evaluate the model discovery capacity of the DiffComp, we pick two well-known PDEs, generate
datasets according to those equations, and attempt to learn them. This task can also be solved using
2

Method

Discovered equation

RuntimeMSE< 1s∼ 10min∼ 5.5min

0.05
9.38 × 10−5
4.01 × 10−5

< 1s
∼ 5min∼ 4min0.04
1.57 × 10−4
2.09 × 10−6

2D heat equation: ut = uxx + uyy
PySINDy
DeePyMod
DiffComp

0.014uxx + 0.014uyy
0.255 + 0.556uxx + 0.606uyy − 0.627u − 0.201uuxx − 0.239uuyy
0.971uxx + 0.981uyy

1D Burger’s equation: ut = 0.05uxx − uux
PySINDy
DeePyMod
DiffComp

0.05uxx − 0.992uux−0.336ux + 0.198uxx − 0.687uux
0.05uxx − 0.994uux

Ross et al. (w1 ∇2 + w2 ∇4 + w3 ∇6 )(u · ∇)q + (w4 ∇4 + w5 ∇6 )q + (u · ∇)∇2 (w6 vx + w7 vy ) [14]1st iteration
w1 ∇2 (u · ∇)q + w2 ∇4 q
∼ 3min
2nd iteration
w1 (u · ∇)∇4 (u · ∇)q + (u · ∇)∇4 q + w3 ∇6 (u · ∇)q + w4 ∇8 q
∼ 2.5min

0.996
0.764

Table 1: Comparison between methods. PySINDy runs almost instantly while DeePyMod and
DiffComp have similar training times, on the order of 10min, as both involve training a neuralnetwork and taking its derivatives. Note that we have not optimized the parameters of PySINDy and
DeePyMod.

existing library-based methods, such as PySINDy [8, 10] and DeepMoD [4]. We also test DiffComp
on a more complex dataset of oceanographic data . Our library contained the following functions:addition, multiplication, and differentiation in the PDE tests, and addition, laplacian and advection in
the oceanographic benchmark. A summary is presented in Table 1. Some of the tests were repeatedwith added noise and the results are shown in Table 2 in the appendix.
The heat equation, ut = uxx + uyy is a diffusion equation describing how a quantity diffuses in a
medium over time. Our model finds the equation ut = 0.971uxx + 0.981uyy . We notice in this case,in the case of Burger’s equation, and in general, that the coefficients found are slightly below theirtrue value. This is due to the effects of the regularization term and the fact that neither the data nor
the network fulfill perfectly the equation. Figure 1a shows snapshots of the solution.
The Burger’s equation ut = νuxx − uux is a second order and non-linear PDE, when ν is zero,
shockwaves can appear. We solve the equation for a value of ν = 0.05 and present the results in Fig.
1b. The equation our model finds is ut = 0.05uxx − 0.994uux .
High resolution oceanic simulations are very expensive, taking months to run sometimes, and low
resolution simulations do not capture all the important interactions. A proposed solution to thisproblem is to use a forcing term, which when added to the low resolution simulation, improves its
quality. Ross et al. [14] compare both symbolic and black box methods to learn the forcing term.
We use a similar methodology, but instead of genetic programming we use DiffComp. Owing to thecomplexity of the data, our pipeline of training a neural network and taking its derivatives did notwork because the neural network failed to converge. However, as [14] have implemented numericallythe differential operations, we use their implementation and therefore do not need to use a network tocompute derivatives.
We adapt the training method used by [14], and incorporate a human in the loop. We use DiffComp
to discover symbolic terms, but instead of relying solely on gradient descent like before, we learn
the terms of z M −1 by gradient descent, and use the least squares method like [14] to get a linear
combination of those terms, that is, the coefficents ξ M in z M = ξ M z M −1 . This is because theweights of the terms have very different orders of magnitude, which a gradient descent algorithm
cannot capture. A human then selects the most meaningful terms of the expression, based on leading
terms and ablation (removing terms and evaluating impact on performance). Those terms are then fed
as added input to DiffComp to reduce the complexity of the composition that has to be discovered.
Additionally, following the loss function of [14], instead of minimising the MSE between û and u,
we maximise the correlation between them, again, because the wildly different orders of magnitude
involved. We run this procedure three times and obtain progressively more complex expressions.
Using as inputs q, u, v we obtain the first expression w1 ∇2 (u · ∇)q + w2 ∇4 q. Then, the threeterms from the first input plus the two terms found as new inputs, we find the second expression
w1 (u · ∇)∇4 (u · ∇)q + (u · ∇)∇4 q + w3 ∇6 (u · ∇)q + w4 ∇8 q. The combined symbolic expression(a sum of those two equation) is the final model. For comparison, Ross et al. [14] find
(w1 ∇2 + w2 ∇4 + w3 ∇6 )(u · ∇)q + (w4 ∇4 + w5 ∇6 )q + (u · ∇)∇2 (w6 vx + w7 vy ).
3(3)

(a) 2D heat equation at t = 0, t = 0.5, and t = 1(b) 1D Burger’s equation, y-axis is the time axis

Figure 1: Ground truth (top), prediction (middle), and error percentage (bottom)(a) 1st expression discovered(b) 2nd expression discovered(c) Combined expression

Figure 2: Comparison between simulation at high resolution, low resolution and low resolution plus
the added forcing term

The third iteration yields an unwieldy expression that is over-fitted. It is shown in the appendix (Eq.
4), along with other evaluation metrics from [14].

4Conclusion

We introduced a differentiable symbolic regressor to be used with a differentiable surrogate, e.g. aneural network, to discover an unknown model directly from data. While DiffComp can also be used
as a pure symbolic regressor, we note that its performance will be lower than state-of-the-art genetic
and tree regressors, such as PySR [7] and gplearn [1]. We tested our method on two differential
equations as well as a real-world benchmark from oceanography, and compared it to other methods.
We found it competitive, with the disadvantage of being slower than PySINDy but the advantage
of able to learn functions compositions of a certain depth, allowing for the discovery of terms thatare not in the original library. Because it uses gradient descent to find the terms, DiffComp canseamlessly find parameters of the terms, like in Burger’s equation, but struggles if the data spans
too many orders of magnitude. Lastly, a few open questions remain; as the gradient dynamics aredifferent from those in neural networks, are optimizers designed for said networks optimal? How tohandle wildly different data scales, like in the forcing term example? How to make DiffComp betterat finding deeper compositions?
Impact statement: This work aims to be a minor step towards more automation in science. As such,
we believe that the societal impact of the method introduced will be minor, but positive. Additionally,at this stage of development we see no ethical concerns.
4

Acknowledgements
This work was partially supported by ServicePublic de Wallonie Recherche under grant n° 2010235 –ARIAC by DIGITALWALLONIA4.AI. Gilles Louppe is recipient of the ULiège - NRB Chair on Big
Data and is thankful for the support of the NRB.

References[1] gplearn: Genetic Programming in Python, with a scikit-learn inspired API. https://github.
com/trevorstephens/gplearn. Accessed: 2022-09-26.
[2] L. Biggio, T. Bendinelli, A. Neitz, A. Lucchi, and G. Parascandolo. Neural symbolic regressionthat scales. CoRR, abs/2106.06427, 2021.
[3] J. Bongard and H. Lipson. Automated reverse engineering of nonlinear dynamical systems.
Proceedings of the National Academy of Sciences, 104(24):9943–9948, 2007.
[4] G.-J. Both, S. Choudhury, P. Sens, and R. Kusters. Deepmod: Deep learning for model discovery
in noisy data. Journal of Computational Physics, 428:109985, 2021.
[5] G.-J. Both and R. Kusters. Fully differentiable model discovery, 2021.
[6] S. L. Brunton, J. L. Proctor, and J. N. Kutz. Discovering governing equations from data by
sparse identification of nonlinear dynamical systems. Proceedings of the National Academy of
Sciences, 113(15):3932–3937, 2016.
[7] M. Cranmer. Pysr: Fast & parallelized symbolic regression in python/julia, Sept. 2020.
[8] B. de Silva, K. Champion, M. Quade, J.-C. Loiseau, J. Kutz, and S. Brunton. Pysindy: A python
package for the sparse identification of nonlinear dynamical systems from data. Journal ofOpen Source Software, 5(49):2104, 2020.
[9] N. Geneva and N. Zabaras. Transformers for modeling physical systems. CoRR, abs/2010.03957,
2020.
[10] A. A. Kaptanoglu, B. M. de Silva, U. Fasel, K. Kaheman, A. J. Goldschmidt, J. Callaham, C. B.
Delahunt, Z. G. Nicolaou, K. Champion, J.-C. Loiseau, J. N. Kutz, and S. L. Brunton. Pysindy:
A comprehensive python package for robust sparse system identification. Journal of Open
Source Software, 7(69):3994, 2022.
[11] P. Lemos, N. Jeffrey, M. Cranmer, S. Ho, and P. Battaglia. Rediscovering orbital mechanics
with machine learning, 2022.
[12] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin,N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala. Pytorch: An imperative
style, high-performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer,
F. d'Alché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing
Systems 32, pages 8024–8035. Curran Associates, Inc., 2019.
[13] M. Raissi, P. Perdikaris, and G. E. Karniadakis. Physics informed deep learning (part ii):
Data-driven discovery of nonlinear partial differential equations, 2017.
[14] A. S. Ross, Z. Li, P. Perezhogin, C. Fernandez-Granda, and L. Zanna. Benchmarking of machinelearning ocean subgrid parameterizations in an idealized model. Earth and Space Science Open
Archive, page 43, 2022.
[15] A. Sanchez-Gonzalez, J. Godwin, T. Pfaff, R. Ying, J. Leskovec, and P. W. Battaglia. Learning
to simulate complex physics with graph networks, 2020.
[16] M. Schmidt and H. Lipson. Distilling free-form natural laws from experimental data. Science,
324(5923):81–85, 2009.
[17] S. Seo and Y. Liu. Differentiable physics-informed graph networks, 2019.

5

Checklist
The checklist follows the references. Please read the checklist guidelines carefully for information on
how to answer these questions. For each question, change the default [TODO] to [Yes] , [No] , or
[N/A] . You are strongly encouraged to include a justification to your answer, either by referencing
the appropriate section of your paper or providing a brief inline description. For example:• Did you include the license to the code and datasets? [Yes] See Section ??.
• Did you include the license to the code and datasets? [No] The code and the data areproprietary.
• Did you include the license to the code and datasets? [N/A]Please do not modify the questions and only use the provided macros for your answers. Note that the
Checklist section does not count towards the page limit. In your paper, please delete this instructions
block and only keep the Checklist section heading above along with the questions/answers below.
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope? [Yes](b) Did you describe the limitations of your work? [Yes](c) Did you discuss any potential negative societal impacts of your work? [Yes](d) Have you read the ethics review guidelines and ensured that your paper conforms tothem? [Yes]2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A](b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [No](b) Did you specify all the training details (e.g., data splits, hyperparameters, how they
were chosen)? [N/A](c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [No](d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)? [Yes]4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [N/A](b) Did you mention the license of the assets? [N/A](c) Did you include any new assets either in the supplemental material or as a URL? [N/A](d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating? [N/A](e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable? [N/A](b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable? [N/A](c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation? [N/A]6

A

Appendix

All the presented in this work were ran on a Thinkpad P1, with an 11th Gen Intel(R) Core(TM)i7-11800H and 16GB of RAM. No GPUs were used. We rerun the experiments from Table 1 withadded noise and show the results in Table 2. We also present further online 4 and offline metrics 3
from Ross et al. [14] to further test the symbolic expression found for the forcing term, and present
the resulting term from the last iteration of training in 4.
MethodDiscovered equation

runtimeMSE< 1s∼ 8min∼ 14min

0.02
9.76 × 10−2
3.43 × 10−5

< 1s
∼ 13min∼ 7min1.60 × 10−5
3.23 × 10−3
2.70 × 10−6

2D heat equation: ut = uxx + uyy
PySINDy
DeePyMod
DiffComp−2.587u3
0.136uxx − 0.879u
0.931(−0.168x + 0.316y)(−1.893ux + 2.441uy − 0.005) + 0.53uxx +1.01uyy − 1.471uxy + 0.06

1D Burger’s equation: ut = 0.05uxx − uux
PySINDy
DeePyMod
DiffComp−0.913uux−0.292ux + 0.205uxx + 0.741uux
0.049uxx − 0.998uux

Table 2: The noise was created by sampling from a gaussian with mean 0 and std equal to 0.01 of thenorm of the initial condition u(0, x). PySINDy’s WeakPDEFIND was used here(a) 1st expression(b) 2nd expression(c) Combined expression

Figure 3: Offline metrics measured between the target forcing term and the discovered symbolic
expression. Note that the correlation was the training target
w1 (u · ∇) ∇2 ∇2 ((u · ∇)(q)) + w12 (u · ∇) w8 (u · ∇) ∇2 ∇2 (q)+ w3 ∇2 w0 (u · ∇) w8 (u · ∇) ∇2 ∇2 (q)+ w11 ∇2 w13 (u · ∇) ∇2 ∇2 (q)+ w4 (u · ∇) w10 ∇2 w13 (u · ∇) ∇2 ∇2 (q)+ w2 (u · ∇) w8 (u · ∇) ∇2 ∇2 (q)+ w5 ∇2 ∇2 ∇2 ∇2 (q)+ w6 (u · ∇) ∇2 ∇2 (q)+ w7 ∇2 w13 (u · ∇) ∇2 ∇2 (q)+ w9 ∇2 ∇2 ∇2 ((u · ∇)(q))7(4)

(a) 1st expression(b) 2nd expression(c) Combined expression

Figure 4: Comparison between the energy budgets of the high resolution simulation, the low resolution
simulation, and the low resolution simulation including the discovered forcing terms. Just like [14]find, we see that despite the high correlation there is no improvement in new low resolution simulation.
We believe that a human in the loop procedure where the human is a domain expert and applied
domain knowledge when selecting the relevant terms will improve the results, as was done in [14]8