Integrating Random Effects in Variational Autoencoders for Dimensionality Reduction of Correlated Data

TL;DR

This paper introduces LMMVAE, a novel variational autoencoder that incorporates random effects to better handle correlated data, improving reconstruction and downstream classification performance.

Contribution

LMMVAE extends traditional VAEs by integrating random effects inspired by linear mixed models, effectively modeling correlations in data.

Findings

Significant reduction in reconstruction error on real and simulated datasets.

Improved negative likelihood loss on unseen data.

Enhanced downstream classification performance.

Abstract

Variational Autoencoders (VAE) are widely used for dimensionality reduction of large-scale tabular and image datasets, under the assumption of independence between data observations. In practice, however, datasets are often correlated, with typical sources of correlation including spatial, temporal and clustering structures. Inspired by the literature on linear mixed models (LMM), we propose LMMVAE -- a novel model which separates the classic VAE latent model into fixed and random parts. While the fixed part assumes the latent variables are independent as usual, the random part consists of latent variables which are correlated between similar clusters in the data such as nearby locations or successive measurements. The classic VAE architecture and loss are modified accordingly. LMMVAE is shown to improve squared reconstruction error and negative likelihood loss significantly on unseen data, with simulated as well as real datasets from various applications and correlation scenarios. It also shows improvement in the performance of downstream tasks such as supervised classification on the learned representations.