DisDiff: Unsupervised Disentanglement of Diffusion Probabilistic Models
Authors: Tao Yang
Yuwang Wang
Yan Lu
Nanning Zheng
Venue: NeurIPS 2023 poster (2023)
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
Tao Yang': Yuwang Wang?!, Yan Lu NNanning Zheng" "f
yt142120stu. x.j.uuxxjuuu xjtu. edu. cn,
National Key Laboratory of Human-Machine Hybrid Augmented Intelligence, National Engineering Rescarch Center for Visual Information and Applications, and Institute of Artificial Intelligence and Robotics, Xi' an Jiaotong University,
Tsinghua University, Shanghai Al Laboratory, Microsoft Research Asia
https:/ https://giihhb....o wttp./t.......o.aooroooo//////// /Dirrryf
Abstract
Targeting to understand the underlying explainable factors behind observations and modeling the conditional generation process on these factors, connect disen- tangled representation learning to Diffusion Probabilistic Models (DPMs) to take advantage of the remarkable modeling ability of DPMs. We propose new task, disentanglement of (DPMs): given pre-trained DPM, without any annotations of the factors, the task 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 repre- sented, and clearly injected into the diffusion process via the sub- aub-gradient fields. To tackle this task, we devise an unsupervised approach named DisDiff, achiev- ing disentangled representation learning in the framework of DPMs. Extensive experiments on synthetic and real-world datasets demonstrate the effectiveness of Disfiff.
Introduction
As one of the most successful generative models, diffusion probabilistic models (DPMs) achieve remarkable performance in image synthesis. They series of probabilistic distributions 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 representation learning [49] based DPM. Diff-AE [33] and PDAE [49] recently proposed methods for representation learning by reconstructing the images in the DPM framework. However, the learned latent representation can only be interpreted relying 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 Al 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
Work done during internships at Microsoft Research Asia. t Corresponding authors: wang-yuwange mail. wwwillssinuu..d..oo nnzheng@ mail.xjtu.edu.cm.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 GAN- based 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 new task: the disentanglement of DPM. Given a pre-trained DPM model, the goal of disentanglement of DPM 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]/mage--ee] pairs [21], supervised pre-trained models [22] such CLIP [34]. One can also flexibly sample partial conditions on the part of the information introduced by the superposition of the sub-grad sub-gradient field, which 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 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 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 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 disen- tangled sub-gradien 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 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, mage--uudde 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 classifier- guided 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-ac [33] proposes a diffusion-basee based auto-encoder for image reconstruction. PDAE [49] uses 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 VAE-based

---

works achieve disentanglement [5. 23, 16, 3] only by the constraints on probabilistic distributions of representations. Locatello 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 new type of inductive bias. Furthermore, DPM can sample non-factorized (non-conditioned) information [18], naturally fitting for disentanglement, this way, achieve disentanglement for DPMs.
3 Background
3.1 Diffusion Probabilistic Models (DPM)
We take DDPM [19] as example. DDPM adopts sequence of fixed variance distributions ((x)( as the forward process to collapse the image distribution p(zo) to (0, These
distributions
((=+/==1) = Naas v 1 1,
(1)
We can then sample x, using the following formula dy ~ N(xx; varzo, (1 (1-a+)), where at - 3 and o M2 (t, i.e., It varzn + /1 dre. The reverse process is fitted by using a sequence of Gaussian distributions parameterized by 0:
po(xt-1laa) = M( petiiplle (o o,)
where ( () is parameterized by Unet en 1), it is trained by minimizing the variational upper bound of negative log-likelihood through:
Lo
E ||e co (xt, t)ll- mo.t,
3.2 Representation learning from DPMs
The classifier guided method [9] uses the gradient of a pre-trained classifier, V log ply zt). to impose condition pre-trained DPM and obtain new conditional DPM: A (zzi pa (2,, t)
or log p(y (xe), o.). Based on the classifier-guided sampling method, PDAE [49] proposes an approach for pre-trained DPM by incorporating auto-encoder. Specifically, given pre-trained DPM, PDAE introduces an encoder Eg to derive the representation by equation 2 = E (zo). They
decoder estimator G (æ,, 2, to simulate the gradient field V log p(2(a)) for reconstructing the input image.
By this means, they create new conditional DPM by assembling the unconditional DPM as the decoder estimator. Similar to regular DPM (1 ( 1) o, G. (14, 2, 2,1), we can use following objective train encoder Eg and the decoder estimator Go:
Vaiy 1 αi
Lxp
E ||e co(x)))) 2g.
01 G
4 Method
In this section, we initially present the formulation of the proposed task, Disentangling DPMs, in Section 4.1. Subsequently, overview of DisDiff is provided in Section 4.2. We then elaborate on the detailed implementation of the proposed Disentangling Loss Scction 4.3, followed by discussion on balancing it with the reconstruction loss in Section 4.4.
4.1 Disentanglement of DPM
We assume that dataset {ar xo ~ p(20)} generated by N underlying ground truth factors Ss, where each E C (1...., N). with the data distribution p(æn). This implies that the underlying factors condition each sample. Moreover, each factor fe follows the distribution p(J"), where p( denotes the distribution of factor Using the Shapes3D dataset example, the underlying
3

---

Object Color Red
f2
fN
Gy
X
Cte 6.23)
floor Color is Graen
Beckgrount Blus
Z
(a) PGM of disentanglement
(b) Image Space
(c) Sampled Images
Figure 1: Illustration of disentanglement of DPMs. (a) is the diagram of Probabilistic Graphical
Models (PGM). (b) is the diagram of image space. (c) is the demonstration of sampled images in (b). Surface indicates the conditional data distribution of a single factor p(xrl/). Different colors correspond various factors. Here, we show three factors: object color, background color, and floor color. Arrows are gradient fields vx, log p(f ( modeled by using a decoder G (z,, t, 2*). The black points are the sampled images shown in (b).
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 (c). The conditional distributions of the factors are {p(xa]/")]c= | C), each of which can be shown as curved surface in Figure (b). The relation between the conditional distribution and data distribution can be formulated as: The conditional distributions of the factors, denoted as {p(xo)5")] E C). can each be represented curved surface in Figure 1 (b). The relationship between the conditional distribution and the data distribution can be expressed as follows:
p(zo)= = f (P())))
f {
..... (M )df
Jdf df
(5)
this relationship holds true because f , Models (PGM), as Figure (a) depicts. A DPM learns model e0(I,, t), parameterized by 0, to predict the noise added to a sample xy. This can then be utilized to derive a score function: express the
S. .......
and xo form a v-structure in Probabilistic Graphical
V., log p(xz) 1//aa- -1 cœ t). Following [9] and using Bayes' Rule, score function of the conditional distribution as:
Vx, logg(z.))"")) Va, (xx)(19)))) + Va, logp(x:),
(6)
By employing the score function of ff), we can sample data conditioned on the factor fo as demonstrated in Figure (b). Furthermore, the equation above can be extended to one conditioned on subset S C C by substituting f with S. The goal of disentangling a DPM is to model Va, log p(=x]_ for each factor c. Based on the equation above, we can learn V.n x, log log (x)))) 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 is unknown. Fortunately, we can employ disentangled representation learning to obtain set of representations {=1ce = I,..., NJ. There are two requirements for these representations: (i) they must encompass all information of zo. which we refer to as the completeness requirement, and (ii) there is bijection, 2" + ʃ must exist for each factor c, which we designate the disentanglenten requirement. Utilizing these representations, the gradient Vz, logp(2" (x)) can serve as an approximation for Vz, log (f))))) In this paper, we propose method named DisDiff solution for the disentanglement of DPM.
4.2 Overview of DisDiff
The overview framework of DisDiff illustrated in Figure 2. Given a pre-trained unconditional DPM on dataset D with factors C, e.g., a DDPM model with parameters 6, po (o(==_1/2;) (x 14 (Kt, 1), ot), our target is to disentangle the DPM in an unsupervised manner. Given an input zo E 'D for each factor E C, our goal is to learn the disentangled representation = simultaneously and its corresponding disentangled gradient field Va, log p(=)().. Specifically, for each representation 2 we employ an encoder Es with learnable parameters to obtain the

---

Encader Eg
Ceooder
Decoder 6
Encader Eg
Druccar
Dierlangidd
Obsentangeed
000k
8002
Disentangling
Representarcom
(a) The networks architecture of DisDiff
(b) The demonstration of disentangling loss
Figure 2: Illustration of DisDiff. (a) Grey networks indicate the pre-trained Unet of DPM c t). different factors by encoder Es (N 3 in the figure), which are then decoded by decoder G to obtain the gradient field of the corresponding factor. We can sample the image under the corresponding condition with the obtained gradient field. (b) We first sample factor and decode the representation 2" to obtain the gradient field of the corresponding factor, which allows us to obtain the predicted zo of that factor. At the same time, can obtain the predicted So of the original pre-trained DPM. We then encode the images into two different representations and calculate the disentangling loss based on them.
Image Ty is first encoded to representations {z1, (a",2% 2²,,... =²)
NN As the random variables =1 ...... = and æ form a common cause structure, it is evident that =1 NN are independent when conditioned on zi- Additionally, this property also holds when conditioning on 6. We subsequently formulate the gradient field, denoted as Vz, log 1ogs(x>)1;)) with 2 {2 E S}. conditioned on a factor subset S C C. as follows:
representations as follows: Eg (zn) = (z)), E3(())........ (=a)} {=
Vx, log logp(=>(1x)))
ΣVE.. V., log 1og
GES
Adopting the approach from [9, 49], we formulate the conditional reverse process as a Gaussian distribution, represented po p= (x4,, 2 In conjunction with Eq. 7, this distribution exhibits shifted mean, which can be described as follows:
M(Fc-i:M + oe Σ V..
(8)
cES
Directly employing Va, log logp(="(z;)) introduces computational complexity and complicates the diffusion model training pipeline, as discussed in [18]. Therefore, utilize network G f), with e S and parameterized by v. estimate the gradient fields Vz, log p(= [x;). To achieve this goal, we first use Eeec Gw 2", .)) to estimate Va, log p(211zz) based on Eq. 7. Therefore, we adopt the same loss of PDAE [49] but replace the gradient field with the summation of (2,, as
VaivI- O₁ G
C, = E le ca( z0.f.6
The aforementioned equation implies that zo can be reconstructed utilizing all disentangled represen- tations 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 V v, log p(= (xx) with decoder (T,, 2", 2t). In the next section, we introduce disentangling loss to address the remaining two conditions.
4.3 Disentangling Loss
In this section, present the Disentangling Loss to address the disentanglement requirement and ensure that Gp 2", serves as an approximation of Va, log logp(==)<;) 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 2 and * where o c, as a necessary condition for disentangling DPM.
Subsequently, we convert the objective of minimizing mutual information into set of constraints applied to the representation space. We denote dn is sampled from the pre-trained unconditioned
5

---

DPM using (0 t), 26 is conditioned on 2 and sampled using ( 2²). Then we can extract the representations from the samples 2g and Tij with Es as 2 = E (fo) and gkle = E (16), respectively. According to Proposition I, we can minimize the proposed upper bound (estimator) to minimize the mutual information.
Proposition 1. For a PGM in Figure (a), we can minimize the upper bound (estimator) for the mutual information between 2" and 2h (where l + c) by the following objective:
min
E
|| I - 11ahic =111 + || - =01
(10)
K,C.zo..5
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. order to minimize the first part, || =* |]. where * c, we need to minimize || a*1 I]. requiring that the k-th (k + c, E C) representation remains unchanged. We calculate follows:
the distance scalar between their k-th representation
(11) We represent the distance between each representation using distance vector d' [d1, d2,...,ccl, As the distances are unbounded and their direct optimization is unstable, we employ the CrossEntropy loss" to identify the index c, which minimizes the distances other indexes. The invariant loss Lin is formulated as follows:
dę = || - 1 -
Lin
E c..
(CrossFntropuu(d.c] c)].
(12)
Variant Loss. In order to minimize the second part, || - 1 || 1* 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:
df = 11 - ==11,
(13)
& = || - =
We adopt a cross-entropy loss to achieve the subjective by minimizing the variant loss Lya. we denote [d]. di.... , d' as d' and [df di
.... d&d as d².
Lvo
E
(CrossEnErrr Eutropput dd" d),c)l.
(14)
5..0.
However, sampling 3i from po(#1-1 1 l 2)) is not efficient in every training step. Therefore, follow [9, 49] to use score-based conditioning trick [41, 40] for fast and approximated sampling. We obtain approximator 66 of conditional sampling for different options of S by substituting the gradient field with decoder Gp:
=)
(15)
nES
To improve the training speed, one can achieve the denoised result by computing the posterior expectation following [8] using Tweedie's Formula:
Dt vI (Te,
S = E[zolree S]
(16)
Var
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:
7d = Al|ão - 2511 1 L. = L, + Taa + Lva).
(17)
where À is Syper-parameter We stop the gradient of ft and 4 for the weight coefficient Ya-
We denote the function torch. nn. www..nn..rrrss CrossEnroppLLLsss as CrossEntropy.

---

Table 1: Comparisons of disentanglement on the FactorVAE score and DCI disentanglement metrics (mean std, higher is better). Dis achieves state-of of off-the-a art performance with a large margin in almost all the cases compared to all baselines, especially on the MPI3D dataset.
Can3D
Shapes3I
MPE3D
Method
FactorVAF score
TXT1
FactorVAF score
DCT
FactorVAF score
DCT
VAE: haved:
FactorVAE B-TCVAB
0.906 0.052 0.835 + 0.082
0.161
0.019
0,840
0.006
0.611
0.082
0.152 0.0 0.025
0.240 + 0,051
0.140 = 0.019
0,873 + 0.074
0.613 = 0.114
0.179 = 0.017
0.237 = 0,056
GAN-bosd:
InfoGAN CR
0.411 1 0.013
0.020 1 0.011
0.587 1. 0.058
0.478 1. 0.055
0.439
0.061
0.241 0.075
Pre-troined GAN-based:
0.039
0.216 + 0.072
0.806 +
0.380
0.196
GS DisCo
0.532 + 0.018
0.209 +: 0.031
0.788 + 0.001
0.284 # 0.004
0.465 0.036
0.229 0.042
0.074
0.877
0.030
Diffusion- basedd:
DisDiff VQ (Ours)
0.976 +: 0.018
0.232 + 0.019
0.902 + 0.043
0.723 +: 0.013
0.617 0.070
0.337 0.037
Experiments
5.1
Experimental Setup
Implementation Details. can be sample in an image space or 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) VQ-version LDM DPM (DisDiff-Ki and DisDiff- VQ). For details of network Gg. we follow Zhang et al. [49] to use the extended Group Normalization [9] by applying the scaling & shifting operation twice. The difference is that learn-able position embedding to indicate c:
AdaGN t, (E(E,GN t 2/.
(18)
where GN denotes group normalization, and [tc. tfl. [ag, 2f{ are obtained from a linear projection: 25 2f = [inearProj(2 Proj(z"), t = linearPros Proj([t p²]). pf is the learnable positional embedding, the feature map of Unet. For more details, please refer Appendix A and B.
Datasets To evaluate disentanglement, we fol- low Ren et al. [36] to popular public datasets: Shapes3D [23], dataset of 3D shapes. MPI3D [12], 3D dataset recorded in a con- trolled 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 real- world datasets, we conduct our experiments on CelebA [29].
Table 2: Ablation study of DisDiff on image en- coder, components, batchsize and token numbers on Shapes3D.
Method
PactorVAF score
DCI
DisDiff-IN IM DisDifi-KL KL
0.783
0.655
0.837
0.660
Disiifi-QQ
0.902
0.723
DisDiff- VQ wo Cin
0.782
0.538
DisDiff-VQ wo L.o
0.810
0.620
Baselines Metrics We compare the per- formance with VAB-based and GAN-based baselines. Our experiment is conducted fol- lowing DisCo. All of the baselines and our method use the same dimension of the repre- sentation vectors, which is 32, referred to in DisCo. Specifically, the VAE-based mod- els include FactorVAE [23] and 3-TCVAE [5]. The GAN-based baselines include InfoGAN- CR [27], GANspace (GS) [14], LatentDiscov- (LD) [43] and DisCo [36]. Considering the
DisDiff-V VQ wo Liis wo detach
0.653
0.414
0.324
0.026
constant weighting loss weighting
0.679
0.426
0,678
0.465
attention condition pas embedding
0.824
0.591
0.854
0.678
orth emhedding
0.807
0.610
latent number N 6
0.865
0.654
latent number N=
0.902
0.723
influence of performance on random seed. We have ten runs for each method. We use two
representative metrics: the FactorVAE score [23] and the DCI [11]. However, since {=9|e E C}
are vector-wise representations, we follow Du et al. [10] to perform PCA as post-processing on the representation before evaluation.
7

---

SRC
SRC
TRT
TRT
Oren
Bangs
Wall
Skin
Floor
Smile
Color
Shape
Hair
Figure 3: The qualitative results on Shapes3D and CelebA. The source (SRC) images provide the representations of the generated image. The target (TRT) image provides the representation for swapping. Other images are generated by swapping the representation of the corresponding factor. The learned factors on Shapes3D are orientation (Oren), wall color (Wall), floor color (Floor), object color (Color), and object shape (Shape). For more visualizations, please refer to Appendix C-H.
5.2
Main Results
We conduct the following experiments to verify the disentanglement ability of the proposed DisDiff model. We regard the learned representation 2" as the disentangled one and use the popular metrics in disentangled representation literature for evaluation. The quantitative comparison results of disentan- glement 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 VAE-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 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-worie data into this paper. We have adopted the TAD and FID metries from [45] to measure the dis- entanglement and generation capabilities of Dis- Diff, respectively. As shown Table 3, Dis- Diff still outperforms new baselines in com- parison.
Table 3: Comparisons of disentanglement on Real- world dataset CelebA with TAD and FID metrics (mean + std). DisDiff still achieves state-of-the-an performance compared all baselines.
worth mentioning that InfoDif- fusion [45] has made significant improvements in both disentanglement and generation capabil- ities on CelebA compared to previous baselines. This has resulted DisDiff's advantage over the baselines CelebA being less pronounced than on Shapes3D.
Model
TAD
FID
3-VAR
0.088 0.043
99.8 + 2.4
InfoVAE
0.000 1. 0.000
77.8
1 1.6
Diffae
0.155 1: 0.010
22.7
2.1
InfoDiffusion
0.299 0.006
23.6
1.3
DisDiff (ours)
0.3054 + 0.010
18.249
2.1
The results are taken from Table 2 in [45], except for the results of DisDiff. For TAD, we follow
the method used [45] to compute 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 [48] to compute TAD, and the official GitHub
repository of [15] to compute FID.

---

TRT
TRT
Oren
Bangs
Wall
Skin
Floor
Smile
Color
Shape
Hair
Figure 4: The partial condition sampling on Shapes3D and CelebA. The target (TRT) image provides the representation of partial sampling images. Images in each row generated by imposing single gradient field of the corresponding factor on the pre-trained DPM. DisDiff samples image condition on only single factor. The sampled image has fixed factor, e.g., the images in the third row have the same background color as the target one. The conditioned factors on Shapes3D are orientation (Oren), wall color (Wall), floor color (Floor), object color (Color), and object shape (Shape).
5.3 Qualitative Results
In order to analyze the disentanglement of DisDiff qualitatively, we swap the representation =" of two images one by one and sample the image conditioned the swapped representation. We follow LDM-VQ to sample images in 200 steps. For the popular dataset of disentanglement literature, take Shapes3D Compared with the VAE-based methods, DisDiff has better image quality. Since there are no ground truth factors for the real-world dataset, take CelebA
example. As Figure 3 shows, DisDiff successfully learned pure factors.
example. As demonstrated in Figure 3, DisDiff also achieves good disentanglement on real-world datasets. Please note that, unlike DisCo [36], Dis 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 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.
DPM type The decomposition of the gradient field 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 is much casier in the latent space of the auto-encoder. Therefore, diffusion types: image diffusion model, e.g., DDIM (DisDiff-IM) KL-version latent diffusion model (DisDiff-KLL KL), and VQ- re-versoon latent diffusion model, e.g., VQ-LDM (DisDiff-Vv VQ). As shown in Table 2, the LDM-version DisDiff outperforms the image-versin version DisDiff as expected.
compare the performance of DisDiff with different
Disentangling Loss Disentangling loss comprises two parts: Invariant loss Lin and Variant loss Lvo- To verify the effectiveness of each part, we train DisDiff-VQ without it. Lin encourages in-variance of representation not being sampled, which means that the sampled factor will not affect the representation of other factors (z*, k * e of generated 26). On the other hand, Leo encourages the representation of the sampled factor (2 of generated 26) to be close to the corresponding one of zg. As shown in Table 2, mainly Cin promotes the disentanglement, and Eya further constrains the model and improves the performance. Note that the disentangling loss is optimized w.r.t. Go but not Ei 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 the encoder as the weight coefficient. In this section, we explore other options to verify its effecttivnness. 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 a different extent.
Condition type & Latent number DisDiff follows PDAR [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 peer-pprraettrr 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 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 DisDiff capable of sampling information exclusively to conditions.
Limitations
Our method is completely unsupervised, and without any guidance, the learned disentangled represen- tations on natural image sets may not be easily interpretable by humans. We can leverage models like CLIP guidance improve interpretability. Since DisDiff is a diffusion-based method, compared and GaNebased, the generation speed is slow, which is also a common problem for
to the DPM-based methods. The potential negative societal impacts are malicious uses.
7 Conclusion
In this paper, we demonstrate new task: disentanglement of DPM. By disentangling a DPM into several disentangled gradient fields, we can improve the interpretability DPM. To solve the task, we build an unsupervised diffusion- based disentanglement framework named DisDiff. DisDiff learns disentangled representation of the input image the diffusion process. 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- aub-gradient field. Applying DisDiff to more general conditioned DPM is direction worth exploring for future work. Besides, utilizing the proposed disentangling method to pre-trate pre-trained conditional DPM makes more flexible.
Acknowledgement
We thank all the anonymous reviewers for their constructive and helpful comments, which have significantly improved the quality of the paper. The work was partly supported by the National Natural Science Foundation of China (Grant No. 62088102).
References
[1] Omri Avrahami, Dani Lischinski, and Ohad Fried. Blended diffusion for text-driven editing of natural images. In Proceedingy of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 18208-18218, 2022.
[2] Yoshua Bengio, Aaron C. Courville, and Pascal Vincent. Representation learning: A review and new perspectives. PAMI, 2013.
10

---

[3] Christopher P Burgess. Irina Higgins, Arka Pal. Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling in beta-vae. (tth::800403599 2018.
[4] Jinkun Cao, Ruiqian Nai, Qing Yang, Jialci Huang, and Yang Gao. An empirical study on disentanglement of negative-frce contrastive learning. arXiv preprint arXiv:2206.047333 2022.
[5] Ricky TQ Chen, Xuechen Li, Roger B Grosse, and David K Duvenaud. Isolating sources of disentanglement in variational autoencoders. In NeurPIS, 2018
[6] Pengyu Cheng, Weituo Hao. Shuyang Dai, Jiachang Liu, Zhe Can, and Lawrence Carin. Club: A contrastive log-ratio upper bound of mutual information. In International conference on machine learning. pages 1779-1788. PMLR, 2020.
[7] 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.
[8] Hyungjin Chung, Byeongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving diffusson models for inverse problems using manifold constraints. Advances in Neural Information Processing Systems, 35: 25683 25696,5666622 2022.
[9] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in Neural Information Processing Systems, 34:8780-8799 2021.
[10] 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.
[11] Cian Eastwood and Christopher K. I. Williams, A framework for the quantitative evaluation of discntangled representations. In ICLR, 2018.
[12] 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: new disentanglement dataset. In NeurlPS, 2019.
[13] Ian Goodfellow. Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communicutions of the ACM, 63(11):
139- 100-144,22 2020.
[14] Erik Härkönen, Aaron Hertzmann, Jaakko Lehtinen, and Sylvain Paris. Ganspace: Discovering interpretable GAN controls. In NeurlPS, 2020.
[15] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by two scale update rule converge to a local nash equilibrium. Advances in neural information processing systems, 30, 2017.
[16] Trina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with constrained variational framework. In ICLR, 2017.
[17] Irina Higgins. David Amos. David Pfau, Sebastien Racaniere, Loic Matthey, Danilo Rezende, and Alexan- der Lerchner. Towards definition of disentangled representations. arXiv preprint 1rXiv:1812.0223381 2018.
[18] Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207. 12598. 2022.
[19] Jonathan Ho, Ajay Jain, and Pieter Abheel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-685/ 2020.
[20] Alexia Jolicoeur-Ne aaicceer-Marttinaau 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.05477 2020.
[21] Bahjat Kawar, Shiran Zada, Oran Lang, Omer Tov, Huiwen Chang, Tali Dekel, Inbar Mosseri, and Michal Irani. Imagic: Text-hased real image editing with diffusion models. arXiv preprint arX0v:2210.09276. 2022.
[22] Gwanghyun Kim, Taesung Kwon, and Jong Chul Ye. Diffusionneli:: Text-guided diffusion models for robust image manipulation. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 2426-2435, 2022.
[23] Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In ICML, 2018
[24] Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv: 2013.
[25] Mingi Kwon, Jaeseok Jeong, and Youngjung Uh. Diffusion models already have a semantic latent space. arXiv preprint arXiv:2210.1090 2022.

---

[26] José Lezama. Overcoming the disentanglement vs reconstruction trade-off via jacobian supervision. In International Conference on Learning Representations, 2018.
[27] Zinan Lin, Kiran Thekumparampil, Giulia Fanti, and Sewoong Oh. Infogan-er and modelcentrality: Seif seff-supervieed model training and selection for disentangling gans. In ICML, 2020.
[28] 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.
[29] 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.
[30] 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. international conference on machine learning, pages 41 14 4124, PMLR, 2019.
[31] Andreas Lugmayr, Martin Danelljan, Andres Romero, Fisher Yu, Radu Timofte, and Luc Van Gool. Repaint: Inpainting using demoising diffusion probabilistic models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11461-11471, 2022.
[32] 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.
[33] Konpat Preechakul, Nattanat Chatthee, Suttisak Wizadwongsa, and Supasorn Suwajanakom. Diffusion autoencoders: Toward a meaningful and decodable representation. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
[34] 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.
[35] Scott E. Reed, Yi Zhang, Yuting Zhang, and Honglak Lee. Deep visual analogy-making In NewrlPS, 2015.
[36] 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.
[37] Chitwan Saharia, William Chan, Saurahh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S Sara Mahdavi, Rapha Gontijo Lopes, et al. Photorcalistic text-to- image diffusion models with deep language understanding. arXiv preprint arXiv:2205. 11487, 2022.
[38] Jascha Sohl-Dia Dickksicst,in, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised lcarning using nonequilibrium thermodynamics. In International Conference on Machine Learning, pages 2256-220 2265. PMLR, 2015.
[39] Jiaming Song. Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020.
[40] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in Neural Information Processing Systems, 32, 2019.
[41] Yang Song, Jascha bohl-Dickstein, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. ICLR, 2020.
[42] 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 discntangled representation learners via multi-stage modelling. CoRR, abs/2010. 13118, 2020.
[43] Andrey Voynov and Artem Babenko. Unsupervised discovery of interpretable directions in the GAN latent space. In ICMI,, 2020.
[44] 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.
[45] Yingheng Wang. Yair Schiff, Aaron Gokaslan, Weishen Pan, Fei Wang, Christopher De Sa, and Volodymyr Kuleshov. Infodiffusion: Representation learning using information maximizing diffusson models. interna- tional conference on machine learning, 2023.
[46] 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-imaa diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1900- 1910, 2023.
12

---

[47] Tao Yang. Xuanchi Ren. Yuwang Wang, Wenjun Zeng, and Nanning Zheng. Towards building group- based unsupervised representation disentanglement framework. In International Conference Learning Representations, 2021.
[48] 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.
[49] Zijian Zhang, Zhou Zhao, and Zhijie Lin. Unsupervised representation learning from pre-trained diffusion probabilistic models. NeurIPS, 2022.
13

---

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 Training procedure
Prepare:
dataset distribution p(zo). pre-trained DPM Eg, pre-trained VQ-VAE, Evq, DvQ. Initialize: encoder Esg decoder Gg.
repeat
2g ~ p(20)
Zzu EvQ
Uniform (0, ......T) T)
I).
Ext Vorzan
VI One
Uniform(1
:.
2 = (T(
C, = - (o t) y yorvi-a o4 Eiec GylEze; ,,,t,
De
Use Zatha t, Eg. EvQ. DvQ and to calculate Lin and Lva by Eq. 12 and Eq. 14. Calculate weight "id by Eq. 17
Update by taking gradient descent step
Vo + + (a(Lin + Lra))
until converged;
Algorithm 2 DDPM Sampling
Prepare: pre-trained DPM Eø, pre-trained VQ-VAE, Evq, DvQ. the sampling factor S C C,
Input: image œ
(x)
2rT (0. I)
for 7 do
if t * then A/
else
E = end if
2ex VI-ai co(2us,))
Eies Golzue, 24, + ore
(x(t-1)
vor
end for return DvQ(2x0)
B Training Details
Encoder We use the encoder architecture following DisCo [36] for fair comparison. The details of the encoder are presented in Table 3, which is identical to the encoder of 3-TCVAE and FactorVAE Table 1.
Decoder We use the decoder architecture following PADE [49]. The details of the decoder are presented Table 4, which is identical to the Unet decoder 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 64 for all datasets. We always set the learning rate 1e 4. We use EMA all model
14

---

Algorithm 3 DDIM Sampling
Prepare: pre-trained DPM co. pre-trained VQ-VAE Evq. DvQ. the sampling factor set S C C, Input: sample sampling sequence {ti)en and t 0,tk =T
(x)r) ZxT ~ (0, I) inferred 2xT of inverse DDIM
for j = K do
6B (2(xtt 2 (1)) co(2ntt Eies vI 11-rr Ont Gpi(Zotyy (1,
Jely
/ co (4) + vI City e0 (0(2ztt 20 ,11)
Eart v 0its
Vayy
end for return Dve(Zuo)
parameters with a decay factor of 0.9999. Following DisCo [36], we set the representation vector dimension to 32.
Table 4: Encoder E architecture used in DisDiff.
Conv 7 64, stride
ReLu
Conv x 4 x 64 x 128, stride
ReLu
Conv x x 128 256, stride = 2
ReLu
Conv 4 x 256 256, stride = 2
ReLu
Conv x 4 x 256 256, stride = 2
ReLu
FC 4096 x 256
ReLu
FC 256 x 256
ReLu
FC 256 x K
Table 5: Decoder (pre-trained DPM) architecture used in DisDiff.
Parameters
Shapes3D Cars3D/ MPI3D
FFHQ
CelebA-HQ
Base channels
16
64
64
Channel multipliers
[ [12,,444
[1,2,3.4]
[1,2,3,4]
Attention resolutions
[ 11 2. 4]
[2,4,8]
[2,4,8]
Attention heads num
8
Model channels
64
224
224
Dropout
0.1
Images trained
0.48M/0.28M/1.032 .03M
130M
52M
scheduler
Linear
Training T
1000
Diffusion loss
MSE with noise prediction
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.
15

---

SRC
TRT
None
None
Azimuth
Grey
None
Car color
Figure 5: The qualitative results on Cars3D. The source (SRC) images provide the representations of the generated image. The target (TRT) image provides the representations for swapping. The remaining rows of images are generated by swapping the representation of the corresponding factor on Shapes3D. DisDiff learns pure factors each representation. The learned factors are None, None, azimuth, grey, and car color.
D Visualizations on MPI3D.
We provide the qualitative results MPI3D in Figure 6. We observe that even such challenging dataset (MPI3D) disentanglement, DisDiff also learns disentangled representation. In addition, 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 learnable matrix A _Cxd, where d is the dimension of direction, have the following equation:
(19) where Q E CxCd is orthogonal matrix, and R d is an upper triangular matrix, which is discarded our method. We use to provide directions in Figure 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 Ir for image editing. The sampling process shown Algorithm 3.
= QR
F Qualitative Results on CelebA-HQ
In order to demonstrate the effectiveness of Orthogonal-constrainn DisDiff, we train on dataset CelebA-HQ. We also conduct our experiments with the latent unit number 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 Figure 11, Orthogonal- constrained DisDiff lears consistent factors for different images. There are more factors than Figure shows. Therefore, we conduct another experiment using another random seed with latent unit number of 10. The results shown in Figure 10.
16

---

SRC
TRT
OB color
Ob shape
3
None
3
None
3
None
BG color
Figure 6: The qualitative results on MPI3D. The source (SRC) images provide the representations of the generated image. The target image provides the representations for swapping. The remaining rows of images are generated by swapping the representations of the corresponding factor on MPI3D. DisDiff learns pure factors by each representation. The learned factors are object color (OB color), object shape (OB shape), None, None, None, and background color (BG color).
nee-trained
éa (x,, ()
Encoder Esp
Disentangling
(x,x,,t,,)))
Decoder G
Encoder F
Figure 7: Illustration of Orthogonallconstrainne DisDiff. In this architecture, we use a powerful decoder to replace the original decoder in DisDiff. Since we train the decoder from scratch, the ability disentanglement on large natural datasets is limited. However, decoder our powerful decoder by feeding set of learnable orthogonal directions to make it output corresponding gradient field. These directions and the decoder compose a powerful decoder that takes directions input and generates the corresponding gradient fields as output.
can use the pre-trained Unet
G CLIP-guided DisDiff
To tackle the disentanglement in real-world settings, we propose replacing the encoder in DisDiff' with powerful encoder, which lets 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 the pre- pre-trained information we can leverage. Therefore, we use the CLIPP text encoder to provide the condition but use the CLIP image encoder to calculate the disentangling loss. Figure shows the obtained new architecture.
17

---

Pre-trained DPM
Image Encoder Ei
Disentangling loss
Decoder Go -
tx
CUP image Encoder Eg
t²(x),t22))
Encoder Egg
Attribute words
Figure 8: Illustration of CLIP-guided DisDiff. this architecture, a powerful encoder replaces the original one in DisDiff. Specifically, adopt the CLIP Image encoder as 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 CLIdde DisD to edit the images by using the inverse ODE of DDIM to infer the ZT.
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 sig- nificantly decrease, which in accordance with our intuition. However, the performance other methods will also decrease under this condition. We conducted experiments using FactorVAE and 3-TCVAE to demonstrate it. We trained our model, FactorVAE and 3-TCVAE on Shapes3D (The number of factors is 6) with 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 3-TCVAE.
Table 6: The disentanglement performance when the latent number is 3 (mean + std, higher is better).
Model
DCI score
Factor VAE score
3-TCVAE
0.098 + 0.043
0.383 + 0.0444 0.4982 0.0049
FactorVAE
0.3113 0.014
DisDiff (ours)
0.3279 + 0.023
0.4986 + 0.0054
Quantitative Results on More Metrics
We experiment with MIG and MED metrics on the Shapes3D dataset, two classifier-free metrics. For MIG, utilized the implementation provided in disentanglement, emmee_liintttttttttt while for MED, used the official implementation available at Github". The experimental results in Table indicate that our proposed method still has significant advantages, especially in MED-topk and MIG.
http://gihit..ooomooooooeerrreet https//ggttbb.oommooocccoottttns _ed
18

---

Table 7: Comparisons of disentanglement on the additional metries: MED, MED-topk, and MIG (mean std, higher better). Except for DisCo and DisDifff other results (MED and MED-topk) are taken from the MED [4]. The missing value is due to the need for publicly released pre-trained models.
Model
MED
MED-topk
MIG
8- B.V
52.5 + 9.4
19.2 1.4
3-TCVAE
53.2 4.9
25.0 0.5
0.406 + 0.175
FactorVAE
55.9
8.2 4.0
0.434 + 0.143
DIP-VAE-I
43.5 1
16.2 = 0.9
DIP-VAE-II
52.6 + 5.2
AnnealedVAE
56.1 + 1.5
MoCo
18.1 = 0.6
MoCov2
13.6 1.7
Barlow Twins
20.0 = 0.3
SimSiam
30.0 2.0
BYOL
19.7 = 1.3
InfoGAN-CR
0.297 + 0.124
LD
0.168 + 0.056
DS
0.356 + 0.090
DisCo
38.0 + 2.7 59.5 + 1.1
54.8 2.1
0.512 + 0.068
DisDiff (Ours)
64.5 1.6
0.5671 + 0.033
K Proof of Proposition
Background: Our encoder transforms the data zy (hereafter denoted ær the encoder keeps fixed for disentangling loss) into two distinct representations z" and z We assume that the encoding distribution follows Gaussian distribution p(2*(zx,, with a mean of 5² However, the distribution p(z[2²) is intractable. We can employ conditional sampling of our DPM gx(x(2") to approximate the sampling from the distribution p(x]="). Given that we optimize the disentangling loss each training step, we implement Eq. 15 and Eq. 16 in the main paper an efficient alternative to sample from qo (z]2²). We further assume that the distribution used to approximate p(=*l2") follows Gaussian distribution (2*1)))) with a mean of A4 E* (D(2))) where represents the operations outlined in Eq. 15 and Eq. 16. We represent the pretrained unconditional DPM using the notation go (x()) R p(x).
Proposition 2. For a PGM: =- + x + 2h, we can minimize the upper bound (estimator) for the mutual information between 2" and 2h (where + c) by the following objective:
min
E
| z* ||
=*1 || |1 =
(20)
Considering the given context, we can present the PGM related to q follows:2 + - zt The first transition is implemented by DPM represented by the distribution go. Subsequently, the second transition is implemented by the encoder, denoted by the distribution p.
Proposition 1.1 The estimator below an upper bound of the mutual information I(2 2²). For clarity, we use expectation notation with explicit distribution.
Test = Ep(ss, (log p(2 |
E Dogg
(21)
Equality is attained if and only if ** and of are independent.
Proof: Lest - I(2",20 *) can be formulated as follows:
Ep( (age:s)loogple 1 [2 - E pp( [log - Ep(es w (o) logp(>*))]
(pp [log p(z x)) Bp(=))p(=)) [log((=*(<0)
(22)
Ep(()) (logp(>*) Ep(sa) [log ((2 12)].
19

---

Using Jensen's 's Inequality, we obtain logp(z*) = log(Ep(=9) (B((Bp(=)\p((=(<>)) [2()* Therefore, the following holds:
> ppt=e)logp(==1<7))
Test I(2 24) EE p(= pog pp( (lg ((==|=))))111 O.
(23)
Since the distribution p==* is intractable, we employ 22 to approximate 12)
la pt [log go (=412]] (pz))) [log go (z"))
(24)
Proposition 1.2: The estimator can be expressed follows:
La -p ||2 1 (u))) ( || 14 + C₁.
(25)
The first term corresponds to dz 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:
log (0 = log -
-E* (D(2))) - C2 -11=* -2hlo| + C2.
(26)
For the second term, obtain:
/ p(z, 2 log go 129 dr
E Ep(z)) [log go (2 (2
/ p(x)) / p(=* log go 1= dd d dx
f mo( / 1(=*)0) log (= :9) dā
(27)
(ay (2 (lo go (=*129]
Eg p ( - E*th (D(=))) + C3
yyy E ( = C3.
Proposition 1.3: If the following condition is fulfilled, the new estimator To can serve upper bound for the mutual information 1(2", ;²),, i... La > 1(2", The equality holds when z² and 2 are independent. We denote g0 (p(2)))))(2>)) as q0(2" 29) 29).
DKL(p(=* (0)00 2²)) < DKu(p(=")p<>))) pp) ,
(28)
Proof:
To I(= 2) can be formulated follows:
tt) flog (2 (= E (x=))) [log go [2 Ep(ze [
Dog 1(og((=*122)) 1ogp(2"))
= Ep( [log 2a 12)) log (0xp(=11>0) [>)] + E (== log)(=*))
-Bpleepoea - log (= + log p( )1 (og -
(og (=
(29)
-Ep(F))((*) log p(zt )p(29)
p(=
E Ep(ce,a)) log (===+
10(=*12<))(>
DDx(p(2")p(=" (2"0""" DKL(p(2 (ac2", 20 (1900((2" (())
If DKi(p(=h z (2) & DrI(p(=*)p(C (1200))00(=) ))),, obtain:
To 1(
(30)
Since the condition is not yet applicable, we will continue exploring alternative conditions.
Proposition 1.4: If (xz(p(2 ( (1900)))))))) 1g 1(25,2 or [/ I(2 2*) e.
e, then, either of the following conditions holds:
20

---

Proof: If DKu(p(e", (00l00"""")) & DKL(p(=")))) ((2 29), then
To 2 1(2
(31)
else:
DxI(o(=")p<>"lll Iloo (9)) < DRI(P(=* : =l0o0(> (
-Buesr,t [ros p(=*(29)
g0(=*1
(32)
ee
((2, max(Dx<(p(2", 29000 2))), DxL( (x)
< (33)
Proposition 1.5: Consider PGM + x + 2h: if D (00
0. it follows that
DKL (xL((<<))))l 0.
Proof: if (oo(())=))))/ppp//)) = O, we have (x(2(2)) = p(x[2²).
log p(=*129)
4o (0
(34)
log I go (p( (p |x)
Ep(ss .*
[log 1] 0.
Ideally, by minimizing (xr(g0 to we can establish either as an upper bound or as an estimator of [22%, Proposition 1.6: we can minimize D (20 (11)))))p))=))>>>>)))))))))))))) by minimizing the following objective:
Ep(= (el ||=== ædle| C4.
(35)
Note that the equation above exactly corresponds to de in Eq. 13, and dg is constant. Proof:
p(=7) / 90(xl22) log 20(x
DKL [x
p(x(20) dxdza.
(36)
We focus on the following term:
f go (0(xl2") log (x(22 dx / go [ogall"))oo - log p(
(37)
p(x]2=)
Applying the Bayes rule, we obtain:
4o (x( (=]2=)))))" log p(=91x)p(x) dx
p(=0)
f (o(())=)) log (x(22) (0gp(=">x)) + logr(=") de
p(r)
[x)] + logp(29) & Ego (ooptt>))lpp11011000))0))) logr(=")
(xl(g)(r12")) (p(x))
(38)
|= |oggex(( exp(-pppppppp
E (x)10 log p(29) + Cs
= E 2d + log + Cs.
We minimize B (z) -
zele| instead.
Note that the loss usually cannot optimized to a global minimum in a practical setting, but how the gap influences To 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(={2=) using (x([²)). Therefore, understand that the model's loss can be regarded as Leat-
21

---

Table 8: Ablation study of DisDiff on image encoder, components, batchsize and token numbers Shapes3D.
Method
SSIM(())
LPIPS(())
MSE())
DCI(1)
Factor VAE(1)
PDAE
0.9830
0,0033
0.0005
0.3702
0.6514
DiffAE
0,9898
0.0015
8.1983e
05
0.0653
0.1744
DisDiff (ours)
0.9484
0.0006
9.9293e 05
0.723
0.902
Comparison with Diffusion-based Autoencoders
We include comparison with more recent diffusion baselines, such PADE [49] and DiffAE [33]. The results are shown in Table 8.
The relatively disentanglement performance of the two baselines can be attributed to the fact that neither 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., Vet [16], betaTCVAE [5], FactorVAE [23]). However, fortunately, disentanglement our model does not lead 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 well-pre-trained generative model helps avoid this trade-off.
22

---

SRC
Pose
Jaw
Bangs
Expose
Long hair
Figure 9: The qualitative results orthogonal constraint on CelebA-HQ (seed 0). The source (SRC) row provides the images of the generated image. The remaining rows of images are generated by shifting the representation along the learned direction of the corresponding factor on CelebA-HQ. DisDiff learns the direction of pure factors unsupervised. The learned factors are listed in the leftmost column. The images are sampled by using the inferred æT.
23

---

SRC
Black hair
iite
Bangs
Pose
Gender & Pose
I
L
Pose 2
Figure 10: The qualitative results orthogonal constraint on CelebA-HQ (seed 1). The source (SRC) row provides the images of the generated image. The remaining rows of images generated by shifting the representation along the learned direction of the corresponding factor on CelebA-HQ. DisDiff learns the direction of pure factors unsupervised. The learned factors are listed in the leftmost column. The images are sampled by using the inferred æT.
24

---

SRC
Expose
Bangs
Long hair
Thin face
Pose
Bangs2
Expression
Long hair2
Bald
Bangs
Less hair
Pale face
Smile
Zoom in
Figure 11: The qualitative results of orthogonal constraint CelebA-HQ (latent number 30). The source (SRC) row provides the images of the generated image. The remaining rows of images are generated by shifting the representation along the learned direction of the corresponding factor CelebA-HQ. DisDiff learns the direction of pure factors unsupervised. The learned factors are listed in the leftmost column. The images sampled by using the inferred ør.
25

---

SRC
Tanned face
Smile
Sad
Red hair
Blond hair
Bald
Old
Woman
Yellow hair
The qualitative results of CLIP-guided DisDiff on FFHQ. The source (SRC) row provides the images of the generated image. The remaining rows of images are generated by shifting the representation along the learned direction of the corresponding factor FFHQ. DisDiff learns the direction of pure factors unsupervised. The learned factors are listed in the leftmost column. The images sampled by using the inferred æT.
Figure 12:
26

---

Young
Surprised
Smile
Sad
Red hair
Blond hair
Bald
Old
Figure 13: The partial condition sampling results of CLIP-guided DisDiff on FFHQ The source (SRC) row provides the images of the generated image. The remaining rows of images are generated by shifting the representation along the learned direction of the corresponding factor on FFHQ. DisDiff learns the direction of pure factors unsupervised. The learned factors are listed in the leftmost column. The images are sampled by using the inferred zr.
27


Reviews:
[{"review_id": "Kd1L5z4p1q", "reviewer": "Reviewer_6mFa", "summary": "This 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.", "questions": "- 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?", "limitations": "There is a section of discussion for limitations.", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Strengths:\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.", "weaknesses": "Weaknesses:\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable comments and suggestions and the encouraging words. We appreciate the time and effort you took to review our paper. Please find below our responses to your concerns and the changes we have made to address them.\n\n\nWe have thoroughly revised Section 4 of our paper. We have made significant changes in the presentation and organization of the section, ensuring that our arguments are clearly articulated and well-supported.  We highly encourage the reviewers to revisit our revised Section 4, as we are confident that the improvements made will provide a clearer understanding of our research. We appreciate the valuable feedback provided by the reviewers and look forward to receiving further comments that will help us refine and strengthen our paper.\n\n\n- Overclaim: We apologize for the oversight in not citing the recent work [1] in ICML 2023. We have now added the reference to this work and revised our claims accordingly in the corresponding paragraphs. We agree that there might have been a timeline issue that led to our unawareness of this work, and we appreciate your understanding. We also acknowledge that our work and [1] can be considered as concurrent works. and we will make this clear in our revised manuscript.\n- Abuse of notations: We apologize for any confusion caused by the notational errors and inconsistencies in our manuscript. We rewrite section 4.3 for better clarity and have carefully revised and corrected these notations. Please refer to the comment “Section 4”. We appreciate your detailed feedback and suggestions, which have helped us improve the clarity of our presentation.\n- More baselines: We agree that including more recent diffusion baselines such as PADE and DiffAE would strengthen our evaluation. We will include these baselines in our experiments and update the results in the paper in the Discussion stage. \n- More metrics for CelebA: We appreciate your suggestion to include the tailored quantitative metric [4] for evaluating disentanglement performance on CelebA. We will add this metric  in the Discussion stage.\n\nFor questions:\n- We propose a modified framework for disentangling the DPM of real world data in Appendix G, I and present the results trained on CelebA HQ and FFHQ. We infer  that one possible explanation for the observed distortion is due to the requirement of  reconstruction and the weak ability of pretrained DPM. The disentanglement process might not be so perfectly captured due to the challenge of the complex real world dataset. Another factor could be the choice of loss functions, which might not be optimal for achieving a perfect balance between disentanglement and image quality. We leave it for future work.\n\n- We have modified the draft in Section 4 for clarity. Eq.6 in our submission is taken from the classifier guidance trick proposed in [1] (Please see equation 14 in it), the role of it  is to build the score function of conditional distribution.  Eq.7 in our submission is taken from the reverse problem proposed in [2] (Please See Equation 10 in it), the role of it is to recover the data sample with the score function. We want to use these equations to sample the recovered samples based on representations.\n\n[1] Diffusion Models Beat GANs on Image Synthesis. \n\n[2] Improving Diffusion Models for Inverse Problems using Manifold Constraints. NeurIPS 2022\n"}, {"review_id": "sWxZT3mZ4O", "reviewer": "Reviewer_ZKsn", "summary": "This 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.", "questions": "- 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.", "limitations": "Section 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.", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- 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.", "weaknesses": "- 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable comments and suggestions. We appreciate your feedback and have made several changes to address your concerns. Please find below our responses to each point you raised.\n\nWe have thoroughly revised Section 4 of our paper. We have made significant changes in the presentation and organization of the section, ensuring that our arguments are clearly articulated and well-supported. We highly encourage the reviewers to revisit our revised Section 4, as we are confident that the improvements made will provide a clearer understanding of our research. We appreciate the valuable feedback provided by the reviewers and look forward to receiving further comments that will help us refine and strengthen our paper.\n\n\n- Overclaim: We apologize for any overclaiming in our previous manuscript. We have now revised the manuscript and properly toned down our claims accordingly.\n- Definitions and mathematical formulation: We have carefully revised the manuscript to provide clearer definitions and improved mathematical formulations. Please refer to comment “Section 4”.We believe that these changes have significantly improved the clarity of our paper and helped address your concerns regarding the understanding of the proposed method. \n- Necessary condition for disentanglement: We agree that the conditional independence of the representations is a necessary, but not sufficient, condition for disentanglement. We have revised the manuscript to clarify this point and provide a more accurate statement: \"We follow prior works in disentangle representation literature that propose a necessary condition that works in practice for disentanglement. For example, the group constraint in [43], the maximization of mutual information in [23, 41]. We propose to minimize the mutual information between as a necessary condition.\"\n- Contradictory claims: We apologize for any confusion. We want to demonstrate that：“ Due to the fact that disentangled representation learning on real-world data is quite challenging. While our method demonstrates effectiveness on synthetic datasets, we find that the learned disentangled representations on natural image datasets may not be easily interpretable by humans. We infer that the possible reasons behind it are the requirement of  reconstruction and the weak ability of pretrained DPM. Therefore, we have also included additional results with a more powerful DPM in the appendix that demonstrate the effectiveness of our method for disentangled generation on real-world data. This shows that our method can be beneficial for disentangled generation, We have now revised the manuscript properly and accordingly.\n\nQuestions:\n- We apologize for any confusion caused by definitions and mathematical formulation. We have carefully revised the manuscript to provide clearer definitions and improved mathematical formulations. Please refer to the comment “Section 4”. We highly recommend you read that part first, we will clarify the questions proposed.\n- “disentangled conditional sub-gradient fields”: We assume that the dataset is generated by N factors, $D=\\lbrace x_0|x_0\\sim p(x)\\rbrace$, there is a one-one mapping: $h: (f^1,f^2,\\dots,f^N)\\rightarrow x_0$, there is a set of conditional distribution $\\lbrace p(x|f^c)|c = 1,\\dots,N\\rbrace$. The disentanglement of DPM is to learn the score function  $\\nabla_{x_t} \\log p(f^c|x_t)$ so that we can sampling from conditional distributions $p(x_t|f^c)$.\n- ”the data sample space is not well-organized“ We mean that the data space is complex, it is difficult to sampling from $p(x|f^c)$, when we have the score function $\\nabla_{x_t} \\log p(x_t)$, by this we want to emphasize the difficulty of the problem.\n- “Cross entropy”: Please note that the cross entropy here we use is the cross entropy loss function in classification problems. The inputs of this function are “logits” and “labels”(GT). In our setting, we regard the distance as logits and indexes as labels.  As comment “Section 4” shows, we use this cross entropy loss as a tool for implementation of minimizing the distance in the Proposition. Since cross entropy loss is a bounded function, the usage of such a loss function has the benefit for training stability.\n\nWe appreciate your constructive feedback, which has helped us improve the quality and clarity of our paper.\n"}, {"review_id": "6l4VIoZyC5", "reviewer": "Reviewer_kfVZ", "summary": "Imprecisions\n- “According to the completeness requirement” add reference\n- line 75-76 PADE-> PDAE\n- line 270-271. Typo: compare with disco. dissdiff -> compared with disco, disdiff\n\nStrengths\n- Achieving disentangled representations in diffusion models is an interesting and useful problem.\n\nWeaknesses\n- Missing derivation of reverse diffusion process, Eq 5, 6, 7.\n- What is the difference between the proposed model and a latent diffusion model?\n- Diff-AE and PDAE should also be considered as baselines in experiments.\n- Image generation experiments against Diff-AE and PDAE should be conducted in addition to disentanglement in order to add robustness to the experiments.\n- Unclear how Figure 3 is generated. Could you describe the process in detail and explain what you mean by “swapping the representation”.\n", "questions": "see summary", "limitations": "see summary", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "see summary", "weaknesses": "see summary", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your valuable comments and suggestions. We appreciate your feedback and have made several changes to address your concerns. Please find below our responses to each point you raised. \n\nWe have thoroughly revised Section 4 of our paper. We have made significant changes in the presentation and organization of the section, ensuring that our arguments are clearly articulated and well-supported. We highly encourage the reviewers to revisit our revised Section 4, as we are confident that the improvements made will provide a clearer understanding of our research. We appreciate the valuable feedback provided by the reviewers and look forward to receiving further comments that will help us refine and strengthen our paper.\n\n- Imprecisions: We have revised the manuscript according to your suggestions and corrected the imprecisions. We have the typos in the main paper.\n- Derivation of reverse diffusion process (Eq. 5, 6, 7 in our submission): Based on your feedback, we have now included the derivation of these equations.Please refer to comment “Section 4”. Eq.5 is taken from [3] (Please see Equation 5 in it). Eq. 6 is taken from the classifier guidance trick proposed in [1] (Please see Equation 14 in it), the role of it  is to build the score function of conditional distribution.  Eq. 7 is taken from the reverse problem proposed in [2] (Please see Equation 10 in it), the role of it is to recover the data sample with the score function. We believe that this addition helps improve the clarity and completeness of our presentation.\n- Difference between the proposed model and a latent diffusion model: The key difference between our proposed model and a latent diffusion model is that the latter is a general DPM without disentangled representations and the ability to sample according to independent factors. Our method is designed to disentangle general DPMs, providing them with these capabilities. We have clarified this distinction in our revised manuscript.\n- Additional baselines and experiments: We appreciate your suggestion to include Diff-AE and PDAE as baselines in our experiments. We will add these baselines and conducted image generation experiments against them in addition to disentanglement in discussion stage. \n- Clarification on Figure 3: We apologize for any confusion regarding the generation of Figure 3. In our revised manuscript, we have provided a detailed explanation of the process. To \"swap the representation,\" we first encode the images to obtain their representations. We then exchange the representations corresponding to the same factor between two different images. Finally, we decode the swapped representations to generate new images.\n\nWe believe that addressing your concerns has significantly improved the quality and clarity of our paper. We appreciate your constructive feedback.\n\n[1] Diffusion Models Beat GANs on Image Synthesis. \n\n[2] Improving Diffusion Models for Inverse Problems using Manifold Constraints. NeurIPS 2022\n\n[3] Unsupervised Representation Learning from Pre-trained Diffusion Probabilistic Models. NeurIPS 2022\n"}]

Author rebuttal:
# Section 4
We assume that a dataset $\mathcal{D} = \lbrace x_0\sim p(x_0)\rbrace$ is generated by $N$ underlying ground truth factors $\mathcal{C} =\lbrace f^c|c=1,\dots,N\rbrace$, where $p(x_0)$ is the data distribution, i.e., each sample is generated by the underlining factors, $h: (f^1, \dots, f^N) \mapsto x_0, \forall x_0 \in \mathcal{D}$ and each factor $f^c \sim p(f^c)$, where $p(f^c)$ is the distribution of factor $c$. The conditional distributions of the factors are $\lbrace p(x_0|f^c)|c=1,\dots,N\ \rbrace$, each of which can be shown as a curved surface in Fig. A (b). 
The relation between the conditional distribution and data distribution can be formulated as:
$$
    p(x_0) = \int p(x_0|f^c)p(f^c)d f^c, (1)
$$
Note that $f^1,\dots,f^N$ and $x_0$ form a v-structure in PGM, as shown in Fig. A (a). A DPM learns a model $\epsilon_\theta(x_t,t)$ to predict the noise added to a sample $x_t$, then this can be used to derive a score function: $\nabla_{x_t}\log p(x_t) = {1/ \sqrt{\alpha_t - 1}}\cdot\epsilon_\theta(x_t,t)$. Follow [6] to use Bayes Rule, we can write the score function of the conditional distribution as:
$$
    \nabla_{x_t}\log p(x_t|f^c) = \nabla_{x_t}\log p(f^c|x_t) + \nabla_{x_t}\log p(x_t), (2)
$$
By using a score function of $p(x_t|f^c)$, we can sampling data conditioned on the factor $f^c$. 
In addition, Eq. 2 can be extended to one conditioned on $\mathcal{S} \subseteq \mathcal{C}$ by replacing $f^c$ with $S$.
The goal of disentangling a DPM is to model $\nabla_{x_t} \log p(x_t|f^c)$ for each factor $c$. 
Base on Eq. 2, we can learn $\nabla_{x_t}\log p(f^c|x_t)$ for a pre-trained DPM. However, $f^c$ is unknown in an unsupervised setting. Fortunately, we can learn a set of  disentangled representations $\lbrace z^c|   c=1,\dots,N\brace$.
There are two requirements for these representations: $(i)$ They include all information of $x_0$,  named as *completeness*, and $(ii)$ there exists one-one mappings: $z^c \mapsto f^c$ for each factor $c$, named as *disentanglement*. With these representations, we can use $\nabla_{x_t}\log p(z^c|x_t)$ to approximate $\nabla_{x_t}\log p(f^c|x_t)$. 

Our target is to disentangle a pre-trained unconditional DPM on dataset $\mathcal{D}$: $p_\theta(x_{t-1}|x_t) = \mathcal{N} (x_{t-1};\mu_\theta(x_t, t), \sigma_t)$ in an unsupervised manner. 
Specifically, given $x_0\in\mathcal{D}$, for each factor $c\in\mathcal{C}$, we can learn the disentangled representation $z^c$ via an encoder $E_\phi$, and the corresponding disentangled gradient field $\nabla_{x_t} \log p(z^c|x_t)$. Since random variables $z^1,\dots,z^N$ and $x_0$ forms a common cause structure, $z^1,\dots,z^N$ are independent when conditioned on $x_0$. We prove that this is also hold for $x_t$. We thus formulate the gradient field $\nabla_{x_t} \log p(z^S|x_t)$, $z^{\mathcal{S}} = \lbrace z^c|f^c \in \mathcal{S}\rbrace$ conditioned on subset $\mathcal{S} \subseteq \mathcal{C}$ as
$$
     \nabla_{x_t} \log p(z^S|x_t) = \sum_{f^c\in S}\nabla_{x_t} \log p(z^c|x_t). (3)
$$
We follow [6,41], to model the conditional reverse process as a Gaussian distribution $p_\theta(x_{t-1}|x_t, z^{\mathcal{S}})$. Together with Eq. 3, the distribution has a shifted mean given by:
$$
     \mathcal{N} (x_{t-1}; \mu_\theta(x_t, t) + \sigma_t\sum_{f^c\in \mathcal{S}}\nabla_{x_t} \log p(z^c|x_t), \sigma_t). (4)
$$
Directly using $\nabla_{x_t} \log p(z^c|x_t)$ complicates the diffusion model training pipeline [14]. We thus use a decoder $G_\psi(x_t, z^c, t), f^c\in \mathcal{S}$ to estimate the gradient fields $\nabla_{x_t} \log p(z^\mathcal{C}|x_t)$. To achieve this goal, we first use $\sum_{f^c\in \mathcal{C}}G_\psi(x_t,z^c,t)$ to estimate $\nabla_{x_t} \log p(z^\mathcal{S}|x_t)$ based on Eq.3.
Therefore, we adopt the loss inPDAE [41] but using summation of $G_\psi(x_t,z^c,t)$ instead as
$$
L_r  = \mathop{\mathbb{E}}\limits_{x_0,t,\epsilon}\|\epsilon - \epsilon_\theta(x_t,t) 
          + \frac{\sqrt{\alpha_t}\sqrt{1-\bar \alpha_t}}{\beta_t}\sigma_t\sum_{f^c\in \mathcal{C}}G_\psi(x_t,z^c,t)\|. (5)
$$

The above equation can result in that $x_0$ can be reconstructed using all disentangled representations $E_\phi(x_0)$, which is exactly the *completeness* requirement of disentanglement. Secondly, to meet the *disentanglement* requirement, each $G_\psi(x_t,z^c,t)$ should approximate $\nabla_{x_t} \log p(z^c|x_t)$, individually. With a learnt $G_\psi(x_t,z^c,t)$, we can sample $x^c_0$ conditioned on the disentangled representation $z^c$.


We introduce the Disentangling Loss for the *disentanglement*, and to constrain that $G_\psi(x_t,z^c,t)$ is an estimation of $\nabla_{x_t} \log p(z^c|x_t)$. 
Since it is hard to find the sufficient condition for disentanglement in unsupervised setting, we follow prior works that propose necessary conditions worked in practice, e.g., the group constraint in [40], the maximization of mutual information in [22,1 in 6mFa]. 
We propose to minimize the mutual information between $z^c$ and $z^k$, $k \neq c$, as a necessary condition to disentangle a DPM. 

In the following, we transfer the minimizing of mutual information to the constraints in the representation space. 
We denote $\hat x_0$ is sampled from the pre-trained unconditioned DPM using $\epsilon_\theta(x_t,t)$, $\hat x^c_0$ is conditioned on $z^c$ and sampled using $G_\psi(x_t,z^c,t)$. 
Then we can extract the representations from the samples $\hat x_0$ and $\hat x^c_0$ with $E_\phi$ as $\hat z^k = E_\phi^k(\hat x_0)$ and $\hat z^{k|c} = E_\phi^k(\hat x^c_0)$, respectively. 
According to Proposition 1 (Proof sketch in Fig. B), we can minimize an proposed variational upper bound to minimize the mutual information.



**Proposition 1** For a PGM in Fig. A (a), an upper bound of the mutual information between $z^c$ and $z^k$ (where $k \neq c$) is as following: 
$$
    \min \mathop{\mathbb{E}}\limits_{k,c,x_0,\hat{x}_0^c}
    \|\hat{z}^{k|c} - \hat{z}^k\|- \|\hat{z}^{c|c} - \hat{z}^c\| + 
    \|z^c - \hat{z}^{c|c}\| (6)
$$

Meta-review: The authors performed additional experiments supporting their work and significantly extended the framework which should be of interest to the wider NeurIPS community.

From the discussion with the reviewers there are no open or controversial points left. The authors exploit an inductive bias which is unique to diffusion models and thus not available in classical approaches. While the specific use and introduction of the inductive bias could have been considered incremental and maybe straightforward, the large scale empirical evaluation on both synthetic and real world datasets are significant contributions and during the rebuttal the authors provided significantly more evidence for their framework (in particular in discussion with reviewer 6mFA). Among these baselines only a more advanced VAE based method joint with more metrics for disentanglement seem missing. For example the considered baselines like beta-VAE are indeed well established but already quite old. While beta-VAE may struggle with the mentioned generation vs. disentanglement tradeoff due to the additional regularization, more modern VAE based methods will not have this problem. However as pointed out by the authors even modern VAE based approaches would not be able to exploit the proposed inductive bias which is specific for diffusion based models. 

While this is overall a borderline paper, the changes deemed necessary from the reviewers perspective seem feasible without an additional iteration and I would thus recommend acceptance while assuming that the authors will keep their word of implementing the proposed changes.
Decision: Accept (poster)