Diffusion probabilistic models (DPMs) have become the state-of-the-art in high-quality image generation. However, DPMs have an arbitrary noisy latent space with no interpretable or controllable semantics. Although there has been significant research effort to improve image sample quality, there is little work on representation-controlled generation using diffusion models. Specifically, causal modeling and controllable counterfactual generation using DPMs is an underexplored area. In this work, we propose CausalDiffAE, a diffusion-based causal representation learning framework to enable counterfactual generation according to a specified causal model. Our key idea is to use an encoder to extract high-level semantically meaningful causal variables from high-dimensional data and model stochastic variation using reverse diffusion. We propose a causal encoding mechanism that maps high-dimensional data to causally related latent factors and parameterize the causal mechanisms among latent factors using neural networks. To enforce the disentanglement of causal variables, we formulate a variational objective and leverage auxiliary label information in a prior to regularize the latent space. We propose a DDIM-based counterfactual generation procedure subject to do-interventions. Finally, to address the limited label supervision scenario, we also study the application of CausalDiffAE when a part of the training data is unlabeled, which also enables granular control over the strength of interventions in generating counterfactuals during inference. We empirically show that CausalDiffAE learns a disentangled latent space and is capable of generating high-quality counterfactual images.

翻译：扩散概率模型（DPMs）已成为高质量图像生成领域的最先进技术。然而，DPMs拥有任意噪声潜变量空间，缺乏可解释或可控的语义信息。尽管已有大量研究工作致力于提升图像样本质量，但针对扩散模型进行表征控制生成的研究仍十分有限。具体而言，基于DPMs的因果建模与可控反事实生成是一个尚未被充分探索的领域。本文提出CausalDiffAE——一种基于扩散的因果表征学习框架，可根据指定因果模型实现反事实生成。我们的核心思想是通过编码器从高维数据中提取具有高级语义的因果变量，并利用反向扩散过程对随机变异性进行建模。我们设计了一种因果编码机制，将高维数据映射至具有因果关联的潜变量因子，并采用神经网络参数化潜变量因子间的因果机制。为强化因果变量的解耦性，我们构建了变分目标函数，并利用先验中的辅助标签信息正则化潜变量空间。此外，我们提出了一种基于DDIM的介入操作驱动的反事实生成流程。最后，针对标签监督受限场景，我们进一步研究了部分训练数据无标签情况下CausalDiffAE的应用，该方法可在推理过程中实现对干预强度的精细控制以生成反事实样本。实验结果表明，CausalDiffAE能够学习解耦的潜变量空间，并生成高质量的反事实图像。