Diffusion PID: Interpreting Diffusion via Partial Information Decomposition

We thank you for your valuable feedback and helpful comments. We address your concerns below.

Q1) Clarification on Eq 2:

We thank you for mentioning this. Unfortunately, our footer note didn't make it through the submission. Eq 2 is correct, and the second equation is derived from the first using the orthogonality principle, as explained below:

$I(X; Y) = \mathbb{E}_{p(x,y)} [\log p(x|y) - \log p(x)]$

$I(X; Y) = \mathbb{E}_{p(x,y)} [ \frac{1}{2} \int \mathbb{E}\_{p(\epsilon)}[ \|\| \epsilon - \hat{\epsilon}\_{\alpha}(x\_{\alpha})\|\| ^2 - \ \| |\epsilon - \hat{\epsilon}\_\alpha (x\_{\alpha} | y) \| \|^2 ] d\alpha]$

By expanding all the squares and re-arranging we get:

$I(X; Y) = \mathbb{E}\_{p(x,y)} \left[ \frac{1}{2} \int \mathbb{E}\_{p(\epsilon)} \left[ \| \| \hat{\epsilon}\_\alpha (x\_\alpha) - \hat{\epsilon}\_\alpha (x\_\alpha | y) \| \|^2 \right] d\alpha \right] + 2 \mathbb{E}\_{p(y)} \left[ \frac{1}{2} \int \mathbb{E}\_{p(x|y), p(\epsilon)} \left[ (\hat{\epsilon}\_\alpha (x\_\alpha) - \hat{\epsilon}\_\alpha (x\_\alpha | y)) \cdot (\hat{\epsilon}\_\alpha (x\_\alpha | y) - \epsilon) \right] d\alpha \right]$

Here,

$+ 2 \mathbb{E}\_{p(y)} \left[ \frac{1}{2} \int \mathbb{E}\_{p(x|y), p(\epsilon)} \left[ (\hat{\epsilon}\_\alpha (x\_\alpha) - \hat{\epsilon}\_\alpha (x\_\alpha | y)) \cdot (\hat{\epsilon}\_\alpha (x\_\alpha | y) - \epsilon) \right] d\alpha \right] \equiv \ominus$

based on the orthogonality principle [1], which states:

$\forall f, \quad \mathbb{E}\_{p(x|y) p(\epsilon)} \left[ f(x\_\alpha, y) \cdot (\hat{\epsilon}\_\alpha (x\_\alpha | y) - \epsilon) \right] = 0$

The term $(\hat{\epsilon}\_\alpha (x\_\alpha | y) - \epsilon)$ represents the error of the MMSE estimator, which is orthogonal to any estimator $f$. Therefore, the second term becomes zero, leading to:

$I(X; Y) = \mathbb{E}\_{p(x,y)} \left[ \frac{1}{2} \int \mathbb{E}\_{p(\epsilon)} \left[ \| \| \hat{\epsilon}\_\alpha (x\_\alpha) - \hat{\epsilon}\_\alpha (x\_\alpha | y) \| \|^2 \right] d\alpha \right]$ which is the 2nd line in Eq 2.

Fig 6: MMSE curves comparing the standard and orthogonal estimators

We also provide graphs comparing the MMSE estimate obtained from the original equation form (the first equation in Eq 2) and the simplified form for varying levels of noise/SNR in Fig 6 in the attached PDF. It can be seen that the original form (dotted line) is more unstable with many zigzag patterns. We also see the orthogonal/simplified form (continuous line) enforces better consistency between the MMSE (blue) and conditional MMSE (red). Thus, this simplification works better in practice.

Q3) Time analysis

Our primary goal in this work was to help improve the interpretability of diffusion models. Although the time complexity of such methods is not of concern usually given that they are used for interpretability of a model, we understand that this is an important aspect. Methodologically, once we have the generated image from diffusion, our method involves running the diffusion model's UNet for denoising to get the MMSE-based log probability estimates and BERT to obtain the terms' probabilities, as required for PID. Thus, our model's time complexity is proportional to that of the diffusion + BERT model. The exact time would depend on the hyperparameters and the computational resources used.

Q4) Public demo and code release

We would like to clarify that we provide code for our method in the supplemental zip, which is visible to the reviewers. This code release will help ensure reproducibility, soundness check and facilitate future research.

We are unable to share a hosted model at this current point in time as we need to maintain anonymity as per the guidelines. We will also release the code and models publicly after the end of the anonymity period. We will also add a public demo for easy usage and testing of our pipeline.

[1] Steven M. Kay. 1993. Fundamentals of statistical signal processing: estimation theory. Prentice-Hall, Inc., USA.

We thank you for your thorough review and insightful feedback. We address your suggestions below and also refer to the global response at the top in a few places.

Q1) Analysis on objects and attributes

We agree that our work could be further improved by providing additional analysis on the interaction between objects and their attributes. We address this by providing new visuals for these types of interactions in Figs 1 and 2 in the attached PDF along with other experiments (CPID) as detailed in the global response above. We will also include these figures in the next revision of the paper.

Q2) Choice of BERT

We would like to point out that BERT was the original language backbone of choice for latent diffusion [1]. They emperically find that BERT can be effectively used to encode semantic information of images in a generation setting. Furthermore, the common methods to compute $p(y)$ involve modeling a distribution of the language's vocabulary in natural language. For this, the only options are: textual databases or language models. Online available databases (ex. Wikipedia:Word frequency, Google Books Ngram Viewer, Corpus of Contemporary American English, etc.), can only be used to obtain the independent probability of each term based on frequency statistics but it remains difficult to obtain an accurate conditional distribution. Moreover, to calculate this over large corpuses online would be computationally very expensive. That's why we used a language model instead. BERT has been the standard language model used in natural language research since its conception due to its excellent performance on various language-based tasks and its accurate model of the vocabulary distribution in both, conditioned and unconditioned cases. We did not feel the need to use LLMs as the context window here is fairly small given that we operate on prompts of moderate length. Thus, this choice does not significantly impact $p(y)$.

[1] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. Highresolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10684–10695, 2022.

We thank you for your comprehensive review and thoughtful comments. We also appreciate your comments for explicitly highlighting the difficulties of adapting PID to diffusion and the challenges of working on the interpretability of models.

Q1) Relationship between PID maps and human intuition

We thank you for acknowledging and drawing attention to the hurdles in this kind of work. Although we show examples that align with human intuition, we also show examples where it's possible that not every relation matches human intuition, as we are exploring the diffusion model's understanding of the world. We also introduce CPID (Conditional PID) as an extension of PID to better account for the context provided by the rest of the prompt. We refer you to the global response above for further details on CPID and the attached PDF for visuals (Figs 1, 2, 3, 4).

Q2) Estimator's uncertainty

We expand further on this query in the global response above and will also add an explanation in the paper's next revision.

Q3) Applications of DiffusionPID

The primary goal of our work was to introduce the concept of PID into the rapidly growing generative field in computer vision. We believe this would lay down the groundwork to further the understanding of the internal concepts learned by diffusion. This would produce more concrete ways to evaluate and progress prompt engineering. PID can identify which elements of a prompt provide unique information, allowing for the refinement of prompts to reduce redundancy and enhance synergy. Our method can also help better understand and counter diffusion's attribute binding problems and learned biases such as those demonstrated in our paper. We also believe incorporating PID-based insights into model training could ensure that the model better captures human-like understanding of concepts and thus make it more human-aligned. One way to address the point regarding the requirement of apriori knowledge of which words to run PID on is that several pairs of terms could be sampled automatically from the prompt (with some filtering of pairs such as based on LLM-based semantic distance between the terms) and run through PID. We agree that the practical application of PID requires further work and is something we plan to explore in the future but believe that this work provides the necessary foundation for this direction.

Q4) PID Diagram

Thank you for raising this point. We have added a figure (Fig 5) to the attached PDF to give more insight into the concept.

Q5) Williams and Beer and Finn C, Lizier JT's work

Yes, as you rightly recognized, one of the major reasons behind building on the Williams and Beer formulation was the pointwise measure. The choice was also based on the observation that the NeurIPS and the AI community previously used this formulation [1]. Furthermore, both works were quite helpful for our paper. We thank you for mentioning the latter and will make sure to cite it in the next revision of the paper.

Q6) Formating and typos

Thank you for bringing this to our attention and we will fix this in the next revision of the paper.

[1] Yu, S., Wickstrøm, K., Jenssen, R., & Principe, J. C. (2020). Understanding convolutional neural networks with information theory: An initial exploration. IEEE transactions on neural networks and learning systems, 32(1), 435-442.

We are deeply grateful for the positive assessment you have given to our work! Will make the required changes in the next revision of the paper.

Thank you again for your suggestions and comments.

We thank you for your detailed review and thoughtful comments. We address your concerns below and will refer to the general response at the top for certain points.

Q1) Novelty of our method and more complex use cases such as an extension to the multi-concept scenario

1. Our approach is the first to adapt PID for diffusion models by developing the required mathematical formulations. We derive a mathematically sound pixel-level PID breakdown in a form that is compatible with and could be integrated into diffusion. Furthermore, we leverage our method to conduct a detailed study of diffusion and to identify its various failure modes, which significantly advances our understanding of different concepts and their interactions in the diffusion model.
2. Building upon the original concept of PID, we also introduce Conditional PID (CPID) for diffusion models to expand the scope and contributions of this work and provide more details in the global response above (Figs 1, 2, 3, 4 in the attached PDF). This also helps address the suggestion on expanding to more terms as CPID takes the rest of the prompt into consideration. We observed that CPID yields superior localized results compared to its PID counterpart in practice as it takes into account the rest of the prompt as context during the image generation process, allowing it to better capture the specific contributions of the two tokens being analyzed. We will include the results in the next revision of the paper as well.
3. Computing the PID and even MI (Mutual Information) between three or more input variables in information theory is a complex process. There is no universally accepted method for defining and calculating the terms of Uniqueness, Redundacy and Synergy. While works such as [1] provide a method to compute the global Redundancy between all input variables, Uniqueness and Synergy remain difficult terms to tackle as these computations need to take all the variables subsets' interactions into consideration. It is generally agreed [1, 2] that multivariate information decomposition is more involved than the bivariate case because it is not immediately obvious how many atoms of information one needs to consider, nor is it clear how these atoms should relate to each other. Even the breakdown provided in [2] for the four-variable case is highly complicated and it is a non-trivial problem to extend these concepts from information theory to more variables.
4. The primary objective of the analysis conducted in the main text was to learn more about the text-to-image diffusion model's understanding of different concepts, which may or may not align with human understanding, and to derive an analysis for specific prompt types. We agree that the examples used in the paper are relatively simple, but they can still shed light on the various interactions and shortcomings of the diffusion model in an easily interpretable form. To address the need for more complex prompts, we provide results on longer prompts with more entities, similar to those used in the diffusion training data, in Figs 1 and 2 in the attached PDF and provide more details on the new experiments in the global response above. We will include the new figures in the next revision of the paper as well.

Q2) Clarification on uniqueness figure 6.5.1 and distinction from single-concept methods

The uniqueness information maps in Fig 6.5.1 tries to show what information that term uniquely contributes on its own to the image generation process. Single-concept methods like DAAM, MI, and CMI highlight the overall information a concept contributes which is a combination of Synergy, Redundancy and Uniqueness. We clarify this in Eq 5 in the main text. It is possible that something can have little Uniqueness, i.e., contributes very little unique information on its own given the rest of the prompt or even just another concept as sufficient context, but still have a contribution through Synergy and/or Redundancy. For instance, in the "calf" vs "field" example, the concept "calf" when interpreted as the calf of an animal, is enough to make the model generate a field. Thus, the "field" concept does not contribute much unique information on its own but it does provide synergistic information in combination with "calf" that helps the model interpret the meaning of "calf" as the animal and not the muscle, which is why we see a high Synergy. As long as there is some other term instead of "field" that can provide the required context to make the diffusion interpret the "calf" as the animal, then removing "field" would still generate the grass field in the image. Similar arguments can be made for the rest of the examples as well. Also, our approach does not change depending on the order of the concepts being processed by DiffusionPID. We will update the figure captions in an upcoming revision to clarify this point. We thank you for bringing this to our attention.

[1] Conor Finn and Joseph T. Lizier. Pointwise information decomposition using the specificity and ambiguity lattices. ArXiv, abs/1801.09010, 2017.

[2] Paul L. Williams and Randall D. Beer. Nonnegative decomposition of multivariate information, 2010.