Paper: Interpretable Diffusion via Information Decomposition (ICLR-2024)

Reviewer comment (Reproducibility):
It is still unclear of the concrete implementation. Without access to the network, but just with API, how to get the mutual information and corresponding visualization.

Author response:
Thanks for your question! We provide our demo code detailing our implementation for image-level and pixel-level MI/CMI in the updated supplementary material. Our information maps don't require the details of the denoiser (the U-Net for Stable Diffusion) but just use the output of the denoiser directly, e.g. predicted noise $\hat \epsilon_\alpha$, to calculate MI/CMI as an integral of MMSE under a bunch of logSNR $\alpha$. Thus, our MI/CMI will benefit from a more powerful denoiser in the future or a mixture of experts. However, attention maps require extracting features from self-attention or cross-attention layers and delicately designing a heuristic way to calculate heatmaps, which is hard to generalize.

Reviewer comment (Novelty):
What is the application field. Attention-based methods are widely applied to image editing and show impressive results. What is the advantage of this decomposition method, compared to attetion.

Author response:
We believe a key strength of our approach is faithful evaluation. On compositional understanding benchmarks, our estimator is significantly better aligned with ground truth compared to other estimators (c.f. the Table in our Scientific Finding section of general response). An application could be image retrieval based on text, for example, along with interpretable maps about why/which parts of an image matched.

Reviewer comment (Presentation):
The paper is poor structured and hard to read.

Author response:
Thank you for your concern! We will update our presentation in our updated version. Meanwhile, we hope that our added information theory preliminaries and contributions can address some of your concerns.

Reviewer comment (Novelty):
The method is relatively trivial, replacing word in the prompt and comparing the difference in pixel level.

Author response:
We are glad that the reviewer found the method so simple that it appears trivial. Mutual information estimation is an extremely challenging problem, with many mathematically complex approaches proposed in the literature. As far as we are aware, ours is the first proposal to estimate MI using diffusion models. However, another diffusion-based method has appeared simultaneously in [1]. We will include this reference in our final version. But we would like to call the reviewer's attention to their final estimator, Eq. 21, which is not nearly as elegant/trivial to understand as ours and is valid only for (average) mutual information but does not generalize to a pointwise version as ours does.

Reviewer comment (Experiments):
Lack of comprehensive comparison with attention based method.

Author response:
We included attention baselines on all tasks that involved word/pixel level localization (Sec. 3.2, 3.3). The main goal of our method, as discussed in the top-level response, is measuring the total strength of relationships between image and text. We show SOTA results for this task in Sec. 3.1. Attention has not, to our knowledge, ever been used as a measure of total dependence between images and prompts because the attention distribution is a normalized distribution over tokens which must always sum to 1.

Reviewer comment (Presentation):
The example in first row of Fig. 1 seems to have difference other than the bear and elephant, why does the heat map only show difference in salient objects? Other examples have similar questions, but Fig. 1 is most obvious.

Author response:
Thanks for the question! A simple reason is just that we threshold the color scale, in addition to your points that these shouldn't exactly correspond.

Reviewer comment (Novelty):
What's the advantage compared to attention based method?

Author response:
We clarify our contributions in the general response. In summary, we propose a mutual information estimator that can be easily obtained from denoising diffusion models. With Stable Diffusion, we find our estimator can provide faithful evaluation for compositional understanding, quantify effect sizes of image edits, and localizing phrases in images.

Reviewer comment (Presentation):
The paper's somewhat inadequate presentation makes comprehension challenging.

Author response:
We appreciate your constructive feedback on the presentation. We have updated Section 2 to include more background on information theory, and will aim to improve presentation throughout.

Reviewer comment (Experiments):
There is insufficient experimentation, which fails to demonstrate how this method could enhance the generation process.

Author response:
We apologize for causing confusion! We would like to invite you to read our Paper objectives and Contributions sections in the general response to clarify that enhancing generation is not the goal. In summary, we believe our core contribution is an information-theoretic estimator that can be obtained directly from a denoising diffusion model. While our approach is not geared towards generation, we nevertheless believe this is an important topic: measuring text-image relationships is a challenging task, and existing approaches are often model- or task-specific. Please refer to our general response for additional discussions and results.

Reviewer comment (Writing):
Inadequate preliminary information on related information theory makes Section 2 difficult to comprehend.

Author response:
We have updated our manuscript with additional background in information theory in Section 2.1 and in Appendix B.

Reviewer comment (Presentation):
Clicking on citations does not redirect to the bibliography, making it challenging to locate the correct reference.

Author response:
Thank you for bringing this to our attention! Indeed, the reference hyperlinks ceased to function when we split the appendix from the main body of our manuscript. We leave this unchanged in our manuscript to prevent excessive triggers of pdfdiff, but will address this issue in subsequent updates.

Reviewer comment (Reproducibility):
Apart from the formations, understanding the implementation of this method is challenging, necessitating more detailed explanations from the authors.

Author response:
Thank you for your interest in understanding our implementation. We have provided code in our supplementary material, and we will release our code on GitHub after the anonymity period.

Reviewer comment (Experiments):
Improved demonstration of how this method enhances the generation process in experiments would be beneficial.

Author response:
Please see our discussion in the second quote block of this response.

Reviewer comment (Experiments):
The authors present this paper as an approach to interpretability for text-to-image diffusion models but there is no clear demonstration of where this could be applied other than better visualizations of text-image relations. It would be informative to discuss (or even demonstrate) whether the mutual information can be exploited to better align text-to-image models with human preferences during inference time or fine-tuning.

Author response:
The main value of our approach is for estimating mutual information for relationship testing, which has many applications and for which we get SOTA results in an image-language similarity matching task in Sec. 3.1. A surprising benefit of this particular mutual information method, in stark contrast to other MI estimators, is that it is possible to interpret the discovered relationships at a word and pixel level.

Reviewer comment (Evaluation):
The broader scope of interpreting the generation process by relating pixels to words is missing from the paper and could greatly benefit it.

Author response:
We invite you to review the table in the Scientific Finding section of our general response. Our INFO metric is more faithful in reflecting the compositional understanding capabilities of diffusion models; it reveals that contemporaneous works underestimate DMs in this department. 

We may presumably adapt our more aligned metric as an unnormalized reward and/or a procedure that filters text-to-image generation results. These filtered results can serve as AI feedback to subsequently improve DMs in compositional understanding via supervised fine-tuning.

Reviewer comment (Evaluation):
How is the attention extracted in Figure 2? The proposed per-pixel mutual information averages the result over all timesteps of the diffusion process. However, the attention seems to be extracted from a single forward pass of the model. That could be disadvantageous as the attention mechanism highlights different scales of features at different timesteps. In section 3.2.2 the authors acknowledge the existence of a "feature pyramid" in the attention but is this again for a single time step of the diffusion process?

Author response:
We appreciate your perspective! While the focus of this work is on evaluating text-image relationships with our INFO metric, we believe INFO may contribute to the mechanistic interpretability of the generation process as well. A large body of work in mechanistic interpretability of language models identifies circuits by attributing a performance metric to certain network components. For LMs, these metrics are often implemented as total probabilities of desired token predictions [1][2]; yet they remain elusive for diffusion models. We hope the INFO metric can serve as a faithful and theoretically-grounded metric for similar works of mechanistic interpretability for diffusion models, a future direction which we identify in the Conclusion section of our manuscript.

Reviewer comment (Theory):
The proposed method is heavily relying on the log-density estimator of McCallester, 2023, but lacks details. Further explanation on the practical behavior of this estimator and how this estimator can be derived, at least in an abstract sense, would help the understanding.

Author response:
Thank you for raising this concern. Our information decomposition builds on ITDiffusion [1] rather than McCallester’s work (clarified in the revision). Unlike the approaches of Song, McAllester, and others, our approach does not require complex differential equation solvers for connecting diffusion and density estimation tasks. We added an appendix on the detailed derivation of the density estimator from [1] in an appendix, to make the paper more self-contained.

Reviewer comment (Presentation):
Rather important figures and tables are in the appendix. There is a fair amount of discussion in the main manuscript with Table 9 and Figure 10, but these are in the appendix, thus making the paper difficult to read.

Author response:
We appreciate your critique! We have moved part of the results from Table 9 and referred example in Figure 10 to our main body.

Reviewer comment (Experiments):
CoCo-IT is the main dataset used for evaluation, but how it has been filtered from CoCo is not described, at least briefly.

Author response:
We described the details of the preprocessing steps for the COCO-IT, COCO100-IT, and COCO-WL datasets in Appendix E.

Reviewer comment (Theory):
Are the “standard definition” and the “orthogonality principle” the full list of the decomposition possibilities? If not, we might need some justification as to why these two are better suited as pointwise information.

Author response:
This is a good point, it is possible to define a whole continuum of different pointwise estimators by simply adding any zero-mean random variable to the estimate. The “standard” definition is special by convention, as it is the difference of two density estimates. The “orthogonal” definition is special because it is the minimum variance unbiased estimator.

Reviewer comment (Evaluation):
Is there a notable relationship between $i^o$ and the classifier-free guidance, considering their similarity in computation? For example, if the guidance scale gets larger, does it make an enhanced visualization or a larger MI between the text and image?

Author response:
This is a great question. The form of the MI and classifier-free guidance formula both include the difference between a conditional estimator and an unconditional one. In this sense, we can interpret the density model with the modified score function using guidance as being one that artificially boosts the mutual information of the original density function. It would be interesting to explore these types of score-based information manipulations further.

Reviewer comment (Reproducibility):
Eq. (5) and (6) contain expectation over the noise, $\epsilon$. But different noises will make dramatically different images. How are the visualizations in Fig. 2 generated, and how are they well aligned with the image?

Author response:
Thanks for your question. The calculation of CMI/MI is just an integral of MMSE differences, and MMSE is an expectation over MSE between gaussian noise and the predicted one, not the noise itself. Additionally, we repeat the denoising on multiple sampled noises.

Reviewer comment (Theory):
The walkthrough of some key results that help construct the proposed information decomposition approach is too brief: it would be nice to have a more in-depth review of these results, for example, the derivation of Eq. (2), the part of diffusion model parameterization that’s related to MMSE, etc. Meanwhile, the experiment section is relatively lengthy: some parts could be moved to the appendix, leaving the contents most closely connected to information decomposition in the main paper.

Author response:
Thank you for your interest in our theoretical results. We have updated Section 2 to provide additional background on information theory and our derivation.

Reviewer comment (Evaluation):
Error bars in Table 3 for the comparison between CMI and attention are not reported.

Author response:
We have updated our Table 3 (now Table 4) in our manuscript with (standard errors) for pixel-level experiments. We observe INFO and attention-based methods to be relatively stable. We will update image-level results in our subsequent updates.

 |           | $i(x; y \| c)$ |  Attention |
 |:-------------------|:--------:|:--------:|
 |     Image-level     |   **0.50**   |   0.31   |
 |      Pixel-level     |   **0.31** ($\pm 0.017$)   |   0.27 ($\pm 0.019$)   |

Reviewer comment (Reproducibility):
Could the authors give a quick example of how $i^o(\mathbf{x};\mathbf{y})$ from Eq. (5) is computed? Under the DDPM parameterization, the noise network predicts the forward noise, thus it makes sense to start with a real image $\mathbf{x}$, sample $\mathbf{\epsilon}$ to obtain $\mathbf{x}_{\alpha}$, and compute the L2 term: this utilizes the forward diffusion. Meanwhile, the goal is to decompose the *generated* samples, which uses the reverse diffusion. Therefore, it’s a bit unclear on how these mutual information estimators are calculated.

Author response:
You say that "the goal is to decompose the generated samples, which uses the reverse diffusion", but we want to strongly emphasize that this is **not** the goal. We want to decompose the information in **any image, real or generated**. Therefore, our measure always uses the "forward diffusion". That is, we always start with an image, add noise, and then measure the error of the denoiser to recover the signal. This doesn't require access to a reverse diffusion process that generates the image.

Reviewer comment (Theory):
The authors mentioned that Eq. (2) would hold for an “optimal denoiser”: what does it mean for a denoiser to be optimal?

Author response:
The 'optimal denoiser' means the denoiser minimizes the MMSE (Eq.1) for recovering $x$ in this noisy channel. We clarified this in the revised text. 

[1] Kong et al. 2023 Information-Theoretic Diffusion ICLR