SlowGuess/NeurIPS_2023_OCR huggingface dataset:
It remains to obtain a recursion for the smoothing densities, which can be shown to be computationally tractable given that the filtering recursion is tractable. The smoothing densities are denoted by
$$
p \left(\mathrm {x} _ {l} \mid \mathrm {y} _ {1: N}\right) = \mathcal {N} \left(\mathrm {x} _ {l}; \xi_ {l}, \Lambda_ {l}\right), \quad \xi_ {N} = \mu_ {N}, \Lambda_ {N} = \Sigma_ {N}. \tag {24}
$$
It can be shown that the posterior process has a backward Markov representation [5], hence the smoothing marginals may be obtained by
$$
p \left(\mathrm {x} _ {l} \mid \mathrm {y} _ {1: N}\right) = \int b _ {l, l + 1} \left(\mathrm {x} _ {l} \mid \mathrm {x} _ {l + 1}\right) p \left(\mathrm {x} _ {l + 1} \mid \mathrm {y} _ {1: N}\right) \mathrm {d x} _ {l + 1}, \tag {25}
$$
where we call  $b_{l,l+1}$  the backwards kernel. Consequently, the problem consists of approximating  $b_{l,l+1}$  such that the marginalization Eq. (25) may be implemented in a computationally frugal manner.
Approximating the backwards kernel For the linear Gaussian case, the backward kernel [46]
$$
b _ {l, l + 1} \left(\mathrm {x} _ {l} \mid \mathrm {x} _ {l + 1}\right) = \mathcal {N} \left(\mathrm {x} _ {l}; \mathrm {G} _ {l} \mathrm {x} _ {l + 1} + \mathrm {v} _ {l}, \mathrm {P} _ {l}\right), \tag {26}
$$
is parametrized by the smoothing gain  $\mathrm{G}_l \in \mathbb{R}^{n \times n}$ , the shift vector  $\mathrm{v}_l \in \mathbb{R}^n$ , and the covariance of the backwards kernel  $\mathrm{P}_l \in \mathbb{R}^{n \times n}$ . Let  $(\cdot)^+$  denote the Moore-Penrose pseudoinverse. Then
$$
\mathrm {G} _ {l} = \Sigma_ {l} \Phi_ {l + 1} ^ {*} \Pi_ {l + 1} ^ {+}, \quad \mathrm {v} _ {l} = \mu_ {l} - \mathrm {G} _ {l} \mu_ {l + 1} ^ {-}, \quad \mathrm {P} _ {l} = (\mathrm {I} - \mathrm {G} _ {l} \Phi_ {l + 1}) \Sigma_ {l} (\mathrm {I} - \mathrm {G} _ {l} \Phi_ {l + 1}) ^ {*} + \mathrm {G} _ {l} \mathrm {Q} _ {l + 1} \mathrm {G} _ {l} ^ {*}. \tag {27}
$$
In the low-rank setting, the backwards kernel is efficiently approximated given the results from above. The approximate smoothing gain is a product of a tall, a small quadratic, and a wide matrix:
$$
\mathrm {G} _ {l} \approx \Sigma_ {l} ^ {1 / 2} \underbrace {\Sigma_ {l} ^ {* / 2} \Phi_ {l + 1} ^ {*} \left(\Pi_ {l + 1} ^ {1 / 2}\right) ^ {+}} _ {=: \Gamma_ {l} \in \mathbb {R} ^ {r \times r}} \left(\Pi_ {l + 1} ^ {* / 2}\right) ^ {+}. \tag {28}
$$
During filtering,  $\Gamma_{l}$  is saved alongside the low-rank factors of the prediction and filtering covariances. We proceed to compute a low-rank representation of the backwards-transition covariance. Consider the following singular value decomposition, truncated at the  $r$ -th largest singular value,
$$
\left(\left(\mathrm {I} - \mathrm {G} _ {l} \Phi_ {l + 1}\right) \Sigma_ {l} ^ {1 / 2} \quad \mathrm {G} _ {l} \mathrm {Q} _ {l + 1} ^ {1 / 2}\right) \approx \widehat {\mathrm {U}} _ {l} \widehat {\mathrm {D}} _ {l} \widehat {\mathrm {V}} _ {l} ^ {*}, \tag {29}
$$
A. vs. KF mean
low-rank dimension
Figure 2: Performance of the low-rank filters on two different settings. In the truly low-rank linear advection problem (a), the RRKF achieves the optimal estimate  $(\diamond)$  when  $r$  exceeds the true rank of the problem. For the spatio-temporal air-quality regression (b), our method is consistently better.
B. vs. KF cov.
A. vs. KF mean
low-rank dimension
low-rank dimension
where  $\widehat{\mathrm{U}}_l\in \mathbb{R}^{n\times r}$ ,  $\widehat{\mathrm{D}}_l\in \mathbb{R}^{r\times r}$  diagonal,  $\widehat{\mathrm{V}}_l\in \mathbb{R}^{r\times r}$ , and  $\mathrm{G}_l$  as given in Eq. (28). Then a rank- $r$  factor of  $\mathrm{P}_l$  is given as
$$
\mathrm {P} _ {l} ^ {1 / 2} = \widehat {\mathrm {U}} _ {l} \widehat {\mathrm {D}} _ {l} \in \mathbb {R} ^ {n \times r}. \tag {30}
$$
As for filtering, in the worst case, the cost for low-rank smoothing scales quadratically with the state dimension  $n$ . The following proposition, which is proved in Appendix A.3, states that this can be further reduced to linear complexity in  $n$ , given the assumptions on the dynamics model are satisfied.
Proposition 3. Given Assumption 1, computing the smoothing density  $p(\mathbf{x}_l \mid \mathbf{y}_{1:N})$  costs  $\mathcal{O}(nr^2 + r^3)$ .
Finally, it is useful to mention that realizations of the backwards process are posterior samples.
This section evaluates the proposed method in different experimental settings. The chosen measure of quality is the distance of the approximate low-rank moments to the exact KF. We measure the mean deviations with the root-mean-squared error (RMSE) and the covariance deviations with the time-averaged relative Frobenius distances. The presented rank-reduced Kalman filter (RRKF) is compared to the ensemble Kalman filter (EnKF) and the ensemble transform Kalman filter (ETKF). All EnKF and ETKF results are given in sample statistics over 20 runs. Section 4.1 shows that for a truly low-rank system, our method recovers the KF estimate up to numerical error as soon as  $r$  exceeds the true rank of the problem. Section 4.2 tests the method on a spatio-temporal GP regression problem with real data. After Section 4.3 evaluates the approximation quality for increasingly low-rank systems in a controlled experimental environment, Section 4.4 verifies the stated asymptotic cost of the method. Finally, a large-scale spatio-temporal regression problem is solved in Section 4.5. In all experiments, we compute the stationary mean and a low-rank factorization of the stationary covariance matrix of the prior and condition the stationary moments on the first measurement in the respective time-series dataset. More details on the experimental setups are given in Appendix D.
4.1 Linear advection model
The proposed algorithm is evaluated in a standard data-assimilation setup in which linear-advection dynamics with periodic boundary conditions are assimilated to a set of simulated data. The setup follows the description in Sakov and Oke [33] and the experimental details are additionally detailed in Appendix D. This problem has rank 51 by construction. Figure 2a shows how the exact KF estimate is recovered by our method for  $r = 51$ , while the ensemble methods converge according to a Monte-Carlo rate. This is particularly clear for the covariance estimate (Fig. 2a, B.).
4.2 London air-quality regression
This experiment uses hourly data from the London air-quality network [23] between January 2019 and April 2019, which amounts to measurements at 72 spatial locations and 2159 points in time. The
Figure 3: How does the rate of spectral decay influence the low-rank approximation? The larger the spatial length scale of a spatio-temporal Matérn model, the faster the spectrum decays and the faster all methods converge to the KF estimate. The first and second row compare the mean estimates and the covariance estimates, respectively. The RRKF estimate is consistently closer and recovers the KF estimate (◇) at  $r = n$ . The cumulative spectrum of the final-step KF covariance is shown in grey.
Figure 4: Best-case and worst-case asymptotic complexities of the low-rank filters. When all Assumptions 1 to 3 are satisfied (A.), the elapsed time grows linearly with the state dimension  $n$ . If parts of the assumptions are not met (B.), the cost scales quadratically with the respective dimension.
data used, together with its processing (except the train-test split), is the same as in the corresponding experiment by Hamelijnck et al. [20]. The model is a spatio-temporal Matérn-3/2 process with prior hyperparameters that maximize the marginal log likelihood. Figure 2b shows that—compared to the ensemble methods—our algorithm is consistently closer to the KF estimate and recovers it at  $r = n$ .
4.3 Spatio-temporal Matérn process with varying spatial lengthscale
Consider a spatio-temporal Gaussian process  $\mathbf{x}(t) \sim \mathcal{GP}(0, k_t \otimes k_x)$  with covariance structure that is separable in time and space. Let both  $k_t$  and  $k_x$  be of the Matérn family with characteristic spatial and temporal lengthscales  $\ell_t, \ell_x$  and output scales  $\sigma_t^2, \sigma_x^2$ . Spatio-temporal Matérn processes can be translated to the formulation in Eq. (1a), as detailed, for instance, by Solin [39]. This setting allows for varying the "low-rankness" by the choice of  $\ell_x$ : the larger  $\ell_x$ , the more information about a spatial point is prescribed by its surroundings. As  $\ell_x$  approaches zero, the spatial kernel matrix gets more and more diagonal, encoding spatially independent states, and causing a barely decaying spectrum. The spectrum for large  $\ell_x$  decays rapidly, making low-rank approximations more accurate for small  $r$ .
In the experiment, we consider a spatial domain  $[0,2]\times [0,2]\subset \mathbb{R}^2$ , which is subsampled at a uniformly-spaced grid with  $\Delta_x = 0.1$  ( $n = 21\times 21 = 441$ ). Noisy observations of the full state trajectory, are drawn from a realization of the prior process. Figure 3 demonstrates that, for increasing spatial lengthscales, all methods converge faster to the true KF estimate. Crucially, however, our method is consistently closer to the optimal estimate and recovers it for  $r = n$ .
Figure 5: Rainfall in Australia. A spatio-temporal GP regression problem with  $n = 30911$  state dimensions, solved with the RRKF with  $r = 1000$ . The columns are three different days in the time series. The rows show the data, mean estimate, and Z-scores, respectively.
4.4 Runtime
To demonstrate the asymptotic computational complexities given in Section 3.3, we investigate the runtimes on two different settings: The first fulfills both Assumptions 1 and 2, whereas the second only fulfills Assumption 2, which means that the cost of prediction is quadratic in the state dimension. Both problems are solved on a temporal grid of size  $N = 100$  with the measurement dimension fixed at  $m = 100$  and the low-rank dimension fixed at  $r = 5$ .
The problem from Section 4.1 serves as the best-case setting. The noise-free linear-advection dynamics amount to multiplication of the state with a circulant matrix, which is implemented efficiently using fast Fourier transform. The measurement operator amounts to array indexing. As subplot A, in Fig. 4 confirms, the computational cost grows linearly with the state dimension. As a worst-case example we employ a spatio-temporal Matérn process with a dense spatial kernel matrix (cf. Section 4.3). In this case, we expect the computation time to scale quadratically with the state dimension because  $\mathrm{BB}^*$  is dense. This is verified by the right plot in Fig. 4 (B.).
4.5 Large-scale spatio-temporal GP regression on rainfall data
As a final experiment, a spatio-temporal GP regression problem is solved on a large gridded analysis data set measuring rainfall in Australia. The data is provided by the Australian Water Availability Project (AWAP) [2, 24]. The number of spatial data points is  $n = 30911$ , which rules out the cubic-in- $n$  Kalman filter. The smoothing posterior is computed at  $N = 40$  time points, from 25 January, 2023 through 5 March, 2023. The total number of observations is thus  $n \cdot N = 1236440$ . The spatio-temporal prior Matérn model is selected by maximizing the marginal log-likelihood of a lower-resolution data set. A set of 6160 spatial points is excluded from the data during filtering and smoothing in order to evaluate spatial interpolation performance. Thus, the measurement dimension is  $m = n - 6160 = 24751$ . Figure 5 shows the low-rank approximation to the GP posterior with  $r = 1000$  low-rank dimensions. The non-observed spatial locations are framed in the horizontal-"S" shape. At measurement locations, the data is described well by the model, while a slight smoothing-out effect can be observed due to the high-frequency features of the model being truncated. At the evaluation points the interpolation aligns with the expectations regarding the simple Matérn model.
The Z-score map shows  $\frac{\xi_l - y_l}{\lambda_l}$ , where  $\xi_{l}$  and  $\lambda_{l}$  denote the smoothing mean and marginal standard deviation and  $y_{l}$  the data point, at time  $t_l$ . This highlights badly calibrated uncertainty estimates, wherever the model under-/overestimates the data and is divided by a small standard deviation, in blue/red. The Z-score distribution receives further attention in Appendix E. All in all, the RRKF achieves high-quality approximate estimates while compressing the state by a factor of  $\frac{30911}{1000} \approx 30$ .
The main limitation of the proposed approximate probabilistic inference scheme is that the truncation to  $r$  dimensions cuts away covariance information, which is not accounted for. As a result, confidence grows as  $r$  shrinks, while estimates should arguably get more uncertain instead with increased compression. In existing low-rank filters, this issue is often counteracted manually by inflating the covariance matrices, on a per-application basis [6, Section 4.4] and calls for a more principled treatment. Finding a way to keep track of that residual uncertainty information is beyond the scope of this paper. It remains unclear how to preserve the stated asymptotic complexities while doing so.
Further, all results herein are stated for linear dynamics and observation models. Extensions to non-linear models could include linearization of the according transitions, or cubature methods [34].
By building upon well-established knowledge about optimal compression, we have proposed an algorithm providing a principled way to balance approximation accuracy against computational cost in high-dimensional state estimation. It combines simple (truncated) singular value decompositions with recent numerical low-rank integrators of large matrix differential equations. We have offered both theoretical and empirical arguments for why it is desirable to use a deterministic algorithm when approximating large-scale probabilistic state-estimation problems.
The authors gratefully acknowledge financial support by the German Federal Ministry of Education and Research (BMBF) through Project ADIMEM (FKZ 01IS18052B), and financial support by the European Research Council through ERC StG Action 757275 / PANAMA; the DFG Cluster of Excellence "Machine Learning - New Perspectives for Science", EXC 2064/1, project number 390727645; the German Federal Ministry of Education and Research (BMBF) through the Tübingen AI Center (FKZ: 01IS18039A); and funds from the Ministry of Science, Research and Arts of the State of Baden-Württemberg. Jörg Nick is funded by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) - Project-ID 258734477 - SFB 1173. Filip Tronarp was partially supported by the Wallenberg AI, Autonomous Systems and Software Program (WASP) funded by the Knut and Alice Wallenberg Foundation. The authors thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting Jonathan Schmidt. The authors also thank Nathanael Bosch and Christian Lubich for many valuable discussions and for helpful feedback.
[1] Adam, V., Eleftheriadis, S., Artemev, A., Durrande, N., and Hensman, J. (2020). Doubly sparse variational Gaussian processes. In Chiappa, S. and Calandra, R., editors, Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pages 2874-2884. PMLR.
[2] Australian Data Archive for Meteorology (ADAM) (2023). Australian Water Availability Project. http://www.bom.gov.au/jsp/awap/index.jsp. [Online; accessed May 2023].
[3] Axelsson, P. and Gustafsson, F. (2015). Discrete-time solutions to the continuous-time differential Lyapunov equation with applications to Kalman filtering. IEEE Transactions on Automatic Control, 60(3):632-643.
[4] Bishop, C. H., Etherton, B. J., and Majumdar, S. J. (2001). Adaptive sampling with the ensemble transform Kalman filter. part i: Theoretical aspects. Monthly weather review, 129(3):420-436.
[46] Tronarp, F., Bosch, N., and Hennig, P. (2022). Fenrir: Physics-enhanced regression for initial value problems. In Chaudhuri, K., Jegelka, S., Song, L., Szepesvari, C., Niu, G., and Sabato, S., editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 21776-21794. PMLR.
[47] Vijayshankar, S., Nabi, S., Chakrabarty, A., Grover, P., and Benosman, M. (2020). Dynamic mode decomposition and robust estimation: Case study of a 2d turbulent boussinesq flow. In 2020 American Control Conference (ACC), pages 2351-2356. IEEE.
[48] Williams, C. and Seeger, M. (2001). Using the Nyström method to speed up kernel machines. In Advances in Neural Information Processing Systems (NeurIPS).