Title: DisDiff: Unsupervised Disentanglement of Diffusion Probabilistic Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: Diffusion Probabilistic Model

Disentangled representation
Abstract: Targeting to understand the underlying explainable factors behind observations and modeling the conditional generation process on these factors, we connect disentangled representation learning to diffusion probabilistic models (DPMs) to take advantage of the remarkable modeling ability of DPMs. We propose a new task, disentanglement of (DPMs): given a pre-trained DPM, without any annotations of the factors, the task is to automatically discover the inherent factors behind the observations and disentangle the gradient fields of DPM into sub-gradient fields, each conditioned on the representation of each discovered factor. With disentangled DPMs, those inherent factors can be automatically discovered, explicitly represented and clearly injected into the diffusion process via the sub-gradient fields. To tackle this task, we devise an unsupervised approach, named DisDiff, and for the first time achieving disentangled representation learning in the framework of DPMs. Extensive experiments on synthetic and real-world datasets demonstrate the effectiveness of DisDiff.

DisDiff: Unsupervised Disentanglement of Diffusion Probabilistic Models

Abstract

Targeting to understand the underlying explainable factors behind observations and modeling the conditional generation process on these factors, we connect disentangled representation learning to Diffusion Probabilistic Models (DPMs) to take advantage of the remarkable modeling ability of DPMs. We propose a new task, disentanglement of (DPMs): given a pre-trained DPM, without any annotations of the factors, the task is to automatically discover the inherent factors behind the observations and disentangle the gradient fields of DPM into sub-gradient fields, each conditioned on the representation of each discovered factor. With disentangled DPMs, those inherent factors can be automatically discovered, explicitly represented, and clearly injected into the diffusion process via the sub-gradient fields. To tackle this task, we devise an unsupervised approach named DisDiff, achieving disentangled representation learning in the framework of DPMs. Extensive experiments on synthetic and real-world datasets demonstrate the effectiveness of DisDiff.

1 Introduction

As one of the most successful generative models, diffusion probabilistic models (DPMs) achieve remarkable performance in image synthesis. They use a series of probabilistic distributions to corrupt images in the forward process and train a sequence of probabilistic models converging to image distribution to reverse the forward process. Despite the remarkable success of DPM in tasks such as image generation [41], text-to-images [37], and image editing [32], little attention has been paid to representation learning [49] based on DPM. Diff-AE [33] and PDAE [49] are recently proposed methods for representation learning by reconstructing the images in the DPM framework. However, the learned latent representation can only be interpreted relying on an extra pre-trained with predefined semantics linear classifier. In this paper, we will refer DPMs exclusively to the Denoising Diffusion Probabilistic Models (DDPMs) [19].

On the other hand, disentangled representation learning [17] aims to learn the representation of the underlying explainable factors behind the observed data and is thought to be one of the possible ways for AI to understand the world fundamentally. Different factors correspond to different kinds of image variations, respectively and independently. Most of the methods learn the disentangled representation

based on generative models, such as VAE [16, 5, 23] and GAN [27]. The VAE-based methods have an inherent trade-off between the disentangling ability and generating quality [16, 5, 23]. The GANbased methods suffer from the problem of reconstruction due to the difficulty of gan-inversion [44].

In this paper, we connect DPM to disentangled representation learning and propose a new task: the disentanglement of DPM. Given a pre-trained DPM model, the goal of disentanglement of DPM is to learn disentangled representations for the underlying factors in an unsupervised manner, and learn the corresponding disentangled conditional sub-gradient fields, with each conditioned on the representation of each discovered factor, as shown in Figure 1.

The benefits of the disentanglement of DPM are two-fold: ( i ) It enables totally unsupervised controlling of images by automatically discovering the inherent semantic factors behind the image data. These factors help to extend the DPM conditions information from human-defined ones such as annotations [49]/image-text pairs [21], or supervised pre-trained models [22] such as CLIP [34]. One can also flexibly sample partial conditions on the part of the information introduced by the superposition of the sub-gradient field, which is novel in existing DPM works. ( ii ) DPM has remarkable performance on image generation quality and is naturally friendly for the inverse problem, e.g., the inversion of DDIM [39], PDAE. Compared to VAE (trade-off between the disentangling ability and generating quality) or GAN (problem of gan-inversion), DPM is a better framework for disentangled representation learning. Besides, as Locatello et al. [30] points out, other inductive biases should be proposed except for total correlation. DPM makes adopting constraints from all different timesteps possible as a new type of inductive bias. Further, as Srivastava et al. [42] points out, and the data information includes: factorized and non-factorized. DPM has the ability to sample non-factorized (non-conditioned) information [18], which is naturally fitting for disentanglement.

To address the task of disentangling the DPM, we further propose an unsupervised solution for the disentanglement of a pre-trained DPM named DisDiff. DisDiff adopts an encoder to learn the disentangled presentation for each factor and a decoder to learn the corresponding disentangled conditional sub-gradient fields. We further propose a novel Disentangling Loss to make the encoded representation satisfy the disentanglement requirement and reconstruct the input image.

Our main contributions can be summarized as follows:

We present a new task: disentanglement of DPM, disentangling a DPM into several disentangled sub-gradient fields, which can improve the interpretability of DPM.

We build an unsupervised framework for the disentanglement of DPM, DisDiff, which learns a disentangled representation and a disentangled gradient field for each factor.

We propose enforcing constraints on representation through the diffusion model's classifier guidance and score-based conditioning trick.

2 Related Works

Diffusion Probabilistic Models DPMs have achieved comparable or superior image generation quality [38, 40, 19, 41, 20] than GAN [13]. Diffusion-based image editing has drawn much attention, and there are mainly two categories of works. Firstly, image-guided works edit an image by mixing the latent variables of DPM and the input image [7, 31, 32]. However, using images to specify the attributes for editing may cause ambiguity, as pointed out by Kwon et al. [25]. Secondly, the classifierguided works [9, 1, 28] edit images by utilizing the gradient of an extra classifier. These methods require calculating the gradient, which is costly. Meanwhile, these methods require annotations or models pre-trained with labeled data. In this paper, we propose DisDiff to edit the image in an unsupervised way. On the other hand, little attention has been paid to representation learning in the literature on the diffusion model. Two related works are Diff-ae [33] and PDAE [49]. Diff-ae [33] proposes a diffusion-based auto-encoder for image reconstruction. PDAE [49] uses a pre-trained DPM to build an auto-encoder for image reconstruction. However, the latent representation learned by these two works does not respond explicitly to the underlying factors of the dataset.

Disentangled Representation Learning Bengio et al. [2] introduced disentangled representation learning. The target of disentangled representation learning is to discover the underlying explanatory factors of the observed data. The disentangled representation is defined as each dimension of the representation corresponding to an independent factor. Based on such a definition, some V AE-based

works achieve disentanglement [5, 23, 16, 3] only by the constraints on probabilistic distributions of representations. Locatello et al. [30] point out the identifiable problem by proving that only these constraints are insufficient for disentanglement and that extra inductive bias is required. For example, Yang et al. [47] propose to use symmetry properties modeled by group theory as inductive bias. Most of the methods of disentanglement are based on VAE. Some works based on GAN, including leveraging a pre-trained generative model [36]. Our DisDiff introduces the constraint of all time steps during the diffusion process as a new type of inductive bias. Furthermore, DPM can sample non-factorized (non-conditioned) information [18], naturally fitting for disentanglement. In this way, we achieve disentanglement for DPMs.

3 Background

3.1 Diffusion Probabilistic Models (DPM)

We take DDPM [19] as an example. DDPM adopts a sequence of fixed variance distributions q ( x t | x t -1 ) as the forward process to collapse the image distribution p ( x 0 ) to N (0 , I ) . These distributions are

We can then sample x t using the following formula x t ∼ N ( x t ; ¯ α t x 0 , (1 -¯ α t )) , where α t = 1 -β t and ¯ α t = Π t t =1 α t , i.e., x t = √ ¯ α t x 0 + √ 1 -¯ α t ϵ . The reverse process is fitted by using a sequence of Gaussian distributions parameterized by θ :

where µ θ ( x t , t ) is parameterized by an Unet ϵ θ ( x t , t ) , it is trained by minimizing the variational upper bound of negative log-likelihood through:

3.2 Representation learning from DPMs

The classifier-guided method [9] uses the gradient of a pre-trained classifier, ∇ x t log p ( y | x t ) , to impose a condition on a pre-trained DPM and obtain a new conditional DPM: N ( x t ; µ θ ( x t , t ) + σ t ∇ x t log p ( y | x t ) , σ t ) . Based on the classifier-guided sampling method, PDAE [49] proposes an approach for pre-trained DPM by incorporating an auto-encoder. Specifically, given a pre-trained DPM, PDAE introduces an encoder E ϕ to derive the representation by equation z = E ϕ ( x 0 ) . They use a decoder estimator G ψ ( x t , z, t ) to simulate the gradient field ∇ x t log p ( z | x t ) for reconstructing the input image.

By this means, they create a new conditional DPM by assembling the unconditional DPM as the decoder estimator. Similar to regular DPM N ( x t ; µ θ ( x t , t ) + σ t G ψ ( x t , z, t ) , σ t ) , we can use the following objective to train encoder E ϕ and the decoder estimator G ψ :

4 Method

In this section, we initially present the formulation of the proposed task, Disentangling DPMs, in Section 4.1. Subsequently, an overview of DisDiff is provided in Section 4.2. We then elaborate on the detailed implementation of the proposed Disentangling Loss in Section 4.3, followed by a discussion on balancing it with the reconstruction loss in Section 4.4.

4.1 Disentanglement of DPM

We assume that dataset D = { x 0 | x 0 ∼ p ( x 0 ) } is generated by N underlying ground truth factors f c , where each c ∈ C = { 1 , . . . , N } , with the data distribution p ( x 0 ) . This implies that the underlying factors condition each sample. Moreover, each factor f c follows the distribution p ( f c ) , where p ( f c ) denotes the distribution of factor c . Using the Shapes3D dataset as an example, the underlying

𝑓 1

𝑓 2

𝑓 𝑁

…

𝑥0

𝑧 1

𝑧 2

𝑧 𝑁

…

②

⑤

⑥

⑦

①

③

④

①

②

③

④

⑤

⑥

⑦

(a) PGM of disentanglement

(b) Image Space

(c) Sampled Images

ObjectColorisRed

Gy(xt,t,z1)

Gy(xt,t,z²)

FloorColor

isGreen

Background

Color is Blue

concept factors include background color, floor color, object color, object shape, object scale, and pose. We illustrate three factors (background color, floor color, and object color) in Figure 1 (c). The conditional distributions of the factors are { p ( x 0 | f c ) | c ∈ C} , each of which can be shown as a curved surface in Figure 1 (b). The relation between the conditional distribution and data distribution can be formulated as: The conditional distributions of the factors, denoted as { p ( x 0 | f c ) | c ∈ C} , can each be represented as a curved surface in Figure 1 (b). The relationship between the conditional distribution and the data distribution can be expressed as follows:

this relationship holds true because f 1 , . . . , f N and x 0 form a v-structure in Probabilistic Graphical Models (PGM), as Figure 1 (a) depicts. A DPM learns a model ϵ θ ( x t , t ) , parameterized by θ , to predict the noise added to a sample x t . This can then be utilized to derive a score function: ∇ x t log p ( x t ) = 1 / √ α t -1 · ϵ θ ( x t , t ) . Following [9] and using Bayes' Rule, we can express the score function of the conditional distribution as:

By employing the score function of p ( x t | f c ) , we can sample data conditioned on the factor f c , as demonstrated in Figure 1 (b). Furthermore, the equation above can be extended to one conditioned on a subset S ⊆ C by substituting f c with S . The goal of disentangling a DPM is to model ∇ x t log p ( x t | f c ) for each factor c . Based on the equation above, we can learn ∇ x t log p ( f c | x t ) for a pre-trained DPM instead, corresponding to the arrows pointing towards the curve surface in Figure 1 (b). However, in an unsupervised setting, f c is unknown. Fortunately, we can employ disentangled representation learning to obtain a set of representations { z c | c = 1 , . . . , N } . There are two requirements for these representations: ( i ) they must encompass all information of x 0 , which we refer to as the completeness requirement, and ( ii ) there is a bijection, z c ↦→ f c , must exist for each factor c , which we designate as the disentanglement requirement. Utilizing these representations, the gradient ∇ x t log p ( z c | x t ) can serve as an approximation for ∇ x t log p ( f c | x t ) . In this paper, we propose a method named DisDiff as a solution for the disentanglement of a DPM.

4.2 Overview of DisDiff

The overview framework of DisDiff is illustrated in Figure 2. Given a pre-trained unconditional DPM on dataset D with factors C , e.g., a DDPM model with parameters θ , p θ ( x t -1 | x t ) = N ( x t -1 ; µ θ ( x t , t ) , σ t ) , our target is to disentangle the DPM in an unsupervised manner. Given an input x 0 ∈ D , for each factor c ∈ C , our goal is to learn the disentangled representation z c simultaneously and its corresponding disentangled gradient field ∇ x t log p ( z c | x t ) . Specifically, for each representation z c , we employ an encoder E ϕ with learnable parameters ϕ to obtain the

𝑥଴

𝑥௧

𝜖ఏ (𝑥 ௧

, 𝑡)

𝐺ట(𝑥௧, 𝑡, 𝑧

ଵ )

𝐺ట(𝑥௧, 𝑡, 𝑧

ଶ )

𝐺ట(𝑥௧, 𝑡, 𝑧

ଷ )

Pre-trained DPM

Encoder

𝐸థ

Decoder

𝐺ట

Disentangled

Representations

𝑧 ଵ

𝑧 ଶ

𝑧 ଷ

Disentangled

Gradient Fields

𝜖ఏ (𝑥 ௧

, 𝑡)

Disentangling loss

𝑧 ଵ

𝑧 ଶ

𝑧 ଷ

𝐺ట(𝑥௧, 𝑡, 𝑧

ଵ )

𝐺ట(𝑥௧, 𝑡, 𝑧

ଶ )

𝐺ట(𝑥௧, 𝑡, 𝑧

ଷ )

Encoder

𝐸థ

Encoder

𝐸థ

Decoder

𝐺ట

(a) The networks architecture of DisDiff

(b) The demonstration of disentangling loss

representations as follows: E ϕ ( x 0 ) = { E 1 ϕ ( x 0 ) , E 2 ϕ ( x 0 ) , . . . , E N ϕ ( x 0 ) } = { z 1 , z 2 , . . . , z N } . As the random variables z 1 , . . . , z N and x 0 form a common cause structure, it is evident that z 1 , . . . , z N are independent when conditioned on x 0 . Additionally, this property also holds when conditioning on ϵ . We subsequently formulate the gradient field, denoted as ∇ x t log p ( z S | x t ) , with z S = { z c | c ∈ S} , conditioned on a factor subset S ⊆ C , as follows:

Adopting the approach from [9, 49], we formulate the conditional reverse process as a Gaussian distribution, represented as p θ ( x t -1 | x t , z S ) . In conjunction with Eq. 7, this distribution exhibits a shifted mean, which can be described as follows:

Directly employing ∇ x t log p ( z c | x t ) introduces computational complexity and complicates the diffusion model training pipeline, as discussed in [18]. Therefore, we utilize a network G ψ ( x t , z c , t ) , with c ∈ S and parameterized by ψ , to estimate the gradient fields ∇ x t log p ( z S | x t ) . To achieve this goal, we first use ∑ c ∈C G ψ ( x t , z c , t ) to estimate ∇ x t log p ( z C | x t ) based on Eq. 7. Therefore, we adopt the same loss of PDAE [49] but replace the gradient field with the summation of G ψ ( x t , z c , t ) as √

The aforementioned equation implies that x 0 can be reconstructed utilizing all disentangled representations when S = C , thereby satisfying the completeness requirement of disentanglement. Secondly, in the following, we fulfill the disentanglement requirement and the approximation of gradient field ∇ x t log p ( z c | x t ) with decoder G ψ ( x t , z c , t ) . In the next section, we introduce a disentangling loss to address the remaining two conditions.

4.3 Disentangling Loss

In this section, we present the Disentangling Loss to address the disentanglement requirement and ensure that G ψ ( x t , z c , t ) serves as an approximation of ∇ x t log p ( z c | x t ) . Given the difficulty in identifying sufficient conditions for disentanglement in an unsupervised setting, we follow the previous works in disentangled representation literature, which propose necessary conditions that are effective in practice for achieving disentanglement. For instance, the group constraint in [47] and maximizing mutual information in [27, 45] serve as examples. We propose minimizing the mutual information between z c and z k , where k = c , as a necessary condition for disentangling a DPM.

̸

Subsequently, we convert the objective of minimizing mutual information into a set of constraints applied to the representation space. We denote ˆ x 0 is sampled from the pre-trained unconditioned

DPM using ϵ θ ( x t , t ) , ˆ x c 0 is conditioned on z c and sampled using p θ ( x t -1 | x t , z c ) . Then we can extract the representations from the samples ˆ x 0 and ˆ x c 0 with E ϕ as ˆ z k = E k ϕ (ˆ x 0 ) and ˆ z k | c = E k ϕ (ˆ x c 0 ) , respectively. According to Proposition 1, we can minimize the proposed upper bound (estimator) to minimize the mutual information.

̸

Proposition 1. For a PGM in Figure 1 (a), we can minimize the upper bound (estimator) for the mutual information between z c and z k (where k = c ) by the following objective:

The proof is presented in Appendix K. We have primarily drawn upon the proofs in CLUB [6] and VAE [24] as key literature sources to prove our claims. As outlined below, we can partition the previously mentioned objectives into two distinct loss function components.

̸

Invariant Loss . In order to minimize the first part, ∥ ˆ z k | c -ˆ z k ∥ , where k = c , we need to minimize ∥ ˆ z k | c -ˆ z k ∥ , requiring that the k -th ( k = c, k ∈ C ) representation remains unchanged. We calculate the distance scalar between their k -th representation as follows:

We can represent the distance between each representation using a distance vector d = [ d 1 , d 2 , . . . , d C ] . As the distances are unbounded and their direct optimization is unstable, we employ the CrossEntropy loss 3 to identify the index c , which minimizes the distances at other indexes. The invariant loss L in is formulated as follows:

Variant Loss . In order to minimize the second part, ∥ ˆ z c | c -z c ∥ - ∥ ˆ z k | c -z k ∥ , in Eq. 10. Similarly, we implement the second part in CrossEntropy loss to maximize the distances at indexes k but minimize the distance at index c , similar to Eq. 12. Specifically, we calculate the distance scalar of the representations as shown below:

We adopt a cross-entropy loss to achieve the subjective by minimizing the variant loss L va , we denote [ d n 1 , d n 2 , . . . , d n N ] as d n and [ d p 1 , d n 2 , . . . , d p N ] as d p .

However, sampling ˆ x c 0 from p θ ( x t -1 | x t , z c ) is not efficient in every training step. Therefore, we follow [9, 49] to use a score-based conditioning trick [41, 40] for fast and approximated sampling. We obtain an approximator ϵ ψ of conditional sampling for different options of S by substituting the gradient field with decoder G ψ :

To improve the training speed, one can achieve the denoised result by computing the posterior expectation following [8] using Tweedie's Formula:

4.4 Total Loss

The degree of condition dependence for generated data varies among different time steps in the diffusion model [46]. For different time steps, we should use different weights for Disentangling Loss. Considering that the difference between the inputs of the encoder can reflect such changes in condition. We thus propose using the MSE distance between the inputs of the encoder as the weight coefficient:

where λ is a hyper-parameter. We stop the gradient of ˆ x 0 and ˆ x c 0 for the weight coefficient γ d .

̸

5 Experiments

5.1 Experimental Setup

Implementation Details . x 0 can be a sample in an image space or a latent space of images. We take pre-trained DDIM as the DPM (DisDiff-IM) for image diffusion. For latent diffusion, we can take the pre-trained KL-version latent diffusion model (LDM) or VQ-version LDM as DPM (DisDiff-KL and DisDiff-VQ). For details of network G θ , we follow Zhang et al. [49] to use the extended Group Normalization [9] by applying the scaling & shifting operation twice. The difference is that we use learn-able position embedding to indicate c :

where GN denotes group normalization, and [ t c s , t c b ] , [ z c s , z c b ] are obtained from a linear projection: z c s , z c b = linearProj ( z c ) , t c s , t c b = linearProj ([ t, p c ]) . p c is the learnable positional embedding. h is the feature map of Unet. For more details, please refer to Appendix A and B.

Datasets To evaluate disentanglement, we follow Ren et al. [36] to use popular public datasets: Shapes3D [23], a dataset of 3D shapes. MPI3D [12], a 3D dataset recorded in a controlled environment, and Cars3D [35], a dataset of CAD models generated by color renderings. All experiments are conducted on 64x64 image resolution, the same as the literature. For realworld datasets, we conduct our experiments on CelebA [29].

Baselines & Metrics We compare the performance with VAE-based and GAN-based baselines. Our experiment is conducted following DisCo. All of the baselines and our method use the same dimension of the representation vectors, which is 32, as referred to in DisCo. Specifically, the VAE-based models include FactorVAE [23] and β -TCVAE [5]. The GAN-based baselines include InfoGANCR [27], GANspace (GS) [14], LatentDiscovery (LD) [43] and DisCo [36]. Considering the

influence of performance on the random seed. We have ten runs for each method. We use two representative metrics: the FactorVAE score [23] and the DCI [11]. However, since { z c | c ∈ C} are vector-wise representations, we follow Du et al. [10] to perform PCA as post-processing on the representation before evaluation.

SRC

TRT

Oren

Wall

Floor

Color

Shape

SRC

TRT

Bangs

Skin

Smile

Hair

5.2 Main Results

We conduct the following experiments to verify the disentanglement ability of the proposed DisDiff model. We regard the learned representation z c as the disentangled one and use the popular metrics in disentangled representation literature for evaluation. The quantitative comparison results of disentanglement under different metrics are shown in Table 1. The table shows that DisDiff outperforms the baselines, demonstrating the model's superior disentanglement ability. Compared with the V AE-based methods, since these methods suffer from the trade-off between generation and disentanglement [26], DisDiff does not. As for the GAN-based methods, the disentanglement is learned by exploring the latent space of GAN. Therefore, the performance is limited by the latent space of GAN. DisDiff leverages the gradient field of data space to learn disentanglement and does not have such limitations. In addition, DisDiff resolves the disentanglement problem into 1000 sub-problems under different time steps, which reduces the difficulty.

We also incorporate the metric for real-world data into this paper. We have adopted the TAD and FID metrics from [45] to measure the disentanglement and generation capabilities of DisDiff, respectively. As shown in Table 3, DisDiff still outperforms new baselines in comparison. It is worth mentioning that InfoDiffusion [45] has made significant improvements in both disentanglement and generation capabilities on CelebA compared to previous baselines. This has resulted in DisDiff's advantage over the baselines on CelebA being less pronounced than on Shapes3D.

The results are taken from Table 2 in [45], except for the results of DisDiff. For TAD, we follow the method used in [45] to compute the metric. We also evaluate the metric for 5-folds. For FID, we follow Wang et al. [45] to compute the metric with five random sample sets of 10,000 images. Specifically, we use the official implementation of [48] to compute TAD, and the official GitHub repository of [15] to compute FID.

TRT

Oren

Wall

Floor

Color

Shape

TRT

Bangs

Skin

Smile

Hair

5.3 Qualitative Results

In order to analyze the disentanglement of DisDiff qualitatively, we swap the representation z c of two images one by one and sample the image conditioned on the swapped representation. We follow LDM-VQ to sample images in 200 steps. For the popular dataset of disentanglement literature, we take Shapes3D as an example. As Figure 3 shows, DisDiff successfully learned pure factors. Compared with the VAE-based methods, DisDiff has better image quality. Since there are no ground truth factors for the real-world dataset, we take CelebA as an example. As demonstrated in Figure 3, DisDiff also achieves good disentanglement on real-world datasets. Please note that, unlike DisCo [36], DisDiff can reconstruct the input image, which is unavailable for DisCo. Inspired by Kwon et al. [25], we generalize DisDiff to complex large-natural datasets FFHQ and CelebA-HQ. The experiment results are given in Appendix E-H.

5.4 Ablation Study

In order to analyze the effectiveness of the proposed parts of DisDiff, we design an ablation study from the following five aspects: DPM type, Disentangling Loss, loss weighting, condition type, and latent number. We take Shapes3D as the dataset to conduct these ablation studies.

DPMtype The decomposition of the gradient field of the diffusion model derives the disentanglement of DisDiff. Therefore, the diffusion space influences the performance of DisDiff. We take Shapes3D as an example. It is hard for the model to learn shape and scale in image space, but it is much easier in the latent space of the auto-encoder. Therefore, we compare the performance of DisDiff with different diffusion types: image diffusion model, e.g., DDIM (DisDiff-IM), KL-version latent diffusion model (DisDiff-KL), and VQ-version latent diffusion model, e.g., VQ-LDM (DisDiff-VQ). As shown in Table 2, the LDM-version DisDiff outperforms the image-version DisDiff as expected.

̸

Disentangling Loss Disentangling loss comprises two parts: Invariant loss L in and Variant loss L va . To verify the effectiveness of each part, we train DisDiff-VQ without it. L in encourages in-variance of representation not being sampled, which means that the sampled factor will not affect the representation of other factors ( z k , k = c of generated ˆ x c 0 ). On the other hand, L va encourages the representation of the sampled factor ( z c of generated ˆ x c 0 ) to be close to the corresponding one of x 0 . As shown in Table 2, mainly L in promotes the disentanglement, and L va further constrains the model and improves the performance. Note that the disentangling loss is optimized w.r.t. G θ but not E c θ . If the loss is optimized on both modules, as shown in Table 2, DisDiff fails to achieve disentanglement. The reason could be that the disentangling loss influenced the training of the encoder and failed to do reconstruction.

Loss weighting As introduced, considering that the condition varies among time steps, we adopt the difference of the encoder as the weight coefficient. In this section, we explore other options to verify its effectiveness. We offer two different weighting types: constant weighting and loss weighting. The

first type is the transitional way of weighting. The second one is to balance the Distangling Loss and diffusion loss. From Table 2, these two types of weighting hurt the performance to a different extent.

Condition type & Latent number DisDiff follows PDAE [49] and Dhariwal and Nichol [9] to adopt AdaGN for injecting the condition. However, there is another option in the literature: cross-attention. As shown in Table 2, cross-attention hurt the performance but not much. The reason may be that the condition is only a single token, which limits the ability of attention layers. We use learnable orthogonal positional embedding to indicate different factors. As shown in Table 2, no matter whether no positional embedding (wo pos embedding) or traditional learnable positional embedding (wo orth embedding) hurt the performance. The reason is that the orthogonal embedding is always different from each other in all training steps. The number of latent is an important hyper-parameter set in advance. As shown in Table 2, if the latent number is greater than the number of factors in the dataset, the latent number only has limited influence on the performance.

5.5 Partial Condition Sampling

As discussed in Section 4.2, DisDiff can partially sample conditions on the part of the factors. Specifically, we can use Eq. 15 to sample image condition on factors set S . We take Shapes3D as an example when DisDiff sampling images conditioned on the background color is red. We obtain images of the background color red and other factors randomly sampled. Figure 4 shows that DisDiff can condition individual factors on Shapes3D. In addition, DisDiff also has such ability on the real-world dataset (CelebA) in Figure 4. DisDiff is capable of sampling information exclusively to conditions.

6 Limitations

Our method is completely unsupervised, and without any guidance, the learned disentangled representations on natural image sets may not be easily interpretable by humans. We can leverage models like CLIP as guidance to improve interpretability. Since DisDiff is a diffusion-based method, compared to the V AE-based and GAN-based, the generation speed is slow, which is also a common problem for DPM-based methods. The potential negative societal impacts are malicious uses.

7 Conclusion

In this paper, we demonstrate a new task: disentanglement of DPM. By disentangling a DPM into several disentangled gradient fields, we can improve the interpretability of DPM. To solve the task, we build an unsupervised diffusion-based disentanglement framework named DisDiff. DisDiff learns a disentangled representation of the input image in the diffusion process. In addition, DisDiff learns a disentangled gradient field for each factor, which brings the following new properties for disentanglement literature. DisDiff adopted disentangling constraints on all different timesteps, a new inductive bias. Except for image editing, with the disentangled DPM, we can also sample partial conditions on the part of the information by superpositioning the sub-gradient field. Applying DisDiff to more general conditioned DPM is a direction worth exploring for future work. Besides, utilizing the proposed disentangling method to pre-trained conditional DPM makes it more flexible.

References

Omri Avrahami, Dani Lischinski, and Ohad Fried. Blended diffusion for text-driven editing of natural images. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 18208-18218, 2022.

Yoshua Bengio, Aaron C. Courville, and Pascal Vincent. Representation learning: A review and new perspectives. PAMI , 2013.

Christopher P Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling in beta-vae. arXiv:1804.03599 , 2018.

Jinkun Cao, Ruiqian Nai, Qing Yang, Jialei Huang, and Yang Gao. An empirical study on disentanglement of negative-free contrastive learning. arXiv preprint arXiv:2206.04756 , 2022.

Ricky TQ Chen, Xuechen Li, Roger B Grosse, and David K Duvenaud. Isolating sources of disentanglement in variational autoencoders. In NeurPIS , 2018.

Pengyu Cheng, Weituo Hao, Shuyang Dai, Jiachang Liu, Zhe Gan, and Lawrence Carin. Club: A contrastive log-ratio upper bound of mutual information. In International conference on machine learning , pages 1779-1788. PMLR, 2020.

Jooyoung Choi, Sungwon Kim, Yonghyun Jeong, Youngjune Gwon, and Sungroh Yoon. Ilvr: Conditioning method for denoising diffusion probabilistic models. arXiv preprint arXiv:2108.02938 , 2021.

Hyungjin Chung, Byeongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving diffusion models for inverse problems using manifold constraints. Advances in Neural Information Processing Systems , 35: 25683-25696, 2022.

Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems , 34:8780-8794, 2021.

Yilun Du, Shuang Li, Yash Sharma, Josh Tenenbaum, and Igor Mordatch. Unsupervised learning of compositional energy concepts. Advances in Neural Information Processing Systems , 34, 2021.

Cian Eastwood and Christopher K. I. Williams. A framework for the quantitative evaluation of disentangled representations. In ICLR , 2018.

Muhammad Waleed Gondal, Manuel Wuthrich, Djordje Miladinovic, Francesco Locatello, Martin Breidt, Valentin Volchkov, Joel Akpo, Olivier Bachem, Bernhard Schölkopf, and Stefan Bauer. On the transfer of inductive bias from simulation to the real world: a new disentanglement dataset. In NeurIPS , 2019.

Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM , 63(11): 139-144, 2020.

Erik Härkönen, Aaron Hertzmann, Jaakko Lehtinen, and Sylvain Paris. Ganspace: Discovering interpretable GAN controls. In NeurIPS , 2020.

Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems , 30, 2017.

Irina Higgins, Loïc Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with a constrained variational framework. In ICLR , 2017.

Irina Higgins, David Amos, David Pfau, Sebastien Racaniere, Loic Matthey, Danilo Rezende, and Alexander Lerchner. Towards a definition of disentangled representations. arXiv preprint arXiv:1812.02230 , 2018.

Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598 , 2022.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems , 33:6840-6851, 2020.

Alexia Jolicoeur-Martineau, Rémi Piché-Taillefer, Rémi Tachet des Combes, and Ioannis Mitliagkas. Adversarial score matching and improved sampling for image generation. arXiv preprint arXiv:2009.05475 , 2020.

Bahjat Kawar, Shiran Zada, Oran Lang, Omer Tov, Huiwen Chang, Tali Dekel, Inbar Mosseri, and Michal Irani. Imagic: Text-based real image editing with diffusion models. arXiv preprint arXiv:2210.09276 , 2022.

Gwanghyun Kim, Taesung Kwon, and Jong Chul Ye. Diffusionclip: Text-guided diffusion models for robust image manipulation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 2426-2435, 2022.

Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In ICML , 2018.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

Mingi Kwon, Jaeseok Jeong, and Youngjung Uh. Diffusion models already have a semantic latent space. arXiv preprint arXiv:2210.10960 , 2022.

José Lezama. Overcoming the disentanglement vs reconstruction trade-off via jacobian supervision. In International Conference on Learning Representations , 2018.

Zinan Lin, Kiran Thekumparampil, Giulia Fanti, and Sewoong Oh. Infogan-cr and modelcentrality: Self-supervised model training and selection for disentangling gans. In ICML , 2020.

Xihui Liu, Dong Huk Park, Samaneh Azadi, Gong Zhang, Arman Chopikyan, Yuxiao Hu, Humphrey Shi, Anna Rohrbach, and Trevor Darrell. More control for free! image synthesis with semantic diffusion guidance. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision , pages 289-299, 2023.

Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV) , December 2015.

Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. In international conference on machine learning , pages 4114-4124. PMLR, 2019.

Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11461-11471, 2022.

Chenlin Meng, Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun-Yan Zhu, and Stefano Ermon. Sdedit: Guided image synthesis and editing with stochastic differential equations. In International Conference on Learning Representations , 2021.

Konpat Preechakul, Nattanat Chatthee, Suttisak Wizadwongsa, and Supasorn Suwajanakorn. Diffusion autoencoders: Toward a meaningful and decodable representation. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR) , 2022.

Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning , pages 8748-8763. PMLR, 2021.

Scott E. Reed, Yi Zhang, Yuting Zhang, and Honglak Lee. Deep visual analogy-making. In NeurIPS , 2015.

Xuanchi Ren, Tao Yang, Yuwang Wang, and Wenjun Zeng. Learning disentangled representation by exploiting pretrained generative models: A contrastive learning view. In International Conference on Learning Representations , 2021.

Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorealistic text-toimage diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487 , 2022.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning , pages 2256-2265. PMLR, 2015.

Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502 , 2020.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems , 32, 2019.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. ICLR , 2020.

Akash Srivastava, Yamini Bansal, Yukun Ding, Cole L. Hurwitz, Kai Xu, Bernhard Egger, Prasanna Sattigeri, Josh Tenenbaum, David D. Cox, and Dan Gutfreund. Improving the reconstruction of disentangled representation learners via multi-stage modelling. CoRR , abs/2010.13187, 2020.

Andrey Voynov and Artem Babenko. Unsupervised discovery of interpretable directions in the GAN latent space. In ICML , 2020.

Tengfei Wang, Yong Zhang, Yanbo Fan, Jue Wang, and Qifeng Chen. High-fidelity gan inversion for image attribute editing. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 11379-11388, 2022.

Yingheng Wang, Yair Schiff, Aaron Gokaslan, Weishen Pan, Fei Wang, Christopher De Sa, and Volodymyr Kuleshov. Infodiffusion: Representation learning using information maximizing diffusion models. international conference on machine learning , 2023.

Qiucheng Wu, Yujian Liu, Handong Zhao, Ajinkya Kale, Trung Bui, Tong Yu, Zhe Lin, Yang Zhang, and Shiyu Chang. Uncovering the disentanglement capability in text-to-image diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 1900-1910, 2023.

Tao Yang, Xuanchi Ren, Yuwang Wang, Wenjun Zeng, and Nanning Zheng. Towards building a groupbased unsupervised representation disentanglement framework. In International Conference on Learning Representations , 2021.

Eric Yeats, Frank Liu, David Womble, and Hai Li. Nashae: Disentangling representations through adversarial covariance minimization. In European Conference on Computer Vision , pages 36-51. Springer, 2022.

Zijian Zhang, Zhou Zhao, and Zhijie Lin. Unsupervised representation learning from pre-trained diffusion probabilistic models. NeurIPS , 2022.

Appendix

A The Training and Sampling Algorithms

As shown in Algorithm 1, we also present the training procedure of DisDiff. In addition, Algorithms 2 and 3 show the sampling process of DisDiff in the manner of DDPM and DDIM under the condition of S .

Algorithm 1 Training procedure

Prepare: dataset distribution p ( x 0 ) , pre-trained DPM ϵ θ , pre-trained VQ-VAE, E V Q , D V Q , Initialize: encoder E ϕ , decoder G ψ . repeat by Eq. 12 and Eq. 14.

Algorithm 2 DDPM Sampling

̸

B Training Details

Encoder We use the encoder architecture following DisCo [36] for a fair comparison. The details of the encoder are presented in Table 3, which is identical to the encoder of β -TCVAE and FactorVAE in Table 1.

Decoder We use the decoder architecture following PADE [49]. The details of the decoder are presented in Table 4, which is identical to the Unet decoder in the pre-trained DPM. DPM We adopted the VQ-version latent diffusion for image resolution of 64 and 256 as our pre-trained DPM. The details of the DPM are presented in Table 4. During the training of DisDiff, we set the batch size as 64 for all datasets. We always set the learning rate as 1 e -4 . We use EMA on all model

Algorithm 3 DDIM Sampling

Prepare: pre-trained DPM ϵ θ , pre-trained VQ-VAE E V Q , D V Q , the sampling factor set S ⊆ C , Input: sample x , sampling sequence { t i } K , and t 1 = 0 , t K = T

i =1 z = E ϕ ( x ) z xT ∼ N (0 , I ) or inferred z xT of inverse DDIM for j = K to 2 do

parameters with a decay factor of 0 . 9999 . Following DisCo [36], we set the representation vector dimension to 32 .

C Visualizations on Cars3D.

We also provide the qualitative results on Cars3D in Figure 5. We can observe that on the artificial dataset Cars3D, DisDiff also learns disentangled representation. In addition, we also provide the results of swapping representations that contain no information.

SRC

TRT

None

None

Grey

None

Car color

D Visualizations on MPI3D.

We provide the qualitative results on MPI3D in Figure 6. We can observe that even on such a challenging dataset (MPI3D) of disentanglement, DisDiff also learns disentangled representation. In addition, we also provide three rows of images swapping representations that contain no information.

E Orthogonal-constrained DisDiff

We follow DisCo [36] and LD [43] to construct learnable orthogonal directions of the bottleneck layer of Unet. The difference is that we use QR decomposition rather than the orthonormal matrix proposed by LD. Specifically, for a learnable matrix A ∈ R C × d , where d is the dimension of direction, we have the following equation:

where Q ∈ R C × d is an orthogonal matrix, and R ∈ R d × d is an upper triangular matrix, which is discarded in our method. We use Q to provide directions in Figure 7. Since the directions are sampled, and we follow DisCo and LD to discard the image reconstruction loss, we use the DDIM inverse ODE to obtain the inferred x T for image editing. The sampling process is shown in Algorithm 3.

F Qualitative Results on CelebA-HQ

In order to demonstrate the effectiveness of Orthogonal-constrained DisDiff, we train it on dataset CelebA-HQ. We also conduct our experiments with the latent unit number of 10 . The results are shown in Figure 9. Since the real-world datasets have many factors that are more than 10 , we also conduct our experiment with the latent unit number of 30 . As shown in Figure 11, Orthogonalconstrained DisDiff learns consistent factors for different images. There are more factors than Figure 9 shows. Therefore, we conduct another experiment using another random seed with a latent unit number of 10 . The results are shown in Figure 10.

SRC

TRT

OB color

Ob shape

None

None

None

BG color

𝜖ఏ (𝑥 ௧

, 𝑡)

Disentangling loss

𝐺 ଵ

(𝑥௧ , 𝑡, 𝑧

ଵ )

𝐺 ଶ

(𝑥௧ , 𝑡, 𝑧

ଶ )

𝐺 ଷ

(𝑥௧ , 𝑡, 𝑧

ଷ )

Encoder

𝐸థ

Encoder

𝐸థ

Decoder

𝐺 ఏ

௖

Pre-trained DPM

G CLIP-guided DisDiff

To tackle the disentanglement in real-world settings, we propose replacing the encoder in DisDiff with a powerful encoder, which lets us leverage the knowledge of the pre-trained model to facilitate the disentanglement. Specifically, we adopt CLIP as the encoder. Since the CLIP is an image-text pair pre-trained model, the multi-modal information is the pre-trained information we can leverage. Therefore, we use the CLIP text encoder to provide the condition but use the CLIP image encoder to calculate the disentangling loss. Figure 8 shows the obtained new architecture.

Disentangling loss

CLIP Image

Encoder

𝐸థ

ூ

Decoder

𝐺 ఏ

௖

CLIP Image

Encoder

𝐸థ

ூ

CLIP Text

Encoder

𝐸థ

்

Attribute words

𝜖ఏ (𝑥 ௧

, 𝑡)

𝐺 ଵ

(𝑥௧ , 𝑡, 𝑧

ଵ )

𝐺 ଶ

(𝑥௧ , 𝑡, 𝑧

ଶ )

𝐺 ଷ

(𝑥௧ , 𝑡, 𝑧

ଷ )

Pre-trained DPM

Figure 8: Illustration of CLIP-guided DisDiff. In this architecture, a powerful encoder replaces the original one in DisDiff. Specifically, we adopt the CLIP Image encoder as our new encoder, which also serves as guidance for ASYRP [25]. By leveraging knowledge learned from pre-trained CLIP, DisDiff can learn meaningful gradient fields. To use the learned knowledge of the language-image pair contrastive pre-training, we utilize attribute words as inputs, where different words correspond to different gradient fields. Finally, we incorporate disentangling loss at the end of the encoder to enforce both the invariant and variant constraints.

H Qualitative Results on FFHQ

We train the CLIP-guided DisDiff on FFHQ to verify its effectiveness. Since the attribute words used in the training of the CLIP-guided DisDiff are unpaired and provided in ASYRP, the learning of these factors is still unsupervised. As discussed in Section 4.2, DisDiff can do partial sampling that is conditioned on the part of the factors. Similarly, as shown in Figure 13, CLIP-guided DisDiff also has such an ability. In addition, as shown in Figure 12, we can also use the trained decoder of CLIP-guided DisDiff to edit the images by using the inverse ODE of DDIM to infer the x T .

I The Latent Number < The Number of GT Factors

When the latent number is smaller than the number of factors, the model's performance will significantly decrease, which is in accordance with our intuition. However, the performance of other methods will also decrease under this condition. We conducted experiments using FactorVAE and β -TCVAE to demonstrate it. We trained our model, FactorVAE and β -TCVAE on Shapes3D (The number of factors is 6) with a latent number of 3, and the hyperparameters and settings in their original paper were adopted for the latter two models accordingly. As shown in Table 6, our model showed a lower degree of disentanglement reduction than these two models and was significantly lower than β -TCVAE.

J Quantitative Results on More Metrics

We experiment with MIG and MED metrics on the Shapes3D dataset, two classifier-free metrics. For MIG, we utilized the implementation provided in disentanglement_lib 4 , while for MED, we used the official implementation available at Github 5 . The experimental results in Table 7 indicate that our proposed method still has significant advantages, especially in MED-topk and MIG.

K Proof of Proposition 1

Background : Our encoder transforms the data x 0 (hereafter denoted as x , the encoder keeps fixed for disentangling loss) into two distinct representations z c and z k . We assume that the encoding distribution follows a Gaussian distribution p ( z c | x ) , with a mean of E c ( x ) . However, the distribution p ( x | z c ) is intractable. We can employ conditional sampling of our DPM q θ ( x | z c ) to approximate the sampling from the distribution p ( x | z c ) . Given that we optimize the disentangling loss at each training step, we implement Eq. 15 and Eq. 16 in the main paper as an efficient alternative to sample from q θ ( x | z c ) . We further assume that the distribution used to approximate p ( z k | z c ) follows a Gaussian distribution q θ ( z k | z c ) , with a mean of z k | c = E k ( D ( z c )) , where D represents the operations outlined in Eq. 15 and Eq. 16. We represent the pretrained unconditional DPM using the notation q θ ( x ) ≈ p ( x ) .

̸

Proposition 2. For a PGM: z c → x → z k , we can minimize the upper bound (estimator) for the mutual information between z c and z k (where k = c ) by the following objective:

Considering the given context, we can present the PGM related to q as follows: z c → x → z k . The first transition is implemented by DPM represented by the distribution q θ . Subsequently, the second transition is implemented by the encoder, denoted by the distribution p .

Proposition 1.1 The estimator below is an upper bound of the mutual information I ( z c , z k ) . For clarity, we use expectation notation with explicit distribution.

Equality is attained if and only if z k and z c are independent.

Proof : I est -I ( z c , z k ) can be formulated as follows:

Using Jensen's Inequality, we obtain log p ( z k ) = log( E p ( z c ) [ p ( z k | z c )]) ⩾ E p ( z c ) [log p ( z k | z c )] . Therefore, the following holds:

Since the distribution p ( z k | z c ) is intractable, we employ q θ ( z k | z c ) to approximate p ( z k | z c ) .

Proposition 1.2 : The new estimator can be expressed as follows:

The first term corresponds to d p k in Eq. 13 of the main paper, while the second term corresponds to Eq. 11 in the main paper.

Proof :

Considering the first term, we obtain:

For the second term, we obtain:

Proposition 1.3 : If the following condition is fulfilled, the new estimator I θ can serve an upper bound for the mutual information I ( z c , z k ) , i.e., I θ ⩾ I ( z c , z k ) . The equality holds when z k and z c are independent. We denote q θ ( z k | z c ) p ( z c ) as q θ ( z k , z c ) .

Proof :

I θ -I ( z c , z k ) can be formulated as follows:

Since the condition is not yet applicable, we will continue exploring alternative conditions.

Proposition 1.4 : If D KL ( p ( z k | z c ) || q θ ( z k | z c )) ⩽ ϵ , then, either of the following conditions holds: I θ ⩾ I ( z c , z k ) or | I θ -I ( z c , z k ) | ⩽ ϵ .

Proof : If D KL ( p ( z k , z c ) || q θ ( z k , z c )) ⩽ D KL ( p ( z k ) p ( z c ) || q θ ( z k , z c )) , then

else:

Proposition 1.5 : Consider a PGM z c → x → z k ; if D KL ( q θ ( x | z c ) | p ( x | z c )) = 0 , it follows that D KL ( p ( z k | z c ) || q θ ( z k | z c )) = 0 .

Proof : if D KL ( q θ ( x | z c ) | p ( x | z c )) = 0 , we have q θ ( x | z c ) = p ( x | z c ) .

Ideally, by minimizing D KL ( q θ ( x | z c ) | p ( x | z c )) to 0, we can establish I θ either as an upper bound or as an estimator of I ( z c , z k ) .

Proposition 1.6 : we can minimize D KL ( q θ ( x | z c ) | p ( x | z c )) by minimizing the following objective:

Note that the equation above exactly corresponds to d p c in Eq. 13, and d n k is a constant. Proof :

We focus on the following term:

Applying the Bayes rule, we obtain:

We can minimize E p ( z c ) E q θ ( x | z c ) ∥ z c -ˆ z c | c ∥ instead.

Note that the loss usually cannot be optimized to a global minimum in a practical setting, but how the gap influences I θ is beyond the scope of this paper. We leave it as future work. Figure 4 in the main paper shows that our model has successfully fitted the distribution of p ( x | z c ) using q θ ( x | z c ) . Therefore, we can understand that the model's loss can be regarded as I est .

L Comparison with Diffusion-based Autoencoders

We include a comparison with more recent diffusion baselines, such as PADE [49] and DiffAE [33]. The results are shown in Table 8.

The relatively poor disentanglement performance of the two baselines can be attributed to the fact that neither of these models was specifically designed for disentangled representation learning, and thus, they generally do not possess such capabilities. DiffAE primarily aims to transform the diffusion model into an autoencoder structure, while PDAE seeks to learn an autoencoder for a pre-trained diffusion model. Unsupervised disentangled representation learning is beyond their scope.

Traditional disentangled representation learning faces a trade-off between generation quality and disentanglement ability (e.g., betaV AE [16], betaTCV AE [5], FactorV AE [23]). However, fortunately, disentanglement in our model does not lead to a significant drop in generation quality; in fact, some metrics even show improvement. The reason for this, as described in DisCo [36], is that using a well-pre-trained generative model helps avoid this trade-off.

SRC

Pose

Bangs

Expose

Long hair

SRC

Black hair

Bangs

Pose

Gender & Pose

Pose 2

ITISHA

FILMAN

Angeles

TILAA

SRC

Expose

Bangs

Long hair

Thin face

Pose

Bangs2

Expression

Long hair2

Bangs

Less hair

Pale face

Smile

Zoom in

SRC

Tanned face

Smile

Sad

Red hair

Blond hair

Old

Woman

Yellow hair

Young

Smile

Sad

Red hair

Blond hair

Old

Peer reviews:
comments: SUMMARY:
Imprecisions
- “According to the completeness requirement” add reference
- line 75-76 PADE-> PDAE
- line 270-271. Typo: compare with disco. dissdiff -> compared with disco, disdiff

Strengths
- Achieving disentangled representations in diffusion models is an interesting and useful problem.

Weaknesses
- Missing derivation of reverse diffusion process, Eq 5, 6, 7.
- What is the difference between the proposed model and a latent diffusion model?
- Diff-AE and PDAE should also be considered as baselines in experiments.
- Image generation experiments against Diff-AE and PDAE should be conducted in addition to disentanglement in order to add robustness to the experiments.
- Unclear how Figure 3 is generated. Could you describe the process in detail and explain what you mean by “swapping the representation”.

STRENGTHS:
see summary

WEAKNESSES:
see summary

QUESTIONS:
see summary

LIMITATIONS:
see summary
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-06-26

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes to disentangle a pre-trained diffusion probabilistic model (DPM) in an unsupervised way to improve interpretability. The author designed two constraints, invariant condition and variant condition, to guide the disentanglement. The proposed method was evaluated on three synthetic datasets and CelebA, a facial image dataset.\n\nSTRENGTHS:\n- The combination of disentangled representation learning and diffusion-based models is an important and challenging problem.\n- Judging solely by the experimental results, the proposed method works to some extent on synthetic datasets.\n\nWEAKNESSES:\n- Due to unclear definitions and poor mathematical formulation, I was unable to completely understand the proposed method.\n- The author claimed that \"the conditional independence of the representations is a necessary condition for disentanglement\". However, it is not sufficient. Disentanglement is more than the independence of representations. Even the definition of disentanglement used in this paper is unclear. \"For the first time achieving disentangled representation learning in the framework of DPMs\" is overclaiming.\n- The author's claims are a bit contradictory. The author claimed that \"disentangling a DPM can improve the interpretability of DPM\" but also admitted that \"Our method is completely unsupervised, and without any guidance, the learned disentangled representations on natural image sets may not be easily interpretable by humans.\" It seems that the author had a goal but failed to achieve it. Is it realistic to consider completely unsupervised disentanglement?\n\nQUESTIONS:\n- What is the mathematical definition of _disentangled conditional sub-gradient fields_? What does \"The data distribution $p(x)$ can be **disentangled** into $N$ independent distributions $\\\\{p(x \\mid f^k) \\mid k = 1, \\dots, N \\\\}$\" mean?\n- What does \"the data sample space is not **well-organized**\" mean?\n- Cross entropy compares probability distributions. However, Eq. (10) uses a \"distance vector\" $d$ and \"index\" $c$. Eq. (12) even uses a difference of distance vectors. The math here is too questionable to make this paper credible. Please revise this part.\n\nLIMITATIONS:\nSection 6 discussed the limitations of unsupervised disentanglement and diffusion-based methods.\n\n\"The potential negative societal impacts are malicious uses.\" is too general and ambiguous.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a new task of unsupervised disentanglement of diffusion probabilistic models (DPMs) and presents an approach named DisDiff to achieve disentangled representation learning in the framework of DPMs. The authors connect disentangled representation learning to DPMs to take advantage of the remarkable modeling ability of DPMs. DisDiff learns a disentangled representation of the input image in the diffusion process and for each factor, DisDiff learns a disentangled gradient field, which brings new properties for disentanglement literature. The proposed method is evaluated on several real-world datasets, and the results show that DisDiff outperforms state-of-the-art methods in terms of disentanglement quality and interpretability. The paper concludes by discussing potential future directions for applying DisDiff to more general conditioned DPMs and pre-trained conditional DPMs.\n\nSTRENGTHS:\nStrengths:\n- Performance: The proposed method is evaluated on several real-world datasets, and the results show that DisDiff outperforms state-of-the-art methods in terms of disentanglement quality and interpretability.\n\n- Clarity: The writing is good overall and the proposed idea in this paper is easy to follow. The authors provide good explanations of the method and the experiments.\n\n- Significance: The proposed method has potential applications in various fields, such as image editing, controllable generation, etc.\n\n- References: The paper provides comprehensive references.\n\nWEAKNESSES:\nWeaknesses:\n- Overclaim: this paper claims that they are the very first work introducing disentanglement tasks for diffusion probabilistic models (DPM). However, a very recent work [1] in ICML 2023 has introduced and studied this problem. I understand that there might be some timeline issue that makes the authors not aware of this work. But it would be better if the authors could properly cite and discuss this work and revise their claims in the corresponding paragraphs.\n- Abuse of notations: there are some wrong notations (might be typos) and abuse of notations, especially in Sec. 4.3. For example, in line 179, it should be $E_\\phi^k(\\hat{x}_0^c)$ and $E_\\phi^k(\\hat{x}_0)$. In lines 180 - 181, the conditioned representation should be $E_\\phi^c(\\hat{x}_0^c)$ and the claim should be that the unconditioned one is closer to $E_\\phi^c(\\hat{x}_0)$ than the conditioned one. I suggest the authors carefully examine their statements instead of letting the reviewers guess the meaning.\n- More baselines: I appreciate that the authors include some of the most important baselines of VAEs and GANs along the line of disentanglement research. However, some more recent diffusion baselines could also be included, e.g., PADE [2] and DiffAE [3] (I understand [1] would be too new to have the open-source codebase to evaluate).\n- More metrics for CelebA: there is a tailored quantitative metric [4] for evaluating disentanglement performance on CelebA. It would be better to have the numbers in the tables.\n- In line 176, fulfil -> fulfill\n\n[1] Wang, Yingheng, et al. \"InfoDiffusion: Representation Learning Using Information Maximizing Diffusion Models.\" arXiv preprint arXiv:2306.08757 (2023).\n\n[2] Zhang, Zijian, Zhou Zhao, and Zhijie Lin. \"Unsupervised representation learning from pre-trained diffusion probabilistic models.\" Advances in Neural Information Processing Systems 35 (2022): 22117-22130.\n\n[3] Preechakul, Konpat, et al. \"Diffusion autoencoders: Toward a meaningful and decodable representation.\" Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022.\n\n[4] Yeats, Eric, et al. \"Nashae: Disentangling representations through adversarial covariance minimization.\" European Conference on Computer Vision. Cham: Springer Nature Switzerland, 2022.\n\nQUESTIONS:\n- Distortion in generated images: although the quantitative metrics (e.g. DCI) are pretty promising, the qualitative results on CelebA still seem a bit distorted with disentangled latent variables. Would there be some training tricks that can be applied to alleviate this issue? Are there any other explanations?\n- In line 159, the authors claim that Eq. (7) is derived from Eq. (6) by using Tweedie's formula, however, in my opinion, this is just a straightforward application of reparameterization trick. Can the authors provide the full derivation of the methodology part?\n\nLIMITATIONS:\nThere is a section of discussion for limitations.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-13"}