Diffusion Model with Cross Attention as an Inductive Bias for Disentanglement

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. This is the first work to reveal the potent disentanglement capability of diffusion models with cross-attention, requiring no complex designs. We anticipate that our findings will inspire more investigation on exploring diffusion for disentangled representation learning towards more sophisticated data analysis and understanding.

Paper

References (44)

Scroll for more · 32 remaining

Similar papers

Peer review

Reviewer Mm6Q6/10 · confidence 3/52024-07-05

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.

Strengths

- The paper is clear and well-written - In general, research disentanglement with diffusion models is promising and a good direction for improving the disentanglement community. - The presented inductive biases are somewhat novel and produce empirically strong disentanglement results with a more straightforward framework.

Weaknesses

- 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

- 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?

Rating

6

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

No significant limitations compared to current disentanglement methods.

Reviewer rNx85/10 · confidence 4/52024-07-08

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.

Strengths

This model delves into a relatively underexplored side of diffusion models and while not particularly original is an interesting specific and simple (in a good way) combination of existing methods. See below for some reservations about the originality (though not a major decision factor here). The paper supports, on quite a small scale, most of its claims. The experiments are mostly well run and analysis and ablation is adequate. I particularly liked the attention visualization figures. The paper is well structured, figures are clear and all in all well presented. I found the language and writing a bit dense however.

Weaknesses

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

See above.

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

These are mostly discussed in the paper.

Authorsrebuttal2024-08-12

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.

Reviewer HVuH7/10 · confidence 3/52024-07-25

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.

Strengths

I find this paper a strong submission. *Novelty*. To the best of my knowledge, utilizing cross-attention with features from an image encoder to achieve feature disentanglement is a novel approach. *Presentation*. The paper is overall well-written and organized. The figures for framework/concept demonstration are also quite clear. *Good intuition and solid experiment*. The proposed method is well-motivated and the paper demonstrates strong empirical performance. Furthermore, the empirical observations validate the functionality of each component in the proposed method, making the paper more sound.

Weaknesses

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

* 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?

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

4

Limitations

Yes

Reviewer rNx82024-08-12

Thank you

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.

Reviewer HVuH2024-08-13

Response to rebuttal

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.

Authorsrebuttal2024-08-13

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.

Program Chairsdecision2024-09-25

Decision

Accept (spotlight)

© 2026 NYSGPT2525 LLC