Neighbour-Driven Gaussian Process Variational Autoencoders for Scalable Structured Latent Modelling
Authors: Xinxing Shi
Xiaoyu Jiang
Mauricio A Álvarez
Venue: ICML 2025 poster (2025)
Keywords: Gaussian process VAE; nearest neighbour
Primary area: probabilistic_methods->gaussian_processes
TL;DR: Propose a scalable inference method for Gaussian process VAEs by leveraging local adjacencies in latent space

Abstract: Gaussian Process (GP) Variational Autoencoders (VAEs) extend standard VAEs by replacing the fully factorised Gaussian prior with a GP prior, thereby capturing richer correlations among latent variables. However, performing exact GP inference in large-scale GPVAEs is computationally prohibitive, often forcing existing approaches to rely on restrictive kernel assumptions or large sets of inducing points. In this work, we propose a neighbour-driven approximation strategy that exploits local adjacencies in the latent space to achieve scalable GPVAE inference. By confining computations to the nearest neighbours of each data point, our method preserves essential latent dependencies, allowing more flexible kernel choices and mitigating the need for numerous inducing points. Through extensive experiments on tasks including representation learning, data imputation, and conditional generation, we demonstrate that our approach outperforms other GPVAE variants in both predictive performance and computational efficiency.

Neighbour-Driven Gaussian Process Variational Autoencoders for Scalable Structured Latent Modelling
Xinxing Shi
Xiaoyu Jiang Mauricio A. Álvarez! 1
Abstract
Gaussian Process (GP) Variational Autoencoders (VAEs) extend standard VAES by replacing the fully factorised Gaussian prior with a GP prior, thereby capturing richer correlations among la- tent variables. However, performing exact GP inference in large-scale GPVAEs is computation- ally prohibitive, often forcing existing approaches to rely on restrictive kernel assumptions or large sets of inducing points. In this work, we pro- pose neighbour-driven approximation strategy that exploits local adjacencies in the latent space to achieve scalable GPVAE inference. By con- fining computations to the nearest neighbours of each data point, our method preserves essential latent dependencies, allowing more flexible ker- nel choices and mitigating the need for numer- ous inducing points. Through extensive experi- ments on tasks including representation learning, data imputation, and conditional generation, we demonstrate that our approach outperforms other GPVAE variants both predictive performance and computational efficiency.
4. Y)
P Z)
3
neighboourhood m(1 = (f 1,1,+11
Figure 1: GPVAE places GP prior over the latent variables
Z to model correlations in the structured data yx). Our approach is to approximate the dense GP prior by leveraging the associated neighbourhood (Xn(t)) Mn(i))
sequential, spatial, or other structured data. In these cases, it is crucial to model correlations among latent variables.
A natural way to introduce such correlations is to adopt a Gaussian Process (GP) prior over the latent representations, resulting in Gaussian Process Variational Autoencoders (GP- VABs) (Casale al., 2018). By encoding the latent vari- ables realisations from GP, GPVAEs enforce structured dependencies through kernels (as shown in Fig. 1). Unfor- tunately. the direct application of GPs leads to an O(N3) computational overhead (N is the number of training sam- ples), making naive approaches prohibitively expensive in large-scale settings.
1. Introduction
Variational Autoencoders (VAEs) (Kingma & WWelling, 2014) have achieved remarkable success in variety of tasks ranging from representation learning (Higgins et al., 2017; Kim & Mnih, 2018) to generative modelling (Sønderby et al., 2016; Razavi et al., 2019). While a conventional VAE assumes a fully factorised Gaussian prior over the latent variables enable amortised inference, this assumption can be overly restrictive in many real-world scenarios involving
To address this limitation, one line of research explores specific kernll structures, such certain Matérn families (Zhu et al., 2023) and low-rank kemelss (Casale et al., 2018). While effective in addressing scalability challenges, these approaches are constrained by their narrow kernel choices. Other scalable GPVAE variants have been proposed, notably using inducing points (Ashman al., 2020; Jazbec et al., 2021), small set of pseudo-points that approximates the GP posterior. These methods. however, may require numerous inducing points in rapidly varying data encounter diffi- culties in optimising the inducing points (Wu et al., 2022). Moreover, fully Bayesian extensions using sampling- based methods offer good uncertainty calibration but often require longer runtime for sufficient samples (Tran et al., 2023).
'Department of Computer Science, University of Manch- ester, Manchester, UK. Correspondence to: Xinxing Shi <xinxing shi @ postgrad.manccesteer. www.mester..c...... Xiaoyu Jiang <xi- aoyu.jiang postgrad manchester.ac...k. Mauricio A. Álvarez www.ralooozaeeeeeeeoooo wwwpeesstaaacncnscccccuuu..
Proceedings of the ndnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
In this work, we propose an efficient neighbour-driven ap- proach to GPVAE training and prediction (as illustrated in Fig. 1). Our key insight is that in many structured datasets- datasets-such such as video frames with temporal adjacency or spatial regions with local patterns-focuscnn on small set of nearest neighbours captures most of the essential cor- relation structures. This intuition aligns with the "first law of geography," which underlies various Nearest Neighbour GP (NNGP) methods (Vecchia, 1988; Stein et al., 2004). Inspired by recent evidence that NNGPs can outperform sparse GPs with inducing points (Hensman et al., 2013) in large-scale geostatistics and machine learning tasks (Datta et al., 2016b; Tran et al., 2021b; Wu et al., 2022; Cao et al., 2023; Allison et al., 2024), we introduce two latent GP prior approximations to neighbour- latent modelling in GPVAEs: (1) Hierarchical Prior Approximation (IIP)): enforces sparsity in dense GP covariance matrices via hi- erarchical mechanism that selects local neighbourhoods in each mini-batch; (2) Sparse Precision Approximation (SPA): faatorsss the GP prior into conditional distributions involving only the nearest neighbours. Both methods sup- port structured latent modelling with flexible kernels while avoiding large sets of inducing points or rigid kernel assump- tions. By integrating these upproximations into mini-batch amortised inference, our approach leverages local correla- tion principles for scalable GPVAE inference, preserving essential latent structures with minimal overhead.
joint Gaussian distribution (Rasmussen & Williams, 2006). The covariance matrix of that joint distribution is determined by the kernel function k: RD IR + R. which encodes how sample points of the imput space x E IR correlate with each other.
Consider N pairs of data points
Y W E RNNK x= 1 E RxD
where y, E R is the n-th observation and x,, E RD corre- sponds to the auxiliary information (e.g., video timestamps or spatial coordinates). We define associated latent vari- ables Z = [ E RN with L latent channels. By the construction of GPVAE, each latent channel z E R follows '' ~ (O, A (X, X)), where kyg is the kernel with parameters vb. The generative process is given by
(z X)
NN' 0, (X,
(I)
N
Po Z) = II
po kn),
(2)
where the latent prior factorises over channels and (>(y,. zn) is modelled by a decoder network with parameters In the subsequent sections, we will use the notation Kxx k(X, X) to represent the covariance matrix evaluated at all pairs of inputs X.
Contributions This paper makes the following contri- butions: (1) We introduce two neighbour-driven GP prior approximations into GPVAE inference that capture essential latent structure while remaining scalable. (2) We present an amortised training procedure that jointly learns the encoder, decoder, and kernel parameters in mini-batch fashion, ac- commodating wide range of kernels without restrictive assumptions large number of inducing points. (3) Our framework delivers competitive performance multiple tasks for latent representation learning, data imputation, and conditional generation. Empirical experiments demonstrate that our approach improves both predictive accuracy and training speed compared to existing GPVAE baselines.
Variational inference for GPVAFs To approximate the intractable posterior p(Z Y,X), X). consider the following variational distribution as in a standard VAE (Kingma & Welling, 2014; Casale et al., 2018):
N
g (Z Y) II W(z. (o(y)), of (3)
where the mean H68 and variance o² are given by en- coder network with parameters Φ. The GPVAE parameters { 0, Φ} can be learned by maximising the following Evi- dence Lower BOund (ELBO):
2. Background
L = Eng ( [log Po Y Z)]
Section 2.1 revisits GPVAEs and the scalability bottleneck from the full GP prior. Section 2.2 summarises the widely used inducing-point formulation for scalable GPVAE infer- ence, while Scction 2.3 reviews NNGP approximations that motivate our approach.
(4)
KL 9g (Z Y) p P X)).
where the KL stands for Kullback-Leible divergence. If Kxx = IN. the GP prior collapses to fully factorised VAE prior and (4) recovers the standard VAE objective. In general, however, the KL term does not factorise across samples due to dense kxx- As a result, (4) cannot be com- puted in mini-batches and incurs ( time complexity, making large-scale inference computationally infeasible.
2.1. Gaussian Process Variational Autoencoder
A GP f(x) ~ (x)m(x), k(x, x')) is a stochastic process for which any finite subset of random variables follows a
2

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
2.2. GPVAEs Based on Inducing Points
3. Neighbour-Driven GPVAE
A line of research addresses the ((N3) bottleneck in GP- VAEs by adopting approximations based on inducing points (Titsias, 2009; Hensman et al., 2013; Jazbec et al., 2021). It uses an alternative variational distribution to eliminate the troublesome ptp [Z] X) in the KL term:
Our approach adopts a similar local-neighbour perspective to NNGPs but applies inside the latent space of GPVAEs. Rather than sparsifying GPs in the observation domain, we build scalable, neighbour- based GP priors within the la- tent dimensions. This adaptation brings the computational benefits of NNGPs to structured latent modelling for high- dimensional and correlated data.
Z Y, X = Pe (z'
X) IIII= Gol2) ya)
(5)
In this section, we introduce two neighbour-driven GP- VAE variants that can exhibit high scalability for large-scale datasets.
Z Y
where Øo is modelled by an encoder, and 9.0 is the nor- malising constant. Equation (5) is further approximated by following the inducing-poiis framework through
3.1. Model Setup
We adopt the standard construction of GPVAE mentioned earlier: a multi-output GP prior on the latent variables Z = [x]] E NRL as in (1) and a decoder network mapping each latent variable zn E IR to its observation yn as in (2). Separate kernels (Ko)= are used across I, latent dimensions to fully exploit GP's expressivity. Since the latent GPs are independent, we omit the superscript for la- tent dimensions We also suppress the subscript for model parameters to make the notations uncluttered.
/ p(Z Zu)
[Z Y,X) ~
where U MDD represents M(< N) inducing locations, Zu is the corresponding inducing variables, and pZ / Zu) is conditional Gaussian. In practice, the aforementioned variational distribution q( (Z Y. X) over inducing points takes a stochastic heuristic form determined by mini-batches. Additionally, the normalising constant Zy.s is subject to a lower bound associated with U. The resulting training objective allows learning the locations of inducing points. We abbreviate the model described above SVGPVAF, and more details are provided in Appendix D.I.
We choose employ the standard encoder from (3) keep the model design straightforward. Employing this conven- tional architecture facilitates direct comparisons with stan- dard VAES in our paper. As we will show in subsequent experiments, this encoder delivers competitive performance. If needed, practitioners can incorporate X by designing an encoder that accepts both X and Y inputs.
2.3. Nearest Neighbour Gaussian Process
NNGPs scale classical GPs by assuming that each data point depends only small set of nearby neighbours, typically determined by spatial or temporal distance. This locality naturally reduces computational costs: rather than dealing with a dense N x N covariance or precision matrix, NNGPs construct block-sparse or banded structures where each row or column involves only the nearest neighbours.
3.2. Neiihbour-Driven Inference
Our idea to impose sparsity on the full GP prior, allowing the KL term of (4) to decompose into manageable terms while preserving essential latent correlations. To achieve this, we develop inference methods using two distinct GP prior approximations inspired by recent NNGPs.
NNGPs have been widely used in geostatistics (Datta al., 2016;;;; Katzfuss & Guinness, 2021; ;; Datta, 2022). Recent works extend NNGPs to general regression tasks, employ- ing variational approaches from inducing points (Tran et al., 2021b; Wu et al., 2022). For example, Tran et al. (2021b) present an NNGP that implements a hierarchical prior with an auxiliary random indicator to determine the selection of inducing variables. In addition, Wu al. (2022) propose variational GP that applies sparse precision approximation for the inducing variable distribution. These approaches have shown advantages over the standard inducing point method, particularly on large datasets characterised by in- trinsically low lengthscales, where many inducing points are required to capture local correlations effectively.
Hierarchical Prior Approximation (HPA) To formulate the inference, we introduce an additional binary random vector w E {0. 1 ~ p(w) to indicate the inclusion N latent variables Z. We then establish a hierarchical structure for the prior based on the selection process
pZZ ()) N (Z 0, D, p(Z, w) (( w)p(w),
(6)
where Dw = diag(w) (w). From (6), the correlations of the unselected variables Z among Z are eliminated, where
I I= = {i w, 0). Then, we set up an amortised varia-
3

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
tional distribution that uses similar form (6) follows!
all points of X). Proceeding with the variational distribution
(( Z Y) defined by (3), obtain the ELBO based on SPA as follows
g( w) = (Z Dw Dwo DY) gZZ, w) - [Z w)p(w);
(Y)
(8)
(9)
LSPA
(ete. [log(yy 2,)]
where u1)) and (-) are the outputs of the encoder. By the construction of (6)-(9) the ELBO for our GPVAE is updated as
2=1
Σ (Z KL (g(2)) ll p(zz) Zn(i)]]
LIPA = Ep(w) (Eg(2 w) logp ( Y Z)
KT. (Z w) pZZ] w)]}
Therefore, the two terms of the ELBO above can be esti- mated by mini-batches from the data
where the first term is the expected log-likeling regliielihood. p(w) specifies how the variables are selected and can be any im- plicit distribution that we can draw samples from. In this work, sampling from p(w) is materialised by the following neighbour- strategy: for each training point yx) s) in a random mini-batch (xx,y:)) yi)e, we pick up the top-H nearest locations the auxiliary point x; in X, whose in- dices denoted as ri)i). For stochastic mini-batch training, we estimate the ELBO under the hierarchical prior by
N
CSPA
Rd
Eg(e) [logp(y |,:]]
I
iEI
N
Byy KL [q(2;) || p(zj Zn(ji)],
JEJ
(13)
where and I are sets of mini-batch indices. We call the proposed model GPVAE-SPA since (12) essentially offers approximation with sparse Cholesky factor for the prior precision matrix K xx (Datta, 2022). A detailed derivation of (13) is provided in Appendix B.2. By setting H = N, we retain all dependencies in the probability chain in (12), thus recovering the full-batch ELBO (4) again. In contrast, setting H = will cause the model to degenerate into conventional VAEs.
N
CHPA
Ex(ar)rs) (y)) [logp(y: 2;)]
iET
(10)
KL (g(n))) P(Z
where the KL divergence breaks down into several terms that involve nower-dimensionss covariance matrices (i.e., H x H). The original full-batch EBLO (4) can be recovered when all elements of w are set to one or equivalently, are set to include all N training points. We refer to our model trained through (10) as GPVAE-HPA. For detailed derivation, please see Appendix B.I.
Sparsity mechanisms and NNGP links HPA and SPA impose neighbour driven sparsity in complementary ways. HPA produces a sparse covariance structure by "switching off" interactions between non -eighbbbgguuiing latent variables through a hierarchical selection variable. In contrast, SPA factorises the GP distribution into chained conditional terms, effectively resulting in sparse precision matrix. Both schemes are valid neighbour-driven approaches suitable for scalable GPVAE inference.
Sparse Precision Approximation (SPA) The inference in this section employs a sparse precision matrix approxima- tion (Vecchia, 1988; Datta, 2022) to latent GPs. Consider factorising the joint distribution p(Z) of a GP by the proba- bility chain rule subject to some ordering
Conceptually, our approach closely aligns with the NNGPs by Tran et al. (2021b) and Wu et al. (2022), which enforce local sparsity on inducing variables. To draw a parallel, the proposed GPVAEs can be seen as applying such a principle to latent NNGPs with each "inducing variable" placed at each data point. We then incorporate decoder network to model the "likelihood" p(Z Y) and an encoder net- work for amortised inference, thereby forming our GPVAF architecture.
p(Z) p(z,) II plzz Zızj1),
(11)
where 1 stands for the collection {2n 1- The SPA is built by imposing conditional independence on (11) by only considering nearest neighbours, leading to
p(Z) p(z1) II p(z, ZnG)),
(12)
3.3. Predictive Posterior
Predictive posterior distributions p(y. x,; Y) are often used in tasks of conditional generation, where an unseen auxiliary x, E RD given. Here, we derive an approximate predictive posterior for the proposed GPVAEs based on neighhouring latent representations. Specifically, computing
where abuse the notation nij)) to represent the indices of Hl nearest neighbours of x, in {xx y (rather than among
Wee have omitted the conditioning variable Y.
4

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
VIDIRO
EVGPVAE
GPVAECuaak
GPVAE-PA
GPVAESPA
VAE
GEVAR-Pwete
B3SSE
GEVABCuaa
SYGPIAE
GEVAR-HDA (oum)
10
GPVAE-SAA (ous)
1028 2009
IPEL 23B unvead timb
Butee 11B ancrammttiee
Iina uxg
M
(b) Reconstruction RMSEs using different numbers of nearest neighhours/indncing points.
(a) Latent trajectory reconstruction of moving ball videos.
Figure 2: Latent representation learning for the moving ball dataset. (a) The leftmost column shows frames overlaid and shaded by time. The orange reconstructed paths are obtained using M/H = 10. (b) The results of the standard VAE and two full-batch baselines are shown with the shaded bands. The error bars and the shades indicate #1 standard deviation.
ply. x., Y) at a new location x, only needs to consider its H nearest neighbours in X (denoted by the index set n())):
Unlike Casale et al. (2018), Pearce (2020) and Fortuin et al. (2020) adopted alternative formulations the variational distribution and applied GPVAE to applications of inter- pretable latent dynamics and time-series imputation, respec- tively, but their models are limited to short sequences.
q(z. Y) p(z. Zn()))9(Zn(()) Y Yn(())
dZn())
ply+1 x.. Y) Re / p(y= 2z) Y) dzz,
As aforemennioned,, SVGPVAE (Jazbec et al., 2021) uses inducing points but may suffer from optimisation issues as the number of inducing points increases. LVAE (Ramchan- dran al., 2021) another inducing -oiitg-poootttttoo based model re- stricted to additive kernels for longitudinal data with discrete instance covariates. MGPVAE (Zhu et al., 2023) exploits state-space representations of one-dimensional Matérn ker- nels to enable Kalman-filte-lle inference. Since the in- ference involves both forward and backward processes, utilising parallel computation to speed up can be chal- lenging. SGPBAE (Tran et al., 2023) treats latent vari- ables, decoder parameters, and kernel parameters in a fully Bayesian fashion and leverages Stochastic Gradient Hamil- tonian Monte Carlo (SGHMC) (Chen et al., 2014). Com- pared to our method, sampling from that model is usually time-consuming. We summarise some recent scalable GP- VAE models in Appendix D.
where qzz Y) Gaussian and can be sampled efficiently in practice give Monte Carlo estimation of ply. Xx, Y).
3.4. Computational Complexity
The computational complexity of our models is twofold: (1) determining the nearest neighbour structure for auxiliary data; (2) Cholesky decomposition of H x H covariance matrices in the KI. term. Locating the H nearest neighbours of each point takes O(HN) in the worst case. Our imple- mentation leverages the Faiss package (Johnson et al., 2019) to efficiently accelcrate the nearest neighbour search on GPUs. After pre-computig the nearest neighbour structure, Cholesky decomposition in (10) and (13) has complexity HH*,,,HH) where N₂ is the mini-batch size and T. is the latent dimensions. Table 15 in Appendix D provides overview of the GP complexity associated with other rele- vant models.
5. Experiments
In this section, we evaluate our models various tasks involving both synthetic and real-world datasets. We begin by examining their ability to learn latent representations in a moving ball dataset. Next, we apply them to rotated MNIST to impute corrupted frames and generate any missing ones. We further conduct a long-sequence conditional generation experiment using MuJoCo action data, along with imputa- tion tasks on geostatistical datasets, one of which is substan- tially larger (N 10²) than those typically seen in previous GPVAEs (N ~ 10).
4. Related Work
The GPVAE was first proposed by Casale et al. (2018) to remove the i..dd assumption of the Gaussian prior in VAEs. The proposed GPVAE uses an encoder and decoder as in standard VAE, which we also adopt in our work. The work has to use a low-rank GP kernel and resort to first-order Taylor series expansion to achieve mini-batch inference, which is overly complicated and computationally inefficient.
5

---

Weighbour-Driver GPVAEs for Scalable Structured Latent Modelling
Tager
2 2
YAE
2 4
1
2
8
2
8
2
0 a 2
B
CA 20
2
2 2
4
2
3
3
3
SYGIVASS
>
>
>
GPVAR
GPVAR Digg
Sanda
2
2
2
2
2
2
7
2 2 2
2 2
2
2
2
2
(M- 90
IVAR
IVA
3
>
SEGEYAS
2 2
2
7
2
VAGEAAD
7-
>-
2-
2 22 222
2 2
2
2
2
Z
2
2
(M (owes)
GIVAESPA
2 22
N
2
CIVARSPA
(b) Generation on an unseen rotated MNIST sequence with missing frames. Figure 3: (a) Corrupted frame imputation with around 60% missing pixels. (b) Missing frames generation on Rotated MNIST with around 60% missing frames. The red boxes indicate the missing frames in that sequence.
(a) Imputation on series with missing pixels.
We benchmark our models against diverse GPVAEs, mainly focusing on scalable models such as SVGPVAE (Jazbec al., 2021), LVAE (Ramchandran et al., 2021), MGPVAE (Zhu et al., 2023), and SGPBAE (Tran al., 2023). Some other baselines, like VAE (Kingma & Welling, 2014), HI- VAE (Nazabal et al., 2020), and NNGP (Wu et al., 2022) are also considered. The performance is assessed using test negative log- likeliiol (NLL) and root mean squared error (RMSE). We report the wall-clock training time to estimate computational efficiency. We strive follow the original baseline settings ensure the fairest comparisons possible. All experimental results report the mean and standard devia- tion from 10 random trials. Further experimental setups can be found in Appendix C.
the latent variables cannot account for correlations between samples. The SVGPVAE begins to learn the trend, but the performance is unsatisfactory due to an insufficient num- ber of inducing points. Our two models are quite close to the full-batch baseline, i.e., GPVAE-Casale, with the same number of nearest neighbours.
Scaling behaviour We further show in Fig. 2b the perfor- mance of our models compared to two full-GP approaches and SVGPVAE across varying numbers of nearest neigh- bours and inducing points. While both GPVAE- Pearce and EPVAE-Casale use fully correlated GP priors, they differ in how they construct their variational distributions GGVAE- Pearce's s setup (5) is more closely aligned with SVGPVAE, while GPVAE-Casale adopts a Slagonal-encoded approach (3) that matches ours. From Fig. 2b, the SVGPVAE requires nearly half the total number of data points in the trajectory to attain performance levels comparable to the full-batch baseline, whereas our models achieve this with even fewer points. The GPVAE-SPA attains good reconstruction quality using only of the total trajectory points. Table 6 shows additional results using different values of H.
5.1. Moving Ball
Our experiments start with latent representation learning on the synthetic moving ball data from Pearce (2020). This dataset consists of back-anddwhiie videos of 30 frames cap- turing the motion of a pixel ball. Video samples are shown in the leftmost column of Fig. 2a. The two-dimensional trajectory of the ball in each video is simulated from GP with a radial basis function (RBF) kernel. We aim to in- fer the ball's trajectory based on the pixel-level frames. In particular, the trajectory is represented by the mean of the latent variable from the encoder.
5.2. Rotated MNIST
Rotated MNIST consists of sequences of handwritten digit images from the MNIST dataset (LeCun et al., 1998), in which each image frame is consecutively rotated at different angles. This classic dataset is designed to assess how well the model can learn latent dynamics and is commonly used in GPVAE literature. Here, we use two versions of the dataset to showcase our models in two distinct scenarios: imputing corrupted frames and generating missing frames.
As each video is relatively short in this dataset and can be grouped into a single batch, conducting full GP inference is feasible. This allows for a direct comparison of our approach with the full-batch baselines, i.e., GPVAE-Casee E-aaale (Casale et al., 2018) and GPVAE-Peare (Pearce, 2020).
Trajectory reconstruction Fig. 2a shows the recon- structed latent trajectories of three test videos using 10 near- est neighbours or inducing points. The RMSE appended to each column directly reflects the reconstruction quality. We can see that the VAE fails to learn the latent dynamics behind the videos, as the fully factorised Gaussian prior over
Imputing corrupted frames This task uses the sequences created by Krishnan et al. (2015). The dataset has 50,000/10,000 training/test sequences, each containing 10 frames. The rotations between two consecutive frames are normally distributed, with around 60% pixels absent ran- domly. Fig. 3a shows a sequence from the dataset.
6

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
Table 1: Performance of different models for tasks of the corrupted and missing frames imputation on Rotated MNIST.
Models
NLL
RMSE
Training Time (slepoch)
VAE (Kingma & Welling, 2014)
0. + 0.038
0.333 + 0.011
75.741 0.660
HI-VAE (Nazabal et al., 2020)
0,114 0.002
0.222 + 0.002
76.969 = 0.594
GPVAE-Diag (Fortuin et al., 2020)
0.094 + 0.000
0.199 + 0.000
127.614 + 2.758
Corrupted
UPVAE-Banded (Fortuin et al., 2020)
0.146 0.001
0.189 = 0.001
768.347 24.297
Frames
MGPVAE (Zhu et al., 2023)
0.090 + 0.001
0.197 + 0.001
110.772 = 0.651
GPVAE-HPA-H5 (ours)
0.095 + 0.000
0.199 + 0.000
137.860 = 0.326
GPVAE-SPA-HS (ours)
0.096 + 0.000
0.202 + 0.000
121.185 0.817
SYGPVAE-M10 (Jazbec et al., 2021)
87,711 + 0.484
6.122 + 0.015
8.280 + 0.111
SVGPVAE-M30 (Jazbec et al., 2021)
77.792 0.592
5.572 0.030
9.180 + 0.188
Missing
LVAE-MIO (Ramchandran et al., 2021)
73.887 = 0.748
5.578 0.022
28.919 0.253
Frames
LVAE-M30 (Ramchandran et al., 2021)
73.545 = 1.035
5.558 = 0.024
29.063 0.122
MGPVAE (Zhu et al., 2023)
72,217 + 0.192
5,496 + 0.005
25.103 + 0.103
GPVAE-HPA-HIO (ours)
71.264 + 1.263
5,521 + 0.046
9.819 + 0.064
GPVAE-SPA-HI0 (ours)
69.538 0.664
5.412 + 0.026
9.358 + 0.083
We follow the settings from Fortuin et al. (2020), where missing pixels are filled with zeros during inference, and the ELBOs are calculated only on the observed pixels in the data. A trade-off parameter added to the ELBO to rebalance the influence of the likelihood and KL terms (Hig- gins et al., 2017). For baselines, we additionally consider GPVAE-Diag and GPVAE-Band proposed in Fortuin et al. (2020), both of which are not amenable to mini-batchin batching. GPVAF-Diag uses an encoder with a diagonal covariance matrix, while GPVAE-Band features a tridiagonal covari- ance. Given that GPVAE-Diag shares the identical structure as our model and employs full batches of sequences, it serves the gold standard for our models. As suggested by Fortuin et al. (2020), all GP-related models employ Cauchy kernels, except MGPVAE using a Matérn-3) 2 kernel. We set H- = 5, which is half the length of a sequence.
ity 0.6, leaving the remaining sequence with varying lengths. This task aims to generate the missing frames at a given timestamp in unobserved sequences. SVGPVAE/LVAE- M10/30 utilise M 10/30 inducing points, and our mod- els use H 10 nearest neighbours for comparison. We adapt LVAE by manually assigning unique IDs to each new sequence during testing.
Fig. 3b and Table 1 provide the experimental results. As shown in Table 1, both our models surpass the baselines regarding NLL, with GPVAE-SPA achieving the lowest RMSE. This is further validated by Fig. 3b. SVGPVAE- M10 exhibits the fastest training time but the poorest per- formance. Our models take slightly longer training time than SVGPVAE-M10, with extra computational overhead caused by additional neighbouring data in a single mini- batch. Although we triple the number of inducing points, which incurs a similar training time to ours, it still cannot outperform ours. I.VAEs, which rely on additive kernels over sequence IDs, narrow the performance gap but still produce higher NLLs than our models. MGPVAE experi- ences longer training times on the long sequences due to its requirement for sequential forward and backward process, which cannot perform time-step mini-batching in parallel.
Imputation results on the test set are presented in Fig. 3a and Table 1. While VAE and HI-VAE achieve shorter training times, they fail to reconstruct the images faithfully due to factorised Gaussian prior. In contrast, our models closely match the performance of GPVAE-Diag, demonstrating strong approximation to the full-batch benchmark. GPVAE- Band achieves the lowest RMSE but incurs the longest training time due to its more complex structure. Although MGPVAE outperforms our model in terms of quantitative metrics, Fig. 3a indicates the generated images do not differ significantly in visual quality, suggesting that the overall performance of our models is comparable to MGPVAE.
5.3. MuJoCo Hopper Physics
The MuJoCo dataset collects physical simulation data from the DeepMind Control Suite (Rubanova et al., 2019). In this experiment, 500 series of 14-dimensional Hopper data are generated, each with 1000 timestamps. The series set is then split into training, validation, and test subsets by 320/80/100 At each timestamp, there is 60% chance that the relevant data features will be completely missing. The task objective is to make posterior predictions at unobserved time steps.
Generating missing frames Following the setup in Zhu et al. (2023), we produce sequences of MNIST frames, where each digit rotates through two full revolutions. This experiment involves 4,000 training and 1,000 test sequences, each with 100 frames. Each frame is dropped with probabil-
7

---

Neighbour-Driver GPVAEs for Scalable Structured Latent Modelling
Table 2: Results of conditional generation at missing timestamps on MuJoCo dataset.
Models
NLL
RMSE
Training Time (s/epoch)
SVGPVAE-M10 (Jazhec et al., 2021)
-1.512
0.032
0.0583 + 0.003
44.000 + 5.477
SVGPVAE-M30 (Jazbec et al., 2021)
-1.071 + 0.175
0.0391 + 0.001
47.893 + 3.170
MGPVAE (Zhu et al., 2023)
-1.708 0.057
0.0398 + 0.002
57.291 + 1.141
GPVAE-HPA-HIO (ours)
-2.335 + 0.032
0.0222 + 0.001
37.000 + 3.793
GPVAE-SPA-HIO (ours)
-1.715 + 0.159
0.0378 + 0.007
39.403 + 2.755
SYGEVAE
SYGUVAE
MMPPAR
GEVARSIPA
GIVAD-SPA
Input
True Layur
SVCPVAE-M100
SCPEAE-M20
RMSE-12207 GPVAR HIA- (12-112)))
RA4E-0094
GPVARSPA SPA. 1120
R3-SE-0 611
Tidksto
X07 Troditew Cariblau latern
Thas
Figure 4: Conditional generation at missing timestamps of a MuJoCo sequence in its 6th, 8th, 10th, and 12th dimensions with 95% credible intervals.
Figure 5: Porosity imputation (i.e., the 1st channel) at the intermediate layer of SPE10 dataset using GPVAE models.
the nearest neighbour principle should be particularly ef- fective. The first dataset, Jura (Goovaerts, 1997), contains hundreds of data points, while the second dataset, SPE10 (Christie & Blunt, 2001), consists of over 10" observations. Additionally, we compare several GP baselines such as ex- act GP (Rasmussen & Williams, 2006), SVGP (Hensman et al., 2013), and VNNGP (Wu et al., 2022). Appendix C.4 provides more experimental settings and additional results.
The results in Table demonstrate that our models consistently outperform both SVGPVAE and MGPVAE in terms of NLL, RMSE, and training time. These results are further supported by the visual comparison presented in Fig. 4. Additional numerical results and illustrative plots can be found in Appendix C.3.
Results
SVGPVAE-MIO shows suboptimal RMSE as an insufficient number of inducing points significantly limits the model capacity. Although increasing inducing points enhances the RMSE performance, SVGPVAE-M30 remains inferior to our models using 10 nearest neighbours. Moreover, increas- ing inducing points in SVGPVAE causes a deterioration in the NLL metric, indicating a reduced capacity to capture predictive uncertainty this experiment. In contrast, we show in Table 11 in Appendix C.3 that GPVAE-HPA and GPVAE-SPA consistently achieve better NLL and RMSE as the number of nearest neighbours increases. These re- sults highlight the advantage of employing neighbour-base baase approximation in the latent space over a strategy relying on inducing points. MGPVAE also demonstrates a poorer performance compared to our proposed models and, notably, requires significantly more computational time than ours.
Jura The Jura dataset provides measurements of three metal concentrations (Nickel, Zinc, and Cadmium) collected in a region of Swiss Jura (x, E R2,yn E R²), The train- ing set includes Nickel and Zinc measurements at all 359 locations, whereas Cadmium measurements are available for only 259 of these sites. We aim to predict Cadmium measurements at the remaining 100 locations.
Table 12 in Appendix C.4.1 compares the performance of our proposed models to competing methods. Our models again surpass the baselines in terms of both RMSE and NLL. The performance of Exact GP and VNNGP is suboptimal due to their inability to capture correlations between output variables. VAE and HI-VAE perform poorly as they do not leverage spatial information. SVGPVAE and SGPBAE exhibit improved performance with an increased number of inducing points; however, they remain inferior to GPVAE- SPA, even using three times the number of inducing points compared to the nearest neighbour count in our models.
5.4. Geostatistical Datasets
We further consider imputation tasks on two multi- dimensional geostatistical datasets of different scales, where
8

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
Table 3: Imputation results of various models on SPEIO dataset.
Models
NLL
RMSE
Training Time (slepoch)
MOGP-M1000 (Hensman et al., 2013)
6.034 + 0.003
1.230 + 0.000
3.820 + 0.021
VNNGP-H20 (Wu et al., 2022)
1.052 + 0.000
0.683 + 0.000
2,490 + 0.084
VAE (Kingma & Welling, 2014)
3,436 + 0.015
1.298 + 0.002
1.416 + 0.020
HI- VAE (Nazabal al., 2020)
0.752 = 0.034
0.647 = 0.008
1.448 + 0.011
SVGPVAE-M100 (Jazbec et al., 2021)
2.025 + 0.020
0.979 + 0.006
4.755 + 0.097
SGPBAE-M20 (Tran et al., 2023)
1.335 + 0.106
0.898 + 0.125
1585.6 = 21.3
0.735 + 0.007
0.638 = 0.007
4.019 + 0.036
GPVAE-SPA-H20
0.736 0.004
0.638 + 0.004
3.365 + 0.117
SPE10 SPE10 is large dataset widely used in petroleum engineering, featuring reservoir properties like porosity and permeability with high spatial helerogeneity, We downsam- ple it to 141,900 data points, with X, E I3 and yn E R4, masking about 50% of the values for imputation. This data size poses a challenge for GP-based models.
Acknowledgements
We thank the reviewers for their insightful comments and constructive suggestions. We also acknowledge the CSF3 the University of Manchester for providing GPU resources. Xinxing is supported by the UoM-CSC Joint Scholarship, and Xiaoyu receives support through the UoM Departmental Studentship for the Department of Computer Science.
Imputation results are summarised in Table 3. We present plots that illustrate the imputation of porosity at the interme- diate layer. In this highly heterogeneous dataset, methods that rely on inducing points, such as Multi-Output GPs (MOGP). SVGPVAE, and SGPBAE, all perform poorly, even using a large number of inducing points (e.g., MOGP- M1000). This is because inducing point methods struggle to effectively manage data that exhibits rapid local changes. In contrast, models employing nearest neighbour approxima- tion (including our models and VNNGP) generally perform better on this dataset, indicating that the nearest neighbour approach helps capture local structure. Importantly, our models outperform VNNGP with 15x fewer trainable pa- rameters, demonstrating that the VAE structure effectively models the latent correlations in the data.
Impact Statement
This paper presents work aimed at advancing machine learn ing. There are potential societal consequences of our work, none of which we feel must be specifically highlighted here.
References
Adam, V., Eleftheriadis, S., Artemev, A., Durrande, N., and Hensman, J. Doubly sparse variational Gaussian processes. International Conference on Artificial Intel- ligence and Statistics, pp. 2874-2886 PMLR, 2020.
Allison, R., Stephenson, A., Pyzer- Knapr-Knapp. E. O., et al. Leveraging locality and robustness to achieve massively scalable Gaussian process regression. Advances in Neural Information Processing Systems, 36, 2024.
6. Conclusion
In this work, we present an efficient neighbour-driven approximation strategy for modelling latent variables in GPVAEs, yielding two new variants: GPVAE-HPA and GPVAB-SPA. By focusing GP interactions nearest neigh- bours in the latent space, both models capture essential corre- lations and facilitate scalable mini-batch training. Empirical results on various tasks ranging from time-series imputation to geostatistical modelling demonstrate consistent accuracy and speed-ups over existing GPVAE baselines.
Ashman, M., So, J., Tebbutt, W., Fortuin, V., Pearce, M., and Turner, R. E. Sparse Gaussian process variational autoencoders. arXiv preprint arXiv:2010.1011 10177, 2020.
Cao, J., Kang, M., Jimenez, F., Sang, H., Schaefer, F. T., and Katzfuss, M. Variational sparse inverse Cholesky approximation for latent Gaussian processes via double kullback-leibler minimization. In International Confer- ence on Machine Learning, pp. 3559-35766 PMLR, 2023.
Limitations and future work While method does not require specific distance metric, we primarily used Fu- clidean distance. Exploring alternatives, such as correlation- based distances (Kang & Katzfuss, 2023), and leveraging advanced manifold-awer metrics in high- high-dimensional data are promising directions for future work.
Casale, F. P., Dalca, A., Saglietti, L.,, Listgarten, J., and Fusi, N. Gaussian process prior variational autoencoders. In Bengio, S., Wallach, H., Larochelle, H., Grauman, K., Cesa- Bianch,,, N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
Chang, P. E., Wilkinson, J., Khan, M. E., and Solin, A. Fast variational learning in state-space Gaussian process models. In IEEE 30th International Workshop Ma- chine Learning for Signal Processing (MLSP). pp. 1-6. IEEE, 2020.
Higgins, I., Matthey, L.., Pal, A., Burgess, C. P., Glorot, X., Botvinick, M. M., Mohamed, S., and Lerchner, A. beta- vae: Learning basic visual concepts with constrained variational framework. In International Conference on Learning Representations (ICIR), 2017.
Chen, T., Fox, E., and Guestrin, C. Stochastic gradient Hamiltonian Monte Carlo. In International Conference on Machine Learning, pp. 1683-1691, PMLR, 2014.
Jazbec, M., Ashman, M., Fortuin, V., Pearce, M., Mandt, S., and Rätsch, G. Scalable Gaussian process variational autoencoders. In International Conference on Artificial Intelligence and Statistics, pp. 3511-3511 PMLR, 2021.
Christie, M. A. and Blunt, M. J. Tenth spe comparative solution project: A comparison upscaling techniques. SPE Reservoir Evaluation & Engineering, 4(04):308- 317, 2001.
Johnson, J., Douze, M., and Jégou, H. Billion-scale similar- ity search with GPUs. IEEE Transactions on Big Data. 7 (3):535-547, 2019.
Kang, M. and Katzfuss, M. Correlation-based sparse in- verse Cholesky factorization for fast Gaussian Caussian-processs process inference. Statistics and Computing, 33(3):56, 2023.
Datta, A. Nearest-neighbor sparse Cholesky matrices in spatial statistics. Wiley Interdisciplinary Reviews: Com- putational Statistics, 14(5),11574, 2022.
Katzfuss, M. and Guinness, J. A general framework for Vec- chia approximations of Gaussian processes. Statistical Science, 36(1), 2021.
Datta, A., Banerjee, S., Finley, A. O., and Gelfand, A. E. Hierarchical nearest-neighbor Gaussian process models for large geostatistical datasets. Journal of the American Statistical Association, 11 111(514)8000-812 2016a.
Kim, H. and Mnih, A. Disentangling by factorising. In Dy, J. and Krause, A. (eds.), Proceedings of the 35th Inter- national Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 2649- 2658. PMLR, 10-15 15 Jul 2018.
Datta, A., Banerjee, S., Finley, A. O., and Gelfand, A. E. On nearest-neighhor Gaussian process models for massive spatial data. Wiley Interdisciplinary Reviews: Computa- tional Statistics, 1(5):162-111, 2016b.
Kingma, D. P. and Welling, M. Auto anto-encoding variational Bayes. In Bengio, Y. and LeCun, Y. (eds.), International Conference on Learning Representations, 2014.
Fortuin, V., Baranchuk, D., Rätsch, G., and Mandt, S. Gp- vae: Deep probabilistic time series imputation. In Interna- tional Conference on Artificial Intelligence and Statistics, pp. 1651-1661, PMLR, 2020.
Krishnan, R. G., Shalit, U., and Sontag, D. Deep Kalman filters. arXiv preprint arXiv:/51 1.005/2,, 2015.
Gardner, J. R., Pleiss, G., Bindel, D., Weinberger, K. Q., and Wilson, A. G. Gpytorch: Blackbox matrix-matris Gaussian process inference with GPU acceleration. In Advances in Neural Information Processing Systems, vol- ume 31, pp. 7587-7597, 2018
LeCun, Y., Bottou, L., Bengio, Y., and Haffner, Gradient- based learning applied document recognition. Proceed- ings of the IEEE, 86(1 1):2278 22244,2222, 1998.
Nazabal, A., Olmos, M., Ghahramani. Z., and Valera, I. Handling incomplete heterogeneous data using vaes, 2020.
Goovaerts, P. Geostatistics for natural resources evaluation, volume 483. Oxford University Press, 1997.
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z.,, Gimelshein, N., Antiga, L., et al. Pytorch: An imperative style, perfpeerformance deep learning library. In Advances in Neural Information Processing Systems, volume 32, pp. 8024-8035, 2019.
Hamelijnck, O., Wilkinson, W., Loppi, N., Solin, A., and Damoulas, T. Spatio-temporal variational Gaussian pro- cesses. Advances in Neural Information Processing Sys- tems, 34:23621-236333 2021.
Pearce, M. The Gaussian process prior vae for interpretable latent dynamics from pixels. In Symposium on Advances in Approximate Bayesian Inference, pp- 1-12. PMI., 2020.
Hensman, J., Fusi, N., and Lawrence, N. D. Gaussian processes for big data. In Uncertainty in Artificial Intelli- gence, pp. 282. Citeseer, 2013.
Hensman, J., Matthews, A. G., Filippone, M., and Ghahra- mani, Z. MCMC for variationally sparse Gaussian pro- cesses. Advances in Neural Information Processing Sys-
Siinonero-Candela J. and Rasmussen, C. E. A unify- ing view of sparse approximate Gaussian process regres- sion. The Journal of Machine Learning Research, 1939-1959, 2005.
tems, 28, 2015.
10

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
Ramchandran, S., Tikhonov, G., Kujanpää, K., Koskinen, M., and I. Lähdesmäki, H. Longitudinal variational autoen- coder. In International Conference Artificial Intelli- gence and Statistics, pp. 1898-3906. PMLR, 2021.
Tran, G.-L., Milios, D., Michiardi, P., and Filippone, M. Sparse within sparse Gaussian processes using neighbor information. In International Conference Machine Learning, pp. 19369910378. PMLR, 2021b.
Rasmussen, E. and Williams, C. K. I. Gaussian Processes for Machine Learning. The MIT Press, 2006. ISBN 026218253X.
Vecchia, A. V. Estimation and model identification for con- tinuous spatial processes. Journal of the Royal Statistical Society Series B: Statistical Methodology, 50(2):297-31 1988.
Razavi, A., van den Oord, A., and Vinyals, O. Generating di- verse high-fidelity images with vq-vae-2. In Wallach, H., Larochelle, H., Beygelzimer, A., d' d'Alché-Buc, F., Fox, E.,, and Garnett, R. (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
Wu, L., Pleiss, G., and Cunningham, J. Variational nearest neighbor Gaussian process. In International Conference Machine Learning. pp. 241 14114-24 PMLR, 2022.
Zhu, H., Balsells-Rodas, C., and Li, Y. Markovian Gaussian process variational autoencoders. In International Con- ference on Machine Learning, pp. 42938-42961, PMLR, 2023.
Rubanova, Y., Chen, R. T., and Duvenaud, D. K. Latent ordinary differential equations for irregulary-ssappler time series. Advances in Neural Information Processing Systems, 32, 2019.
Särkkä, S. and Solin, A. Applied stochastic differential equations, volume 10. Cambridge University Press, 2019.
Särkkä, S. et al. Recursive Bayesian inference on stochustic differential equations. Helsinki University of Technology, 2006.
Solin, A. et al. Stochastic differential equation methods for spatio-temporal Gaussian process regression. PhD Thesis, 2016.
Sønderby, C. K., Raiko, T., Maaløe, I., Sønderby, S. K., and Winther, O. Ladder variational autoencoders. Advances in Neural Information Processing Systems, 29, 2016.
Springenberg, J. T., Klein, A., Falkner, S., and Hutter, F. Bayesian optimization with robust Bayesian neural net- works. Advances in Neural Information Processing Sys- tems, 29, 2016.
Stein, M. L., Chi. Z., and Welty. L. J. Approximating likelihoods for large spatial data sets. Jourmal of the Royal Statistical Society Series B: Statistical Methodology, 66 (2):255229) 22755-296.222 2004.
Titsias, M. Variational learning of inducing variables in sparse Gaussian processes. In International Conference on Artificial Intelligence and Statistics, pp. 567-574. PMLR, 2009.
Tran, B.-H., Rossi, S., Milios, D., Michiardi, P., Bonilla, E. V., and Filippone, M. Model selection for Bayesian toencoders. Advances in Neural Information Processing Systems, 34:19730-19742, 2021a.
Tran, B.-H., Shahbaba. B., Mandt, S., and Filippone, M. Fully Bayesian autoencoders with latent sparse Gaus- sian processes. In International Conference on Machine Learning, pp. 34409-34440 34430. PMLR, 2023.

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
A. Acronyms
We list all acronyms appearing in the paper in the following Table for reference.
Table 4: Acronym summary
Acronym
Meaning
GP
Gaussian Process
VAE
Variational AutoEncoder
GPVAE
Gaussian Process Variational AutoEncoder
NNGP
Nearest Neighbour Gaussian Process
SVGP
Sparse (Stochastic) Variational Gaussian Process
SVGPVAE
Sparse Variational Gaussian Process Variational AutoEncoder Longitudinal Variational AutoEneoder
LVAE
MGPVAE
Markovian Gaussian Process Variational AutoEncoder
SGPBAE
Sparse Gaussian Process Bayesian AutoEncoder
HPA
Hierarchical Prior Approximation
SPA
Sparse Precision Approximation
NLL
Negative Log-Likelihood
RMSE
Root Mean Squared Error
B. ELBO Derivation
This section details the ELBO derivation of (10) and (13) in Section 3.2. Throughout, we use subscripts denote the learnable parameters of the distributions: & represents the encoder parameters, 0 denotes the decoder parameters, and 1b is associated with the latent GP parameters.
B.1. GPVAE-HPA ELBO
We begin by outlining the derivation for GPVAE-HPA. In each training iteration, randomly sample mini-batch B (xx with the index set Z. For each training point (x) E B, look for the И locations X.(i) E / in X = |x= such that nii) are indices of top H nearest points to X,. This neighbour-driven strategy materialises the sampling binary indicative vector w for latent variables. Then, the ELBO given by
/ q(z w) log
Y
Z)p(2 Z w)p(w) dw
IYg >
oZZ w)
- p(w)
{ q(Z w) log pgp[] Y Z) dz q(Z w) log p(Z w)
dz dw (Z w)
Bptw) (By(z)m) log Z) KT. (4 w) || p(Z w]]}
Bp(w)
2.))
KL (Z w) p(Z w]]
N
Eg(z)) [logply, 22)]
KL [g(Z(()) || p(Zz())))
= CHPA.
N
REI
The following part provides a detailed derivation of each term the above-mentioned ELBO. Specifically, The ELBO CHPA using a hierarchical prior given by
CHPA
E (m.ls)) [log
z;)]
tre-parameterization trick)
T
IEI
KL (Zn(e)
Y xn()) |I Pe (Zn(i) Xn(i))
(closed form)
I
IET
KL|gs [[v]]
12

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
Expected log-likelihod (if the decoder models Gaussian distributions with mean μl and variance of in the MuloCo experiment in Section 5.3.)
E (xill))) [log Pa z:]] 15
(xrllrr)) [log/Vxx [zo(z)), (z())))
k
(y (z))) log 2TT (2.) (2.))
209 (2:)
2
Here, we slightly abuse notation by writing z; (z) ys) to indicate re- parameterizze samples. The superscript A: represents the k-th entry of a vector.
KL divergence
KL
190
|1 Pul = Σ KL (N(zh 1 Eg) 11 (1 I |/}
1=1
(uls m4/) (2,) (ul' m/) + Ur -- + log EX1 log - - -nn()]]]
where 1 represents the latent dimension and Σn is diagonal.
M1g 1 (Ynt))
Σg = og (Yn()):
n1 = 1/s (Xn()))
Σ, = K! (Xn(1)) Xn(s))
B.2. GPVAE-SPA ELBO Next, we derive the ELBO for GPVAE-SPA. Following the same notation as in the main text, we have
logp(Y) 2 (()) log p(Y
Z) j=== p(z))
Zn(i)) dz
((()
IT, N j=1 ((z))
= - (()) Z) dZ
(Z)) log
dz
/=11 p(z) Zn(ii)
-/ (ZZ))og((Y Z) Σ
/ a ((Zne)) log plz; Zn(z))
((2,)
dz dZnne))
N
Σ By(xs) [logply: x:)]
Σ j=l
(Z(zne))) KL [a(z)) p(z Zn(z))]
In every training iteration, randomly sample a mini-batch of training data indices I and mini-batch of inducing indices J for stochastically maximising the following ELBO
N
N
CSPA
Bns(aily)) log loV 2,)]
E
[KL(qa(z) 1ys) (z Zn Xag)))}.
I
|T
I.T.
JET
Expected log- likeliool is the same as in LHPA-
KL divergence
EgglZoon (Zacn (Yn)))
(KL yz) ll (z Zn(i), X,t)))}}
L
E
yo (Ylcn))aaan) (at))
KLlga(zz) 1ys)
pol (z²) Znetii> Xncsrll
KL[gal [Po]
13

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
First, given a data index j and the corresponding neighbour set Zn(z), we have
KL l9o |I pvl = KL [A(z))
My:
Eng
Mp Ep]]
(
m4/)2
Σ
+ log Erpp - log 12/1 -
2 1=1
Ep
Σ
(14)
1
up) + + En + log Epp
(
(1gg 2
L
2
En
where 2, E RR, Znez) E RloO)|xL and both pn and En are diagonal:
Mg = mo(y) Eg = =3](y);;
[ n ml (x²z + m/,))) K' KK m(i)m(()) [ZnC]) m' (enn))))
21pp = - m(n) K! K- (o)))))))k(..
K Kno))s"" Ep - diag (2)))
n()..(t))
Then, calculate the expectation w.r.t.
En(z) II (z (n(o))) II (Zhen Pleyy Shar))
where V'())) 1//(X n()) and She() diag(s)er)) 02 (Yng)). Therefore,
E (aas)) [(A - 1 = bb D' She b' miys + { m (x,) + b² bb (xGss(mn)) m' (Xn(j)) -11
where b/ D' K!- expression.
m())m())) Ky (n))ss The above expression can be plugged into (14) derive the KI. term of the ELBO
C. Experimental Details and Additional Results
Our implementation is open-sourced at https::
com/ http:///www....ii//////////o-aaa
For fair comparisons, most scalable models (including our models, SVGPVAE, MGPVAE, SGPBAE, VAE, HI-VAE and GP models) are implemented PyTorch (Paszke et al., 2019) and GPyTorch (Gardner et al., 2018). We use the official code for the baselines GPVAE-Diag and GPVAE-Bant from Fortuin al. (2020) and modify the code make the running process comparable with others. The experiments are run NVIDIA A100-SXM4 V100-SXM2 GPU of a high- perforeeformncc cluster. We use the modern similarity search package, Faiss (Johnson et al., 2019), for nearest neighbour searches. Most training time is estimated on an NVIDIA RTX-4090 GPU, except for the missing pixel imputation task, which is tested on an RTX-2080-Ti due to software compatibility.
While the ELBO of GPVAE-SPA can theoretically be computed using two separate mini- batches, in practice, we follow Wu al. (2022) to simplify the process by using the same mini-batches of data each iteration. The experimental details are provided in the following sections.
C.1. Moving Ball
Experimental settings All GP-related models use RBF kernels with the lengthscale initialised to 2. In each epoch, we generate 35 videos using distinct local seeds and train for 25,000 epochs. Our models use the same multilayer perceptron (MLP) structures and training settings in Jazbec et al. (2021) and Pearce (2020). Two independent GPs control the two-dimensional variable in the latent space of the GPVAE. We summarise the experimental settings in Table 5.
In each training epoch, 35 videos with 30 frames each are simulated for parameter learning. The decoder yields independent Bernoulli distributions for 1024 pixels of each frame. When testing, we derive the latent trajectory by least- -quares projection and compute the sums of squared errors on a testing batch of 35 videos. The MSE for one experimental trial is obtained by averaging the squared errors over 10 video batches. The RMSE is then the square root of the MSE
14

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
Table 5: Experimental settings for the moving ball experiment.
Setting
Value
Videos in each epoch
35
Frames in each video
30
Frame size
32x 32
Encoder (MLPs)
1024 + 500
Latent dimensionality
2
Decoder (MLPs)
2 - 500 - 1024 + Sigmoid
Activation function
tanh
Optimizer Training epochs
Adam, Ir = 0.001
25000
Additional results We provide RMSEs for the two proposed models in Table 6, with the number of nearest neighbours H ranging from 3 to 15. The results confirm that the reconstruction improves with more nearest neighbours, which enables the models to leverage more information from data in this experiment. We also present the reconstructed latent trajectories in Fig. 6 with 5 or 15 inducing points/nearest neighbours.
Table 6: Additional experimental results in terms of RMSE with different numbers of nearest neighbours.
H
3
7
10
15
GPVAE-HPA
17,400 + 3.691
8. 3.149 + 3.294
4.959 + 0.250
2.978
0.101
1.709 + 0.146
GPVAE-SPA
3,507 0.318
1.905 + 0.148
1.289 + 0.065
1.079 + 0.097
1.245 = 0.063
VAE
SVGPVAE
CPVAE
SPVAESPA
VIDBO
SVGPVAE
CPVAE
CPVAE
CPVAESPA
grvw
(b) Latent trajectory reconstruction with H 15 or M = 15. Figure 6: Additional illustration of latent representation learning for the moving ball series using different numbers of nearest neighbours/inducing points. (a) GPVAE-SPA has already achieved good performance using only 5 nearest neighbours. (b) When using 15 nearest neighbours/inducing points, all models except VAE achieve good reconstruction quality, but ours still show better RMSE performance.
(a) Latent trajectory reconstruction with H 5 or M =
C.2. Rotated MNIST
The RMSES reported in the following experiments are obtained through the encoder mean, computed at missing pixels and averaged. The NLL is evaluated using 20 samples from the latent variables. Specifically,
NLL = 1--logg log /
p(Y Z)g(Z) dz ~ - llog- 1 Ex(y Z.)
5
= log S - logsumexpp=s..-. log Z.),
where S is the number of samples.
15

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
Table 7: Experimental settings for the Healing MNIST pixel imputation task.
Setting
Value
Training/test sequences
50,000/10,000
Frames per sequence
10
Frame size
28 x
Encoder (Convs+ MLPs)
Conv(1, 256, ks=3) - Conv(256, 1, ks=3) 256
256 - 512
Latent dimensionality
256
Decoder
256 256 - 256 256 784 + Sigmoid
Activation function
ReI.U
Optimizer
Adam, /r = 0.0005
Mini-batch of sequences
50
Training epochs
Trade-off parameter
HPA=1.5; SPA=1.0
C.2.1. CORRUPTED FRAME IMPUTATION
The dataset, also called Healing MNIST, simulates real-world situations where medical data is often incomplete (Krishnan al., 2015). For this task of missing pixel imputation, follow the settings in Fortuin al. (2020), where Cauchy kernels and convolutional layers (Conv) are used as components of the model structure. The MGPVAE baseline uses Matérn- kernel. The lengthscale and outputscale of all kernels are initialised to 2 and I, respectively. The settings are listed in Table 7. Fig. 7 illustrates additional imputation results for other digits in the dataset.
Input
Iaput
13
VAE
VAE
6
6
6
6
6
6
6
6
6
6
III VAE
HI VAE
6
6
6
6
6
6
6
6
6
GPVAE Diag
GPVAE Diag
6
6
6
6
6
6
6
6
6
6
/ult
(full
GPVAE Band
GPVAE Band
6
6
6
6
6
6
6
6
6
6
(fufl
0
0
D
2
7
6
6
6
6
6
6
6
6
6
6
GPVAE
3
3
3
3
GPVAE
6
6
6
6
6
6
6
6
6
6
(murs)
(oura)
GPVAE.SPA
3
3
3
3
GPVAE.SPA (Irra)
6
6
6
6
6
6
6
6
6
6
Figure 7: Additional imputation results on the Healing MNIST dataset.
C.2.2. MISSING FRAME GENERATION
Experimental settings We employ similar architecture to Zhu et al. (2023), with convolutional layers in the encoder and deconvolutional (Deconv) layers in the decoder. We use a Matérn- 3 kernel for MGPVAE and RBF kernels for the others (with the exception LVAE, which uses an additive kernel incorporating distinct sequence IDs). The lengthscale and outputscale of all kernels are initialised to 2 and I, respectively. In SVGPVAE, we select the initial locations of the inducing points be evenly distributed across the entire time range.
Additional results In addition to the H = 10 results detailed in the main text, we evaluate our models using varying numbers of nearest neighbours. The metrics for NLL and RMSE are presented in Table 9. This experiment demonstrates that increasing the number of nearest neighbours can enhance the visual quality of the generated images. Fig. illustrates some additional results the task of generating missing frames based on the timestamps.
16

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
Table 8: Experimental settings for the missing frame generation experiment on Rotated MNIST.
Setting
Value
Training/test sequences
4,000/1 0000
Frames per sequence
100
Frame size
28
28
Encoder (Convs MLPs)
Conv(I, 32, ks=3) Conv(32. 64, ks=3)
32
Latent dimensionality
16
16 + 1568 De2 64, ks=3)
Decoder
Deconv(64, 32, ks=3) - Deconv(32, 1, ks=3) Sigmoid
Activation function
ReLU
Optimizer
Adam, lr 0.0005
Mini-batch of sequences
50
Mini-batch of frames
20
Training epochs
300
Trade-off parameter |
HPA=1.5; SPA=1.0
Table 9: Our models'
performance on the task of missing frames generation with different numbers of nearest neighbours.
H
10
20
GPVAE-HPA
101,741 + 0.712
71.264 + 1.263
69.637
1.027 69.516
NLL
GPVAE-SIE
70.865 0.656
69.538 + 0.664
0.543
GPVAE-HPA
6.669 = 0.033
5,521 + 0.046
5.444 + 0.043
RMSE
GPVAE-SPA
5.472 0.029
5,412 + 0.026
5.425 +0.020
C.3. MuJoCo Hopper Physics
Experimental settings All models use 15-dimensional latent space and Matérn- - kernels. We adopt the same encoder and decoder architectures and training configurations as Zhu et al. (2023). For each latent GP, we initialise the lengthscale and outputscale of the Matérn- 3 kernel to 50 and 1, respectively. All models are trained for 500 epochs. At test time, we draw 20 latent samples to compute the NLL and RMSE. These metrics are computed over the entire sequence for predictive performance on the test set. Additional setups are presented in Table 10.
Table 10: Experimental settings for the MuJoCo Hopper Physics experiment.
Setting
Value
Encoder (MLPs)
14 > 32 + 30
Latent dimensionality
15
Decoder (MLPs)
15 16 14
Activation function Optimizer
ReLU
Adam, Ir 0.001
Mini-batch of sequences
16
Mini-batch of frames
32
Training epochs
500
Table 11: Additional NLL and RMSE results for varying numbers of nearest neighbours on MuJoCo dataset.
H
10
20
GPVAE-HPA
-2.322 + 0.061
-2.335 + 0.032
-2.332 + 0.021
NLL
GPVAE-SPA
-1.700 + 0.237
-1.715 + 0.159
-1.835 + 0.125
GPVAE-HPA
0.0225 0.001
0.0222 0.001
0.0221
0.000
RMSE
GPVAE-SPA
0.0392 0.011
0.0378 0.007
0.0339 0.005
17

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
7
7
SVGPVAE (M
SVGPVAE IM
IVAE
IVAE
MGPVAE
GPVAE-HPA (II oura)
GPVAE-SPA (M oura]
Input
4
4
H
SVGIVAE
4
(M
SVGPVAE
4
(M
IVAE (M
4
IVAE
4
4
MGPVAE
4 11
GPVAE.HPA (If ours)
4 u
GPVAE-SPA o0rs)
Figure 8: Missing frames generated from two unseen MNIST sequences. The red boxes indicate the missing frames.
Additional results We provide additional NLL and RMSE results for proposed models in Table 11, with Hl E (5, 10, 20}. These results confirm that the predictive performance increases with the number of nearest neighbours, consistent with our early findings. Fig. 11 shows additional plots for test MuJoCo sequence across its 1st, 2nd, 3rd, 4th, 5th, 7th, 9th, 11th 13th, and 14th dimensions.
C.4. Geostatistical Datasets
C.4.1. JURA
Table 12 presents the results on the Jura dataset discussed in the main text. We also illustrate the prediction in Fig. 9,
Experimental settings We follow the experimental setting in (Tran et al., 2023), in which RBF kernels are used, and the lengthscale and outputscale are initialised to 0.1 and More details are in Table 13. SGPBAE (Tran et al., 2023) adopts a stochastic encoder, for which the input size is 6 instead of It uses Adaptive SGHMC (Springenberg al., 2016) with learning rate Ir 0.001, and the encoder is trained by Adam with learning rate Ir 0.001. After 500 burn-in iterations, 50 samples collected every 20 iterations. The exact GP and VNNGP are applied to the Cadmium variable only (because they do not model multi-output data natively), while all other models utilise information from all 3 variables.
C.4.2. SPE10
The SPE10 dataset is derived from the Tenth Comparative Solution Project in petroleum engineering. widely used as a benchmark for assessing how different numerical and geostatistical methods perform in heterogeneous reservoir conditions,
18

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
Table 12: Imputation results of various models on the Jura dataset.
Models
NLL
RMSE
Exact GP (Rasmussen & Williams, 2006)
1.305
0.724
VNNGP-HIO (Wu et al., 2022)
1.387 = 0.002
0.700
0.000
VAE (Kingma & Welling, 2014)
2.015 = 0.344
0.994
0.121
HI-VAE (Nazabal al., 2020)
1.082 = 0.079
0.710
0.028
SVGPVAE-M10 (Jazbec al., 2021)
1.032 = 0.025
0.678 0.016
SVGPVAE-M20 (Jazbec et al., 2021)
1.002 = 0.040
0.661 0.026
SVGPVAE-M30 (Jazbec et al., 2021)
0.981 = 0.062
0.650 + 0.039
SGPBAE-MIO (Tran et al., 2023)
2.171 = 0.826
1.089 = 0.253
SGPBAE-M20 (Tran et al., 2023)
2.005 = 0.781
1.061 = 0.244
SGPBAE-M30 (Tran al., 2023)
1.435 = 0.316
0.913 0.124
GPVAE-HPA-H10
1.230 + 0.184
0.678 + 0.054
GPVAE-SPA-HI0
0.939 = 0.066
0.584 = 0.023
Table 13: Experimental settings for the Jura experiment.
Setting
Value
Encoder (MLPs)
3 + 20 + 4
Latent dimensionality
2
Decoder (MIPss
- 5 +
Activation function
ReLU
Optimizer
Adam, Ir = 0.001
Mini-bated size
100 (except exact GP, which uses all 259 training data) 2000 for exact GP and VNNGP, 300 for others HPA=1.8, 8, 1.0 for others
Training epochs
Trade-off parameter /
The original dataset contains 220 x 60 x 85 grid cells. Each grid cell represents a discrete unit (a rectangular cuboid) in the 3D reservoir model. Its typical properties are porosity, which indicates the fraction of pore space in the cell, and permeability (often provided in æ, y, and z directions). These properties collectively determine the flow behaviour of fluids (oil, water, or gas) within the reservoir.
Experimental settings In this experiment, we downsample the cells by factor of 2, resulting in 110 x 30 x 43 locations. Each value of the dataset has a 0.5 probability of being dropped. The experimental settings are listed in Table 14.
All models use Cauchy kernels with the lengthscale and outputscale initialised to 2 and The likelihood noise is fixed 0.25 except for GP models, whose likelihood noise is initialised to 0.01. For SVGPVAE, the inducing locations uniformly initialised in the data space. For SGPBAE, the SGHMC step size is set to 0.0001, and the steps along the chain are 100. After 1500 burn-in iterations, we keep 100 samples. The performance of SGPVAE is evaluated using 100 samples. Between each sample, the stochastic encoder is updated for 50 iterations.
Table 14: Experimental settings for the SPE10 experiment.
Setting
Value
Encoder (MLPs)
4 256 - 64 3
Latent dimensionality
3
Decoder (MLPs)
3 - 64 256 - 4
Activation function
ReLU
Optimizer
Adam, (r = 0.001
Mini-batch size
1000
Training epochs
500
Trade-off parameter 3
HPA=1500, 0.2 for others
19

---

Neighbour-Driver GPVAEs for Scalable Structured Latent Modelling
Additional results
We plot the imputation results of the models some layers in Fig. 10.
Figure 9: Predictive mean for Cadmium on the Jura dataset test split (100 locations).
Ispa
MOGP
YIOGGNIEE
YAE
JIVAS
KYC
ACPRAK NOS
OIPAR
Lorer
SYGPVARXXI
SOUDAENII
GI APA
Lapur
SYGPPAR
KGPRAR MOS
OPYAR HPA Raian
OPEAR SOA HRR)
Figure 10: Additional imputation results of SPEIO. (top) The permeability along the y direction at the 26th layer; (middle) The permeability along the x direction at the 41st layer; (bottom) The porosity at the 22nd layer.
20

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modellngg
SVGPVAE 10)
SVGPVAE
SVGPVAE 30)
MCPVAE
GPVAE-HPA o0.sss))
GPVAE-SPA 10, ours)
d d=1
031
d ==2 2
011
1.80
d 3
1.8
d
5
d
7
1.00
0.4
d 9
2.5
1.50
0.89
d
11
13
1.90
d= 14
2.00
Predictive Mean
95% Predictive Coufidence Interval
Ground Truth
Figure 11: 95% posterior credible intervals for test MuloCo sequence for all other dimensions with 1000. The red dots mark unobserved data points.
21

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelligg
D. Scalable Gaussian Process Variational Autoencoders
give an overview the related models in the following Table 15.
Table 15: A summary of relevant models. N, M, and H are the number of data points, inducing points, and nearest neighbours. Ng is the mini-batch size. P represents the feature dimensions of low-rank representations. is the dimensions of the state variables.
Arbitrary kernel
Standard VAE encoder
Models
Mini-batching'
GP complexity
Reference
VAE
/ O(NP2 + pd) ((NS)
Kingma & Welling (2014)
GPVAE-Casale
Casale et al. (2018)
GPVAE-Pearce
/
Pearce (2020)
GPVAE-Diag
O(NS)
X
Fortuin et al. (2020)
GPVAE-BIS Bandd
(NN))
X
X
Fortuin et al. (2020)
SVGPVAE
O(N + M3) N)
X
Jazbec et al. (2021)
MGPVAE
x
Zhu et al. (2023)
SGPBAE
O(NM + M
Tran et al. (2023)
GPVAE-HPA
H*) (N
this work
GPVAE-SPA
/
this work
Here means mini aani-batching along the "time" dimension.
In this section, we summarise recent scalable GPVAE models used in our experiments, including SVGPVAE (Jazbec et al., 202 (Section D.I), MGPVAE (Zhu et al., 2023)(Section D.2), and SGPBAE (Tran et al., 2023)(Section D.3). We believe this section will be helpful for practitioners interested in trying these models. Although SGPBAE more accurately categorised as an extension of the Bayesian autoencoder (Tran et al., 2021a) rather than a VAE, included here for completeness. For more detailed information about these models, please refer the related papers.
SYGPVAE
Replacing exact GPs with Sparse Variational GP (SVGP) techniques (Titsias, 2009; Hensman al., 2013) to enhance the scalability of GPVAEs is natural approach. However, such methods cannot be directly integrated into GPVAEs, they either require the entire dataset to be loaded into memory challenge, Jazbec et al. (2021) propose a novel approach that leverages specific parameterisation of the distribution of inducing variables based on Pearce (2020). This formulation enables the training objective to accommodate mini-batching. amortization, and the use of arbitrary kernels.
not suited for amortization. In response to this scalability
In Pearce (2020), they use the variational distribution from (5), and the ELBO given by:
N
COPYAE-Peare
Eq (an lym)
logpo(y,. 2n,)
Σ log (z vn)
+ Σ log o (Y, 2m1
where Zo.0 (Y) J (x
The posterior (z yn) can be regarded derived from "pseudo" GP regression with input-observation pairs (X.rk and noise oz (Y). We denote the latent dataset as (X.Y Y = (Y), == o3' (Y)). This training objective will require O(N*) time complexity due to log Z/,o (Y,X) and do (Zn Yn) terms without any approximation. Jazbee al. (2021) approximate the encoder's output distribution o using SVGP's posterior:
(z(z) / p(z Zu)g(Zu) dZu
(15)
= K Uon, Kxx + KRuKob (A Kuu)
where are M inducing locations and the variational distribution gZZu) = NN p. A). Minimizing KL.gg. ||
22

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
equivalent to maximizing the (inner) lower bounds Lи proposed by Hensman al. (2013)²
log Z0 2 LH = Σ log N (0e KK,,, 03)
Jas + tr(AA:)) (l)))Zu)) |l (uu))
(16)
2
where ie and o, are from the latent dataset, kys = X Kx Kx.uKUu Kux: and A, KüoKus K The practical training objective for SVGPVAE given by replacing t0.0 with ys from (15) and lowering bound log % with CH from (16) simultaneously:
N E E (z..) logpolyn (,,)
ELBO
Σ log ( yn)
+ Enn
(17)
Mini-batching in ELBO Take mini-batch of size No, X, c C, C Y, and the encoder ia creates a mini-hatch of the latent dataset (X,, Yo. op.. SVGPVAE computes stochastic estimates for μu, A of the variational distribution (Zu) = (Zu pe, A) at each latent channel I:
N
Σx KUU + Kux, diag(o," 2)x,,,
No
14. = Kuu(E)) (oo)(E")KKuu "Kux, (
N
(18)
No
A = ()u(E)) KKU..
These estimators converge to the true values for N₁ - N. Σ/ is an unbiased estimator for Σ' but this doesn't hold for p4 and A Then, the training objective SVGPVAE for a mini-batch Xs, Y, based on the above Monte Carlo estimators given by
N
N Bg (z.))
LSVGPVAB
logpolyn zn. Σ log qo (z', yn)
+ Σ CH.
No
Prediction After training, generating y. for an unseen follows the procedure:
Encode (possibly all) training sample to get {uys (Y) lo (Y))
Compute Eny
R MMM p'y RM and A'y R MxM using the estimators (18) for I {1. L}.
Compute posterior predictive distribution:
q(zz Y, x.), / p(z, Zu)N(Zu (N. dZu.
which is Gaussian with mean and variance
H._ Kx, Kx,u(EN) - KKx diag(o)"(Y))uu
0²' Kx,x, - Kx. .Ku (Kuo - -- K Kul
Compute z, = + a. e, with sampled from a standard normal distribution, and generate y. ~ poly. 2.).
D.2. MGPVAE
This work is inspired by the fact that some miedimensional GPs with Matérn kernels can be written linear Stochastic Differential Equations (SDEs) (Särkkä & Solin, 2019), which has an equivalent discrete linear state space representation. For Gaussian likelihood, the linear discrete-time representation enables Kalman filtering and smoothing that computes the posterior distributions in linear time w.r.t series length. This work places such discrete GP representations on the latent space of GPVAE. However, the GPVAE decoder is a non-linear function of latent variables. They further apply site-based approximation (Chang et al., 2020) for the non-linear likelihood terms to enable analytic solutions for the filtering and smoothing procedures.
We suppress channel index here for simplicity.
23

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
D.2.1. SDE REPRESENTATIONS FOR MARKOVIAN GPs
GP 2~ GP(O, k) with a Markovian kernel (e.g., Matérn- 2 can be written with an Itō SDE of latent dimension d (Solin et al., 2016; Hamelijnck al., 2021):
ds; Fsydt + LdBr, 2t Hs:,
(19)
where F E Rddd is the feedback matrix, IR the noise effect, E xx the emission matrix. B, is an e-dimensional (correlated) Brownian motion with spectral density matrix Qe E Rexe such that B B, ~ (Q..
Suppose the initial state St, (mo, Po) with the stationary state mean and covariance. Given all timestamps ( .tr.
we solve for each with initial time 1,-1, then the the linear SDE (19) admits:
de 1"
St elti-ta St.
(
Je
Therefore, can derive a corresponding discrete-time solution, allowing the recursive updates of 5 given St, (Särkkä al., 2006):
St: = A1,+ 1St, + q. Sto ~N Po),
(20)
, eA, q, NO Qt,+1),,
where A, t++1 - t.. Provided that Po exists and known, Q1,+1 can be easily obtained in closed form (Solin et al., 2016):
Q c Po A,, 10
Note that d, F.F. H.H.H mo Po. Qu in (19) and (20) depends on the kernel's type and parameters.
D.2.2. MARKOVIAN GPVAES
MGPVAEs utilize the linear SDE form of Markovian GPs in the latent space and transform it into a discrete State-space model (SSM) with non-conjugate measurements (i.e., non-linear likelihood modelled by a decoder). Specifically, given N (possibly unevenly spaced) time steps {41, .... (N}. for the I-th latent dimension,
s11 A A4.441s + d
s/ ~ (( P)). of ~ N(O. (1++1): 2 H's y: z po zz),
where we abuse the timestamp subscripts, using s; E dL to denote the state variable at time 4, 2, E R² is the latent variable of the GPVAE, and / represents the parameters of the decoder. The L latent Markovian GPs are independent, and we will suppress the channel in subsequent expressions keep the notation uncluttered.
Model Provided SSM with system states {s: [s and observations (Y.) the forward process is described by the joint distribution
N
p(s: - II p((s)) (Y) = II
(Y, s:).
i=1
f=1
For = 0, O,....AY 1, the factors of the above distribution are defined by
p(s s;) -s A1,418; Qi,41),
ply: s;) (Y: Hs,).
Again, the transition matrices {A,i+1, Q and the emission matrix H are determined by the Markovian kerel. The measurement p(y; s;) is modeled by a non-linear decoder network, so the smoothing distribution p(s; | Yi:) no longer analytical. Like Pearce (2020), this work utilises pseudo observations yI.N and pscudo noise variances V1.V from the output of the encoder to substitute the measurement model p(y: s:) with
plý: s:)) N((;; Hs,, V.)
24

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
such that the linear Gaussian relationship is maintained and the standard Kalman filter and the RTS smoother can be applied. The smoothing distribution conditioned on the pseudo-observations
p(s II Myy Hs,, V.) P(51.N VI.N) = J (((NN))) II N Hs;, v.)ds
(($1:N),
will be used as a variational approximation g(S1.r) to the true posterior p(S yi:s) later.
Variational Inference This work employs variational inference to approximate the intractable posteriors. The variational distribution ((81:N) is proposed to approximate the posterior p($1:N YI:N) over S1:N E RNxxL Note that z, is a deterministic linear transformation of s; so the stochasticity arises from S1.N- The variational (($1.N) is learned by maximizing the following ELBO:
N
MMGPAA
((s) logpo st) - KL (g((NN)))) ll p(S
N
log II NWO; Hs; v.) log / p(su.) II N Hs, v.) ds
(x)) log py ss) - Bo(s
1=1
Σ
(e.) log (yg |s) "'E yas,)) log WS; Hs, v4) log p((II))) II N Hs,, ds
(x))) plog s:) - log Ny; Hs,, v.)] + logp(ý1.N).
Prediction Computing the state-space posterior distribution at a newly introduced time point 1, relies on the smoothed dis- tributions obtained the adjacent time points. More concretely, let and 6. denote the immediate predecessor and successor time points, respectively, and the predictive q(sr.) )) can be obtained as ((...) = J p(st, st_, St,, (S St_ 8t_ dss, ds... More details can be found in Appendix A.I of Adam al. (2020). Using the relationship zr. Hs.. can obtain the predictive posterior (2,,)) and q(y.).
D.3. SGPBAE
This work presents fully Bayesian autoencoder (BAE) that treats the GP parameters, the latent variables, and the decoder parameters in Bayesian fashion (i.e., SGPBAE), To carry out scalable inference, the authors (1) use the fully independent training conditionals (FITC) (Quinonero-Canella & Rasmussen, 2005), (2) adopt stochastic Hamiltonian Monte Carlo (SGHMC) (Chen et al., 2014), and (3) employ an amortized stochastic network as the encoder to learn to draw samples from the posterior of the latent variable. Additionally, this encoder avoids making strong assumptions about the form of the posterior by producing samples from the (approximate) implicit posterior. However, MCMC usually requires a sufficient number of iterations to reach the stationary distribution (convergence), and determining whether it has converged can be challenging. Besides, manual tuning of many parameters adds complexity to the process.
The prior over the latent variables Z = [2'|[_ E N in this paper is from L independent GPs F [F]]__ 1-1 E NxL equipped with M inducing points and additive Gaussian noise. Specifically, the variables at the I-th channel have the following joint distribution:
GP Prior
1.
(2' p², u' SS (b) = (zo rr, (Dp((' ut (,,
(f ut [8,))= N ([
fl
Kxx
Kxsi
1) po(ed (f( =f()))=(()) Naa
u' 0,
Ksxie
Kss
where = |u
-- E MXL and S MXD are inducing variables and locations. vé contains the kernel parameters (c.g.. lengthscale, variance) and the noise variance o3. A fully Bayesian treatment in this work leads to a lognormal prior py (+) over si and a uniform prior (S) over S.
25

---

Neighbour-Driven GPVAEs for Scalable Structured Latent Modelling
Independent Conditionals The overall joint distribution should be decomposed observations sample from the posterior over all the latent variables using SGHMC. The authors further impose independence in the conditional distribution (Quinonero-Candela & Rasmussen, 2005):
( u',S, 16) N (r' | [ ss u', diag Kxxiv
Kxsjø K sS|v Ksx(e
Decoder neural network p(Y 0) with weights/biases regarded as random variables 3. Therefore, defining V {Vv U, 0, the potential energy function is given by
Σ log / p(z' | ())((( u S. ar logp(Y 2,
U Z) = log p(v: 1 U,
= (1og);(0) + log pe (S) + P S) + log
N (=)
s lix,, Kx,, sk. Kks. + + log an 0)
N
log (2/ Kx, lx. - Kx,s K + to,) + log plyn Zz. 0)
= log (
V
B
BEB
U),
where ($, 2) is stochastic estimate from a mini-batch. In practice, u' is sampled by whitening the prior, i.e., u' Lv,LL LL = Kss, v ~ (0, (Hensman et al., 2015).
Encoder is an implicit stochastic network concatenating random noise and Y. the encoder is a multilayer perceptron (MLP), the authors concatenate the random seeds and the input into long vector. The dimension of the random seeds is the same as that of the input. If the encoder a CNN, they spatially stack the random seeds and the input.
Sampling by SGHMC SGHMC (Chen al., 2014) uses noisy but unbiased estimation of the gradient VŪ(@) computed from mini-batch of the data, where group the sampled variables into @ ( Z}. Introducing auxiliary momentum variables r, the discretized Hamiltonian dynamics are then updated as follows:
A = -,M-Ir Ar =
N(0 -B))).
where is the step size, M an arbitrary mass matrix that serves as precondition, C user-defined friction ma- trix, and the estimate for the noise of the gradient evaluation. This work adopts an adaptive version of SGHMC (Springenberg et al., 2016), where these hyperparameters are automatically adjusted during a burn-in phase. After this period, these hyperparameters are fixed. Specifically, the updates of the hyperparameters are:
M-1
= diag V
7
Alo
- + 1-1
AT
--8 - 11
Age
- + -
= Invo:
C al.
Here, a is a momentum coefficient. By substituting = nn 2r. the discretized Hamiltonian dynamics becomes
Δθ = v, Av = av -v
II)
TTo remain notation consistency. The original paper uses &o to represent the decnder's 'ss parameters and 0 kernel parameters instead.
26


Reviews:
[{"review_id": "tM9qDAWF0O", "reviewer": "Reviewer_d16V", "summary": "Variational Autoencoders are deep generative models the learn low dimensional latent representations or high dimensional data, e.g. images in pixel space. When we have extra auxiliarly information, such as images from video and we also know the timestamps of the frames, instead of inferring independant latent representations for each image, we can use the timestamps to force structure of of the series of latent representations and the Gaussian Process Prior VAE is a class of model that uses the auxiliary infomration for inform a correlated GP prior over the corresponding latent representations of the high dimensional datapoints.\n\nHowever GPs have $O(N^3)$ computational cost where $N$ is the number of data points in the GP. As a result, recent works has proposed to use Sparse GPs, each prediction is made based an a small set of pseudo points reducing costs as compared to using the full dataset. This work proposes to use nearest neighbor GPs, each prediction is made based on using the nearest points in the dataset.\n\nThe authos", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for your valuable feedback. We provide a detailed response below.\n\n---\n### 1. Mathematical notations and clarification\n\nConceptually, our approach closely aligns with the NNGPs proposed by [1], where a hierarchical mechanism enforces local sparsity on inducing variables. To draw a parallel, our GPVAE-HPA can be seen as applying such an idea to latent NNGPs with each “inducing variable” placed at each data point (similar to [2]). We then incorporate a decoder network to model the \"likelihood\" from $Z$ to $Y$ and an encoder network for amortised inference on $Y$, thereby forming our GPVAE architecture.\n\nWe acknowledge that using $q(Z|Y)$ directly in $L_{HPA}$  (page 3) may cause confusion. We will replace it with $E_{p(w)}[q(Z|w,Y)]$, thus rewriting $L_{HPA}=E_{p(w)}[E_{q(Z|w,Y)}\\log p(Y|Z)-KL[q(Z|w)|p(Z|w)]]$, which leads to an ELBO similar to Eq(10) of [1]. Here, $p(w)$ is any implicit distribution from which we can sample. In our implementation, $p(w)$ reflects a nearest-neighbour selection strategy akin to Eq (11) of [1]. We will clarify this connection with $L_{HPA}$ in the main text and Appendix B.1 and correct the minor notational issues in Eq(8) as $q(Z|w)=\\mathcal{N}(Z|D_w\\mu(Y),D_w\\sigma^2(Y)D_w)$. We believe these revisions will both enhance the clarity of our derivation and highlight how GPVAE-HPA naturally builds upon established results for NNGPs.\n\nIn our setup, the prior $p(Z|w)$ and the approximate posterior $q(Z|Y,w)$ in the KL share the same degenerate dimensions, meaning those components are effectively “turned off” in both distributions. It is reasonable to define their \"KL contribution\" as zero. Thus, the overall KL is effectively computed only on the non-degenerate subspace, ensuring the KL remains well-defined.\n\n---\n### 2. Longitudinal VAE (LVAE) as baseline\n\nThe LVAE is also an inducing-point-based model, originally designed for longitudinal data with discrete “instance” IDs (e.g., a patient seen over multiple visits). This model embeds the instance ID into an additive kernel. Our experiments introduce *new* sequences at test time and do not explicitly provide IDs. Nevertheless, we have run additional experiments to address your request. We adapt LVAE by manually assigning unique IDs to each new sequence during testing. We presented the results of LVAEs with different numbers of inducing points $M$ from 10 random trials.\n\nThe following table shows the performance of LVAEs on the Rotated MNIST *Missing Frame* experiment:\n|Model/Metric|NLL|RMSE|Training time(s/epoch)|\n|-|-|-|-|\n|LVAE-M10|73.887$\\pm$0.748|5.578$\\pm$0.022|28.919$\\pm$0.253|\n|LVAE-M30|73.545$\\pm$1.035|5.558$\\pm$0.024|29.063$\\pm$0.122|\n|GPVAE-HPA-H10(ours)|**71.264$\\pm$1.263**|**5.521$\\pm$0.046**|**9.819$\\pm$0.064**|\n|GPVAE-SPA-H10(ours)|**69.538$\\pm$0.664**|**5.412$\\pm$0.026**|**9.358$\\pm$ 0.083**|\n\nThis table demonstrates that although LVAEs can get moderate results in terms of NLL and RMSE, our models consistently outperform LVAEs with fewer parameters and faster training.\n\nThe following table illustrates the performance of LVAEs on the MuJoCo experiment:\n|Model/Metric|NLL|RMSE|Training time(s/epoch)|\n|-|-|-|-|\n|LVAE-M10|-0.003$\\pm$0.316|0.175$\\pm$0.021|**35.564$\\pm$0.342**|\n|LVAE-M30|-0.814$\\pm$0.201|0.107 $\\pm$ 0.023|35.951$\\pm$0.443|\n|GPVAE-HPA-H10(ours) |**-2.335$\\pm$0.032**|**0.022$\\pm$0.001**|37.000$\\pm$3.793|\n|GPVAE-SPA-H10(ours) |**-1.715$\\pm$0.159**|**0.038$\\pm$0.007**|39.403$\\pm$2.755|\n\nThe table shows LVAEs exhibit substantially lower performance than our models with comparable training time. \n\nWe will clarify the model differences in the revised manuscript and include relevant results in Section 5 of our paper and the appendix.\n\n---\n### 3. Including GPVAE-Pearce\n\nWe included GPVAE-Pearce [3] as a *full GPVAE* baseline in Section 5.1. Specifically, we plot its performance with a dashed blue line in Figure 2(b) alongside GPVAE-Casale [4] to represent another “best-case” full-GP approach. Although both GPVAE-Pearce and GPVAE-Casale use fully correlated GP priors, they differ in how they construct their variational distributions—GPVAE-Pearce’s setup is more closely aligned with SVGPVAE [5], while GPVAE-Casale adopts a diagonal-encoder approach that matches ours. In our experiment, GPVAE-Casale slightly outperformed GPVAE-Pearce, so Figure 2(a) highlights GPVAE-Casale alone for clarity, whereas Figure 2(b) includes both baselines to show the overall trend.  We will clarify this point in the final revision to avoid confusion about whether GPVAE-Pearce was tested.\n\n\n[1] Sparse Within Sparse Gaussian Processes using Neighbor Information (ICML 2021)\n\n[2] Variational Nearest Neighbor Gaussian Process (ICML 2022)\n\n[3] The Gaussian Process Prior VAE for Interpretable Latent Dynamics from Pixels (AABI 2020)\n\n[4] Gaussian Process Prior Variational Autoencoders (NeurIPS 2018)\n\n[5] Scalable Gaussian Process Variational Autoencoders (AISTATS 2021)", "claims_and_evidence": "They authoprs claim their method has higher accuracy and quicker traiing times, I feel both of which are intuitive and justified by experiments.", "methods_and_evaluation_criteria": "## Hierarchical Prior Approximation ELBO Derivation\nI beleive I undertsand the intention of the approximation, however to me it seems the notation may have an error or two.\n- The prior Equation 6 and 7, $p(Z|\\textbf{w}) = \\mathcal{N}(Z|0, D_W K_{XX}D_W)$ suggests that a masking vector $\\textbf{w}\\in \\\\{0, 1\\\\}^N$ sets some of the prior covariance rows/cols to 0, effectively the corresponding dimensions of $\\textbf{Z}$ have a prior density that is a point mass at $z_i=0$, they become fixed constants and prior density is infinite everywhere (Note $det(D_wK_{XX}D_w)^{-1} = 0^{-1}=\\infty$).\n- (minor typo) Equation 8 gives the approximate posterior in Eq 8, assuming the masking vector $\\textbf{w}$ is known whould also have mean 0 for the masked elements, $q(Z|\\textbf{w}) = \\mathcal{N}(Z|D_W\\mu(Y), D_W K_{XX}D_W)$\n- on P3, $L_{HPA}$ contains the term $q(Z|Y)$ which is not defined, though it is stated to be the same as Equation 4, the traditional VAE, presumably the unmasked approximate posterior. I beleive this to be an minor error, marginalizing \n$\\int_\\textbf{w}q(Z|\\textbf{w})p(\\textbf{w})d\\textbf{w}$ yields a mixture of gaussians, a point mass and the prior, (see [1], page 3, left col, bottom).\n- on P3, $L_{HPA}$ contains the term $KL(q(Z|\\textbf{w})||p(Z|\\textbf{w}))$, if some dimensions of $Z$ are point masses then $p(Z|\\textbf{w})$ both denstiies are infinite everywhere.\n\n- Eq 10, the final ELBO $L_{HPA}$ seems intuitive and valid, the encoder should learn a good reconstruction as well as the neihbourhood should adhere to the local GP prior based on neighbour points. I believe a similar result can be acheived via a different root as the masking $\\textbf{w}$ seems to make things a bit messy in my view and I beleive it does not have the same effect as simply removing dimensions (which is the desired goal and matcheas with Eq 10).\n\nPresumably one could derive an ELBO f\n\n\n\n[1] [Sparse within Sparse Gaussian Processes using Neighbor Information, Tran et.al.](https://arxiv.org/pdf/2011.05041)", "theoretical_claims": "There are no theoretical proofs. See Methods for comment on deriving the ELBO $L_{HPA}$", "experimental_designs_or_analyses": "The paper contains many experiments from the literatire as well as some new ones.\n\n- page 5, section 5.1, it is stated that GP-VAE Pearce is included but I don't see it anywhere, presumably in such a setting, the full expensive model of Pearce 2019 is the best case scenario and provides an upper bound on how any GP-VAE could perform in such a small scale use case?\n\nI beleive the Longitudinal Variational Autoencoder should be added as a baseline for time series experiments [Longitudinal Variational Autoencoder Ramchandran et.al. AISTATS 2021](https://proceedings.mlr.press/v130/ramchandran21b.html)", "relation_to_broader_scientific_literature": "This work nicely builds upon GP-VAE, SGP-VAE and provides an intuitive step forward showing expected performance gains.", "essential_references_not_discussed": "[Longitudinal Variational Autoencoder Ramchandran et.al. AISTATS 2021](https://proceedings.mlr.press/v130/ramchandran21b.html)", "other_strengths_and_weaknesses": "Overall I am positive about the paper. It is intuitive and simple and clean. My main concerns are cleaning up some maths", "other_comments_or_suggestions": "Not at this time", "questions_for_authors": "- Can the derivation of Equation 10 be done by starting firectly from assuming we have a minibtach, and ignoring the hierachical approach?\n- Are Pearce 2019 included in the results? If not can they be added\n- Can the authors include the Longitudinal VAE where approriate?", "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": "", "extra_scores": {}}, {"review_id": "fVAEcH2QD3", "reviewer": "Reviewer_B5r9", "summary": "This paper successfully reduces the computational cost of Gaussian Process Variational Autoencoder (GPVAE) by incorporating the Nearest Neighbour Gaussian Process (NNGP). The experimental results demonstrate that the proposed method achieves high performance while reducing computational costs, compared to existing approaches aimed at making GPVAE more scalable.\n\n## update after rebuttal\nWhile I will keep my score as Weak Accept, I would like to note that my overall opinion leans toward acceptance of this paper.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your thoughtful feedback. We hope to address your concerns below.\n\n---\n### 1. Selecting neighbours in the data space vs. the latent space\n\n**A natural extension from NNGPs with local correlation**\n\nOur approach follows the principle of NNGPs, which typically rely on adjacency in the *input coordinates* $X$ (spatial, temporal, etc.) for large-scale GPs. Extending this principle to GPVAE, we assume that points close in $X$ are likely to have correlated latent representations under the GP prior, which generally holds well in practice. The GP prior *encourages* points that are near in $X$ to stay close in the latent space. This regularisation enforces correlation in the latent variables such that points with neighbouring inputs are unlikely to end up arbitrarily far apart in the latent space.\n\n**Possible extensions with complex kernels**\n\nWe recognise that data points far in data space can sometimes be mapped near in latent space, especially when the encoder is strongly non-linear. In those cases, one might extend the kernel to be non-stationary or incorporate periodic terms. This more flexible kernel design can learn that certain points—though far in raw input coordinates—are effectively neighbours according to the kernel values. In principle, one could then pick neighbours by kernel similarity rather than Euclidean distance. For the tasks we explored, standard kernels and data-space adjacency already perform well.\n\nIn summary, using data-space neighbours offers a straightforward, well-tested approach inspired by NNGPs.\n\n---\n### 2. Explaining performance gains over full GPVAE in some cases\n\n**Proposed method vs. GPVAE-Diag**\n\nA *full* GP prior can be a stronger regulariser, but it also complicates optimisation, making converging to a good solution harder. By focusing on local correlations rather than forcing a single large covariance across all data points, our method “loosens” the constraints on $Z$ and may achieve better results in some tasks. \n\n**Proposed method vs. GPVAE-Banded**\n\nBoth GPVAE-Diag [1] and our model share a standard encoder leading to a *diagonal* posterior (Eq(3)). By contrast, GPVAE-Banded [1] employs a *tridiagonal* encoder covariance, which introduces extra parameters (+~50%) to capture between-sample correlations more explicitly. Although this may yield slightly better performance on *short* sequences, it does not scale easily to longer sequences. Our design focuses on a diagonal encoder with a neighbour-based GP prior, achieving a more balanced trade-off between accuracy and computational feasibility.\n\n---\n### 3. Further discussion on some experimental results & hyperparameter guidelines\n\n**Short corrupted vs. long missing frames: different experimental settings**\n\n- The *Corrupted Frames* experiment (following [1]) features a large dataset (60,000 sequences) but *short* sequences (10 frames each). MGPVAE handles short sequences relatively efficiently because its sequential computations do not accumulate large overhead when the sequence length is small. Meanwhile, since the image transformations are relatively smooth from one frame to the next (no missing frames), MGPVAE’s Markov structure enjoys an advantage in capturing local frame-to-frame correlations.\n- By contrast, the *Missing Frames* experiment (following [2]) uses *longer* sequences (100 frames each), and each frame has a 60% chance of being dropped. This creates irregular, sometimes large gaps between observed frames, which presents two challenges for MGPVAE: (1) Given a longer sequence, the forward–backward process needs more runtime to propagate information; (2) The model may accumulate more error at large time gaps where many frames are missing in uneven patterns. Our neighbour-driven method does *not* rely on strictly sequential processes and can have faster training and better predictive performance in this setting.\n- HI-VAE can train faster as it does *not* learn inter-sample dependencies through a GP prior, but this may come at the cost of accuracy. \n\n**Guidelines for the number of neighbours $H$**\n\nThe number of nearest-neighbours $H$ directly influences how local or global the GP prior is. In practice, picking $H$ in the same (or smaller) range as the number of inducing points used by sparse-GP baselines is a good starting point: smaller $H$ reduces runtime but may weaken correlation modelling, while larger $H$ may improve accuracy at the expense of more computation. We recommend beginning with a moderate $H$ based on (1) estimated correlation length scales and (2) available computational budget, then adjusting based on validation performance.\n\n---\n### 4. Request for Pseudo-Code\n\nWe completely agree that providing pseudo-code would help readers reproduce our approach. We will include a succinct algorithm pseudo-code in the appendix of the final version. \n\n[1] GP-VAE: Deep Probabilistic Time Series Imputation (AISTATS 2020)\n\n[2] Markovian Gaussian Process Variational Autoencoders (ICML 2023)", "claims_and_evidence": "The claim that \"existing GPVAEs are computationally expensive\" is clear and well-supported by evidence. The proposed method, which addresses this issue using NNGP, is reasonable and well-motivated.", "methods_and_evaluation_criteria": "The paper evaluates the proposed method using various datasets and evaluation metrics.", "theoretical_claims": "I have reviewed the theoretical aspects of the paper, and they appear to be valid. However, I have one concern, which I will outline in the Questions Section.", "experimental_designs_or_analyses": "The paper presents extensive experiments on multiple datasets, with various discussions. I have a few concerns, which I will outline in the Questions Section.", "relation_to_broader_scientific_literature": "The references provided appear to be sufficient.", "essential_references_not_discussed": "I believe the necessary references are cited appropriately.", "other_strengths_and_weaknesses": "The idea of using NNGP to reduce the computational cost of GPVAE is simple yet effective.", "other_comments_or_suggestions": "Including a pseudo-code for the proposed method would make the paper easier to understand.", "questions_for_authors": "The use of NNGP to reduce the computational cost of GPVAE is a simple yet effective approach, and it appears to achieve high performance experimentally. I am generally in favor of accepting this paper, but I have the following concerns.\n\n1. If I understand correctly, the nearest neighbors are selected in the data space. However, since the Gaussian Process is applied to the latent variables, I find this slightly counterintuitive. Why did you choose to find neighbors in the data space rather than the latent space? Also, is the distance in the data space preserved in the latent variable space, especially in cases where the latent space exhibits structural transformations? (For example, two points that are far apart in the data space might be close in the latent space, and vice versa.)\n\n- If computational cost were not a concern, I would expect that GPVAE would achieve the best performance. However, in Figure 2, the proposed method slightly outperforms GPVAE in reconstruction error, and in Table 1, the NLL follows the order: GPVAE-Diag < Proposed Method < GPVAE-Banded. What is the reason for this?\n\n- While the proposed method achieves consistently high performance, some results are worth discussing. In Table 1 (Corrupted Frames), MGPVAE achieves better performance and lower computational cost than the proposed method, whereas in Missing Frames, the proposed method outperforms others. In Table 3, the proposed method slightly outperforms HI-VAE in performance, but HI-VAE is faster in training time. Given this situation, to what extent can hyperparameter tuning (i.e., increasing or decreasing the number of neighbors) improve the performance?\n\n- Since the number of neighbors is a hyperparameter that trades off between performance and computational cost, do you have any guidelines for setting this value?\n\nI would appreciate your response.", "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": "", "extra_scores": {}}, {"review_id": "uwBhkWy76w", "reviewer": "Reviewer_Van9", "summary": "The authors propose a new approximation for Gaussian Process Variational Autoencoders to improve computational efficiency. A GP-VAE is a VAE for the structured data case (e.g. where we have both images Y, and auxiliary information X, such as time, or location) that replaces the independent Gaussian prior over the latents with a GP prior (with the kernel operating over X, giving functions from e.g. time or location, to latents), better utilising correlations between datapoints. The authors use the nearest-neighbour kernel approximation from the GP literature, which only keeps the covariances from the set of H nearest neighbours for each points, leading to a block-sparse kernel matrix, which has similar computational complexity to an inducing-point GP. They propose and evaluate two different ways of combining nearest-neighbour GPs with the GP-VAE framework, and find that their model outperforms other approximate GP-VAE models with similar time complexity.\n\n### Update after rebuttals\nI have maintained my score of \"accept\" - please see rebuttal comment.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We genuinely appreciate your insightful and encouraging feedback and highlighting aspects of our paper that might have gone unnoticed. Below, we hope to address the points in the review.\n\n---\n### 1. Using auxiliary information $X$\n\nWe acknowledge that our non-GP baselines, VAE and HI-VAE, do *not* use the auxiliary information $X$ - yet they are widely adopted as baselines in the GPVAE literature (e.g., [1,2,3,4]) to underscore the benefits of imposing a GP prior. \n\nActually, as demonstrated by several GPVAE works (e.g., GPVAE-Casale [6],  SVGPVAE [2], SGPBAE [1]), providing low-dimensional $X$ directly as input to a standard VAE (e.g., in a *conditional VAE* fashion [5]) rarely outperforms GPVAE models on the similar tasks we explored in the paper. As you suggested, we also implemented a Conditional VAE (CVAE) baseline for our *Missing Frames* task, following the design in [6]. Specifically, we processed the scaled timestamp $t$ into $[\\sin(t), \\cos(t)]$ and used an MLP to produce an additional time feature, then stacked it to the vector after the encoder and before the decoder. The missing frames were generated based on average latent representations and new time features. Below are the results comparing CVAE to our models from 10 random trials:\n| Model/Metric| NLL| RMSE| Training Time(s/epoch)|\n|-|-|-|-|\n|CVAE|213.891$\\pm$13.121|10.276$\\pm$0.189| **4.736$\\pm$0.033**|\n|GPVAE-HPA-H10(ours)|**71.264$\\pm$1.263**|**5.521$\\pm$0.046**|9.819$\\pm$0.064|\n|GPVAE-SPA-H10(ours)|**69.538$\\pm$0.664**|**5.412$\\pm$0.026**|9.358$\\pm$0.083|\n\nAs shown, although CVAE trains faster, it performs significantly worse, aligning with previous findings that *directly* injecting low-dimensional $X$ into a standard VAE generally underperforms GPVAE. These results illustrate that the GP prior provides a more effective mechanism to correlate latent variables based on $X$.  \n\n---\n### 2. Clarify baseline differences across experiments\n\nWe omitted certain baselines for practical reasons:\n\n**VAE / HI-VAE in Missing Frames**\n\nIn the *Corrupted Frames* setting, partially corrupted images can still be passed to the encoder at test time. However, the *Missing Frames* experiment requires generating entire frames from the latent structure with the new timestamp $x_*$. Since VAE and HI-VAE do not incorporate $X$, they cannot directly generate fully missing frames. Hence, we excluded them in the *Missing Frames* baselines.\n\n**GPVAE-Diag / GPVAE-Banded for Longer Sequences**\n\nBoth GPVAE-Diag and GPVAE-Banded [3] were designed mainly for relatively *short* sequences. In our trials with 100-frame sequences, these models showed prohibitively long training times. Since the *Missing Frames* experiment was specifically designed to focus on scalable models suitable for longer and irregular sequences, we chose to exclude methods that do not scale in practice. In future revisions, we will clarify these reasons in the paper.\n\n---\n### 3. Typographical issue, code release, and paper structure\n\nWe appreciate you noting the boldface inconsistency in Table 3—this was indeed a formatting mistake, and we will correct it. Additionally, as stated in our appendix, we plan to release all source code for our experiments, including our re-implementations of numerous baselines in PyTorch for fairness and ease of comparison. We hope this effort could facilitate further research on GPVAE-based methods.\n\nWe agree that Table 15 in Appendix D is quite informative. We will try to include it in the main paper (subject to page limits). We also acknowledge your comments about Section 2.2 on inducing points. Although we included it for context (many large-scale GP methods rely on inducing points), we will streamline or move some details to the appendix to keep our main text more focused on neighbour-driven approaches.\n\n---\n### 4. Differences between HPA and SPA\n\nConcretely, HPA (Hierarchical Prior Approximation) imposes sparsity in the *covariance* matrix by “switching off” non-neighbouring latent variables through a hierarchical selection variable. This yields a sparse *covariance* structure. SPA (Sparse Precision Approximation) factorises the GP distribution into chained conditional terms, effectively resulting in a *sparse precision* matrix rather than a sparse covariance matrix. In practice, both are valid neighbour-driven approaches that reduce computational cost. We will further clarify these differences in the final manuscript.\n\n---\nWe hope these clarifications address the points you raised. \n\n[1] Fully Bayesian Autoencoders with Latent Sparse Gaussian Processes (ICML 2023)\n\n[2] Scalable Gaussian Process Variational Autoencoders (AISTATS 2021)\n\n[3] GP-VAE: Deep Probabilistic Time Series Imputation (AISTATS 2020)\n\n[4] The Gaussian Process Prior VAE for Interpretable Latent Dynamics from Pixels (AABI 2020)\n\n[5] Learning Structured Output Representation using Deep Conditional Generative Models (NeurIPS 2015)\n\n[6] Gaussian Process Prior Variational Autoencoders (NeurIPS 2018)", "claims_and_evidence": "- The authors' method indeed appears to improve task performance compared to other approximate GP-VAE methods.\n- For comparisons against non-GP VAE methods, the authors at one point state \"VAE and HI-VAE perform poorly as they do not leverage spatial information\"; it is not clear to me whether this means these methods were not provided X at all, which would obviously result in worse performance.", "methods_and_evaluation_criteria": "- The proposed solution of using nearest-neighbour GPs in a GP-VAE is very sensible, given the existing literature on nearest-neighbour GPs, and makes sense for the proposed VAE task. Given that inducing point GPs have already been tried in this domain, this paper fills in an obvious gap in the literature.\n- The authors test their method on a good range of tasks, all of which fall under the \"structured latent modelling\" umbrella.\n- The two proposed variants look quite similar at a high level, but sometimes result in moderately different performance. It would have been useful to explain why two different variants were proposed, compare and contrast the two methods in the paper after both have been presented. In particular, there are subtle differences between the ELBOs of the two methods that would be worth highlighting.", "theoretical_claims": "- I did not verify the authors' derivations for their method, which were consigned to the appendix. However, the resulting formulae looked sensible, and did seem to support the time-complexities stated by the authors.", "experimental_designs_or_analyses": "- The experiments are sensible, and cover a good variety of datasets.\n- The authors compare against many other baseline methods, which is good. It is not always clear why some baselines are only used in certain experiments but not others; for example, in Table 1, the baselines are different between the \"Corrupted Frames\" experiment and the \"Missing Frames\" experiment, and I'm not sure why (perhaps I missed an explanation in the paper).\n- All experiments report means and standard deviations over 10 random trials, including the baselines (which they ran themselves), which is excellent. Additionally, they also report the wall-clock time of all the experiments, which is important given that the point of the approximation is computational efficiency.", "relation_to_broader_scientific_literature": "- Nearest neighbour GPs were already proposed in earlier papers\n- GP-VAEs were already proposed in earlier papers\n- The novelty in this paper stems from combining nearest-neighbour GPs with GP-VAEs. Though this does not seem to have required a huge amount of novel theoretical work, it is an obvious gap in the literature, and this is a well written and executed attempt to fill this gap.", "essential_references_not_discussed": "The references provided seem suitable.", "other_strengths_and_weaknesses": "### Strengths\n- The paper is very well written. I particularly liked the fact that the authors made it very clear how different models related to each other, with sentences such as \"In contrast, setting H = 0 will cause the model to degenerate into VAEs.\"\n\n### Weaknesses\n- The proposed methods are not always that much better than other baselines. For example, in Table 3, the HI-VAE baseline performs as well as their method, despite not utilising the spatial information.", "other_comments_or_suggestions": "- The table comparing time complexities in Appendix D (Table 15) is very nice and would be useful to have somewhere in the main text, if space allows.\n- Section 2.2 (in the background section) talks about inducing point GPVAEs, but as I understood it, we were not using this framework in this paper. If this is not necessary background to understand this paper's methods, it may be best to remove it to avoid adding irrelevant and confusing details.", "questions_for_authors": "- Why are the baselines different between the two experiments in Table 1?\n- How are you performing bolding in the tables? In Table 3, you have bolded HI-VAE's training time, though this is not the lowest training time on that table.\n- Do the non-GP VAE baselines have access to X at all, e.g. as arguments to the encoder / decoder?", "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": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposes a method for improving the computational efficiency of Gaussian Process Variational Autoencoders (GP-VAEs) by incorporating the Nearest Neighbour Gaussian Process (NNGP) approach. The reviewers praised the paper for its clarity and the novelty of combining NNGPs with GP-VAEs, noting that it fills an obvious gap in the literature. The main criticisms centered around methodological choices, the selection of baselines, and the interpretation of results, with reviewers seeking further justification and clarification on these aspects.

The authors responded by addressing the concerns raised, providing additional explanations and justifications for their approach. The reviewers indicated that the authors' rebuttal has addressed their concerns, particularly regarding the methodological choices and experimental design, leading to a more positive assessment of the paper. Therefore, the reviewers unanimously agreed to accept the paper. We would still recommend that the authors take the reviewers' feedback into account when preparing the camera-ready version.
Decision: Accept (poster)