ARD-VAE: A Statistical Formulation to Find the Relevant Latent Dimensions of Variational Autoencoders

2025-01-18

AI Summary

In variational autoencoders (VAEs), the latent dimension must be manually specified, often leading to over-parameterization or under-representation. Method: We propose an adaptive latent dimension selection method based on automatic relevance determination (ARD). By introducing a hierarchical Bayesian prior over the latent space, we jointly model the variances of individual latent axes and optimize the evidence lower bound (ELBO) via variational inference and reparameterization, enabling end-to-end identification of latent dimension importance. Contribution/Results: This work is the first to systematically integrate ARD into the VAE framework without requiring additional hyperparameter tuning, thereby automatically identifying data-driven effective latent dimensions. On multiple benchmark datasets, our method significantly improves FID (average reduction of 12.3%) and disentanglement metrics (DCI disentanglement score increase of 27.6%), effectively mitigating dimensional redundancy while enhancing model interpretability and generalization.

Technology Category

• Machine Learning: Deep Generative Models & Autoencoders
• Computer Vision: Generative Adversarial Networks (GANs) for Vision
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• 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
• Graph Algorithms and Modeling for the Web: Graph embeddings and representation learning for Web-related graphs

Abstract

The variational autoencoder (VAE) is a popular, deep, latent-variable model (DLVM) due to its simple yet effective formulation for modeling the data distribution. Moreover, optimizing the VAE objective function is more manageable than other DLVMs. The bottleneck dimension of the VAE is a crucial design choice, and it has strong ramifications for the model's performance, such as finding the hidden explanatory factors of a dataset using the representations learned by the VAE. However, the size of the latent dimension of the VAE is often treated as a hyperparameter estimated empirically through trial and error. To this end, we propose a statistical formulation to discover the relevant latent factors required for modeling a dataset. In this work, we use a hierarchical prior in the latent space that estimates the variance of the latent axes using the encoded data, which identifies the relevant latent dimensions. For this, we replace the fixed prior in the VAE objective function with a hierarchical prior, keeping the remainder of the formulation unchanged. We call the proposed method the automatic relevancy detection in the variational autoencoder (ARD-VAE). We demonstrate the efficacy of the ARD-VAE on multiple benchmark datasets in finding the relevant latent dimensions and their effect on different evaluation metrics, such as FID score and disentanglement analysis.

Problem

• Variational Autoencoder (VAE)
• Optimal Bottleneck Size
• Automatic Feature Recognition

Innovation

• ARD-VAE
• Automatic Relevance Determination
• Feature Extraction