Energy-Based Modelling for Discrete and Mixed Data via Heat Equations on Structured Spaces

Thank you for your review. We first want to comment on the weaknesses of our paper mentioned in your review:

The authors extend the formalism of Energy Discrepancy to the case of including discrete state in addition to continuous features. Whether or not this justifies an entire publication can be debated

Here is why we hope that this work is interesting to the community: Currently, discrete EBM training methods are hard to tune and don’t translate to mixed state spaces. We want to introduce a simple performant training method to address these limitations.

To do so, we introduce a graphically informed noising process which may also be of great interest in other domains, e.g. discrete diffusion modelling, too. The theoretical understanding of the diffusion allows automated tuning of the time parameter $t$ as a function of the state space size (Theorem 2) which is a critical new contribution to this work substantially different from prior work and important to tabular data, where state space sizes typically vary across between features.

Our theoretical analysis is supported by the experimental results that show good performance, thus unlocking a new tool for density estimation on non-continuous domains.

[...] it is not clear to me if sampling the EBM trained in such way need MCMC to be generative ?

To use EBM as a generative model, one needs to solve two challenges: Training the model and sampling from the model. This paper addresses the training part, only. As you say, sampling methods have to be used to obtain synthetic samples from the model. We employ a fairly standard method for sampling with interleaved updates of the numerical and discrete state. The implementation can be found in the appendix (Algorithm 2).

Importantly, tuning a sampler at inference time is much easier than tuning it for a stable training algorithm. Indeed, we were unable to implement the MCMC-based training for tabular datasets, while sampling from our model with MCMC was unproblematic.

It should be noted that EBMs are particularly useful in non-generative downstream tasks like out of distribution detection and calibration of classifiers, because unlike other generative models they output the unnormalised data density, explicitly. For this reason we focussed on contributing to better EBM training.

They should also check that the trained model is at equilibrium (the generated samples corresponds to the equilibrium properties of the model).

As stated in lines 78-81, ED offers the theoretical guarantee that $p_{\theta^\ast}(\mathbf x) = p_\mathrm{data}(\mathbf x)$ after training converged. Hence, if we run an long-run MCMC sampler on our learned model with Metropolis Hastings acceptance step, the produced samples are at equilibrium by construction (The MH acceptance guarantees detailed balance which is equivalent to samples being at equilibrium.)

The same guarantees does not exist for MCMC-based training approaches because short-run MCMC has to be used in practice introducing biases in the optimisation, see e.g. [1]. For empirical evidence of these favourable properties, see figure 2, 7, 8 which contrast ED and CD results.

More importantly, the comparison for the tabular dataset is only done at the level of the AUC curves. Can at least the authors compare the frequency and correlations amongst the generated samples and the true ones ?

Thanks for your suggestion. We include two new metrics here: single-column density similarity and pair-wise correlation similarity. These two metrics can measure the similarity of the frequency of single columns and correlations of pair-wise columns between the generated and the true table. They can be computed by using the open-source API: SDMetrics. It shows that the proposed ED approaches can outperform baselines or achieve comparable performance in most datasets.

single column  density similarity score

pair-wise correlation similarity score

The authors said that their work is one of the first dealing with tabular data, I at least find this one dealing with EBM and tabular dataset: [...] Also, this article https://arxiv.org/abs/2405.15376 deals with generative and categorical variables for EBM.

Thank you for these interesting references! We missed the first reference in our literature research. Unfortunately, we are using different benchmark data sets so a direct comparison is difficult. Furthermore, RBM architectures tend to be more constrained than deep NN architectures, which has the advantage of stabilised training at the cost of flexibility. This means that ED and the referenced work have slightly different advantages and drawbacks.

The second article is very interesting. It is concurrent work to ours and was published after the NeurIPS deadline, so we were not aware of it.

[1] Nijkamp et al. Learning non-convergent non-persistent short-run MCMC toward energy-based model

[...]For small timesteps, the loss function seems very closely related to pseudolikelihood estimation [...]

Thank you for raising this point. It seems that pseudo likelihood can be derived from energy-discrepancy as follows: Define $q(\mathbf y \vert \mathbf x) = \frac{1}{d} \sum \delta( y_i, \ast)$ which masks exactly one entry of the data vector. Then, energy discrepancy is given for a sampled perturbation $y = x_{\neg I}$ as

$$U_\theta(\mathbf x) + \log \sum_{x} q(\mathbf x_{\neg i} \vert \mathbf x) \exp(-U_\theta(\mathbf x)) = -\log \frac{\exp(-U_\theta(\mathbf x))}{\sum_{s \in \{1, 2, \dots, S\}}\exp(-U_\theta(\mathbf x_1, \dots, \mathbf x_i = s, \dots, \mathbf x_d))}= -\log p_\theta(\mathbf x_i \vert \mathbf x_{\neg i})$$

Hence, this specific ED loss function is indeed a MC approximation of pseudo-likelihood. Energy discrepancy is appealing because it is more general and more tunable through the choice of $t$ and $M$. We will discuss this connection in our revised version.

Since these methods have mixed continuous and discrete data can the authors do a separate benchmark that only models the discrete columns?

The performance on continuous data is compared in prior work [9]. In this work, we benchmark ED on various discrete data sets. For a separate benchmark on discrete/continuous variables in tabular data we are missing baselines that perform the same experiment, and an implementation of comparative experiments was not possible in the given time frame.

I was confused by the statement that method is sensitive to the assumption that the data distribution is positive on the whole space. Why is this more of an issue for ED than other EBM training techniques? [...]

Other training techniques for EBMs are adaptive, i.e. the negative samples are created by attempting to self-sample from the model. If the sampler is well-tuned, this produces smoothed estimates of the data distribution (see e.g. Figures 2, 7 to see the smoothing effect compared to ED). Thus, these methods produce biased estimates which can be advantageous when training on image data, where the data support is small compared to the ambient space. ED produces more accurate estimates, but tends to oversaturate quickly when the diffusion produces uninformative negative samples.

I would also like to see more comparisons of CD with large numbers of MCMC steps v.s. ED with large numbers of samples.

• Table 2 reports results for CD with 40 MCMC steps. CD with larger numbers of MCMC steps is rarely used due to the cost of training.
• Similarly, ED with large number of samples would loose the key advantage of ED of being cheaper to compute than CD or MLE with MC approximated normalisation constant.

A major difficulty with increasing the number of MCMC steps in CD is that the sampler likely needs to be retuned. For all these reasons, we chose to compare CD with a typical number of MCMC steps (e.g. 40 steps) and ED with M = 32, i.e. maximising the parallelisation capabilities of our GPU.

Finally, here are the references mentioned in the rebuttal and in the following comments.

[1] Du et al. Compositional Visual Generation with Energy Based Models, NeurIPS 2020

[2] Glaser et al. Maximum Likelihood Learning of Unnormalized Models for Simulation-Based Inference

[3] Grathwohl et al. Your Classifier is Secretly an Energy Based Model and You Should Treat it Like One, ICLR 2020

[4] Grathwohl et al. Oops I took a Gradient: Oops I Took A Gradient: Scalable Sampling for Discrete Distributions, ICML 2021

[5] Xu et al. Modeling Tabular Data using Conditional GAN, NeurIPS 2019

[6] Kotelnikov et al. TabDDPM: Modelling Tabular Data with Diffusion Models, ICML 2023

[7] Campbell et al. A Continuous Time Framework for Discrete Denoising Models, NeurIPS 2022

[8] Lou et al. Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution, ICML 2024

[9] Schroeder et al. Energy Discrepancies: A Score independent Loss for Energy-Based Models, NeurIPS 2023

[10] Lyu S. Unifying non-maximum likelihood learning objectives with minimum kl con- traction., NeurIPS 2011

[11] Gao et al. Learning energy-based models by diffusion recovery likelihood. ICLR 2020

[12] Luo et al. Training Energy-Based Models with Diffusion Contrastive Divergences

[13] Lee et al. GUIDING ENERGY-BASED MODELSVIA CONTRASTIVE LATENT VARIABLES, ICLR 2023

[14] Foster et al. A Unified Stochastic Gradient Approach to Designing Bayesian-Optimal Experiments, AISTATS 2020

[15] van den Oord et al. Representation Learning with Contrastive Predictive Coding

[16] Le Cun et al. A Tutorial on Energy-Based Learning

Thank you for your thorough review and apologies that we could not fit all responses in the rebuttal.

We now want to respond to the questions left out in the rebuttal. We follow the chronological order of your review and the rebuttal.

It was unclear why small timesteps were introduced and why Euler integration was needed.

We don't use the Euler discretisation in practice and our method is simulation-free. For the uniform perturbation, the exact solution is given by

$$q_t(y \vert a) = \left(\frac{1}{S} + \frac{S-1}{S}e^{-St}\right) \delta(y, a) + \left(\frac{1}{S} - \frac{1}{S}e^{-St}\right) \sum_{k \neq a}\delta(y, k)$$

This is equivalent to equation (7) via the time rescaling $t \leftarrow (1/S(1-e^{-St}))\in [0, 1/S)$ We are going to replace equation (7) with the closed form solution.

[...] since the noise process is applied independently across the dimensions, it should scale with O(S^3). [...] Was the point that for some problems S^3 is too big [...]?

Thank you for this comment! We missed that one could diagonalise any graph Laplacian in $O(S^3)$ time before running the training which extends the generality of our suggested method significantly. We consider this for the revision of this work.

[...] the authors say that adding uniform noise [...] “introduce too much noise to inform the EBM about the correlations in the data set.” Can the authors make this statement more precise in the text? [...] it is not clear why this isn’t solved by taking smaller timesteps [...]

The intuition is the same as our earlier response on informative perturbations. We will revise the statement in the paper. The grid perturbation always changes exactly one entry of the data vector which is the main different to the uniform perturbation with small $t$. We observed empirically that sometimes this is beneficial, e.g. for the binary image data the grid perturbation outperforms the heat perturbation (a Bernoulli distribution) with a small time step. For tabular data, one can also condition the negative samples on the column that was perturbed, so that the parameter update is only informed by the change of a single column. We have not quantified the benefit of this analytically.

For the related works, I felt the “Contrastive loss functions” paragraph needs more discussion. [...]

We will revise our discussion of related contrastive loss functions. Without giving proofs, we want to give a brief summary of the connections we see:

The following papers are based on the same functional as energy discrepancy:

• ED is equivalent to KL contractions [10]. However, [10] does not turn the KL contraction into a training loss, and does not explore the influence of the perturbing distribution.
• ED for Gaussian perturbation is equivalent to the diffusion recovery likelihood functional [11]. However, [11] uses a large number of noise scales, an MCMC scheme, no stabilisation, and only one negative sample to produce a training algorithm, so methodologically it sits between ED and CD training approaches.
• ED is equivalent to the concurrent work Noise Contrastive Divergence [12]. However, [12] approximates the loss with a variant of score matching, thus turning it into a different loss in practice.
• ED is strongly related to the concurrent work Contrastive Latent Energy Learning [13]. Again, the loss function implemented in practice by CLEL is very different as it only considers a single negative sample obtained from a latent code.
• If the perturbation satisfies the detailed balance relation, ED becomes equivalent to the contrastive divergence loss [9].
• As remarked by your review, pseudo-likelihood methods can be seen as a special case of ED.

The following works discuss a similar training objective as ED:

• The stabilised ED loss is equivalent to the prior contrastive estimation (PCE) bound [14] because both losses are based on approximations of KL divergences. However, PCE is not used to learn a probabilistic model but to find an optimal Bayesian experimental design, a very different learning task.
• Similarly, the stabilised ED is structurally similar to InfoNCE [15], a loss used in representation learning.
• The stabilised ED is related to the log loss [16]. However, the log-loss was introduced without connection to learning the EBM as a probabilistic model, introducing meaningful ways to obtain contrastive samples, or providing theoretical guarantees, and we could not find experimental results for this loss.

It should be noted that none of these references apply to discrete or mixed data like our work.

Thank you for your review and your questions! Due to lack of space we focus on questions of most interest for all reviewers. Citations and responses to the rest can be found in the comment.

Significance: [...] Are there better baselines to compare against [...]?

EBMs are appealing because they output the unnormalised data density explicitly, and allow more flexible modelling than e.g. normalising flows. This unlocks a unique set of downstream applications, see [1, 2, 3].

Given our interest in using EBMs as a tool in downstream tasks, we are not trying to claim SOTA compared to all generative models. For binary discrete image data, our method is competitive with other EBM training approaches like [4]. For tabular data we compare with [5] due to lack of EBM baselines. Prompted by your comment we found that a tabular DDPM gives a stronger baseline [6] and we will add this reference to our table.

Questions

[...] What are examples of processes [q] that are more and less informative?

In principle, the negative samples can be drawn from an arbitrary noise distribution. If, for example, $q$ is independent of the data, the contrastive term becomes just the log-normalisation constant up to a constant, i.e.

$$\log \sum_{ x_-\in X} q( y) \exp(-U_\theta( x_-)) = \log q( y) + \log \sum_{ x_-\in X} \exp(-U_\theta( x_-))$$

which is an example of a non-informative $q$ distribution. In this case, ED becomes the log likelihood in theory, but the normalisation constant $Z_\theta$ needs to be approximated from thousands of samples to be low-variance.

An informative noising distribution $q$ achieves two things.

• $\sum_{ x_-\in X} q( y\vert x_-)\exp(-U_\theta( x_-))$ can be computed from a small number of samples with lower variance than the full normalisation $Z_\theta$ because the integrand is localised around $y \sim q( y \vert x_+)$.
• similar to GAN training, the model converges prematurely if the negative samples $x_-$ are trivial to distinguish from data $x_+$.

I am confused about the connections to the heat equation [...] It was unclear why small timesteps were introduced and why Euler integration was needed.

We don't use the Euler discretisation in practice and our method is simulation-free, i.e. for any time $t$ we can sample exactly from the transition density $q_t(b|a)$ using the closed form solution in Proposition 1. For the uniform perturbation, the exact solution is equivalent to equation (7) via the time rescaling $t' := (1/S(1-e^{-St}))\in [0, 1/S)$, (see comments). We are going to replace equation (7) with the closed form solution.

[...] Is there something that the heat equation view is really buying us? [...] “why not just use large times since this is maximum likelihood?”

Any CTMC could be used to define a noising process, but this leaves us with many choices for the rate function. In previous work using CTMC [7, 8], the structure on the discrete space is ignored and only the uniform and absorbing diffusion is discussed.

We choose the rate function as the graph Laplacian which defines the smallest possible transitions in the graph. For small $t$, the corrupted data points are then with highest probability the nearest neighbours of the data point in the graph.

However, a small corruption may miss global distribution properties like mixture weights [9]. Theorem 1 tells us that for large $t$, ED performs MLE which captures these global properties at the cost of higher variance of the estimated loss. The heat equation allows us to trade-off the variance and better statistical properties of the loss.

To support our intuition empirically, we include new experimental results in the attached PDF, Figure 1, showing that a large $t$ requires increasing $M$ to achieve accurate energy estimates.

The authors make a point of saying that ED is much more efficient that CD and point to timing experiments in the appendix. However, it appears that authors are only reporting timing for M=4 samples when in the paper M=32 is used. [...]

$M=4$ is often sufficient to obtain good results, see table 8. For the main text of this paper we used $M=32$ to maximise our GPU usage. In fact, ED can be computed in parallel, while CD requires sequential sampling algorithms that cannot be parallelised. We provide the training time for $M=32$ below:

It shows that ED with $M=32$ is still more performant than CD-5. We will change the corresponding table in the revision.

[...] 1) since the baseline methods were published in 2019 are there more sensible baselines to compare with? [...] 2)

As mentioned earlier, there is a better baseline that uses a diffusion model [6] which outperforms our approach on the tabular data benchmark. We provide the experimental results for detailed comparison:

TabED refers to the best results from our proposed methods. We remark that, while the DDPM outperforms the GAN and EBM, we are interested in tasks that cannot be addressed by diffusion models.

I would like more honest discussion throughout the paper of the relative strength and drawbacks of their method. [...]

We will revise the paper keeping your comment in mind. In brevity, the biggest strength of ED are improved theoretical guarantees over CD or robust behaviour when tuning an MCMC sampler is difficult (e.g. tabular data). The biggest drawback of ED is currently its scalability to data more high-dimensional than MNIST.

Thank you for your review and your helpful comments.

Despite the method's solid contributions and experimental design, the motivations behind each step and their presentations are not very clear, making it hard to follow. For instance, in Section 3.1, the paper discusses different structured and unstructured categorical values, introducing the four types {cyc, ord, unif, abs}. However, it is not clear why these specific structures are chosen. Are they meant to cover all categorical values comprehensively, or are they the most common in tabular data? Providing a clearer rationale would help readers understand the choices made.

Thank you for letting us know that the presentation is not clear in this point, so we would like to provide additional background on this topic: To create the Energy Discrepancy loss we have to introduce a noising process that produces contrastive negative samples which the model can learn from. In principle, the noising processes {unif, abs} already cover all categorical values and this is the approach taken for discrete diffusion models, see e.g.

[1] Campbell et al. A Continuous Time Framework for Discrete Denoising Models, NeurIPS 2022

[2] Kotelnikov et al. TabDDPM: Modelling Tabular Data with Diffusion Models, ICML 2023

[3] Lou et al. Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution, ICML 2024

In this work, we observe that many discrete types of data have additional structure such as seasonality/periodicity (e.g. month of year) or ordering/grading (e.g. age, RGB pixel values, salary) that give us more fine-grained control over the noising process than the diffusions in [1, 2, 3]. This generalises previous diffusions while remaining simulation-free, i.e. we can compute and sample from the noising distribution for any $t$ in closed form.

The intuition for respecting the structure of the underlying discrete space is that we seek negative samples that are close to the data support to provide a good training signal (This is similar to the discriminator training in GANs).

If the discrete space of interest has a more complicated graph structure, the process either needs to be simulated from eigen decompositions of the graph Laplacian or one can resort to the simpler uniform or absorbing perturbation. Since all our benchmark data sets didn't require an involved graphical structure, we stick to the more common ordinal or cyclical structure which generalises previous noising processes [1, 2, 3].

The scalability of the proposed method in such scenarios is a significant concern. An analysis or discussion on how the method handles large categorical values would be beneficial. This could include potential modifications or considerations to ensure that the method remains efficient and practical when applied to datasets with large categorical variables.

Thank you for your suggestions. Could you elaborate what you mean by large categorical values in case we misunderstand your question?

Scalability is indeed a valid concern for energy-based models in general. Most scalable continuous EBMs such as [4], which trains EBMs on image net 32x32, have other problems, e.g. over-smoothing the estimated densities due to missing theoretical guarantees, and are limited to the continuous domain. This limits their applicability in some key downstream tasks like reliable out of distribution detection.

For discrete data sets, alternative training methods for EBMs produce comparable results to our method. Compared to MCMC based training methods like contrastive divergence [5], the energy discrepancy approach requires less tuning and runs faster since the loss can be computed in a single parallel pass through the network and does not rely on sequential steps. Furthermore, these works have yet to be extended to mixed state spaces, and many tabular data sets of interest have a comparably small number of features.

We are also interested in better scalability in future work. However, this will likely require a combination of EBMs with other generative modelling approaches such as a variational decoder model as in [6]. This method has already been established for the energy discrepancy training method. Secondly, the energy function may be learned with the energy discrepancy loss at various noise scales as in [7], thus improving training stability and learning a sampler alongside the EBM. Finally, EBMs may be an interesting way to modify a pre-trained base model as in [8]. While ED can likely be used to modify [6, 7, 8], this is not the main contribution of our work.

[4] Du et al. Implicit Generation and Generalization in Energy-Based Models, NeurIPS 2019

[5] Grathwohl et al. Oops I took a Gradient: Oops I Took A Gradient: Scalable Sampling for Discrete Distributions, ICML 2021

[6] Pang et al. Learning Latent Space Energy-Based Prior Model, NeurIPS 2020

[7] Gao et al. Learning Energy-Based Models by Diffusion Recovery Likelihood, ICLR 2021

[8] Deng et al. Residual Energy-Based Models for Text Generation, ICLR 2020

What’s more, I strongly recommend moving these algorithms from the appendix into the main body of the paper. This would make the paper easier to follow and more accessible to readers who need to understand the detailed workings of the method.

Thank you for this suggestion. We will include the algorithms in the main text for the revision of this paper.

Thank you for your helpful comments and questions.

Although the energy discrepancy method has already been proposed and published in previous work, I found the justification for the method slightly confusing while reading this paper. What is Theorem 1 exactly saying? (see questions)

How should I interpret the left side of the equation in Theorem 1, and the fact that the right side approaches zero with large enough t? How does this link ED to maximum likelihood, exactly?

We want to address both your question and the weaknesses mentioned in your review. Theorem 1 is meant to guide the hyperparameter choice of $t$. Another way of stating Theorem 1 is as follows:

$$\mathrm{ED}{q_t}(p_\mathrm{data}, U_\theta) - \mathcal L_{\mathrm{MLE}}(\theta) - z_t = \mathcal O(e^{-\lambda t})$$

Here, $z_t$ is a constant independent of $\theta$, so the optimisation landscapes of energy discrepancy estimation and maximum likelihood estimation in $\theta$ align at an exponential rate, except for a shift $z_t$ which does not affect the optimisation.

Practically speaking, this property implies that large time parameters yield losses with favourable statistical properties as the MLE has the lowest variance among all unbiased estimators. The most striking example of this is global mode awareness, i.e. as shown in prior work [1], energy discrepancy losses are capable of estimating the mixture weight distributions in multi-modal data sets while being more tractable than maximum likelihood estimation. However, this comes at the cost of higher variance of the loss function due to the sample approximation used to compute the loss function in practice.

Since arbitrarily large $t$ are impractical, the parameter needs to be tuned. For this reason we study the opposite limit $t\to 0$ in the following sections. Similar to the late stages of GAN training, our goal is that negative samples obtained from the noising process remain close to the data support to provide a good training signal. The infinitesimal analysis relates the properties of ED for $t\to 0$ to the geometry of the underlying discrete space. The perturbed samples produced are then with highest probability the nearest neighbours of the data point in the graph. This makes it possible to get fine-grained control about the properties of the ED loss with a single time parameter.

While these results guide our understanding of the influence and tuning of $t$ in the algorithm, Theorem 1 is not needed to justify ED as a loss function. Instead, the justification was given in prior work and is reiterated in lines 78-81: One can show that after convergence, for a sufficiently rich family of neural networks and arbitrary amounts of data, $p_\theta(\mathbf x) = p_\mathrm{data}(\mathbf x)$ which holds for all $t$.

The loss also is, in practice, approximated with something slightly different than the proposed loss, which seems conceptually a bit confusing. However, this is not a major concern given that the base method has been proposed and published in previous work.

The loss used in practice can be justified by observing that for any $w>0$:

$$\lvert \mathcal L_{w, M, t}(\theta) - \mathrm{ED}{q_t}(p_\mathrm{data}, U_\theta)\rvert \xrightarrow{M, N\to\infty} 0$$

Thus, the estimated loss function is asymptotically unbiased, and empirically we observe that the loss performs similarly to our expectations from the theory. It should be noted that most EBM training approaches need some form of approximation and stabilisation. MCMC-based training of deep EBMs usually use short run MCMC, introducing biases into the derived parameter update, and use a regulariser $\mathcal L_{\mathrm{reg}}(\theta) = U_\theta(\mathbf x_+^i)^2 + U_\theta(\mathbf x_-^i)^2$ which alters the desired MLE loss. In comparison, ED is easier to justify, and indeed we observe better performance in most applications.

To further support our theoretical analysis empirically, we conduct discrete density estimation experiments with different $w,M$ (see Figure 2 in the attached PDF for illustration). Despite using varying $w$, our approach can still learn a faithful energy landscape using a sufficiently large $M$, verifying the theoretical analysis that $\mathcal L_{w, M, t}(\theta) \rightarrow \mathrm{ED}{q_t}(p_\mathrm{data}, U_\theta)$ when $M, N \rightarrow \infty$.

What is Avg. Rank in Table 1?

We rank each of the 7 methods in table 1 according to their AUC scores from 1 to 7. The average rank is the average ranking of the method across the used data sets.

[1] Schroeder et al. Energy Discrepancies: A Score independent Loss for Energy-Based Models, NeurIPS 2023