PaperlayerView exact revision on arXiv

Source-linked AI summary

Latent mixed-effect models for high-dimensional longitudinal data

Priscilla Ong, Manuel Haußmann, Otto Lönnroth, Harri Lähdesmäki

arXiv:2409.11008v1cs.LGstat.APstat.ML

TL;DR

Longitudinal data can be high-dimensional, nonlinear, and time-varying, while VAE independence assumptions and GP prior costs leave modeling challenges. LMM-VAE uses linear mixed models to structure VAE priors and accommodate covariates. Its variants show competitive prediction results and test-MSE patterns across the described Rotating MNIST and Physionet evaluations.

  • Problem

    Longitudinal data contain complex covariates and repeated measurements, while standard VAEs assume independent observations and scalable incorporation of many covariates into longitudinal priors remains underexplored.

    Source: §1 · pp. 1–2

  • Method

    LMM-VAE uses linear mixed models to parameterize VAE priors, with variants supporting grouped observations and predictive training.

    Source: §1 · p. 2, §4.2 · p. 5, §4.3 · p. 5

  • Results

    LMM-VAE variants are competitive with LVAE on Physionet prediction, while Rotating MNIST test MSE decreases or remains stable as basis functions increase.

    Source: §7.3 · p. 8, §7.2 · p. 8

  • Takeaways & Limitations

    LMM-VAE offers an efficient, globally parameterized approximation to GP prior VAEs while retaining the ability to add basis functions.

    Source: §7.2 · p. 8

    Identifiability requires sufficiently diverse covariates, an injective decoder, and additional technical conditions; it holds only up to rotations and translations.

    Source: §5 · p. 6

Abstract

from arXiv · show

Modelling longitudinal data is an important yet challenging task. These datasets can be high-dimensional, contain non-linear effects and 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.

1. Introduction

Longitudinal data combine repeated measurements with rich covariates, but existing approaches face scaling and modeling challenges. LMM-VAE uses linear mixed models to parameterize VAE priors and is presented as a scalable, interpretable alternative connected to GP-based methods.

1. Introduction

  • Longitudinal datasets can be high-dimensional, nonlinear, and time-varying, while existing methods scale poorly for such data (Zipunnikov et al., 2014; Ramchandran et al., 2021).VAEs also assume independent and identically distributed observations, so they do not directly capture longitudinal correlations.Source: §1 · p. 1
  • GP prior VAEs impose latent correlation structure but have cubic complexity in sample count and may struggle to incorporate many auxiliary covariates (Fortuin et al., 2020; Casale et al., 2018; Ramchandran et al., 2021).The passage also notes that prior work has limited coverage of scalable longitudinal models with large covariate sets.Source: §1 · pp. 1–2
  • LMM-VAE models the VAE prior with a linear mixed model and basis functions, accommodating arbitrary covariates and both shared and random effects.The proposal is described as scalable and as simplifying training in standard deep-learning frameworks.Source: §1 · p. 2
  • The authors present LMM-VAE as competitive with commonly used GP-based methods and connect its theory to GP prior VAEs.They frame LMMs as a practical and theoretical alternative, particularly for high-dimensional covariates.Source: §1 · p. 2

2. Related Work

Related work extends VAE expressivity and structures latent representations, while GP prior VAEs are especially relevant to longitudinal data. The discussion highlights GP scaling and covariate-integration challenges that motivate LMM-VAE.

2. Related Work

  • Other VAE work expands posterior expressivity, uses informative priors, or targets disentangled latent spaces (Fortuin et al., 2020; Rezende and Mohamed, 2015).This paper instead incorporates auxiliary covariates into the prior, including time-varying side-profile information.Source: §2 · p. 2
  • Structured VAEs combine neural networks with probabilistic graphical models; LMM-VAE extends structured latent representations to arbitrarily many longitudinal covariates (Johnson et al., 2016).Sauty and Durrleman (2022) use a mixed-effect longitudinal VAE whose parametric structure is confined to disease progression and requires internal MCMC sampling.Source: §2 · p. 2
  • GP prior VAEs use GP expressiveness and smoothness for flexible, robust priors, but GP learning scales cubically with sample count.Approximations include Taylor methods and inducing points, whose optimization can complicate training and may reduce prior expressiveness (Casale et al., 2018; Ramchandran et al., 2021; Jazbec et al., 2021; Bauer et al., 2016; Hensman et al., 2013).Source: §2 · pp. 2–3
  • Integrating high-dimensional covariates into GP priors remains difficult: some approaches assume low-dimensional inputs, while kernel construction or training can limit performance (Jazbec et al., 2021; Casale et al., 2018; Ramchandran et al., 2021).The cited methods use low-dimensional covariates, object-and-view kernels, or additive kernels, respectively.Source: §2 · p. 3
  • Linear and linear mixed models are established tools for low-dimensional biomedical and longitudinal data and are also used with high-dimensional brain imaging.Table 1 is identified as a comparison of related methods.Source: §2 · p. 3, Table 1 · p. 3

3. Background

The background defines linear models and mixed effects, then describes VAE inference and conditional priors. Nonlinear basis functions extend linear covariate effects, while CVAEs condition the generative model and encoder on auxiliary information.

3. Background

  • A linear model represents a latent response as a covariate-weighted sum plus noise, written z = Ax + ϵ.The model uses a covariate vector x and coefficient matrix A; the residual is normally distributed with equal variance.Source: §3 · p. 3
  • Nonlinear effects for continuous covariates can be represented by basis functions such as x², sin(x), or cos(x), while categorical covariates can be one-hot encoded.The passage assumes each covariate is continuous or binary after encoding.Source: §3 · p. 3
  • LMMs combine shared effects with random effects specific to subsets of instances, such as patient-specific offsets or subgroup interactions.The covariates are divided into shared-effect and random-effect sets.Source: §3 · pp. 3–4
  • A VAE specifies a generative model pω(y, z) = pθ(y|z)pφ(z) and uses an encoder to approximate the intractable posterior.The encoder’s parameters are optimized jointly with the generative-model parameters using an evidence lower bound.Source: §3 · p. 4
  • CVAEs condition the generative model and encoder on auxiliary covariates, and GP prior VAEs use a GP for the conditional latent prior.The background gives pω(y, z|x) = pθ(y|z)pφ(z|x) as the GP prior VAE form.Source: §3 · p. 4

4. Linear Mixed Model VAEs

LMM-VAE incorporates linear mixed models into VAE priors, supporting arbitrary auxiliary covariates, shared and random effects, and efficient learning through global parametrization. The section presents ordinary and structured variants, plus a predictive training routine for conditional generation.

4. Linear Mixed Model VAEs

  • LMM-VAE places an LMM in the VAE prior, accommodating arbitrary auxiliary covariates and shared and random effects through global parametrization.The paper proposes two variants, collectively called LMM-VAE.Source: §4 · p. 4

4.1. Ordinary LMM-VAE

  • The ordinary variant models latent variables with a standard LMM and supports minibatch SGD through global parametrization.Its approximate posteriors are Gaussian; the first ELBO expectation can be estimated by Monte Carlo, while remaining terms are analytically tractable under Gaussian priors.Source: §4.1 · p. 4

4.2. Structured LMM-VAE

  • The structured variant shares latent variation across observations in each group, which the paper assumes corresponds to a single instance such as a patient.Its encoder must account for this shared variation.Source: §4.2 · p. 5

4.3. Predictive training for LMM-VAE

  • GSNN-based predictive training sets the variational approximation equal to the prior and trains the prior for data generation.The paper says this approach could yield competitive generation performance when auxiliary covariates are sufficiently informative.Source: §4.3 · p. 5

5. Identifiability

The paper applies identifiability conditions for conditional latent-variable models to both LMM-VAE variants. Under the stated assumptions, the model is identifiable up to rotations and translations.

5. Identifiability

  • Both LMM-VAE variants meet the stated constraints when the decoder is injective and the other required technical conditions hold.The conditions include L < D and sufficiently diverse covariates with at least kL + 1 distinct values.Source: §5 · p. 6
  • The identifiability result requires the observation model to decompose as y = f(z) + ϵ, with ϵ drawn from p(ϵ).The model also requires sufficiently diverse covariates and additional technical conditions summarized in Appendix D.Source: §5 · p. 6
  • Under these conditions, (f, T, λ) is identifiable up to rotations and translations.The result follows the conditional latent-variable structure attributed to Khemakhem et al. (2020).Source: §5 · p. 6

6. GP prior VAEs as LMM-VAEs

The paper connects LMM-VAEs to GP prior VAEs through finite basis expansions and a spectral representation. This connection offers a GP-prior implementation within LMM-VAE and retains a simpler training procedure through global parametrization.

6. GP prior VAEs as LMM-VAEs

  • Adding basis functions can make LMMs arbitrarily complex while keeping the model linear in its parameters.The paper links this construction to the result that GPs correspond to Bayesian linear regression with infinitely many basis functions.Source: §6 · p. 6
  • A stationary covariance can be approximated with a finite basis expansion using regular or random Fourier frequencies.The paper constructs the features from sine and cosine terms using the spectral density.Source: §6 · p. 6
  • A reduced-rank GP prior VAE with M Fourier features can be implemented directly by LMM-VAE by setting z = Aϕ(x).The paper also notes a connection between LMM-VAE and LVAE.Source: §6 · p. 6
  • The spectral representation and global parametrization let LMM-VAE simplify training of GP-connected models through straightforward SGD optimization.The paper contrasts this with the cubic sample scaling of GP learning described in its related-work discussion.Source: §6 · p. 6

7. Experiments

Experiments assess imputation, future prediction, and time-based interpolation by MSE on Health MNIST, Rotating MNIST, and Physionet Challenge 2012.

7. Experiments

  • The evaluation compares methods across three tasks, using similar encoder and decoder architectures for consistency.Source: §7 · pp. 6–7

7.1. Health MNIST

  • LMM-VAE imputation MSE remains robust across linear-model parameterizations and outperforms GP-based baselines.Source: §7.1 · p. 7, Table 2
  • LMM-VAE has lower conditional test MSE than GPP-VAE and SVGP-VAE, and is competitive with LVAE when covariates match.Source: Figure 3 · p. 7, §7.1 · p. 7

7.2. Rotating MNIST

  • On Rotating MNIST interpolation, sLMM-VAE matches or improves on GP counterparts, while GSNN training performs better on deterministic data.Source: Table 3 · p. 8, §7.2 · p. 8
  • On Rotating MNIST, test MSE decreases or remains stable as Fourier basis functions are added.The basis expansion is described as an efficient approximation to GP-prior VAEs that preserves global parametrization and efficient training.Source: Figure 4 · p. 8, §7.2 · p. 8

7.3. Physionet Challenge 2012

8. Conclusion and Future Work

LMM-VAE models high-dimensional longitudinal data with numerous covariates, and the paper reports that auxiliary covariates matter empirically. Its analysis connects the model to GP-based VAEs and identifies future extensions.

8. Conclusion and Future Work

  • LMM-VAE is presented as a method for high-dimensional longitudinal data that scales to large datasets with numerous covariates.The conclusion also describes it as novel.Source: §8 · p. 8
  • Empirical evaluation highlights the importance of modelling auxiliary covariates and the model’s potential for longitudinal analyses.The conclusion does not claim that this potential is established for every longitudinal task.Source: §8 · pp. 8–9
  • Theoretical analysis connects LMM-VAE to GP-based VAEs as a reduced-rank approximation, with basis functions adaptable to different modelling tasks.Future work includes counterfactual reasoning and alternative neural network-based priors.Source: §8 · p. 9

A.1. sLMM-VAE encoder architecture

The structured LMM-VAE amortizes inference while accounting for shared latent variation, then uses a factorized Gaussian variational posterior to obtain a tractable ELBO. Its formulation shares a linear model across encoder and generative model.

A.1. sLMM-VAE encoder architecture

  • The sLMM-VAE encoder accounts for shared latent-space variation by sharing linear model parameters A with the generative model.The encoder maps each observation y_n to variational parameters, including μ_n and σ_n².Source: §A.1 · p. 12

A.2. sLMM-VAE ELBO formulation

  • A factorized Gaussian variational posterior yields a tractable sLMM-VAE ELBO, with z_n formed from A x_n and group-specific noise.The model may treat A and θ as hyperparameters by dropping their prior distributions.Source: §A.2 · p. 12
  • The formulation may perform competitively when auxiliary covariates sufficiently characterize the output.The noise variable ε is identical for observations from the same group, and A is learned as part of the decoder rather than only as part of the prior.Source: §A.2 · p. 12

Appendix B. Gaussian stochastic neural network training objective for LMM-VAEs

The appendix relates the Gaussian stochastic neural network objective to the ELBO, describes LMM-VAE as a linear-basis route to GP-prior VAEs, and discusses its identifiability under conditional-VAE assumptions.

Appendix B. Gaussian stochastic neural network training objective for LMM-VAEs

  • Aligning the variational approximation with the prior makes the KL term zero, so the GSNN objective is the ELBO when q equals the prior.The objective supports conditional generation when covariates are sufficiently informative about the corresponding output.Source: Appendix B · pp. 12–13
  • GSNN training resembles regression: a covariate-conditioned linear layer with noise feeds a decoder trained by expected log-likelihood.The objective does not estimate latent variables because it is designed for the generative task only.Source: Appendix B · p. 13

Appendix C. GP prior VAEs as LMM-VAEs

  • LVAE additive latent components can be represented by linear models over Fourier features, connecting GP-prior VAEs to LMM-VAE.Each additive component uses its corresponding covariate subset and Fourier features.Source: Appendix C · p. 14

Appendix D. Identifiability

  • LMM-VAE inherits conditional-VAE identifiability, and the paper links this property to interpretable coefficients and data-generation mechanisms.The stated motivation is that identifiability supports parameter interpretation for downstream longitudinal analyses.Source: Appendix D · p. 14
  • Under the theorem’s conditions and a sufficiently rich variational posterior, the paper states an asymptotic identifiability guarantee.The theorem characterizes identifiability up to specified transformations.Source: Appendix D · p. 15

D.1. Clinical Trial Data and Health MNIST

The authors assess identifiability using latent-space correlations alongside predictive metrics on clinical-trial data and Health MNIST. They report competitive LMM-VAE results across multiple latent dimensions, while noting that correlation alone can mislead.

D.1. Clinical Trial Data and Health MNIST

  • Mean correlation coefficient compares true latent variables with estimated posteriors, but computing it requires ground-truth latents unavailable in realistic practice.The authors describe comparing correlations across training runs as a common approximation and pair it with a predictive metric.Source: §D.1 · p. 15
  • LMM-VAE reports competitive results across multiple latent dimensions on clinical-trial data and Health MNIST.At least one other model serves as a baseline in each experimental setup.Source: Table 5 · p. 16, Table 6 · p. 16, §D.1 · p. 16
  • Run-to-run latent correlation can be perfect at a shared bad local optimum, so identifiability assessment remains conditional on predictive performance.The authors evaluate a variant of MCC alongside NLL or MSE.Source: §D.1 · p. 16

E.1. Dataset Description

The appendices describe synthetic and clinical evaluation settings and provide supplementary illustrations of predictions, id-level variation, and basis-function trends.

E.1. Dataset Description

  • Health MNIST includes age-related shifts, disease-dependent rotations, noisy observations, and 25% masked pixels, with the final 15 timepoints of 100 subjects held out for testing.The dataset contains 1,300 instances, each with 20 timepoints.Source: §E.1 · p. 16

E.3. Additional illustrations of LMM-VAE’s predictions

  • Health MNIST illustrations show test-set reconstructions and predictions for two biological-sex groups alongside noisy inputs and original images.The illustrated oLMM-VAE uses 16 latent dimensions.Source: Figure 7 · p. 17, §E.3 · p. 17

F.1. Dataset Description

  • Rotating MNIST contains 400 digit instances rotated through 16 angles; test sequences use four consecutive angles, and validation ids also occur in training and testing.The revised split is intended to evaluate learned id-specific parameters.Source: §F.1 · pp. 17–18

F.3. Additional illustrations on Interpretability

  • Rotating MNIST trajectories separate shared angular rotation from id-level variation, while different noise realizations correspond to different digit identities.The first two basis functions are sin(θ) and cos(θ); conditioning on an instance-specific ϵ(k) makes the latent mapping deterministic in the auxiliary covariates.Source: §F.3 · p. 19

G.1. Dataset Description

  • The patient dataset contains 3,599 training instances, each with at least 15 observations, and withholds the last 10 timepoints of 1,200 patients for testing.Train and validation sets contain 3,599 and 3,598 unique instances, respectively.Source: §G.1 · p. 19

F.4. Further experiments for Rotating MNIST

  • MSE generally decreases as the number of basis functions increases in the reported 32-dimensional oLMM-VAE and sLMM-VAE experiments.The sLMM-VAE VI experiment also reports this trend at latent dimension 32; θ and A are deterministic in the GSNN experiments.Source: §F.4 · p. 19, Figure 12 · p. 21, Figure 13 · p. 21

H.1. Dataset Description

  • The clinical-trial dataset comprises 480 subjects and 6,605 observations; preprocessing combines laboratory and vital-sign trajectories with demographic, adverse-event, and medication covariates.Timepoints with under 70% of outcome information and patient trajectories shorter than five observations were excluded.Source: §H.1 · p. 22

Appendix J. Model Architectures

  • The appendices direct readers to separate architecture tables for Health MNIST, Rotating MNIST, Physionet, and Project Data Sphere experiments.The Rotating MNIST architecture differs from the cited prior architecture in using a sigmoid final activation instead of ELU.Source: Appendix J · p. 23, Table 10 · p. 24, Table 11 · p. 24
Loading 2409.11008v1…