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

2024-12-22

AI Summary

Traditional VAEs assume independent observations, limiting their ability to model intrinsically correlated data—such as spatiotemporal or clustered structures. To address this, we propose REM-VAE, a VAE framework integrating random effects. Its core innovation is the first incorporation of linear mixed-model principles into VAEs: it explicitly disentangles fixed effects (individual-specific latent variables) from random effects (cluster-level correlations) and employs a learnable covariance structure to capture inter-sample dependencies. Accordingly, we reformulate the ELBO objective to incorporate a structured prior over random effects. Evaluated on synthetic benchmarks and multiple real-world spatiotemporal and population datasets, REM-VAE achieves significant improvements: −12.3% reduction in reconstruction error, −9.7% decrease in negative log-likelihood, and an average +4.8% gain in downstream classification accuracy. These results demonstrate REM-VAE’s effectiveness and generalization advantage in modeling correlated data.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Reasoning under Uncertainty: Relational Probabilistic Models
• Computer Vision: Diffusion Models for Vision

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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.

Problem

• Variational Autoencoder (VAE)
• Intrinsic Correlation
• Performance Enhancement

Innovation

• LMMVAE
• Linear Mixed Models
• Variational Autoencoders