Sculpting Latent Spaces With MMD: Disentanglement With Programmable Priors

2025-10-13

AI Summary

This work identifies a fundamental limitation of KL-divergence-based prior regularization in variational autoencoders (VAEs): it fails to reliably enforce posterior aggregation toward a factorized Gaussian prior, resulting in entangled latent representations. To address this, we propose a programmable prior framework grounded in the Maximum Mean Discrepancy (MMD), enabling flexible modeling of complex, semantically aligned priors within the VAE architecture. Furthermore, we introduce an unsupervised Latent Predictability Score (LPS) to quantitatively assess disentanglement. Experiments on CIFAR-10 and Tiny ImageNet demonstrate that our method achieves state-of-the-art mutual information-based disentanglement performance while preserving high-fidelity reconstructions—thereby resolving the inherent reconstruction–disentanglement trade-off prevalent in conventional disentanglement approaches.

Technology Category

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

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

Learning disentangled representations, where distinct factors of variation are captured by independent latent variables, is a central goal in machine learning. The dominant approach has been the Variational Autoencoder (VAE) framework, which uses a Kullback-Leibler (KL) divergence penalty to encourage the latent space to match a factorized Gaussian prior. In this work, however, we provide direct evidence that this KL-based regularizer is an unreliable mechanism, consistently failing to enforce the target distribution on the aggregate posterior. We validate this and quantify the resulting entanglement using our novel, unsupervised Latent Predictability Score (LPS). To address this failure, we introduce the Programmable Prior Framework, a method built on the Maximum Mean Discrepancy (MMD). Our framework allows practitioners to explicitly sculpt the latent space, achieving state-of-the-art mutual independence on complex datasets like CIFAR-10 and Tiny ImageNet without the common reconstruction trade-off. Furthermore, we demonstrate how this programmability can be used to engineer sophisticated priors that improve alignment with semantically meaningful features. Ultimately, our work provides a foundational tool for representation engineering, opening new avenues for model identifiability and causal reasoning.

Problem

• Addressing KL divergence's failure in disentangled representation learning
• Introducing MMD-based framework to sculpt latent spaces programmatically
• Achieving state-of-the-art independence without reconstruction trade-offs

Innovation

• Uses Maximum Mean Discrepancy for latent sculpting
• Introduces Programmable Prior Framework for disentanglement
• Enables explicit latent space engineering without reconstruction trade-off