Scalable marginalization of correlated latent variables with applications to learning particle interaction kernels

Given a set of latent variables in a model, do we fit a model with a particular set of latent variables, or do we integrate out the latent variables when making predictions? Marginalization of latent variables is an iconic feature of the Bayesian analysis. The art of marginalization in statistics can at least be traced back to the De Finetti’s theorem [12], which states that an infinite sequence $\{X_{i}\}^{\infty}_{i=1}$ is exchangeable, if and if only if there exists a random variable $\theta\in\Theta$ with probability distribution $\pi(\cdot)$, and a conditional distribution $p(\cdot\mid\theta)$, such that

Marginalization of nuisance parameters for models with independent observations has been comprehensively reviewed in [7]. Bayesian model selection [8, 4] and Bayesian model averaging [42], as two other examples, both rely on the marginalization of parameters in each model.

For spatially correlated data, the Jefferys prior of the covariance parameters in a Gaussian process (GP), which is proportional to the squared root of the Fisher information matrix of the likelihood, often leads to improper posteriors [6]. The posterior of the covariance parameter becomes proper if the prior is derived based on the Fisher information matrix of the marginal likelihood, after marginalizing out the mean and variance parameters. The resulting prior, after marginalization, is a reference prior, which has been studied for modeling spatially correlated data and computer model emulation [39, 46, 28, 20, 37].

Marginalization of latent variables has lately been aware by the machine learning community as well, for purposes of uncertainty quantification and propagation. In [29], for instance, the deep ensembles of models with a scoring function were proposed to assess the uncertainty in deep neural networks, and it is closely related to Bayesian model averaging with a uniform prior on parameters. This approach was further studied in [64], where the importance of marginalization is highlighted. Neural networks with infinite depth were shown to be equivalent to a GP with a particular kernel function in [38], and it was lately shown in [32] that the results of deep neural networks can be reproduced by GPs, where the latent nodes are marginalized out.

In this work, we study the marginalization of latent variables for correlated data, particularly focusing on scalable computation. Gaussian processes have been ubiquitously used for modeling spatially correlated data [3] and emulating computer experiments [50]. Computing the likelihood in GPs and making predictions, however, cost $\mathcal{O}(N^{3})$ operations, where $N$ is the number of observations, due to finding the inverse and determinant of the covariance matrix. To overcome the computational bottleneck, various approximation approaches, such as inducing point approaches [54], fixed rank approximation [10], integrated nested Laplace approximation [48], stochastic partial differential equation representation [33], local Gaussian process approximation [15], and hierarchical nearest-neighbor Gaussian process models [11], circulant embedding [55], many of which can be summarized into the framework of Vecchia approximation [59, 27]. Scalable computation of a GP model with a multi-dimensional input space and a smooth covariance function is of great interest in recent years.

The exact computation of GP models with smaller computational complexity was less studied in past. To fill this knowledge gap, we will first review the stochastic differential equation representation of a GP with the Matérn covariance and one-dimensional input variable [63, 22], where the solution can be written as a dynamic linear model [62]. Kalman filter and Rauch–Tung–Striebel smoother [26, 45] can be implemented for computing the likelihood function and predictive distribution exactly, reducing the computational complexity of GP using a Matérn kernel with a half-integer roughness parameter and 1D input from $\mathcal{O}(N^{3})$ to $\mathcal{O}(N)$ operations. Here, interestingly, the latent states of a GP model are marginalized out in Kalman Filter iteratively. Thus the Kalman filter can be considered as an example of marginalization of latent variables, which leads to efficient computation. Note that the Kalman filter is not directly applicable for GP with multivariate inputs, yet GPs with some of the widely used covariance structures, such as the product or separable kernel [5] and linear model of coregionalization [3], can be written as state space models on an augmented lattice [19, 17]. Based on this connection, we introduce a few extensions of scalable marginalization for modeling incomplete matrices of correlated data.

The contributions of this work are twofold. First, the computational scalability and efficiency of marginalizing latent variables for models of correlated data and functional data are less studied. Here we discuss the marginalization of latent states in the Kalman filter in computing the likelihood and making predictions, with only $\mathcal{O}(N)$ computational operations. We discuss recent extensions on structured data with multi-dimensional input. Second, we develop new marginalization techniques to estimate interaction kernels of particles and to forecast trajectories of particles, which have wide applications in agent-based models [9], cellular migration [23], and molecular dynamic simulation [43]. The computational gain comes from the sparse representation of inverse covariance of interaction kernels, and the use of the conjugate gradient algorithm [24] for iterative computation. Specifically, we reduce the computational order from $\mathcal{O}((nMDL)^{3})+\mathcal{O}(n^{4}L^{2}M^{2}D)$ operations in recent studies [34, 13] to $\mathcal{O}(Tn^{2}MDL)+\mathcal{O}(n^{2}MDL\log(nMDL))$ operations based on training data of $M$ simulation runs, each containing $n$ particles in a $D$ dimensional space at $L$ time points, with $T$ being the number of iterations in the sparse conjugate gradient algorithm. This allows us to estimate interaction kernels of dynamic systems with many more observations. Here the sparsity comes from the use of the Matérn kernel, which is distinct from any of the approximation methods based on sparse covariance structures.

The rest of the paper is organized below. We first introduce the GP as a surrogate model for approximating computationally expensive simulations in Section 2. The state space model representation of a GP with Matérn covariance and temporal input is introduced in Section 3.1. We then review the Kalman filter as a computationally scalable technique to marginalize out latent states for computing the likelihood of a GP model and making predictions in Section 3.2. In Section 3.3, we discuss the extension of latent state marginalization in linear models of coregionaliztion for multivariate functional data, spatial and spatio-temporal data on the incomplete lattice. The new computationally scalable algorithm for estimating interaction kernel and forecasting particle trajectories is introduced in Section 4. We conclude this study and discuss a few potential research directions in Section 5. The code and data used in this paper are publicly available: https://github.com/UncertaintyQuantification/scalable_marginalization.

We briefly introduce the GP model in this section. We focus on computer model emulation, where the GP emulator is often used as a surrogate model to approximate computer experiments [52]. Consider a real-valued unknown function $z(\cdot)$, modeled by a Gaussian stochastic process (GaSP) or Gaussian process (GP), $z(\cdot)\sim\mathcal{GP}(\mu(\cdot),\sigma^{2}K(\cdot,\cdot))$, meaning that, for any inputs $\{\mathbf{x}_{1},\ldots,\mathbf{x}_{N}\}$ (with $\mathbf{x}_{i}$ being a $p\times 1$ vector), the marginal distribution of $\mathbf{z}=(z(\mathbf{x}_{1}),...,z(\mathbf{x}_{N}))^{T}$ follows a multivariate normal distribution,

where $\bm{\mu}=(\mu(\mathbf{x}_{1}),...,\mu(\mathbf{x}_{N}))^{T}$ is a vector of mean or trend parameters, $\sigma^{2}$ is the unknown variance and ${\mathbf{R}}$ is the correlation matrix with the $(i,j)$ element modeled by a kernel $K\,(\mathbf{x}_{i},\mathbf{x}_{j})$ with parameters $\bm{\gamma}$. It is common to model the mean by $\mu(\mathbf{x})=\mathbf{h}(\mathbf{x}){\bm{\beta}}\,$, where $\mathbf{h}(\mathbf{x})$ is a $1\times q$ row vector of basis function, and $\bm{\beta}$ is a $q\times 1$ vector of mean parameters.

When modeling spatially correlated data, the isotropic kernel is often used, where the input of the kernel only depends on the Euclidean distance $K(\mathbf{x}_{a},\mathbf{x}_{b})=K(||\mathbf{x}_{a}-\mathbf{x}_{b}||)$. In comparison, each coordinate of the latent function in computer experiments could have different physical meanings and units. Thus a product kernel is often used in constructing a GP emulator, such that correlation lengths can be different at each coordinate. For any $\mathbf{x}_{a}=(x_{a1},\ldots,x_{ap})$ and $\mathbf{x}_{b}=(x_{b1},\ldots,x_{bp})$, the kernel function can be written as $K(\mathbf{x}_{a},\mathbf{x}_{b})=K_{1}(x_{a1},x_{b1})\times...\times K_{p}(x_{ap},x_{bp})$, where $K_{l}$ is a kernel for the $l$th coordinate with a distinct range parameter $\gamma_{l}$, for $l=1,...,p$. Some frequently used kernels $K_{l}$ include power exponential and Matérn kernel functions [44]. The Matérn kernel, for instance, follows

where $d_{l}=|x_{al}-x_{bl}|$, $\Gamma(\cdot)$ is the gamma function, $\mathcal{K}_{\nu_{l}}(\cdot/\gamma_{l})$ is the modified Bessel function of the second kind with the range parameter and roughness parameter being $\gamma_{l}$ and $\nu_{l}$, respectively. The Matérn correlation has a closed-form expression when the roughness parameter is a half-integer, i.e. $\nu_{l}={2k_{l}+1}/{2}$ with $k_{l}\in\mathbb{N}$. It becomes the exponential correlation and Gaussian correlation, when $k_{l}=0$ and $k_{l}\to\infty$, respectively. The GP with Matérn kernel is $\lfloor\nu_{l}-1\rfloor$ mean square differentiable at coordinate $l$. This is a good property, as the differentiability of the process is directly controlled by the roughness parameter.

Denote mean basis of observations $\mathbf{H}=(\mathbf{h}^{T}(\mathbf{x}_{1}),...,\mathbf{h}^{T}(\mathbf{x}_{N}))^{T}$. The parameters in GP contain mean parameters $\bm{\beta}$, variance parameter $\sigma^{2}$, and range parameters $\bm{\gamma}=(\gamma_{1},...,\gamma_{p})$. Integrating out the mean and variance parameters with respect to reference prior $\pi(\bm{\beta},\sigma)\propto 1/\sigma^{2}$, the predictive distribution of any input $\mathbf{x}^{*}$ follows a student t distribution [20]:

with $N-q$ degrees of freedom, where

with ${\mathbf{h}}^{*}(\mathbf{x}^{*})=\left({{\mathbf{h}(\mathbf{x}^{*})}}-\mathbf{H}^{T}{{\mathbf{R}}}^{-1}\mathbf{r}(\mathbf{x}^{*})\right)$, $\hat{\bm{\beta}}=\left(\mathbf{H}^{T}{{\mathbf{R}}}^{-1}\ \mathbf{H}\right)^{-1}\mathbf{H}^{T}{{\mathbf{R}}}^{-1}\mathbf{z}$ being the generalized least squares estimator of $\bm{\beta}$, and $\mathbf{r}(\mathbf{x}^{*})=(K(\mathbf{x}^{*},{\mathbf{x}}_{1}),\ldots,K(\mathbf{x}^{*},{\mathbf{x}}_{N}))^{T}$.

Direct computation of the likelihood requires $\mathcal{O}(N^{3})$ operations due to computing the Cholesky decomposition of the covariane matrix for matrix inversion, and the determinant of the covariance matrix. Thus a posterior sampling algorithm, such as the Markov chain Monte Carlo (MCMC) algorithm can be slow, as it requires a large number of posterior samples. Plug-in estimators, such as the maximum likelihood estimator (MLE) were often used to estimate the range parameters $\bm{\gamma}$ in covariance. In [20], the maximum marginal posterior estimator (MMPE) with robust parameterizations was discussed to overcome the instability of the MLE. The MLE and MMPE of the parameters in a GP emulator with both the product kernel and the isotropic kernel are all implemented in the RobustGaSP package [18].

In some applications, we may not directly observe the latent function but a noisy realization:

where $z(.)$ is modeled as a zero-mean GP with covariance $\sigma^{2}K(.,.)$, and $\epsilon(\mathbf{x})\sim\mathcal{N}(0,\sigma^{2}_{0})$ follows an independent Gaussian noise. This model is typically referred to as the Gaussian process regression [44], which is suitable for scenarios containing noisy observations, such as experimental or field observations, numerical solutions of differential equations with non-negligible error, and stochastic simulations. Denote the noisy observations $\mathbf{y}=(y(\mathbf{x}_{1}),y(\mathbf{x}_{2}),...,y(\mathbf{x}_{N}))^{T}$ at the design input set $\{\mathbf{x}_{1},\mathbf{x}_{2},...,\mathbf{x}_{N}\}$ and the nugget parameter $\eta=\sigma^{2}_{0}/\sigma^{2}$. Both range and nugget parameters in GPR can be estimated by the plug-in estimators [18]. The predictive distribution of $f(\mathbf{x}^{*})$ at any input $\mathbf{x}^{*}$ can be obtained by replacing $\mathbf{R}$ with $\mathbf{\tilde{R}}=\mathbf{R}+\eta\mathbf{I}_{n}$ in Equation (4).

Constructing a GP emulator to approximate computer simulation typically starts with a “space-filling” design, such as the Latin hypercube sampling (LHS), to fill the input space. Numerical solutions of computer models were then obtained at these design points, and the set $\{(\mathbf{x}_{i},y_{i})\}^{N}_{i=1}$ is used for training a GP emulator. For any observed input $\mathbf{x}^{*}$, the predictive mean in (4) is often used for predictions, and the uncertainty of observations can be obtained through the predictive distribution. in Figure 1, we plot the predictive mean of a GP emulator to approximate the Branin function [56] with N training inputs sampled from LHS, using the default setting of the RobustGaSP package [18]. When the number of observations increases from $N=12$ (middle panel) to $N=24$ (right panel), the predictive error becomes smaller.

The computational complexity of GP models increases at the order of $\mathcal{O}(N^{3})$, which prohibits applications on emulating complex computer simulations, when a relatively large number of simulation runs are required. In the next section, we will introduce the state space representation of GP with Matérn covariance and one-dimensional input, where the computational order scales as $\mathcal{O}(N)$ without approximation. This method can be applied to problems with high dimensional input space discussed in Section 4.

Collective motions with particle interactions are very common in both microscopic and macroscopic systems [35, 36]. Learning interaction kernels between particles is important for two purposes. First, physical laws are less understood for many complex systems, such as cell migration or non-equilibrium thermodynamic processes. Estimating the interaction kernels between particles from fields or experimental data is essential for learning these processes. Second, simulation of particle interactions, such as ab initio molecular dynamics simulation, can be very computationally expensive. Statistical learning approaches can be used for reducing the computational cost of simulations.

For demonstration purposes, we consider a simple first-order system. In [34], for a system with $n$ interacting particles, the velocity of the $i$th particle at time $t$, $\mathbf{v}_{i}(t)=d\mathbf{x}_{i}(t)/dt$, is modeled by positions between all particles,

where $\phi(\cdot)$ is a latent interaction kernel function between particle $i$ and all other particles, with $||\cdot||$ being the Euclidean distance, and $\mathbf{u}_{i,j}(t)=\mathbf{x}_{j}(t)-\mathbf{x}_{i}(t)$ is a vector of differences between positions of particles $i$ and $j$, for $i,j=1,...,n$. Here $\phi(\cdot)$ is a two-body interaction function. In molecular dynamics simulation, for instance, the Lennard-Jones potential can be related to molecular forces and the acceleration of molecules in a second-order system. The statistical learning approach can be extended to a second-order system that involves acceleration and external force terms. The first-order system as (21) can be considered as an approximation of the second-order system. Furthermore, the interaction between particles is global, as any particle is affected by all other particles. Learning global interactions is more computationally challenging than local interactions [51], and approximating the global interaction by the local interaction is of interest in future studies.

One important goal is to efficiently estimate the unobservable interaction functions from the particle trajectory data, without specifying a parametric form. This goal is key for estimating the behaviors of dynamic systems in experiments and in observational studies, as the physics law in a new system may be unknown. In [13], $\phi(\cdot)$ is modeled by a zero-mean Gaussian process with a Matérn covariance:

Computing estimation of interactions of large-scale systems or more simulation runs, however, is prohibitive, as the direct inversion of the covariance matrix of observations of velocities requires $\mathcal{O}((nMDL)^{3})$ operations, where $M$ is the number of simulations or experiments, $n$ is the number of particles, $D$ is the dimension of each particle, $L$ denotes the number of time points for each trial. Furthermore, constructing such covariance contains computing an $n^{2}LM\times n^{2}LM$ matrix of $\phi$ for a $D$-dimensional input space, which takes $\mathcal{O}(n^{4}L^{2}M^{2}D)$ operations. Thus, directly estimating interaction kernel with a GP model in (22) is only applicable to systems with a small number of observations [34, 13].

This work makes contributions from two different aspects for estimating dynamic systems of interacting particles. We first show the covariance of velocity observations can be obtained by operations on a few sparse matrices, after marginalizing out the latent interaction function. The sparsity of the inverse covariance of the latent interaction kernel allows us to modify the Kalman filter to efficiently compute the matrix product in this problem, and then apply a conjugate gradient (CG) algorithm [24, 21, 49] to solve this system iteratively. The computational complexity of computing the predictive mean and variance of a test point is at the order of $\mathcal{O}(TnN)+\mathcal{O}(nN\log(nN))$, for a system of $n$ particles, $N=nMDL$ observations, and $T$ is the number of iterations required in the CG algorithm. We found that typically around a few hundred CG iterations can achieve high predictive accuracy for a moderately large number of observations. The algorithm leads substantial reduction of computational cost, compared to direct computation.

Second, we study the effect of experimental designs on estimating the interaction kernel function. In previous studies, it is unclear how initial positions, time steps of trajectory and the number of particles affect the accuracy in estimating interaction kernels. Compared to other conventional problems in computer model emulation, where a “space-filling” design is often used, here we cannot directly observe the realizations of the latent function. Instead, the output velocity is a weighted average of the interaction kernel functions between particles. Besides, we cannot directly control distances between the particles moving away from the initial positions, in both simulation and experimental studies. When the distance between two particles $i$ and $j$ is small, the contribution $\phi(||\mathbf{x}_{i}(t)-\mathbf{x}_{j}(t)||)\mathbf{u}_{i,j}(t)$ can be small, if the repulsive force by $\phi(\cdot)$ does not increase as fast as the distance decreases. Thus we found that the estimation of interaction kernel function can be less accurate in the input domain of small distances. This problem can be alleviated by placing initial positions of more particles close to each other, providing more data with small distance pairs that improve the accuracy in estimation.