Sculpting Latent Spaces With MMD: Disentanglement With Programmable Priors
This addresses the challenge of unreliable disentanglement in machine learning, offering a foundational tool for representation engineering that could impact model identifiability and causal reasoning, though it is incremental in improving upon existing VAE methods.
The paper tackles the problem of learning disentangled representations by showing that the standard KL divergence penalty in VAEs fails to enforce the target distribution, and introduces a Programmable Prior Framework using MMD to sculpt latent spaces, achieving state-of-the-art mutual independence on datasets like CIFAR-10 and Tiny ImageNet without reconstruction trade-offs.
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.