XDDPM: EXPLAINABLE DENOISING DIFFUSION PROB- ABILISTIC MODEL FOR SCIENTIFIC MODELING

Authors: qianru  zhang
chenglei  yu
yudong  yan
xiangyu  kuang
yi  ma
yuansheng  cao
siu-ming  yiu
tailin  wu
Venue: ICLR_cc_2024_Workshop_AI4DiffEqtnsInSci
Published: 2024-02-09 12:27:50

Abstract: In recent years, diffusion models have emerged as powerful tools for generatively predicting high-dimensional observations across various scientific and engineering domains, including fluid dynamics, weather forecasting, and physics. Typically, researchers not only want the models to have faithful generation, but also want to explain these high-dimensional generations with accompanying signals such as measurements of force, currents, or pressure. However, such explainable generation capability is still lacking in existing diffusion models. Here we introduce Explainable Denoising Diffusion Probabilistic Model (xDDPM), a simple variant to the standard DDPM that enables the generation of samples in an explainable manner, focusing solely on generating components that are pertinent to the given signal. The key feature of xDDPM is that it trains the denoising network to exclusively denoise these relevant parts while leaving non-relevant portions noisy. It achieves this by incorporating an Information Bottleneck loss in its learning objective, which facilitates the discovery of relevant components within the samples. Our experimental results, conducted on two cell dynamics datasets and one fluid dynamics dataset, consistently demonstrate xDDPM's capability for explainable generation. For instance, when provided with force measurements on a jellyfish-like robot, xDDPM accurately generates the relevant pressure fields surrounding the robot while effectively disregarding distant fields.

Main paper: 
 introduction: 
 Generative models like diffusion models (Ho et al., 2020) are valuable for complex tasks in science and engineering, learning from high-dimensional distributions to predict dynamics in fluid dynamics (Cachay et al., 2023) , weather forecasting (Price et al., 2023) , molecular dynamics (Wu et al., 2022) , and more. However, interpreting the increasingly complex generated output (often 10 3 to 10 6 dimensions) is challenging. Researchers seek to understand the relationship between generated components and accompanying low-dimensional signals, such as pressure measurements in fluid Explainable generation tasks pose significant challenges in linking relevant components of highdimensional generated samples to low-dimensional signals. Existing attribution methods, such as Schulz et al. (2020) ; Selvaraju et al. (2017) , are designed for classification tasks with prepared datasets, making them unsuitable for dynamic generation tasks without available low-dimensional signals at inference time. Additionally, the high dimensionality (up to 10 3 to 10 6 ) of the samples further complicates the identification of relevant components. To the best of our knowledge, no diffusion model research has addressed this crucial question.We present xDDPM, an innovative approach to address the task of explainable generation. Unlike previous methods, xDDPM trains the denoising network to selectively denoise the signal-relevant parts of the sample, allowing the non-relevant parts to remain noisy. This explicit association between generation and the given signal is achieved by incorporating an Information Bottleneck (IB) module into the existing DDPM method (Ho et al., 2020) . The IB module learns to identify the relevant components for the signal, which are then used to modify the noise target during training. Consequently, xDDPM's denoising network can generate samples containing only the relevant signal parts while keeping the non-relevant parts noisy during inference.We make the following contributions: 1) Introducing xDDPM, a diffusion-based model for explainable generation. It incorporates an Information Bottleneck to modify the denoising target, generating samples with only signal-relevant components. 2) Providing valuable datasets, including cell dynamics and fluid dynamics, with over 45,000 videos and 1 million images. These datasets facilitate research in explainable generation. 3) Conducting extensive experiments, outperforming seven state-of-the-art baselines on the datasets. For implementation and results, visit https://anonymous.4open.science/r/xDDPM. xDDPM enables interpretable scientific modeling and sets a new performance standard. 
related work: 
 Diffusion Probabilistic Models (DM). Ho et al. (2020) ; Kingma et al. (2021) are leading techniques for density estimation and high-quality sample generation. They leverage image-like data using a UNet neural backbone Ronneberger et al. (2015) ; Ho et al. (2020) ; Dhariwal & Nichol (2021) . A reweighted objective in training improves synthesis quality. Another approach for image generation is GAN-based methods, such as infoGAN Creswell et al. (2018) ; Chen et al. (2016) , which generates realistic samples and maximizes mutual information between latent variables and output. However, these methods lack explanations for generated samples. Denosing Diffusion Probabilistic Model (DDPM) for AI for Science and Explanation for Artificial Intelligence are put into Appendix.  
method: 
 In this section, we detail our method of Explainable Denoising Diffusion Probabilistic Model (xDDPM). We first introduce the problem setup in Section 3.1. In Section 3.3, we detail our xDDPM method, including its learning objective, model architecture, and inference method.Consider that we have high-dimensional data samples1 {x(n) } N n=1 of the variable X ∈ R D . Accom- panying each x (n) is the low-dimensional signal s (n) ∈ R d for signal variable S.Here, instead of the standard task of learning the distribution2 P X (x) for generating X, we are interested in learning a distribution P W (w) for generating the explainable variable W , such that W is as faithful to X as possible while containing exclusively relevant features to S. Concretely, we want to learn a distribution P W (w) for sampling W , where W = g(X, S) ∈ R D has the same dimension as X, such that the relevant features of W w.r.t. S obey the same distribution as X and the non-S-relevant features obey a Gaussian distribution. Note that during inference time, we do not have the accompanying signal S for generating W . Thus, it is impossible to first generate X and then use attribution methods to find which features of X are relevant to S and then produce W . Therefore, the model for generating W must be learned at training time.The DDPM method. To tackle the above task, we build upon the recent advances of Denoising Diffusion Probablistic Models (DDPMs) Ho et al. (2020) . DDPM is an elegant method to learn highdimensional distributions P X (x) given data samples {x (n) } N n=1 . Concretely, DDPM consists of a forward process that adds t steps of Gaussian noise to the data sample3 x 0 to obtain a noisy sample x t :x t = √ ᾱt x 0 + √ 1 -ᾱt ϵ,where ϵ ∼ N (0, I) has the same dimension as x 0 and ᾱt is a predefined schedule. Since the summation of Gaussian is also a Gaussian, the t steps of adding Gaussian noise is equivalent to adding a single Gaussian √ 1 -ᾱt ϵ. After T steps of forward process, x T approximates a standard Gaussian distribution. In the reverse process, the DDPM learns a denoising model ϵ θ (x t , t) that aims to revert the forward process:x t-1 = 1 √ α t x t - √ 1 -α t √ 1 -ᾱt ϵ θ (x t , t) + σ t η, t = T, ...1.Here α t , σ t are pre-defined schedule and η ∼ N (0, I). To learn ϵ θ (x t , t), DDPM use the denoising objectiveEQUATIONwhich aims to predict the added noise ϵ based on the noisy samplex t = √ ᾱt x 0 + √ 1 -ᾱt ϵ.For more details about DDPM, see Appendix A.2.The xDDPM method. Although the above DDPM can learn complicated distribution P X (x) given data samples, it is insufficient to learn the distribution P W (w) for the explainable variable W since there are no ground-truth data W to learn the distribution from. Our key insight is that the relevance discovery of X to the signal S is essentially finding the minimal sufficient information contained within X for predicting S, and the Information Bottleneck Tishby et al. (2000) provides the exact principle and technique we need for extracting such minimal sufficient information. Specifically, we consider a noisy representation X S of X where X S only contains the relevant features of X w.r.t. S. To learn such a representation, we employ the following Information Bottleneck (IB) objective:EQUATIONHere I(•; •) denotes mutual information, and β is a hyperparameter. By minimizing the above IB objective, it encourages X S to contain as much information as possible for predicting S, while retaining as much little information as possible about X, thus encouraging X S to contain the minimal sufficient information of X for predicting S.We define X S as X multiplied with a continuous mask M ∈ [0, 1] EQUATIONThe first term provides an upper bound for I(X; X S ) and the second term provides an upper bound for -β • I(X S ; S) in Eq. 2. Here the encoder p φ (X S |X), the prior distribution q φ (X S ), and the decoder q φ (S|X S ) are all parameterized by learnable neural networks. Minimizing the above L VIB objective encourages finding an encoder p φ (M |X) for the mask M so that the X S = X ⊙ M approximately extracts the minimal necessary information of X for S.Given the above IB module, how can we encourage the DDPM to learn to generate the explainable variable W relevant to S? Note that the objective L ϵ θ in Eq. 10 encourages the denoising network ϵ θ to denoise all features of x t . Empowered with the mask M found by IB, our xDDPM method instead only denoise the features deemed relevant by M . Concretely, we introduce the following modified denoising objective:EQUATIONHere M ∼ q φ (M |X) is obtained by the VIB. For xDDPM, the above denoising objective and the VIB objective are jointly optimized, constituting the full xDDPM objective:EQUATIONHere the expectation is taken w.r.t. (X, S) ∼ p(X, S) (from data), X S ∼ p φ (X S |X), and ϵ ∼ N (0, I). During training, the VIB objective L VIB helps discover the mask M that indicates the relevant features of X to S. For features that are found relevant, the corresponding mask elements in M are near 1, so the xDDPM's denoising objective L ϵ θ ,M reverts to the DDPM objective L ϵ θ . On the other hand, for the features that are deemed irrelevant to S, the corresponding mask elements in M are approximately 0, and the corresponding features in the noise target ϵ ⊙ M are 0. In other words, the L ϵ θ ,M trains ϵ θ to predict zero noise on these irrelevant features.Architecture design. Here we detail the neural network architecture for xDDPM. We use grid-based data (e.g., image and videos) to illustrate the architectural design 4 . For the denoising network ϵ θ , we use the standard choice of U-Net Ronneberger et al. (2015) . For modeling the distribution p φ (X S |X), we first use a U-Net that takes as input X and returns the estimation of mean µ M,φ (X) and logit ξ M,φ (X) for the mask M as its feature maps. Then we use the reparameterization trick Kingma & Welling (2013) to represent p φ (X S |X) with the Gaussian N (X S ; µ X S ,φ (X), σ 2 X S ,φ (X)), where the mean and the variance are conditioned on X as follows:µ X S ,φ (X) = clamp [0,1] (µ M,φ (X)) ⊙ X σ 2 X S ,φ (X) = softplus(ξ M,φ (X) ⊙ X)We employ clamp [0,1] to make sure that the mask mean stays between 0 and 1, and softplus(x) = log(1 + e x ) to ensure that the variance is non-negative. For the prior term q φ (X S ) in L VIB (Eq. 3), there are multiple options. The most general way is to use a mixture of full Gaussians where the mixing weight and the Gaussian mean and covariance matrix are learnable. Instead, we find that the simplest choice of letting q φ (X S ) be a diagonal Gaussian N (X S ; 0, I) works quite well. For q φ (S|X S ), we assume that S obeys a diagonal Gaussian N (S; µ S,φ (X S ), I) where µ S,φ (X S ) can be a U-Net encoder followed by a Multilayer Perceptron (MLP). Thus, log q φ (S|X S ) reduces to a standard MSE loss. Combined together, L VIB in Eq. 3 reduces toL VIB = E (X,S)∼p(X,S) 1 2 β||µ S,φ (X S ) -S|| 2 2 + 1 2 1 + D j=1 log σ 2 X S ,φ (X) -µ 2 X S ,φ (X) -σ 2 X S ,φ (X) j (6)Here in the second term, the j denotes the jth feature for the variable, and it sums over D dimensions since X S ∈ R D . For detailed derivation, see Appendix A.8.Remark on learning speed. Since xDDPM only needs to learn to denoise the relevant features which contain much less information than the full variable, it is much more efficient to train than the original DDPM. Empirically, we find that xDDPM typically converges around twice as fast as the DDPM method. This demonstrates the added benefits of how xDDPM for improving learning speed. 
experiments: 
 In this section, we aim to answer the following questions: (1) How is the explanation ability of our method xDDPM compared with other state-of-the-art methods (e.g., SmoothGrad and GuidedGrad-CAM) on three datasets (wetting, tension and fluid) in terms of three different metrics (namely IoU, Sensitivity-n and Correlation)? (2) How does each component of our method xDDPM affect its explanation capability? (3) What is the efficiency of xDDPM compared with other baselines?We provide results and experiment analysis in supplementary. 
conclusion: 
 Existing diffusion models lack explainability for high-dimensional observations and accompanying signals.To address this, we introduce xDDPM, an Explainable Denoising Diffusion Probabilistic Model. By training the denoising network to focus on relevant components and leaving non-relevant parts noisy, xDDPM achieves explainable generation. Incorporating the Information Bottleneck aids in discovering the relevant components. Experimental results on cell dynamics and fluid dynamics datasets consistently demonstrate xDDPM's explainable generation ability. In Appendix A.12, we also discuss broader impacts and limitations of our method. In summary, xDDPM holds significant potential in bridging the gap between accurate generation and interpretability within diffusion models, critical for modeling tasks across science and engineering. 2022), enabling the analysis of protein-protein interactions and gene regulatory networks. In the field of physics, the DDPM has been applied to study the diffusion of particles in complex systems, such as the spread of heat in materials. Furthermore, in the domain of chemistry, the DDPM has been employed to understand the diffusion of molecules and reactions in chemical systems. These studies highlight the versatility and effectiveness of the DDPM in capturing and analyzing diffusion dynamics across various scientific disciplines Xu et al. (2022) . Ongoing research aims to further explore its potential for solving complex problems in AI for Science.Explanation for Artificial Intelligence. Explanation, also known as attribution discovery, is an ever-evolving research field that has seen the development of various methods aimed at understanding the significance and contribution of different input features. 2020) is to incorporate the information bottleneck as an explanatory component to shed light on fixed and trained neural networks using standard stochastic gradient methods. These attribution techniques like Schulz et al. (2020) necessitate that the task at hand involves classification and their objective is to associate the decision made for a particular class with the significance of pixels in a provided dataset. As far as we know, our research is the first to estimate the amount of information utilized for attribution purposes in diffusion models applied to scientific modeling. 
 
 Appendix: The Denoising Diffusion Probabilistic Model (DDPM) Ho et al. ( 2020) consists of two essential processes: the forward process (or diffusion process) and the reverse process. Let's focus on describing the forward process first. The forward process in a diffusion model approximates the posterior distribution q(x 1:T |x 0 ), which represents the sequence of latent variables x 1:T given an initial value x 0 . This approximation is achieved by iteratively applying a Markov chain that adds Gaussian noise gradually over time. Specifically, the forward process is represented as follows:EQUATIONHere, x t denotes the latent variable at time step t, and x t-1 is the variable at the previous time step. The distribution q(x t |x t-1 ) is modeled as a Gaussian distribution with mean µ t (x t-1 ) and variance β t I, where β t is the variance parameter at time step t, and I represents the identity matrix.The mean µ t (x t-1 ) can depend on the previous latent variable x t-1 and is typically modeled using neural networks or other parameterized functions. By sequentially applying the distribution q(x t |x t-1 ) for each time step, starting from the initial value x 0 , we obtain an approximation of the posterior distribution q(x 1:T |x 0 ) that captures the temporal evolution of the latent variables via q(x 1:T |x 0 ) := T t=1 q(x t |x t-1 ).To describe the reverse process, let's consider a diffusion model with T time steps. Given an observed data point x T at the final time step, the goal is to generate a sample from the initial distribution p(x 0 ). The reverse process in a diffusion model can be formulated as follows: (1) Initialization: Set x T as the observed data point. (2) Iterative Sampling: Starting from t = T -1 and moving backwards until t = 0, sample x t from the distribution p(x t |x t+1 ), where p(x t |x t+1 ) represents the reverse diffusion process.The distribution p(x t |x t+1 ) in the reverse process is typically modeled as a Gaussian distribution, similar to the forward process. However, the mean and variance parameters are adjusted to account for the reverse direction. The specific form of p(x t |x t+1 ) is defined as follows:p θ (x 0:T ) := p(x T ) T t=1 p θ (x t-1 |x t ) (8) p(x t |x t+1 ) := N (x t ; µ θ (x t+1 , t + 1), θ (x t+1 , t + 1))By iteratively sampling from the reverse process, we can generate a sequence of latent variables x 0:T that follows the reverse diffusion process. This reverse sequence represents a sample from the initial distribution p(x 0 ). The reverse process is crucial for training the diffusion model. During training, the model learns to approximate the reverse process by minimizing the discrepancy between the generated samples and the observed data points. This training procedure ensures that the model captures the underlying data distribution and can generate realistic samples.The optimization objective of the diffusion model is conducted via the following negative log likelihood:E[-log p θ (x 0 )] ≤ E q [-log p θ (x 0:T ) q(x 1:T |x 0 ) ] = E q [-log p(x T ) - t≥1 log p θ (x t-1 |x t ) q(x t |x t-1 ) ] =: L (9) A.3 METHODThe DDPM method. To tackle the above task, we build upon the recent advances of Denoising Diffusion Probablistic Models (DDPMs) Ho et al. (2020) . DDPM is an elegant method to learn highdimensional distributions P X (x) given data samples {x (n) } N n=1 . Concretely, DDPM consists of a forward process that adds t steps of Gaussian noise to the data sample5 x 0 to obtain a noisy sample x t :x t = √ ᾱt x 0 + √ 1 -ᾱt ϵ,where ϵ ∼ N (0, I) has the same dimension as x 0 and ᾱt is a predefined schedule. Since the summation of Gaussian is also a Gaussian, the t steps of adding Gaussian noise is equivalent to adding a single Gaussian √ 1 -ᾱt ϵ. After T steps of forward process, x T approximates a standard Gaussian distribution. In the reverse process, the DDPM learns a denoising model ϵ θ (x t , t) that aims to revert the forward process:x t-1 = 1 √ α t x t - √ 1 -α t √ 1 -ᾱt ϵ θ (x t , t) + σ t η, t = T, ...1.Here α t , σ t are pre-defined schedule and η ∼ N (0, I). To learn ϵ θ (x t , t), DDPM use the denoising objectiveEQUATIONwhich aims to predict the added noise ϵ based on the noisy samplex t = √ ᾱt x 0 + √ 1 -ᾱt ϵ.For more details about DDPM, see Appendix A.2. xDDPM method. At inference time, xDDPM generates the explainable variable W using the same procedure as DDPM:EQUATION)Algorithm 1 xDDPM training 1: repeat 2: x 0 ∼ p(x 0 ) 3: ϵ ∼ N (0, I) 4: Take gradient descent step of L xDDPM w.r.t. θ and φ:∇ (θ,φ) [||ϵ ⊙ M -ϵ θ ( √ ᾱt x 0 + √ 1 -ᾱt ϵ, t)|| 2 2 + λL VIB ] 5: until convergedstarting with x T ∼ N (0, I), and we have W := x 0 as the last sample. We see that at inference, the irrelevant features are not denoised and remain Gaussian because the ϵ θ is trained to predict 0 for these features. In this way, we can generate the explainable variable W in inference time where only the relevant features of S are denoised.Taken together, the modified denoising objective L ϵ θ ,M (Eq. 4) and the IB module for discovering the relevant features (Eq. 3) constitute the key innovation of xDDPM. We provide the algorithm for training as Alg. 1.Architecture design. Here we detail the neural network architecture for xDDPM. We use grid-based data (e.g., image and videos) to illustrate the architectural design6 . For the denoising network ϵ θ , we use the standard choice of U-Net Ronneberger et al. (2015) . For modeling the distribution p φ (X S |X), we first use a U-Net that takes as input X and returns the estimation of mean µ M,φ (X) and logit ξ M,φ (X) for the mask M as its feature maps. Then we use the reparameterization trick Kingma & Welling (2013) to represent p φ (X S |X) with the Gaussian N (X S ; µ X S ,φ (X), σ 2 X S ,φ (X)), where the mean and the variance are conditioned on X as follows:µ X S ,φ (X) = clamp [0,1] (µ M,φ (X)) ⊙ X σ 2 X S ,φ (X) = softplus(ξ M,φ (X) ⊙ X)We employ clamp [0,1] to make sure that the mask mean stays between 0 and 1, and softplus(x) = log(1 + e x ) to ensure that the variance is non-negative. For the prior term q φ (X S ) in L VIB (Eq. 3), there are multiple options. The most general way is to use a mixture of full Gaussians where the mixing weight and the Gaussian mean and covariance matrix are learnable. Instead, we find that the simplest choice of letting q φ (X S ) be a diagonal Gaussian N (X S ; 0, I) works quite well. For q φ (S|X S ), we assume that S obeys a diagonal Gaussian N (S; µ S,φ (X S ), I) where µ S,φ (X S ) can be a U-Net encoder followed by a Multilayer Perceptron (MLP). Thus, log q φ (S|X S ) reduces to a standard MSE loss. Combined together, L VIB in Eq. 3 reduces toL VIB = E (X,S)∼p(X,S) 1 2 β||µ S,φ (X S ) -S|| 2 2 + 1 2 1 + D j=1 log σ 2 X S ,φ (X) -µ 2 X S ,φ (X) -σ 2 X S ,φ (X) j (12)Here in the second term, the j denotes the jth feature for the variable, and it sums over D dimensions since X S ∈ R D . For detailed derivation, see Appendix A.8.Remark on learning speed. Since xDDPM only needs to learn to denoise the relevant features which contain much less information than the full variable, it is much more efficient to train than the original DDPM. Empirically, we find that xDDPM typically converges around twice as fast as the DDPM method. This demonstrates the added benefits of how xDDPM for improving learning speed.A.4 EXPERIMENTS Datasets. Our datasets include two cell dynamics datasets and a fluid dynamics dataset. Detailed information about these datasets, including specific statistics, can be found in Figure 2 : The Sensitivity-n values of all methods were evaluated on two datasets, specifically tension and wetting. The range of masked pixels varied from 10% to 90%. The definition of Sensitivity-n can be found in Eq. 14. It was observed that the Sensitivity-n values did not significantly differentiate between the baseline methods. However, our proposed method xDDPM outperformed the others for most masking percentages, demonstrating that our method can discover relevant pixels that are highly correlated to the prediction of the signal. is generated through simulations with systematically modified key parameters to observe wetting behavior. In contrast, the tension dataset examines how the shape of a system evolves under specific grid and domain conditions, achieved by adjusting parameters like tension parameters. As an example, in Figure 3 (b), each trajectory X consists of cell states across 20 time steps. We see that the cell shape gradually becomes more rounded as time progresses within each trajectory due to the surface tension. In addition, we have accompanying signal S of cell areas or circumferences. Our focus extends beyond learning a probabilistic model for generating cell trajectory videos. We are equally interested in identifying the which video pixels corresponds to the signal of area or circumferences. For the Fluid dataset, please refer to Appendix A.9 for a detailed explanation. As exemplified in Figure ?? (b), our interest lies not only in predicting the pressure field conditioned on the boundary but also in discerning which components of the pressure field contribute to the force measurements acting on the boundary. 2021), we have adopted Intersection over Union (IoU) and Sensitivity-n as the evaluation metrics for all the methods in our study. These metrics are widely used in the field and provide valuable insights into the performance of the approaches. The definition of IoU is shown as follows:EQUATIONIoU is a value between 0 and 1, where a higher value indicates a better overlap between the predicted and ground truth regions. The rationale behind utilizing IoU is to calculate the intersection over union between the generated images and the ground-truth cases, considering the corresponding signal variances such as cell areas or circumferences.The Sensitivity-n metric, introduced by Ancona et al. ( 2017), evaluates attribution methods by randomly masking network inputs and quantifying the correlation between the masked attribution and the corresponding decrease in classifier score. For regression task, Sensitivity-n computes the Pearson correlation coefficient between the attribution values and the associated deviation in prediction when masking a set M n of n randomly selected pixel indices:EQUATIONHere, R i represents the relevance of the indexed pixel i, S(x) denotes the prediction from the decode network with clean inputs x, and S(x [x Mn =0] ) represents the prediction value from the decode network after randomly setting the masked pixels to 0. Intuitively, by masking random pixels, the prediction S(x [x Mn =0] ) will deviate from the prediction S(x) that is obtained without masking the pixels. Sensitivity-n quantifies how the reduction of the prediction performance S(x) -S(x [x Mn =0] ) 2 is correlated with the total relevance i∈Mn R i (x) a method assigns to the masked pixels. A good attribution method should assign high relevance values to the pixels such that when those pixels are masked, the prediction performance drop significantly, producing a high Sensitivity-n value.Baselines. We performed comprehensive experiments, comparing our method xDDPM with several state-of-the-art baselines on three distinct datasets. These baselines include Random Schulz et al. Methods that generate cells more similar to the ground truth exhibit higher explanatory abilities. In this regard, xDDPM demonstrates superior explanatory ability as the generated images at each time step closely resemble those of the ground truth.(2020), Gradient Baehrens et al. (2010) , Saliency Simonyan & Zisserman (2014 ), GuidedBP Springenberg et al. (2014 ), GuidedCAM Selvaraju et al. (2017 ), SmoothGrad Smilkov et al. (2017 ), and GuidedGrad Ribeiro et al. (2016) . To help reproduce the experiment results of all methods, we also provide detailed implementations of all baselines in Appendix A.10.We evaluate all methods using three metrics: IoU, Sensitivity-n (50% pixel masking), and Correlation (Table 1 ). Visual results of state-of-the-art baselines and our proposed xDDPM are presented in Figures 3 and ?? . Additional experiments involve masking pixels in generated samples to assess explanation capability (Figure 2 ). Sensitivity-n average values are provided in Table 4 (Appendix A.10). From these results, we make the following observations.Results. Based on the results presented in Table 1 , we observed that our method consistently outperforms all other approaches across all metrics on the three datasets. This can be attributed to several key factors. Firstly, our method xDDPM employs the Information Bottleneck mechanism to guide the discovery of relevant features related to the signal variable. By emphasizing the generation of relevant parts in the samples, our method xDDPM effectively improves the relevance of the generated samples to the cell areas, circumferences, and pressures of the fluid dataset. This is achieved by maximizing the mutual information between the diffused samples and the signal variable. In comparison, the baselines generate samples without such emphasis. Secondly, our proposed mask mechanism, applied to the Gaussian noise, plays a crucial role in generating samples that closely resemble the ground truth samples. Furthermore, we observe that Gradient, Saliency, and SmoothGrad outperform the other methods on all three datasets. These baselines leverage gradient-based techniques to capture the relevant parts related to the signal variable successfully.In Figure 2 , we present the Sensitivity-n values obtained by all methods on images from the wetting and tension datasets. We see that our method, xDDPM, gives a much higher Sensitivity-n values than the other baselines in most conditions (masking percentage). This shows that our method has the strongest ability to identify relevant pixels that are highly correlated to the prediction of the signal.Visualization comparison. Figures 3 and ? ? display generated samples from the best baseline, Gradient, and our proposed method. In Figure 3 , mixed colors denote non-relevant parts to the signal (cell areas, circumference, and pressures) as identified by the methods. We see that: (1) Both our method and Gradient can generate samples with relevant parts to cell areas, but our method achieves even higher relevance and explanation ability. (2) In Figure ??, our method correctly discovers the relevant parts of pressure field (the front of the jellyfish) that contribute to the force on the boundary, while the best baseline fails to generate relevant parts. This demonstrates the effectiveness and general applicability of our method. Figure 4 : Ablation study of xDDPM on the wetting and tension datasets. We considered four variants for our method, which are as follows: (1) "-M which removes the continuous mask used for Gaussian noise; (2) "-I 1 (representing I(X S ; S))" which eliminates the mutual information between the representations of data samples and the signal variable; (3) "-I 2 (representing I(X; X S ))" which removes the mutual information between the data samples and their corresponding representations; (4) "-L VIB " which excludes the Information Bottleneck mechanism that guides our method xDDPM. We notice that each component of xDDPM contributes the final performance. We conducted experiments to assess the impact of each component in our method, presented in Figure 4 . The four variants considered are: (1) "-M " (eliminates the continuous mask for Gaussian noise), (2) "-I 1 (rep. I(X S ; S)" (removes mutual information between data sample representations and the signal variable), (3) "-I 2 (rep. I(X; X S )" (eliminates mutual information between data samples and their representations), and (4) "-L VIB " (excludes the Information Bottleneck mechanism guiding our method). By comparing these variants with the original xDDPM, we evaluate the significance of each component in improving performance. From the analysis in Figure 4 , we observe that all the components of our method are essential to its performance, further validating the effectiveness of our approach. Additionally, we note that the Information Bottleneck (IB) mechanism (-L VIB ) demonstrates the most substantial impact on the effectiveness of our method.In this section, we examine the training time required for all methods and present the results in Table 2 . To ensure fairness, all methods are implemented in Python 3.9, PyTorch 2.1 (GPU version).The experiments are conducted on a server equipped with 48 cores and 8 Nvidia GeForce H800 GPUs. Our observations reveal that our method, xDDPM, achieves competitive model efficiency compared to the baselines across the three datasets with outstanding performance. This validates that our method xDDPM is well-suited for application on large-scale datasets, while maintaining a stable explanatory ability. Among the baselines, SmoothGrad stands out as the slowest due to its approach of enhancing gradient-based attribution maps by averaging gradients across multiple inputs, albeit with competitive performance.Published as a conference paper at ICLR 2024In this section, we provide the deduction progress for the following formula:L VIB = E (X,S)∼p(X,S) 1 2 1 + D j=1 log σ 2 X S ,φ (X) -µ 2 X S ,φ (X) -σ 2 X S ,φ (X) j + 1 2 β||µ S,φ (X S ) -S|| 2 2 (15)Assuming that the prior q φ (X S ) and the posterior approximation p φ (X S |X) are both Gaussian distributions, we proceed with the following notation. Let D represent the dimensionality of X S . Through reparameterization trick, the variational mean and standard deviation evaluated on X are µ X S ,φ (X) and σ 2 X S ,φ (X), respectively. Consequently, we have the following relationship:EQUATIONAnd:EQUATIONTherefore, we can have the following deduction result:-D KL (q φ (X S )||p φ (X S )) = q φ (X S )(log p φ (X S ) -q φ (X S ))dX S (20)EQUATIONThus, β log q φ (S|X S ) = β log(N (S; µ S,φ (X S ), I) is deduced to a standard MSE loss, Thus, we can obtain the loss L VIB = E (X,S)∼p(X,S)1 2 1 + D j=1 log σ 2 X S ,φ (X) -µ 2 X S ,φ (X) -σ 2 X S ,φ (X) j + 1 2 β||µ S,φ (X S ) -S|| 2 2 .A.9 FLUID DATA GENERATIONWe employ the Lily-Pad simulator Weymouth (2015) for generating both the training and testing datasets. The resolution of the 2D flow field is configured to be 128 × 128. It's noteworthy that in the context of Lily-Pad, the flow field is assumed to extend infinitely. The head of the jellyfish remains fixed at the coordinates (25.6, 64). The representation of its two wings takes the form of identical ellipses, characterized by a fixed ratio of 0.15 between the shorter and longer axes. At every instant, symmetry is maintained across the central horizontal line defined by y = 64. To delineate the wing boundaries, we meticulously sample a total of M = 20 points along each wing. The pivotal parameter governing the jellyfish's control in this 2D experiment is the opening angle of the wings. This angle is defined as the deviation between the longer axis of the upper wing and the horizontal line. It serves as the crucial control signal, denoted as w.Each trajectory originates with the widest possible opening angle and proceeds along a periodic cosine curve with a period of T ′ = 200. Trajectories are distinguished by variations in their initial angle, angle amplitude, and phase ratio, denoted as τ -which represents the ratio between the closing duration and the entire pitching duration. For each trajectory, the initial angle w 0 is generated through a two-step process. Initially, a random mean angle w (m) is sampled within the range of [20 (m) , 60 • -w (m) )]. The resultant initial angle w 0 is then computed as w 0 = w (m) + w (a) , constrained within the range of [10 • , 60 • ]. Meanwhile, the phase ratio τ is randomly chosen from the range of [0.2, 0.8]. The opening angle w t at step t follows a specific pattern: it decreases from w (m) + w (a) to w (m) -w (a) as t advances from 0 to τ T ′ , and then it increases from w (m) -w (a) to w (m) + w (a) as t progresses from τ T ′ to T ′ . Beyond this point, w t exhibits periodic variations for t > T ′ . This configuration aligns with previous studies on the propulsive performance of jellyfish Kang et al. (2023) . Each trajectory is simulated for a total of 600 simulation steps, equivalent to 3 periods. To conserve space, only the segment of the trajectory spanning from T ′ = 200 to 3T ′ = 600 steps is saved, with a step size of 10. This decision is made as the simulation from t = 0 to T ′ = 200 is primarily intended for initializing the flow field. Consequently, each trajectory is stored as a sequence comprising T = (600 -200)/10 = 40 discrete steps.In addition to tracking the positions of the wing boundary points and the opening angles w, we incorporate an image-like representation of the wing boundaries. This alternative representation contains valuable spatial information that can be more efficiently assimilated alongside the PDE states (fluid field) through convolutional neural networks. For each trajectory, this image-like boundary representation aligns seamlessly with the shape of the PDE states. At each time step, the boundaries of the two wings are combined and transformed into a tensor of dimensions [3, 64, 64] . Within each grid cell of this tensor, three distinct features are included: a binary mask indicating whether the cell resides within a boundary (marked as 1) or within the fluid (denoted as 0), and a relative position (∆x, ∆y) representing the distance from the cell center to the nearest point on the boundary. This representation enhances the compatibility between boundary information and PDE states. For every trajectory, we retain data on PDE states, opening angles, boundary points, boundary masks, offsets, and force data. These components are specified as follows:• PDE states u: These have a shape of [ T , 3, 64, 64] , representing the fluid field states for each time step, including velocity in the x and y directions and pressure. To conserve space, we downsample the resolution from 128 × 128 to 64 × 64.velocity: [ T , 2, 64, 64].pressure: [ T , 1, 64, 64].• opening angels w: they have a shape of [ T ] . For each step, we save the opening angle in radians.• boundary points: shape [ T , 2, M, 2]. With a shape of [ T , 2, M, 2], we record the boundary points for both the upper and lower wings. Each wing comprises M = 20 points, and each point has 2 coordinates. To ensure compatibility with the downsampling of states, the coordinates in the x and y directions are scaled down to half (64/128) of their original values.• boundary mask and offsets b: They have a shape of [ T , 3, 64, 64] . For each time step, this includes a mask indicating the merged wings along with the half coordinates of boundary points and offsets in both the x and y directions. The resolution is 64 × 64, matching that of the PDE states.mask: [ T , 1, 64, 64].offsets: [ T , 2, 64, 64].• force: it has a shape of [ T , 2]. For each step, the simulator outputs the horizontal and vertical force from fluid to the jellyfish. The horizontal force is regarded as a thrust to jellyfish if positive and a drag otherwise.We generated a total of n = 45, 000 trajectories, each distinguished by varying parameters such as w (a) , w (m) , and τ . Each trajectory occupies approximately 2MB of storage space, contributing to an overall dataset size of around 100GB. To create training samples, we employed sliding time windows that encompassed T = 20 consecutive time steps of both states and boundaries. This configuration of all methods were assessed on two datasets, specifically tension and wetting. The range of masked pixels varied from 10% to 90%. The definition of Sensitivity-n can be found in Equation 10. It was observed that the Sensitivity-n values did not exhibit significant differentiation among the baseline methods, as depicted in Figure 2 . We have averaged the results shown in Figure 2 in this table. However, our proposed method xDDPM consistently outperformed the others for most masking percentages. This demonstrates that our method can effectively identify relevant pixels that have a strong correlation with the signal's prediction. • Guided Grad-CAM (Guided Gradient-weighted Class Activation Mapping) combines the concepts of Guided Backpropagation and Grad-CAM to highlight important regions in an input image for a specific class prediction. The formulation for GuidedGradCAM is as follows:EQUATION• Saliency is an attribution method used to understand the importance of input features for a given model prediction. The formulation for Integrated Gradients is as follows:EQUATIONInterpolate(x ′ , x, α) is a function that linearly interpolates between the baseline x ′ t and the actual input x t at a given scale α. The baseline x ′ t is the mean of x t Considering all of the aforementioned baselines, we have constrained the range of M to be clamped within the interval [0.7, 1.0]. This adjustment aids the baselines in achieving improved performance across all datasets.To ensure fairness and consistency in our experiments, we made several design choices regarding the configuration of our model. The hidden dimension of the U-Net neural network was set to 64, providing an appropriate level of complexity for the task at hand. To balance the contribution of the Information Bottleneck loss (L VIB ), we assigned a weight λ of 0.1 in Eq. 5. Additionally, we set the value of β to 1 for all datasets, promoting a suitable trade-off between reconstruction accuracy and information preservation. For the Gaussian diffusion process, we performed 1000 diffusion steps to allow for comprehensive information exchange. The channel size multiplier of the U-Net neural networks was set to [1, 2, 4, 8] , ensuring effective feature extraction across various scales. The number of channels for the two cell dynamics datasets and the fluid dataset was set to 20 and 42, respectively, accommodating the unique characteristics of each dataset. Three datasets were standardized to an image size of 32 × 32 pixels. To facilitate convergence, we adopted a learning rate of 8 × 10 -5 .A.12 BROADER IMPACTS AND LIMITATIONS Our approach, xDDPM, expands the implementation of explanatory capabilities in denoising diffusion probabilistic models within the AI for science domain. This advancement enables efficient utilization (g) SmoothGrad (h) GuidedGrad (i) xDDPM (ours) Figure 5 : We present visualizations comparing our proposed method, xDDPM, with all baselines, on the Fluid dataset. The central subfigure represents the actual boundary of a jelly-like robot, which serves as a condition for generating the fluid pressure fields, denoted as X. The other subfigures depict the generated pressure field X by our xDDPM and other baselines. The blended colors indicate the irrelevant parts of the pressure fields to the force measurement signal S, as identified by the methods, while the smoother colors represent the relevant parts. We observe that our xDDPM accurately identifies that the enclosing region of the pressure fields is most relevant to the force measurement on the boundary. Furthermore, the generated pressure field exhibits consistent angles with respect to the given boundary. In contrast, the other baselines fail to generate a consistent pressure field and struggle to identify the relevant pixels associated with the signal. on extensive biological and physical datasets, offering significant implications across diverse scientific and engineering fields. In the realm of biology, diffusion models aid in the modeling of protein structures and drug interactions. Additionally, in materials science, diffusion models contribute to the customization of material microstructures and properties. In the physical sciences, these models are employed to simulate the trajectories of moving objects. We believe that, endowed with the explanation generation capability of xDDPM, we can help gain more insight into the modeling of such systems.Our approach combines the strengths of diffusion models and introduces the capability of generating samples with enhanced relevance to signals such as cell areas, cell circumferences, and forces of solid boundaries. However, there are currently certain limitations that need to be addressed. Firstly, our method xDDPM is currently exclusively applied to regression tasks and has not been adapted for classification tasks. It is an exciting future work to adapt it to the classification task, thus improving its generality. Additionally, there is room for improvement in terms of the scalability of our method, particularly when it comes to handling high-resolution samples (e.g., images or videos with 1024 × 1024 resolution). We intend to address these limitations in our future endeavors.In conclusion, our approach makes the first step to imbue diffusion models with explanatory capabilities. Through the integration of the Information Bottleneck (IB) mechanism, our model, xDDPM, empowers the generation of samples that showcase heightened relevance to signal variances. We are confident that xDDPM will play a pivotal role in enhancing modeling techniques for a wide array of scientific applications.Here we use capital letters (e.g., X) to denote random variables and lowercase letters (e.g., x) to denote their instances. The lowercase letters with superscript (e.g., x (n) ) denote data samples.Sometimes, we may be interested in learning a condition distribution P X|C (x|c) for optional condition variable C. For notation simplicity, we ignore such conditions and only add it as needed.Here the subscript t in xt denotes the denoising step.Note that our method is fully general and can also deal with other types of input such as graph and sequence.Here the subscript t in xt denotes the denoising step.Note that our method is fully general and can also deal with other types of input such as graph and sequence. 


Decision: This paper proposes to Explainable Denoising Diffusion Probabilistic Model (xDDPM) that enables the generation of samples in an explainable manner. The concerns raised by the reviewer PPM1 are valid, including missing experiments. However I notice that more information is included in the appendix. Although we have a page limit of 4, I strongly command the author to re-arrange the paper more properly and include more important content in the main text instead of leaving them in the appendix.

Reviews: review: title: title: Relevant topic but diffuse and brief results and evaluation
main_review: main_review: ### Summary:
An explainable DDPM is introduced to learn relevant components in three different tasks. Compared to other baselines, xDDPM is superior, yet several points remain unclear to me, as detailed below.

Overall, the paper organization and structure would benefit largely from a focus on nailing down the task concretely (what is the model supposed to produce) and describing how xDDPM accomplishes it.

### Strengths:
- Link to code and resources provided (not carefully checked, though)
- Nice to see how xDDPM finds the 

### Weaknesses
- Given that "Random" is superior to many other methods, it appears unclear how suitable the other methods are in solving this task. Aren't there other methods in the literature that are better suited and allow for a stronger competition?
- The same argument also holds for the chosen metrics. For example, how does IoU measure whether the method extracts relevant components $S$ of the signal $X$? Also, it is unclear what goes into the metrics. Do you evaluate the generated mask $M$ or the actual output of your method?
- Unclear what price xDDPM introduces to the performance. That is, how accurate is xDDPM compared to DDPM in, e.g., fluid flow forecasting?
- Experiments is extremely vague/short and results section is missing.
- It remains unclear to me, how the method performs and what it is actually supposed to produce as output.

### Questions:
1. How does the forcing around the jelly fish-like robot look like in a normal DDPM?
2. How is mutual information, denoted as $I(\cdot;\cdot)$ defined?
3. What does $S$ represent intuitively? The pixels that change the most, or parts in the signal that effect other parts, or something completely different? Maybe you can provide some handy examples.
4. Wouldn't a baseline already perform well that computes parts in the scene that change least?

### Minor comments:
- Figures seem to lap over the left margin. Maybe correct for improved presentation.
title: title: Original and interesting approach to explainability in diffusion models
main_review: main_review: Summary:
The authors propose to incorporate an Information Bottleneck criterion into the training of diffusion models, to facilitate the generation of samples that contain an explanatory signal within the high-dimensional data space. They define the explanatory signal through a mask that focuses the denoising process only on those regions that are 'relevant'. The proposed IB approach is compared against several other attribution methods as a means to generate these masks.

Pros:
- Interesting and original framing of the question how to obtain explainable samples from diffusion models.
- Theoretically well-motivated objective, with both qualitatively and quantitatively convincing results.
- Scientifically relevant datasets.
- Motivation and theoretical background were presented in a way that was easy to follow.

Cons:
- The application of the baselines is left somewhat ambiguous, are they incorporated into the DDPM training instead of IB?
- The text would benefit from a more in-depth discussion of the experimental setup, perhaps in favor of a less theoretical perspective on IB?
- The utility of the method for scientific application is hard to judge from the provided experiments, could the authors present a case study of how they intend to utilize this method to address scientific questions?