Variational Learning of Disentangled Representations

Disentangled representations separate factors that are shared across conditions from those that are condition-specific. Such separation is needed for generalization to new domains, treatments, patients, or species. A dominant line of work pursues this goal through variational formulations. While these approaches achieve partial disentanglement, they often exhibit three common limitations: they either do not remove all condition-specific information from the shared representation, allow the shared representation to become uninformative, or impose independence assumptions that do not reflect the underlying generative process. In this work, we introduce DisCoVR, a variational framework that addresses these limitations. Its objective is aligned with the probabilistic structure of the data-generating process, and includes an adversarial term that prevents condition-specific information from being encoded in the shared representation. DisCoVR reconstructs the data from both shared and condition-specific representations, ensuring that each remains informative, and uses a structured prior that further reinforces the informativeness of both representations. We show that across synthetic, image, and single-cell RNA-sequencing datasets, DisCoVR achieves stronger disentanglement compared to previous approaches.

Machine-learning models are often trained on data collected under different conditions, such as patients, sites, treatments, or species. In these settings, generalization to new conditions requires separating structure that is shared across conditions from variation that is specific to each condition. Existing variational methods only partially achieve this: the shared representation can still contain condition information, the condition-specific representation can become uninformative, or the model assumptions can misalign with the data-generating process. We introduce DisCoVR, a variational framework that learns two representations: one condition-invariant and one condition-aware. The method uses a probabilistic objective matched to the assumed data-generating structure, two reconstruction paths to keep both representations informative, a structured prior based on class-wise averages of the invariant representation, and an adversarial term that discourages condition information from leaking into the invariant representation. The resulting optimization is formulated as a max-min game and has a unique equilibrium under the stated assumptions. Across synthetic data, and real-world applications, DisCoVR reconstructs the data well while more cleanly separating shared structure from condition-specific effects than previous approaches.