Improving Probabilistic Diffusion Models With Optimal Diagonal Covariance Matching

Thanks for your comments and the favourable on our presentation! Here are our responses to your concerns.

The idea seems both incremental and way-too simple, which is not a fundamental change in the covariance estimation that would lead to a revolutionized/big change in the performance and/or efficiency. It will have a limited impact on the community.

ANSWER:

Thank you for your valuable feedback. While our method is simple, it is also highly effective:

1. Compared to other diffusion models using heuristic or state-independent isotropic variance, our approach consistently outperforms in both FID and likelihood.
2. Compared to previous methods for diagonal covariance learning, some achieve better FID but worse likelihood (e.g., SN-DDPM), while others achieve better likelihood but worse FID (e.g., NPR-DDPM). In contrast, our method strikes a superior trade-off between likelihood and FID.

We believe the proposed method has the potential for a broader impact on the community. While FID is a widely used metric for generation quality, many applications of generative models—such as out-of-distribution (OOD) detection, lossly compression, or AI4Science problems—place greater importance on metrics like diversity or likelihood. Our method consistently improves likelihood and diversity without compromising FID, all while enhancing efficiency. We will clarify this point further in the revised version.

Moreover, the main insight of our paper is to show that improving covariance estimation can lead to better performance in both sample quality (FID) and density estimation (likelihood) with fewer NFEs. We believe this insight will inspire future research. For example, more advanced covariance approximation methods, such as low-rank approximations, as suggested by the reviewer, represent a promising direction for future work.

The experimental results seem not extraordinary. It does not fundamentally improve the performance (generation quality), which is understandable since covariance is more about sampling efficiency and sample diversity. However, even in terms of the sampling diversity and efficiency, such as indicated by Table 4, the proposed method does not significantly beat strong baselines.

ANSWER: Thank you for your comment. While the proposed method does not significantly outperform strong baselines, we emphasize that it achieves a better trade-off between FID and diversity/likelihood with fewer NFEs. We have included an additional visualization to better highlight this benefit on CIFAR10 (CS) (please refer to this figure for details).

Moreover, as shown in Figure 4, although our method achieves similar FID performance to DDIM, it notably enhances the diversity measured by recall. Additionally, as suggested by reviewers mSjs and tC1T, we expect our work will inspire future research to explore more advanced approximations of the optimal covariance, such as lower-rank or block-diagonal approximations. While these approaches may introduce additional computational costs, we believe they have the potential to achieve better performance. Developing efficient training algorithms for such advanced covariance matching methods is a promising direction for future work.

Moreover, there are no error bars or uncertainty in the experimental results.

ANSWER: Thanks for raising this point, we put the error bar result in Table 9 in the Appendix. We will further highlight this in the main text to improve the rigorousness.

Thanks for the detailed response. I raised my score. Please include the mentioned important experimental results and visualization in the main paper. Please open-source your code if the paper is accepted.

Thanks for your favour in our work! Below are our answers to your questions. Please feel free to leave additional comments if you have any further concerns or would like to discuss them further.

The authors show that the proposed OCM improves covariance estimation for DDPM and DDIM, which results in improved sample quality and better efficiency. There are many other stochastic and deterministic samplers with different sampling schedules, does the improved covariance helps with all these samplers, e.g. EDM [1] stochastic and deterministic samplers?

ANSWER: Thank you for your constructive comment. Directly applying the improved covariance to EDM’s stochastic and deterministic samplers is not straightforward because these samplers rely on ODE and SDE solvers. However, our proposed method can be naturally extended to continuous time, allowing it to integrate into the EDM framework when used with DDPM and DDIM samplers.

Specifically, EDM adopts the variance-exploding forward process $p(x_t|x_0) = \mathcal{N}(x_t; x_0, \sigma_t^2 I)$. Since this forward process is Gaussian, it also has an associated optimal covariance in the backward process. Using our proposed optimal covariance matching objective, we can learn the diagonal covariance for this process. During inference, the DDPM or DDIM samplers can leverage the learned covariance for sampling. However, it remains an open question whether this approach can outperform the original EDM sampler or how to integrate the improved covariance into ODE/SDE-based samplers. These are promising directions for future research.

Any insights or intuitions why comparing to baselines, learning the covariance via OCM objective gives better covariance estimation over all noise levels?

ANSWER:

Thank you for your insightful question. We believe it is a very important question, which we should discuss more in the paper to provide more intuitions. Below, we provide a detailed hypothesis to explain why learning covariance via the OCM objective outperforms baseline methods across all noise levels.

I-DDPM: This method jointly learns the mean and covariance by maximizing the evidence lower bound (ELBO) of the likelihood. We hypothesize that this joint optimization introduces challenges because the covariance estimation depends heavily on the mean. When the mean is suboptimal during training, the ELBO becomes overly loose, resulting in an additional bias in the covariance estimation. In contrast, our method simplifies this process by directly regressing the target covariance, which avoids the complexities of optimizing a loosely bound and leads to a less biased learning procedure.

SN-DDPM: This method estimates diagonal covariance directly from noisy empirical data. While it improves FID in some cases, it struggles with likelihood estimation. A key issue here is the quadratic term in their model (as discussed in Equation 24 of their paper), which amplifies estimation errors. This effect is particularly pronounced at low noise levels, where score estimation is less accurate. Our method, by directly targeting covariance estimation through the diagonal Hessian, mitigates such issues and provides a more stable estimation framework.

NPR-DDPM: This approach achieves strong likelihood results by learning the residual of the diagonal covariance using a pre-trained score network. However, its FID performance is weaker. We hypothesize that this is due to the inherent difficulty of learning the residual covariance, a more complex task than directly regressing the diagonal Hessian, as in our method.

We hypothesize two reasons why our method achieves smaller estimation error:

1. Our approach directly regresses the diagonal covariance, which is simpler and more effective than estimating covariance through a loose lower bound.
2. Learning the diagonal covariance from the score helps reduce the high variance inherent in noisy data, compared to directly learning covariance from the noisy data. This potentially leads to improved performance.

In summary, the improvement of our method may be attributed to the fact that regressing the diagonal Hessian directly from the pre-trained score is a more effective strategy for optimal covariance estimation compared to other baselines.

Thanks for the clarification.

I also read other reviews and your responses, that's very helpful and addresses some of my concerns which I was about to ask. Also thanks for the detailed replies, I really appriciate it. I've increased my score accordingly.

My last comment may be less relavant to the paper, but I'm still interested to see your thoughts. I feel like that could be promising to explore.

We appreciate the reviewer's constructive feedback on our paper. Regarding the weaknesses pointed out, we would like to provide the following clarifications:

The paper doesn't seem to propose anything radically new; the introduction of sensible approximations of the covariance of the time reversal has been explored quite widely in recent years and so the contribution of this paper is the particular approach to doing that which, while well motivated and seemingly having good properties, does not seem particularly radical.

ANSWER:

We appreciate the reviewer’s acknowledgement of the motivation behind our method and its promising properties. We would like to emphasize that our contribution lies in introducing a new learning objective aimed at reducing the approximation error in diagonal covariance estimation. This innovation enables our method to achieve an optimal trade-off between FID and likelihood metrics with fewer NFEs, which has been a persistent challenge for existing methods.

For example, some approaches prioritize better FID at the expense of likelihood (e.g., SN-DDPM), while others improve likelihood but compromise FID (e.g., NPR-DDPM). Our method addresses this trade-off effectively, offering a balanced improvement in both metrics.

To explicitly highlight this advantage, we have included a new trade-off plot that varies the number of network evaluations on CIFAR-10 (CS). Please refer to this Figure for illustration. We will add this plot to the revised version to emphasize the benefits of our approach.

Moreover, we attribute the main insights of our research to the observation that better covariance approximation can lead to improving performance. We hope this insight can inspire future research on more advanced covariance approximation techniques, such as low-rank and block-diagonal approximation.

There are some odd choices made presentationally at times which to my mind make the paper a bit less readable and hence a bit less accessible.

ANSWER: Thank you for pointing these out. We have revised them based on your suggestions. We greatly appreciate these suggestions that make the paper more readable.

Thanks a lot for your valuable comments. Your suggestions are very helpful in further improving the work, and we are refining the manuscript accordingly.

In general, the paper is good. However the results miss visualization (in the core text). A toy example would be welcome for visualization as well and easy results interpretation/comparison. The performances reported do not exhibit high improvement over other methods. Although it shows to be challenging! Maybe other performances should be reported to strengthen the defense of this method over others (the authors often point efficiency as a main goal, but no convincing observations in this favor

ANSWER: Thanks for your valuable comments. We recognize the importance of clear visual demonstrations. In Section 5.1 and Figure 3, we have already provided a toy example using a mixture of Gaussians, where we compare covariance estimation accuracy and DDPM/DDIM sampling efficacy and show that our method archives lower estimation error and better sampling performance. However, we understand this could be elaborated further. In the revised version, we will include additional visualization to enhance clarity and facilitate result interpretation and comparison.

While our method may not outperform all strong baselines in every setting, we emphasize that its key advantage is achieving the best FID-likelihood trade-off with fewer steps. Following your suggestion, we have included an additional visualization on CIFAR10 (CS) to better highlight this benefit (please refer to this figure for details). Furthermore, the central insight of our paper is that improved covariance estimation can lead to better performance in both sample quality (FID) and density estimation (likelihood). We hope this insight will inspire future research. For example, more advanced covariance approximation methods, such as low-rank approximations, as suggested by the reviewer, represent a promising direction for future work.

The title is a bit misleading. You do not make optimal covariance matching, only diagonal matching. I also found the choice of diagonal only matching a little bit skipped. I do not say that it's a bad approximation, but it has its limitations. You should motivate more 'why' diagonal is already a good approximation (more than the complexity motivation).

ANSWER: Thanks for raising this point. In the revision, we will clarify the use of diagonal covariance more clearly in the title/abstraction to prevent potential misinterpretation. Our choice to match the diagonal covariance is primarily driven by considerations of complexity and scalability, aligning it with the architectures commonly employed in mainstream structures such as iDDPM. Our primary goal is to demonstrate that the proposed objective of learning the state-dependent diagonal approximation is better than the state-independent isotropic covariance or other diagonal covariance learning methods, which are common techniques used in diffusion sampling. Additionally, we added an experiment on comparing full covariance and the learned diagonal in the toy MoG setting (please check the DDPM and DDIM results), where we found that the performance of the true full covariance and learned diagonal are similar.

However, we do agree with the reviewer that the covariance structure will be more important in the high dimensional case, where the sparse approximation like low-rank approximation $L_\theta (x_t) L_\theta (x_t)^T$ or block-diagonal approximation can give more flexible structure information and potentially improve the effectiveness of the optimal covariance approximation. However, how to find a scalable loss to efficiently learn $L_\theta$ or the block-diagonal for large dimensional data is a challenging task which could not be easily integrated into the proposed loss function (see more discussion in the comparison to the setting in the inverse problem below). We thus only focus on how to improve the diagonal covariance estimation accuracy, leaving more flexible approximation as a promising direction for future research. We want to thank the reviewer again for pointing out this exciting future direction, we will add a relevant discussion in our revised paper.

Concerning Rademacher samples, it would be nice to see a metric (like FID and IS) evolving with M (the number of samples) to strengthen your claim line 199.

ANSWER: Thanks for the great suggestion. Notably, the training loss defined in Equation 11 is unbiased, meaning that even with a single Rademacher sample, it can accurately recover the true diagonal covariance under optimal training.

To validate our argument, we include an ablation study examining the effect of varying the number of Rademacher samples $M$. Specifically, we conduct experiments on CIFAR10 (CS) using the proposed OCM-DDPM method and evaluate performance based on FID and NLL. The empirical results, presented in the following tables, indicate that a larger M (e.g. M=3) can give a small improvement in FID. This improvement is likely due to the reduced gradient estimation variance during training with a larger M. However, in most cases, different values of $M$ yield consistent performance.

This is practically desirable as, in practice, setting $M$=1 allows for efficient training while maintaining strong performance, While larger M may offer marginal gains in FID, the trade-off between computational efficiency and performance improvement will be further discussed in our revised version. We sincerely thank the reviewer for this insightful suggestion!

FID results:

NLL results:

Furthermore, we are happy to provide additional insights into the OCM objective. It is noteworthy that a similar approach is also employed in denoising score matching. Specifically, given the denoising score identity

we can learn a neural network $s_\theta(x_t)$ to approximate $\nabla\_{x_t} \log p(x_t)$ by minimising the following loss

in which the inner expectation can be estimated using a large number of Monte Carlo samples to reduce variance; however, this method is biased and impractical. Similar to Theorem 2 in our paper, this loss function can be reformulated into the following form:

This objective, commonly referred to as denoising score matching, is unbiased and works well with a single Monte Carlo sample $x_0 \sim p(x_0 | x_t)$.

it would be nice to show more extensive comparison with other covariance matching methods. Either that target full covariance or more costly diagonal estimation. Insist more on the advantage of your estimator concerning the estimation speed/accuracy ratio. In the same trend as table 4. Maybe you can have a small toy problem (MoG, ...) for which it is better to interpret results (and visualize them) + having access to GT covariance.

ANSWER: Thanks for the great suggestion. We compare our method to other covariance matching approaches using a toy problem in Figure 3, where the data distribution is a two-dimensional mixture of Gaussians (MoG). In this scenario, we have access to the true covariance matrix. Figures 3(a) and 3(b) illustrate the $L_2$ error of the estimated diagonal covariance across different noise levels, demonstrating that our method achieves the lowest error for all $t$ values.

To demonstrate the sampling efficacy of the learned covariance, we also perform DDPM and DDIM sampling for the MoG problem. As shown in Figures 3(c) and 3(d), our method outperforms others with fewer generation steps. Additionally, we consider including a baseline that samples using the true full covariance (please click the anonymous link to see the DDPM and DDIM results). The results show that our method achieves performance comparable to the full covariance. However, it is worth noting that while the diagonal approximation and full covariance yield similar performance in this two-dimensional toy problem, we believe that approximating the full covariance would achieve better performance in high-dimensional problems, albeit at the cost of increased computational complexity.

In equation (10), you still need $H(\tilde{x})$ which can be large as you mention. Do you model it explicitly for training the covariance network or do you model it implicitly using jvp/vjp benefits? This should be discussed. How does it slow down training compared to classical denoising-score matching or other methods?

ANSWER: Thanks for raising this. We compute the Hessian $H(\tilde{x})$ implicitly using the Jacobian-vector product. Specifically, for a given sample, we first evaluate its score using the score network $s(x_t)$, then perform a dot product $s(x_t) \cdot v$ to obtain a scalar. Next, we use auto-diff to compute $H(\tilde{x}_t)v$. Finally, we apply the element-wise product to obtain $v\odot H(v_t)v$. This approach allows us to avoid explicitly computing $H(\tilde{x})$ which significantly reduces the training memory and cost. We will make it clear in the revision.

It is important to note that we only train a small additional network, meaning the training cost is comparable to that of classical denoising-score matching. However, since we need to compute the Jacobian-vector product, the computational cost of our method is higher than that of other baselines. Specifically, we summarize the average wall-clock time per iteration for each method below:

Can the introduction of an additional network (meaning additional approximation) lead to cumulative errors? This is not discussed at all, but the network not being well aligned with the covariance of the trained diffusion will surely lead to poor performances even if the diffusion model was good.

ANSWER:

Thank you for this insightful comment. The introduction of the additional diagonal covariance network is indeed a widely adopted approach in popular diffusion frameworks such as I-DDPM and SN-DDPM. We want to highlight that the introduced network is actually used to learn to reduce the covariance estimation error with the following points:

1. Comparison to DDPM/DDIM: In models like DDPM and DDIM, covariance is typically chosen heuristically to be the lower or upper bound of the true variance. In contrast, learning the covariance through an additional network can reduce the approximation error compared to just using the bound.
2. Comparison to Independent Isotropic Variance: The learned $x_t$-dependent diagonal covariance is significantly more flexible than independent isotropic variance, allowing lower estimation error.
3. Comparison to Other Diagonal Covariance Learning Methods: We show that our proposed objective leads to lower estimation errors compared to other methods for diagonal covariance learning. This reduction in error directly translates to improved generative and likelihood results.

Thus, the additional network is introduced not to add approximation error but to actively learn and reduce it, outperforming other heuristic or less flexible methods. Sorry for the misleading on this, we will make this point clear in the revised paper.

Spotted typos

ANSWER: We sincerely thank the reviewer for identifying the typos, which we will address in the revised version. We deeply appreciate the time and thoughtful feedback provided, as it has significantly contributed to improving our work. We kindly request the reviewer to consider our additional explanations and let us know if they address the concerns raised, potentially influencing the assessment of our work. Please do not hesitate to highlight any other points that you believe may be relevant.

Sorry, I forgot to ask.

What about figure 1? Will you change it or put something more explicit? It's still hard to catch the message without GT comparison. Also I pointed in my Spotted typos a problem with $M$ and $K$. But if I understand correctly, $K$ is the number of timesteps. Then what is the error? The caption of Figure 1 or the $K$ which should be $M$?

From what I see, I imagine that you decrease the noise as you go to the right of the figure. Then it's not with different Rademacher samples but different noise levels. Am I missing something here?

Thanks for your super quick reply. We will include your suggestions in the revision. While we are still working on it, the revised version will be uploaded soon. Thanks again for your suggestion. It has been incredibly helpful in improving our work!

Thanks for the comment. It offers a very insightful suggestion! Our implementation aligns with the architecture you described. As shown in Equation 20, we use a BaseNet as the shared backbone and add an additional head to output $diag(\mathbb{V}[x|x_t])$. The backbone and the mean head $\mathbb{E}[x|x_t]$ are derived from the pretrained diffusion model (see Table 6 in the Appendix for details), with their parameters fixed. We only train the diagonal Hessian head. This architecture also aligns with the design used in the baselines.

As you pointed out, this approach benefits from weight sharing. Additionally, it preserves the original prediction of the mean $\mathbb{E}[x|x_t]$ and the additional inference time is negligible (see Table 4). Furthermore, since the diagonal Hessian head is relatively lightweight, its training cost is even lower than that of training the original diffusion model (see the wall clock comparison in our previous discussion).

Thanks for your revised version!

I have a last question that came to mind.

Why don't you use the same network to output both $\mathbb{E}[x\mid x_t]$ and $diag(\mathbb{V}[x\mid x_t])$ as you would do for classical parameterized Gaussian model? This could have the benefit of sharing weights, reducing even more the computational overhead of the covariance estimation and potentially limiting mismatch between predicted diagonal and diagonal from the score.

However, I agree that an additional network allow to plug your method to an existing pre-trained denoiser, adding external component (the second network) that will not change the original denoiser architecture. I also agree that this could slow down a bit the training if the shared backbone is large. However, at inference, I only see advantages. Can you comment on that, especially for diffusion model trained from scratch?