Learning normalized image densities via dual score matching

Thank you for your review, which will help us to improve the paper. We ran several experiments to address your questions, which we answer below.

Other datasets: We applied our model to CelebA images at 80x80 resolution. Our model achieves a NLL of 1.94 bits/dim. We are not aware of reported NLLs on CelebA (ImageNet64 has been a de facto standard), but it is reasonable that the value is lower than for ImageNet64 due to reduced diversity in the images. Additionally, we verified that the histogram of log probabilities (as in Figure 3) is well-approximated by a Gumbel distribution (apart from a few outliers on the low-probability edge, due to corrupted images). Over the 20k test images, the range from highest to lowest probabilities is 17dB/dim. This is half that of ImageNet64, which is also reasonable given the lower diversity in image content. We also ran our method on CIFAR-10, obtaining an NLL of 3.16 bits/dim, and a similarly-skewed log p distribution (with a range of values spanning 23 dB/dim). We hope that these additional results alleviate the reviewer’s concerns. We believe that most photographic image datasets should lead to qualitatively similar results regarding the experiments of Section 3.

Downstream applications of energy-based models: OOD detection based on probability values is notoriously difficult, due to the counterintuitive behavior of probability densities in high-dimensional spaces. As we show in Figure 3, the highest-probability images are not typical, and could be considered outliers (this is particularly apparent in the highest-probability image, which is blank). The following paper demonstrates that generative models can assign higher probability density values to images from other datasets (e.g., MNIST) than to those in their training set (e.g., FashionMNIST): Nalisnick, E., Matsukawa, A., Teh, Y. W., Gorur, D. & Lakshminarayanan, B. Do Deep Generative Models Know What They Don’t Know? (2019).

Nonetheless, probability values can be used to detect certain degradations such as additive noise. We verified that adding noise to the reference images in Figure 3 decreases their probability (this experiment is similar to the third panel of Figure 5, except that the value of $t$ given to the model is 0). In addition, one application of our model is that it can be used to estimate the variance of the noise $t$ added to an image $y$ when it is unknown. Indeed, the value of $t$ that maximizes $p(y|t)$ provides an accurate estimate of the true value of $t$. We will incorporate figures demonstrating these two experiments in the paper. Finally, we highlight that our analysis of the learned energy model revealed novel statistical and geometric properties of image densities learned from data, which we believe to be of general interest to the NeurIPS community.

On theoretical guarantees of denoising score matching: We briefly touch on this subtle point in Appendix A.3. Obtaining theoretical guarantees requires establishing a Poincaré inequality on the joint distribution $p(y,t)$, which Figure 1 suggests could be true, but such inequalities are notoriously challenging to prove and have only been obtained in restricted settings. However, even if this inequality was shown to hold, in practice there is no guarantee that the trained network generalizes. Proving that this is the case is an even bigger open problem. So one always needs to resort to empirical verifications, which are complicated by the fact that the “true probability density” is not known for images. This is the goal of the strong generalization test of Figure 2, which shows that energy models trained on distinct images compute nearly identical probabilities. Our energy-based model is the first for which such a stringent test has been performed. (We also conduct a more direct verification on synthetic data in Figure 1.) We believe these tests to give as much credibility as can be realistically obtained to the other empirical findings of Figures 3 to 5.

Thank you for the detailed review, which will help us significantly improve the paper.

Clarifying contributions: We believe that the primary contributions of the paper are twofold: A combination of architecture and objective that allow efficient training of an accurate normalized energy model. We model the energy directly (rather than the score) so that energy computation is done in a single forward pass. The purpose of the time score objective is to ensure that the energy values are accurate and normalized (see below). They go hand in hand: in our setting, time score matching is meaningless for score models (see below). Analysis of the model (trained on ImageNet) to elucidate properties of natural image densities. We believe these observations to be novel and of value to the community (as mentioned by Reviewer hMws). These two contributions are not disconnected. While most of the experiments in 2. could in principle be done with any generative model that can compute probabilities, they are much more easily done when the log probability is explicitly available with a single forward pass. Furthermore, the dimensionality measurements require access to the distribution of noisy images. Our energy model is the only one that provides access to the density of clean or noisy images in a single forward pass. We will incorporate these points into the introduction.

Novelty: Beyond our combination of energy architecture and dual score matching objective which is novel (see below for time score matching), we are not aware that any of the results in Section 3 have appeared in the literature. Could the reviewer provide references that establish similar results?

Purpose of time score matching: While the time score is not used at test time to evaluate the energy, it plays a crucial role during training to ensure that the learned energy (rather than the “space” score) is accurate. This is demonstrated in Figure 1: the energy model trained with single score matching does not learn the ground truth energy, while dual score matching succeeds. The reason is that single score matching does not sufficiently constrain the values of the energy when the data is multimodal: one can add different constants to the energy in disconnected modes without affecting the score. Matching time scores fixes these unwanted degrees of freedom. This is because the joint distribution $p(y,t)$ is unimodal, contrary to $p(x)$, as visualized in the right panel of Figure 1. Matching its score, which contains both the “space” and time components, thus constrains the energy up to a single additive constant, as opposed to one per mode of $p(x)$. We will clarify this central point in the paper, which was not sufficiently explained.

Novelty of time score matching: While time score matching has been proposed before, our particular objective is novel. In Choi et al., the time score matching objective uses a second derivative in time, which is more computationally expensive than our regression-based objective. The conditional expectation identity (4), and the resulting objective (5), are novel compared to this prior work. However, after submission of the paper we became aware of the concurrent work “Density Ratio Estimation with Conditional Probability Paths” (Hanlin et al., 2025) that derives the same identity and objective. We will incorporate citation and discussion of this into our paper.

NLL calculation: The reviewer is mistaken. The fact that our NLL calculation is fast (a single forward pass of the model $U_\theta$) is specifically because of our choice parameterization and objective (this was the motivation for these choices). Our energy model can compute the probability of 50k images in 12s on an A100 GPU, whereas a diffusion model requires upwards of 3h20 as it requires at least 1000 forward (+ backward) passes per image (assuming 100 time steps and 10 noise samples for computing the divergence with Hutchinson trick or the MSE as explained in Appendix A.2). Appendix A.2 provides a literature review of other methods of estimating energies, not a description of our method. We do not rely on any algorithm from Karczweski et al. (who do not consider EBMs).

Normalization: Normalization is computed only once post-training, by calculating the constant $c = \mathbb E_y [U_\theta(y,t) | t=t_{\max}] + \frac{d}{2}\log(2\pi e t_{\max})$. Test-time normalized energies are then obtained with a single forward pass by calculating $U_\theta(y,t) - c$.

Conservative diffusion models: The goal of the paper is not to obtain better diffusion models by imposing them to be conservative. Rather, our primary goal is to learn the energy, a byproduct of which is a conservative score model. We merely point out that due to our choice of architecture, there are virtually no differences between the original score network and the gradient of the learned energy.

Computational drawback: We do not understand the question. We mentioned the limitation of the added computational cost of the double backpropagation in the main text. To be more precise, the training time for a score model for 1M steps on ImageNet with an A100 GPU was 1d8h, while that of our energy model was 5d. We will mention this in the paper and apologize for the imprecise statement.

Approximations: As in any empirical work, there are approximations involved. We will update the text to make these more explicit. It is common practice in energy-based models to resort to approximate normalization as exact normalization is typically impossible (continuous integrals can only be calculated by discretization unless a closed form is available). The purpose of Appendix A.3 is not to prove theoretical guarantees (we say this explicitly) but to provide a sketch of what would be needed to establish them for future research.

Terminology: We apologize for the confusion. As an abuse of language, we called the gradient of the energy the score rather than negative score. Thus, there is no typo in eqs (2) and (3): $U$ is the negative log probability, not the log probability (see equation (1)). However, we disagree with the second remark on terminology. Log probability is equal to the negative energy up to an additive constant $\log p(x) = -U(x) -\log Z$, not the normalized exponentiated negative energy. We used “normalized energy” to refer to $-\log p(x)$, a term that for instance is used in the review by Song and Kingma (2021).

Thank you for the author’s timely rebuttal. It was overall quite helpful, I have to say. There were several things unclear in the original paper, many of which were clarified and I now have a better idea of the proposition.

The good parts:

• The clarification regarding the “purpose of time score matching” was the most helpful. And, as I can see you realised too that, it wasn’t adequately explained.
• The clarification on “NLL calculation” was also quite helpful. I think you should consider re-explaining this properly in the paper. Specifically, the comparison with the “score model” style NLL computation. In your case, it is almost trivial to compute NLL because what you are modelling is pretty much an approximation of NLL itself.
• Thank you for clarifying the terminology. This was specifically confusing to me because, throughout the paper, you used both terms “normalised energy” and NLL in different places. I could not comprehend that they are same thing (as you clarified in the rebuttal). I say more about this below.

The still not-so-good parts:

• Using the term “normalised energy”, IMO, is quite confusing, specially when you are referring to simply NLL — why not just say NLL ! Because, there can be multiple ways to normalise energy $U(\cdots)$. Only one specific normalisation leads to NLL. I can come up with an arbitrary value of $c = c_0$ for which $U(\cdots) - c_0$ refers to the same underlying probability and score.
• The statement that you are “learning normalised energy model” is also confusing. You are not learning normalised model directly. You are simply learning the “normalised energy model” and then deriving a normalisation constant. Moreover, the constant is not theoretically guaranteed to be exact — its an approximation too. By your logic, one can claim that score based method also learn a “normalised energy model” because NLL can be derived from the learned score model too (yes, its expensive to compute, but possible).
• Finally, it is clear that the model training takes a lot longer. Typically, in practice, such double-backpropagation is not very popular as a solution. So, what you are getting isn’t really for free.

Overall, I feel this paper may be a good analytical tool. But, considering computational disadvantage, I am not sure how much actual impact it can make. Also, the paper lack significant quantitative evaluation on downstream application; instead focuses on low-level analysis.

Thank you for the kind comments! We especially appreciate your positive evaluation of our analysis of the learned energy-based model. We answer your questions below.

Clarification of Figure 1: Single score matching fails on multimodal distributions because the objective does not constrain the absolute energy levels of separated modes. Indeed, one can add an arbitrary constant to the energy that is different in each mode without changing the score in each mode. Of course, this requires “stitching” together the different constants, which adds a term to the score, but the transition can occur in-between the modes where there is no data to constrain the scores. We hope this clarifies the issue: we will update the text to be more explicit.

Clarification of Figure 3: The purple-to-yellow-colored arrows in Figure 3 correspond to the images on the right of Figure 3 (as mentioned in the caption), not the figures in Appendix. The brown and green arrows respectively correspond to a uniform noise image in [0, 1] and a constant image of intensity 0.5 (not shown).

Code: We will publish code to reproduce all experiments, as well as trained model checkpoints.

Scaling to higher-resolution images: The main bottleneck in increasing image resolution is the memory cost of double backpropagation. This is shared by all energy models that are trained with a score-based objective. It can be alleviated with standard techniques such as sliced score matching (Song et al., 2019), which here corresponds to projecting the DSM residual $\nabla_y U_\theta - (y-x)/t$ on a random direction $v$ and minimizing its square. This enables the use of forward-mode auto-differentiation to compute the scores, as only inner products $\langle \nabla_y U_\theta, v\rangle$ are required. This can be further combined with auto-differentiation software optimizations. For instance, the ICLR blog post “How to compute Hessian-vector products?” (Dagréou et al., 2024) reports up to 3x memory reduction with the use of optimized implementations (JIT compilation, etc) in the related setting of Hessian-vector products.

Thank you for your review. We have run several experiments to answer your questions, which will help us to improve the paper.

Computational complexity of the training process: Training an energy model with DSM+TSM losses converges in the same number of iterations as for DSM loss alone. The additional time complexity of the TSM loss is negligible: real world training time is 4d16h for DSM and 5d for DSM-TSM, for 1M steps on an A100 GPU.

Higher-resolution images: We have now applied our model to CelebA images at 80x80 resolution, and find that it achieves a NLL of 1.94 bits/dim. We are not aware of reported NLLs on CelebA nor other higher-resolution datasets (ImageNet64 has been a de facto standard), but it is reasonable that the value is lower than for ImageNet64 due to reduced diversity in the images. Additionally, we verified that the histogram of log probabilities (as in Figure 3) is well-approximated by a Gumbel distribution (apart from a few outliers on the low-probability edge, due to corrupted images). Over the 20k test images, the range from highest to lowest probabilities is 17dB/dim. This is half that of ImageNet64, which is also reasonable given the lower diversity in image content. Finally, we mention that scaling dual score matching to higher-resolution images is bottlenecked by the memory cost of double backpropagation. This can be alleviated with standard techniques such as sliced score matching (Song et al., 2019), which is beyond the scope of this work.

Comparison with other EBMs: We will add a comparison with the Routing Transformer (Roy et al., 2021) that achieves a NLL of 3.43 bits/dim (compared to our value of 3.36). But we were unable to find an energy-based model (i.e., a model which directly learns the energy) that reports NLL on ImageNet64 - all models we examined were auto-regressive, flow-based, or compute variational lower bounds, while older models consider lower-resolution datasets. We also were unsure what the reviewer meant by “training stability”? Our training procedure is quite stable, as it relies on two regression objectives. In addition, these objectives have the same minimum, so that there are no adversarial instabilities (as in e.g. GANs).

Normalization of losses with dimensionality: The factors $1/d$ and $1/d^2$ in equation (8) set the relative scales of the two objectives. This is equivalent to minimizing $t \ell_{\mathrm{DSM}} + \lambda t^2 \ell_{\mathrm{TSM}}$ with $\lambda = 1/d$.

Extension to conditional models: The proposed method is straightforward to extend to conditional generative models. Given conditioning information $c$, the network computes the conditional energy function $U(y,t,c) = -\log p(y|t,c)$. The expression of the DSM and TSM losses are unchanged, as well as the energy architecture $U_\theta(y,t,c) = \frac12 \langle s_\theta(y,t,c), y\rangle$. We will mention this in the paper, and we thank the reviewer for the suggestion.

FID: The generative abilities of the score $\nabla_y U_\theta$ derived from our energy model $U_\theta$ are equal to that of the base score network $s_\theta$. FID thus measures the performance of the base diffusion generative model, which is orthogonal to the focus of this work (energy-based models).

Contributions: We would like to note that the reviewer summary omits several aspects of the paper that we believe are important contributions, regarding analyzed properties of the learned density model. We show that images with highest vs lowest probability have starkly different visual features, and that the high-probability images are not the most typical images, as intuition might suggest (Figure 3). We also show that the support of some (but not all) natural images can be locally described as a low-dimensional manifold. We believe these observations to be novel and of value to the NeurIPS community (as mentioned by Reviewer hMws).