Learning Diffusion Priors from Observations by Expectation Maximization

Thank you for your review and the legitimiate concerns you have raised.

• W1 (Experiments) We follow your suggestion and benchmark MMPS against previous methods (DPS and $\Pi$GDM). We invite you to consult the global rebuttal regarding these additional experiments.

  Concerning Figures 1 and 2, we use Eq. (20) to compute the covariance $\mathbb{V}[x \mid x_t]$, which is tractable in this toy experiment. We do not see how this would be misleading.

  We also note that we never materialize the Jacobian in the experiments of Section 5. Instead, we leverage the CG method to solve the linear system in Eq. (22), which only requires a cheap vector-Jacobian product per iteration.

• W2 (Identifiability) Indeed, we forgot to explain how our method would behave when $p(x)$ cannot be uniquely identified from the observations. We propose to replace lines 298-300 with the following paragraph

  Finally, as mentioned in Section 6, empirical Bayes is an ill-posed problem in that distinct prior distributions may result in the same distribution over observations. In other words, it is generally impossible to identify "the" ground-truth distribution $p(x)$ given an empirical distribution of observations $p(y)$. Instead, for a sufficiently expressive diffusion model, our EM method will eventually converge to a prior $q_\theta(x)$ that is consistent with $p(y)$, but generally different from $p(x)$. In future work, we would like to follow the maximum entropy principle, as advocated by Vetter et al. [36], so as not to reject any possible hypothesis.

  We emphasize that this identifiability issue is a limitation of the problem itself, and not of our method.

We believe that this rebuttal addresses most of your concerns and, therefore, kindly ask you to reconsider your score.

Thank you for your answer.

Regarding Figures 1 and 2 I do not see what is their point since you use the true covariance matrix. A more convincing should include the approximation obtained using the Jacobian of the denoiser and then showing that with this approximation one can get decent results, comparable to what you obtain using the true covariance matrix.

Thank you for clarifying your concern here. When the denoiser is trained optimally, that is when $d_\theta(x_t, t) = \mathbb{E}[x \mid x_t]$, there is no approximation in Eq. (21). Tweedie's formula using the Jacobian of the denoiser gives the covariance matrix exactly. Instead of training a denoiser for this toy problem, we assume an optimal denoiser, which ensures that the results are not biased by a choice of parameterization.

This is only a minor weakness however. I leave my score unchanged.

Given that we have addressed your concerns, we are genuinely surprised by your decision to keep your rating unchanged (weak accept). We would greatly appreciate it if you could describe how to improve our submission.

Thank you for your review and the legitimate concerns you have raised.

• W1 Indeed, we are not the first to use the covariance $\mathbb{V}[x \mid x_t]$ to improve the approximation of the likelihood score. We believe that Finzi et al. [24] were the first to propose it, shortly followed by Boys et al. [25]. As explained in section 6, the difference of our approach resides in the use of the entire covariance matrix, without the need to materialize it (which is intractable) thanks to the conjugate gradient method. Instead, Boys et al. [25] use a row-sum approximation of the covariance's diagonal, which voids some of the benefits of using the covariance as illustrated in Figure 1 an 2. In addition, when the covariance is not diagonal, the row-sum approximation may result in null or negative values, which leads to total failures (NaNs) in our experiments.

  We thank you for bringing STSL [1] to our attention. Are we correct in understanding that STSL uses the trace of the covariance/Hessian to guide the sampling, but does not justify it with a Gaussian approximation of $p(x \mid x_t)$? If so, we are not sure how STSL is similar to the idea of TMPD [25] and/or MMPS.

  In the manuscript we refer a few times to MMPS as a "new" posterior sampling scheme. We propose to replace these occurences with the word "improved". Would this be satisfactory?

• W2 Materializing the Jacobian $\nabla_{x_t}^\top d_\theta(x_t, t)$ would indeed result in high time and memory complexity, as mentioned at lines 142-144. However, we never do. Instead, we leverage the CG method to solve the linear system in Eq. (22), which only requires a cheap vector-Jacobian product per iteration.

• W3 In our experiments, $\sigma_t$ ranges between $10^{-3}$ and $10^2$. This figure can be easily translated to the VP SDE with the relation $\bar{\alpha}_t = \frac{1}{1 + \sigma_t^2}$. The divergence should also be scaled accordingly. We note that using the covariance $\mathbb{V}[x | x_t]$ is orders of magnitude better than heuristics, especially in low noise regimes.

• W4 Yes, this equation simply rewrites $v = (\Sigma_y + A \mathbb{V}[x \mid x_t] A^\top)^{-1} (y - A \mathbb{E}[x \mid x_t])$ and has a solution as long as $\Sigma_y + A \mathbb{V}[x \mid x_t] A^\top$ is invertible, which is the case if $\Sigma_y$ is SPD.

• W5 We follow your suggestion and benchmark MMPS against previous methods. We invite you to consult the global rebuttal regarding these additional experiments.

• W6 This is not true. Section 5.3 is dedicated to a comparison with GSURE-Diffusion [82] on the accelerated MRI experiment. We note that Kawar et al. [82] do not provide the identifiers of the scans they use for evaluation in their manuscript or code.

• W7 Indeed we provide these proofs only for completeness. We propose to add the following sentence in Appendix B for clarity

  We provide proofs of Theorem 1 for completeness, even though it is a well known result [62-65].

We believe that this rebuttal addresses most of your concerns and, therefore, kindly ask you to reconsider your score.

We sincerely apologize for the inconvenience caused by the writing of our manuscript. We believe that most of your questions stem from the lack of clarity of Section 4.2, where we fail to state that MMPS is not bound to the EM context and can be applied to any diffusion prior. We propose to clarify this matter by replacing lines 108-112 in Section 4.2 with the following

To sample from the posterior distribution $q_\theta(x \mid y) \propto q_\theta(x) \, p(y \mid x)$ of our diffusion prior $q_\theta(x)$, we have to estimate the posterior score $\nabla_{x_t} \log q_\theta(x_t \mid y)$ and plug it into the reverse SDE. In this section, we propose and motivate a new approximation for the posterior score. As this contribution is not bound to the context of EM, we temporarily switch back to the notations of Section 2 where our prior is denoted $p(x)$ instead of $q_\theta(x)$.

Diffusion posterior sampling Thanks to Bayes' rule, the posterior score $\nabla_{x_t} \log p(x_t \mid y)$ can be decomposed into two terms [17, 18, 21-25, 53]

$\nabla_{x_t} \log p(x_t \mid y) = \nabla_{x_t} \log p(x_t) + \nabla_{x_t} \log p(y \mid x_t) \, .$

As an estimate of the prior score $\nabla_{x_t} \log p(x_t)$ is already available via the denoiser $d_\theta(x_t, t)$, the remaining task is to estimate the likelihood score $\nabla_{x_t} \log p(y \mid x_t)$.

We now answer your questions, in light of this clarification.

• Q1 We describe the pipeline in plain text in Section 4.1 and as an algorithm in Appendix A.

• Q2 The measurement matrix $A$ is defined by the specific instruments that are used to gather the observations. If the configuration or environment of the instruments changes, the measurement matrix $A$ may also change. The prior $p(A)$ is therefore the empirical distribution of $A$ for the observations $y$ we have access to. To paraphrase, we do not specify $p(y, A)$, it is imposed upon us by the task at hand.

  We propose to clarify this matter in the manuscript with the following sentence at line 82

  For example, if the position or environment of a sensor changes, the measurement matrix $A$ may also change, which leads to an empirical distribution of pairs $(y, A) \sim p(y, A)$.

• Q3 You are right. The main difference is the parameterization of the prior with a diffusion model. At each step we use the current diffusion prior $q_\theta(x)$ to generate samples from the posterior(s) $q_\theta(x \mid y)$ via the proposed MMPS method.

• Q4 Indeed, MMPS can be used to solve any linear inverse problem with a diffusion prior. However, this prior does not need to be the distribution of clean data. It can be any prior, including our intermediate priors $q_{\theta_k}(x)$.

• Q5 We describe the pipeline in plain text in Section 4.1. We unfortunately do not have the space to move Algorithm 1 into the main text.

• Q6 We do not use a pre-trained diffusion model. As you correctly concluded, $q_k(x) := q_{\theta_k}(x)$ is the current diffusion prior parameterized by the current denoiser $d_{\theta_k}(x_t, t)$.

• Q7 We show in Section 4 that in the context of EB, the EM step is equivalent to minimizing the KL bewteen $\pi_k(x)$ and $q_{\theta_{k+1}}(x)$. For diffusion models, this KL minimization is generally conducted by denoising score matching or equivalent formulations (e.g. ELBO) using samples from $\pi_k(x)$.

We believe that this rebuttal addresses most of your concerns and, therefore, kindly ask you to reconsider your score.

Thank you for your review and the pertinent questions you have asked.

• W1 Although we agree that a benchmark with previous empirical Bayes methods would be valuable for the statistical inference community, we don't think it would be relevant to justify our work. First, our method is based on the established EM algorithm, which has stood the test of time. Second, our goal is to train diffusion models from observations as they are best-in-class for modeling high-dimensional distributions and proved to be remarkable priors for Bayesian inference. However, as explained in Section 1 and 6, previous empirical Bayes methods are not applicable to diffusion models. These methods are also typically bound to low-dimensional latent spaces. For these reasons, it is challenging to design a benchmark between our work and these previous EB methods that would be both fair and informative. Instead, we choose to benchmark our work against similar methods in the diffusion model literature.

• Q1 This assumption is actually the same as the Gaussian approximation of Eq. (17). Indeed, assuming that $p(x \mid x_t)$ is Gaussian, we have (Bishop [67])

  $\mathbb{E}[x \mid x_t, y] = \mathbb{E}[x \mid x_t] + \mathbb{V}[x \mid x_t] A^\top (\Sigma_y + A \mathbb{V}[x \mid x_t] A^\top)^{-1} (y - A \mathbb{E}[x \mid x_t])$

  but Tweedie's formulae also gives

  $\mathbb{E}[x \mid x_t, y] = x_t + \Sigma_t \nabla_{x_t} \log p(x_t \mid y) = \mathbb{E}[x \mid x_t] + \Sigma_t \nabla_{x_t} \log p(y \mid x_t)$

  Therefore, we have

  $\Sigma_t \nabla_{x_t} \log p(y \mid x_t) = \mathbb{V}[x \mid x_t] A^\top (\Sigma_y + A \mathbb{V}[x \mid x_t] A^\top)^{-1} (y - A \mathbb{E}[x \mid x_t])$

  which is equivalent to Eq. (20) since $\mathbb{V}[x \mid x_t] = \Sigma_t \nabla_{x_t} \mathbb{E}[x \mid x_t]^\top$.

  Finzi et al. [24] provide a detailed analysis of the Gaussian approximation of Eq. (17) and prove (Theorem 2) that moments of order $n > 2$ converge to 0 at a rate of $\sigma_t^n$.

• Q2 Our method does not make any assumptions on the covariance $\Sigma_y$. It can be anisotropic and/or non-diagonal. In fact, it is always possible to rewrite the likelihood $\mathcal{N}(y \mid Ax, \Sigma_y)$ as $\mathcal{N}(\Sigma_y^{-1/2} y \mid \Sigma_y^{-1/2} A x, I)$. However, you are right that the signal-to-noise ratio has an impact on our method. We expect larger noise levels to slow down the convergence of the EM algorithm, but lead to an equivalent stationary distribution, under the assumption of infinite data. The rationale is that it is possible to identify all the moments of $p(Ax)$ given $p(y)$ regardless of $\Sigma_y$.

• Q3 The assumption of a linear Gaussian forward process $p(y \mid x)$ does not "restrict the performance" of our method but limits its applicability. However, many real-life and scientific problems can be formalized with linear Gaussian forward processes. The accelerated MRI experiment is a good example. The measurement matrix $A$ and covariance $\Sigma_y$ are defined by the specific instruments that are used to gather MRI scans. If the configuration or environment of the instruments changes, the measurement matrix $A$ and covariance $\Sigma_y$ may also change. The prior $p(A)$ is therefore the empirical distribution of $A$ for the observations $y$ we have access to. To paraphrase, we do not specify $p(y, A)$, it is imposed upon us by the task at hand.

  We propose to clarify this matter in the manuscript with the following sentence at line 82

  For example, if the position or environment of a sensor changes, the measurement matrix $A$ may also change, which leads to an empirical distribution of pairs $(y, A) \sim p(y, A)$.

• Q4 There are no experiments on the improvement per iteration of the CG method currently, but we have conducted additional experiments to benchmark MMPS for this rebuttal and find that increasing the number of CG iterations improves image quality/sharpness, but only marginally and with rapidly diminishing returns. We invite you to consult the global rebuttal regarding these additional experiments.

We believe that this rebuttal addresses most of your concerns and, therefore, kindly ask you to reconsider your score.

Thank you for your in-depth review and the legitimate concerns you have raised.

• W1 Indeed, this is one of the limitation of our method, which we mention in Section 3. However, we would like to mention that in our pipeline we start each training step with the previous parameters, which reduces the cost of each training step. Overall, the entire EM procedure for the corrupted CIFAR-10 experiment takes around 4 days for 32 EM iterations on 4 A100 GPUs (see Appendix C), which is similar to the training of AmbientDiffusion [77].

• W2 Our EM method is indeed sensitive to the quality of posterior samples, as mentioned in Sections 4.1, 4.2 and 7. We propose to add the following sentence at line 291.

  [...] sensitive to the quality of posterior samples. In fact, we find that previous posterior sampling methods [21, 22] lead to disappointing results, which motivates us to develop a better one.

• W3 We appreciate that you recognize MMPS as a valuable contribution in its own right. We follow your suggestion and benchmark MMPS against previous methods (DPS and $\Pi$GDM). We invite you to consult the global rebuttal regarding these additional experiments.

• W4 We follow your suggestion and repeat the corrupted CIFAR-10 experiment at more corruption levels (0.25, 0.50 and 0.75). We invite you to consult the global rebuttal regarding these additional experiments.

  We note that the other experiments of AmbientDiffusion [77] also use a diagonal measurement $A$ (a mask), which is why we conduct the accelerted MRI experiment, where the measurement is a more challenging undersampled Fourier transform.

• W5 Indeed, a strengh of our method is that it can handle a wide variety of measurement types, or even learning from several sources with different measurement types. Our three experiments present different measurement types: a linear projection $A \in \mathbb{R}^{5 \times 2}$, a random binary masking, and an undersampled Fourier transform. We note that Gaussian blur, as you suggest, is equivalent to masking the high frequencies of an image, which is similar to the undersampled Fourier transform in the accelerated MRI experiment.

• Concerning W6 and W7, we thank you for bringing these concurrent works to our attention. We will discuss them in our related work section. Are we correct in understanding that in "Consistent Diffusion Meets Tweedie" the goal is to fine-tune a pre-trained diffusion model using data corrupted by isotropic Gaussian noise?

• W8 Indeed, we forgot to explain how our method would behave when $p(x)$ cannot be uniquely identified from the observations. We propose to replace lines 298-300 with the following paragraph

  Finally, as mentioned in Section 6, empirical Bayes is an ill-posed problem in that distinct prior distributions may result in the same distribution over observations. In other words, it is generally impossible to identify "the" ground-truth distribution $p(x)$ given an empirical distribution of observations $p(y)$. Instead, for a sufficiently expressive diffusion model, our EM method will eventually converge to a prior $q_\theta(x)$ that is consistent with $p(y)$, but generally different from $p(x)$. In future work, we would like to follow the maximum entropy principle, as advocated by Vetter et al. [36], so as not to reject any possible hypothesis.

  We emphasize that this identifiability issue is a limitation of the problem itself, and not of our method.

We believe that this rebuttal addresses most of your concerns and, therefore, kindly ask you to reconsider your score.