The inverse Kalman filter

Dynamic linear models (DLMs) or linear state space models are ubiquitously used in modeling temporally correlated data [37, 5]. Each observation $y_{t}\in\mathbb{R}$ in DLMs is associated with a $q$-dimensional latent state vector $\bm{\theta}_{t}$, defined as:

where $\mathbf{F}_{t}$ and $\mathbf{G}_{t}$ are matrices of dimensions $1\times q$ and $q\times q$, respectively, $\mathbf{W}_{t}$ is a $q\times q$ covariance matrix for $t=2,\dots,N$, and the initial state vector $\bm{\theta}_{1}\sim\mathcal{MN}(\mathbf{b}_{1},\mathbf{W}_{1})$ with $\mathbf{b}_{1}=\mathbf{0}$ assumed herein. The dependent structure of DLMs is illustrated in Fig. 1(a), where each observation $y_{t}$ is associated with a single latent state $\bm{\theta}_{t}$.

DLMs are a flexible class of models that include many widely used processes, such as autoregressive and moving average processes [26, 27]. Some Gaussian processes (GPs) with commonly used kernel functions, including the Matérn covariance with a half-integer roughness parameter [12], can also be represented as DLMs [38, 13], with closed-form expressions for $\mathbf{F}_{t}$, $\mathbf{G}_{t}$ and $\mathbf{W}_{t}$. This connection, summarized in the Supplementary Material, enables differentiable GPs within the DLM framework, allowing efficient estimation of smooth functions from noisy observations.

The Kalman filter (KF) provides a scalable approach for estimating latent state and computing likelihood in DLMs, which scales linearly with the number of observations [16], as reviewed in Appendix A. In particular, the KF enables efficient computation of $\mathbf{L}^{-1}\mathbf{u}$ for any $N$-dimensional vector $\mathbf{u}$ in $\mathcal{O}(q^{3}N)$ operations, where $\mathbf{L}$ is the Cholesky factor of the covariance matrix $\bm{\Sigma}=\mbox{cov}[\mathbf{y}_{1:N}]=\mathbf{L}\mathbf{L}^{T}$, summarized in Lemma 6 in the Appendix.

Computational challenges remain for high-dimensional state space models and scenarios where KF cannot be directly applied, such as the dependent structure shown in Fig. 1(b), where each observation $y_{i}$ can be associated with multiple latent states. This interaction structure, introduced in Section 3, is common in physical models, such as molecular dynamics simulation [28]. One way to overcome these challenges is to efficiently compute $\bm{\Sigma}\mathbf{u}$ for any $N$-dimensional vector $\mathbf{u}$ and utilize optimization methods such as the conjugate gradient (CG) algorithm for computing the predictive distribution [15]. This strategy has been used to approximate the maximum likelihood estimator of parameters in GPs [32, 25]. Yet, each CG iteration requires matrix-vector multiplication, which involves $\mathcal{O}(N^{2})$ operations and storage, making it prohibitive for large $N$.

To address this issue, we introduce the inverse Kalman filter (IKF), which computes $\mathbf{L}\mathbf{u}$ and $\mathbf{L}^{T}\mathbf{u}$ with $\mathcal{O}(q^{3}N)$ operations without approximations, where $\mathbf{L}$ is the Cholesky factor of any $N\times N$ DLM-induced covariance $\bm{\Sigma}$ and $\mathbf{u}$ is an $N$-dimensional vector. The complexity is significantly smaller than $\mathcal{O}(N^{2})$ in direct computation. The latent states dimension $q$ is not larger than 3 in our applications, and it includes many commonly used models such as twice differentiable GPs with a Matérn covariance in (15), often used as a default choice in GP surrogate models.

The IKF algorithm can be extended to accelerate matrix multiplication and inversion, a key computational bottleneck in many problems. We integrate the IKF into a CG algorithm to scalably compute the predictive distribution of observations with a general covariance structure:

where $\bm{\Sigma}_{j}$ is a DLM-induced covariance matrix, $\mathbf{A}_{j}$ is sparse, and $\bm{\Lambda}$ is a diagonal matrix. This structure appears in nonparametric estimation of particle interactions, which will be introduced in Section 3. Both real and simulated studies involve numerous particles over a long time period, where conventional approximation methods [34, 19, 3, 7, 10] are not applicable. This covariance structure also appears in other applications, including varying coefficient models [14] and estimating incomplete lattices of correlated data [31]. We apply our approach to the latter in Section S5 and provide numerical comparisons with various approximation methods in Section S9 of the Supplementary Material.

In this section, we introduce an exact algorithm for computing $\bm{\Sigma}\mathbf{u}$ with $\mathcal{O}(q^{3}N)$, where $\mathbf{u}$ is an $N$-dimensional real-valued vector and $\bm{\Sigma}=\mbox{cov}[\mathbf{y}_{1:N}]$ is an $N\times N$ covariance matrix of observations induced by a DLM with $q$ latent states in Eqs. (1)-(2). This algorithm is applicable to any DLM specified in Eqs. (1)-(2). Denote $\bm{\Sigma}=\mathbf{L}\mathbf{L}^{T}$, where the Cholesky factor $\mathbf{L}$ does not need to be explicitly computed in our approach. We provide fast algorithms to compute $\mathbf{\tilde{x}}=\mathbf{L}^{T}\mathbf{u}$ and $\mathbf{x}=\mathbf{L}\mathbf{\tilde{x}}$ by Lemma 1 and Lemma 2, respectively, each with $\mathcal{O}(q^{3}N)$ operations. Detailed proofs of these lemmas are available in the Supplementary Material. In the following, $u_{t}$, $x_{t}$, $\tilde{x}_{t}$ denote the $t$th element of the vector of $\mathbf{u}$, $\mathbf{x}$, and $\mathbf{\tilde{x}}$, respectively, and $L_{t^{\prime},t}$ denotes the $(t^{\prime},t)$th element of $\mathbf{L}$ for $t$ and $t^{\prime}=1,2,\dots,N$.

The algorithm in Lemma 1 is unconventional and nontrivial to derive. In contrast, Lemma 2 has direct connection to the Kalman filter. Specifically, the one-step-ahead predictive distribution in step (ii) of the Kalman filter, given in Appendix A, enables computing $\mathbf{L}^{-1}\mathbf{x}$ for any vector $\mathbf{x}$ (Lemma 6). Lemma 2 essentially reverses this operation, computing $\mathbf{L}\mathbf{\tilde{x}}$ for any vector $\mathbf{\tilde{x}}$, leading to the term Inverse Kalman filter (IKF).

The IKF algorithm, outlined in Algorithm 1, sequentially applies Lemmas 1 and 2 to reduce the computational cost and storage for $\bm{\Sigma}{\mathbf{u}}$ from $\mathcal{O}(N^{2})$ to $\mathcal{O}(q^{3}N)$ without approximations. This approach applies to all DLMs, including GPs that admit equivalent DLM representations, such as GPs with Matérn kernels where roughness parameters $0.5$ and $2.5$ correspond to $q=1$ and $q=3$, respectively. Details are provided in the Supplementary Material.

While Algorithm 1 can be applied to DLM regardless of the noise variance, we do not recommend directly computing $\bm{\Sigma}\mathbf{u}$ in scenarios when the noise variance $V_{t}$ is zero or close to zero in Eq. (1). This is because when $V_{t}=0$, $Q_{t}$ in Eq. (27), the variance of the one-step-ahead predictive distribution, could be extremely small, which leads to large numerical errors in Eq. (28) due to unstable inversion of $Q_{t}$. To ensure robustness, we suggest computing $\mathbf{x}=(\bm{\Sigma}+V\mathbf{I}_{N}){\mathbf{u}}$ with a positive scalar $V$ and outputting $\mathbf{x}-V{\mathbf{u}}$ as the results of $\bm{\Sigma}{\mathbf{u}}$. Here, $\bm{\Sigma}+V\mathbf{I}_{N}$ can be interpreted as the covariance matrix of a DLM with noise variance $V_{t}=V$. Essentially, we introduce artificial noise to ensure that $Q_{t}$ computed in KF is at least $V$ for numerical stability. The result remains exact and independent of the choice of $V$.

We compare the IKF approach with the direct matrix-vector multiplication of $\bm{\Sigma}\mathbf{u}$ in noise-free scenarios in Fig. 2. The experiments utilize $z_{t}=z(d_{t})$ with $d$ uniformly sampled from $[0,1]$ and the Matérn covariance with unit variance, roughness parameter $2.5$, and range parameter $\gamma=0.1$ [12]. Panels (a) and (b) show that IKF significantly reduces computational costs compared to direct computation, with a linear relationship between computation time and the number of observations. IKF achieves covariance matrix-vector multiplication for a $10^{6}$-dimensional output vector in about 2 seconds on a desktop, making it highly efficient for iterative algorithms like the CG algorithm [15] and the randomized log-determinant approximation [29]. Fig. 2(c) compares the maximum absolute error between robust IKF with $V=0.1$ and non-robust IKF with $V=0$. The non-robust IKF exhibits large numerical errors due to instability when $Q_{t}$ approaches zero, while robust IKF remains stable by ensuring $Q_{t}$ is no smaller than $V$, even with near-singular covariance matrices $\bm{\Sigma}$.

Lemmas 1 and 2 facilitate the computation of each element in $\mathbf{L}$ using DLM parameters and enable the linear computation of $(\mathbf{L}^{T})^{-1}\tilde{\mathbf{x}}$. These results are summarized in Lemmas 3 and 4, respectively, with proofs provided in the Supplementary Material.

The IKF algorithm is motivated by the problem of estimating interactions between particles, which is crucial for understanding complex behaviors of molecules and active matter, such as migrating cells and flocking birds. Minimal physical models, such as the Vicsek model [35] and their extensions [1, 9], provide a framework for explaining the collective motion of active matter.

Consider a physical model encompassing multiple types of latent interactions. Let $\mathbf{y}_{i}$ be a $D_{y}$-dimensional output vector representing the $i$th particle, which is influenced by $J$ distinct interaction types. For the $j$th type of interaction, the $i$th particle interacts with a subset $p_{i,j}$ of particles rather than all particles, typically those within a radius $r_{j}$. This relationship is expressed as a latent factor model:

where $\bm{\epsilon}_{i}\sim\mathcal{MN}(0,\sigma^{2}_{0}I_{D_{y}})$ denotes a Gaussian noise vector. The term $\mathbf{a}_{i,j,k}$ is a $D_{y}$-dimensional known factor loading that links the $i$th output to the $k$th neighbor in the $j$th interaction, with the unknown interaction function $z_{j}(\cdot)$ evaluated at a scalar-valued input $d_{i,j,k}$, such as the distance between particles $i$ and $k$, for $k=1,\dots,p_{i,j}$, $i=1,\dots,n$, and $j=1,\dots,J$. For a dataset of $n_{p}$ particles over $n_{\tau}$ time points, the total number of observations is $\tilde{N}=nD_{y}=n_{p}n_{\tau}D_{y}$.

An illustrative example of this framework is the Vicsek model, a seminal approach for studying collective motion [35] (see Fig. S1(a) in the Supplementary Material). In this model, the 2-dimensional velocity of the $i^{\prime}$th particle at time $\tau$, $\mathbf{v}_{i^{\prime}}(\tau)=(v_{i^{\prime},1}(\tau),v_{i^{\prime},2}(\tau))^{T}$, has a constant magnitude $v=||\mathbf{v}_{i^{\prime}}(\tau)||$, where $||\cdot||$ denotes the $L_{2}$ norm, and $v_{i^{\prime},1}(\tau)$ and $v_{i^{\prime},2}(\tau)$ are the components of the velocity along two orthogonal directions for $i^{\prime}=1,\dots,n_{p}$ and $\tau=1,\dots,n_{\tau}$. The velocity angle $\phi_{i^{\prime}}(\tau)=\tan^{-1}(v_{i^{\prime},2}(\tau)/v_{i^{\prime},1}(\tau))$ is updated as

where $\epsilon_{i^{\prime}}(\tau)$ is a zero-mean Gaussian noise with variance $\sigma^{2}_{0}$. The set of neighbors $ne_{i^{\prime}}(\tau-1)$ includes particles within a radius of $r$ from particle $i^{\prime}$ at time $\tau-1$, i.e., $ne_{i^{\prime}}(\tau-1)=\{k:||\mathbf{s}_{i^{\prime}}(\tau-1)-\mathbf{s}_{k}(\tau-1)||<r\}$, where $\mathbf{s}_{i^{\prime}}(\tau-1)$ and $\mathbf{s}_{k}(\tau-1)$ are 2-dimensional position vectors of particle $i^{\prime}$ and its neighbor $k$, respectively, and $p_{i^{\prime}}(\tau-1)$ denotes the total number of neighbors of the $i^{\prime}$th particle, including itself, at time $\tau-1$.

The Vicsek model in Eq. (14) is a special case of the general framework in Eq. (13) with one-dimensional output $y_{i}=\phi_{i^{\prime}}(\tau)$, i.e. $D_{y}=1$, with the index of the $i$th observation $i=(\tau-1)n_{p}+i^{\prime}$ being a function of time $\tau$ and particle $i^{\prime}$. The Viscek model contains a single interaction, i.e. $J=1$, with linear interaction function $z(d)=d$, where $d$ being the velocity angle of the neighboring particle, and the factor loading being $1/p_{i^{\prime}}(\tau-1)$.

Numerous other physical and mathematical models of self-propelled particle dynamics can also be formulated as in Eq. (13) with a parametric form of a particle interaction function [1]. Nonparametric estimation of particle interactions is preferred when the physical mechanism is unknown [17, 23], but the high computational cost limits its applicability to systems with large numbers of particles over long time periods.

In this work, we model each latent interaction function $z_{j}(\cdot)$ nonparametrically using a GP. By utilizing kernels with an equivalent DLM representation, we can leverage the proposed IKF algorithms to significantly expedite computation. This includes commonly used kernels such as Matérn kernel with a half-integer roughness parameter [12], often used as the default setting of the GP models for predicting nonlinear functions, as the smoothness of the process can be controlled by its roughness parameters [11].

For each interaction $j$, we form the input vector $\mathbf{d}^{(u)}_{j}=[d^{(u)}_{1,j},\dots,d^{(u)}_{N_{j},j}]^{T}$, where ‘u’ indicates that the input entries are unordered, $d^{(u)}_{t,j}=d_{i,j,k}$ with $t=\sum^{i-1}_{i^{\prime}=1}p_{i^{\prime},j}+k$ for any tuple $(i,j,k)$, $k=1,\dots,p_{i,j}$, $i=1,\dots,n$, $j=1,\dots,J$, and $N_{j}=\sum^{n}_{i=1}p_{i,j}$. The marginal distribution of the $j$th factor follows $\mathbf{z}_{j}^{(u)}=[z_{j}(d_{1,j}^{(u)}),\dots,z_{j}(d_{N_{j},j}^{(u)})]^{T}\sim\mathcal{MN}(\mathbf{0},\,\bm{\Sigma}_{j}^{(u)})$, where the $(t,t^{\prime})$th entry of the $N_{j}\times N_{j}$ covariance matrix $\bm{\Sigma}_{j}^{(u)}$ is $\mbox{cov}[z_{j}(d_{t,j}^{(u)}),z_{j}(d_{t^{\prime},j}^{(u)})]=\sigma^{2}_{j}c_{j}(d_{t,j}^{(u)},d_{t^{\prime},j}^{(u)})=\sigma^{2}_{j}c_{j}(d)$, with $d=|d_{t,j}^{(u)}-d_{t^{\prime},j}^{(u)}|$, $\sigma^{2}_{j}$ and $c_{j}(\cdot)$ being the variance parameter and the correlation function, respectively. We employ the Matérn correlation with roughness parameters $\nu_{j}=0.5$ and $\nu_{j}=2.5$. For $\nu_{j}=0.5$, the Matérn correlation is the exponential correlation $c_{j}(d)=\exp(-d/\gamma_{j})$, while for $\nu_{j}=2.5$, the Matérn correlation follows:

where $\gamma_{j}$ denotes the range parameter.

By integrating out latent factor processes $\mathbf{z}_{j}$, the marginal distribution of the observation vector $\mathbf{y}=(\mathbf{y}^{T}_{1},\dots,\mathbf{y}^{T}_{n})^{T}$, with dimension $\tilde{N}=nD_{y}$, follows a multivariate normal distribution:

where $\mathbf{A}_{j}$ is a sparse $\tilde{N}\times N_{j}$ block diagonal matrix, such that the $i$th diagonal block is a $D_{y}\times p_{i,j}$ matrix $\mathbf{A}_{i,j}=[\mathbf{a}_{i,j,1},\dots,\mathbf{a}_{i,j,p_{i,j}}]$ for $i=1,\dots,n$. The posterior predictive distribution of the latent variable $z_{j}(d^{*})$ for any given input $d^{*}$ is also a normal distribution:

with the predictive mean and variance given by

where $\bm{\Sigma}_{y}=\sum^{J}_{j=1}\mathbf{A}_{j}\bm{\Sigma}^{(u)}_{j}\mathbf{A}^{T}_{j}+\sigma^{2}_{0}\mathbf{I}_{\tilde{N}}$, $\bm{\Sigma}^{(u)}_{j}(d^{*})=[c_{j}(d^{(u)}_{1,j},d^{*}),\dots,c_{j}(d^{(u)}_{N_{j},j},d^{*})]^{T}$, and $\mathbf{r}$ represents additional system parameters.

The main computational challenge in calculating the predictive distribution (17) lies in inverting the covariance matrix $\bm{\Sigma}_{y}$, where $\tilde{N}$ and $N_{j}$ range between $10^{5}$ and $10^{6}$ in applications. Though various GP approximation methods have been proposed [34, 30, 8, 2, 22, 10, 3, 18], they typically approximate the parametric covariance $\bm{\Sigma}_{j}^{(u)}$ rather than $\bm{\Sigma}_{y}$, thus not directly applicable for computing the distribution in Eq. (17).

To address this, we employ the CG algorithm [15] to accelerate the computation of the predictive distribution. Each CG iteration needs to compute $\bm{\Sigma}_{y}\mathbf{u}$ for an $\tilde{N}$-dimensional vector ${\mathbf{u}}$. To employ the IKF, we first need to rearrange the input vector $\mathbf{d}^{(u)}_{j}$ into a non-decreasing input sequence $\mathbf{d}_{j}=[d_{1,j},\dots,d_{N_{j},j}]^{T}$. Denote the covariance $\mbox{var}[\mathbf{z}_{j}]=\bm{\Sigma}_{j}$ with $\mathbf{z}_{j}=[z_{j}(d_{1,j}),\dots,z_{j}(d_{N_{j},j})]^{T}$. The covariance of the observations can be written as a weighted sum of $\bm{\Sigma}_{j}$ by introducing a permutation matrix for each interaction: $\bm{\Sigma}_{y}=\sum^{J}_{j=1}\mathbf{A}_{\pi,j}\bm{\Sigma}_{j}\mathbf{A}^{T}_{\pi,j}+\sigma^{2}_{0}\mathbf{I}_{N}$, where $\mathbf{A}_{\pi,j}=\mathbf{A}_{j}\mathbf{P}^{T}_{j}$ with $\mathbf{P}_{j}$ being a permutation matrix such that $\bm{\Sigma}_{j}=\mathbf{P}_{j}\bm{\Sigma}^{(u)}_{j}\mathbf{P}^{T}_{j}$. After this reordering, the computation of $\bm{\Sigma}_{y}\mathbf{u}$ can be broken into four steps: (1) $\mathbf{u}_{j}=\mathbf{A}_{\pi,j}^{T}\mathbf{u}$, (2) $\mathbf{x}_{j}=\bm{\Sigma}_{j}\mathbf{u}_{j}$, (3) $\mathbf{\hat{u}}_{j}=\mathbf{A}_{\pi,j}\mathbf{x}_{j}$, and (4) $\sum^{J}_{j=1}\mathbf{\hat{u}}_{j}+\sigma^{2}_{0}\mathbf{u}$. Here $\mathbf{A}_{j}$ is a sparse matrix with $N_{j}D_{y}$ non-zero entries. The IKF algorithm is used in step (2) to accelerate the most expensive computation, with all computations performed directly using terms in the KF algorithm without explicitly constructing the covariance matrix.

We refer to the resulting approach as the IKF-CG algorithm. This approach reduces the computational cost for computing the posterior distribution from $\mathcal{O}(\tilde{N}^{3})$ operations to pseudolinear order with respect to $N_{j}$, as shown in Table S1 in the Supplementary Material. Furthermore, the IKF-CG algorithm can accelerate the parameter estimation via both cross-validation and maximum likelihood estimation, with the latter requiring an additional approximation of the log-determinant [29]. Details of the CG algorithm, the computation of the predictive distribution, parameter estimation procedures, and computational complexity are discussed in Sections S3, S4, S6, and S7 of the Supplementary Material, respectively.

Several future directions are worth exploring. First, computing large matrix-vector products is ubiquitous, and the approach can be extended to different models, including some large neural network models. Second, it is an open topic to scalably compute the logarithm of the determinant of the covariance in Eq. (3) without approximation. Third, the new approach can be extended to estimate latent functions when forward equations are unavailable or computationally intensive in nonlinear or non-Gaussian dynamical systems, where ensemble Kalman filter [6, 33] and particle filter [20] are commonly used. Fourth, the proposed approach can be extended for estimating and predicting multivariate time series [21] and generalized factor processes for categorical observations [4].