Approximate Latent Force Model Inference

Jacob D. Moss∗
Computer Lab
University of Cambridge
Cambridge, UK
jm2311@cam.ac.uk

Felix L. OpolkaComputer Lab
University of Cambridge
Cambridge, UK
flo23@cam.ac.uk

Bianca Dumitrascu
Computer Lab
University of Cambridge
Cambridge, UK
bmd39@cam.ac.uk

Pietro LiòComputer Lab
University of Cambridge
Cambridge, UK
pl219@cam.ac.uk

Abstract
Physically-inspired latent force models offer an interpretable alternative to purely
data driven tools for inference in dynamical systems. They carry the structureof differential equations and the flexibility of Gaussian processes, yielding interpretable parameters and dynamics-imposed latent functions. However, the existing
inference techniques rely on the exact computation of posterior kernels which are
seldom available in analytical form. Applications relevant to practitioners, suchas diffusion equations, are hence intractable. We overcome these computational
problems by proposing a variational solution to a general class of non-linear andparabolic partial latent force models. We demonstrate the efficacy and flexibility of
our framework by achieving competitive performance on several tasks.

1

Introduction & BackgroundA vast variety of phenomena in the sciences, from systems of interacting particles to the aerodynamics
of aircraft, are represented by differential equations. Machine learning approaches such as neuralordinary differential equations (ODEs) [Chen et al., 2018] and non-parametric ODEs [Heinonen et al.,2018] have become popular black-box models. However, the lack of an explicit equation means thatthese methods lose intepretability and the ability to infer latent terms. On the other hand, Latent ForceModels (LFMs) [Lawrence et al., 2007, Alvarez et al., 2009] are a hybrid paradigm combining thestructure and extrapolation of explicit models with the flexibility of Gaussian processes (GPs). Thisfusion is a powerful approach to problems in biophysics [López-Lopera et al., 2019, Moss and Lió,
2020] and stochastic control [Särkkä et al., 2018, Cheng et al., 2020]. However, they often require a
intractable integration of the kernel function, severely limiting their usefulness.
Addressing these drawbacks, we introduce Alfi, a flexible variational framework for parameter
estimation in ordinary and partial differential equation (PDE)-based LFMs. The main contributions
are: (i) a gradient-matching step followed by fine-tuning with the full solution, enabling quadratictime inference in LFMs; and (ii) inference of latent forces from partially known dynamics in temporaland spatio-temporal settings without onerous kernel derivations.
Latent Force Models (LFMs) are differential equation models which leverage explicit dynamics toinfer latent forcing terms [Lawrence et al., 2007, Alvarez et al., 2009]. The differential equation, h,∗Correspondence to Jacob Moss jm2311@cam.ac.uk.

Fourth Workshop on Machine Learning and the Physical Sciences (NeurIPS 2021).

Figure 1: Computational graphical model for Alfi. Step (1) involves pre-estimating the parameters andlatent forces of the differential equation in Eq. 1 by gradient-matching with the spline-interpolated
derivative. Step (2) involves the more expensive forward solution of the differential equation andbackpropagation through the ODE or PDE solver for fine-tuning the parameters.
parameterised by Θ, enforces a structural relationship between P outputs, y(x) ∈ RP , N inputs,x ∈ RN , and L unobserved latent forces, f (x) ∈ RL . GP priors are assigned to the latent forces, fi :fi ∼ GP(0, κi (x, x0 )), and these forces can be mixed by some response function G(f ).
differential

differential equation}|z }| { z{D y(x) = h y, x; Θ, G(f1 (x), ..., fL (x)) ,(1)

where D is a differential operator, for example an nth order (partial) derivative. The covariance
between outputs, κy,y0 (t, t0 ), is analytically tractable if G is a linear operator. In these cases,
maximum marginal likelihood yields the differential equation parameters and inference can be carriedout with standard posterior GP identities (see Rasmussen and Williams [2005]).

2

Approximate Latent Force Models

We propose Alfi, an approximate latent force inference framework significantly expanding the classof latent force models. It supports any differentiable kernel, and is agnostic to the functional formof the response to latent forces, G. An example of a non-linear intractable response is a saturating
function such as the Hill equation for modelling reaction kinetics in biophysical problems:Qwji
i fi (t)G(f (t); wj ) = Q(2)w
ji + e−wj0
i fi (t)where wji are interaction weights, incorporating cooperation/competition between the forces fi .
2.1

Variational Objective

Let Ŷ ∈ RN ×P represent the forward solution of a differential equation at N points given parameters
Φ, a differentiable transformation (G) of latent forces, and input X ∈ RN ×D . The input dimensions
are D = 1 (e.g. time) for ODEs, or D > 1 (e.g. time and space) for PDEs. Our formulation Lextends to arbitrary output dimensions, P . Let F = fi> i be our set of independent latent forceswhere fi ∼ GP(0, κΘ (x, x0 )) with kernel hyperparameters Θ. We assume a Gaussian likelihood,p(Y|F, Φ) = N (Y|Ŷ, σ 2 ), and seek to optimise parameters by marginalising over the latent forces:GP
Z z likelihood
L z}| { Y}| {p(Y|Φ, Θ) = p(Y|F, Φ)p(fi |X, Θ) df .

(3)iThis marginal is intractable when G is non-linear, so we derive an expectation lower-bound (ELBO)
for the marginal likelihood using a variational approximation. We add M inducing points, ui ∈ L
RM ×D , for each latent force: U = u>i i . Due to the conditional independence across the latent
forces given any mixing parameters (e.g. Eq. 2), a factorised variational distribution is used:2

QL
q(u) = i N (mi , Ci ) with variational parameters mi ∈ RM and Ci ∈ RM ×M . This yields the
ELBO we seek to to optimise (the full derivation is in Appendix C):
log p(Y|Φ, Θ) = Eq(F) [log p(y|F, Φ)] −L
X

KL(q(ui )||p(ui )),

(4)

i=1

R
where q(F) = p(F|U)q(U)dU and we approximate the expected log-likelihood with Monte Carlosampling. The variational posterior is p(fi∗ |ui ) = N (fi∗ |m∗i , S∗i ) where m∗i = K∗M K−1
M M mi and−1S∗i = K∗∗ + K∗M K−1
(S
−
K
)K
K
[Hensmanetal.,2015].
i
M
M
M
∗
MM
MM
2.2

Optimisation of Variational Objective

At each epoch, a batch of the latent forces is sampled from the variational posterior using thereparameterisation trick, enabling pathwise gradient estimation. Parameters are pre-estimated in thefirst few epochs, followed by fine-tuning with the forward solution. With e differential equation
evaluations, the computational complexity is O(N M 2 + e). Pseudocode along with sample codesnippets are in Appendix B and E. We now detail these steps.
Step (1): Parameter pre-estimation by gradient matching. The forward solution of an ODE or
PDE is the most computationally expensive component of our algorithm. As such, we develop a
pre-estimation step to initialise the variational and differential equation parameters. We implementa gradient-matching algorithm whereby the differential equation is evaluated, with no calls to thesolver, at all training points and compared with the interpolated gradient of the data at those points.
We chose natural cubic splines for this purpose since they are linear in the number of points; more
expensive alternative interpolations are discussed in Appendix B.
Step (2): Forward solution fine-tuning. This step involves executing an appropriate differential
equation solver to generate samples from the LFM. For ODEs, we employ an ODE solver, throughwhich the gradients are backpropagated. PDE-based LFMs are more involved since they are muchharder to solve than in the ODE case. The GP component uses an anisotropic kernel where eachdimension gets a separate lengthscale, since the latent force is likely to vary across time and spaceat different scales. We solve the PDE using the finite element method (FEM), which first requires
converting it into its weak form, which is then solved on a discretised spatial mesh. The weak formis written in the domain-specific Unified Form Language (UFL) [Alnæs et al., 2014]. Alfi uses the
FEniCS package [Logg et al., 2012, Alnæs et al., 2015] for solving the variational problems. The
gradients can be determined with reverse-mode algorithmic differentiation using the adjoint equations,computed using the dolfin-adjoint package [Mitusch et al., 2019].

3

Experiments

Starting with a simulation of predator-prey dynamics, we demonstrate the ability to infer the prey
dynamics without knowledge of prey abundance. We also investigate PDE-based LFMs with thespatio-temporal production of protein in fruit fly embryos with a reaction-diffusion equation.
ODE System: The Lotka-Volterra equations are a system of non-linear first-order differential
equations used to model predator-prey dynamics. The population of predators, v, is related to the
population of prey, u, with the following pair of equations:du= αu − βuv,
dtdv
= δuv − γv,
dt(5)

where α, δ are growth rates; β, γ are decay rates; and for clarity u = u(t) and v = v(t).
Assuming the population of prey is unknown and positive, we assign a Gaussian process prior to thefunction u(t) ∼ GP(0, κ(t, t0 )) and a softplus response, G(•) = log (1 + e• ), meaning the solution
is intractable. While the the ubiquitous radial basis function (RBF), a universal smooth function
approximator, would provide a good interpolation, a more sensible kernel would be periodic:
2 sin2 (π|t − t0 |)/p
κp (t, t0 ) = exp
,
(6)
`2
3

where p is the period parameter, determining the interval of time after which the function repeats. As
seen in Figure 2, the periodic kernel enables realistic extrapolation beyond the training time range.
As shown in Appendix G.2, an RBF kernel yields inferior results.
PDE System: Reaction-diffusion equations are second-order parabolic PDEs ubiquitous in the naturalsciences. We take a model of gap gene protein production in Drosophila [López-Lopera et al., 2019]:
reactionzdiffusion

}|

{

z

}|
{
∂u(x, t)∂ 2 u(x, t),(7)= Su(x, t) − λy(x, t) + D∂t∂x2
where D is the diffusion constant(s), S is the production rate of the driving mRNA (the latent force)and λ is the decay rate. The output, y, is the protein concentration. We assign an anisotropic
RBF kernel prior over the latent force u. Dirichlet boundary conditions are assigned such thaty(x = 0, t) = 0 and y(x = l, t) = 0 where 0 and l are the spatial boundaries.
We start with the weak form for this PDE by multiplying by a test function v ∈ V and integrating thesecond derivatives by parts over the spatial domain, Ω.
ZZyv + ∆tD∇x y · ∇x v d~x = (y n + ∆tSu − ∆tλy)v d~x,
(8)
Ω

Ω

a(y, v) = Ln+1 (v)

(9)

ELBO

where a(y, v) is a bilinear form, v is known as a test function belonging to some Banach space V , and function parameters have been
omitted (i.e. u = u(x, t)). For completeness, the full derivation andtranscription into UFL can be found in Appendix D and E.2.

0.4

0.2
Pre-estimation
Without
0.0

The results are shown in Table 1, and the full qualitative outputs
−0.2
0
2
4
6
8
10
are in Appendix G.3. We note that when using the parameter preTime (min)estimation procedure, the convergence is significantly sped up andFigure 3:Pre-estimation:
the ELBO starts at a lower initial value, as shown in Figure 3.
gradient-matching improves
convergence.

4

Discussion

Alfi enables approximate inference in linear, non-linear, ordinary, and partial latent force models. The
output of Alfi is a conditional Gaussian process, meaning the resolution, uniformity, and range ofevaluation points is arbitrary. In addition, the error of the result is tunable using the tolerances in thesolver. When compared with analytical LFMs, our method overcomes the complex derivations ofintegrated kernels, which are either intractable or involve complicated convolutions [Lawrence et al.,2007]. Since FEM is limited by the curse of dimensionality, so will Alfi’s performance on PDEs.
Regardless, the severity is minimal since many physical problems are in 3D or 4D.
Ethical considerations: There could be sensitive applications (e.g. medical time-series) where blind
faith in inferred forces can lead to detrimental effects on the patient. This is a consideration inherent
in all models, however in this case the latent nature could be more pernicious.
3.0

0.5

Predator population

Prey population

Predator
Prey
TargetPrediction
Target
1

0
0

20

Time

0

2

Prey population

Figure 2: Alfi with a periodic kernel on the Lotka-Volterra task. Orange crosses indicate the training
datapoints. Blue and green lines are used for predator and prey respectively, orange line for unseen
prey ground truth. Left: the output dynamics with the training data as crosses along with the latent
dynamics. Right: the phase space (u, v) alongside the unseen ground truth.
4

Table 1: Results of Alfi versus the closest model in the literature, GP-mRNA [López-Lopera et al.,2019]. Alfi obtains a comparable performance in most cases, although the higher latent force inference
error is an avenue for further research. For metrics explanation see Appendix G.1.
Q2 [%] ↑mRNA
Gap Proteinµ±σµ±σ

68 − CA±σ [%]mRNA
Gap Protein
µ±σ
µ±σ

Model

Trunk
Gap
Gene

GP-mRNAkr
kni
gt

90.5 ± 2.0
81.1 ± 2.5
91.2 ± 1.9

90.8 ± 0.6
88.7 ± 0.8
92.3 ± 0.6−3.7 ± 6.7
18.3 ± 6.3−0.7 ± 3.3

21.8 ± 4.0
35.3 ± 5.3
25.8 ± 1.8

Alfi*

kr
kni
gt

82.4 ± 1.2
77.2 ± 2.6
72.4 ± 3.3

95.7 ± 0.2
86.9 ± 2.9
86.5 ± 1.2

3.3 ± 3.5
−1.5 ± 1.6−3.9 ± 4.4

19.4 ± 2.0
12.4 ± 9.1
5.8 ± 14.3

References
Martin S Alnæs, Anders Logg, Kristian B Ølgaard, Marie E Rognes, and Garth N Wells. Unified
form language: A domain-specific language for weak formulations of partial differential equations.
ACM Transactions on Mathematical Software (TOMS), 40(2):1–37, 2014.
Martin S. Alnæs, Jan Blechta, Johan Hake, August Johansson, Benjamin Kehlet, Anders Logg, Chris
Richardson, Johannes Ring, Marie E. Rognes, and Garth N. Wells. The fenics project version 1.5.
Archive of Numerical Software, 3(100), 2015. doi: 10.11588/ans.2015.100.20553.
Mauricio Alvarez, David Luengo, and Neil D Lawrence. Latent force models. In Artificial Intelligenceand Statistics, pages 9–16. PMLR, 2009.
Kolja Becker, Eva Balsa-Canto, Damjan Cicin-Sain, Astrid Hoermann, Hilde Janssens, Julio R
Banga, and Johannes Jaeger. Reverse-engineering post-transcriptional regulation of gap genes indrosophila melanogaster. PLoS Comput Biol, 9(10):e1003281, 2013.
Ben Calderhead, Mark Girolami, and Neil D Lawrence. Accelerating bayesian inference overnonlinear differential equations with gaussian processes. In Advances in neural information
processing systems, pages 217–224. Citeseer, 2009.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential
equations. arXiv preprint arXiv:1806.07366, 2018.
Li-Fang Cheng, Bianca Dumitrascu, Michael Zhang, Corey Chivers, Michael Draugelis, Kai Li, andBarbara Engelhardt. Patient-specific effects of medication using latent force models with gaussian
processes. In International Conference on Artificial Intelligence and Statistics, pages 4045–4055.
PMLR, 2020.
Frank Dondelinger, Dirk Husmeier, Simon Rogers, and Maurizio Filippone. Ode parameter inference
using adaptive gradient matching with gaussian processes. In Artificial intelligence and statistics,
pages 216–228. PMLR, 2013.
Jacob R Gardner, Geoff Pleiss, David Bindel, Kilian Q Weinberger, and Andrew Gordon Wilson.
Gpytorch: Blackbox matrix-matrix gaussian process inference with gpu acceleration. arXiv
preprint arXiv:1809.11165, 2018.
Jouni Hartikainen, Mari Seppanen, and Simo Sarkka. State-space inference for non-linear latent force
models with application to satellite orbit prediction. arXiv preprint arXiv:1206.4670, 2012.
Pashupati Hegde, Markus Heinonen, Harri Lähdesmäki, and Samuel Kaski. Deep learning withdifferential gaussian process flows. arXiv preprint arXiv:1810.04066, 2018.
Markus Heinonen, Cagatay Yildiz, Henrik Mannerström, Jukka Intosalmi, and Harri Lähdesmäki.
Learning unknown ode models with gaussian processes. In International Conference on Machine
Learning, pages 1959–1968. PMLR, 2018.
5

James Hensman, Alexander Matthews, and Zoubin Ghahramani. Scalable variational gaussian
process classification. 2015.
Patrick Kidger, James Morrill, James Foster, and Terry Lyons. Neural controlled differential equations
for irregular time series. arXiv preprint arXiv:2005.08926, 2020.
Diederik P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling.
Improving variational inference with inverse autoregressive flow. arXiv preprint arXiv:1606.04934,
2016.
Neil D Lawrence, Guido Sanguinetti, and Magnus Rattray. Modelling transcriptional regulation using
gaussian processes. Advances in Neural Information Processing Systems, 19:785, 2007.
Anders Logg, Kent-Andre Mardal, Garth N. Wells, et al. Automated Solution of Differential
Equations by the Finite Element Method. Springer, 2012. ISBN 978-3-642-23098-1. doi:
10.1007/978-3-642-23099-8.
Andrés Felipe López-Lopera, Nicolas Durrande, and Mauricio Alexander Alvarez. Physicallyinspired gaussian process models for post-transcriptional regulation in drosophila. IEEE/ACM
transactions on computational biology and bioinformatics, 2019.
Sebastian K Mitusch, Simon W Funke, and Jørgen S Dokken. dolfin-adjoint 2018.1: automatedadjoints for fenics and firedrake. Journal of Open Source Software, 4(38):1292, 2019.
Jacob Moss and Pietro Lió. Gene regulatory network inference with latent force models, 2020.
Alexander Norcliffe, Cristian Bodnar, Ben Day, Jacob Moss, and Pietro Lio. Neural ode processes.
arXiv preprint arXiv:2103.12413, 2021.
Carl Edward Rasmussen and Christopher K. I. Williams. Gaussian Processes for Machine Learning(Adaptive Computation and Machine Learning). The MIT Press, 2005. ISBN 026218253X.
Simo Särkkä, Mauricio A Álvarez, and Neil D Lawrence. Gaussian process latent force models forlearning and stochastic control of physical systems. IEEE Transactions on Automatic Control, 64(7):2953–2960, 2018.
Michalis K Titsias, Neil D Lawrence, and Magnus Rattray. Efficient sampling for gaussian processinference using control variables. In NIPS, pages 1681–1688. Citeseer, 2008.
Wil Ward, Tom Ryder, Dennis Prangle, and Mauricio Alvarez. Black-box inference for non-linear
latent force models. In International Conference on Artificial Intelligence and Statistics, pages
3088–3098. PMLR, 2020.
Philippe Wenk, Alkis Gotovos, Stefan Bauer, Nico S Gorbach, Andreas Krause, and Joachim M
Buhmann. Fast gaussian process based gradient matching for parameter identification in systems
of nonlinear odes. In The 22nd International Conference on Artificial Intelligence and Statistics,
pages 1351–1360. PMLR, 2019.

Checklist
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] See Section 4.
(c) Did you discuss any potential negative societal impacts of your work? [Yes] See
Section 4(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? [Yes] See Section 4
6

(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)? [Yes] In the
supplementary material(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they
were chosen)? [Yes] In Appendix H(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [Yes](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] In Appendix F
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? [Yes](b) Did you mention the license of the assets? [Yes] In the Supplementary Materials, where
the datasets are included.
(c) Did you include any new assets either in the supplemental material or as a URL? [Yes](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? [Yes] The data used does not contain personallyidentifiable or offensive information.
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]7

A

Comparison

Machine learning methods for differential equations typically assume the equation is unknown, andfollow either a GP or neural network approach. For example, differential GPs [Hegde et al., 2018],
generalise the discrete layers of a deep GP to a sequence of infinitesimal steps in which the input iswarped through a vector field defined by an stochastic differential equation (SDE). However, the SDE
acts only as a mechanism for constructing a model rather than as an actual physical model. A similar
approach also treats the ODE as unknown using a Gaussian process and uses the full solution alongwith sensitivity equations to extract gradients [Heinonen et al., 2018]. This model, however, only uses
the conditional mean of the Gaussian process, whereas Alfi also captures the covariance. The neuralnetwork methods, popularised by Chen et al. [2018], represent the ODE function as a neural network,where the weights are optimised by backpropagation through an ODE solver or via adjoint equations.
Some recent extensions add stochasticity in the form of initial condition distributions [Norcliffe et al.,2021]. Unknown dynamics models lack the ability to infer latent forces, and therefore cannot becompared directly with LFMs.
Latent force models (LFMs) are between the explicit and unknown differential equation paradigms;
the equation depends on a set of latent forces which are inferred using the explicit dynamics. There
are several other approaches for overcoming the intractable posterior in non-linear LFMs, namely the
Laplace approximation [Lawrence et al., 2007] and MCMC [Titsias et al., 2008]. Hartikainen et al.
[2012] show the connection between non-linear latent force models and non-linear SDEs, thereby
enabling an approximation using Gaussian filtering and smoothing. Neural network approximations
have also been attempted, where an inverse autoregressive flow [Kingma et al., 2016] is used to model
the variational posterior [Ward et al., 2020].

B

Training Scheme

Natural cubic splines: Our pre-estimation procedure is similar to neural controlled differential
equations Kidger et al. [2020] only in the choice of spline. Our interpolant is not fed through asolver–it is used only for gradient estimation. There are many possible interpolation approaches touse instead of cubic splines. One popular method is using a Gaussian process [Calderhead et al.,2009, Dondelinger et al., 2013, Wenk et al., 2019] with a differentiable kernel (and mean) function.
Gaussian processes with such kernels are closed under differentiation, and the covariance terms
2
K(t,t0 )can be found as cov(ẏj (t), ẏj (t0 )) = ∂ ∂t∂t
. Cubic splines are linear in the number of points, in0contrast to GPs which are cubic (or quadratic for non-naive implementation [Gardner et al., 2018]).
Since cubic splines are used in pre-estimation, only the 1st and 2nd derivatives are continuous. If themodel has higher order derivatives, a higher order spline can be used or the higher-order derivatives
ignored—only an approximation is desired in this step.
Natural gradient optimisation: optimising the variational parameters using gradient descent inparameter space assumes Euclidean geometry. Variational parameters exist in distribution space, so weshould use a more appropriate geometry within which we calculate distance, such as KL-divergence.
Natural gradient optimisation takes steps in this space.
Algorithm 1: Alfi Scheme: Steps (1) and (2)Input: x, y
initialization;for i ← 1 to T do
latent_forces = variational_posterior(x);if preestimation thentarget ← spline_interpolate_gradient(x, y) ;// target is gradientpred ← de_func(initial_conditions, latent_forces);
elsetarget ← y ;// target is solutions
pred ← solve(de_func, initial_conditions, latent_forces);end
loss ← elbo(pred, target)
end

8

C

ELBO Derivation

For completeness, we include here the derivation of the ELBO in Equation 4.
ZZ
log p(Y|Θ, Φ) = log
p(Y|F)p(F|U)p(U)dUdF,
ZZ
p(Y|F)p(F|U)p(U)p(F|U)q(U)
= log
dUdF
p(F|U)q(U)We now apply Jensen’s inequality:
ZZ
p(Y|F)p(U)≥p(F|U)q(U) log
dUdF,
q(U)
Z

Z
p(U)
= q(U)
p(F|U) log p(Y|F)dF + log
dU,
q(U)


= Eq(U) Ep(F|U) [log p(Y|F)] − KL(q(U)||p(U)),= Eq(F) [log p(Y|F)] −L
X

KL(q(ui )||p(ui )),

(10)
(11)

(12)
(13)
(14)
(15)

i

D

Heat Equation Weak Form Derivation

We go through the complete instructive derivation for the heat equation’s variational form below.
∂t y(x, t) = Su(x, t) − λy(x, t) + D∇2x y(x, t)(16)

The first step in FEM is to find the variational form. We henceforth omit function parameters (e.g.
u = u(x, t)). First, we take the finite backwards difference as follows ∀n = 0, 1, 2, ...:
y n+1 − y n= Su − λy + D∇2x y∆t
We rearrange to find a bilinear form on the left-hand-side and a linear form on the right:∂t y n+1 ≈y n+1 + ∆tλy n+1 − ∆tD∇2 y n+1 = y n + ∆tSu

(17)

(18)

For brevity we now drop the n + 1 indices (y = y n+1 ). The next step in FEM is to multiply both sides
by a test function v ∈ V and integrate over the spatial domain Ω. The Dirichlet boundary conditions
are defined such that y(x = 0, t) = 0 and y(x = l, t) = 0 where 0 and l are the spatial boundaries.
ZZ(yv + ∆tλyv − ∆tD∇2x yv)d~x = (y n + ∆tSu) v d~x
(19)
Ω

Ω

We can integrate the second derivatives by parts:
Z
Z∆tD∇2x yv d~x = D [∇x y · v]l0 −∆tD∇x y · ∇x v d~x| {z }
Ω
Ω

(20)

0

where [∇x y · v]l0 vanished due to boundary conditions. This yields the variational (weak) form:ZZ
yv + ∆tλyv + ∆tD∇x y · ∇x v dx =Ω(y n + ∆tSu)v d~x

(21)

Ω

a(y, v) = Ln+1 (v)(22)

This weak form can then be written in UFL almost verbatim.

E

Code

E.1

Differences between an ODE- or PDE-based LFMBoth OrdinaryLFM() (for ODE-based LFMs) and PartialLFM() (for PDE-based LFMs) inheritfrom the VariationalLFM() class which provides methods for sampling from the variational GP.
9Depending on which class is used, the practitioner must implement the ode_func or pde_func
function, which are used by the solvers to calculate the forward solution, and the pre-estimator(PreEstimator) to calculate data gradients.
An example LFM for a transcriptional regulation model from is written below, demonstrating how
simple Alfi makes constructing LFMs:
class TranscriptionLFM(OrdinaryLFM):
def __init__(self, num_outputs, gp_model, config: VariationalConfiguration, **kwargs):
super().__init__(num_outputs, gp_model, config, **kwargs)
self.decay_rate = Parameter(torch.rand((self.num_outputs, 1)))
self.basal_rate = Parameter(torch.rand((self.num_outputs, 1)))
self.sensitivity = Parameter(torch.rand((self.num_outputs, 1)))
# optionally add any constraints...
def odefunc(self, t, h):f_latents = self.f
f_latents = self.G(f_latents) # optionally add non-linearity
dh = self.basal_rate + self.sensitivity * f_latents - self.decay_rate * hreturn dh
E.2

Unified Form Language for PDEs

The PartialLFM requires an additional module called fenics_module. This is defined in Unified
Form Language (UFL) [Alnæs et al., 2014]. An example UFL snippet for the reaction-diffusion
equation from Section 3 is shown below:
mesh = UnitIntervalMesh(20)V = FunctionSpace(mesh, 'P', 1)u = TrialFunction(V)v = TestFunction(V)a = (1 + dt * decay) * u * v * dx
a = a + dt * diffusion * inner(grad(u), grad(v)) * dx
L = (u_prev + dt * sensitivity * latent_force) * v * dx
bc = DirichletBC(V, 0, 'on_boundary')y = Function(V)solve(a == L, y, bc)

FTime profiling

As shown in Figure 4, the convergence of the variational LFM has a faster or comparable training
time to the analytical solution, as shown in Figure 4b for the 1D case.

G

Additional Results

G.1

Metrics

We use the Q2 = 1 − SMSE criterion, where SMSE is the mean-squared-error (MSE) normalised by
the target variance, overcoming the sensitivity of MSE to the scale of the target values [Rasmussen
and Williams, 2005]. In addition, to give more insights into the more challenging spatio-temporalexperiments, we also quote coverage interval accuracy deviation (68% − CA±σ ), which measures
the proportion of targets within 1 standard deviation of the predictive distribution.
10

LaFoMo

0.4
0.3

Loss

ELBO

0.2
Pre-estimation
Without

0.1
0.0

6

1.0

4

0.8

2

0.6

0

0.4
0.0

−0.1
−0.2
0

2

4

6

8

Lawrence et al. (2007)

0.1

0.2

Time (min)

0.3

0.0

0.2

0.4

Time (min)0.6(b) Convergence times of Alfi compared with theanalytical model in Lawrence et al. [2007]. We
found that increasing the learning rate for the lattermodel lead to instabilities in the cholesky decompositions. Both were run on an AMD Ryzen 5600X
CPU.

10

Time (min)(a) Latent Force Model. Training times for thereaction-diffusion PDE LFM with and without preestimation. Training was run on a Titan Xp GPU.

Figure 4: Time profiling.
G.2

Lotka-Volterra

In Figure 5, the results for the Lotka-Volterra task are shown when using an RBF kernel instead of
the periodic kernel.
3.0

0.5

0

Predator population

Prey population

Predator
Prey
Target20

Time

Prediction
Target1002Prey population

4

Figure 5: Alfi with an RBF kernel on the Lotka-Volterra task. This demonstrates inferior extrapolation
performance compared with a periodic kernel, since a periodic kernel has additional inductive biases.
G.3

Drosophila

In Figure 6, the results for the reaction-diffusion experiment with the Drosophila embryogenisis
dataset are shown.

H

Hyperparameters

Hyperparameter optimisation is typically carried out with maximum marginal likelihood, or a lowerbound thereof. The hyperparameters were initial set to reasonable values: lengthscales of kernels
were set to the smallest distance between datapoints, and rates to a uniform random number between
0 and 1. For the comparison with GP-mRNA [López-Lopera et al., 2019], we kept the experimental
conditions the same: the rates were fixed to the values reported by Becker et al. [2013], and thelengthscales were learnt by maximising the ELBO.

11

Latent (Prediction)

95

x

3

54

54

97

Output (Prediction)95

3

Output (Target)

95

3

2

2

2

1

1

1

1

0

0

0

33
95

54

33
95

97

54

0
33
95

97

54

97

3

3

3

3

2

2

2

2

1

1

1

1

0

0

0

33
95
3

54

97

33
95
3

54

97

0
33
95

3

54

97

3

2

2

2

2

1

1

1

1

x

33
95

3

2

97

x

33
95

Latent (Target)

95

0
33

0
33

54

97

t

0
33

54

97

54

t

0
33

97

t

54

97

t

Figure 6: Alfi on the Drosophila reaction-diffusion task. First and second columns: the latent mRNA
concentration alongside ground truth. Third and fourth columns: protein concentration alongside
ground truth.

12