CoInD: Enabling Logical Compositions in Diffusion Models

W1: Formal Definition of Non-Uniform and Partial Support: We defined both non-uniform and partial support scenarios in Appendix D.4.

W2: More Complex Causal Graphs: There are two causal graphs, the true underlying one and the observed one.

Underlying Causal Graph: Compositionality is formally defined over independently identifiable attributes (see Wiedemer et al., 2024 and references therein). As such, the causal graph we use in the paper satisfies the requirements for compositional generation, which are $C_i\rightarrow X$ and $C_i$ and $C_j$ affect $X$ separately (L120-124). The dependent attribute scenario alluded to by the reviewer does not satisfy the formal definition of compositionality. In other words, one cannot achieve the defined (full) compositionality if attributes are dependent.

Having said that, CoInD can be trivially adapted to any known causal graphs between the attributes. To achieve logical compositionality while respecting the underlying causal relations, Eq. (4) can be modified to preserve the known causal relations. For example, for a three-attribute ($C_1$, $C_2$, and $C_3$) case, if $C_1$ and $C_2$ are causally related to each other and $C_3$ is independent of other attributes, Eq. (4) will be,

$L_{\text{comp}} = W_2\left(p(X\mid C), \frac{p_\theta(X\mid C_1, C_2)p_\theta(X\mid C_3)p_\theta(C_1, C_2)p_\theta(C_3)}{p_\theta(C)p_\theta(X)}\right)$.

One can apply CoInD as we did in the main paper to achieve logical compositionality while breaking the underlying causal relations.

Observed Causal Graph: In all experiments, the observed causal graph is more complex, with dependencies between attributes. For example, in the Shapes3d dataset, we considered six attributes that could be varied. These six attributes are correlated in partial and orthogonal support settings, resulting in a complex causal graph.

W3: Experiments on Complex Real-World Dataset: We applied CoInD to generate face images with logical compositions on two attributes, “gender” and “smiling,” in two settings: (1) training the model from scratch and (2) fine-tuning an existing text-to-image like SD1.5. We used the CelebA face dataset. We added experimental details in Sec. 5.2 and results in Appendix E.

Additionally, we demonstrated how CoInD can be employed to explicitly and independently control the strength of each attribute in the generated image (Appendix E.2). This is feasible only since CoInD explicitly learns the disentangled scores of $p(X|C_i$).

Also, please see the common response “1. Experiments on Real-world Face Dataset”.

Q1: Early Stopping: Thank you for the suggestion. We tried early stopping by Karras et al. (2022) for the partial support scenario on the Colored MNIST dataset. However, the model performance is poor, as we infer from our preliminary results in the table below.

Both Composed GLIDE and CoInD show performance deterioration when early stopping is used. We hypothesize that for unseen compositions, guidance throughout the learning is required for accurate compositions. Karras et al. (2022), unlike our problem of interest, focus on generating seen attribute combinations, so the early stopping idea may not be relevant for compositional generation. We believe a more careful study is warranted.

Q2: Choice of Hyperparameter $\lambda$: We added a detailed discussion on choosing $\lambda$ in Appendix C.3. $\lambda$ is selected such that the gradients from $L_\text{score}$ and $L_\text{CI}$ are balanced in magnitude. We did not use cross-validation or Bayesian optimization as they are prohibitively expensive while training diffusion models. Please refer to the common response “5. What is the criterion used to select the hyperparameter λ?” for more details.

Q3: Continuous Attributes: Yes, CoInD can natively handle continuous features, simply by using a continuous feature as the conditional input to the learned score function. In fact, in the Shapes3D experiment, we treat the discrete orientation labels as a continuous label in our experiments. In Appendix G.5, we demonstrate CoInD’s ability to generate images corresponding to the interpolation between two samples of continuous attributes observed during training.

Q4: CMI instead of JSD: We use JSD as a diagnostic metric for evaluation rather than as an objective to optimize. In scenarios where the labels are discrete, we employ JSD to measure the distance between the joint likelihood with the product of the marginals and not to learn to enforce conditional independence. The choice of JSD is motivated by the fact that it is symmetric. We will explore if CMI can be used as a metric for Shapes3D where the attributes are continuous. In summary, the choice between JSD and CMI only affects the evaluation of CoInD but not its training.

W1/Q1: Link between failure of diffusion and violation of independence: We apologize for any misunderstanding created in our case study in Sec. 3. In the case study, we noted the CS drop between joint $p_\theta(X \mid C)$ and marginal sampling $\frac{p_\theta(X)}{p_\theta(C)}\prod_i \left(\frac{p_\theta(X\mid C_i)p_\theta(C_i)}{p_\theta(X)}\right)$. This drop indicated the violation of conditional independence, i.e., $C_i \not\perp\not\perp C_j\mid X$. This, in turn, is due to learning incorrect marginals $p_\text{train}(X|C_i)$.

Please also refer to the common response for a detailed explanation “3. How is the failure of standard diffusion models linked to the violation of independence assumptions?” We have also restructured Sec. 3 to clarify this and included detailed proofs in Appendices B.1 and B.2.

W2: Making proofs more accessible: Thank you for the suggestion. To make the proofs more accessible, we added detailed derivations with intermediate steps in Appendices B (Proofs for Claims) and C (Practical Considerations). To summarize, the proof involves the following steps:

Step 1: We start with 2-Wasserstein distance between the joint likelihood $p_\text{train}(X\mid C)$ and the product of marginals $\frac{p_\theta(X)}{p_\theta(C)} \prod_i \frac{p_\theta(X \mid C_i)p_\theta(C_i)}{p_\theta(X)}$.

Step 2: We apply triangle inequality on this distance w.r.t. the learned joint likelihood $p_\theta(X\mid C)$.

Step 3: Kwon et al. (2022) showed that the 2-Wasserstein distance between two distributions is upper bounded by the expectation of the $L_2$ norm of the difference in their scores (Eq. 16 in Kwon et al., (2022)). We apply this result to each term to get their corresponding upper bounds.

W3/Q2: Diverse Datasets and More Quantitative Insights: The submitted paper already provides quantitative insights on two synthetic datasets under four different support settings through 2 metrics (CS and JSD).

In the revised paper, we added an experiment with a more realistic CelebA dataset (Tab. 2) and compared the effectiveness with CS, FID, and Shannon entropy metrics.

In summary, the revised paper provides the following quantitative insights:

1. Conditional independence: JSD in Figs. 4a, 6a, and Tab 2 demonstrate that baseline methods have higher JSD than CoInD. We conclude that CoInD respects conditional independence, $C_i\not\perp\not\perp C_j\mid X$, better than the baselines.
2. Fine-grained attribute control: Fig 11 shows that CoInD also provides fine-grained control compared to baselines (Sec. 5.2 and App. E.2).
3. Diversity of unconditioned attributes: Shannon entropy in Fig. 5 confirms that CoInD is more diverse in unconditioned attributes.

If the reviewer has further specific suggestions for quantitative insights, we will happily include them in the paper.

If our responses addressed your initial comments, please consider raising the score.

Q1: Effect of Large $\lambda$: Yes, a high λ value gives perfect conditional independence and also affects how well $\mathcal{L}_\text{score}$ is optimized. Fig. 8 shows that CS of CoInD on Colored MNIST with partial support drops for $\lambda > 1$. A model trained with a very large $\lambda$ value, thus, ignores the attribute conditioning and instead generates unconditional images which is a trivial, yet undesirable, solution to enforce conditional independence. Also, please see the common response “5. What is the criterion used to select the hyperparameter λ?” for more details on our hyperparameter selection.

Q2: Dependent Attributes in Data Generating Process: CoInD is flexible to the underlying causal relations between the attributes and can be modified according to the goal, i.e. whether to break known dependence for conditional generation or to improve logical compositionality while maintaining the known dependence relations. Please see to the common response “4. What issues will the proposed algorithm face if the attributes are actually dependent on the true data-generating process?” for a more detailed answer.

Q3: Scalability of CoInD w.r.t. Number of Attributes: Our experiments show composition with 6 attributes in Shapes3D. CoInD achieves a high CS even in this case, demonstrating its scalability. Increasing the number of attributes to investigate the scaling law of CoInD requires retraining the diffusion model, which we could not finish in the limited review time. We will add this to the next revision of the submission.

Q4: Orthogonal Support and CS Measurement:

Orthogonal support (Fig. 3): Yes, CoInD can generate such unseen images. We show one such example of an unseen combination in Fig. 1(d) [color: cyan digit: 4].

CS measurement: Yes, on average, 55% of the images contain the desired attributes correctly. Conformity score (CS) is the average across all unseen attribute combinations. We added the CS for each attribute combination for the partial support scenario in Colored MNIST in Fig. 16 in App. G.3.

Overall, Fig 4a of the paper shows that CoInD outperforms the baselines by a huge margin on the unseen attributes. We reproduce this table from the main paper to emphasize the performance gains.

While the baselines achieve a mere 22% CS score, CoInD achieves 55%, representing a significant improvement.

Q5: Failed Image Generations: Thank you for the suggestion. We added failed generated images in Appendix G.2 for Colored MNIST, Shapes3d, and CelebA datasets. CoInD sometimes leaks attributes from the closest seen attribute combination.

If our responses addressed your initial comments, please consider raising the score.

W1: Experiments on Real-World Image Dataset: We applied CoInD to generate face images with logical compositions on two attributes, “gender” and “smiling,” in two settings: (1) training the model from scratch and (2) fine-tuning an existing text-to-image like SD1.5 (Appendix E.3). We used the CelebA face dataset. We added experimental details in Sec. 5.2 and results in Appendix E.

Additionally, we demonstrated how CoInD can be employed to explicitly and independently control the strength of each attribute in the generated image (Appendix E.2). This is feasible only since CoInD explicitly learns the disentangled scores of $p(X|C_i$).

Also, please see the common response “1. Experiments on Real-world Face Dataset”.

W2: Motivation Examples for CoInD: Thank you for your suggestion. We have added image editing as an application that could benefit from logical compositionality in Sec. 1 (Introduction). More applications and an example are provided in the common response “2. What are the practical applications of composition generation?”.

W3: Conformity Score Definition: We added a discussion of conformity score in the experimental section (Sec. 4). We have also made a minor correction to the definition of conformity score at L203:

$\text{CS}(g) = E_{p(C)p(U)}\left[\prod_i 1(C_i, \phi_{C_i}(g(C, U))\right]$

where $C$ is the desired attribute combination at test time, $U$ denotes the stochastic factors in the generative model $g$, and $\phi_{C_i}$ denotes the attribute classifier that predicts the value assumed by $C_i$ from a generated image. $1(\cdot, \cdot)$ is the indicator function. $p(C)$ distribution corresponds to that of the seen combinations for the uniform scenario and of unseen combinations for other scenarios. We updated the PDF accordingly in both Sec. 3 (case study) and 4 (experimental evaluation) to describe what CS truly measures.

W4: Effectiveness of CoInD for Orthogonal Support (Fig. 3): Yes, the proposed approach is effective for the orthogonal support scenario shown in Figure 3. Fig. 6(a) and Tab. 2 provide quantitative evidence that CoInD outperforms the baselines in this scenario.

W5: Numbers in Line 417: We modified the lines (L424-425 in the revised paper) to emphasize CoInD's improvement over the baselines.

W6: Examples of open and closed set compositions: By open-set composition, we believe the reviewer meant generating images where attributes take values never seen in the training set. For example, in the Colored MNIST experiment, this corresponds to generating a digit in a new 11th color that has never been observed in the training set. We have added this example to L147-149 in Sec. 2. As we mentioned in Sec. 2, the open-set composition is beyond the current scope of our paper and remains an exciting avenue for future work. Also, please see the common response “6. Open set vs closed set composition”.

W1: Experiments on Real-World Images: Thanks for a very great suggestion! We applied CoInD to generate face images with logical compositions on two attributes, “gender” and “smiling,” in two settings: (1) training the model from scratch (Sec 5.2), and (2) fine-tuning an existing text-to-image like SD1.5 (Appendix E.3). We used the CelebA face dataset. We added experimental details in Sec. 5.2 and results in Appendix E.

Additionally, we demonstrated how CoInD can be employed to explicitly and independently control the strength of each attribute in the generated image (Appendix E.2). This is feasible only since CoInD explicitly learns the disentangled scores of $p(X|C_i$).

Also, please see the common response “1. Experiments on Real-world Face Dataset”.

W2: Practical Applications of NOT Operation: NOT composition is a valuable tool in controllable image generation. For example, suppose one wants to generate an image of a cow. The resulting image is that of a cow in a pastoral background. NOT composition allows the user to have more control over the background attribute. For creative purposes, the user may ask to generate an image with the attributes “cow, NOT pasture”.

Please also refer to the common response “2. What are the practical applications of composition generation?” for applications and examples of logical compositionality.

W3: Why Logic Operations on Attributes with Multiple Values: Our work aims to facilitate controllable image generation. Restricting the model’s choice in some attributes is valuable for controllable image generation. For instance, when generating thousands of images, it is desirable to have diversity in attribute values. However, through this logical formulation, we can afford users explicit control over this diversity.

A real-world example is “generate an image of a butterfly with [blue OR orange OR white] color.” Since butterflies come in several colors, using a logical composition with multiple choices is better than excluding several other possible colors to get limited diversity.

W4: Motivation for and processing of Logical Expressions: Motivation for Logical Expressions: Logical operations are preferred for controllable generation since they (1) offer fine-grained control over the desired attributes and (2) can be interpreted directly. Also, see “2. What are the practical applications of composition generation?” in our common response.

Processing Logical Expressions: We use the probability formulation of the logical expression provided by Li et al. (2023) and Du et al. (2020). Here, logical expressions are expressed in terms of score functions of likelihood. Eqs. (1) and (2) in the preliminary section show these formulae for AND and NOT operations. L162-172 in Sec.1 in our paper discusses these in detail.

The outputs of diffusion models directly give the score functions of likelihood terms. To obtain images, we substitute these outputs in Eq. (1) and (2) and sample according to Langevin dynamics given in Eq. (11). Please refer to Appendix A for a detailed description of the preliminaries of diffusion models.

W5: Is AND Enough For Logical Compositions: No, AND is not sufficient. One must combine primitive operations: OR, NOT, AND to support arbitrarily complex logical operations. For example, the NAND operation is NAND(X, Y) = NOT (X AND Y).

If our responses addressed your initial comments, please consider raising the score.

W4: Generate Attributes with Non-Uniform Probability: Yes, CoInD can generate attributes with non-uniform probability as demonstrated by the above Gaussian support experiment. If only some attributes are conditioned in the logical expression, as we showed in the discrete Gaussian example above, the unconditioned attributes take their true conditional distribution $p_\text{true}(C_{\text{uncond}} \mid C_{\text{cond}})$.

W5: Incorrect Marginals and Violation of Conditional Independence: Yes, the limitation we refer to in the paper is about the violation of conditional independence, but this is caused since the marginals are learned incorrectly. Please refer to the common response “3. How is the failure of standard diffusion models linked to the violation of independence assumptions?” for a detailed answer. We have also updated Sec. 3 in the main paper accordingly and have added Appendices B.1 and B.2 to explain the same.

W6: Open-Set Composition: Yes, open-set composition is beyond the scope of this paper as we aim to achieve controllable generation with known attributes. Please refer to the common response “6. Open set vs closed set composition”.

However, we showed that logical compositionality can be achieved by fine-tuning an existing model (i.e., without training from scratch). In Appendix E.3, we fine-tuned Stable Diffusion v1.5 to generate face images where we controlled gender and smile attributes.

Q1 Definition of conformity score: CS measures how many of generated images have attributes (given by the classifier) that satisfy the desired logical composition. Formally, it is defined as,

$\text{CS}(g) = E_{p(C)p(U)}\left[\prod_i 1(C_i, \phi_{C_i}(g(C, U))\right]$

where $C$ is the desired attribute combination at test-time, $U$ denotes the stochastic factors in the generative model $g$, and $\phi_{C_i}$ denotes the attribute classifier that predicts the value assumed by $C_i$ from a generated image. $1(\cdot, \cdot)$ is the indicator function. $p(C)$ distribution corresponds to that of the seen combinations for the uniform setting and of unseen combinations for other settings. These images are generated by passing the expected attributes in the form of logical compositions. We have updated Sec. 3 (case study) of the PDF accordingly.

Q2: General Logical Formulae: Yes, CoInD is applicable for general logical formulae. We demonstrated the capability of CoInD for AND, NOT, and OR compositions since any logical composition can be expressed through compositions of these primitives. See examples of general logical formulae in Fig 1d.

Q3: Implied Claim or Assumption: Yes, this is an additional assumption made in L123. We have clarified this in L120. Apart from the independence of the attributes $C_i$’s, we also assume that they affect the image $X$ separately, i.e. their effects are disentangled. As a result, $C_i$’s are conditionally independent given $X$. This assumption of an invertible mixing function is commonly used in causal representation learning.

If our responses addressed your initial comments, please consider raising the score.

W1: Soundness and Clarity of Equations: First, we list our loss function and the variations we consider. The equation numbers correspond to the revised PDF.

Eq. (7): $L_{\text{CI}} = E_{p(X, C)}\| \nabla_X\log p_\theta(X\mid C) - \nabla_X\log p_\theta(X) - \sum_i \left[\nabla_X\log p_\theta(X\mid C_i) - \nabla_X\log p_\theta(X)\right]\|_2^2$ (CI loss)

$L_{\text{CI}} = E_{p(X, C)}E_{j,k\sim\mathcal{U}(n)}\| \nabla_X\log p_\theta(X\mid C_j, C_k) - \nabla_X\log p_\theta(X\mid C_j) - \nabla_X\log p_\theta(X\mid C_k) + \nabla_X\log p_\theta(X)\|_2^2$ (scalable CI loss)

Eq. (8): $L_{\text{comp}} \leq K_1\sqrt{L_{\text{score}}} + K_2\sqrt{L_{\text{CI}}}$ (theoretical bounds)

Eq. (9): $L_{\text{final}} = L_{\text{score}} + \lambda L_{\text{CI}}$ (CoInD loss)

The purpose of Eq. (8) was to show that an upper bound of the distance between the true distribution $p(X\mid C)$ and its parameterized causal factorization $\frac{p_\theta(X)}{p_\theta(C)} \prod_i \frac{p_\theta(X\mid C_i)p_\theta(C_i)}{p_\theta(X)}$ consisted of the score term $L_{\text{score}}$ and the conditional independence term $L_{\text{CI}}$.

We found directly optimizing Eq. 8 to suffer from some instability due to larger gradient magnitudes and needs more careful learning rate scheduling. With such tuning, optimizing Eq. 8 gives similar conclusions as Tables 3 and 4 in Appendix C.2. For larger scale experiments we preferred Eq. 9 since we could leverage optimization hyperparameters (learning rate, learning rate schedule, etc.) of existing well-tuned implementations of diffusion models.

The conditional independence in Eq. (7) is linear w.r.t. the number of attributes $n$. However, due to Hammon and Sun, (2006), conditional independence can be approximated with pairwise conditional independence for large $n$. Therefore, using our scalable CI loss, we enforce conditional independence between a pair of attributes $(C_j, C_k)$, randomly sampled every batch. Refer to Appendix C for a detailed derivation.

W2: Uniform Distribution and Conditional Independence: Yes, the attributes being uniform and being independently conditioned on $X$ are two different ideas. CoInD’s objective is to enforce conditional independence between the attributes irrespective of the nature of their distributions. We used four scenarios — uniform, non-uniform, partial, and orthogonal supports — to derive and verify the effect of the learned conditional independence on logical compositionality in diffusion models. These scenarios were used because they (1) clearly depicted the underlying causal model that we wished to study, and (2) are widely seen in practice. Similar scenarios were considered in Wiedemer et al., (2024) (Sec. 3.1, Fig. 3) and Schott et al., (2020) (Sec. 2, Fig. 2).

CoInD does not require $p(C_i)$ to be uniform. We have added the following scenario to the Colored MNIST experiment to demonstrate this point:

Let $C_1$ and $C_2$ be color and digit in Colored MNIST where both $C_1$ and $C_2$ follow a non-uniform categorical distribution (similar to discrete Gaussian). During inference, we task the models with generating images corresponding to arbitrary attribute combinations. Refer to Appendix G.1. The following table compares the performance of CoInD against the baselines.

Conclusions:

1. Similar to our previous results on Colored MNIST, CoInD achieves better CS on all tasks with around 5-10% and 10-40% improvement in CS over Composed GLIDE and LACE respectively.
2. The learned $p_\theta(C_i\mid C_j)$ is similar to the true distribution $p_\text{true}(C_i\mid C_j)$ as shown in Figure 14(c) in Appendix G.1.

W3: Not Uniform After Logical Conditioning: Yes, it is true that when conditioned on a logical expression, the marginal probability $p(C_i\mid\text{logical expression})$ will not be uniform even if the true marginal distribution of the attributes $p(C_i)$ is uniform. This is expected since the purpose of conditioning is to restrict the space of generated images, and thus by definition, it will no longer be uniform. It is not clear to us why and in what form this represents an “issue”.

Fine-grained Adjustment of Attributes in Generated Images using CoInD (Sec. 5.2, App. E.2)

We evaluate CoInD and Composed GLIDE on the task of adjusting the amount of “smiling” in the image by adjusting a scaling factor $\gamma$ that varies the amount of “smiling” in an image. Refer to Eq. (34) for the exact formulation.

Conclusions:

1. CoInD achieves nearly $2\times$ the CS of Composed GLIDE on increasing $\gamma$.
2. CoInD achieves better (lower) FID than Composed GLIDE.
3. Fig. 11 in the Appendix qualitatively compares the samples generated by CoInD and the baseline for varying $\gamma$. CoInD allowed us to vary the amount of “smiling” in the generated images without affecting the gender-specific attributes. However, Composed GLIDE associated the smile attribute with the gender attribute due to their association in the training data. Hence, it generated images with gender-specific attributes such as long hair, which explains Composed GLIDE's worse FID scores.

Real-world face image generation by fine-tuning a T2I model (App. E.3)

Table: Evaluation of CoInD on face image generation when T2I is fine-tuned

Conclusion:

1. CoInD outperforms Composed GLIDE in CS and FID on joint sampling and AND composition tasks. The quantitative results are given below. For more results and details, refer to Appendix E.3.

2. [BxVf, oaF6] What are the practical applications of composition generation?

Many practical scenarios desire explicit control over many attributes of the generated images. For example, suppose you wish to generate a face image of an old man without a beard. This requirement can be expressed as a logical relation in terms of known attributes as “man” AND “old” AND NOT “beard.” Furthermore, logical compositionality is a generalization of the widely studied conditional generation.

Apart from the use cases mentioned above and those discussed in prior work (Liu et.al. (2023), Du et al., (2020), Nie et al., (2021)), we demonstrate that explicitly learning the disentangled scores of $p(X|C_i$) enables diffusion models to afford fine-grained control over the attribute’s strength. For example, we control the amount of smiling in the generated face images through the expression in Eq. (34) in Appendix E.2. These results offer a new direction for fine-grained attribute control in diffusion models.

Code and Reproducibility: We have released the complete code to replicate all experiments along with checkpoints that enable others to reproduce our experimental results.

Simulation Results: Reviewer D8GZ requested a simulation study. To enhance inner workings of our approach, we now present an analysis in Section F.2 where we explain why CoInD works using 2D Gaussian simulated dataset. We selected a 2D Gaussian because its score function can be derived in closed form.

Clarifications: During the review process, we identified additional relevant work and have incorporated these references into the revised manuscript. Additionally, to address potential confusion among readers, we have refined the problem statement and clarified our underlying assumptions. In response to Reviewer [D8GZ, 5mnE]'s concerns regarding the failure mode (RQ1), we have rewritten Section 3 to improve clarity for the readers.

Additionally we would like to clarify questions raised by the meta reviewer:

Meta Assumption about Conditional Independence: The conditional independence assumption is a common modelling assumption applicable in many scenarios. However, it is rather not a meta assumption per se. For example, when modelling velocity, acceleration, and force of an object on Earth, these attributes are not conditionally independent. Therefore, we want to explicitly state to settings where CoInD is applicable.

Clarification of Independence Statement: The statement “Last, we assume that C₁, …, Cₙ affect X independently of each other” was recognised as informal and potentially confusing. We have revised this description to provide a more formal explanation of our assumption regarding the independent contributions of the components. We also slightly modified causal graph to reflect these assumptions.

Unfaithful Distributions and Deterministic Maps: ``Authors consider an unfaithful distribution in the sense that features are independent given X although they are d-connected ….. Especially, when these features are deterministic maps from X, such as reading off whether someone is wearing eyeglasses or not from an image." Although “deterministic maps” is one particular case to state of more general case, where the underlying generative model has a Jacobian of at least rank m ( number of attributes ) [1]. We have modified the problem definition to reflect meta reviewer’s comments.

Conditionals $p(C_i|X)$ simply become indicators: While the true conditionals $p(C_i|X)$ reduce to indicator functions, the learned approximations $p_\theta(C_i|X)$ are obtained via an implicit classifier within the diffusion model are not encouraged to learn these indicator functions. Importantly, the diffusion model is not incentivised to learn the correct $p_\theta(C_i|X)$ , and as commonly observed, machine learning models fail to capture the true indicator function due to confounder shift. In the partial support setting of Colored MNIST the diffusion model relies on color information to guide the digit generation. The following results illustrate this behaviour:

The decrease in accuracy of implicit classifier $p_\theta(\text{digit}|X)$ accuracy in the unseen set compared to the training set indicates that the inherent classifier in the diffusion model relies on spurious correlations. Diffusion models are not immune to these spurious correlations our method addresses this issue by correcting these spurious correlations, rather than attributing the improvement to any “magic” of the diffusion process or its non-deterministic nature.

We appreciate the valuable feedback from AC and the reviewers, which has significantly contributed to the clarity of our work.

[1] Thaddäus Wiedemer, Prasanna Mayilvahanan, Matthias Bethge, and Wieland Brendel. Compositional generalization from first principles. In Advances in Neural Information Processing Systems, 2024