Latent mixed-effect models for high-dimensional longitudinal data
Authors: Priscilla Ong
Manuel Haussmann
Otto Lönnroth
Harri Lähdesmäki
Venue: Accepted by TMLR ()

Abstract: Modelling longitudinal data is an important yet challenging task. These datasets can be high-dimensional, contain non-linear effects and feature time-varying covariates. Gaussian process (GP) prior-based variational autoencoders (VAEs) have emerged as a promising  approach due to their ability to model time-series data. However, they are costly to train and struggle to fully exploit the rich covariates characteristic of longitudinal data, making them difficult for practitioners to use effectively. In this work, we leverage linear mixed models (LMMs) and amortized variational inference to provide conditional priors for VAEs, and propose LMM-VAE, a scalable, interpretable and  identifiable model. We highlight theoretical connections between it and GP-based techniques, providing a unified framework for this class of methods. Our proposal performs competitively compared to existing approaches across simulated and real-world datasets.

Published in Transactions on Machine Learning Research (05/2025)
Latent mixed-effect models for high-dimensional longitudinal data
Priscilla Ong
Saneiila.onggaga.ggoo
Deportment of Computer Science
Aalto University
Manuel Haußmann"
haussmammmmmaaa..see dk
Department of Mathemalics and Compuler Science University of Southern Denmark
Otto Lönnroth
otto. otto.lonnroth@oaalto.
Department of Computer Science
Aalto University
Harri Lähdesmäki
horri harrillahdemmaki@oaltt,
Deportment of Computer Science
Aallo Universily
Reviewed on OpenReview: https: // openrevie.. net/ forum? id= 7A96yteeF3
Abstract
Modelling longitudinal data is important. yet challenging task. These datasets be high-dimensional, contain non-linear effects and feature time-varying covariates. Gaussian process (GP) prior-based variational autoencoders (VAEs) have emerged promising approach due to their ability to model time-series data. However, they are costly to train and struggle to fully exploit the rich covariates characteristic of longitudinal data, making them difficult for practitioners to use effectively. In this work, we leverage linear mixed models (LMMs) and amortized variational inference to provide conditional priors for VAEs, and propose LMM-VAN a scalable, interpretable and identifiable model. We highlight theoretical connections between it and GP-based techniques, providing unified framework for this class of methods. Our proposal performs competitively compared to existing approaches across simulated and real-world datasets.
1 Introduction
Longitudinal datasets, typically containing repeated measurements of individuals over time, have applications in numerous fields such as education, psychology, the social sciences, and the biomedical field. Longitudinal study designs particularly useful for revealing associations between explanatory covariates and response variable, such as the relationship between risk factors and discase progression (Diggle 2002). but require appropriate statistical tools that can account for correlations both within subjects as well as across subjects. Analysis of univariate or low-dimensional longitudinal data dominated by various linear mixed models (LMMs), also known as multi-level models, along with other additive modelling approaches. However, existing methods scale poorly to relevant longitudinal datasets, such as electronic health records, which are often high-dimensional (Zipunnikov et al. 2014) and contain non-linear effects, time-varying covariates, and missing values RRanchadrann al. 2021).
Variational autocncoders (VAEs) are a class of models commonly used for representation lcarning and generative modelling [Rezende et al. 2014 Kingma & Welling 2014). Nevertheless, they cannot be directly
Work done the Department of Computer Science, Aalto University.

---

Published in Transactions on Machine Learning Research (05/2025)
Latent Space
3
3
3
Encoder
Decoder
3
Time-series data
Time-series
reconstruction
Lincar Mixed Model (LMM)
shared and random effects -non-lincar basis functions
Overview of LMM- VAE. The latent space is modulated by auxiliary covariates parametrized by a linear mixed model (LMM).
Figure 1:
applied to longitudinal data as they assume that observations are independent and identically distributed (ii.d), thereby failing to capture correlations between samples. This limitation has been partially addressed by recently proposed Gaussian process (GP) prior VAEs that use auxiliary covariates to impose correlation structure on the samples' latent representations. However, such models have cubical complexity with respect to (w.r.t) the number of samples, necessitating model simplifications (Fortuin et al. 2020) or specific GP approximations to be applied as part of the VAE model fitting (Casale et al. 2018 Ramchandran et al. 2021). Another challenge concerns fully exploiting the rich auxiliary covariates available for modelling. While conditional VAËs (CVAEs) (Sohn et al. 2015) can easily incorporate any number of covariates, limited work has tackled the problem of finding an appropriate time-series or longit udinal model that can effectively scale to include large number of covariates into the VAE prior. Given the potential upside in model performance, addressing the model's 's scalability, specifically to include more covariates, is a fruitful avenue of rescarch.
Chosing specific parametrization for the prior is a model selection problem. Besides using a sufficiently rich model class, including appropriate covariates is integral |George 2000) for modelling the latent space effectively. Different from previous work, we propose modelling the prior using a linear mixed model with appropriately chosen basis functions. While simple, this model class is scalable, vastly simplifies the training procedure within standard deep learning frameworks, and enjoys several advantages discussed below.
Contributions. We propose the Linear Mized Model VAE (LMM-VAB) which capable of handling high-dimensional data and large dataset sizes, modelling an arbitrary number of covariates in the prior, can be adapted to problems of different complexity via basis functions, and enjoys the advantages of being interpretable and identifiable by way of parametrization. We demonstrate that LMM-VAEs are competitive against commonly used GP-based mods for integrating auxiliary covariates into the prior. We also highlight. theoretical connections between our proposed model and GP prior VAËs, bridging the gap between our work and existing methods. From both a practical and theoretical staudpoint, LMMs show promise as an alternative VAE prior, especially in the presence of high dimensional covariates. Section compares previous methods with proposed LMM-VAE. See Figure 1 for visual summary.
2
Related Work
The literature extending VAEs' modelling capacity is vast. It spans from extending the expressivity of variational posterior distributions, c.g., using multivariate Gaussian with tridiagonal inverse covariance structure (Fortuin et al. 2020) or normalizing flows (Rzende & Mohamed, 2015). lo enabling more informative priors compared to the standard normal used in plain VAEs (Kingma & Welling 2019). Alternative research efforts focus on endowing the model with desirable characteristics, such as a disentangled latent space (Higgins)
et al. 2016 Zhao al. 2019). Our work enhances the latent space representation by incorporating auxiliary covariates into the prior, such time-varying side-profile information.
2

---

Published in Transactions on Machine Learning Research (05/2025)
Table 1: Comparison related methods.
MODEL
PRIOR
COVARIATES
MINIBATCHING
IDENTIFIABLE
REFERENCES
CVAE
(LI.D
GAUSSIAN
ARBITRARY
1 UNKNOWN UNKNOWN UNKNOWN UNKNOWN
X
SOHN ET AL.
2015
CPP-VAE
CP
LIMITED
PSEUDO
CASALE ET AL.
2018
SVGP-VAE
GP
LIMITED
V
JAZREC ET Al..
222T
IVAR
GP
ARBITRARY
V
RAMCHANDRAN FT AL..
2021
LONGITUDINAL VAE
LM
LIMITED
V
SAUTY & DURRLEMAN
2022
LMM-VAE
LMM (OR GP)
ARBITRARY
OUR WORK
Structured Variational Autoencoder (SVAE). We build upon the ideas of Johnson et al. 2016)) who combine neural networks and probabilistic graphical models to provide structured latent representations. LMM-VAE's enhancements are designed to induce structure using arbitrarily many covariates characteristic of longitudinal data. Incorporating such structure is important in many real-world applications, e.g., Sauty & Durrleman ) characterize the progression of Alzheimer's disease via a mixed-effect longitudinal model in the VAE setting. However, their model's parametric structure is confined to disease progression and requires an internal procodure of Markov chain Monte Carlo sampling. LMM-VAE, in contrast, provides a generalizod rendition of this approach.
Gaussian Process Variational Autoencoder. The family of GP prior VAEs is most relevant as related work in modelling the prior. Within the GP-VAE framework, there have been multiple developments (Casale
ct al. 2018 Ashman et al. 2020 Ramchandran et al. 2021 Jazbec et al. 2021 Zhu et al. 2023) ). where the
GP's expressiveness and smoothness are leveraged to enable flexible, yet robust VAE priors. We focus on those GP-VAEs that are most compatible with modelling longitudinal data, i.e. repeated measurements with auxiliary covariates, including individual-specifie information.
While GP-based priors (Casale et al., 2018: Jazbec et al. 2021 Ramchandran et al. 2021 can model longitudinal data, using a GP model component comes with host of difficulties. Firstly, learning scales cubically with respect to the number of samples (Rasmussen & Williams, 2006), which limits the model's scalability, Efforts to reduce training complexity include approximating the GP priors via Taylor approximation
(Casale al. 2018) or through inducing points (Ramchandran et al.] 2021 Jazbec et al. 2021 p. For the latter, optimizing the inducing point locations is challenging (Bauer al. 2016 and may complicate training
due to the coupled learning of both the latent variables and the inducing points (Titsias 2009 Hensman et al. 2013). Categorical inputs, which are typically modelled in longitudinal set-ups, may exacerbate this issue as they are not amenable to gradient-based optimization. In addition, unless appropriately designed and implemented, these approximations may result in diminished expressiveness of the GP-priors.
Secondly, modelling all available covariates with a GP is non-trivial. Jazbec et al. (2021) 's approach may be ill-suited for this task it assumes low-dimnnsional covariates. Casall al. (2018) posit that the auxiliary information can be represented by an object and view kernel, where the da comprises of objects in different views. However, constructing these kernels from high dimensional side-profile information consisting of continuous and categorical covariates remains unclear. Meanwhile, [Ramchandran et al. 2021) propose using additive kernels including all covariates, with each component or pair implementing a specific kernel. However, performance gains may be limited by the challenges of training the GPs as a VAE prior.
Lincar Models. For certain modelling tasks, lincar models (LM) (Galton 1886 Fisher 1925) are a practical and competitive alternative highly parametrized models. For low-dimensional data, LMs and LMMs have become the standard tool of analyzing various biomedical and longitudinal data (Mundo et al, 2022h. LMs are also applied in the context of high-dimensional data, such as when applied to structural and functional brain imaging data (Schulz et al. 2020) or when performing dimensionality reduction using principal component. analysis (PCA) or its generative counterpart, probabilistic PCA (PPCA).
However, the assumption of a strictly lincar relationship between covariates and response variables may be overly restrictive in certain contexts. Instead of enforcing lincarity throughout the model, we impose this constraint only within the latent space while allowing for more complex, nonlinear mappings between the
3

---

Published in Transactions on Machine Learning Research (05/2025)
latent space and the response variables. Moreover, relax the strict lincarity in latent space via the basis function expansion. By combining the advantages of LMMs with the flexibility of neural networks, LMM-VAE extends the traditional LMM framework, providing a more powerful and adaptable approach for modelling complex data.
3 Background
Linear Models. Consider a pair (x, where æ = (#1,...,,xq) xT)" = 2, is Q-dimensional covariate vector and 2 R_ is I-dimensional response variable. The linear model (I.M) for (æ, 2) is
a1 + +.... aqio + Ax + +,,
(1)
where a, E R4, A (aI,..., aQ) RLXQ and e OO. assuming equal variance. For N pairs of covariates and response variables, {(Em: 2n zn)
| n= the lincar model is given as
Z = AX + E,
where Z = (Z1,..., 2n), X =
m.) and E (c1, ex) such that
N
(Z] A II ((z n=1
(2)
Without loss of generality (w.loo.g.); we assume that each covariate is either continuous binary since categorical covariates can be hot encoded into binary values. To model non-linear effects in Z, we follow common practice and extend the linear model with non-linear basis functions o(-) for continuous covariates, x' (<x), such as x' =2² x'= = sin(x), x' = cos(x), etc. We elaborate on this in Section 6
Longitudinal datasets consist of repeated measurements of instances (e.g. patients) over time and are commonly modelled using linear mixed models (LMM) (Laird & Ware 1982). While standard LMs are designed model effects that are shared across all instances (fixed effects), M. can simultaneously model both fixed and random effects. Random effects reprosent individual- or group-specific variations, accounting for differences across individuals groups, whereas fixed effects have consistent influence across all observations. Modelling both types of effects captures individual- and group-level variations while maintaining generalizability across observations, resulting more structured and meaningful latent space representations. W.l.o.g, we assume that the covariates are ordered as
While
x [1 Tli... FS, IS41.... TS
where = S RG with the first S covariates modelling shared effects, while the remaining R model random effects. For example, 2s may include the age gender of individual that we would like to model shared effect. Similarly, (patients), which can be used to model e.g. instance-specific random offset terms. Covariates æn may also include other types of variables, such as interaction terms specifying random effects for arbitrary subgroups of individuals, We write the LMM as
may include binary covariates corresponding to the identity of all instances
(a)...., ag) Ug (as+I aS+R) OR+6 +
Au
Asag
AR
+ c= = (As.AAR) +
shared effects
random effects
Thus, we use the same compact notation shown in Eqs. equation I and equation 2 for both LMs and LMMs.
Thus far, our current formulation does not fully align with conventions commonly used in the LMM literature. To improve clarity and align more closely with conventions in the LMM literature, we include simple example in Appendix A that formalizes the problem in the context of LMM- VAE.
Identifiabilit depends the CVAF parametrization.
4

---

Published in Transactions on Machine Learning Research (05/2025)
In
Pm
A
yn
Figure 2: L.MM-VAE, plate diagram with probabilistic priors on all parameters. Shaded and blank circles refer observed and latent variables, respectively. For more detailed illustration of both fixed and random effects, refer to Appendix B
Variational Autoencoders. We assume observations y E =y=R² and latent variables E TR,, where L D. Given a datasct Y (y1... .y.......M. containing observations, we assume that Y generated by latent. variables Z (21,..., ,N) and write the joint. generative model for single observation (y. z) pe (y) where = {0, 40. For vanilla latent variable models, is typically assumed to have the standard normal prior, i.c., ~ (O(, I). where I is the L-by-L identity matrix. We are typically interested in inferring the posterior Pw (=ly), which is generally intractable due to the need for computing the evidence p(y) = (p)))))
VAEs (Rezende Mohamed 2015 Kingma Welling 2014 rely on amortized variational inference, by specifying parameterized approximation qu (2)g)) known as the encoder, to the intractable true posterior (zly). Its parameters optimized jointly with w by maximizing evidence lower bound (ELBO)
logp=(Y)> 2 Egg [log pa (1 - KL (go (go( (ZlY) (z)),,
where KT denotes the Kullback-Leibler divergence. Conditional VAEs (CVAEs) (Shnn et al. 2015 additionally condition the model on additional side information, i.e,, auxiliary covariates, E X. CVALs typically assume (y, = x) extended similarly by augmenting by the covariates, ge (=ly,,,))) x). Their ELBO objective is obtained by conditioning the probabilitics of the generative model and the encoder with covariates X = (x1,.... ÆN).
either z(ar) = po ())
The encoder of CVAE can be
GP prior VAEs also belong to the family of CVAEs and assume a joint prior pw =|x) = po (2)=)) where P (x) is a GP.
4 Linear Mixed Model VAEs
We propose structured VAE by incorporating LMMs into the prior of the VAE, and propose two variants of the model, which we collectively refer to as LMM-VAE, This model can accommodate an arbitrary number of auxiliary covariates in the prior, model both shared and random effects, and allow efficient model learning via the global parametrization.
4.1 Ordinary SMM--VEE
The first model variant, which we call ordinary LMM-VAE (OLMM-VAE), assumes that a standard LMM is used to model the latent variable. Assuming a prior on matrix A and an optionally probabilistic decoder, we formulate the generative model for sample y with covariates as
0 ~ p(8)
(3)
A ~ p(A)
(4)
z 2 x ~ N(z)An, T
(5)
y|2,0 0 ~ p(y|z,0), 0),
(6)
where 8 parameterizes a decoder fo y as well the Gaussian (or other) likelihood model. If needed, can define different priors for shared and random effects, p(A) = p( p(As)p(AR).
5

---

Published in Transactions on Machine Learning Research (05/2025)
We assume a factorizing variational posterior (o 0, Z = II, gz(2n (x)g()))(()), where all approximate posteriors are Gaussian, and & denotes an encoder network that maps individual observation Yn into parameters of the Gaussians, mean pn and diagonal covariance o The ELBO is given as
log logp(Y(X) 2 (F (og Dogp(yn/
(A)A) [KL (go vn) p(z, 2))))
KL (g(4) || p(4)) KL (0()) || p(0))
It is straightforward optimize this ELBO with mini- sase-bbtthh--asee stochastic gradient descent (SGD) as the parametrization of LMM-VAE model is global.
The first. expectation can be approximated via Monte Carlo sampling, while the remaining terms are analytically tractable assuming Gaussian priors.
For the n-th observation, we write the LMM prior mean for an arbitrary dimension I of z as 2 = a + e, where ay denoles the 1-th row of A. Using the independence of the latent dimensions in both prior and posterior, the second term of the ELBO becomes
E ((A) [KL (4o (zly) ll p(z[A,
Σ By(a) KL (20(21))) |l (()((.)))))
(m axx) 2
1 of o²
1 log of o²
Bg(a)
20²
where μa (Mi,..., ur.)² and o² (o ,ot oZ)" are the parameters of go (zlg). We simplify the numerator of the first term further
IE [M4 - ) (μ E
[ T x)² + var(a.)) (x x),
where o is the Hadamard product, and the expectations with respect to ((A) are analytically tractable.
We could alternatively treat the generative model parameters 0 and as deterministic hyperparameters, which case the LMM-VAE model reduces to Equations (5) and (6. The reconstruction and regularization terms in the ELBO then simplify to (=.llv)) [log p(y, En, 8) and KL (9.(2x,l))) p(=,, A.,,), respectively.
4.2 Structured LMM-VAE
In structured LMM-VAE (SLMM-VAE), impose further structure and inductive bias into LMM-VAE for datasets that do not require observation specific variation in the latent space. Instead of having N I for each E (1,..., N), assume that variation in the latent space is shared for observations arising from the same group. Suppose that the N observations are divided into K groups. Throughout this work, we assume that cach group consists of observations from a single instance (i.c., a patient), but groups could be defined more generally. W.l.ogg, we assume that the N observations are ordered such that the first nn observations, 4 = [L,...,n.]) belong to the first instance, the next n2 observations, I2 {ni + 1,..., nh n2), to the second instance, and so on. We define the generative model for K groups consisting of Σk n = observations, as
p(0).
A p(A),
(k) N(0, d.I), for k = 1,..., K,
Z|A,X AX + E.,
(7)
Y|Z,0 II p 12.. 12. 8),
(8)

---

Published in Transactions on Machine Learning Research (05/2025)
where
E" = (E (1) ,
e (K)
(K)
Th times
TA times
(Ei,.... Ex) E 1
Note that Eq. oquation 7 specifics that all samples from the k-th group share the same latent variation c(A), L.e. Zn Am, + ((k for all la- The encoder for the amortized variational inference needs to be designed such that the shared latent space variation is accounted for.
4.3 SLMM-VAE encoder architecture
Let & denote the encoder network for the LLMM-VAE model from Section 1.1 that maps individual observation yn into parameters of the variational approximation, fn and o By allowing the generative model to share the linear model parameters A with the encoder, we define amortized variational parameters 4 (h) and or (8)2 for Gaussian approximation (Y,
A, as follows:
1 (R)
1
Σ Axn
Pn.
The
nET
(10)
(k)
Σ (p - (Ax pn))²
nET
where Mn is obtained from yo with vi, and (1) is computed element-wise. By this formulation, all observations belonging to the 4thh group would share the same amortized variational parameters, thereby enforcing more structured representations at the individual level.
4.4 Predictive training for LMM-VAE
As last variation, we present a model-agnostic training routine that could yicld improved predictions for conditional generalion. First introduced by (Sohn et al. (2015), the Gaussian stochaslic neural network (GSNN) relies on training a predictive model, where the variational approximation set to be identical to
the prior, Le. (z|x,, y) = p(z(2r). The prior is then directly trained for the task of data generation, and the ELBO loss for the GSNN model is given as
1
(OSNN(Y|X,AL
log (ynn/
(11)
S ==1
where
Zn,s Axnn + Cn,s
(12)
En,s ~((
(13)
where use S Monte Carlo samples to estimate the log- likeiiieioood. GSNN-based training objective for the SAE model variant can be defined similarly. When the auxiliary covariates are sufficiently informative, using GSNN-based training could yield competitive performance for data generation. See Appendix D for further discussion of the GSNN loss.
5 Identifiability
Drawing Khemakhem et al. (2020)'s 's work, LMM-VAF's design conditional latent variable model guarantecs it the property of identifiability. Khemakhem et al. (2020) require the following structure
(r(
where (z)x(z)x) belongs to the exponential family,
p2) II
my m; exp
Tre (=)))is(a))
(14)

---

Published in Transactions on Machine Learning Research (05/2025)
with being the base measure, T the sufficient statistics, and ((x) the corresponding parameters. ((y(z) needs to be decomposable into y (=) + with ~ p(c), and I. < D. Additionally, we require æ to be sufficiently diverse, such that it consists of at least KL + 1 distinct values, where k is the number of sufficient statistics, an injective function f, as well as additional technical conditions summarized in Appendix the model identifiable in (S, U,T.A.A up to rolations and translations. Given our choice of priors in equation and equation X and assuming an injective decoder, both our oLMM-VAE as well as our sLMM-VAE, variation fulfill the constraints. See Appendix G for detailed discussion and results on the Project Data Sphere dataset.
G Then,
Motivation. Identifiability ensures that the true joint distribution over observed and latent variables, to rotations and translations, can be estimated (Khemakhem et al. 2020p. In other words, identifiability implies consistency of parameter estimation, up to certain transformations. Although an asymptotic theoretical guarantee, identifiability
desirable property because enables a meaningful interpretation of model coefficients and the data generation mechanism. In real-world applications utilizing longitudinal datasets, downstream analyses following model fitting would allow practitioners to gather insights and understand how model predictions are generated, based the estimated parameters. These properties useful because many real-world applications of longitudinal data seek to address questions that extend beyond mere model fitting.
6 GP prior VAEs as LMM-VAES
As noted in Section 3 LMs and LMMs can be made arbitrarily complex with additional basis functions, while maintaining lincarity with respect to the paramcters. This basis function extension highlights a usefuul connection between LMM-VAEs and GP prior VAEs. This connection builds upon a well-known result. that
GPs correspond to Bayesian linear regression with infinitely many basis functions (see e.g. (Rasmussen & Williams 2006)). ). Here, we use the spectral domain representation, but similar constructions could be derived for other basis, such eigenfunctions from the Mercer's theorem (Rasmussen & Williams 2006 or Laplace eigenfunctions (Solin & Särkkä 2019).
Based on the spectral domain representation, a univariate stationary covariance function k(r), where r = x x-2', can be approximated using a finite basis function expansion (Rasmussen & Williams
2006
0o des
M
A(r)= = 2
e(w)
cos(wmr)) = k(r),
M
where s(w) is the kernel's spectral density at frequency w, the imaginary number, σ² the kernel variance, M the number of terms in the approximation, and wm are either regular or random Fourier frequencies (with distribution proportional to s(w)). To construct the basis functions, we the trigonometric identity cos(u- cos(u) cos(v) + sin(u) sin(v) (Tompkins & Ramos 2018), and represent k(r) using the feature map
(xx) = (cos(w)z)...... (os((/ sin(wix)....,,inn
Following
Hensman et al. 2018 the approximate GP is then given as
f(x) ~ GP 00, 2 ((1)) gP (0, k(z - x'))
with an equivalent parametric expression
f(x) = ap(x),
where a is row vector as before, p(aa))= -N, S) with S = diag(s(wr)..... s(wM)) for regular Fourier features, and p(a")= = N (0, // for random Fourier features.
8

---

Published in Transactions on Machine Learning Research (05/2025)
Consider GP prior VAE with stationary covariance conditioned with single continuous covariate Existing GP-VAE models assume the prior for the latent space factorizes across the dimensions, i.e., p(=|r) = II (P(O, kx x')). For the reduced-rank approsimation of a GP prior with M Fourier features, the GP prior VAE can be directly implemented by LMM-VAE by setting 2 = AΦ(x), where
A (a :=........ at))
p(A) p(a/) defined above, and letting o + O. A similar connection between the LVAE Ramchandran et al. 2021) and LMM-VAE, can also be established (see Appendix
While connected to the GP-VAEs, LMM-VAE simplifies the overall training procedure. It achieves this via the spectral representation and global parameterization allowing for straightforward SGD optimization.
7 Experiments
We evaluate LMM- VAE on three tasks: (1) missing value imputation, (2) future prediction, and (3) time-based interpolation, measuring its performance using the Mean Squared Error (MSE). Our primary comparison is with the GP prior VAEs, which are most similar to our approach in modelling the prior through regression techniques. For consistency, we employed similar encoder and decoder architectures across all methods (see Appendix M for further details).
We rely on Health MNIST and Rotating MNIST, which are two datasets derived from MNIST LeCun
et al. 2010), and the Physionet Challenge 2012 dataset (Silva et al. 2012). See Appendices H and K for hyperparameters and further details on the experimental set-up not discussed in the main paper. Results in bold are within one standard deviation of the best mean per experimental sct-up. As a consistency check, also compared the performance of LMM and LMM-VAE simplified, low-dimensional set-up, which is presented in Appendix P
7.1 Health MNIST
Following prior work (Ramchandran et al. 2021), ; we generate a longitudinal dataset with missing pixels by augmenting the MNIST dataset. This dataset replicates various characteristies present in real medical data, where each snapshot of the time series corresponds to the health state of patient. The trajectory is designed to evolve non-linearly with time. We select the digits '3' and '6' to represent two biological sexes. There are Q = 6 covariates describing the dataset, which are age, id, diseasePresence,, diseaseAge, sex, and location. Further details on the dataset and train-validation-tess splits are in Appendix H.1
By design, GPP-VAE and SVGP- VAE are parametrizod by two kernels describing the object and view. As such, we train these two models using the id and age covariates per respective kernel. We parametrize LVAE using the optimal set-up as reported in their paper, ie. id, age, sexxage, and
RRanchanraan et al. 2021). We train three LMM-VAE variants to demonstrate different paramoterizations, summarizing the covariates for each configuration in Figure 3 Note that the color-coding in Table 2 follows this mapping.
Missing Value Imputation. We report missing value imputation performance based on the training set. in Table 2 LMM-VAE's imputation MSE remains robust across the different linear model parameterizations and outperforms the Ga-baaedd baselines.
Future Prediction. We further evaluate our proposed model by its ability to generate new data, based the test MSEs for conditional generation per model configuration. The predicted trajectories of two instances are illustrated in Appendix H.3
By design, the data generation mechanism is a complex function of several key covariates. As shown in Figure 3 LMM-VAE's ability to model these covariates leads to a lower conditional test MSE comparod GPP- VAE and SVGP-VAE In addition, LMM-VAE's sensitivity to different prior parametrizations highlights the importance of including relevant covariates that could reveal information about the data generation
9

---

Published in Transactions on Machine Learning Research (05/2025)
Table 2: Imputation MSEs on Health MNIST.
MODEL (dim(z) 32)
MSE +
GPP-VAE (CASALE ET AL..
2018
0.021
SVGP-VAL [JABBEC ET AL.
20211
0.015=t0.0012
LVAE [RAMCHANDRAN ET AL.
2021
0.018=0.0006
OAM,-VAE, VI (Ouns)
0.002:0.0000
OLMM- VI (Ouns)
1.002::0.000
OLMM-VAE, VI (OURS)
0.002 0000
0.035
0.030
0.025
0.020
0.015
0.010
fid. age, sex age, locat on age)
0.005
rid. age, sex age, diseasePresence diseaseAge)
age, age,
age, dseasePresence diseaseAge)
0.000
GPRVAE
SVGINAAE
LVAE
(Cassle
(aabbe (Ra macandrann
Latent dimension
2018)
20211
2021;
Figure Predictive test MSEs on the Health MNIST datasot with latent dimension dim(z) 32.
process. Compared to IVAE, Figure 3 inddicates that LMM-VAE reports competitive MSE when the covariates used for training are identical Interestingly, the flexibility provided by GP-prior may not translate inlo significant performance gain.
7.2 Rotating MNIST
The Rotating MNIST is another popular dataset used to benchmark CVAE models (Casale et al. 2018 Jazbec et al. 2021 Ramchandran et al. 2021 This dataset is composed of 400 unique instances of the digit. 3' rotated through 16 angles uniformly distributed in [0, 2m). Our treatment of the train-validation--oss splits
differs from that originally featured in (Casale et al. 2018 Jazbec et al.) 2021 which we elaborate on in Appendix
Interpolation. In this section, we evaluate the models by their ability to interpolate four consecutive angles that are unseen for test subjects during training. The data generation mechanism depends only on the digit instance and angular information, which can be easily modelled by the GP priors. Consequently, GP-based VAEs are competitive baselines. Due to the deterministic nature of the data generation process, GSNN training yiclds superior performance. Still, sLMM-VAE also able to match, and improve upon, the performance of its GP counterparts (see Table 3)
In their basic form, LMMs may have limited modelling capacity. However, as discussed in Section 0 we can make LMM- VAE arbitrarily complex to suit the modelling task at hand. Using basis functions defined by regular Fourier frequencics (sin and cos with increasing Trequencies wm) for both LMM-VAE variants, we show in Figure that test MSE for Rotating MNIST decreases or remains stable as the number of basis functions increases. This aligns with the theoretical considerations in Section 6 that adding basis functions LMM-VAE makes it strong, efficient approximation for GP prior VAEs, while preserving global parametrization and efficient training.
Improving performance with basis functions.
2LVA contains additional interaction between id and age; see Ramchandran
2021
10

---

Published in Transactions on Machine Learning Research (05/2025)
Table 3: Interpolation test MSEs on Rotating MNIST.
MODEL (dim(z) 16)
MSE +
CVAE (SOHN ET AL. 2015
0.036=000003
GPP-VAE CASALE ET AL.
2018
0.017=0.0005
SVGP-VAL VAE [JAZBEC AL.
2021
0.018=0.0000
IVAE (RAMCHANDRAN ET AL.
2021
0.026:00.0007
LLMM-VAE, VI (Ouns)
0.027 20.0003
SLMM-VAE, VI (OURS)
0.017#0.0000
OLMM-VAE, GSNN (OURS)
0.014 0.014:00. 0.000 0003
SLMM- VAE, GSNN (OURS)
0.014 0.0004
0.045
0.020
OLMM-VAE
0.018
test MSE
0.040
0.016
0.035
0.014
0.030
0.012
SLMM-VAE
0.025
0.010
2
4
12
2
8
12
number of basis functions
Figure 4: The test MSE decreases as the number of included basis functions increases (dim(z) = 16).
7.3 Physionet Challenge 2012
Future Prediction. We also evaluate LMM-VAE on real-world medical data from the Physionet Challenge 2012 (Silva et al. 2012). The full dataset contains approximately 12000 patients monitored in the Intensive Care Unit for 48 hours. We modelled 37 different attributes, such as heart rate and body temperature. The auxiliary covariates used are time of measurement, id, age, gender, icu type, and mortality. Given the high dimensionality of the auxiliary covariates, we focus on LMM-VAE and IVAF, the two model classes designed handle arbitrarily many covariates. As shown in Table 4 both LMM-VAL and LVAE effectively extract information from the auxiliary covariates for future prediction. The LMM-VAE variants trained with GSNN demonstrate competitive performance compared to IVAF, with further improvements observed as the number of basis functions incrcases in Figure These results underscore their potential for real-world healthcare applications.
0.52
OLMM-VAE
0.52
SLMM-VAE
test MSE
0.50
0.50
0.48
0.48
8
number of basis functions
Figure 5: The test MSE decreases the number of basis functions for the auxiliary covariate time increases. Plots are generated via GSNN training.
11

---

Published in Transactions on Machine Learning Research (05/2025)
Table 4: Test MSEs on the Physionet Challenge 2012.
MODEL (dim(z) 30)
MSE +
MEAN PREDICTION
0.902=#0.0000
VAE (KNGMM & WELLING 2014
0.905=:0.0030
LVAE RAMCHANDRAN ET AL. 2021
0.480=0.0349
OLM,-VAE, VI (Ouns)
0.618
SLMM-VAE, VI (Ouns)
0.0049
OLMM-VAE, GSNN (OURS)
0.486::0 0009
SLMM-VAE, GSNN (OURS)
0.488::0.0m 0.009
8 Conclusion and Future Work
We present LMM-VAE, a novel method for modelling high-dimensional longitudinal data that scales to large datasets with numerous covariates. Empirical evaluation reveals the importance of modelling auxiliary covariates and the potential of using LMM-VAF for longitudinal analyses. Theoretical analysis demonstrates connections to the Ga-baaedd VAEs. This connection provides solid foundation to adapt LMM-VALL to different modeling tasks by adaptively incorporating basis functions, well establishes LMM-VAE's reduced-rank approximation method for GP prior VAEs. Avenues for future
potential to be viewed research include extending our method to facilitate counterfactual reasoning, and incorporating alternative priors for the latent variables.
Acknowledgements
We thank the Aalto Science-IT project for the generous computational resources. This work is supported by the Research Council of Finland (decision number: 359135).
12

---

Published in Transactions on Machine Learning Research (05/2025)
References
Matthew Ashman, Jonathan So, Will Tebbutt, Vincent Fortuin, Michael Pearce, and Richard E. Turner. Sparse gaussian process variational autocncoders. ArXiu, 2020.
Douglas Bates, Martin Mächler, Ben Bolker, and Steve Walker. Fitting linear mixed-effects models using lmc4. Journal of Statistical Software, 67(1):1-48, 2015. doi: 10.18637 22118637jss.v067.000
Matthias Bauer, Mark van der Wilk, and Carl Edward Rasmussen. Understanding probabilistic sparse gaussian process approximations. In Advonces in Neural Information Processing Systems, 2016.
Francesco Paolo Casale, Adrian Dalca, Luca Saglietti, Jennifer Listgarten, and Nicolo Fusi. Gaussian process prior variational autoencoders. In Advances in Neural Information Processing Systems, 2018.
Peter Diggle. Analysis of longitudinal data. Oxford university press, 2002.
Ronald A. Fisher. Statistical Methods for Research Workers. Oliver and Boyd, 1925.
Vincent Fortuin, Dmitry Baranchuk, Gunnar Ractsch, and Stephan Mandt. Gp-vac: Deep probabilistic time series imputation. In Proceedings of the Twendy Third International Conference on Arlificial Intelligence and Statistics, 2020.
Francis Galton. Regression Lowards mediocrity in hereditary stature. Journal of the Anthropological Instilute, 1886.
Andrew Gelman. Multilevel (hierarchical) modeling: What it and cannot do. Technometrics, 2006.
Edward I. George. The variable selection problem. Journal of the American Statistical Association, 2000.
Angela Green, Katherine Reeder-Hayes, Robert Corty, Ethan Basch, Mathew Milowsky, Stacie Dusetzina, Antonia Bennett, and William Wood. The project data sphere initiative: Accelerating cancer research by sharing data. The oncologist, 2015.
James Hensman, Nicolò Fusi, and Neil D. Lawrence. Gaussian processes for big data. In Proceedings of the Twenty-Ninth Conference Uncertointy in Artificial Intelligence, 2013.
James Hensman, Nicolas Durrande, and Arno Solin. Variational fourier features for gaussian processes. Journal of Machine Learning Research, 2018.
Irina Higgins, Loic Matthoy, Arka Pal, Christopher P. Burgess, Xavier Glorot, Matthew M. Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with a constrained variational framework. In International Conference on Learning Representations, 2016.
Aapo Hyvärinen and Hiroshi Morioka. Unsupervised feature extraction by time-contrastive learning and nonlinear ica. In Proceedings of the 30th International Conference on Neural Information Processing Systems, 2016.
Metod Jazbec, Matt Ashman, Vincent Fortuin, Michael Pearce, Stephan Mandt, and Gunnar Rätsch. Scalable gaussian process variational autocncoders. In Proceedings of The 24th International Conference on Artificial Intelligence and Slatistics, 2021.
Matthew J Johnson, David K Duvenaud, Alex Wiltschko, Ryan P Adams, and Sandeep R Datta. Composing graphical models with neural networks for structured representations and fast inference. In Aduances in Neural Information Processing Systems, 2016.
Tyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoencoders and nonlinear ica: A unifying framework. In Proceedings of the Twenty Third International Conference Artificial Intelligence and Slatistics, 2020.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference Learning Representations (ICLR), 2015.
13

---

Published in Transactions on Machine Learning Research (05/2025)
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In 2nd International Conference Learning Represenlations, 2014.
Diederik P. Kingma and Max Welling, An introduction variational autoencoders. Foundalions and Trends in Machine Learning, 2019.
Nan M. Laird and James H. Ware. Biometrics, 1982.
Yann LeCun, Corinna Cortes, and Christopher J.C. Burges. Mnist handwritten digit database, 2010.
llya Loshchilov and Frank Hutter. Decoupled weight decay regularization, 2019. URL https://arxiv.org/
abs/1711 121/1 05101
Ariel I. Mundo, Timothy J. Muldoon, and John R. Tipton. Generalized additive models to analyze non-linear trends in biomedical longiludinal data using r: Beyond repeated measures anova. and linear mixed models. RioRriv 2022.
Tony Plate. Accuracy versus interpretability in flexible modeling: Implementing tradeoff using gaussian process models. Behaviormetrika, 2000.
Assaf Rabinowicz and Saharon Rosset. Cross-validation for correlated data. Journal of the Americon Statistical Association, 2022.
Siddharth Ramchandran, Gleb Tikhonov, Kalle Kujanpää, Miik Koskinen, and Harri Lähdesmäki, Lon- gitudinal variational autocncoder. In International Conference on Artificial Intelligence and Statistics, 2021.
C. E. Rasmussen and C. K. 1. Williams. Gaussian Processes for Machine Learning. MIT Press, 2006.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Proceedings of the 32nd International Conference on Machine Learning, 2015.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In Proceedings of the 31st International Conference on Machine Learning, 2014.
Benoît Sauty and Stanley Durrleman. Progression models for imaging data with longitudinal variational auto encoders. In Medical Image Computing and Computer Assisted Intervention, 2022.
Marc-Andre Schulz, B. Yeo, Joshua Vogelstein, Janaina Mourao-Miranada, Jakob Kather, Konrad Kording, Blake Richards, and Danilo Bzdok. Different scaling of linear models and deep learning in ukbiobank brain images versus machine-learning datasets. Nature Communications, 2020.
Tkaro Silva, George Moody, Roger Mark, and Leo Anthony Celi. Predicting mortality of icu patients: The plysioonet/computing in cardiology challenge 2012. Predicting Mortality of 1CU Patients: The /eoompuin/ in Cardiology Challenge, 2012.
Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R. Garnett (eds.), Advances in Newral Information Processing Systems, 2015.
Arno Solin and Simo Särkkä. Hilbert space methods for reduced-rank gaussian process regression. Slatistics and Computing, 2019.
Michalis Titsias. Variational learning of inducing variables in sparse gaussian processes. In Proceedingss of the Twelth International Conference on Artificial Intelligence and Statistics, 2009.
Anthony Tompkins and Fabio Ramos. Fourier feature approximations for periodic kernels in time-series modelling. In Proceedings of the Thirty-Second AAAT Conference on Artificial Intelligence and Thirtieth Innovative Applications of Artificial Intelligence Conference and Eighth AAAI Symposium on Educational Advances in Artificial Intelligence, 2018.
14

---

Published in Transactions on Machine Learning Research (05/2025)
Shengjia Zhao, Jiaming Song, and Stefano Ermon. Infovae: Balancing learning and inference in variational autoencoders. Proceedings of the AAAT Conference on Arlificial Intelligence, 2019.
Harrison Zhu, Carles Balsell-RRodas and Yingzhen Li. Markovian gaussian process variational autoencoders.
In Proceedings of the 40th International Conference Machine Learning. 2023.
Vadim Zipunnikov, Sonja Greven, Haochang Shou, Brian S. Caffo, Daniel S. Reich, and Ciprian M. Crainiceamu. Longitudinal digh-dimensional principal components analysis with application to diffusion tensor imaging of multiple sclerosis. The Annals of Applied Statistics, 2014.
15

---

Published in Transactions on Machine Learning Research (05/2025)
APPENDIX
A Connecting LMM Set-ups to LMM-VAE
To enhance clarity, we present simple example that links the formulation of LMM set-ups to the context of LMM-VAE.
Consider a medical dataset where we aim to model the latent space using covariates such as patient ID, age, and time. For simplicity, assume the dataset. includes three unique patients, and we apply one-hot encoding (OHE) represent patient ID.
In the AMM-VAE framework, the response variable is the latent variable, E R', We categorize the covariates follows: the effect of patient ID is modelled random effect, while that for age and time are treated fixed effects. The L-dimensional latent variable for patient p at time point can then be modelled as:
2t
(p)
As
time age
.AR 0
E
time age
AR)
I
+
0
Axx + E.
Here, the fixed and random effects matrices, As and AR, respectively. represents the noise error term, and o) the vector of covariates (such as time, age, and one-hot encoded patient ID) patient p at timepoint
We stack the observations for all time points of patient p as:
Zp AX, + Ep.
and more generally, the entire dataset can be represented as
A
which we arrive at following our formalism in the paper. This approach provides clear structure for modelling the latent space while accounting for both fixed and random effects through global parametrization, using covariates that are relevant to the dataset.
B Hierarchical Plate Model
C
sLMM-VAE
C.1 sLMM-VAE ELBO formulation
Assuming a factorizing variational posterior q(A, A, Ex = -T go (c(A)) Y, A, (A))())(()) where all distributions are Gaussians, we oblain tractable ELBO similarly to that of oLMM-VAE in the main text
logn(V(X)2
F. F= [logp(yn(0, 2.)]
k= nEh
EyCA) KL ( A, (0((())) ((((()))))
(KL()(A) p(A)) KL (g(0) |(0)),
16

---

Published in Transactions on Machine Learning Research (05/2025)
(k)
(k)
2c (R)
Z
F/
1k
K
9
Figure 6: LMM-VAE's hierarchical plate diagram, which illustrates both fixed and random effects.
Latent Space
3
3
3
Encoder
Decoder
3
(sin(8),.cos(h)) cos(6:)) cos(()))
(sin(())
Figure 7: Diagram of SLMM-VAF.
where 2n, = Ax, + c()) As with OLMM-VAE, we could drop prior distributions p(0) and p(A) and treat the generative model parameters as perrparaameters. We refer this model class as the structured LMM- VAE (sLMM- VAF)..VAE). The associated diagram for sLMM-VAE can be found in Figure
The formulation of sLMM-VAE could yield competitive performance when the auxiliary covariates sufficiently characterize the output. This intuition is outlined below:
The noise variable identical for observations arising from the same group, thus reducing stochas- ticity.
2. The linear model A is now explicitly accounted for and learned as part of the decoder. Note that in the original oLMM-VAE, A is only learned as part of the prior.
D Gaussian stochastic neural network training objective for LMM-VAES
We credit Sohn et al. (2015) for introducing the Gaussian stochastic neural network (GSNN) training objective in the context of conditional generation. We expand on the motivation behind this approach. First, the ELBO objective consists of the well-known roconstruction and regularization (the KL divergence) terms. This means that at training time the data y reconstructed from the variational approximation, whereas test time, the data is generated from the prior. The key insight from Sohn et al. (2015) to ensure consistency
17

---

Published in Transactions on Machine Learning Research (05/2025)
between training- time reconstruction and test-time generation. To achieve this, the authors propose aligning the variational approximation with the prior, i.e., setting (o(Znlyn) = p(zn A, Fn). Substituting this into the ELBO results in the KL divergence becoming exactly zero, reducing the ELBO to the GSNN training objective as shown in in Eq. I1 in Section 4.4 In essence, the GSNN training objective corresponds to the standard ELBO when the variational approximation equals the prior.
The GSNN training framework is suitable for conditional generation when the covariates an are rich enough such that they can sufficiently guide the generation of the corresponding data sample yn. In other words, when the covariates Zn sufficiently informative in characterizing Yn. the GSNN training objective explicitly trains the model to generate accurate/ recurate/high-qualitt data samples yv. leading to competitive performance. Results from our empirical evaluation supports this statement.
GSNN can also be understood from the regression view point: if the amount of variability in the latent variables approaches zero, then GSNN training objective would effectively reduce regression. To gain further intuition about what the training procedure does for the LMM-VAE model, we can explicitly re-write the GSNN loss as
N log log 2n, 0)
LGSNN (M|X,
8=1 ==1
log (yn Axx, + En,a 0).
5
We now observe that the trained network is composed of a first linear layer conditioned En, where there is some noise injection (can also be thought of as regularizalion), followed by another network parameterizing the decoder. In other words, we are now learming a mapping from In + yn via a neural network composed of a linear layer and a second network. When we rewrite the loss in the aforementioned manner, there is a strong resemblance to the task of regression (with some regularization in the first layer). That is, we are explicitly training the linear model together with the decoder network to maximize the expected log- logelikelihood of conditioned on This stands in contrast to variational inference, where we learn indirectly by primarily minimizing A) [KL( (g (z,, (z.l/y.)] (r,) p(z,, [,, x.))]]
Additionally, we would like to point out that limitation of the GSNN objective is that it does not provide any estimate of the latent variables because it is designed for the generative task only.
E
Further characterization of LMM-VAE
E.1
Additional comparison with SVGP-VAE
The approxination quality of LMM-VAL improves as the number of basis functions increases. Similarly, the sparsity of SVGP-VAE is determined by the number of inducing points, with larger set generally leading to richer approximation. In ablation study, we analyze the effect of varying the nunber of inducing points in SVGP-VAR using the rotating MNIST experiments, and include the results in Table 5) We observe an improvement in Test MSE as the number of inducing points increases, leading to a consistent, albeit gradual, enhancement.
A similar trend observed with SLMM-VAL Figure 4 where test performance improves and stabilizes as more basis functions are added. These results indicate alignment in trends across the different GP approximation methods. While SVGP-VAE and sLMM- VAE deliver comparable and competitive results, the GSNN variants of LMM-VAE outperform both, achieving the strongest test performance in the rotating MNIST experiments.
E.2
LMMs vs LMM-VAES
To better characterize LMM-VAE's performance, we evaluate its behaviour in two low-dimensional sottings: one where classical linear mixed models (LMMs) are expected to perform well, and another where their performance is likely to be limited.
18

---

Published in Transactions on Machine Learning Research (05/2025)
Table 5: Test MSE as a function of inducing points for SVGP-VAE on Rotating MNIST.
MODEL (dim( x) = 16)
No. OF INDUCING POINTS
MSE +
16
0.019
32
0.018
SVGP-VAE JAZBEC ET AL. 2021
48
0.017:10.00m
61
0.01 00000077
We simulate low-dimensional synthetic dataset based on latent variable model. A univariate latent. variable is defined for each subject across time, incorporating subject-specific random intercept and fixed effect for time. This latent variable drives three observed variables y E R via predefined linear measurement functions, with additive Gaussian noise. The resulting dataset captures structured variation across subjects and over time, providing a controlled environment for model comparison. The experiments were conducted as a function of the number of subjects, each observed at 20 timepoints. For validation and lesting, the last timepoints from 20 subjects were held out, with the remaining data used for training.
We compare two modelling approaches: classical lincar mixed model (LMM) and variant of OLMM-VAE, where the observed response is modelled as linear transformation of the latent variable. Based on the lest MSE in the first column of Figure S both models exhibit similar trends with respect to dataset size (i.e., number of subjects). This outcome aligns with expectations, as both models are well-suited to linear generative processes. However, the oLMM- VAE achieves lower overall test MSE, suggesting improved generalization.
To further evaluate the OLMM-VAE's flexibility, we generate a second dataset. variant which y now non-linear functions of the latent variable. We train both the LMM and oLMM-VAE to this new dataset. and report the results in the second column of Figure In this set-up, the oLMM- VAE is extended to incorporate non-linear decoder, enabling it to learn flexible mapping from the latent space to the observed outputs. With this added capacity, the oLMM-VAL yiclds substantially lower test MSE than the classical LMM, highlighting its advantage in capturing non-linear structure.
E.2.1 Experimental details
LMM was implemented via the Ime4 library in R (Bates al., 2015). following the ap- http:///////eeeeeee/eee/uaeeaee/
proach described in https::/macctthebbio.. artp:dorrrrbbbbiiiiiaaaaaaeeeeen
notes html## to fit multivariate mixed models.
Mcanwhile, the VAM was employed and trained using GSNN training. For the lincarly generated dataset, the decoder was parameterized as a linear function and trained with learning rate of 0.001. For the non-linearly generated data, the decoder consisted two hidden layers with Tanh nonlinearities and hidden sizes of 16 and 8, trained with a learning rate of 0.01. Across both set-ups, we use a latent dimension of 1, set og as 0.005, and use the AdamW Optimizer (Loshchilov Hutter 2019). We monitored the validation loss and applied strategy akin to early stopping, saving the model weights with the lowest validation loss. The oLMM-VAE was trained for maximum of 2000 epochs.
F GP prior VAEs as LMM-VAES
Here, we outline how LVAE can be implemented as LMM-VAE, Assuming the j-th additive latent component in IVAE model depends on a specific subset of covariates (0)) E 20) X and that the corresponding Fourier features are (z(())) The linear basis function approximation for the j-th additive latent component is then z (0) (z)())) and the complete model can be written as
R
R
z = Σ #=1
0 = Σ
1 (2)(1)))
(15)
1=1
where R is the number of additive kernels in IVAE model and p(A) 11 p( ((A))).
19

---

Published in Transactions on Machine Learning Research (05/2025)
Linearly generated data
Non-linearty generated data
0.0050
0.0045
Test EStMSE
0.08
0.0040
0.04
0.0035
0.02
0.0030
50
75
100
125
150
175
200
75
100
125
150
175
200
Number
Number Subjects
LMM
OLMM-VAE
Figure 8: Test MSEs, averaged across dimensions, as function of the number of subjects, based on simulated data. The first column describes the set-up in which is generated linear function of the latent variable, while the second column describes that in which y is non-linear function the latent variable.
G Identifiability
Given that our model belongs to the family of conditional VAEs, LMM-VAE inherits the property of being identifiable. This result was established by Khemakhem et al. 22020), who rely on auxiliary covariates æ constrain the latent space via conditional prior p(=(æz).
Formulation. In general, an unconstrained latent-variable model of the form
p 2) = p(z))=)((z),
is not identifiable. Instead Khemakhem et al. 2020 assume, adapted to our notation, that
ply, z (z) p,(0l=))r,>((//)
where py(()z) f(2) + e, with e p(c), and f 2 is injective mapping, c.g., neural net. The conditional prior
dimn(=)
(r.((z|a)
II mx(2;) exp
Σ (1)(
is a fully factorized member of the exponential family of distributions. Here, Tig are the sufficient statistics, and X(.) some function æ, e.g., a neural net.
Relevant to our current discussion are Definitions and of Khemakhem et al. (2020), which define our notion of identifiability. We cite here in full for completeness.
Definition (Khemakhem al. 2020). Let be equivalence relation on O. We say that p(æ) ((z))))(() is identifiable up to if
po(x:) p=(x) 0. The elements of the quotient space 0/~ are called the identifiability classes.
Definition 2 (Khemakhem et al. 2020). Let be the equivalence relation on θ defined as follows:
(f.T.)) (F. 3AAC (z) AT +c Yo X,
where A a dim(z)k x dim(z)k matrix and e is a vector. If is invertible, we denote this relation by ~ If A is a block permutation matriz, we denote by ~p.
Given these, the main theorem
20

---

Published in Transactions on Machine Learning Research (05/2025)
Theorem 1 (Khemakhem et al,, 2020). Assume that observe data sampled from generative model defined according to the model specified above with porameters (S.T.A) A)) Assume the following holds:
(i) The set fa E (lp)(())= =0 has measure zero, where Pp is the characteristic function of p(z).
(ti) the mizing function f is injective.
(tti) the sufficient statistics Tig are differentiable almost everyahere and (Ti5)1s5sk are linearly independent on any subset of X of measure greater than zero.
(it) there exist dim(z)k + distinct points x." ,...., dim(()) such that the matrix
L = -
((xdiim()))) {(x²)),
of size dim(z)k x dim(z)k is invertible,
then the parameters (f.T,)) ---identifiable
Picking p(=) = _V(=00, o and suitable choice of architecture for all that remains to be shown is that our prior a mean-field exponential family distribution. We have that for a normal prior p(A) = -N(vec(Aj00,
p(=[0)= = / x))A)dA
/ N (F
dim(=)
II Ex(a,00,5-1) W (2ila,z,0)) i=1
dim(z)
II
N (210,02 + xT
i.e, p(=|æ) is an exponential family distribution. Assuming a rich enough variational posterior g(z) we can then guarantee asymptotic identifiability given Theorem 4 by Khemakhem et al. 2020))
G.1 Clinical Trial Data and Health MNIST
For brief insight into LMM-VAE's identifiability, we perform series of experiments involving a randomized control trial (RCT) dataset based colorectal cancer treatment Green et al. (2015). and Health MNIST. Additional information about the RCT data is disclosed in Appendix
Illustrating Identifiability. A measure to evaluate identifiability is the mean correlation coefficient (MCC) between the true latent variables and their estimated posterior (Hyvärinen & Morioka 2016 Khemakhem ct al. 2020). Computing the MCC necessitates access to the ground truth for correlation comparisons, an unrealistic assumption in practice. A common approximation method involves comparing the correlation between different training runs, together with a second metric evaluating the latent space's predictive performance. While desirable, obtaining identifiability is still conditioned on good predictive performance, latent space that collapses to the same bad local optimum per run would still result in perfect correlation.
We compute variant of LMM-VAE's MCC alongside predictive metric (NLL MSE). We include at least one other model per experimental set-up to serve baseline. For both real world clinical trial dala and Health MNIST, LMM-VAE reports competitive results across multiple latent dimensions (see Tables and respectively).
We use Khemakhem lib/metrics. Py
2020 implementation to compute the MCC: httpp/////.tta.....o.//t//////aa
21

---

Published in Transactions on Machine Learning Research (05/2025)
Table MCC for oLMM- VAE and VAE on the clinical trial data.
MODEL
LATENT DIMENSION
NLL +
MCC
VAE
10.0
0.47:10.03
20
OLMM-VAR (OURS)
-14.142.9
0.50+0.03
VAE
-10.9+7.5
0.48=+0.5
18
OLMM-VAE (OURS)
100.716.7
0.51 0..1:
VAE
99.91
0..1
15
OLMM-VAR (OURS)
11..611.4
0.53
Table 7: MCC for LMM and VAE on Health MNIST, We use the orange parametrization for ML VAE.
MODEL
LATENT DIMENSION
MSF +
MCC f
VAE
1.05910.000
0.591...3
32
OLMM-VAE (OURS)
0.023+0.000
0.58+00.01
VAE
0.058+0.002
0.57+0.02
16
OLMM-VAE (OURS)
0.003
0.5610.02
VAE
0.056
0.62.10.04
8
OLMM-VAE (OURS)
0.025:t0.001 0.60+0.03
H Health MNIST
H.1 Dataset Description
We rely on the dataset construction script from Ramchandran et al. 2021
To generate shared age-related effect, we gradually shift all digit instances towards the right corner over time. Half of the instances of '3' and '6' assumed to be healthy (diseasePresence 0), while the other half are inflicted with disease (diseasePresencce = 1). The diseased instances are rotated across 20 timepoints, with the rotation degree determined by the time to disease diagnosis (discaseAge).
We further include a binary noise covariate location, which is randomly assigned to each unique instance, and apply a random rotational jitter each data point to simulate noisy observations. Additionally, we mask out 25% of each image's pixels (to assess imputation capabilities).
In total, there are 1300 unique instances present in the dataset, where 650 correspond to the biological sex Male, and the remaining 650 correspond to Female. Each of these unique instances have sequence length of 20. We withhold the last 15 timepoints of 100 subjects to construct the test set. The first five timepoints of these aforementioned subjects are included in the training set. The remaining dataset is then randomly split to construct the train and validation scts, in an approximate ratio of 85 : 15.
H.2 Experimental details
We took the implementation of the baseline SVGP-VAF Jazbec et al. (0021) from https://githubb.oo ratschlab/SVCP--V2 GPP-VAE Casale et al. (2018) from https: / thub. Sttm//pcaaaale/GPVAE and com
To perform experiments, we use the default hyperparameter setting specified in the respective repositories. Across all baselines, the architecture used for the experiment can be found in Appendix M
LVAE Ramchandran et al. (2021) from
For experiments regarding VAE, we use Adam optimizer Kingma Ba (2015) learning rate of 0.001.
We monitor the loss the validation set and employ a strategy similar in spirit to early stopping, where
22

---

Published in Transactions on Machine Learning Research (05/2025)
save the weights of the model with the optimal validation loss. LMM-VAE was allowed to run for maximum of 2500 epochs. We define d= 1.
For all experiments, we report the mean and standard deviation obtained across five runs.
H.3 Additional illustrations of LMM-VAE's predictions
We visualize two sets of trajectories corresponding to individuals in Figure
Time from disease event
t-5
t-4
t-3
t-2
t-1
to
t1
t3
15
t6
t7
t8
t
data
3
2
3
3
3
3
22
3
3
3
?
LMM-VAE
reconstruction
3
3
3
3
3
3
3
3
3
3
3
3
3
3
3
LMM-VAE
prediction
3
3
3
3
3
3
3
3
3
3
3
3
3
3
true
image
3
3
3
3
3
3
3
3
3
3
3
3
3
data
6
6
3
402
V
LMM-VAE
reconstruction
6
6
6
6
6
6
6
6
9
9
9
9
9
9
6
LMM-VAE prediction
6
6
6
6
6
6
6
6
6
6
6
6
6
6
true image
6
6
6
6
6
6
6
6
v
V
V
V
V
Figure 9: We illustrate 2 sets of images here, each corresponding to different biological sex. Per image
set, we visualize the reconstructions and predictions obtained on the test set by oLMM- VAE with 16 latent dimensions in the second and third rows. The noisy data along with the original uncorrupted images are also depicted in the first and last rows, respectively.
H.4 Supplementary tables for Health MNIST
The tables containing the experimental results for Health MNIST are described in Table S and Table g
Table 8: Imputation MSFs for LMM-VAE on the HealthMNIST Dataset.
MODEL
LATENT DIMENSION IMPUTATION MSE L
LMM-VAE
1.00200.0000
LMM-VAE
16
0.002+0.0000 0000
LMM-VAE
0.00241.0000 0000
LMM-VAE
0.002+0 0000
LMM-VAE
32
0.002
LMM-VAE
0.002 0.0000000000 0000
Rotating MNIST
1.1 Dataset Description
We revise the train/vallin aeain/valddationtteen splits to reflect our interest evaluating the learned id-specific parameters.
This method of evaluation not new Gelman (2006): Rabinowicz & Rosset 2022 Our modification results
23

---

Published in Transactions on Machine Learning Research (05/2025)
Table 9: Predictive Test MSES for LMM-VAB on the HealthMNIST Dataset.
MODEL
LATENT DIMENSION PREDICTIVE MSE 4
LMM-VAE
0.0284
LMM-VAE
16
0.0183+0.0001
LMM-VAE
0.018510.0011
LMM-VAE
0.028511.0000
LMM-VAE
32
0.0177+0.0013
LMM-VAE
0.0177
in the validation set comprising of ids that are also present in the train and test sets, which is not the case in the original set-up (Casale et al. 2018 Jazbec al. 2021).
We generate 400 instances of the digit '3' and rotate each of these unique instances through 16 angles spaced uniformly in 00 2T). This results in 6400 images in the entire dataset. To construct the test set, we consider 80 instances at random, and take subsequences corresponding to four consecutive angles of the aforementioned digits. The train and validation sets are then randomly constructed based on the remaining images, in an approximate ratio of 80 : 20.
1.2
Experimental details
For the baselines CVAE (Sohn et al., 2015 : GPP-VAE (Casale et al.) 2018 and SVGP-VAE (Jazbec et al. 2021 we rely on the implementation from httpps//// tttb///gttuuuu...eeeg thub. wtt..//ottaaaaa//eeeeaaao.
The default experimental sct-up is used, with a slight modification to the final activation layer (for purposes of standardizing the architectures across models). See Appendix M for the exact archilecture used.
To perform experiments for LVAE, we use the experimental set-up from htttps://github com/SidRama/ VAE (Ramchandran et al., (2021). Similar the other basclines, the actual architecture usod for the experiment can be found in Appendix M
For experiments regarding LMM-VAB, of 0.001. We also use an exponential learning rate scheduler with a step size of 500. We monitor the loss on the validation set and employ a strategy similar in spirit to early stopping, where we save the weights of the model with the optimal validation loss. LMM-VAE was allowed to run for a maximum of 2000 epochs. We also define o, 1. 1.
the Adam optimizer (Kingma & Ba 2015 and learning rate
For all experiments, we report the mean and standard deviation obtained across four runs.
1.3 Additional illustrations on Interpretability
Illustrating Interpretability. As additive model, LMM-VAE enjoys the benefit of being interpretable, which has merited interest in domains such as healthcare Plate (2000). We showcase this property with the sLMM- VAE variant on Rotating MNIST in Figure 10. Conditioned on an instance-specific e ((k) the latent space mapping is fully deterministic as a function of the auxiliary covariates.
inp the rotation angle by 0, the first two basis functions for the Rotating MNIST are æ (sin(0), "o. For latent dimension I, the latent space z/ α.æ +&; under consideration. We illustrate both shared and random effects in Figure I0 The vertical shifting of
(k) where k denotes the digit instance
the graph from the mcan function, as induced by (k) corresponds to id-level variation. Simultancously, the rate of angular rotation remains consistent for the different digit ids, constituting a shared effect.
Additional illustrations of how different noise realizations correspond to different digit identities can be found in Appendix L33
Additional illustrations (Figure 11 Figure 12 Figure 13 further illustrate how different noise realizations
correspond to different digit identities.
24

---

Published in Transactions on Machine Learning Research (05/2025)
4
in
mean
id: 323
2
id: 326
0
-2
0 1
2
3
radians
Figure 10: The dashed red line represents the expected function the 1st latent dimension w.r.t angular information. We obtain the trajectories of two specific instances (id: 323 and 326) by adding f (323) and (326) respectively. Note that by moving along one these deterministic trajectories, we get images of the same instance rotated through different angles.
2.5
3.
5.0
5.0
meen
Figure 11: Visualisation of each the 16 latent dimensions with their corresponding mean plot (represented by the dashed red line). In addition, the mean of the digit instance (id: 323) also indicated here (solid blue line). The shaded region corresponds to the uncertainty in the function draw of id:323 standard deviation is represented here).
1.4 Further experiments for Rotating MNIST
Here, we provide additional illustrations (Figure 14 Figure lead to improvement in performance.
15 highlighting how additional basis functions
J Physionet Challenge 2012
J.1 Dataset Description
In total, there are 3599 unique instances present in the dataset for training, This palient subset is selected such that each instance is represented by at least 15 observations. We withhold the last timepoints of 1200 patients to construct the test set. The train and validation on the remaining observations, in an approximate ratio of 80 20. The resulting train and validation sets comprise of 3599 and 3598 unique instances respectively.
are then randomly constructed based
25

---

Published in Transactions on Machine Learning Research (05/2025)
-2
radians
id: 323
id: 326
Figure 12: Visualisation of each of the latent dimensions with the addition of the mean of id-specific noise.
In this plot, unique ids are represented.
id: 323
5
E
id: 326
3
Figure 13: Test predictions of 2 different digit instances corresponding to Figure 12 The generated digits are rotated through the same angles. However, they assume different identities with different noise realizations along each latent dimension.
0.0200
0.0200
0.0175
0.0175
0.01.5
Test MSE
0.0125
Test TestMSE MSE
0.01.00
0.0100
0.0075
0.0075
0.0050
OLMM -VAE
0.0050
OLMM-VAE
0.0025
SLMM-VAE
0.0025
SLMM-VAE VAE
0.0000
2
12
12
Number of basis functions
Number basis functions
Latent dimension of 32. Right: Latent dimension of 16. In general, MSE generally decreases the number of included basis functions increases for latent dimension of 32. We assume 9 and A are deterministic in this set of experiments.
Figure 14: Plot generated via GSNN training of oLMM-VAE and LAM.-VAE. Left:
J.2 Experimental details
We use the Adam optimizer Kingma & Ba 2015) and learning rate of 0.001, with an exponential learning rate scheduler. We monitor the loss the validation set and employ strategy similar in spirit to early
26

---

Published in Transactions on Machine Learning Research (05/2025)
0.01850
SLMM-VAE
0.01825
0.01800
Test MSE
0.01775
0.01750
0.01725
0.01700
0.01675
12
Number basis functions
Figure 15: Plot generated via VT of SLMM-VAR for latent dimension of 32. In general, MSE generally decreases as the number of included basis functions increases. We assume 0 and A deterministic in this set of experiments.
stopping, where we save the weights of the model with an optimal validation loss. LMM-VAE was allowed to run for maximum of 1000 epochs.
To perform experiments for IVAE, we use the experimental set-up from https:// com/ SidRama/ Bongitudinal-Ai (Ramchandran et al. 2021). Similar to the other baselines, the actual architecture used for the experiment can be found in Appendix M
For all experiments, we report the mean and standard deviation obtained across five runs.
K Project Data Sphere
K.1 Dataset Description
We use medical time series dataset derived from a randomized control trial for the treatment of colorectal cancer (datasct identifier: Colorec SanfiU 2007 131). This dataset is taken from the open data sharing platform 'Project Data Sphere' (Green et al. 2015
We performed the following pre-processing steps:
The following source domains were selected: Laboratory measurements (Im), adverse event information (ae), vital signs (vs), concomitant. medication (em), and demographic information (dm).
Information on the domains were formulated to longitudinal samples Y comprising of (lb, vs), and auxiliary covariates X comprising of (dm, ae, cm). Only measurement timepoints in ae and em with known start and end time were considered.
3. Timepoints with less than 70% information Y were excluded.
4. The 10 most common ae and cm in X were considered.
5. Patients with trajectory of less than 5 observations were excluded.
6. Each feature in Y was standardised to have zero mean and unit variance.
The final pre-processed datasct used for modelling is composed of 480 subjects with N = 6605 total observations, y E 3030, and æ E 222 Following one-hot encoding of the patient ids, we have æ E 503
The approach to constructing the train-validatiiontess sets similar to that for Health MNIST. For the test set, we consider the sequence of visits following the 5th visit (inclusive) of 100 patients. The first visits of these aforementioned patients are included in the train set. The remainder of the dataset is then split amongst the train and validation sets.
27

---

Published in Transactions on Machine Learning Research (05/2025)
L
Identifiability Experiments
L.1 Clinical Trial Data (Project Data Sphere)
The actual architectures used for LMM-VAE and the vanilla VAE can be found in Table
For experiments regarding the VAE and LMM- VAE, we rely on the Adam optimizer Kingma & Ba 2015 and a learning rate of 0.001. We also use an exponential learning rate scheduler with a step size of 500. We monitor the loss on the validation the weights of the model with the optimal validation loss. LMM-VAE was allowed to for maximum of 2000 epochs. We also define 0, = 1.
and employ a strategy similar in spirit to carly stopping, where we save
For all experiments, we report the mean and standard deviation obtained across five runs.
L.2 Health MNIST
The actual architectures used for LMM-VAE and the vanilla VAE can be found in Table 14
For experiments regarding the VAE and LMM-VAE: VAE, we use the Adam optimizer Kingma & Ba 2015 and a learning rate of 0.001. We also use an exponential learning rate scheduler with a step size of 500. We monitor the loss on the validation set and employ a strategy similar spirit to early stopping, where we save the weights of the model with the optimal validation loss. LMM-VAE was allowed run for maximum of 2000 epochs. We also define o=
For all experiments, report the mean and standard deviation obtained across five runs.
M Model Architectures
This section outlines the different model architectures used for the different experiments. Table 10 contains the neural network architecture used in the Health MNIST experiments, which follows [amchandran et al. (2021). Table 11 describes the architecture used for the Rotating MNIST experiments. This architecture is almost identical to that reported in Casale et al. (2188), apart from final activation layer of sigmoid instead of the ELU to suit the modified pre-processing pipeline. Table 12 details the architecture used for the Physionet Challenge experiments. Finally, Tables 13 and for the experiments on identifiability, where the decoder is composed of fully connected layers and ELU Sigmoid activation.
M.1 Identifiability Experiments
The architecture used for the RCT dataset is described in Table 13 while that for Health MNIST is
Table 14
28

---

Published in Transactions on Machine Learning Research (05/2025)
Table 10: Neural Network Architecture used for the model in the Health MNIST experiments.
Hyperparameter
Value
Dimensionality of input
36 x 36
Number of convolution layers
2
Kernel size
3 x
Stride
2
Pooling
Max Pooling
Inference Network
Pooling kernel size
2 x 2
Pooling stride
2
Number of feedforward layers
2
Width of feedforward layers
300, 30
Dimensionality of latent space
L
Activation function of layers
ReLU
Dimensionality of input
L
Number of transposed convolution layers
2
Kernel size
4 x
Generative Network
Stride
2
Number of feedforward layers
2
Width of feedforward layer
30, 300
Activation function of layers
ReLU, Sigmoid
Table 11: Neural Network Architecture used for the model in the Rotating MNIST experiments.
Hyperparameter
Value
Dimensionality of input
28 x 28
Number of convolution layers
Kernel size
x 3
Stride
Inference Network
Number of feedforward layers
Width of feedforward layers Dimensionality of latent space
128 T. ELU
Activation function of layers
Dimensionalitty of input
Number of convolution layers + Upsampling
Kernel size
3 3
Generative Network
Stride
Number of feedforward layers
Width of feedforward layers
128
Activation function of layers
ELU, Sigmoid
Table 12: Neural Network Architecture used for the model in the Physionet Challenge experiments.
Hyperparameter
Value
Dimensionality of input
37
Number of feedforward layers
2
Inference Network
Number of hidden units per layer
64, 32
Dimensionality of latent space L
Activation function of layers
ELU
Dimensionality of input Number of feedforward layers
L
2
Generative Network
Dimensionality of input Activation function of layers
32, 64
ELU
29

---

Published in Transactions on Machine Learning Research (05/2025)
Table 13: Neural Network Architecture used for the model in Project Datasphere's experiments.
Hyperparameter
Value
Dimensionality of input
30
Number of feedforward layers
2
Inference Network
Number of hidden units per layer
25, 23
Dimensionality of latent space L
Activation function of layers
ELU
Dimensionality of input
T.
Number of feedforward layers
2
Generative Network
Dimensionality of input Activation function of layers
23, 25
ELU
Table 14: Neural Network Architecture used for the model Health MNIST's experiments.
Hyperparameter
Value
Dimensionality of input
36 x 36
Number of convolution layers
Kernel size
Stride
Pooling
Max Pooling
Inference Network
Pooling kernel size
2 2
Pooling stride
2
Number of feedforward layers
Width of feedforward layers
300, 30
Dimensionality of latent space
Activation function of layers
ReLU
Dimensionality of input Number of feedforward layers
Generative werw
Dimensionality of input Activation function of layers
30, 300, 1000, 1200
ELU, Sigmoid
30


Reviews:
[{"review_id": "vMF31iFyLN", "reviewer": "Reviewer_hND8", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n- Well motivated problem.\n- A nice related work section and a generally solid introduction to the area.\n- Strong experimental results.\n\n**Weaknesses**\n- I find the exposition unclear at times.\n- More discussion/comparison with SGP-VAE is needed.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "We would often like to model longitudinal data in order to impute missing values or predict new values, for example. However, such data might be very high dimensional, and so it is more effective to instead model the data through a low dimensional latent space—typically achieved with VAEs. In order for VAEs to be able to model the correlations between samples within the same task (i.e. the state of the features representing a patient's health at different times, where the time feature is termed an *auxiliary feature*), the priors over the elements of the latent variable $z$ are endowed with some way to condition on the auxiliary variable. In previous work, this has been done by using GP priors. Since GPs are a kernel method and therefore suffer when the number of auxiliary features is high, the authors propose to use linear mixed models instead, which are more scalable in the number of auxiliary features. Importantly, their approach maintains scalability in the number of datapoints too, and is easy to train unlike existing sparse GP VAE approaches. They present two variants of their model, and propose two objective functions. They also open the door to non-linearising their linear models via the kernel trick. The proposed methods exhibit solid performance compared to the benchmarks across all experiments.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I use this section to request changes to the manuscript as well as outline my concerns in general.\n\n1. From the paper, I do not feel like I understand what the *random* features are or how they work in the model. It would be great if a clearer explanation between the role of these vs. the *shared* features could be given, perhaps by way of more/better examples.\n2. There are certain passages that are either in an appendix and need to be brought into the main text, or vice versa. I outline what I find to be the key offenders of this as follows:\n    - The ELBO derivation at the bottom of page 5 should be in an appendix, with only the key result remaining in the main text. \n    - An intuitive explanation of the sLMM-VAE encoder that sits in the main text would be great. \n    - Section 5 (Identifiability) is meaningless to the reader without a convincing motivation for an identifiable model in the first place---bring (at least part of) the subsection on 'Motivation' of appendix E into the main text.\n3. Non-linearising the LMM with a fixed number of basis functions is, in spirit, the same thing as using a sparse GP (with, here, the advantage that inference is performed in parameter space). However, does this not mean you run into the same *curse-of-dimensionality* problems that GPs/kernel methods have for large numbers of features---does the kernel trick here not limit the scalability of the LMM-VAE in the number of auxiliary variables (the very problem you seek to address)? Perhaps I have misunderstood something important here.\n4. In the future work section, you mention using NNs for prior means. Would this not destroy the nice interpretability that the (identifiable and linear) LMM-VAE enjoys with respect to the auxiliary variables?\n5. You explore the effect of varying the number of basis functions in the kernelised LMM-VAE a little bit, but never compare to varying the number of inducing points in the SGP-VAE. Further, I cannot find anywhere in the paper where the number of inducing points in the SGP-VAE is reported. It would be good to ensure the degree of sparsity is similar in order to be confident that the comparison is indeed a fair one.\n6. Minor issue, but appendix D has the wrong title.", "broader_impact_concerns": "N/A.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "JFQf85F2zU", "reviewer": "Reviewer_ZB6M", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n\nThe paper tackles a well-known problem in the context of VAEs and representation learning, where usually temporal structure of longitudinal is considered as a second priority (in the sense that usually the focus has been so long centered on reconstruction, OoD and uncertainty quantification). In that sense, I find the paper interesting and kind of in a good direction when reviewing previous SOTA works and problems. One of the strengths of the work comes from the fact of considering longitudinal data that is multimodal (either regression or classification problems, that could include categorical variables for instance). However I do believe that the paper do not develop well this side of the topics and do not link itself to some good literature on that sort of problems written during the last years. A good style of writing and in general, clarity, are there — in that sense I did not detect any problem on understanding what the authors are proposing and what they want to do.\n\n**Weaknesses:**\n\nTo me, the most surprising aspect while reading the paper is the lack of reference or mention to principal component analysis (PCA), its probabilistic version (PPCA), the GP adaptation of such vast literature, super well-known as Gaussian process latent variable models (GPLVMs) and dozens or other variants that have been around the community in the last years under names like GP-decoders or similar titles.\n\nIn that regard, it is hard for me to believe that VAE models updated with sorts of PPCA structure have not been studied since their apparition +10y ago. In fact, the linear consideration of both encoders and decoders have been around since the beginning for the understanding and analysis of such generative models. \n\nEven if we assume that the LMM-VAE is not related to PPCA (I’m very sure it is), its connection with GPs and basis functions is trivial… or at least, there is not novelty there for anyone with a little background on probabilistic linear models, GPs and the weight/function-space view duality of the previous ones.\n\nLast but not least, the experiments are certainly limited and show important flaws. To name two of them: I do not entirely believe the results or conclusions drawn from the tables and figures that show improvements of one order of magnitude wrt well-established and rigorous methods in the side of longitudinal-structured VAEs. Why I say this? Well, the question that quickly comes to mind is if that performance with dim($z$)$=32$, dim($z$)$=30$ or dim($z$)$=16$ and zero variance on the MSE accuracy for the Rotated/Health versions of the MNIST dataset really responds to a great method, or is just an oddly tuned experiment that worked well for such dimensions. I kindly ask to have a bit of more rigour on these sort of experiments, particularly to show readers the complete and clear performance of the method for different latent space structures and dimensions. The computational cost would be also a topic to consider here, but lets let it out of the discussion for now..", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "Motivated by the problem of poor scalability of VAEs to real-world longitudinal datasets, that usually contain non-linear effects, time-varying covariates and missing values, the authors introduce the linear mixed model VAE (LMM-VAE). Thus, the driving idea is to make improvements in terms of interpretability, identifiability and last, make the model available to the introduction of basis functions (i.e, of periodic nature), such that the model can capture certain types of structures in the temporal/longitudinal data. In practice, the model is a VAE, where one assumes the latent variable is a noisy linear projection of some additional auxiliary covariate $X$, given observed outputs $Y$. The experiments reflect some improvement w.r.t. more complex and well-established SOTA methods, mostly in the context of GP-VAEs and sequantially-structured VAEs.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I do believe the work need a thorough revision, and I am certainly concerned about two principal issues (I would like to see this solved somehow):\n\n- *1)* Is there a connection with PPCA? Are LMM-VAEs actually something similar to them? Could authors review that side of the literature appropriately and give a bit more of clarity and wide vision on that side?\n\n- *2)* Given the experiments considered, could the authors make a thorough study of the behaviour of the method in the context of latent space dimensionality, structure and number of parameters involved?", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "PjtqBsjl7d", "reviewer": "Reviewer_xWFG", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "On the whole the paper is easy to read and follow, and I think the contribution is both valid and valuable to the community. That said, I have some comments:\n\nSo in the interests of transparency, some of my work falls at the intersection of machine learning, statistics, and importantly, empirical research in the medical, health, and social domains. As such, I am quite familiar with the need for models which account for the hierarchical nature of data. One of the key motivations for this is to account for the shared variance for (e.g.) multiple measures from the same patient. If one fails to sufficiently deal with this, the associated statistical inference is bias (underestimated standard errors). Whilst there is also a more general case that accounting for nested structure helps the model more closely align with the data generating process (DGP), and therefore the model is more predictive, the case for mitigating bias for statistical inference is extremely important for drawing conclusions about e.g. the efficacy of treatments. With this in mind, I have the following comments:\n\n1) The paper could benefit from a strong/simple motivating example which is carried through the methodology. This is particularly important in my view, because one of the key audiences for this paper could include people working in the applied domains, and tying the approach to a more conventional example would be particularly helpful given that the authors use a formalism which is not quite the same as is usually seen in LMM literature (i.e. y_{i,j} =… where i is the individual and j is the cluster/timepoint). I would recommend the authors bring this example back up briefly at each stage of the exposition, to tie each proposition back to something tangible. \n2) Relatedly, Figure 2 is fine, but doesn’t really help one represent the problem the authors are trying to deal with - there is no representation of the nested dimension (e.g. time) in this figure. Perhaps it would be worth including a second figure here (or modifying the current figure), showing more explicitly how the DGP contains a hierarchical structure. In my view, this wouldn’t be unusual, and honestly could even be argued to be one of the main purposes of PGMs. \n3) I’m curious/surprised about the choice of evaluations. In some ways, I feel like there was a strange mixture of conventional computer vision with health/medical (‘Time from disease event’ on top of some rotating MNIST images was particularly artificial). I think the image example is nice but should either (a) be replaced with an image example which is more relevant (medical imaging?) and/or (b) focus more on the clinical trial and physio examples.\n4) Relatedly, the evaluation metrics used make sense, but would there be a way to assess the extent to which key statistical parameters are better estimated with this method? As mentioned, biased standard errors are one of, if not <the> most important motivation for mixed modelling. If I were an empirical researcher with time series data, I’d want to use the authors’ method because of the flexible function form, but be a bit lost as to whether it actually does the job in terms of facilitating statistical inference on <simple> tasks which I might otherwise fall back on a conventional LMM analysis. Relatedly, where is the benchmark here for simple/linear tasks - if I use LMM-VAE instead of, say, the lmer package in R, how do they compare?\n5)  Could the authors comment on the Table 6 MCC results - why do the authors think the VAE is so close to the oLMM-VAE?\n\n\nMinor:\n\nClarify that LMMs are also known as, hierarchical linear models, multilevel models, mixed effects models (fixed and random effects).\nAdd equation numbers to equations on page 4.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors deal with an important gap in the literature - how to handle non-i.i.d. / clustered / longitudinal data whilst making minimal assumptions regarding the functional form. They propose a combination of the LMM with variational inference. The literature on linear mixed models is extensive, and as described by the authors, has wide applicability given the importance of accounting for the impact on variance associated with samples arising from shared processes (e.g. patients from the same hospital, measures from the same patient). Their proposed methods (oLMM-VAE and sLMM-VAE) are compared for their efficacy at missing value imputation, forecasting, and interpolation.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The requested changes can be inferred readily from my list of comments above.", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Reviewer hND8 particularly appreciated the clarity of the final version, ultimately recommending acceptance. Reviewer ZB6M valued the idea but suggested the novelty is perhaps incremental when viewed from the lens of probabilistic PCA and GP-LVMs. Reviewer xWFG found the contribution valuable but requested further evaluation to situate LMM-VAE within the landscape of classical LMMs. (See above for details)
Decision: Accept (poster)