Can Diffusion Models Disentangle? A Theoretical Perspective

Liming Wang (Massachusetts Institute of Technology) · Muhammad Jehanzeb Mirza (Massachusetts Institute of Technology) · Yishu Gong (Takeda) · Yuan Gong (Massachusetts Institute of Technology) · Jiaqi Zhang (Chongqing University) · Brian Tracey (Takeda) · Katerina Placek (Takeda) · Marco Vilela (Analog Devices) · Jim Glass (Massachusetts Institute of Technology)

diffusion models disentangled representationsfinite-sample global convergencegaussian mixture modelsidentifiability conditionsindependent subspace modelslatent variable modelsmultiple viewsnon-invertible mixing processespartial labelsstochastic mixing processesstyle guidance regularizationsubspace recoverytheoretical framework weak supervision

This paper presents a novel theoretical framework for understanding how diffusion models can learn disentangled representations with commonly used weak supervision such as partial labels and multiple views. Within this framework, we establish identifiability conditions for diffusion models to disentangle latent variable models with \emph{stochastic}, \emph{non-invertible} mixing processes. We also prove \emph{finite-sample global convergence} for diffusion models to disentangle independent subspace models. To validate our theory, we conduct extensive disentanglement experiments on subspace recovery in latent subspace Gaussian mixture models, image colorization, denoising, and voice conversion for speech classification. Our experiments show that training strategies inspired by our theory, such as style guidance regularization, consistently enhance disentanglement performance.