Deep Kernel Learning of Nonlinear Latent Force Models
Authors: Jacob Moss
Jeremy England
Pietro Lio
Venue: Accepted by TMLR ()

Abstract: Scientific processes are often modelled by sets of differential equations. As datasets grow, individually fitting these models and quantifying their uncertainties becomes a computationally challenging task. Latent force models offer a mathematically-grounded balance between data-driven and mechanistic inference in such dynamical systems, whilst accounting for stochasticity in observations and parameters. However, the required derivation and computation of the posterior kernel terms over a low-dimensional latent force is rarely tractable, requiring approximations for complex scenarios such as nonlinear dynamics. In this paper, we overcome this issue by posing the problem as learning the solution operator itself to a class of latent force models, thereby improving the scalability of these models. This is achieved by employing a deep kernel along with a meta-learned embedding of the output functions. Finally, we demonstrate the ability to extrapolate a solution operator trained on simulations to real experimental datasets, as well as scaling to large datasets.

Published in Transactions on Machine Learning Research (10/2024)
Deep Kernel Learning of Nonlinear Latent Force Models
Jacob Moss
jmM311 @eamm.m... uk
Computer Lab University of Cambridge
Jeremy England
GSK
Petach Tikea, Israel
Pietro Lió
Computer Lob
University of Cambridge
Reviewed OpenReview: https: openreview. httpe:///penneeveel
net/ forum? id= CWJIII4466
Abstract
Scientific processes are often modelled by sets of differential equations. As datasets grow, individually fitting these models and quantifying their uncertainties becomes a computa- tionally challenging task. Latent. force models offer a mathhematically-gooundedd balance be- tween data-driven and mechanistic inference in such dynamical systems, whilst accounting for stochasticity in observations and parameters. However, the required derivation and com- putation of the posterior kernel terms over low-dimensional latent force is rarely tractable, requiring approximations for complex scenarios such as nonlinear dynamics. In this paper, we overcome this issue by posing the problem as learning the solution operator itself to class of latent force models, thereby improving the performance and scalability of these mod- els. This achieved by employing a deep kernel along with a meta-learned embedding of the output functions. Finally, we demonstrate the ability extrapolate solution operator trained on simulations real experimental datasets, as well scaling to large datasets.
1 Introduction
Differential equations are mathematical models that describe the change function with respect one or more variables, such as time. They play central role in the natural and social sciences, providing a way to model and understand complex systems and phenomena, such the growth and decline of pop- ulations [Burghes 1975), morphogenesis (Turing. 1990). the dynamics of biochemical reactions (Thomas et al. 1976 Schoeberl et al., 2002 Barenco et al. 2006], and so on. They provide a rigorous, well-studied, and mathematically grounded method of making historic and future predictions in complex systems. In a machine learning context, the modelling power of differential equations make them excellent inductive biases. A popular method of incorporating these equations within a Bayesian machine learning setting is the latent force model (LFM), introduced by Lawrence et al. (2006) model network of genes regulated by common protein. LFMs assume that the underlying dynamics of system be modelled in terms of low-dimensional latent force, typically with Gaussian process prior, within system of differential equa- tions. The involvement of non-parametric Gaussian process within an interpretable parametric system results in a powerful framework for drawing mechanics--cosstrainel inferences in noisy, high-dimensionall and nonlinear dynamics. has since been extended to embryogenesis (López-Lopera et al. 2019), where the latent force represents mRNA concentration; patient interventions (Cheng et al., 2020), with latent forces corresponding to different treatments; and movement segmentation (Alvarez et al. 2010b However, there are significant computational challenges hindering the usability of these models
large datasets.

---

Published in Transactions on Machine Learning Research (10/2024)
Inferring the latent force requires computing the posterior distribution, the covariance functions of which are determined by the differential equations describing the model dynamics. These are analytically tractable only for small set of scenarios. The remaining cases, typically nonlinear dynamical systems, requires some approximation for the posterior. There are various approximations, such as using an ODE or PDE solver Moss et al. 2021), deep GPs (McDonald & Álvarez, (201), and filtering (Hartikainen & Sarkka, 2012. Ward et al., 2020). Crucially, however, these existing works operate on small batascts-ofter a single parameterisation not considering the scaling of the approach to multiple independent tasks of the same LFM. How do we extend the original model considering one network of genes to many thousands? Currently, important use-cases such as fitting multiple ODFs (Santra 2018) or performing Bayesian model selection in dynamical models (Babtie et al. 2014) with LFMs are too costly. Increasingly large datasets are only underlining the importance of solving the challenge of scaling up LFMs to work in the multi-task setting.
In this paper, we propose a novel method addressing these scalability issues, resulting 250x to 15, 000x faster latent force inferences when compared with existing approximations. We propose a deep learning approach to solve a general class of LFMs, avoiding differential equation solving sleps and variational ap- proximations by instead learning the dynamics in deep kernel (Wilson et al. 2016). Our framework makes use of neural network represcntations of sets, such as the Transformer (Vaswani et al. 2017) and the Fourier neural operator (Li et al. 2020), in order to produce a function embedding for each task. Given a task consisting of only the input mesh, for example time, and the observed functions' embedding, our model infers the associated latent force with Gaussian process conditioning. This makes our approach much faster than training an LFM on individual tasks. This method can model complex nonlinear dynamics and pro- vides solutions even to multivariate problems such as partial differential equations which were previously computationally infcasible for large datascts.
2 Preliminaries
Gaussian processes Gaussian processes are stochastic processes often used priors for latent functions in Bayesian machine learning models that map from inpuls E RR to predictions f(x) E R. A GP prior
f k(z,,a")
(1)
is described by its mean function m(x) and its kernel function s(x, The mean function is usually set to 0 for standardised data. The kernel function may have a set of hyper-parameters 0, such as the lengthscale in an RBF kernel, kRRF(* x¹) = exp(- -11 - - al|2)) Under this prior, any finite collection of points f(X) for inputs X [x), Ø2. .,.....,........................ NN is normally distributed: f~ ((m((X), x(X, X)). In this paper, we denote the covariance Kf n((X, x). With Gaussian likelihood for observations y, meaning y -Nff.o)) the posterior distribution for training data X, y is analytically tractable and given by
fl fly fly~N (M~N(M(,XX)MXXX + o'1l-''0 "11l"y,,, n(x, x') k(x,, X [K(X, X) + +o'I xxx(xxxxxe)
'x(xxa'))
Morcover the marginal likelihood has a closed form expression enabling optimising the kernel hyper- parameters 0 by maximising the marginal likelihood using maadient-based optimisation and is given by
p(y) Ny 00,((, + +²DI)
Deep Kernel Learning Deep kernel learning presented by Wilson et al. 2016) constitutes an attempt to combine the representation learning capabilities of deep neural norks with the non-parametric nature of Gaussian processes. A neural network is used to map an input x into a latent space yielding vector NN(x) E PR This representation is then fed into a base kernel k(s-) (such as an RBF kernel) to yield the covariance between inputs x(NN(x), NN(22)).
Latent Force Models LFMs incorporate explicit dynamics of differential equations the kernel functions of Gaussian processes (GPs) in order to infer latent forcing terms (Lawrence et al.) [006 Alvarez et al. 2009). The latent force captures the underlying process and structure in the data, while being unobserved
2

---

Published in Transactions on Machine Learning Research (10/2024)
Simulated functions
lotest dinendions
Prediction and latent force inference
Set Encoder oobten 7.
I, MER" *x4
kiXu))
Deep Kernel
true function training point
timepoints
Nu C. CR
NN
RDF
extrapolation
f EER Training only
Training Objective loge(t()) + log
Simulated latent
R
Figure 1: Schematic of DKLFM. First, a dataset of latent force tasks created by sampling the latent force and differential equation parameters and solving the forward solution. The simulated functions are embedded by aggregating the output state of an encoder. A deep kernel is learned to represent the convolution operator of an arbitrary LFM. For training tasks, the model minimises the loss in Equation O with access to simulated latent force data. For test tasks, the latent force is unobserved and inferred via the cross-covariance only, as in a typical LFM scenario. The diagram shows one task; in reality, we train over batches of tasks.
and shared amongst the outputs. The differential equation, 9, parameterised by O enforces mechanistic relationship between the vector-valued output function, h())= E R", and an unobserved, scalar-valued latent force, f(x) E R. with D-dimensional input æ E TRP. This input often time but can be any variable. An ordinary differential equation (ODE) involves derivatives a function of one independent variable, describing how this function changes (D = 1). A partial differential equation (PDE) involves derivatives with respect multiple variables, used model milti-dimensional phenomena like heat and fluid dynamics (D 1).
dh(ar) = (er h G(f(z))),
dæ
or equivalently for a PDE. The force can be transformed by some response function G(-). A GP prior is assigned to the latent force, I ~ GP(O, k(æ, æ²)), which naturally accounts for biological noise and enables non-linear expressivity through kernels. Some LFM literature considers multiple forces but we do not cover this due the identifiability issues they pose. An analytical expression for the covariance between outputs, Kin; is possible under the necessary condition that G is a linear operator. In these cases, maximum marginal likelihood yields the differential equation parameters and inference can be carried out with standard posterior GP identities (see Rasmussen & Williams (2005) Approximations are required where G is non-linear.
3 Deep Kernel Learning of Latent Force Models
In this scction, we present the Deep Kernel Latent Force Model (DKLFM): novel approach to multi-task dynamical modelling. We first detail the problem setting and derive our objective function, and finally discuss any design choices in our approach.
3.1 Model Formulation
The model setup is summarised in Figures I and 2 illustrating the generative process. Our input x is transformed by the latent force f, which is further manipulated by set of differential equations to yield the output functions h. We aim to solve the inverse problem of inferring given noisy observations of h. We start by illustrating our approach with model used later in this paper (Barenco et al. 2006). The time derivative of the mRNA, h(x), for P genes is related to its regulating transcription factor protein, f()) by
basal rate
response decay term
dh(x) dæ:
ge (æ, h, G(J()))) b +8 ((f()))) ch(x)
(4)
where x E R time and s, e RS are the base transcription rates, scnsitivitics to the transcription factor, and docay rates of the genes respectively. G function, for example a nonlinearity enforcing positivity saturation term enforcing limits on the latent force. The ODE parameters thus O {b, c...
3

---

Published in Transactions on Machine Learning Research (10/2024)
y
Figure 2: Graphical model illustrating the generative process from latent forces to output functions.
At training time, we assume a simulated datasct of N tasks, each consisting of a different latent force function and setting of parameter values. The n-th task is the set (Xn, Yn: Jn), where X, E R denotes T observed D-dimensional imput points, e.g. temporal (D 1) and apatio-tempore (D > 1). We have noisy observations, Y, RTT,, of the P-dimensional solutions to the differential equation at each of these T imput points. The solutions are assumed to follow the GP h with realisations hn where the P dimensions are flattened. The observations Yn are similarly condensed into the blocked vector yn- Since we carry out training on a simulated dataset generated by the LFM, we also have access to latent force observations, fn E IR.. For clarity, we henceforth omit the subscript task index and can be assumed that we are dealing with each task independently. These tasks are split into train and test sets, At inference time, our objective is infer the latent force and the model makes use only of the output function observations.
We assign GP prior to ʃ and use Gaussian likelihood, i.e.
f GP (P(mm)()mmp (,,((,)))) y ~ Nh, I2),,
(prior) (likelihood)
where my is the mean function of the GP prior for which we learn constant output, and Ky is the kernel.
For each task, the distribution over output realisations h is implicitly determined via their joint distribution with the latent function outputs, f. In practice we make a modelling assumption that this joint is multivariate Gaussian. Inferences are then made using GP conditioning on the output observations for test tasks. This results in the joint distribution
N ~
(
[ [4
KJs
Kfh Kmn
(joint)
Knf
where the mean vectors and covariance matrices obtained as described below. This assumption of multivariate Gaussian joint distribution holds only for linear transformations of the latent force (i.e. where G is linear function). Nevertheless, we demonstrate in Section 4.3 that this approximation yiclds strong quantile similarities to the true empirical distribution even for nonlinear transformations.
Deep kernels In an LFM, the kernel function is derived from solving set of differential equations. In the case of a general non-linear equation, the kernel has no closed-form solution. We instead approximate the kernel with neural network of sufficient capacity, mapping the imputs to latent representation vectors of size L.a before feeding them into base kernel, K : Te x Ia + R (iilson et al. 2016). This base kernel provides an additional inductive bias, for example ensuring smoothness in the latent space with an RBF kernel or periodicity with a periodic kernel. To map to this latent space, we construct two separate networks NN, : D+L, + R'a, NN, D+, + Rhe for the latent and output functions respectively, where L, is the size of the task representation, r E PxL,, discussed further in the next section. We also nood common network, NN. RLa + RLa, which maps both representations onto the same latent space. This common network helps to obtain informative cova-cova between latent force and output functions: need to map both to common latent space. For simplicity, select simple MLP for all networks. We found that incorporating skip connections greatly improved the performance whilst reducing overfitting, To ralustrate for task n, the cross-function covariances, Kyn and Knh, are computed as follows. The deep kernel receives a concatenation of the task representation with the inputs, 2n = Tn + En.
K fh = # (NN,(NN))..),,, NNN, (=())), KEn = K (NN, NNN,
(5)
The same concatenation occurs for the latent covariance Kşf and for the mcan functions, e.g. H2 = my (zn). The covariance matrices are therefore of shape Kss E RRTX, Kfx E RTXT.P. and KAn E PXP-X-PP
In some experiments, particularly periodic scenarios, helped to concatenate the input mesh after applying the neural network. This granted the periodic kernel access to both the latent vector and input mesh. For example, based on Equation Kft Fierrioicc (I,, < (NN(NN)2)))) sr NN.(NN, (N))),(<,))).
4

---

Published in Transactions on Machine Learning Research (10/2024)
The deep kernel weights and the base kernel and mean function hyperparameters optimised jointly by maximising the marginal likelihood of the output functions and latent force. This has the closed form:
f(v.f) = p(u. (.) dh dh / plulh))((h)))) f)(()
= =v(f) / Noyl h od I) PAIs+ Knis) dh =N(f H4,, K PMy. KNys + r²1).
(6)
where in practice the negative log marginal likelihood is minimised and log.N(y) Pnlfs Knys + 0²1) = -2 - MMns) (Knlf + +021) -x -Mus)) - log KNis 0²1
Here, Mays and Kulf are defined as
PMl = Ph + Knf KiiKit f, KNs Kinh Ki,Kij Khh-
(7)
(8)
At inference time, we receive an arbitrary input x". At these input locations, we define f as the inferred and unobserved latent force using the posterior predictive distribution, defined by
Pi-ly = Ms. + K K K fxly Kj.f- K for Kon hh Kafe.
(9)
(10)
Note that using exact GP inference in this approximate model. If the input space is very large, then the matrix inversion can become computationall challenging and variational approximation may be easily interchanged here. We will now discuss specific components of the model in more detail.
3.2 Task Representation
The generative process relies both the latent force and set of task-specifie parameters relating to the dynamics equations, such as reaction and decay rates. By tweaking these parameters, it is possible that latent forces from different tasks could produce the same or similar output functions. Moreover, tasks on the same input mesh would result in identical covariance matrices, limiting cross-task utility. The GP therefore nceds accoss to a task representation in order to generalise across tasks. We denote this r, = emb(æ,, Vn), defined by emb PxT + RPL,, where the embedding size, L,, hyper-parameter and each of the P output functions are treated independently. Note that it does not observe any latent force data. Instead, this representation is used by the deep kernel to learn the relationship between the output functions and the latent force in its latent space. This separates inference of latent forces from the dynamics parameters, whilst enabling the task representation to be computed for tasks where no latent. force observations exist.
This embedding must contain the dynamics information usually captured by the differential equations and their parameters, and should be invariant to input resolution in order to maintain the flexibility of Gaussian processes. This also enables super-resolution inference; test cases can be at an arbitrary resolution higher than the training data, as we demonstrate Scction 4.2 To that end, we explored two different encoders in our research: Fourier neural operator (Li et al. 2020 and a Transformer (Vaswani et al., 2017).
The Fourier neural operator is effectively an MLP in the Fourier domain, and originally developed to solve partial differential equations (PDEs) due to the intricate relationship between differential calculus and Fourier analysis, their mesh-invariance,, and fast solving speeds. However, a severe limitation is that the Fourier transform requires the input to be a regularly spaced. On the other hand, the Transformer is invariant resolution and regularity of observations. We implemented the decoder only consisting of lincar layer followed by self-attention layers. Sinusoidal positional encoding enables the modelling of an irregular mesh. The weights trained by maximising the marginal likelihood in Equation
5

---

Published in Transactions on Machine Learning Research (10/2024)
Thain
Test tacks
Outpats mRNNA count
Time
Latons DiLNN count
pelatios
Figure Training and test transcriptional regulation tasks. Notice that even for test tasks, the learned variance encapsulates most ground truth. Test tasks do not have access to latent force data.
Outputs
Latent force
Outputs
Latent force
mRNA count
leak
Predator
64 61 64 14 Time
Time
Time
Time
(a) DKLFM infers the protein concentration of tran- scription factor p53. ground truth was published by [areenco et al. (2006). The model was trained only
(b) DKLFM infers the predator-prey relationship in Lotka- Volterra setup. The model has only been trained within the time range denoted by the grey shading, and extrapolates the periodic nature beyond.
on simulations of Equation
Figure DKLFM extrapolation both tasks and input domain.
4 Experiments
In this section investigate the performance of DKLFM two ODF-based LFMs and PDE-based LFM. Given that this is the first multi-task model for latent force models, analyse the performance real, experimentaall--eriieel datasets not in the synthetic training distribution. We compare our approach with two models from the literature on solving latent force problems. Alfi is a variational approximation with a strong mechanistic prior, mcaning the model constrained by the diffcrential equations. For lincar 9, Alfi resorts to the exact solution. The DeepLFM involves dynamics-informmn random features which are composed with each layer. This results in only a mid-strength mechanistic prior, since the deep representation is not fully constrained by dynamics. Our approach has weak mechanistic prior since it is not encoded directly into the model, but rather it is learnt.
4.1 Nonlinear Ordinary Differential Equations
The first ODE model the similar to the original application of latent force models [Lawrence et al., 2006): the biological process of transcriptional regulation. We validate additionally on experimentally-deeivee data from Barenco al.) (2006), where cancer cells were subject to ionising radiation and the concentration of mRNA was measured via microarray at different timepoints. The data contains transcript counts for five targets of the transcription factor p53 over seven timepoints and three replicates. We also consider the paired Notka-Volterra equations, which govern predator-prey dynamics and exhibit periodic solutions.
Latent Force Setup The transcriptional regulation task is defined by the ODE in Eq. 4 where the exact solution is only tractable when the response function the identity. In this case, we set G to the softplus function, G(f(z)) +xg(1+ (xp(f(=))) to ensure positive protein abundance. We start by sampling parameters for Equation from empirical distribution of parameters learnt by running the Alfi Mosss

---

Published in Transactions on Machine Learning Research (10/2024)
Table I: Comparison to baseline models for the transcriptional regulation ODE and the the reaction diffusion PDE. For the DKLFM, we train on a dataset of 256 and 384 tasks for the ODE and PDE tasks respectively, Alft and DeepLFM optimise each task independently. Results are averaged over 20 instances. DKLFM- is ablation the common component of the deep kernel is removed by setting NN. to the identity. DKLFM-b is an ablation of the Fourier embeddings, replacing them with 4-layer MLP. Training was an NVIDIA GeForce RIX 4090 GPU. Due to the differentiable PDE-solving package used by Alft not being GPU-compatible, the PDE tasks were fit. on an AMD Ryzen 5600x CPU. NLL is the negative log-likelihood, and the time column corresponds to the inforence time per-task.
TASK
MODEL
MSE (LATENT)
MSE 4 (OUTPUT)
NLL 4
TIME (s)
MECHANISTIC
ODE
Alfi (exact)
0.117
0.0155
-1.29
3.27
Strong
ODE
DeepLFM 0.108
0.0332
1.42
12.6
Mid
ODE
DKLFM
0.0028
-2.39
1.18 X 10 10-2
Weak
ODE
DKLFM-a 0.898
0.0027
8.64
1.01 X 10
Weak
ODE
DKLFM-b
0.321
0.0075
-2.07
9.21 X 10-3
Weak
PDE
Alfi (approx) 0.0886
0.0215
-0.727
> 600
Strong
PDE
DeepLFM
0.356
0.547
96.7
Mid
PDE
DKLFM 0.131
1.92 X 10 10-7
-1.96
0.0523
Weak
PDE
DKLFM-a
0.998
3.01 x 10-
-0.799
0.0606
Weak
et al.
2021 latent force inference package on the p53 network of genes experimentally measured by Barenco et al. (2006). Next, the latent force is sampled from GP prior with RBF kernel, and the ODE is solved numerically using any differentiable solver, yielding a single task.
The Notka-Volteree tasks are defined by the equations:
du(æ) cru(x) - (u(x))(x) dx
do(x) yu())o(()) 60(x),
(11)
dr
where u(x) and v(x) are prey and predator populations respectively as a function of time, with growth rates a and 7, and decay rates 3 and 8. We use a periodic kernel for x in Eq. 5 for this task in order to capture the periodic nature of the Lotka- Volter- solutions, thus improving temporal extrapolation. We assume that we seek to infer the predator concentration from the abundance of prey; i.e. we take the predator population be the latent force. We simulated dataset of Lotka- Voltera solutions corrosponding to different sampled rales, a, 3 U(005, 1.5), using 4th-order Runge- Kulla solver. This is slightly different regime to the transcriptional regulation model, where a Gaussian process was sampled for each datapoint. In this case, we validate our model's ability to infer latent force that was not explicitly generated by a Gaussian process.
In both cases, Gaussian-distributee and split into training, validation, and test tasks. Figure demonstrates that DKLFM can infer distributions over latent forces for the task of transcriptional regulation. We then apply the model trained on the simulated dataset to a real microarray dataset from Barenco et al. (2006). and show our inferred transcription factor concentration alongside the unobserved ground truth in Figure 4a Next, we demonstrate the intra-task extrapolation in Figure 4b where the input has been extended into the past and future. Finally, in Table I we compare our results with closest models from the literature, finding lower errors and computation times.
random noise is added to the latent forces. We generate 500 instances
4.2 Partial Differential Equations
PDE-based LFMs are the multivariate extension of ODEs and are significantly harder to solve. This is illustrated by the absence a method capable of exactly solving all classes of PDEs. Numerical solvers typically operate on a mesh and thus suffer the curse of dimensionality. Here, demonstrating the ficxibility of DKLFM, we fit reaction diffusion equations with very moderate dataset of 384 low-resolution tasks. The test tasks then be inferred at much higher resolution compared training time.
7

---

Published in Transactions on Machine Learning Research (10/2024)
Latent Force Setup Reaction diffusion equations have many uses; in this paper we look at the bio- logical process of Drosophila embryogenesis (formation of the fruit-fly embryo). The spatiotemporal RNA
expression, h((,t) of gap genes is measured using a reaction diffusion PDE from López-Lopera et al. (0019)
ah(a,t) t) = sf(art) s((e,) Ah(a,t) + d Ph(zz,t) 01
(12)
Ox²
Here, 8 is the production rale of the driving mRNA, [(x, 1) is the latent force, X is the decay rale and d is the diffusion rate. Notice that the latent force is 2-dimensional; DKLFM can take any multivariate input.
In order to simulate a datasct from Equation 12 we implemented the Green's function approximation from et al. (2019). This approximation gives the full covariance matrix, including cross-covariance, between latent force and outputs, and is faster in this direction than Alfi. Since the joint covariance matrix is singular due to repeated inputs, sampling is implemented using the eigendecoompoitiion rather than Cholesky, We generate 448 tasks in this fashion, taking less than two hours on an AMD Ryzen 5600x, of which 384 are used for training. From an empirical inspection of the gap gene dataset from [Bekkrr et al. (2013). we uniformly sampled production rates in the range [0.2, 1.0], decay rates in the range [0.01, 0.4], and diffusion rales in the range (0.001, 0.1]. For the latent force, we sampled the two lengthscales (corresponding to spatial and temporal dimensions) in the range [0.1, 0.4] since both dimensions are normalised to [0, 1].
We show that we can learn a general solution operator for PDE tasks in Figure 5 invariant to imput resolution. In Table D we show how our framework compares against single-instance models. Alfi tends to obtain very accurate results due to backpropagating the loss through the solver. However, DKLFM beats Alfi in output MSE with a greatly reduced computational burden, and using a modest dataset of 384 tasks. DKLFM also achieves a competitive error for the latent function.
4.3 Performance and Approximation Cost
The utility of LFMs for large scientific datasets is limited by their lengthy training times. In genomics, a realistic scenario is where a bioinformatician will want to train an LFM on several thousand genes; for example, RNA velocity (La Manno et al., 2018) solves a splicing kinctics ODE on the entire human genome. This limits the use of the available approximations. Moreover, the Gaussian joint distribution used in this work is model assumption needing verification. We therefore qualitatively demonstrate the uncertainty in our approximated posterior predictive distribution using quantile-quantile plots.
Our framework, however, solves many LFMs simultaneously rather than optimising a single instance. An analysis of the relationship between error and training set size is therefore key to finding the point our error rate drops to the level of solving an individual instance. If this point is less than or similar to the number of instances in a typical use-case, then it is computationally preferable to generate a simulated dataset of this size rather than to train individual LFMs. This is because the training dataset required to reach this performance becomes smaller than the evaluation set, making it more time-efficient to use the DKLFM. For this study, we compare against Alfi, an accurate nonlinear LFM approximation defined in Moss et al. (2021). We chose the ODE task, since the PDE solver in Alfi is too computationally intensive for this comparison. In Figure G we confirm our hypothesis by plotting the MSE versus datasct size for our model, and horizontal lines are the mean MSE for Alfi over a subset of 64 tasks.
In order to analyse the posterior uncertainty of DKLFM, we use quantile-quantile plots showing samples from our posterior compared with real, simulated samples. We do this for three different tasks: 1. the nonlinear transcriptional regulation ODE model in Section 4.1 2) a logistic growth model with sinusoidal nonlinearity; and 3) the reaction-liffusion experiment from Section 4.2 The logistic growth model measures the increase of a resource, for example a population or a plant, as it reaches its maximum value. The rate of increase in size or quantity of the resource at time 2 E RR, h((x), is expressed with the ODE
dh(æ) rh((x) (1 (1- h(xx) dæ
K (z)(()))
(13)
where r is the growth rate, K is the carrying capacity (maximum resource size). B is the response factor to the temperature, (ax), and G is the cosine function. This equation was chosen since it can result in
8

---

Published in Transactions on Machine Learning Research (10/2024)
Observed Outputs
Space
Predicted Outputs
(super resolution)
Space
-2
2
Predicted Latent Ground Latent Force (super- )
Spucee
Spacee
Timze
Timte
T Timee
Time
Time
Figure 5: DKLFM trained on synthetic reaction diffusion datasct. The first column is training example and the next four are test cases, where the latent force is not observed. The embedding size was increased to 96 to account for the increase dimensionality. The model was trained with 21 21 spatiotemporal grid. prediction time, 40 x 40 grid was used to illustrate the super-resolution capability. Each pair of plots vertically shares the same colorbar enforce the same scale and accurately demonstrate inference accuracy.
multi-modal distribution due to the sinusoidal response to temperature. The QQ-plots shown Figure i demonstrate that our assumption that the output function can be distributed as a multivariate Gaussian leads to robust uncertainty quantification. Even for the edge of the logistic growth model, the associated QQ-plot exhibits close fit to the unimodal Gaussian distribution.
5 Related Work
Differential equation-based inference in dynamical systems with Gaussian processes introduced in Lawrence et al. (2006) and Alvarez et al. (2009). These approaches derive the kernel functions by solv- ing the convolution integral of a base kernel with a linear operator corresponding to the ODE solution. When the dynamics becomes nonlinear, the Laplace approximation was used for the marginal likelihood. The benefit of this approach that entirely non-parametric and enforces strong mechanistic behaviour. The primary issue is the requirement of analytically solving the specific ODE, which not possible for some equations. Moreover, the first and second derivatives of the nonlinearity also required. Also in line with

---

Published in Transactions on Machine Learning Research (10/2024)
0.6
Ousput MSE (eurs)
0.06
Laient MSB (ours)
Output MSE (AIS)
0.3
0.05
(Alf)
Ontpot MSB
Latent MSE
0.00
0.4
0.03
0.3
0.31
0.2
0.00
0.1
300
Nomber of training tasks
Figure 6: We plot MSEs against training dataset size for the genetic regulation LFM. This demonstrates the point at which becomes more economical to use DKLFM rather than single-task models. At around 200 instances, the performance of DKLFM matches that of a single-task model optimised by Alfi.
Transcriptomics
Growth
Reaction Diffusion
0.3
Quantiles of approximated data
Quantiles of approximated data
1.75
1,50
Quantiles of approximated
0.2
1:25
1.00
0.0
0.75
0.50
0.25
Quantiles of true data
Quantiles of true data
Quantiles of true data
Figure 7: Quantile-quantile plots demonstrating the impact of our model approximation of a multivariate Gaussian join distribution in Section 3.1 We plot sorted samples at various timepoints linearly spaced across the imput domain for three tasks. The transcriptomics task exhibits some deviation from an exact quantile fit due to the nonlinear transformation of the latent force. The reaction diffusion task involves only linear transformations of Gaussian process resulting in a very tight close quantile fit.. The logistic growth model can result in multinodal (and therefore non-Gaussian) distribution.
this method is the switching LFM (Alvarez et al. 2010a which segments time and switches between a number of latent forces, with only single force being active in each interval.
Filtering approaches have also been investigated (Hartikainen & Sarkka, 2012 Särkkä et al. 2018 Ward et al.) 2020). They typically employ state-space model for approximating the postcrior of a non-lincar LFM. For example, Ward et. al. (2020) use autoregressive flows to construct joint density of the state using variational inference, bypassing complex kernel derivations. This results in flexible model, scalable to multiple latent forces and output functions. However, the model exhibits the over-confidence prevalent in such black-box variational approaches.
Alfi Moss et al. 2021) avoids the complex derivations kernel functions by sampling the latent force from the GP prior and gradient matching to pre-ostimate reasonable paramoters. An ODE or PDE solver then used fine-tune the parameters with the solution of the equations. Backpropagating through a solver is far too computationally intensive for multi-task setting, rendering this approach impractical for large datasets.
10

---

Published in Transactions on Machine Learning Research (10/2024)
McDonald & Alvarez (2021 tackles non-linear and non-stationary dynamics by constructing deep GP
(Damianou & Lawrence 2013). At each layer, an RBF kernel is convolved with the Green's 's function of the ODE. This deep representation enables the modeling of a wider range of tasks than a standard LFM. It is, however, not directly applicable to PDEs or to a multi-task sctting.
There also an adjacent area of research applying deep kernels for dynamical modelling. For example, Botteghi et al. (2022) denoise observations of the output functions via variational autoencoders, and learn neural network representation of dynamics. In addition, Chen et al. (2022) propose meta-learning framework for deep kernels consisting of a bilevel optimisation problem. The inner level consists of learning subset of parameters which are adapted on individual tasks. At the outer level, the remaining subset of parameters are trained to yield the best loss on average over training tasks after the inner level is optimised. These approaches do not, however, infer latent forces. In this work, we treat the meta-network as an entirely separate component which is trained end-to-end with the deep kernel network.
6 Conclusion
We have introduced novel framework for latent force models by leveraging the expressive power of deep kernels combined with a learned task representation. Where standard LFMs require an optimisation loop to find kernel parameters, our approach only requires GP conditioning on observations at prediction time, enabling extremely fast latent. force inference. Specifically, this involves inverting a T x T matrix with O(T²) computational complexity. If the input size is too large, a technique such as variational inducing points reduces the computational complexity to O(T with M inducing points. DKLFM is therefore an ezact inference proba bilistic model: the first of its kind for learning the approximate solution operator for an arbitrary nonlincar LFM. We achieve this by learning deep kernel corresponding to the differential equation by training a simulated dataset of tasks. The embedding of each task's observations are interpreted the task representation, containing information such as rate parameters. At test time, this representation is combined with an arbitrary poutt-ossiibl of higher resolution to the training data-in order to compute the latent forces using the Gaussian process conditioned on observations.
The inference performance of DKLFM is reliant on learning a good cross-covariance between latent forces and observations. We hypothesise that this is why this model does not exhibit the tendency for over-confidence away from training data in predictions commonly found in deep kernel learning and related approaches. Since we learn the same deep kernel over dataset of tasks, this has proved not to be an issue. We demonstrate in two ablations the importance of both the common component of the deep kernel, NN,, as well as the Fourier embeddings in lcarning a good representation of the dynamics.
Limitations Speed is limited by generating the training dataset; however, this is easily parallelised. More- over, as in our PDE experiment, the generative direction, i.e. going from parameterisation to an instance is much faster than the inverse problem of inferring the latent force from the output functions. We envisage these models being used analogously to large language models (Shanahan 2022), where user can fine-tune a pretrained DKLFM to make latent force inferences on their dataset. This would be particularly useful for computationally intensive simulations.
The original LFM is strongly mechanistic, deriving the covariance function from strict dynamics equations. While this may be more robust in the presence of lots of noisc, is overly rigid for real-world tasks. Our approach is weakly mechanislic: dynamics are not imposed, but rather parametrically learnt from the data and paired with nonparametric GP to condition unseen data. As with AlphaFold (Jumper et al. 2021 combining biophysical priors and data-driven approaches may be more appropriate for complex problems.
Despite not being an issue for the examples covered in this paper, our model approximation can lead to incorrect confidence predictions. Due to modelling the joint as a multivariate Gaussian, and therefore h as a GP, the DKLFM theoretically cannot model transformations of the latent force that lead to non- non-Ganssian distributions over the output, h, such as those with multiple modes. However, Section obtain close quantile fits even for multimodal distribution.
shows we still
11

---

Published in Transactions on Machine Learning Research (10/2024)
Further work DKLFM explicitly treats the uncertainty in the latent forces and output functions. In active learning context, we can query input points where the output function latent forces have high- uncertainty. This enables experiment design, for example to determine an appropriate coarseness for a time-couree experiment. Furthermore, we have currently considered one latent force per task. Multiple forces lead to identifiability issues, where many combinations of latent forces would solve the same LFM.
12

---

Published in Transactions on Machine Learning Research (10/2024)
References
Mauricio Alvarez, David Luengo, and Neil D Lawrence. Latent force models. In Artificial Intelligence and Statistics, pp. 9 16. PMLR, 2009.
Mauricio Alvarez, Jan Peters, Neil Lawrence, and Bernhard Schölkopf. Switched latent force models for movement segmentation. In J. Lafferty, C. Williams, J. Shawe-Taylor, R. Zemel, and A. Culotta (eds.), Advances in Newral Information Processing Systems, volume 23. Curran Associates, Inc., 2010a. URL https: pttp::////oeeeegg...m.m neurips eeeurips.ccpappee.. filess/paper/210/EEi. (o00ffo4rradd2//37ccc55555d
pdf.
Mauricio Alvarez, Jan Peters, Neil Lawrence, and Bernhard Schölkopf. Switched latent force models for movemnett segmentation. Advances in neural information processing systems, 23, 2010b.
Ann C Babtie, Paul Kirk, and Michael PH Stumpf. Topological sensitivity analysis for systems biology. Proceedings of the National Academy of Sciences, 111(622:18507-18812,, 2014.
Martino Barenco, Daniela Tomescu, Daniel Brewer, Robin Callard, Jaroslav Stark, and Michael Hubank. Ranked prediction of p53 targets using hidden variable dynamic modeling. Genome biology, 7:1 18, 2006.
Kolja Becker, Eva Balsa-Cannn Damjan Cicin-Sain, Astrid Hoermann, Hilde Janssens, Julio R Banga, and Johannes Jacger. Reverss--ngineering tost-rranncriptionnl regulation of gap genes in drosophila melanogaster. PLoS computational biology, 1000),, 2013.
Nicolò Botteghi, Mengwu Guo, and Christoph Brune. Deep kernel learning of dynamical models from high- dimensional noisy data. Scientific reports, 12(1):21530, 2022.
David N Burghes. Population dynamics introduction to differential equations. International Journal of Mathematical Education in Science and Technology, 6(3):265-2P) 276, 1975.
Wenlin Chen, Austin Tripp, and José Miguel Hernández-Lobato, Meta-learning adaptive deep kernel gaussian processes for molecular property prediction. arXiv preprint XXXi:2205.00700 2022.
Li-Fang Cheng, Bianca Dumitrascu, Michael Zhang, Corey Chivers, Michael Draugelis, Kai Li, and Barbara Engelhardt. Patient-specific effects of medication using latent force models with gaussian processes. Tn International Conference on Artificial Intelligence and Statistics, pp. 4045 4055. PMLR, 2020.
Andreas Damianou and Neil D Lawrence. Deep gaussian processes. In Artificial intelligence and statistics, pp. 207-215. PMLR, 2013.
Jouni Hartikainen and Simo Sarkka. Sequential inference for latent force models. arXiv:1202.3P7 2012.
arXiv preprint
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidck, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nalure, 596(7873):583 589, 2021.
Gioele La Manno, Ruslan Soldatov, Amit Zeisel, Emelic Braun, Hannah Hochgerner, Viktor Potukhov, Katja Lidschreiber, Maria E Kastriti, Peter Lönnerberg, Alessandro Furlan, al. Rna velocity of single cells. Nature, 560(77199:4444499,, 2018.
Neil Lawrence, Guido Sanguinetti, and Magnus Rattray. Modelling transcriptional regulation using gaussian processes. Advances in Neural Information Processing Systems, 19, 2006.
Zongyi Li, Nikola Kovachki, Kamyar Aviwadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXfv:2010008888 2020.
Andrés F López-Lopera, Nicolas Durrande, and Mauricio A Alvarez. Physically-inspired gaussian process models for post-transeriptional regulation in drosophila. JEEE/ACM transactions on computational biology and bioinformaties, 18(2):656 666, 2019.
13

---

Published in Transactions on Machine Learning Research (10/2024)
Thomas McDonald and Mauricio Álvarez. Compositional modeling of nonlinear dynamical systems with ode-based random features. Advances in Neural Information Processing Systems, 11138899138877 2021.
Jacob D Moss, Felix L Opolka, Bianca Dumitrasou, and Pietro Lió. Approximate latent force model inference. arXiv preprint Xrx::22099 11851, 2021.
Carl Edward Rasmussen and Christopher K. T. Williams. Gaussian Processes for Machine Learning (Adaptive Computation and Machine Learning). The MIT Press, 2005. ISBN 026218253X.
Tapesh Santra. Fitting mathematical models of biochemical pathways to steady state perturbation response data without simulating perturbation experiments. Scientific Reports, 8(1):11679, 2018.
Simo Särkkä, Mauricio A Alvarez, and Neil D Lawrence. Gaussian process latent force models for learning and stochastic control of physical systems. IEEE Transactions on Automatic Control, 64(():2953-22)) 2960, 2018.
Birgit Schoeberl, Claudia Eichler-Jonsson, Ernst Dieler Gilles, and Gertraud Müller. Computational model- ing of the dynamics of the map kinase cascade activated by surface and internalized egf receptors. Nature biotechnology, 20(4):370-375, 2002.
Murray Shanahan. Talking about large language models. arXiv preprint arXiv:2212.033: 03551, 2022.
René Thomas, Anne-Marie GATHOYE, and Lucie Lambert. A complex control circuit: Regulation of immunity in temperate bacteriophages. European Journal of Biochemistry, 71(1):211-227, 1976.
Alan Mathison Turing. The chemical basis of morphogenesis, Bulletin of mathematical biology, 52:153 197, 1990.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Atlention all you need. Advances neural information processing systems, 30, 2017.
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, pp. 3088-3099 PMLR, 2020.
Andrew Gordon Wilson, Zhiting Hu, Ruslan Salakhutdinov, and Eric P. Xing. Deep kernel learning. In Proceedings of the 19th International Conference on Artificial Intelligence and Statistics, volume 51, 2016.
14

---

Published in Transactions on Machine Learning Research (10/2024)
A Appendix
A.1 Pseudocode
Algorithm Training procedure of DKLFM
for each task do
f- emb( Vn.)
> Compute representations
Zn Tn, @ Dn
D. Concatenate inputs
K f = r. (NN, (NNN, ( (Zn) )), (NN,NN,
Latent force covariance Cross-function covariance
K fh (NN, (NN (Zn)
(NN, (2))))
Kihh (NNG((NI) (Zn)
NN.(NNN (2,))
> Output covariance
PAIS 1A Knf K if
KNls = KRN Knf KJj 2- Kth
p(y. f) [ p K (1))((y)) Mal Kulf + 021)
backpropagate Io
end for
Algorithm 2 DKLFM inference
for each task do f emb(#n, Un)
Compute representations >> Concatenate inputs
b En Fx <- emb( (a; y%)
2 Tn a x
KA (NN, (z,, NN. (NN, (z))))
> Output covariance P Output covariance
K y-h INNOONNN
(NN,
K fsh (NN (2,,
(NN
Cross-function covariance
M f p0 + K Jon hh
K K Ky - K Ki An Kif-
Myrly = Myn + K Kin
K Kully K you+ yo K hh K Kay end for
return NW faly K (Hgglg K ky-ly)
15


Reviews:
[{"review_id": "ooQFdWI9vN", "reviewer": "Reviewer_gFy4", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper motivates using deep kernel learning (DKL) to improve some shortcomings of LFM. The proposal is illustrated with a model of the derivative of the mRNA (eq.  4). This emphasis on models for genetics is also adopted in the experimental part. \n\nThough the paper can contribute to the application of the DKL+LFM to genetics, there are some inconsistencies in the delivery of the contribution. \n\n- Initially, the contribution is presented as a DKL extension to LFM. Theoretically/conceptually, I am not sure about the significance of such a contribution, given that the advantages of DKL are known in the general case and should apply to this particular case too\n- The paper stated that the introduction of DKL will address issues with tractability, scalability and performance. However, this is tested only in a particular dataset type, meaning that the claim about the superiority of the proposed model, in the general case, is vaguely supported. I think that the general applicability of the proposal is debatable. \n- The results (Table 1) show a dramatic improvement in terms of performance & computational cost. To give the reader more insight into these impressive results, perhaps the authors should elaborate on such differences and discuss the limitations of the proposal (besides the ones in line 313) regarding the cases/problems where its performance is not as appealing as that in Table 1\n- From my lack of expertise in models/applications to genetics models, perhaps the contribution should be presented as an application of LFM to this community rather than a contribution to the LFM/GP community.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The article aims to improve over standard latent force models by bypassing intractability and enhancing scalability and performance.  The authors do so through deep kernel learning, which allows them to avoid solving the diff eqs or resorting to variational approximations. The proposal is validated experimentally against Alfi and DeepLFM on different datasets", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please address my comments above. Also: \n- Why incorporating DKL into LFM is a challenge, and thus, why should it be considered a contribution?\n- How general is the proposed approach? Will it work for datasets of different nature to those considered here?\n- Refer to the results in Table 1: what are the shortcomings (or the price to pay) for such an impressive performance figures?\n- Does the introduction of DKL hinder interpretability in the proposed model? is that relevant in the considered applications?", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "zWiHBie0bQ", "reviewer": "Reviewer_SrVB", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "*Strengths*:\n- paper is well written\n- method is clearly explained, achieves good performance, and appears much more tractable/scalable than existing methods", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a new deep kernel method for learning latent force models, which is applicable to a wide array of problems that make use of differential equations.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "*Content*:\n- (suggested, not critical) Offer more systematic discussion of the complexity and scalability of the proposed method compared to existing methods. It's mentioned that there's a speed up of 250x to 15000x, but that's on a few specific tasks, right? What do you expect the pattern to be more generally, and is the evidence these few applications, or theoretical results about complexity, or something else? The complexity of one part of the approach is briefly mentioned in the discussion, but it's not clear how this relates to the complexity of the other parts of the approach or how it compares to the existing methods\n- (suggested, not critical) The abstract mentions that the proposed method is faster than existing methods, but the results also show that the performance is better, right? Could be worth also mentioning this in the abstract.\n\n*Formatting*: \n- remove the strange bars appearing in the right-hand margin on most pages\n- fix the width of the table on p. 7\n- the text in many of the figures is too small to be legible; remove any unnecessary text and make sure what remains is legible", "broader_impact_concerns": "n/a", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "pNIxV2Lkgf", "reviewer": "Reviewer_pBCA", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "S: The main idea is sensible and strong: adding DNNs to LFM is a great idea!\n\nS: The results are good\n\nW: The paper is written for LFM audience, and I struggled a lot to understand the paper due to numerous ambiguities, implicit assumptions and sloppiness in the mathematical presentation, and also in terms of claims, research setting, motivation and problem description. The research setting is barely defined (ODEs are first described half-way in the paper), the probabilistic model needs more exposition, and the embedding part is barely even defined. The presentation clarity needs to be significantly improved and clarified such that the method and contributions are un-ambiguous, precise and rigorous. Otherwise the paper is well written and flows well.\n\nW: The results indicate the the system can learn or solve simple ODEs, and has competetive performance two few baselines. The main issue is no ablation comparison to a standard LFM without embeddings or deep kernels, and the paper needs to show how these two contributions improve.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes to couple kernel input-warping to the LFM model, which interpolates PDE solutions from an observed dataset of related PDEs. The main contribution is adding a neural network inside the kernels to enrich the otherwise quite rigid kernel model. “Deep kernels” are known to be a useful approach, and adding it to LFM models is a good idea and has clear potential for dramatic empirical improvements.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "To be accepted, the paper needs to\n\n(i) Define its claims, open problem, and contribution rigorously; and define the applications or usecases of the approach [demonstrate interest and claims]\n\n(ii) Improve the clarity of the mathematical presentation\n\n(iii) Add ablations to explain how the contributions (embedding + NN_) improve [claim evidence]\n\nI'm looking forward to author clarifications and revisions. I hope my critique below is not taken as malice: I was genuinely confused.\n\n\nRunning comments reading the paper for first time\n\n- Sec 2 LFM description is confusing. It is not clear what is the nature of x (do we have T scalar values or one T-dimensional vector; is this a multi-output system?), what is PxT (is this T vectors of size P, or one matrix?), or what the force does (why it is T → T?). Is is not clear what is the \\cal{D} or what is the differential over. It is unclear what kind of process this describes (is this PDE or ODE or something else?). It’s also unclear where the Dy goes or what we do with it, or why do we care about it. If we do regression, why do we care about differentials? The paper would become much clearer by starting with a definition of ODE/PDE and defining that the dy/dt = h( t, x, f(x,t) ) is the main interest in this paper, and then tying all subsequent math to this master equation. Most of my confusions below stem from the discrepancy between the overly \"static\" probabilistic perspective and the PDE presentation of the experiments, and not defining the PDE setting in the beginning of the paper.\n- I’m confused why the latent force is scalar in 3.1 even though we have P outputs. Isn’t it quite limiting to have scalar force? Can’t we have vector force?\n- It’s also a bit confusing what the X,Y,f,h represent. How should we understand these? How come we are observing the latent forces? It would be good to have an example to concretise what these all mean in some simple usecase. These do not connect with PDEs currently.\n- It’s not clear what “inference time” means. What do you “infer”? Is this network predictions or parameter posteriors, or something else?\n- I think the above means that you don’t use the observed latents during training. But why not? Surely you want to use all the observations to train the model.\n- It’s not clear what “h” means or represents. What is its domain and image? What does it do?\n- What is K_ff or K_hf? Please define these kernels and their domains. I think these are cross-function kernels, but this is not made clear.\n- It’s really strange the the model definition section does not define what kind of function you are learning. Can you define something like y = h( f(x) ) or maybe it’s y = h(x, f(x)) or maybe y = h(x) + f(x). It’s also not clear if this is regression or time-series fitting or something else. What task are you solving? I’m not really sure since the paper does not define what kind of function they estimate. I also can’t see where is the differential D or g() here. Did we forget about them? We just talked about them in eq 3. If they don’t appear in the paper, why did we introduce them? [Looking back, I think the system in this paper is dy/dt = h(t,x,f..), and I did not guess this from first read]\n- mu_f is undefined in eq 4. Is this a prior? (But prior is m_f…). Same issue with mu_h, what is it?\n- What does K_hf mean in eq 5? I don’t understand what kernel is this, or between what. Is the idea that h and f are measured at different input points? Isn’t this a cross-function kernel? Please be rigorous and give some exposition to this.\n- Earlier it was stated that you don’t use f during inference time, but eq 5 seems to use it. Can you explain?\n- Why do you want to infer the f* at testing time? What does f mean? Why do we care about? I thought that the idea is to estimate the output y during testing time (or perhaps the h). What kind of predictions do we want to do, and can you give some example usecase that would concretise these concepts?\n- In eq 4 we define the GP h|f. Doesn’t this lead to the h and f being the same function? If you evaluate eq 5 for h at the same locations as the f, you would just get h=f+mu, ie equal up to a constant. I wonder what is the underlying assumption on how h and f relate? I don’t really see how can you cross-interpolate between two function like this: this looks like just interpolating the function f to new locations. I think this already comes down to the (joint) equation, which looks like a joint of a single GP, but I guess it’s some kind of cross-joint between two GPs. I’m not convinced this is correct or you can do this if you mean to have two different GPs. I don’t think two different GPs are jointly Gaussian. (h has not been defined so I’m not sure what kind of object it is). Is the idea that h=f? I’m quite confused.\n- In sec 3.2. it is mentioned that we are interested in “reactions” and “decay rates”. Ok, so it seems that we are now talking about some kind of physico-chemical temporal systems, perhaps ODEs or SDEs or PDEs. This was not clear from earlier presentation, where nothing like this has been defined. The sec 3.1. does not have time anywhere!! You should link all sections of the paper so that they all describe the temporality and differentiality of the system, and hopefully start from the du(x,t)/dt definitions.\n- It’s unclear what the tasks are. I guess we are now talking about some chemical reaction systems, so are different tasks different chemical systems with different parameters or different initialisation, or both? Please be rigorous.\n- Earlier m_f was defined as prior for f. I guess that the inputs are x. Now we have m_f(z). Since we concatenate, the z is larger in dimension than x. I think you are redefining m_f here. Please make this explicit.\n- Why is K_fh in eq 7 bold and capital? Is this a matrix-variate kernel, or perhaps a kernel matrix? It does not look like it: I think it’s just a scalar kernel value. The notation needs to match the tensor sizes (or please define what the notation means).\n- I’m quite confused what $h$ and $f$ mean in the context of eq 7. Suddenly we map the same input z to both h and f through the two networks. Err… what? I thought that h takes as input the f, which comes from x. Surely you can’t just skip this and go directly x → h and x → f. I’m quite confused what these things mean, or how should we interpret them. How should we conceptualise this?\n- I’m quite confused what the embeddings NN_f and NN_h and NN_c mean or what is their motivation. What are we trying to achieve here? Is this a feature extraction, or dimensionality reduction? What are the input and output sizes of these networks? Why do we have the NN_c? Isn’t it odd to map two different embedding with same network?\n- I also now realise that I’m not sure what x_n means. I think it’s just R^D vector, but it might also be something else, since the discussion kinda implies towards set embeddings.\n- I’m quite confused of the discussion around embeddings and transformer. I’m not sure what we are achieve here, or what is the motivation for this stuff, or what is our goal or problem.\n- I don’t see where the “emb” goes in the equations. Currently it’s nowhere. I don’t understand why the emb takes as input PBTD. Why this choice? I think Y is of dim (B,T,P) and X is of dim (B,T,D) and f is of dim (B,T) How do we then go from these things to the (P,B,T,D)? What is the input to the emb? Why is the output (P,B,L)? Why not (P,L) or (B,L)?\n- The fourier description is way too vague, and I have no idea where these things go in the equations, or what terms are we talking about. I guess this is about NN_f = Fourier or NN_h = Fourier? Or is this about emb=fourier? Likewise for transformer: not clear what parts of the model it operates over.\n- The discussion here mentions meshing multiple times. I wonder where does this go in the equations. Is this just about the choice of X? But isn’t X given us by the data?\n- It is suddenly mentioned that you are mesh-invariant. Where did this come from? Is this just about having a GP? Isn’t every GP mesh-invariant? I’m not sure what kind of argument you are making here.\n- The paper mentions cross-covariance, which was not defined earlier. Can you make this more rigorous: what is the cross part?\n- The paper also mentions that kernels come from some PDE definitions. Can you define this? Currently all kernels seem to be undefined.\n- I’m not sure where the NN : R^L1 → R^L2 go. Does this replace NN_f, NN_c, NN_h? But the equation (periodic) still has these separate..\n- Ok, so looking at eq 8 it seems that y is a function of time. I would not have guessed this based on the basic definitions in sec 3. In Sec 3 the Y (and other stuff) is represented like doing standard regression, and the notion of temporality is not present. Here we then suddenly have full-blown ODEs emerging. I also would not have guessed that the task refers to having different forcings. Please define the research setting earlier in the paper.\n- What is the h in the example systems? What is y, h, f in the LV equations? It would be good to define all of these (eg. in a table). What do you observe in the example systems, what is data?\n- The LV mentions a periodic kernel. But on what or between what?\n- I don’t understand what are the different tasks in the LV. Is this different forcing choices or different initialisations or something else?\n- I’m not sure what you are doing in the experiments. For instance in the LV, what do you learn? I think you learn the forcings, but isn’t this the task? Do you want to learn the time derivative (but the probabilistic equations don’t have time derivatives…)? Or do you want to just interpolate the observed states? Do you run some ODE forward in time? Which one? I’m really struggling to connect the underlying system, the probabilistic model, the data, the assumptions you make, the knowns/unknowns, and the goal you want to achieve.\n- I’m also not sure what you are comparing your method against. What problem do you want to solve? Which methods do you want to improve? Why do you compare against these two methods?\n- I think your learning task is to learn the forcing f(t) while also ensuring the solution y(t) is also accurate. Why do you want to do this? Can you demonstrate why this is an interesting and significant problem? What are the applications of this?\n- In the PDE problem you generate many system solutions under different forcings (I assume). I think you are trying to learn the solution map by using data of related solutions. But you already know the true system (eq 10). If you already can generate 448 runs, why can’t you then generate the 1 test run you actually want to get without doing the 448 runs at all?\n- Looking at fig 5, I think your approach is that you simply extrapolate the solution y_test(x,t) and forcings f_test(x,t) from a bunch of {y_train} and {f_train}, without running the PDE solver. [It would be great to define this as an equation.] I think this type of approach can be useful, but it requires likely a lot of observations of related systems, and will fail unless your dataset contains examples that had similar forcing to f_test and similar initial states to y_test. You also require the knowledge of the true system equation. Can you then explain what is the benefit or usecase for this approach? Why or when should we use this?", "broader_impact_concerns": "No issues", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes to use deep kernel learnings (i.e. standard kernels over a learned embedding of inputs) in the context of latent force models. Despite its simplicity, the method is quite effective empirically, and the reviewers appreciated the improved manuscript. I recommend acceptance.
Decision: Accept (poster)