XFACTORS: Disentangled Information Bottleneck via Contrastive Supervision
Alexandre Myara, Nicolas Bourriez, Thomas Boyer, Thomas Lemercier, Ihab Bendidi, Auguste Genovesio
TL;DR
XFactors addresses the challenge of disentangled representation learning under weak supervision by partitioning the latent space into a residual subspace $\mathcal{S}$ and factor-specific subspaces $\mathcal{T}_i$, with each factor $y_{f_i}$ encoded in $\mathcal{T}_i$ via InfoNCE contrastive losses. The approach combines the Information Bottleneck principle with a direct-sum latent decomposition, enforcing Gaussian priors on $\mathcal{S}$ and the aggregate $\bigoplus_i \mathcal{T}_i$ while maximizing mutual information between each $\mathcal{T}_i$ and its corresponding label and minimizing cross-information between $\mathcal{S}$ and $\mathcal{T}$. Training uses two encoders $\psi_s$ and $\psi_t$ to produce $\mathcal{S}$ and $\mathcal{T}$, a decoder $\phi$ for reconstruction, and a composite loss $\mathcal{L} = \mathcal{L}_{\text{reco}} + \beta_s \mathcal{L}_{\text{KL}}^{\mathcal{S}} + \beta_t \mathcal{L}_{\text{KL}}^{\mathcal{T}} + \sum_i \lambda_i \mathcal{L}_{\text{InfoNCE}}^{(i)}$. Empirically, XFactors achieves state-of-the-art disentanglement scores across multiple datasets, supports controlled factor swapping via latent replacement, scales with increasing latent capacity, and demonstrates competitive results on CelebA, with code available at the provided repository. A key limitation is reconstruction quality due to the standard VAE reconstruction–KL trade-off, suggesting future work to integrate stronger generative heads. Overall, the method offers a scalable, stable, and controllable route to targeted disentanglement without adversarial training or extensive classifiers, with broad potential applications in robust counterfactual generation and batch-effect handling.
Abstract
Disentangled representation learning aims to map independent factors of variation to independent representation components. On one hand, purely unsupervised approaches have proven successful on fully disentangled synthetic data, but fail to recover semantic factors from real data without strong inductive biases. On the other hand, supervised approaches are unstable and hard to scale to large attribute sets because they rely on adversarial objectives or auxiliary classifiers. We introduce \textsc{XFactors}, a weakly-supervised VAE framework that disentangles and provides explicit control over a chosen set of factors. Building on the Disentangled Information Bottleneck perspective, we decompose the representation into a residual subspace $\mathcal{S}$ and factor-specific subspaces $\mathcal{T}_1,\ldots,\mathcal{T}_K$ and a residual subspace $\mathcal{S}$. Each target factor is encoded in its assigned $\mathcal{T}_i$ through contrastive supervision: an InfoNCE loss pulls together latents sharing the same factor value and pushes apart mismatched pairs. In parallel, KL regularization imposes a Gaussian structure on both $\mathcal{S}$ and the aggregated factor subspaces, organizing the geometry without additional supervision for non-targeted factors and avoiding adversarial training and classifiers. Across multiple datasets, with constant hyperparameters, \textsc{XFactors} achieves state-of-the-art disentanglement scores and yields consistent qualitative factor alignment in the corresponding subspaces, enabling controlled factor swapping via latent replacement. We further demonstrate that our method scales correctly with increasing latent capacity and evaluate it on the real-world dataset CelebA. Our code is available at \href{https://github.com/ICML26-anon/XFactors}{github.com/ICML26-anon/XFactors}.