Diffusion Model with Cross Attention as an Inductive Bias for Disentanglement

Authors: tao  yang
cuiling  lan
yan  lu
nanning  zheng
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-08 19:33:15

Abstract: Disentangled representation learning strives to extract the intrinsic factors within observed data. Factorizing these representations in an unsupervised manner is notably challenging and usually requires tailored loss functions or specific structural designs. In this paper, we introduce a new perspective and framework, demonstrating that diffusion models with cross-attention can serve as a powerful inductive bias to facilitate the learning of disentangled representations. We propose to encode an image to a set of concept tokens and treat them as the condition of the latent diffusion for image reconstruction, where cross-attention over the concept tokens is used to bridge the interaction between the encoder and diffusion. Without any additional regularization, this framework achieves superior disentanglement performance on the benchmark datasets, surpassing all previous methods with intricate designs. We have conducted comprehensive ablation studies and visualization analysis, shedding light on the functioning of this model. We anticipate that our findings will inspire more investigation on exploring diffusion for disentangled representation learning towards more sophisticated data analysis and understanding.

Main paper: 
 introduction: 
 Disentangled representation learning strive to uncover and understand the underlying causal factors of observed data [1, 13] . This is believed to possess immense potential to enhance a multitude of machine learning tasks, facilitating machines to attain better interpretability, superior generalizability, controlled generation, and robustness [33] . Over the years, the field of disentangled representation learning has attracted significant academic interest and many research contributions. Numerous methods, encompassing Variational Autoencoders (VAE) based techniques (such as β-VAE [14, 2] , FactorVAE [19] ), Generative Adversarial Networks (GAN) based approaches (such as InfoGAN [5] , InfoGAN-CR [23] ), along with others [37, 28] , have have been proposed to advance this field further.Figure 1 : Average attention map across all time steps in stable diffusion. We draw inspiration from the process of text-to-image generation using a diffusion model with cross-attention. Utilizing the highly 'disentangled' words as the condition for image generation, the cross-attention maps observed from the diffusion model exhibit a strong text semantic and spatial alignment, indicating the model is capable of incorporating each individual word into the generation process for a final semantic aligned generation. This leads us to question whether such a diffusion structure could be inductive to disentangled representation learning.Originally, Variational Autoencoders (VAEs) are conceived as deep generative probabilistic models, primarily focusing on image generation tasks [20] . The core idea behind VAEs is to model data distributions from the perspective of maximizing likelihood using variational inference. Subsequent research has revealed that VAEs possess the potential to learn disentangled representations with appropriate regularizations on simple datasets. To enhance disentanglement, a range of regularization losses have been proposed and integrated within the VAE framework [14, 19, 21] . Similarly, GANs have incorporated regularizations to enable the learning of disentangled features [5, 23, 44] . Despite significant progress, the disentanglement capabilities of these models remain less than satisfactory, and the disentangled representation learning is still very challenging. Locatello et al. demonstrate that relying solely on regularizations is insufficient for achieving disentanglement [24] . They emphasize the necessity of inductive biases on both the models and the data for effective disentanglement. A fresh perspective is eagerly anticipated to shed light on this field.Recently, diffusion models have surfaced as compelling generative models known for their high sample quality [36] . Drawing inspiration from the evolution of VAE-based disentanglement methods, we are intrigued by the question of whether diffusion models, also fundamentally designed as deep generative probabilistic models, possess the potential to learn disentangled representations. Obtaining a compact and disentangled representation for a given image from diffusion models is non-trivial. Diffusion Autoencoder (Diff-AE) [26] and PDAE [42] move a step forward towards using diffusion models for representation learning by encoding the image into a feature vector, incorporating this into the diffusion generation process. However, these representations have not exhibited disentanglement characteristics. What inductive biases are essential for the learning of disentangled representations? Could we have a diffusion-based framework possessing such inductive biases? Notably, in text-to-image generation, a conditional diffusion model integrates the 'disentangled' text tokens through cross attention, demonstrating the ability to generate semantically aligned images [30, 36, 12] . Interestingly, the observed cross-attention map reveals that different words have their corresponding spatial regions of high affinities, exhibiting strong semantic and spatial alignment as illustrated in Figure 1 . These disentangled representations of 'word's could potentially contribute to a more streamlined generation process. Inspired by this, we wonder whether such diffusion structure with cross attention can act as an inductive bias to facilitate the disentangled representation learning.In this paper, we endeavour to investigate this question and explore the potential of diffusion models in disentangled representation learning. We discover that the diffusion model with crossattention can serve as a strong inductive bias to drive disentangled representation learning, even without any additional regularization. As illustrated in Figure 2 (a), we employ an encoder to transform an image into a set of concept tokens, which we treat as 'word' tokens, acting as the conditional input to the latent diffusion model with cross attention. Here, cross-attention bridges the interaction between the diffusion network and the image encoder. We refer to this scheme as EncDiff. EncDiff is powered by two valuable inductive biases, i.e., information bottleneck in diffusion, and cross attention for fostering 'word' (concept token) and spatial alignment, contributing to the disentanglement. Experimental results on benchmark datasets demonstrate that EncDiff achieves excellent disentanglement performance, surpassing all the previous methods with elaborate designs. Comprehensive ablation studies show that the strong disentanglement capability is mainly attributed to 1) the diffusion modelling and 2) the cross-attention interaction. Visualization analysis provides insights into the effectiveness of the disentangled representations. The KL divergence between the data distribution q(x t-1 |x t , x 0 ) and the Gaussian prior distribution N (0, I) under four different variance (β) schedules: cosine, linear, sqrt linear and sqrt. The results have been normalized by the number of dimensions.EQUATIONWe have four main contributions.• We uncover that the diffusion model with cross-attention can serve as a strong inductive bias for enabling disentangled representation learning. • We introduce a simple yet effective framework, EncDiff, powered by a latent diffusion model with cross attention and an ordinary image encoder, for disentangled representation learning. • This framework inherently incorporates two valuable inductive biases: the information bottleneck in diffusion, and cross attention, fostering concept token and spatial alignment. We analyze that the diffusion process inherently possesses the time-varying information bottlenecks. • Without additional regularization or specific designs, our framework achieves state-of-the-art disentanglement performance, even outperforming the latest methods with more complex designs.We anticipate the new perspective will illuminate the field of disentanglement and inspire deeper investigations, paving the way for future sophisticated data analysis, understanding, and generation. 
related work: 
 Disentangled Representation Learning Disentangled Representation Learning endeavours to train a model proficient in disentangling the underlying factors of observed data [1, 13, 33] . A plethora of methods have been proposed to augment the generative models of VAEs and GANs, endowing them with disentanglement capability. These methods primarily rely on probability-based regularizations applied to the latent space. To improve disentanglement, most approaches focus on how to regularize the original VAE. For instance, this includes weighting the evidence lower bound (ELBO) as in β-VAE [14] , or introducing different terms to the ELBO, such as mutual information (InfoVAE [43] , InfoMax-VAE [29] ), total correlation (Factor-VAE [19] ), and covariance (DIP-VAE [21] ). Locatello et al. [24] demonstrate that relying solely on these regularizations is insufficient for achieving disentanglement. DIP-VAE [21] introduces a regularizer on the expectation of the approximate posterior over observed data, by matching the moments of the distributions of latents. In this work, we analyze and identify two valuable inductive biases of diffusion that promote the disentangled representation learning in our framework. Inductive biases on both the models and the data are necessary. In this paper, we investigate the disentanglement capability of diffusion models and demonstrate that diffusion models with cross-attention can serve as a powerful inductive bias for disentanglement.Diffusion Models Diffusion models have emerged as a powerful new family of deep probabilistic generative models [36] , surpassing VAEs and GANs for image generation and many other tasks.Diffusion models progressively perturb data by injecting noise and then learn to reverse this process for the generation. A question arises as to whether diffusion models can effectively serve as disentangled representation learners. It is challenging to obtain a compact yet disentangled representation of an image from a diffusion model. Diff-AE [26] and PDAE [42] investigate the possibility of using DPMs for representation learning, whereby an input image is autoencoded into a latent vector. However, these representations do not manifest disentangled characteristics. SlotDiffusion [35] and LSD [17] integrate diffusion models into object-centric learning, where diffusion acts as a improved slot-toimage decoder, and slot attention is still the key for promoting object-centric learning. Moreover, they aim to learn object-wise representation but still cannot disentangle the factors/attributes of an object/a scene. Limited research has explored disentangled representation learning by leveraging diffusion models. InfoDiffusion [34] encourages the disentanglement of the latent feature of Diff-AE [26] by introducing mutual information and prior regularization, similar to InfoVAE [43] . CL-Dis [18] introduces a VAE to guide diffusion model to learn disentangled representation. DisDiff [40] employs a pre-trained diffusion model for disentangled feature learning. DisDiff adopts an encoder to learn the disentangled representations and a decoder to learn the sub-gradient field for each disentangled factor. It requires multiple decoders to predict these sub-gradient fields for all the factors and complex disentanglement losses, resulting in a costly and intricate process. The recent work SODA [16] leverages a diffusion model to generate novel view of image for representation learning, revealing the capability of capturing visual semantics to diffusion model. Is it necessary to impose these complicated regularizations upon diffusion models as [40] ? Does a strong inductive bias facilitating disentanglement already exist within diffusion models? In this paper, we endeavour to answer these questions and illustrate that a simple framework driven by a diffusion model without any additional regularization is capable of achieving superior disentanglement performance.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: Please see Section 5 in the main paper.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
method: 
 We aim to investigate the potential of diffusion models in disentangled representation learning. We propose a simple yet effective framework, EncDiff, that exhibits strong disentanglement capabilities, even without additional regularizations. We analyze and identify two valuable inductive biases, i.e., information bottleneck in diffusion, and the cross attention for fostering 'word' (concept token) and spatial alignment, thereby promoting disentanglement. We elaborate on the framework design in subsection 3.1 and the analysis in subsection 3.2, respectively.Figure 2 (a) illustrates the flowchart. It consists of an image encoder that transforms an input image into a set of concept tokens and a diffusion model that serves as the decoder to reconstruct the image. Cross-attention is employed as the bridge for the diffusion network and the image encoder.Image Encoder For a given input image I, the image encoder τ ϕ aims to provide a set of concept tokens S = {s 1 , • • • , s N }, which act similarly to the word embeddings in the prompts for text-toimage generation in the latent diffusion models (LDMs) [30] . Without loss of generality, we use an ordinary CNN network as the image encoder to obtain concept tokens. Following the design of the encoder in VAE [20] , we use a fully connected layer to transform the feature map into a feature vector. We treat each dimension of the encoded feature vector as a disentangled factor and map each factor to a vector (i.e., concept token) by non-shared MLP layers (as illustrated by Figure 3 ).We follow LDM [30] to construct our diffusion model in the latent space, which demonstrates superior generation ability. LDM is one of the most popular diffusion models, proposing to conduct diffusion denoising in the latent space. To condition the concept tokens during image generation, cross-attention is used to map these tokens into the intermediate representations of the U-Net in the diffusion model. This is accomplished by using the cross-attention defined as Attention(Q, K, V ) = softmax QK T √ d•V , where the spatial feature in the intermediate feature map in diffusion model serves as a query, the concept tokens act as keys and values.We conduct an end-to-end training of the encoder and the diffusion model, utilizing the optimization objective of reconstructing noise, which is a methodology aligned with that employed in LDM [30] .We analyze that there are two crucial inductive biases in diffusion models: the information bottleneck in diffusion, and the cross-attention interaction. We analyze that the diffusion process inherently possesses the time-varying information bottlenecks.β-VAE [14] and AnnealVAE [2] utilize the original Kullback-Leibler (KL) divergence in VAEs to enhance the disentanglement capability, where the KL divergence constraint plays a role of an information bottleneck. Here, we analyze the presence of an information bottleneck mechanism that promotes the disentanglement in diffusion models. Without loss of generality, our analysis focuses on the diffusion model in image latent space [30] . The analysis also holds in pixel space.Within the framework of the latent diffusion model, we add Gaussian noise to an image latent x 0 over T steps according to a variance schedule β 1 , • • • , β T . This process yields a sequence of noisy samplesEQUATIONLet α t = 1 -β t and ᾱt = t i=1 α t . x t can be obtained using the following equation [15, 36] :EQUATIONwhere the noise ϵ is sampled from a Gaussian distribution N (0, I).The diffusion model optimizes a network (e.g., U-Net) ϵ θ to predict the noise from the noisy input x t and the conditioning input S (concept tokens), with the loss function defined asEQUATIONHere, we omit the weighting terms of loss function for simplicity. The latent x 0 can be reconstructed based on the predicted noises.Let's analyze the inherent information bottleneck at each time step t in the reverse diffusion process.In the reverse diffusion process, the reverse conditional distribution is tractable as [15] :q(x t-1 |x t , x 0 ) = N (x t-1 |μ t , βt I),where μt = √ ᾱt-1EQUATIONWe formulate the Kullback-Leibler (KL) divergence C t between q(x t-1 |x t , x 0 ) and the Gaussian prior distribution p(x t-1 ) = N (0, I) at step t -1 as:EQUATIONwhere n denotes the number of dimension of signal x (i.e., x 0 , x t , x t-1 ).In Figure 2 (b), we present a plot illustrating the KL divergence C t under different variance (β) schedules, including linear, sqrt linear, cosine [25, 15] , and sqrt schedules [6, 30] . We can see that as the time step t decreases, the KL divergence C t increases, indicating the information carried by x t-1 increases and leading to an increasingly looser information bottleneck over x t-1 . According to [2, 14] , such a time-varying information bottleneck may play an important role in promoting disentanglement. Different variance (β) schedules results in different KL divergence curves. We found that these different variance schedules lead to different disentanglement performance (see Subsection 4.4).Actually, optimizing the loss of conditional diffusion model as in (3) is equivalent to push the reverse conditional distribution p θ (x t-1 |x t , S) to approach q(x t-1 |x t , x 0 ) (see the explanation in Appendix B). By this means, the information bottlenecks Ct = D KL (p θ (x t-1 |x t , S)||p(x t-1 )) over x t-1 tend to approach C t = D KL (q(x t-1 |x t , x 0 )||p(x t-1 )) in training for all the time steps. According to theorem in Appendix C, the information bottleneck over latent x t-1 is transferred to the condition S (i.e., concept tokens). Intuitively, this is because x t-1 is controlled by the concept tokens and the network parameters θ, as indicated by p θ (x t-1 |x t , S). The concept token representations S are learnable, and the information bottleneck is transferred to and imposed on S. With time-varying information bottlenecks, the diffusion process encourages a range of different information capacities on S during the diffusion process, promoting the disentanglement of concept tokens.The information bottleneck described above shares a certain similarity with the optimization objective in AnnealVAE [2] . The minimizing objective of AnnealVAE is expressed as follows:EQUATIONwhere ϕ, φ are the parameters of the encoder and decoder; the latent representation and data are denoted as z, x, respectively. C is a handcrafted constant used to control the information bottleneck of the latent space. Different factors to be disentangled may contain different amounts of information.Instead of using a constant during training, AnnealVAE dynamically allocates larger amount of information (larger C) to the latent units as the training iteration increases. So that different factors can be learned at various training stages.In diffusion models, where the KL divergence characterizes the information amount, we observe that the information amount varies in reverse diffusion steps, see Figure 2 (b).The information bottleneck sheds a light on disentangling, acting as an inductive bias for diffusion. However, using the information bottleneck still only has theoretical feasibility. Its effectiveness also relies on the structure design of diffusion model. We believe that the cross-attention design in conditional diffusion model is crucial for disentanglement, serving as another effective inductive bias.Our objective is to train an encoder that obtains a set of concept tokens taking an image as input under the guidance of the diffusion model. We take the output of the encoder as the condition of the U-Net of diffusion model for image generation. We incorporate the concept tokens into diffusion model through cross-attention. Intuitively, a spatial position of an image is related to several concepts, e.g., object color and shape in Shapes3D. Each spatial feature is composed of several related concept-based representations. Interestingly, cross-attention in the U-Net play a similar role, where each spatial feature servers as the query, and the learned concept tokens are used as the keys and values to refine the query. In subsection 4.4, we validate the necessity of the two inductive biases leading to disentanglement.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Justification: We provide the information of computer resources used in Appendix D.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
experimental setup: 
 Implementation Details The trainable parts are the encoder and diffusion model. We employ the popular diffusion structure of latent diffusion [30] by default. Without loss of generality, following [30] , we use the VQ-reg to avoid arbitrarily high-variance latent spaces and sample images in 200 steps. We adopt the cosine as the variance (β) schedule in the diffusion model by default. By default, we use a CNN encoder for the image encoder to obtain a set of disentangled concept tokens. We use a CNN encoder similar to that used in [40] . We denote our scheme as EncDiff.Training Details During the training phase of EncDiff, we maintain a consistent batch size of 64 across all datasets. The learning rate is consistently set to 1 × 10 -4 . We adopt the standard practice of employing an Exponential Moving Average (EMA) with a decay factor of 0.9999 for all model parameters. The training hyper-parameters follows DisDiff [40] and DisCo [28] . For each concept token, we follow DisDiff [40] to use a 32 dimensional representation vector. We train EncDiff on a single Tesla V100 16G GPU. A model takes about 1 day for training.Datasets To evaluate the disentanglement performance, we utilize the commonly used benchmark datasets: Shapes3D [19] , MPI3D [10] and Cars3D [27] . Shapes3D [19] consists of a collection of 3D shapes. MPI3D is a dataset of 3D objects created in a controlled setting. Cars3D is a dataset consisting of 3D-rendered cars. For real-world data, we conduct our experiments using CelebA, a dataset of celebrity faces with attributes. Our experiments are carried out at a 64×64 image resolution, consistent with previous studies [19, 4, 28, 40] . Baselines & Metrics We compare the performance of our method with VAE-based, GAN-based, and diffusion-based methods, following the experimental protocol as in DisCo [28] . The VAE-based models we use for comparison are FactorVAE [19] and β-TCVAE [4] , while the GAN-based baselines include InfoGAN-CR [23] , GANspace (GS) [11] , LatentDiscovery (LD) [32] and DisCo [28] . Each method utilizes scalar-valued representations. DisDiff [40] uses vector-valued representations.EncDiff has two kinds of representations simultaneously. We focus on the scalar-valued in the main paper. We follow DisDiff to set N to 20. For these vector-valued representations, we follow [7, 40, 38] to perform PCA as a post-processing on the representation before evaluation. To assess the potential variability in performance due to random seed selection, we have fifteen runs for each method for reliable evaluation, reporting the mean and variance. Regarding evaluation metrics, we adopt two representative metrics, the FactorVAE score [19] and the DCI [8] .We compare the disentanglement ability of our EncDiff with the state-of-the-art methods. Table 1 shows quantitative comparison results of disentanglement under different metrics. We can see that EncDiff achieves the best performance on all the datasets except Cars3D, showcasing the model's superior disentanglement ability. EncDiff achieves superior performance by leveraging the strong inductive bias from the diffusion model with cross-attention without using any additional regularization losses. EncDiff also outperforms InfoDiffusion [34] and DisDiff [40] by a significant marginal, even though DisDiff uses complex disentanglement loss and inference the decoder multiple times for prediction sub-gradient fields. On the Cars3D dataset, the quantitative evaluation is not so reliable because some factors, such as color and shape, are not included in the labels. From the visualization in Figure 6 of Appendix E, we can see that EncDiff achieves superior disentanglement compared to DisDiff despite the lower FactorVAE score. Moreover, we have conducted experiments on real-world dataset CelebA. Table 2 shows the comparisons of disentanglement performance and generation quality in terms of TAD and FID metrics (mean ± std). EncDiff achieves the state-of-the-art performance on both aspects compared to all baselines.In addition, our EncDiff achieves superior reconstruction quality (see Appendix G for more details).Visualization Analysis on the Disentanglement We qualitatively examine the disentanglement properties of our proposed method. We interchange the concept tokens (factors) of the learned representation of two distinct images and observe the generated images conditioned on these exchanged representations. For illustration purposes, we focus on the widely used Shapes3D dataset from the disentanglement literature and show the results in Figure 4 . We can see that our EncDiff successfully isolates factors. Notably, in comparison to VAE-based methods, EncDiff delivers superior image quality quantitatively (please refer to Appendix H).Visualization of Learned Cross-Attention Maps As mentioned in Section 1, our model draws inspiration from the alignment between 'word' tokens (disentangled representations) and spatial features. The alignment is demonstrated by the learned cross-attention maps. We verify whether our learned concept tokens present the disentangled characteristics by visualizing the alignment of concept tokens with spatial positions through cross-attention maps. As depicted in Figure 5 , the results of our model exhibit exemplary alignment between concept tokens and spatial positions. Distinct concept tokens are associated with varying attended spatial regions, corresponding to different semantics that are comprehensible by humans, such as the region of "Floor" and "Color" for the images from Shapes3D. In our framework, in order to analyze and understand the key factors contributing to the advancement of disentangled representation learning, we conduct ablation studies covering three aspects in the design: 1) whether to use diffusion as the decoder; 2) whether to use cross-attention as the bridge for interaction;3) The influence of different variance (β) schedules. We conduct these ablation studies on the Shapes3D dataset. Please see Appendix H for more ablation studies.To validate whether the use of a diffusion model as an inductive bias for disentangled representation learning is crucial, we employ a network structure similar to the U-Net in our used diffusion model as the decoder, utilizing reconstruction l 2 loss to optimize the entire network. We designed a variant (EncDec w/o Diff) of EncDiff to have an autoencoder-like structure, by reusing the image encoder as encoder and the lower half of the U-Net structure as decoder for reconstruction. In contrast to EncDiff, we discard the multiple step diffusion process and only run once feedforward inference to get the reconstruction. If the autoencoder's performance drops significantly, this will provide evidence about the importance of the diffusion process instead of the U-Net architecture. Specifically, we remove the encoder part of the U-Net and the skip connection between it and the decoder part of the U-Net. We then feed the U-Net decoder with a randomly initialized learnable spatial tensor to maintain the structure of the decoder U-Net. Similarly to EncDiff, the encoded disentangled representations are input to the decoder through cross-attention (CA). We refer to this scheme as EncDec w/o Diff. Table 3 shows the results. The performance of EncDiff with diffusion significantly outperforms EncDec w/o Diff by 0.46 and 0.79 in terms of FactorVAE score and DCI, respectively. This indicates that inductive bias from diffusion modelling is crucial for achieving effective disentanglement.To incorporate the image representation to the diffusion model as a condition, we use cross-attention by treating each disentangled representation as a conditional token (similar to the use of 'word' token in text-to-image generation in stable diffusion model [30] ). As an alternative, similar to that in Diff-AE [26] and InfoDiffusion [34] , we use adaptive group normalization (AdaGN) to incorporate the representation vector (by concatenating the concept tokens) to modulate the spatial features. We name this scheme as EncDiff w/ AdaGN. Table 3 presents the results. We can see that EncDiff w/ AdaGN is inferior to EncDiff, with a significant decrease of 0.33 in terms of DCI. Cross-attention facilitates the alignment of each concept token with the corresponding spatial features, akin to the alignment of the 'word' token to spatial features in the text-to-image generation. In contrast, AdaGN did not efficiently promote disentanglement.We investigate the influence of the different variance (β) schedules, including including linear, sqrt linear, cosine [25, 15] , sqrt schedules [6, 30] , on the disentanglement performance. From Table 4 , we can see that distinct schedules result in different performance, demonstrating the influence on disentanglement of different information bottleneck schedules. Note that the FactorVAE scores are all very high and cannot well reflect the performance. We prefer to use DCI metric here for evaluation. We can see that the cosine schedule performs the best and we adopt it by default. The linear schedule approaches that of the sqrt linear in terms of the curve shape, please see Figure 2 (b) and achieves the similar performance in terms of DCI.Scalar-valued vs. Vector-valued Manners We treat each dimension of the encoded feature vector as a disentangled factor, followed by a mapping to concept token (vector) for each factor. Another design alternative is to directly split the feature vector into N chunks, with each chunk being a concept token, similar to DisDiff [40] . We name this vector-valued design and refer to it by DisDiff-V. may serve a bottleneck role and contribute to the disentanglement. We think that the vector-based representation potentially extracts more information and hence enforces a looser bottleneck than scalar-valued representation. Note that the more information encoded, there is a higher probability that the encoded information is correlated, which is contradictory for disentanglement. Therefore, the tighter bottleneck from scalar-valued representation leads to (a slightly) better performance.Results on Pixel Space As stated in Section 3.2.1, our analysis is applicable in pixel space as well.In order to verify this, we trained EncDiff directly in pixel space on the Shapes3D dataset, which we denote the scheme as EncDiff pixel. The results of the disentanglement analysis are presented in Table 6 . The performance of our framework in pixel space remains robust, indicating that the operation in latent space is not a critical factor for achieving effective disentanglement.We compare the computational complexity of Diff-AE [26] , DisDiff [40] , and our EncDiff in terms of the parameters (Params.), floating-point operations (FLOPs), and inference time (seconds/sample) for sampling an image. As shown in Table 7 , our EncDiff demonstrates much higher computational efficiency than Diff-AE and DisDiff. 
limitations: 
 Our method operates in a fully unsupervised manner and exhibits strong disentanglement capability on simple datasets. Similar to other disentanglement-based methods [14, 19, 5, 23, 37] , obtaining satisfactory performance on complex data remains a challenge. As a diffusion-based method, the generation speed of EncDiff is faster than DisDiff [40] . However, it is still slower compared to VAE-based and GAN-based methods. More effective sampling strategies, as employed in DPM-based methods, could be utilized for accelerating in the future. 
conclusion: 
 This paper unveils a fresh viewpoint, demonstrating that diffusion models with cross-attention can serve as a strong inductive bias to foster disentangled representation learning. Within our framework EncDiff, we reveal that the diffusion model structure with cross-attention can drive an image encoder to learn superior disentangled representations, even without any regularization. Our comprehensive ablation studies demonstrate that the strong capability is mainly attributed to diffusion modelling and cross-attention interaction. This work will inspire further investigations on diffusion for disentanglement, paving the way for sophisticated data analysis, understanding, and generation. 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes]Justification: Please see appendix A.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [Yes]Justification: When we release our models, we ask the users to follow the user guidelines to assure the positive usage.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: We have cited the original papers of the datasets. We have checked the licenses of used datasets and followed them.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. 
 
 Appendix: This paper presents work whose goal is to advance the field of disentanglement learning. Our research is designed to be a positive force for innovation purpose. Viewed from a societal lens, the potential negative impacts is the malicious use of the models. This highlights the critical necessity of incorporating ethical considerations in the utilization of our method for responsible AI.B Optimization of the Reverse Conditional Probability p θ (x t-1 |x t , S)Two distinct approaches exist for decomposing the loss function, Variational Lower Bound(VLB), of the diffusion model. The derivations presented herein closely follow the methodology outlined by Ho et al. [15] . Without loss of generality, we could incorporate the conditional input denoted by S to the diffusion.For the first decomposition alternative, we can derive the following equations:EQUATIONBy applying Bayes' Rule and the Markov property of the diffusion process, we conclude thatEQUATIONThis leads toEQUATIONFor the second alternative, the loss function is derived to be E x0,ϵ,t ∥ϵ θ (x t , t, S) -ϵ∥ [15] .The optimization of conditional diffusion model by predicting the noises through minimizing E x0,ϵ,t ∥ϵ θ (x t , t, S) -ϵ∥ is thus equivalent to minimizing the second term of (9), i.e., Σ t>1 D KL (q(x t-1 |x t , x 0 )||p θ (x t-1 |x t , S)), which pushes the reverse conditional probability p θ (x t-1 |x t , S) to approach q(x t-1 |x t , x 0 ).Theorem C.1. The Kullback-Leibler divergence is invariant under a differentiable mapping f , i.e .where x = f (S) is a differentiable function between x and S and p(x) and q(x) are the probability density functions of the probability distributions P and Q, respectively.Proof : The Kullback-Leibler divergence (KL divergence) is defined asD KL (p(x)|q(x)) = p(x) log p(x) q(x) dxAccording to the change of variable theorem, we haveEQUATIONwhere |f ′ (x)| denotes the Jacobian of f (x), p(S) is the corresponding distribution of p(x) under the mapping f . q(S) is the corresponding distribution of q(x) under the mapping f . Combine the two equations, we haveD KL (p(x)|q(x)) = p(f (x))|f ′ (x)| log p(f (x))|f ′ (x)| q(f (x))|f ′ (x)| dx = p(f (x)) log p(f (x)) q(f (x)) |f ′ (x)|dx.(11) Considering the property of integral that for S = f (x) we have dS = |f ′ (x)|dx. We then have the following:EQUATIONThis indicates that the information bottleneck on x is transferable to the input S of the function.For our EncDiff (scalar-valued), following the approach of [22] , each dimension of the encoded feature vector is treated as a disentangled factor. Each factor is then mapped to a vector (i.e., concept token) using non-shared MLP layers, as illustrated in Figure 3 . For vector-valued EncDiff (EncDiff-V), inspired by DisDiff [40] , we partition the feature vector into N (e.g., 20) chunks, referred to as concept tokens, to encode different factors.Image Encoder Architecture To ensure an equitable comparison, we employ the encoder architecture, consistent with DisCo [28] and DisDiff [40] . The encoder specifications are detailed in Table 8 . Diffusion U-Net Architecture The diffusion architecture adheres to the design principles of latent diffusion [30] and DisDiff [42] . Table 9 provides a detailed overview of the network structure, similar to the structure of Latent Diffusion Probabilistic Model.We pretrain VQ-VAE in diffusion. Then the diffusion network and our image encoder are jointly trained.We present qualitative results for Cars3D in Figure 6 . Notably, on the synthetic dataset Cars3D, EncDiff demonstrates the acquisition of disentangled representations, presenting better disentanglement ability than DisDiff [40] ). We can see that our EncDiff can capture the factors of "Color", "Azimuth", and "Shape" while DisDiff failed to capturing them. Furthermore, we include three rows of images demonstrating the manipulation of representations lacking informative content, denoted as "None".We present the qualitative outcomes for MPI3D in Figure 7 . Remarkably, on the challenging disentanglement dataset MPI3D, EncDiff showcases its ability to obtain disentangled representations.We showcase the visualization of attention maps for our model's disentangled representations on the Cars3D dataset, as depicted in Figure 8 . Similar to EncDiff, the attention maps provide insights into the acquisition of disentangled representations in the synthetic setting of Cars3D. Notably, our attention maps reveal distinct alignments between concept tokens and spatial features, demonstrating the disentangled characteristics learned by our model.For the MPI3D dataset, Figure 8 demonstrate the qualitative outcomes of attention map visualization.In this challenging disentanglement scenario, EncDiff excels in acquiring disentangled representations. The attention maps further illustrate the alignment between concept tokens and spatial positions, affirming the model's ability to disentangle complex features in MPI3D. Autoencoding Reconstruction Quality To investigate the autoencoding reconstruction quality of EncDiff, we conduct the same quantitative experiments with PDAE, DisDiff and Diff-AE. We follow them to evaluate the reconstruction quality using averaged SSIM, LPIPS, and MSE. The results are shown in Table 10 . It is evident that EncDiff outperforms DisDiff on all metrics. EncDiff achieves the state-of-the-art performance of SSIM and LPIPS with a strong disentanglement capability. Measuring on Scalar-valued vs. Vector-valued Representations in our EncDiff For EncDiff, each dimension of the scalar-valued representation is mapped to a representation vector, resulting in two representations in EncDiff: a scalar-valued one and a mapped vector-valued one. From Table 11 , we can see that these two representations have similar performance. Influence of the Number of Concept Tokens Similar to other disentanglement methods, the number of disentangled latent units influences performance. To study such an effect on EncDiff, we train our model with the following number of tokens: 5, 10, 15, 20 and 30, respectively. As shown in Table 12 , when the number of tokens is fewer than the number of ground truth factors, the performance significantly drops. With the use of more tokens, the performance improves. In accordance with the setup proposed by [39] , EncDiff adopts the default setting of 20 tokens for a fair comparison with other methods [39, 40, 28] . We are wondering whether additional regularization can further promote the disentanglement. To validate this, we conducted the following experiments to investigate the effectiveness of incorporating two types of constraints: sparsity and orthogonality, respectively.We explored the orthogonality constraint proposed in [3] , which enforces orthogonality from a group theory perspective. We adapted their method of transforming representations using Euler encoding to enforce orthogonality within EncDiff. We adopted the official implementation on github to modify EncDiff, denoted as EncDiff with [3] . We integrated the sparsity regularization terms proposed in [9] into EncDiff to facilitate disentanglement, denoted as EncDiff with [9] . We take the techniques proposed in [41] , which involve matrix decomposition and matrix exponentiation to construct orthogonal matrices. We replaced the scalar MLP mappings with a series of learnable orthogonal vectors to ensure orthogonality in the representations, denoted as EncDiff with [41] .The results are shown in Table 13 . We observed that the regularization can slightly improve the performance further on our EncDiff. For simplicity, we do not incorporate any regularization on all other results.Table 13 : Ablation study on the additional regularization over EncDiff. We use a CNN encoder by default.EncDiff w/[3] 0.999 ± 0.000 0.965 ± 0.040 EncDiff w/ [41] 0.999 ± 0.002 0.972 ± 0.029 EncDiff w/ [9] 0.999 ± 0.001 0.975 ± 0.031 EncDiff 0.999 ± 0.001 0.969 ± 0.030Influence of the Image Encoder Architecture To investigate the influence of encoder design, we use a powerful encoder to replace the CNN encoder used in the EncDiff. We adopt a transformer encoder with a set of learnable tokens as the disentangled representations, as introduced in [39] , which are refined through cross-attention. We refer to this scheme EncDiff w/Trans. For fair comparison, the model size of the encoders is similar. As shown in Table 14 , EncDiff w/Trans is comparable to EncDiff. When considering the performance gap between DisDiff [40] and our EncDiff, the influence of encoder structures is small and is not the key factor for influencing disentangling capabilities. A similar phenomenon is observed in vector-valued one. We also perform ablation study related with EncDiff-V to validate the effects of the two inductive bias, the inefficacy of additional regularization. Similar trends as EncDiff are observed.Ablation on Two Inductive Bias of EncDiff-V In alignment with the main paper, we also conducted an experiment to assess the effectiveness of the two inductive biases in EncDiff-V. We adopt the same decoder used in Section 4.4 to substitute the diffusion in EncDiff-V. We denote this model as EncDec-V w/o Diff. On the other hand, to study the effectiveness of cross-attention, we use the same conditional decoder of EncDiff w/ AdaGN in EncDiff-V, denoted as EncDiff-V w/ AdaGN. The performance of both of these two models drops significantly, as indicated by the results in Table 15 .Table 15 : Ablation study on the influence of the two different inductive biases of EncDiff-V. For EncDec-V w/o Diff, we replace the diffusion model with a decoder while cross attention is preserved for the interaction. For EncDiff-V w/ AdaGN, we replace the cross attention with AdaGN.EncDec-V w/o Diff 0.682 ± 0.092 0.246 ± 0.073 EncDiff-V w/ AdaGN 0.956 ± 0.039 0.520 ± 0.083EncDiff-V 0.999 ± 0.000 0.900 ± 0.045Besides Shapes3d, we conducted the main ablation study on another dataset MPI3D. The results are shown in Table 16 and 17 below. We can observe that the trends are consistent with that on Shapes3D (see Table 3 ). The computational complexity of the VAE-based and GAN-based methods are also listed as shown in Table . The computational complexity and inference time VAEs and GANs still have strength, but diffusion has much better generation and disentangling ability. Among the diffusion models, our EncDiff has better generation and disentangling ability but less computational cost and inference time. Q K V Q K VA photo of Dog In order to discuss the disentanglement of EncDiff on real-world data, a newly designed architecture is introduced in Figure 9 , named EncDiff(SD).In this architecture, we use a strong pretrained model (stable diffusion v1.4) to replace the original diffusion in EncDiff. Motivated by DreamBooth [31] for customized representation learning, we assign several latent units for each object for inverting and learning object semantic representation. As shown in Figure 9 , in order to disentangle these latent units, we use a set of non-shared MLPs to map these latent units into concept tokens. Different from EncDiff, the instance for disentangling is not an image but the semantics of objects. The target is to disentangle concepts or properties (e.g., color, long-hair, big-eared) from the inverted objects (dog).Diffusion LossFigure 10 shows sampling process of EncDiff (SD), we denote the image on the left as image1 and image on the right as image2. We use a prompt "A <token from image1, token from image2> dog is playing a blue ball" to sample image. The <token from image1> encodes the color of the dog. The <token from image2> encodes the type of the dog. We can sample images of new objects that combine color of dog1 and type of dog2. We take the new prompt as the condition of Stable Diffusion, so that we can use diffusion as an inductive bias for disentangling. The process above is the same as EncDiff. Firstly, to disentangle the objects from the background, we replace the Diffusion Loss with DisenBooth Loss. Secondly, to learn multiple concepts, we follow Custom Diffusion to finetune only the KV layer in cross attention. Lastly, for the ease of learning disentangled concept with a few images, we take an image from each object to construct a training batch. The results are shown in Figure 11 , 12, 13. Our model also demonstrates the ability to disentangled semantic factors on real and complex data. Justification: We introduce a new perspective and framework, demonstrating that diffusion models with cross-attention can serve as a powerful inductive bias to facilitate the learning of disentangled representations.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.Justification: We have read the NeurIPS Code of Ethics and confirm that we follow that.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 


Decision: The paper presents a novel approach to learning disentangled representations. There is a consensus in the reviews that this paper brings a valuable contribution to the field. Some concerns where raised regarding the emprical evaluation, specifically regarding the number of datasets and the selection of models used for comparison. This is a common concern with diesntanglement papers, particularly when they are contrasted to supervised learning papers. In this context, however, developing empirical evaluation with the similar scale (data sets) is significantly more challenging.  Overall, the authors present a solid and broad evaluation and have constructively adopted suggestions from the reviews.

Reviews: review: paper_summary: paper_summary: The paper proposes a novel method that utilizes a concept-extracting image encoder and the cross-attention mechanism in conditional diffusion models for achieving the learning of disentangled representations. Comprehensive experiments, visualizations and ablation studies confirm the effectiveness of the proposed method.
strength_weakness: strength_weakness: I don't see any major flaws in the paper.

*Typos and writing related*. 
*  Line 174, it would be better to briefly describe why this information bottleneck promotes disentanglement instead of just cite other papers.
* Line 282, "utilizing reconstruction l2 loss is used to optimize the entire network." utilizing and used are redundant. 
* In 4.1, Implementation Details does not mention details about the diffusion model, readers can mistakenly think that diffusion models are not trained.
* I don't understand the first ablation study "Using Diffusion as Decoder or Not". The designed experiment seems to remove the upper half of a U-Net model, why does this provide evidence about the importance of diffusion? 
* Table 6 only list computational complexity of diffusion-based method, it would be good to also include other baseline methods compared in Table 1.
questions: questions: * The current method is trained end-to-end and the authors mention that achieving disentanglement in complex data is still hard. Is it possible to fine-tune existing pre-trained LDM using the proposed approach for achieving better disentanglement in complex data? 
* In the third ablation study, why does scalar-valued perform better? It seems that the vector-based can potentially extract more information and hence perform better? Can the author further elaborate here?
limitations: limitations: Yes
paper_summary: paper_summary: This paper proposes a representation learning method which employs an encoder as part of a latent diffusion model. During training the encoder takes the target clean image and encodes it into a compressed representation vector. This vector is then used to condition the denoising UNet as it tries to denoise the noisy observation. The conditioning is done with cross attention where queries come from the UNet inner layer activations and key and values come from the conditioning vector.
The method is trained end to end on some simple datasets and is shown to learn disentangled representations, unsupervised. Results are compared to some existing disentanglement methods with favourable results. It is also qualitatively shown that at least on simple datasets the method works well - each factor in the learned representation learns a different source of variability in the data - colour, orientation etc.
strength_weakness: strength_weakness: All in all this is a nice paper which suffers from a few (sometimes major) weaknesses. 

Disentanglement is a thorny subject when it comes to larger scale experiments. Not just technically, but because when data is more complex the actual definition of what are the disentangled factors becomes blurrier and blurrier. In that sense the paper is inherently in trouble - because the paper focuses so much on disentanglement there is almost no point of me criticising the small scale of experiments, but I will still do it. Using such small datasets really does take away from the potential strength of the paper - these toy-ish datasets are good to get a general idea of how a method works, but I don't think that it can be the final experimental set in today's standards. Gains on "disentanglement" metrics don't mean much in my opinion, and the differences between different model on such simple data is, I think, negligible.

That being said, even for this specific problem setup, there are still issues at hand - the use of "latent" diffusion, as opposed to pixel diffusion, is both surprising, not explained in the paper and not analyzed properly. I would expect to see results directly on pixels, taking into account the potential of the LD encoder to perform much of the work of disentanglement. No appropriate experiments is shown in the paper, and the LD use is taken for granted.

Finally, there are some missing works and baselines that at least should have been discussed, if not compared to in the paper: DIPVAE (Variational inference of disentangled latent concepts from unlabeled observations. 2018), and more recently SODA (https://arxiv.org/abs/2311.17901)
questions: questions: See above.
limitations: limitations: These are mostly discussed in the paper.
paper_summary: paper_summary: The paper introduces a new approach using diffusion models with cross-attention to improve the learning of disentangled representations in images. By encoding an image into concept tokens and using cross-attention to connect the encoder and the U-Net of the diffusion model, the authors show that the diffusion process creates information bottlenecks that promote disentanglement.
strength_weakness: strength_weakness: - The ablation studies are conducted only on one dataset. I am concerned that the results could be inconsistent between different languages. 

- The model does not compare itself to Diff-AE[24] because it does not explicitly disentangle the data. However, to show improvement, it is essential to compare the models to see if they are better since they are similar in many aspects. Diff-AE does show disentanglement qualities in the paper. In addition, a comparison of high-quality Diff-AE and the suggested method could be interesting. 

- The resource comparison between the method and other competitive methods is unclear and lacks empirical results.
questions: questions: - It is unclear to me if section 3.1 is claimed to be the background or contribution of the paper?

- In Eq.4, the author introduces the reverse diffusion process. However, the equation describes the probability of $x_t$ given $x_{t-1}$, which is, as far as I know, the forward process. Is there a typo maybe?
limitations: limitations: No significant limitations compared to current disentanglement methods.

Comments: title: Response to rebuttal
Thank you
Response
comment: Thank you very much for your great efforts in reviewing our paper and responses, and the recognition of our work! We appreciate your valuable feedback and will incorporate these good suggestions into our revision.
I thank the authors for the rebuttal.

After reading it along with other reviews, I'll keep my original score since I believe this is a strong submission. In the meantime, due to the somewhat simplified experimental setting, I'll not raise the score further.
Thank you very much for your great efforts in reviewing our paper and responses. We appreciate your thoughtful and constructive feedback  and will incorporate the suggestions into our revision.
Thank you taking the time to answer my concerns.

Beyond the basic limitations of disentagnlement papers most of my concerns have been answered and I am raising my score.
Thank you very much for your great effort and valuable feedback!
Thank you for your kind response. 
My concerns have been addressed, and I will raise my score.