Energy-Based Diffusion Language Models for Text Generation

## Abstract

Abstract Despite remarkable progress in autoregressive language models, alternative generative paradigms beyond left-to-right generation are still being actively explored. Discrete diffusion models, with the capacity for parallel generation, have recently emerged as a promising alternative. Unfortunately, these models still underperform the autoregressive counterparts, with the performance gap increasing when reducing the number of sampling steps. Our analysis reveals that this degradation is a consequence of an imperfect approximation used by diffusion models. In this work, we propose Energy-based Diffusion Language Model ( EDLM ), an energy-based model operating at the full sequence level for each diffusion step, introduced to improve the underlying approximation used by diffusion models. More specifically, we introduce an EBM in a residual form, and show that its parameters can be obtained by leveraging a pretrained autoregressive model or by finetuning a bidirectional transformer via noise contrastive estimation. We also propose an efficient generation algorithm via parallel important sampling. Comprehensive experiments on language modeling benchmarks show that our model can consistently outperform state-of-the-art diffusion models by a significant margin, and approaches autoregressive models’ perplexity. We further show that, without any generation performance drop, our framework offers a 1.3 × \times × sampling speedup over existing diffusion models.

## 1 Introduction

In recent years, autoregressive (AR) models have achieved remarkable advances in modern large language modeling (Vaswani et al., ; Wolf et al., ; Brown et al., ; Chowdhery et al., ; Touvron et al., ), showing impressive results in tasks such as expert-level coding (Roziere et al., ) and reasoning (Wei et al., ; Trinh et al., ). Despite the considerable progress achieved, the left-to-right generative paradigm exhibits several long-standing drawbacks. For example, training and sampling of autoregressive models require a fixed ordering of the sequence, constraining their generation flexibility (Hoogeboom et al., ; Shih et al., ). Furthermore, at test time, the generation is heavily conditioned on previous generations, which leads to accumulated error a.k.a exposure bias (Bengio et al., ). As a result, alternative generative paradigms beyond left-to-right generation are actively being explored.

Discrete diffusion models (DMs) have recently emerged as a promising competitor for discrete sequence generation (Austin et al., ; Dieleman et al., ; Gat et al., ). Unlike AR models, discrete DMs conduct generation by progressively decoding the full sequence in parallel starting from a fully masked sequence, offering great potential in bidirectional controllable generation and sampling acceleration. Most recently, great efforts have been invested in improving their performance by extending discrete diffusions to operate in continuous time (Campbell et al., ; ; Sahoo et al., ; Shi et al., ).
Unfortunately, despite their promise, these models still underperform AR models by a large margin, resulting in limited practical usage.
Developing principled and powerful discrete DMs remains an open problem.

In this work, we take a closer look at the existing discrete diffusion models and reveal a vital mismatch issue between the training and sampling distribution in its current form. Specifically, the existing discrete diffusion models aim to predict all missing tokens in parallel at each intermediate diffusion step, but the denoising joint distribution is simply parameterized as a product of token-wise independent distributions. As a result, intermediate denoising steps ignore sequence-level correlations, which results in serious accumulated decoding errors and prevents users from efficient parallel decoding with fewer timesteps. To this end, in this paper, we propose Energy-based Diffusion Language Model (EDLM), an unnormalized energy-based model (EBM) that learns to jointly denoise the full sequence at each diffusion step. Our key innovation is to learn an EBM for each denoising distribution $p({\mathbf{x}}_{t-1}|{\mathbf{x}}_{t})$, where the energy directly operates on the sequence level and captures the correlation between tokens. The fundamental challenge of our framework is to develop efficient training and sampling methods for the unnormalized model. To overcome these challenges, we design the EBM in a novel residual form

$$ $p_{\textnormal{{EDLM}}}({\mathbf{x}}_{t-1}|{\mathbf{x}}_{t})=p_{\textnormal{ diffusion}}({\mathbf{x}}_{t-1}|{\mathbf{x}}_{t})\exp(-E({\mathbf{x}}_{t-1},{ \mathbf{x}}_{t})),$ (1) $$

applied over pretrained diffusion models $p_{\textnormal{diffusion}}$. Such formulation enjoys several distinct advantages. First, we analytically show that the EBM parameters can be easily obtained by leveraging pretrained autoregressive models or finetuning from bidirectional transformers via noise contrastive estimation, bypassing expensive maximum likelihood training. Furthermore, the framework corrects the decoding error and enables fast generation by conducting efficient important sampling in parallel over samples from the diffusion proposal distribution. Importantly, we highlight that, when leveraging pretraiend AR as the energy function, our approach can be interpreted as parallel sampling from pretrained language models using diffusion models as the proposal distribution.

EDLM can be viewed as a new family of discrete generative models that marries energy-based and diffusion-based models. The novel combination addresses the training and sampling distribution mismatch problem, enjoys better generation quality with less accumulated error, and improves sampling efficiency by reducing the number of sampling steps. We further provide a formal estimator to calculate the perplexity of EDLM, allowing our model to be compared against other models in a standard way. We conduct comprehensive experiments on two common language modeling benchmarks to evaluate the performance of our proposed method. Results show that on the perplexity metric, EDLM can consistently achieve state-of-the-art performance among diffusion-based counterparts, and approaches or matches AR models. On the generation quality metric, compared over the most competitive diffusion baseline, EDLM shows up to $49\%$ generative perplexity gain with the same number of sampling timesteps, and can achieve up to 1.3$\times$ sampling speedup when keeping the same sampling performance.

## 2 Related Work

Diffusion models are powerful models with surprising results in generating high-quality images (Sohl-Dickstein et al., ; Song & Ermon, ; Ho et al., ; Dhariwal & Nichol, ). These models were originally designed for continuous data, with both forward and backward reverse processes parameterized as Gaussian Markov chains. In recent years, several methods have been developed to extend the diffusion framework to generate discrete data (Austin et al., ; Hoogeboom et al., ), with Campbell et al. (); Zhao et al. () further extending the framework to model discrete data in continuous time. Most recently, several concurrent works (Lou et al., ; Sahoo et al., ; Shi et al., ) archived impressive progress in large-scale language modeling by scaling the model and simplifying the training and sampling processes. Additionally, Campbell et al. (); Gat et al. () developed flow matching methods for discrete data, which rely on a formulation similar to that of diffusion models. Despite the considerable progress in the area, we observe that all these approaches are subject to the training and sampling distribution mismatch issue, where the learned joint denoising distribution is simplified as independent distributions for each token. This problem results in a significant accumulated decoding error during the parallel sampling process and prevents users from conducting fast sampling with a small number of denoising timesteps. Notably, Lezama et al. () notices a similar issue (focusing on diffusion models in the image generation domain) and proposes to learn a corrector network to correct the independent decoding error. However, this work concentrates on the image generation domain, and the correction process needs to be run sequentially, which takes additional inference time. In contrast, we propose an energy-based denoising process that directly captures the correlation at each denoising step and enables efficient parallel decoding via an importance sampling scheme.

## 3 Discrete Diffusion Models

General Discrete Diffusion. Assume our data ${\mathbf{x}}$ lives in a finite discrete space of size $m-1$. In this paper, we augment the categorical space with an additional mask state with index $m$. In a general discrete diffusion framework (Austin et al., ), the diffusion process is defined as a Markov chain $q({\mathbf{x}}_{t}|{\mathbf{x}}_{t-1})=\text{Cat}({\mathbf{x}}_{t};Q_{t}{
\mathbf{x}}_{t-1})$, which repeatedly multiplies ${\mathbf{x}}$ with matrices $Q_{t}$ over $T$ discrete time steps.
Given these transitions, the marginal distributions at each timestep can be written in closed-form as $q({\mathbf{x}}_{t}|{\mathbf{x}})=\text{Cat}({\mathbf{x}}_{t};\bar{Q}_{t}{
\mathbf{x}})=\text{Cat}({\mathbf{x}}_{t};Q_{t}\cdots Q_{1}{\mathbf{x}})$. Such forward process can also be viewed as an interpolation between a clean data sample ${\mathbf{x}}$ and a reference distribution $\text{Cat}(\cdot;\pi)$ induced by $\bar{Q}_{T}$:

$$ $\displaystyle q({\mathbf{x}}_{t}|{\mathbf{x}})=\text{Cat}({\mathbf{x}}_{t}; \alpha_{t}{\mathbf{x}}_{0}+(1-\alpha_{t})\pi),$ (2) $$

where $\alpha_{t}\in[0,1]$ is a strictly decreasing function w.r.t $t$, with $\alpha_{0}\approx 1$ and $\alpha_{1}\approx 0$.

In the continuous time limit, for two arbitrary times $0\leq s\leq t\leq 1$, the transition distributions can be written as $q({\mathbf{x}}_{t}|{\mathbf{x}}_{s})=\text{Cat}({\mathbf{x}}_{t};\alpha_{t|s}{
\mathbf{x}}_{s}+(1-\alpha_{t|s})\pi)$, where $\alpha_{t|s}=\alpha_{t}/\alpha_{s}$ (Zhao et al., ; Shi et al., ). This implies that during each diffusion step $s\to t$ the token will jump to a sample from the prior distribution $\pi$ with a probability of $(1-\alpha_{t|s})$.
The forward process allows us to compute many distributions in closed form. One particular distribution of interest is the reversal of the forward process given ${\mathbf{x}}_{0}$, that is, the posterior distribution given by

$$ $\displaystyle q({\mathbf{x}}_{s}|{\mathbf{x}}_{t},{\mathbf{x}}_{0})=\text{Cat} \left({\mathbf{x}}_{s};\frac{[\alpha_{t|s}{\mathbf{x}}_{t}+(1-\alpha_{t|s}){ \bm{1}}\pi^{\top}{\mathbf{x}}_{t}]\odot[\alpha_{s}{\mathbf{x}}_{0}+(1-\alpha_{ s})\pi]}{\alpha_{t}{\mathbf{x}}_{t}^{\top}{\mathbf{x}}_{0}+(1-\alpha_{t}){ \mathbf{x}}_{t}^{\top}\pi}\right).$ (3) $$

Masked Diffusion. In this paper, we focus on masked (i.e., absorbing state) diffusion models, where the target distribution is set as $\pi={\mathbf{m}}$. At each diffusion step $t$ each token transitions to the ‘masked’ state ${\mathbf{m}}$ with some probability. Under such masking framework, the forward marginals (eq. 2) are given by
$q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})=\alpha_{t}{\mathbf{x}}_{0}+(1-\alpha_{t})
{\mathbf{m}}$, and the posterior (eq. 3) can be simplified as

$$ $\displaystyle q({\mathbf{x}}_{s}|{\mathbf{x}}_{t},{\mathbf{x}}_{0})=\begin{ cases}\text{Cat}({\mathbf{x}}_{s};{\mathbf{x}}_{t})&{\mathbf{x}}_{t}\neq{ \mathbf{m}},\\ \text{Cat}\left({\mathbf{x}}_{s};\frac{(1-\alpha_{s}){\mathbf{m}}+(\alpha_{s}- \alpha_{t}){\mathbf{x}}_{0}}{1-\alpha_{t}}\right)&{\mathbf{x}}_{t}={\mathbf{m} }.\end{cases}$ (4) $$

Diffusion models aim to learn a backward model $p_{\theta}(x_{s}|x_{t})$ to approximate the reversal of the forward process.
Leveraging Equation 4 we can parameterize the model $p_{\theta}({\mathbf{x}}_{s}|{\mathbf{x}}_{t})$
as

$$ $\displaystyle p_{\theta}({\mathbf{x}}_{s}|{\mathbf{x}}_{t})=q({\mathbf{x}}_{s} |{\mathbf{x}}_{t},{\mathbf{x}}_{0}={\bm{\mu}}_{\theta}({\mathbf{x}}_{t},t))= \begin{cases}\text{Cat}({\mathbf{x}}_{s};{\mathbf{x}}_{t}),&{\mathbf{x}}_{t} \neq{\mathbf{m}},\\ \text{Cat}\left({\mathbf{x}}_{s};\frac{(1-\alpha_{s}){\mathbf{m}}+(\alpha_{s}- \alpha_{t}){\bm{\mu}}_{\theta}({\mathbf{x}}_{t},t)}{1-\alpha_{t}}\right),&{ \mathbf{x}}_{t}={\mathbf{m}},\end{cases}$ (5) $$

where ${\bm{\mu}}_{\theta}=p_{\theta}({\mathbf{x}}_{0}|{\mathbf{x}}_{t})$ is known as ${\mathbf{x}}_{0}$ predictor, since it predicts the mean of the distribution over ${\mathbf{x}}_{0}$ given ${\mathbf{x}}_{t}$. However, as explained next, the practical implementation of this ${\mathbf{x}}_{0}$ predictor results in a mismatch between training and sampling distributions.

Problem Statement. Existing discrete DMs learn a denoising distribution ${\bm{\mu}}_{\theta}=p_{\theta}({\mathbf{x}}_{0}|{\mathbf{x}}_{t})$ to match the true reversal $q({\mathbf{x}}_{0}|{\mathbf{x}}_{t})$.
Let ${\mathbf{x}}$ denote the tokens of a full sequence in this section.
In practice, the model is parameterized as a factorized denoising model:

$$ ${p}_{\theta}({\mathbf{x}}_{0}|{\mathbf{x}}_{t})=\Pi_{i}{p}_{\theta}({\mathbf{x }}_{0}^{i}|{\mathbf{x}}_{t})=\Pi_{i}{\bm{\mu}}_{\theta}^{i}({\mathbf{x}}_{t},t ),\text{ where }{\bm{\mu}}_{\theta}({\mathbf{x}}_{t},t)=\begin{cases} \operatorname{softmax}(f_{\theta}({\mathbf{x}}_{t},t))&{\mathbf{x}}_{t}={ \mathbf{m}},\\ {\mathbf{x}}_{t}&{\mathbf{x}}_{t}\neq{\mathbf{m}},\end{cases}$ (6) $$

and the predictor ${\bm{\mu}}_{\theta}$ predicts each token in ${\mathbf{x}}_{0}$ independently.
This factorization enables us to conduct an efficient denoising step $p_{\theta}({\mathbf{x}}_{s}|{\mathbf{x}}_{t})$ by first sampling all ${\mathbf{x}}_{0|t}$ tokens from $p_{\theta}({\mathbf{x}}_{0}|{\mathbf{x}}_{t})$ in parallel and then masking certain tokens according to the forward $q({\mathbf{x}}_{s}|{\mathbf{x}}_{t},{\mathbf{x}}_{0|t})$.
However, this parameterization ignores dependencies between tokens in the sequence, a fundamental limitation which implies that ${p}_{\theta}({\mathbf{x}}_{0}|{\mathbf{x}}_{t})$ can never match the exact backward $q({\mathbf{x}}_{0}|{\mathbf{x}}_{t})$. As a result, this parallel sampling introduces accumulated errors, as the factorized denoising step ${p}_{\theta}$ does not match the original generative model $p_{\theta}$ for the joint distribution of the elements of ${\mathbf{x}}_{0}$.

## 4 Method

In this section, we formally introduce Energy-based Diffusion Language Model (EDLM). Our work aims to design a new family of energy-based discrete generative models that address the fundamental mismatch problem between ${p}_{\theta}({\mathbf{x}}_{0}|{\mathbf{x}}_{t})$ and $q({\mathbf{x}}_{0}|{\mathbf{x}}_{t})$ in existing models. We first describe the general EBM formulation in Section 4.1, and then elaborate on how to obtain the energy function by either leveraging pretrained AR models or finetuning via noise contrastive estimation in Section 4.2. Then we discuss how to estimate the likelihood (perplexity) of EDLM in Section 4.3, and finally introduce the efficient parallel sampling algorithm in Section 4.4

### 4.1 Residual Energy-based Models

Discrete DMs can be viewed as learning a conditional ${\mathbf{x}}_{0}$ predictor for each denoising step $t$. For example, original discrete DMs learn ${\bm{\mu}}({\mathbf{x}}_{t},t)$ to directly predict independent distributions for each token in ${\mathbf{x}}_{0}$.
In our framework, given diffused data ${\mathbf{x}}_{t}$ at timestep $t$, we introduce the generative denoising kernel as an unnormalized density:

$$ $\displaystyle p_{\theta,\phi}({\mathbf{x}}_{0}|{\mathbf{x}}_{t})={{\bm{\mu}}}_ {\theta}({\mathbf{x}}_{0}|{\mathbf{x}}_{t})\frac{\exp(-{\bm{E}}_{\phi}({ \mathbf{x}}_{0},{\mathbf{x}}_{t},t))}{Z_{\phi}({\mathbf{x}}_{t})}$ (7) $$

where ${\bm{\mu}}_{\theta}$ is the pretrained diffusion model, ${\bm{E}}_{\phi}$ is the energy introduced to capture the correlation in the ${\mathbf{x}}_{0}$ sequence, and $Z_{\phi}({\mathbf{x}}_{t})$ is a normalizing factor known as the partition function.
In the following text, we use $p_{\theta,\phi}$ to denote the joint model, ${\bm{E}}_{\phi}$ the (residual) energy function, and keep the pretrained ${\bm{\mu}}_{\theta}$ fixed.
Computing the partition function is intractable since it requires summing over the whole ${\mathbf{x}}$ space, which is exponential in the sequence length. In our language modeling experiments, the vocabulary size is around $50$k and the generation length is $1024$ tokens, resulting in a space of size around $50,000^{1024}$.
We aim to design solutions to efficiently train the parameters of the energy function so that the joint model distribution gets close to the true reversal $q({\mathbf{x}}_{0}|{\mathbf{x}}_{t})$, while avoiding computing the partition function.

### 4.2 Implementation of Energy Function

Training energy-based models with the intractable partition function is a long-standing challenge in machine learning (Hinton, ; Carreira-Perpinan & Hinton, ; LeCun et al., ). Typical maximum likelihood estimation training requires approximation of the participation function using Markov chain Monte Carlo (MCMC) sampling, which is computationally infeasible for high-dimensional data.
We aim to find efficient ways to train the energy function’s parameters for the ${\mathbf{x}}_{0}$ predictor $p_{\theta,\phi}({\mathbf{x}}_{0}|{\mathbf{x}}_{t})$.
In the following, we describe two methods to do this. One involves taking pretrained autoregressive language models as energy functions without any training, and efficiently running sampling by taking all tokens as inputs in parallel. The second solution involves fine-tuning the pretrained diffusion model via noise contrastive estimation, where the model is parameterized with bidirectional transformers and potentially captures richer correlations.

#### 4.2.1 Leverage Pre-trained Autoregressive Models

Let $p_{\textnormal{AR}}({\mathbf{x}}_{0})$ be an autoregressive model trained over clean samples ${\mathbf{x}}_{0}$, i.e., $p_{\textnormal{AR}}({\mathbf{x}}_{0})=\Pi_{i}p_{\textnormal{AR}}({\mathbf{x}}^
{i}_{0}|{\mathbf{x}}^{<i}_{0})$. Since AR models are only defined over clean samples ${\mathbf{x}}_{0}$, it is not directly clear how to leverage them in denoising tasks, since the models are never trained on diffused data ${\mathbf{x}}_{t}$. Our key insight is that, in absorbing discrete diffusion, the diffused ${\mathbf{x}}_{t}$ is the same as ${\mathbf{x}}_{0}$ with certain tokens masked according to the forward process. More specifically, letting $\bar{{\mathbf{x}}}_{0}={\mathbf{x}}_{0}[{\mathbf{x}}_{t}\neq{\mathbf{m}}]$ and ${\mathbf{x}}_{0}/\bar{{\mathbf{x}}}_{0}$ denote the set of ${\mathbf{x}}_{0}$ tokens corresponding to unmasked and masked components in ${\mathbf{x}}_{t}$, we observe that we can induce the denoising transitions using

$$ $p_{\textnormal{AR}}({\mathbf{x}}_{0}\,|\,{\mathbf{x}}_{t})=p_{\textnormal{AR}} ({\mathbf{x}}_{0}\,|\,\bar{{\mathbf{x}}}_{0})=\frac{p_{\textnormal{AR}}({ \mathbf{x}}_{0}/\bar{{\mathbf{x}}}_{0},\bar{{\mathbf{x}}}_{0})}{p_{\textnormal {AR}}(\bar{{\mathbf{x}}}_{0})}=\frac{p_{\textnormal{AR}}({\mathbf{x}}_{0})}{p_ {\textnormal{AR}}(\bar{{\mathbf{x}}}_{0})},$ (8) $$

where the normalizing constant is given by $p_{\textnormal{AR}}(\bar{{\mathbf{x}}}_{0})=\sum_{{\mathbf{x}}_{0}/\bar{{
\mathbf{x}}}_{0}}p_{\textnormal{AR}}({\mathbf{x}}_{0}/\bar{{\mathbf{x}}}_{0},
\bar{{\mathbf{x}}}_{0})$. This partition function also involves a sum over the ${\mathbf{x}}_{0}$ space and is intractable to compute. However, noting that in the reversal of masked diffusion (Equation 5) unmasked tokens $\bar{{\mathbf{x}}}_{0}$ are fixed, we have that $p_{\textnormal{AR}}({\mathbf{x}}_{0}\,|\,{\mathbf{x}}_{t})\propto p_{
\textnormal{AR}}({\mathbf{x}}_{0})$,(^1^11This property does not require the autoregressive structure of $p_{\textnormal{AR}}$, but only relies on the fact that during the diffusion reversal process unmasked tokens do not change.) which can be computed efficienctly.

Unlike the original diffusion models (Equation 6) that decode each token independently, AR models factorize the sequence distribution and take into account dependencies between tokens. Therefore, it is reasonable to assume AR models can provide a better approximation of the diffusion posterior than the factorized approximation used by diffusoin models, i.e., $p_{\textnormal{AR}}({\mathbf{x}}_{0}\,|\,{\mathbf{x}}_{t})\approx q({\mathbf{x
}}_{0}\,|\,{\mathbf{x}}_{t})$. The key challenge is that sampling the autoregressive model $p_{\textnormal{AR}}({\mathbf{x}}_{0}\,|\,{\mathbf{x}}_{t})$ at each denoising step is computationally expensive. However, it is efficient to evaluate $p_{\textnormal{AR}}({\mathbf{x}}_{0})$, which is proportional to the posterior $p_{\textnormal{AR}}({\mathbf{x}}_{0}\,|\,{\mathbf{x}}_{t})$, and thus can be used in our residual energy-based formulation.

Recall that our residual energy model (eq. 7) uses $p_{\theta}({\mathbf{x}}_{0}\,|\,{\mathbf{x}}_{t})\frac{\exp{(-{\bm{E}}_{\phi}(
{\mathbf{x}}_{0},{\mathbf{x}}_{t}))}}{Z_{\phi}({\mathbf{x}}_{t})}$ to approximate $q({\mathbf{x}}_{0}\,|\,{\mathbf{x}}_{t})$. By taking $p_{\textnormal{AR}}({\mathbf{x}}_{0}\,|\,{\mathbf{x}}_{t})$ as the proxy of $q({\mathbf{x}}_{0}\,|\,{\mathbf{x}}_{t})$, we approximate the optimal residual energy as

$$ $\displaystyle{\bm{E}}_{\phi}({\mathbf{x}}_{0},{\mathbf{x}}_{t})$ $\displaystyle=-\log{q({\mathbf{x}}_{0}}|{\mathbf{x}}_{t})+\log{p_{\theta}({ \mathbf{x}}_{0}\,|\,{\mathbf{x}}_{t})}-\log Z$ (9) $\displaystyle\approx-\log{p_{\textnormal{AR}}({\mathbf{x}}_{0}|{\mathbf{x}}_{t }})+\log{p_{\theta}({\mathbf{x}}_{0}\,|\,{\mathbf{x}}_{t})}-\log Z$ $\displaystyle q({\mathbf{x}}_{0}\,|\,{\mathbf{x}}_{t})\approx p_{\textnormal{ AR}}({\mathbf{x}}_{0}|{\mathbf{x}}_{t})$ $\displaystyle=-\log{p_{\textnormal{AR}}({\mathbf{x}}_{0}})+\log{p_{\theta}({ \mathbf{x}}_{0}\,|\,{\mathbf{x}}_{t})}+\underbrace{\log p_{\textnormal{AR}}( \bar{{\mathbf{x}}}_{0})-\log Z}_{\text{New Partition Function}}$ See Equation 8 . $$

For common EBM model inference, such as MCMC, we only require energies up to a constant. Therefore, as shown in the equation above, with a pretrained AR model we can simply define the energy function as $-\log{p_{\textnormal{AR}}({\mathbf{x}}_{0}})+\log{p_{\theta}({\mathbf{x}}_{0}
\,|\,{\mathbf{x}}_{t})}$. The fact that AR language models can evaluate likelihoods in parallel enables the development of efficient sampling algorithms. We leverage this fact to develop a sampling method based on self-normalized importance sampling using $p_{\theta}({\mathbf{x}}_{0}\,|\,{\mathbf{x}}_{t})$ as the proposal distribution (see Section 4.4). Importantly, we highlight that this energy formulation translates to sampling from the AR language model as the target distribution using the denoising distribution as the proposal distribution, providing a novel way to conduct parallel sampling from pretrained AR language models via importance sampling.

Carry-Over (CO) Parameterization. Note that in diffusion reversal (eq. 4), unmasked tokens are always directly carried over from ${\mathbf{x}}_{t}$ to ${\mathbf{x}}_{0}$, i.e., $q({\mathbf{x}}_{0}|{\mathbf{x}}_{t})=\text{Cat}(\cdot;{\mathbf{x}}_{t})$ for ${\mathbf{x}}_{t}\neq{\mathbf{m}}$. Recent progress in discrete diffusion models show that directly parameterizing this property into denoising model may improve the generation performance. In practice this is implemented by substituting the output of the ${\bm{\mu}}_{\theta}$ network to simply copy the unmasked tokens of ${\mathbf{x}}_{t}$. Similarly, in our AR EBM model, we can also carry over the unmasked tokens in each denoising step, by directly setting $p_{\textnormal{AR}}({\mathbf{x}}_{0}^{i}|{\mathbf{x}}_{0}^{<i})={\mathbf{x}}_{
t}^{i}$ for ${\mathbf{x}}_{t}\neq{\mathbf{m}}$. This “carry-over autoregressive” (coAR) EBM implementation allows simply computing $p_{\textnormal{coAR}}({\mathbf{x}}_{0}\,|\,{\mathbf{x}}_{t})=\frac{p_{
\textnormal{coAR}}({\mathbf{x}}_{0})}{p_{\textnormal{coAR}}(\bar{{\mathbf{x}}}
_{0})}$ by just $p_{\textnormal{coAR}}({\mathbf{x}}_{0})$, as we always carry over those tokens from ${\mathbf{x}}_{t}$ and therefore $p_{\textnormal{coAR}}(\bar{{\mathbf{x}}}_{0})=1$. Such practice enables us to evaluate the exact denoising likelihood without estimating the partition function. We will introduce how to calculate the data likelihood based on denoising distribution in Section 4.3.

#### 4.2.2 Training with Noise Contrastive Estimation

We can also train the parameters of the residual energy function using Noise Contrastive Estimation (NCE) (Gutmann & Hyvärinen, ), specifically its conditional version (Ma & Collins, ). First, NCE training relies on contrastive samples from the data distribution and a noise distribution, where the noise distribution needs to be close to the data distribution. Second, it involves computing the likelihoods of these samples by model distribution and noise distribution. In our approach, given a pair of clean data ${{\mathbf{x}}}_{0}$ and diffused data ${\mathbf{x}}_{t}$, we set true posterior $q(\hat{{\mathbf{x}}}_{0}|{\mathbf{x}}_{t},{{\mathbf{x}}}_{0})$(^2^22In this subsection, we use notation $\hat{{\mathbf{x}}}_{0}$ instead ${{\mathbf{x}}}_{0}$ to distinguish the generated data and clean data ${\mathbf{x}}_{0}$.) as positive distribution and the denoising distribution $p_{\theta}(\hat{{\mathbf{x}}}_{0}|{\mathbf{x}}_{t})$ as the negative distribution. Thanks to the residual energy formulation from eq. 7, the log-odds reduce to $\log p_{\theta,\phi}-\log p_{\theta}=-{\bm{E}}_{\phi}$, and the objective is simplified into the binary classification objective

$$ $\displaystyle{\mathcal{L}}_{\text{NCE}}$ $\displaystyle(\phi;\theta)=-\mathbb{E}_{{{\mathbf{x}}}_{0}\sim p_{\text{data}} ,{\mathbf{x}}_{t}\sim q({\mathbf{x}}_{t}|{{\mathbf{x}}}_{0})}\Big{[}$ (10) $\displaystyle\mathbb{E}_{{\mathbf{x}}_{+}\sim q(\hat{{\mathbf{x}}}_{0}|{ \mathbf{x}}_{t},{{\mathbf{x}}}_{0})}\log\frac{1}{1+\exp({\bm{E}}_{\phi}({ \mathbf{x}}_{+},{\mathbf{x}}_{t},t))}+\mathbb{E}_{{\mathbf{x}}_{-}\sim{p}_{ \theta}(\hat{{\mathbf{x}}}_{0}|{\mathbf{x}}_{t})}\log\frac{1}{1+\exp(-{\bm{E}} _{\phi}({\mathbf{x}}_{-},{\mathbf{x}}_{t},t))}\Big{]},$ $$

where ${\mathbf{x}}_{+}$ is positive data from the true posterior and ${\mathbf{x}}_{-}$ is negative data from diffusion model, given the diffused sequence ${\mathbf{x}}_{t}$. Here, the true posterior $q(\hat{{\mathbf{x}}}_{0}|{\mathbf{x}}_{t},{{\mathbf{x}}}_{0}):={{\mathbf{x}}}_
{0}$ is defined as simply recovering the true data. Training the energy function can be viewed as training a conditional classifier to discriminate the real text and text generated by the denoiser used by the diffusion model. Intuitively, the training objective attempts to capture the correlation in ${\mathbf{x}}_{0}$ generations, assigning negative energy to real data and positive energy to data produced by the denoiser network. As a result, the joint model $p_{\theta,\phi}$ acts as a corrected denoising distribution able to take into account correlations between tokens. A detailed pseudo code for the NCE training process is provided in Algorithm 2.

### 4.3 Evaluation with Rao-Blackwellized Likelihood Bounds

A common protocol for evaluating discrete generative models on language modeling is perplexity (PPL), which relies on computing log-likelihoods for a hold-out test dataset. In this section, we explain how to estimate the log-likelihood with EDLM.
First, with the energy-based $x_{0}$ predictor $p_{\theta,\phi}({\mathbf{x}}_{0}|{\mathbf{x}}_{t})$, the step denoising model $p_{\theta,\phi}({\mathbf{x}}_{s}|{\mathbf{x}}_{t})$ is given by

$$ $\displaystyle p_{\theta,\phi}({\mathbf{x}}_{s}|{\mathbf{x}}_{t})$ $\displaystyle=q({\mathbf{x}}_{s}|{\mathbf{x}}_{t},{\mathbf{x}}=p_{\theta,\phi} ({\mathbf{x}}_{0}|{\mathbf{x}}_{t}))=\begin{cases}\text{Cat}({\mathbf{x}}_{s}; {\mathbf{x}}_{t}),&{\mathbf{x}}_{t}\neq{\mathbf{m}},\\ \text{Cat}\left({\mathbf{x}}_{s};\frac{(1-\alpha_{s}){\mathbf{m}}+(\alpha_{s}- \alpha_{t})p_{\theta,\phi}({\mathbf{x}}_{0}|{\mathbf{x}}_{t})}{1-\alpha_{t}} \right),&{\mathbf{x}}_{t}={\mathbf{m}}.\end{cases}$ (11) $$

For the simplified reverse process in Equation 5, the log-likelihood can be simply computed as a form of Rao-Blackwellization. The Rao-Blackwellized likelihood bound can be estimated by the discrete-time diffusion loss of finite $T$: ${\mathcal{L}}_{\text{diffusion}}=\sum_{i=1}^{T}\mathbb{E}_{q}[D_{\mathrm{KL}}(
q({\mathbf{x}}_{s_{i}}|{\mathbf{x}}_{t_{i}},{\mathbf{x}})\|p_{\theta,\phi}({
\mathbf{x}}_{s_{i}}|{\mathbf{x}}_{t_{i}}))]$. Under EDLM, the loss can be written as:

$$ $\displaystyle{\mathcal{L}}_{\text{diffusion}}$ $\displaystyle=\sum_{i=1}^{T}\mathbb{E}_{q}\left[\frac{\alpha_{t_{i}}-\alpha_{s _{i}}}{1-\alpha_{t_{i}}}\left(\log p_{\theta}({\mathbf{x}}_{0}|{\mathbf{x}}_{t _{i}})-{\bm{E}}_{\phi}({\mathbf{x}}_{0},{\mathbf{x}}_{t_{i}},{t_{i}})-\log Z_{ \phi}({\mathbf{x}}_{t_{i}})\right)\right],$ (12) $$

where $p_{\theta}({\mathbf{x}}_{0}|{\mathbf{x}}_{t_{i}})=\langle{\bm{\mu}}_{\theta}({
\mathbf{x}}_{t_{i}}),{\mathbf{x}}\rangle$ is a simple cross-entropy between ${\bm{\mu}}_{\theta}$ predicted logits and the clean data.
We can extend the objective Equation 12 to the continuous limit by taking $T\to\infty$, which induces the following negative evidence lower bound (NELBO) ${\mathcal{L}}_{\infty}$:

$$ $\displaystyle{\mathcal{L}}_{\infty}=\mathbb{E}_{q}\int_{t=0}^{t=1}\frac{\alpha _{t}^{\prime}}{1-\alpha_{t}}\left(\log p_{\theta}({\mathbf{x}}_{0}|{\mathbf{x} }_{t})-{\bm{E}}_{\phi}({\mathbf{x}}_{0},{\mathbf{x}}_{t},t)-\log Z_{\phi}({ \mathbf{x}}_{t})\right){\textnormal{d}}t.$ (13) $$

where $\alpha_{t}^{\prime}$ denotes the derivative of $\alpha_{t}$ w.r.t $t$. Such ELBO is tighter than the discrete-time version, and is invariant to the noise schedule. However, this bound requires the partition function $\log Z_{\phi}({\mathbf{x}}_{t})=\log\sum_{{\mathbf{x}}_{0}}p_{\theta}({\mathbf
{x}}_{0}|{\mathbf{x}}_{t})\exp(-{\bm{E}}_{\phi}({\mathbf{x}}_{0},{\mathbf{x}}_
{t}))=\log\mathbb{E}_{{\mathbf{x}}_{0}\sim p_{\theta}}\exp(-{\bm{E}}_{\phi}({
\mathbf{x}}_{0},{\mathbf{x}}_{t}))$ and thus is intractable to compute. We use two estimators based on the work (Nowozin, ; Deng et al., ) to estimate the partition function $\log Z_{\phi}$.

###### Theorem 1 .

Given diffused data ${\mathbf{x}}_{t}$ at timestep $t$, let $\log{\mathcal{Z}}_{n}$ denote the empirical estimation of $\log Z_{\phi}({\mathbf{x}}_{t})=\log\mathbb{E}_{{\mathbf{x}}_{0}\sim p_{\theta
}}\exp(-{\bm{E}}_{\phi}({\mathbf{x}}_{0},{\mathbf{x}}_{t}))$ with $n$ samples ${\mathbf{x}}_{0}^{(i)}\sim p_{\theta}(i=1,\cdots,n|{\mathbf{x}}_{t})$: $\log{\mathcal{Z}}_{n}=\log\frac{1}{n}\sum_{i=1}^{n}\exp(-{\bm{E}}({\mathbf{x}}
_{0}^{(i)},{\mathbf{x}}_{t}))$. Then $\forall\epsilon>0$, $\exists N>0$ such that $\forall n>N$ we have

$$ $\log Z_{\phi}-\epsilon<\mathbb{E}[\log{\mathcal{Z}}_{n}]<\log Z_{\phi}<\mathbb {E}[(2n-1)\log{\mathcal{Z}}_{n}-2(n-1){\mathcal{Z}}_{n-1}]<\log Z_{\phi}+\epsilon.$ (14) $$

This can be used to estimate lower and upper bounds of the partition function. We note, however, that the bounds only hold asymptotically when $n$ is sufficiently large. In practice, we follow Nowozin () to use the leave-one-out strategy to estimate ${\mathcal{Z}}_{n-1}$, which is proven to yield an estimator with lower variance. In addition, we want to recall that in our autoregressive energy function instance (section 4.2.1), we introduce the carry-over parameterization, where the joint energy model is self-normalized, which enables us to exactly compute the ELBO without estimating $Z_{\phi}$.

###### Theorem 1 .

### 4.4 Efficient Generation via Importance Sampling

Sampling from EDLM with each denoising step operated by the joint model is non-trivial. Naive approaches relying on Gibbs sampling (Gelfand, ; Hinton, ) were originally applied to binary inputs and are not scalable to large dictionaries with energy functions parameterized as large transformer models. In this paper, we resort to self-normalizing importance sampling (Owen, ; Hammersley, ). With our joint model as the product of the diffusion model $p_{\theta}$ and residual energy function ${\bm{E}}_{\phi}$, and given intermediate diffusion data ${\mathbf{x}}_{t}$ at timestep $t$, we can conduct efficient parallel sampling by: 1) sampling multiple ${{\mathbf{x}}_{0}}$ predictions $\{{\mathbf{x}}^{i}_{0}\}_{i=1}^{k}$ from the diffusion denoiser $p_{\theta}({\mathbf{x}}_{0}|{\mathbf{x}}_{t})$; 2) feeding samples into energy function in parallel to compute energies; and 3) resampling a single ${\mathbf{x}}_{0}$ from the pool $\{{\mathbf{x}}^{i}_{0}\}_{i=1}^{k}$ according to the energy values. The sampled ${\mathbf{x}}_{0}$ is then fed into the posterior formulation $q({\mathbf{x}}_{s}|{\mathbf{x}}_{t},{\mathbf{x}}_{0})$ to perform one reverse step (i.e. one-step denoising).

Figure: Algorithm 1 Denoising via Importance Sampling

Importance Sampling Window. While importance sampling introduces additional computation, it yields a significant reduction in the parallel decoding error. Therefore, the method allows us to conduct diffusion sampling with fewer denoising steps, reducing the overall sampling wall-clock time. To further accelerate the sampling speed, we introduce the concept of importance sampling window length ${\textnormal{w}}\in[0,1]$, which sets the timestep for stopping importance sampling, i.e., we only conduct importance sampling in the time window $t\in[1-{\textnormal{w}},1]$. Interestingly, in our experiment, we notice that the energy-based importance sampling during the early stage of denoising sampling contributes more to the generation quality improvement. We conclude that this phenomenon is due to the fact that during the early stage of sampling the diffusion model is prone to make more errors in independent ${\mathbf{x}}_{0}$ prediction, since there is little information on the full sequence. This encourages us to explore sampling with a short time window w for higher efficiency, which we discuss in detail in section 5.3.

Detailed pseudo code for the sampling procedure is provided in Algorithm 1. In practice, we observe that a relatively small importance sampling size can already yield significant performance gain, though theoretically we would only recover exact samples from the joint model distribution asymptotically when the number of importance samples goes to infinity.

## 5 Experiments

This section presents the results achieved by our method on different language modeling tasks. We compare our EDLM against existing diffusion models, reporting the widely adopted metrics perplexity and generative perplexity. We provide the experimental setup in section 5.1, describe our results in section 5.2, and provide additional ablation studies in section 5.3.

### 5.1 Experimental Setup

Datasets. We use two text datasets: 1) Text8 (Mahoney, ), a relatively small-scale, character-level text modeling benchmark extracted from English Wikipedia, and 2) OpenWebText, an open-source replica of the unreleased WebText (Gokaslan & Cohen, ) dataset used to train GPT-2.

Baselines. We compare EDLM against state-of-the-art discrete diffusion models and a transformer AR model (Vaswani et al., ). Discrete diffusion baselines include Discrete Diffusion Model (D3PM) (Austin et al., ), Score Entropy Discrete Diffusion model (SEDD) (Lou et al., ), Masked Diffusion Language Model (MDLM) (Sahoo et al., ). We notice other recent works such as discrete flow matching (Campbell et al., ; Gat et al., ) and continuous-time masked diffusion (Shi et al., ), but they are either not open-sourced or focus on non-language settings. Therefore, we take their concurrent work with similar formulation MDLM (Sahoo et al., ) as the representative for comparison. On the small-scale Text8 benchmark, we additionally evaluate other discrete generative models including Plaid (Gulrajani & Hashimoto, ), Bayesian Flow Network (Graves et al., ), Any-order Autoregressive Models ARDM (Hoogeboom et al., ) and MAC (Shih et al., ), and flow-based methods IAF/SCF (Ziegler & Rush, ), AR Argmax Flow (Hoogeboom et al., ), Discrete Flow (Tran et al., ), and Multinomial Diffusion (Hoogeboom et al., ). Note that, these methods are not scaled up yet to the large dictionary and long sequence length, and thus we only involve them in the small-scale benchmark.

**Table 1: Bits Per Character (BPC) on Text8 test set.**
| Method | BPC ($\downarrow$) |
| --- | --- |
| Autoregressive |  |
| Transformer AR | 1.23 |
| IAF/SCF | 1.88 |
| AR Argmax Flow | 1.39 |
| AR Discrete Flow | 1.23 |
| Any-order Autoregressive |  |
| ARDM | $\leq$ 1.43 |
| MAC | $\leq$ 1.40 |
| Continuous Diffusion |  |
| Plaid | $\leq$ 1.48 |
| BFN | $\leq$ 1.41 |
| Discrete Diffusion |  |
| Mult. Diffusion | $\leq$ 1.72 |
| D3PM Uniform | $\leq$ 1.61 |
| D3PM Absorb | $\leq$ 1.45 |
| SEDD Absorb | $\leq$ 1.41 |
| MDLM | $\leq$ 1.40 |
| EDLM (Ours) | $\leq$ 1.24 |

Metrics. We follow the convention (Shi et al., ) to use the common protocols Bits Per Character (BPC), Perplexity (PPL), and Generative Perplexity (Gen PPL) to evaluate generative sequence models. For a sequence of length $L$, BPC metric is given by ${-\frac{1}{L}\sum_{i=1}^{L}\log_{2}p({\mathbf{x}}_{i})}$ and PPL is given by $\exp{(-\frac{1}{L}\sum_{i=1}^{L}\log p({\mathbf{x}}_{i}))}$, which can be viewed as the average number of tokens the model is uncertain of. BPC and PPL are calculated based on model likelihoods on true data from the test set. Gen PPL instead consists of likelihoods calculated by another large oracle model on data generated by the evaluated models. Intuitively, PPL and BPC similarly evaluate the likelihood modeling capacity, while Gen PPL evaluates generation quality and consistency. To compute Gen PPL we generate $2048$ samples of $1024$.

Implementation Details. For all models including our methods and baselines, we follow the common practice of using standard 12-layer transformers similar to GPT2-small scale (Radford et al., ; Shi et al., ). Our proposed EDLM combines two models, the diffusion model $p_{\theta}$ and the energy function ${\bm{E}}_{\phi}$. For all experiments, we use pretrained MDLM (Sahoo et al., ) as the diffusion model $p_{\theta}$. For AR-based EBM (see section 4.2.1, named EDLM-AR or EDLM-coAR when using carry-over), we directly leverage the pretrained AR model as the energy function. For NCE fine-tuned EBM (see section 4.2.2, named EDLM-NCE), we finetune the pretrained MDLM, with the energy function computed by projecting the mean-pooled last token embeddings down to a single scalar value. Note that, MDLM relies on transformers with bidirectional attention while AR only imposes casual attention. Therefore, we expect that EDLM-NCE can capture more sequence information than EDLM-AR.

### 5.2 Results

Text8. Following the previous common practice (Austin et al., ; Lou et al., ), we evaluated all models on short text chunks of length $256$.
We follow the same dataset splits and transformers model size to parameterize the denoising models.
Results are summarized in Table 1, where we report the standard bits-per-character metric for the Text8 test set.
In this small-scale experiment, we did not notice significant differences between EDLM-AR and EDLM-NCE, and therefore only report one result as representative.
As shown in the table, EDLM outperforms all previous diffusion models, whether in discrete or continuous diffusion formulation. We also outperform the any-order autoregressive models, which also generate sequences with flexible decoding order and therefore have a strong connection to discrete diffusion models.
Importantly, for the first time, our diffusion model even approaches the performance of transformer AR and AR Discrete Flow, both relying on autoregressive modeling. For this experiment, due to the limited dictionary and length size, we did not observe meaningful signals in generation quality metrics (Gen PPL). We leave detailed generation quality and diversity evaluation to the large-scale OpenWebText experiments next.

**Table 2: Test perplexities ($\downarrow$). Left part: results evaluated on the OpenWebText test set; Right part: zero-shot results on unseen datasets. All perplexities for diffusion models are upper bounds.**
|  | OpenWebText | PTB | Wikitext | LM1B | Lambada | AG News | Pubmed | Arxiv |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| AR | 17.56 | 82.05 | 25.75 | 51.25 | 51.28 | 52.09 | 49.01 | 41.73 |
| SEDD | 24.56 | 100.09 | 34.28 | 68.20 | 49.86 | 62.09 | 44.53 | 38.48 |
| MLDM | 23.83 | 95.26 | 32.83 | 67.01 | 47.52 | 61.15 | 41.89 | 37.37 |
| EDLM-NCE (Ours) | 21.52 | 93.21 | 30.77 | 63.19 | 46.92 | 60.02 | 41.80 | 36.63 |
| EDLM-AR (Ours) | 20.49 | 89.67 | 29.24 | 60.80 | 49.70 | 57.27 | 45.90 | 38.38 |
| EDLM-coAR (Ours) | 17.58 | 89.73 | 28.31 | 60.23 | 50.04 | 57.94 | 46.31 | 39.02 |

**Table 3: Generative perplexities on unconditional text generation. Llama2, Llama3, and GPT2 ($\downarrow$) correspond to Generative perplexities evaluated by different oracle models.**
| Method | Timesteps | Llama2$\downarrow$ | Llama3$\downarrow$ | GPT2$\downarrow$ | Entropy |
| --- | --- | --- | --- | --- | --- |
| Data | - | 7.0 | 9.4 | 14.7 | 7.7 |
| Autoregressive | 1024 | 22.9 | 40.3 | 35.7 | 8.1 |
| SUNDAE | 200 | 29.5 | 45.1 | 34.7 | 5.2 |
| Ssd-LM | $>$10000 | 73.3 | 203.1 | 99.2 | 4.8 |
| D3PM Absorb | 1024 | 692.3 | 754.9 | 842.3 | 7.6 |
| SEDD | 256 / 512 / 1024 / 2048 | 36.1 / 32.5 / 27.3 / 23.1 | 65.0 / 54.3 / 43.7 / 36.2 | 64.8 / 52.2 / 41.5 / 33.7 | 7.8 / 7.7 / 7.6 / 7.5 |
| MDLM | 256 / 512 / 1024 / 2048 | 37.2 / 30.6 / 27.6 / 23.9 | 66.8 / 52.6 / 44.6 / 37.6 | 66.8 / 52.4 / 42.6 / 34.9 | 7.9 / 7.8 / 7.6 / 7.5 |
| EDLM-AR (Ours) | 256 / 512 / 1024 / 2048 | 34.7 / 26.8 / 19.6 / 14.6 | 62.2 / 44.4 / 28.8 / 20.8 | 62.1 / 42.0 / 25.5 / 17.9 | 7.9 / 7.6 / 7.2 / 6.9 |
| EDLM-NCE (Ours) | 256 / 512 / 1024 / 2048 | 35.7 / 26.3 / 19.0 / 14.6 | 62.9 / 44.1 / 28.8 / 20.7 | 61.7 / 42.5 / 25.5 / 17.7 | 7.9 / 7.6 / 7.3 / 6.9 |

OpenWebText. We report the perplexity results of all models trained on the large-scale OpenWebText dataset in Table 2. We evaluate the model capacity for both the OpenWebText test set and seven out-of-domain unseen datasets, used to validate the models’ ability in zero-shot generalization. These unseen datasets include Our Penn Tree Bank (PTB) (Marcus et al., )), Wikitext (Merity et al., ), LM1B, Lambada (Paperno et al., ), AG News (Zhang et al., ), and Scientific Papers (Pubmed and Arxiv) (Cohan et al., ).
As shown in the results, we again observe that EDLM outperforms existing diffusion methods by a significant margin, and also approaches the AR baseline.
For EDLM-NCE and EDLM-AR, since the energy function is unnormalized, we use the upper bound estimator in Theorem 1 to estimate the upper bound for negative log-likelihood, which induces the upper bounds of perplexities shown in the table. For EDLM-coAR, as discussed in Section 4.2.1, we overcome the estimation of the partition function when combined with carry-over and thus compute the exact ELBO for likelihoods.

We also compare our method against prior non-autoregressive generative models for generation quality. All models are trained on OpenWebText, and results are presented in Table 3. We additionally incorporate Step-unrolled Denoising Autoencoder (SUNDAE) (Savinov et al., ) and Ssd-LM (Han et al., ) into the comparison. These two methods are also highly related diffusion methods for text generation which cannot be assessed by log-likelihood. In this experiment, we set the importance sampling window ${\textnormal{w}}=1$, i.e. use importance sampling at every denoising step, to evaluate the full generation ability of EDLM. Additionally, EDLM-AR shows similar results either with or without the carry-over trick, therefore we only report one method for simplicity. As shown in Table 3, our method again consistently outperforms all baselines on generative perplexity for all numbers of denoising timesteps, while keeping reasonable diversity based on the entropy metric. We highlight that our method does not hurt the entropy (diversity), since it achieves similar entropy as MDLM when keeping similar Gen PPL. For example, by comparing EDLM with 512 timesteps and MDLM with 1024 steps, we can see they show roughly the same Gen PPL and entropy score, but EDLM requires significantly fewer sampling steps to reach that level of performance.

Figure: (a) Generative Perplexity vs. Time.
Refer to caption: x1.png

### 5.3 Analysis and Ablation Study

Improved Sampling Efficiency. In previous experiments, we show that with the same number of timesteps sampling, EDLM can achieve significantly better generation results with less accumulated error. In this part, we highlight that EDLM can also enable better sampling efficiency when keeping the same generative quality, by using reduced denoising steps. As introduced in section 4.4, our key observation is that the energy-based importance sampling correction during the early stage of denoising sampling contributes more to the generative perplexity improvement. Therefore, to accelerate the sampling speed, in this experiment, we set the importance sampling window as ${\textnormal{w}}=0.2$, which empirically yields the best generation quality with the same wall-clock sampling time budget. ${\textnormal{w}}=0.2$ means we use importance sampling for $t\in[0.8,1]$, and just use the original diffusion model ${\bm{\mu}}_{\theta}$ for denoising for $t\in[0,0.8)$. We provide additional results with other window size in Section C.2. We test AR baseline and diffusion-based models with $[512,768,1024]$ denoising steps, and plot the curve of corresponding metric vs. wall-clock time for generating a single $1024$-length sentence in Figures 1(a) and 1(b). As shown in the figure, EDLM achieves a better sampling quality vs time trade-off compared against MDLM. Specifically, when achieving a similar generative perplexity as the AR baseline (35.7), we can see from the figure that EDLM only takes $\sim$13 seconds while the MDLM baseline takes $\sim$17 seconds, indicating a $\sim$1.3$\times$ acceleration without any performance drop. This speedup further highlights the advantage of our energy-based denoising formulation and importance sampling scheme, where we can conduct sampling with reduced denoising steps and correct the decoding errors by parallel importance sampling.

Effect of Importance Sampling. We further study the effect of importance sampling hyperparameters in this section. We fix the number of denoising timesteps as $1024$, and investigate the generation performance with varying importance sampling size and window length. We refer readers to Section 4.4 and Algorithm 1 for details of these two hyperparameters. Again, we found that EDLM-NCE and EDLM-AR exhibit similar trends in the study, and therefore show one study result here. The ablation study results are summarized in Figure 1(c), from which we highlight two key observations: 1) Row-wise comparison indicates that the sampling quality is not sensitive to the importance sampling size. However, we emphasize that the conclusion is based on sizes smaller than 16, as a larger size will lead to out-of-memory issues on the GPU. 2) Column-wise study show that a longer importance sampling window can consistently improve the sampling quality. However, we note here that when setting the window as 0 0, the sampling decays to the original MDLM sampling, and the $1024$ steps GPT-2 Gen PPL is $42.6$ (see Table 3). Therefore, a $0.2$ window already significantly reduces the Gen PPL to $\sim 30$ while a longer window only marginally further improves the results. This phenomenon suggests that major decoding error happens during early-stage denoising sampling, and a short window can greatly overcome the errors while maintaining high sampling efficiency, as shown in our efficiency study above.

## 6 Conclusion

In this paper, we introduced Energy-based Diffusion Language Model (EDLM), which integrates energy-based models with discrete diffusion models to address the limitations of parallel text generation. By leveraging a residual energy-based approach, EDLM effectively reduces the mismatch between training and sampling distributions in discrete diffusion models, resulting in improved generation quality and efficiency. Through experiments on both small and large language modeling benchmarks, EDLM demonstrates state-of-the-art performance among diffusion models and approaches the quality of autoregressive models, while offering significant sampling speedup. These results highlight the potential of energy-based approaches in advancing discrete generative modeling, setting the stage for further exploration of parallel generation techniques.

## Appendix A Proof

To make the paper self-contained, we incorporate the relevant theoretical proof for Theorem 1 from (Nowozin, ; Deng et al., ). Our key difference is correcting a minor misargument from estimation on $Z_{\phi}$ to estimation on $\log Z_{\phi}$.

###### Theorem 1 .

Given diffused data ${\mathbf{x}}_{t}$ at timestep $t$, let $\log{\mathcal{Z}}_{n}$ denote the empirical estimation of $\log Z_{\phi}({\mathbf{x}}_{t})=\log\mathbb{E}_{{\mathbf{x}}_{0}\sim p_{\theta
}}\exp(-{\bm{E}}_{\phi}({\mathbf{x}}_{0},{\mathbf{x}}_{t}))$ with $n$ samples ${\mathbf{x}}_{0}^{(i)}\sim p_{\theta}(i=1,\cdots,n|{\mathbf{x}}_{t})$: $\log{\mathcal{Z}}_{n}=\log\frac{1}{n}\sum_{i=1}^{n}\exp(-{\bm{E}}({\mathbf{x}}
_{0}^{(i)},{\mathbf{x}}_{t}))$. Then $\forall\epsilon>0$, $\exists N>0$ such that $\forall n>N$ we have

$$ $\log Z_{\phi}-\epsilon<\mathbb{E}[\log{\mathcal{Z}}_{n}]<\log Z_{\phi}<\mathbb {E}[(2n-1)\log{\mathcal{Z}}_{n}-2(n-1){\mathcal{Z}}_{n-1}]<\log Z_{\phi}+\epsilon,$ (15) $$

###### Proof.

Following Eq. 35 in Nowozin (), we have $\mathbb{E}[\log{\mathcal{Z}}_{n}]$ as

$$ $\mathbb{E}[\log{\mathcal{Z}}_{n}]=\log Z_{\theta}-\frac{\gamma_{2}}{2\gamma^{2 }}\frac{1}{n}+\frac{1}{3\gamma^{3}}\frac{\gamma_{3}}{n^{2}}-\frac{1}{4\gamma^{ 4}}(\frac{3}{n^{2}}\gamma_{2}^{2}+\frac{1}{n^{3}}(\gamma_{4}-3\gamma_{2}^{2})) +\frac{1}{5\gamma^{5}}(\frac{10}{n^{3}}\gamma_{3}\gamma_{2}+\frac{1}{n^{4}}( \gamma_{5}-10\gamma_{3}\gamma_{2}))+o(n^{-3})$ (16) $$

where $\gamma=\mathbb{E}[\log{\mathcal{Z}}_{n}]$, $\gamma_{k}=\mathbb{E}[(\log{\mathcal{Z}}_{n}-\gamma)^{k}]$.

By omitting the higher-order expansion, we equivalently have
$\mathbb{E}[\log{\mathcal{Z}}_{n}]=\log Z_{\theta}-\frac{\gamma_{2}}{2\gamma^{2
}}\frac{1}{n}+o(n^{-1}),$
which indicates that $\lim_{n\to\infty}\mathbb{E}[\log{\mathcal{Z}}_{n}]=\log Z_{\theta}$. Then, we have

- •
$\forall\epsilon>0$, $\exists N_{1}>0$ such that for $n>N_{1}$, $\mathbb{E}[\log{\mathcal{Z}}_{n}]>\log Z_{\theta}-\epsilon$.
- •
Since $\lim_{n\to\infty}n(\log Z_{\theta}-\mathbb{E}[\log{\mathcal{Z}}_{n}])=\lim_{n
\to\infty}\frac{\gamma_{2}}{2\gamma^{2}}+o(1)=\frac{\gamma_{2}}{2\gamma^{2}}>0$, so $\exists N_{2}>0$ such that for $n>N_{2}$, $\log Z_{\theta}>\mathbb{E}[\log{\mathcal{Z}}_{n}]$.

In summary, we have proved that $\log Z_{\theta}-\epsilon<\mathbb{E}[\log{\mathcal{Z}}_{n}]<\log Z_{\theta}$.

Again, by omitting higher-order expansions using Equation 16, we can have another equivalent form.
$\mathbb{E}[\log{\mathcal{Z}}_{n}]=\log Z_{\theta}-\frac{\gamma_{2}}{2\gamma^{2
}}\frac{1}{n}+\frac{c}{n^{2}}+o(n^{-2})$
where $c$ is a constant, and therefore $\mathbb{E}[(2n-1)\log{\mathcal{Z}}_{n}-2(n-1)T_{n-1}]=(2n-1)\mathbb{E}[\log{
\mathcal{Z}}_{n}]-2(n-1)\mathbb{E}[T_{n-1}]=\log Z_{\theta}+\frac{\gamma_{2}}{
2\gamma^{2}}\frac{1}{n}+o(n^{-1})$. This indicates $\lim_{n\to\infty}\mathbb{E}[(2n-1)\log{\mathcal{Z}}_{n}-2(n-1)T_{n-1}]=\log Z_
{\theta}$. Therefore, we have

- •
$\forall\epsilon>0$, $\exists N_{3}>0$ such that $\forall n>N_{3}$ $\mathbb{E}[(2n-1)\log{\mathcal{Z}}_{n}-2(n-1)T_{n-1}]<\log Z_{\theta}+\epsilon$.
- •
Since $\lim_{n\to\infty}n(\mathbb{E}[(2n-1)\log{\mathcal{Z}}_{n}-2(n-1)T_{n-1}]-\log Z
_{\theta})=\lim_{n\to\infty}\frac{\gamma_{2}}{2\gamma^{2}}+o(1)>0$, so $\exists N_{4}>0$ such that when $n>N_{4}$ we have $\mathbb{E}[(2n-1)\log{\mathcal{Z}}_{n}-2(n-1)T_{n-1}>\log Z_{\theta}$.

Given all the results above, we have that $\forall\epsilon>0$, for $\forall n>N$ that $N=\max\{N_{1},N_{2},N_{3},N_{4}\}$

$$ $\log Z_{\theta}-\epsilon<\mathbb{E}[\log{\mathcal{Z}}_{n}]<\log Z_{\theta}< \mathbb{E}[(2n-1)\log{\mathcal{Z}}_{n}-2(n-1)T_{n-1}]\log Z_{\theta}+\epsilon$ $$

∎

###### Theorem 1 .

###### Proof.

## Appendix B Algorithm

We provide the detailed pseudo code for noise contrastive training of the EBM in Algorithm 2.

Figure: Algorithm 2 Noise Contrastive Estimation Training

## Appendix C Additional Experiment Details and Results

### C.1 Additional Experiment Details

We provide additional experiment setup details in this section.

Text8.
We follow all the common practices in Austin et al. (); Campbell et al. () to conduct Text8 experiments, which has a dictionary size of $28$ with $26$ lowercase letters, a white-space token, and a mask token. We follow the standard dataset split and train MDLM using a standard 12-layer transformer architecture. the transformer also has the same number of heads (12) and hidden dimension (784) as in Austin et al. (). The model is trained text chunks of length 256 for 1 million steps with batch size 512. In the NCE setting, we then finetune from this MDLM with a pooling layer and an additional scalar energy prediction head. We observe the finetuning procedure to converge fast in just $10,000$ steps. For both MDLM training and EDLM-NCE finetuning, we follow previous work to use the cosine learning rate schedule with a linear warm-up of $2000$ steps. We set the channel-wise dropout rate as $0.05$ and conducted optimization with AdamW and a learning rate of $0.0003$. We similarly adopt a weight decay factor of $0.03$. The NCE finetuning process can be done on 4 GPUs for less than 4 hours.

OpenWebText.
We follow the standard data split in (Sahoo et al., ) to leave a validation split with the last 100k docs as the validation set.
We tokenize OpenWebText with the GPT2 tokenizer, with a vocabulary size of around 50K. All models are trained with sequences wrapped to a length of $1,024$ and additionally set eos as the first and last token of every batch.
All the architectural choices are kept the same with the Text8 experiment, where we use transformers with 12 layers, a hidden dimension of 768, 12 attention heads, and a timestep embedding of 128 when applicable. Word embeddings are not tied between the input and output.
Other training details are also kept the same, i.e., we use the AdamW optimizer with a batch size of 512, learning rate $0.0003$ with a linear warm-up of $2500$ steps. We train all models for 1M steps with the dropout rate reduced to $0.1$. Again, the NCE finetuning of the energy function from the pretrained diffusion model is efficient and can converge in $400,000$ steps.

### C.2 Additional Experiment Results

We provide additional experimental results in this section.

Behavior of Energy Function. We first provide a closer look at the behavior of energy function across different denoising steps. Specifically, for each timestep $t$, we draw diffused data from $q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$, and then sample positive data ${{\mathbf{x}}}_{+}\sim q(\cdot|{\mathbf{x}}_{t},{\mathbf{x}}_{0})$
and negative data ${{\mathbf{x}}}_{-}\sim p_{\theta}(\cdot|{\mathbf{x}}_{t})$, similar to Algorithm 2. Then we study the energy value for $E_{\phi}({\mathbf{x}}_{+},{\mathbf{x}}_{t})$ and $E_{\phi}({\mathbf{x}}_{-},{\mathbf{x}}_{t})$.

We summarize all plots for the energy w.r.t timestep in Figure 2. We sample 16 negative samples for the visualization. Specifically, in Figure 2, in the first row, we plot the energy of positive samples and the average energy of negative samples, to see whether the energy function can differentiate true and generate data; in the second row, we plot the maximum and minimum energy value of the 16 negative samples, to see whether energy function can differentiate better and worse samples in generated data; and in the third row, we plot the effective sampling size (ESS) for energies of the 16 negative samples. Formally, let ${\mathbf{e}}^{i}$ denote the energy for $i$-th negative sample, we first normalize the energies to $\hat{{\mathbf{e}}}^{i}=\frac{{\mathbf{e}}^{i}}{\sum_{i=1}^{16}{\mathbf{e}}^{i}}$, and then ESS is given by

$$ $\text{ESS}=\frac{(\sum_{i=1}^{16}\hat{{\mathbf{e}}}^{i})^{2}}{\sum_{i=1}^{16}( \hat{{\mathbf{e}}}^{i})^{2}},$ (17) $$

which intuitively helps quantify the degeneracy of the sample weights in importance sampling. A low ESS indicates that a few samples dominate the weights, which can lead to poor estimation quality. Different columns of Figure 2 correspond to results for different EDLM parameterization. From the first two rows, we can see all EDLM implementations can assign meaningful energies to true data, good generated sample, and bad generated sample. For ESS, it indicates that the NCE energy function is generally better than AR ones for importance sampling. However, we note that we did not observe that EDLM-NCE is clearly better than EDLM-AR, and we conclude that this is mainly because our ESS analysis is conducted with 16 samples instead of an extremely large number. We recall here that GPU memory bound the sampling batch size to 16, and therefore large ESS analysis in practice is not useful.

Figure: (a) EDLM-NCE
Refer to caption: x4.png

Figure: (a) Generative Perplexity vs. Time.
Refer to caption: x13.png

Additional results with different sampling window length ${\mathbf{w}}$. We provide additional sampling results with different importance sampling window sizes ${\mathbf{w}}$. All results are provided in Figure 3. As shown in the figure, especially the reference line of AR in Figure 3(a), we can see that generally ${\mathbf{w}}\leq 0.6$ can all lead to better efficiency, among which ${\mathbf{w}}=0.2$ achieve the best speedup.

## Appendix D Generated Samples from EDLM

### D.1 Generated Samples from EDLM -NCE with 1024 timesteps.

…a lot. I remember when you was a kid, you felt like you weren’t a human. And I said, “You didn’t wanna get into that.” “Yeah, that’s just physically [grown-up], I kind of weighed down with that.\n\nWell, yeah, trying to help him, but it was something to give. Not the little stuff. It was tough to get from an older person. It just kinda won. That’s the way with things like that.\n\nButyou had to learn something else, more of appreciation of it, watching guys play [inside out] as you grew. Everything they showed up you knew was always in the way they wanted, as just as your parents wanted you.\n\nAnd[my mom] when I was yelling at her: “But you get the opportunity to work something for each other, and you can cook us meals as well.” “Yep.\n\nWehad to find some way out. It got borderline ridiculous, but we found a way.\n\nHeand I speak and we talk a lot. That’s really the thing about us. Just the two dogs at the right time.\n\n“Not funny man. Just like we are, baby. We’re not stupid.”\n\n“What happened like when you’re getting bumps?”\n\n“No, I cant́ see it from that side of the eye. I wish I could. I cant́ see it on my face. Yeah, I dont́ really. I picked it up and it’s not the right way. But the way I got it.”\n\nMarcusCousins\n\nIthink I’m not confused. He was so close to me. But I’m never confused either.\n\n“Are you playing the same type of game now?”\n\n“Always the same way. You play it at the same time. You don’t have the same way before you play the game. And I got used to it. He had a good season. He was a defensive superstar.\n\n“Sometimes we would go through the type of moves when I’s saying things, be like, ‘Oh, he’m cheating on me,’ or something like that because we just know like that. Lord, you know, stuff. That’s part of this life. Like my dad and I had a little bit, and that got you sometimes it didn’t work to the liking. [You] would physically get him.\n\nWedid, but he always demanded that I catch up to him. And when I tell him, “Catch me if you say something stupid, you’re not stupid.” If he does say stupid something, you know, I’m going to respond to it.”\n\nIt’s like “Whatś going to do catch ńít? If he’s gonna catch it, thatś my one.”\n\nAnythinglike that?\n\n“Itś funny because a lot of people have a playing style that plays the player they are. He’s a caricature of what it is, but he’s a kind guy and it affects the game.”\n\n“That’s a good one.\n\n“His skills that he play, can he tell you?”\n\n“Yes. He’s an unbelievable dude.\n\nYeah, he’s not the guy that everybody goes to play with, just out on the team. He plays stuff because he takes the game seriously. But you see a player that goes and develop that’s easy to feed off of, and that’s good. They win games, too, by being a teammate, in that environment.\n\nButthat’s all.\n\nHonestly, it hasn’t been really easy for me. I played 35 percent of the time and these guys ask a lot about me [laughs]. I told you guys those guys. I know that when I was young I would not have made it. I was just born that way when I was there. So dont́ call me the only man that I really liked or the backbone of the team, ever. That guy is the person in the locker.\n\nWhenyou’re all of the guys in your team, you don’t have to make credit for what you do.\n\nButI definitely miss Tom or Cole. Everybody does that to me. I love Chad Mac. It’s the Strict style. Heś always teaching me on being an open guy.\n\nThatkind of stuff, what I like the most. If I lose a lot of games,…

### D.2 Generated Samples from EDLM -AR with 1024 timesteps.

…have to make you better. But all you have to do, you can go talk about that with him. If you want it to accomplish that from his position as the coach, then you definitely feel that as a manager will have a tremendous benefit as a player, and that’s very important, that is, why hockey. That’s what you want to achieve as an organization. So, when you get to manage, where the truth is that you’re better as a player and the more points you score and the more passes you have to make and make plays more often, you feel that it’s enhanced your effectiveness as a player.\n\nRNS: And you can do that?\n\nJPP: Totally. I just think you can. So, you’ve gotta know, if you’re going to play that role of a manager, you’ll be the one that knows where to go. And that’s how great it is, every year.\n\nRNS: If it’s difficult as a manager in your last one, did you find yourself right about that decision? Did you feel better?\n\nQ“Well, I sat there the other day thinking, ‘OK, I have to do this job to get a career in that position.’ So when I do that: one-thousand up, no down. I could not go on. I didn’t do what I need to say about that and I can’t help but be proud of the success that I have created.” And I think, “You just gotta realize that you’re not supposed to truly have your own voice.” “I can’t do that’ ‘I let the people out of my decisions,’ then I think you are one of that.\n\nButyou can go through that. Just to live in this is to live in having to really utilize this and offer this to you. So, if you’re to move out of this job, where you want to move out of this job when you are a player, where maybe when you gotta be manager, when you gotta do business, or when you are a hockey player, you decide which is different. So we’re the same positions but we talk about how to do it.\n\nAndjust, more and more we’ve said: that’s how you’re gonna get there. If you want that manager role, you should do this. Because you’ve got to be passionate about your passion for your side, frontend your team up and do everything that takes to build this team as an organization. And they’re players. So you don’t get to go around and have to do that. If you don’t do that, you’ll have a lot of control of your players on the team. You also’ve got to help them and show them what they’re doing, to be involved. To take care of and make sure that their players succeed. So, you have to be productive within the team. That’s two things.\n\nRNS: So what level of coaching is interesting and what do you hope to get out of all of those roles and positions?\n\nJPP: Yeah, you never know. And you know how they’re actually putting these things. So when you start a coach, you’re in to young players, you’re in with multiple players every day, you coach to the players who are currently in a role in an organization, and obviously that helps build them. So basically, your really good coaching means you’ll wear your jersey, be the best player you can be from your experiences. I always tell people after you’re done doing the things, you trust me that it’s the most important to take place next.\n\nSo, when I started coaching, I do not think that was a great leap because that was my first time in the NHL. But I would like to tell the coaches, “You’re gonna stick with your team, you’ll lead the team. Let’s be with the best,” and I assume, the next coach, you’ll be the one and you’ll be the one and that season you’ll be the one. And that’s what you gotta think about is that if you’re more involved and in unique ways, if you’re learning from your experience and will be successful as importantly, just, you know, put their teams on the table. That’s part of what we are trying to do. And I’d say I’d like to put people on that front that are strong and fit for the job. So that it…

### D.3 Generated Samples from EDLM -coAR with 1024 timesteps.

…to explain that to someone who has not shared information with Google itself and not the powers that be. And Kim didn’t.\n\n“I never wanted to do this because before it, I didn’t know anything. Finally, you feel as if you do.”\n\nIlaughed.\n\n“It didn’t feel right.”\n\n“But yeah, we stayed together for a while,” Kim said, “and when we had time for us to, we were talking.”\n\nIwas almost in silence. “I’m not able to act out that way,” she said, “so I just kept having my time. One day when I was traveling and rehearsing with my own band, we had a really funny meeting. It’s funny, it was just plain funny.”\n\n“That day, we were talking to other girls friends and friends and I had my first year-long meeting with social media guys. As a women’s boss.”\n\n“I feel like different company here,” she said. “I mean, it’s basically an out-of-work company.”\n\n“I had to pick up my car at the same time when we first met.”\n\n“Oh yeah. It was weird, because I did a lot with her, and she didn’t know my love for her while I was in her way. Odd. She was sitting right right next to me at my cousin’s house, hanging out with friends.”\n\nWhenshe was speaking to North, she felt in touch with friends and family in Marlboro where I lived, where she raised me and how we still hadn’t familiar with each other since as far back as I went and saw how she had got confidence over him when I was 18. I didn’t, too, but she said he took everything else into account. Even with Wests death this morning, I can still have a private conversation with Kim with everyone. Some conversations on the phone and for the other girls.\n\nKimleaned in very talkily and he said to her, “We don’t like you much,” but he wasn’t trying to make California’s jealous. “You’re the one if you’re engaged to somebody and if you see someone else, it’s probably never good for you, so don’t go out to anything that’s not good for you.” And that’s all too true.\n\nIfirst asked Kim West about his faith in him as a woman’s friend, what he thought of moving to Marlboro and how that goes him down.\n\n“I’m a guy that is always doing something great,” West had said. “And then I’ll come up short, do well myself, and I go, ‘I really want someone on this,’” she said.\n\n“That’s sad,” Jörg smiled, “this kind of guy who’s not that smart and not do well. I mean, you looked the fuck you down and you figured out that’s more important to me than you thought you were, basically. You wanted to do something.”\n\n“It’s sad,” he said. “That was all I had my years trying to figure out about you so that I could link you and spend time with somebody who’s gonna roll.”\n\n“I could just read his shit and die,” she said, gesturing to West’s letters. “And that’s how I was left behind. It’s always struck me that I didn’t know how I wanted to spend with this group or another.”\n\n“Had you dated her in the past?” she asked, asking if he had more of a problem with her? “No, not really,” he said. “I’m a nice guy and I would talk back to anyone and that is why I actually talk to people.”\n\n“That would be a really creepy place to be in,” he told her, adding, “it was like…the first time I’ve lost a friend.”\n\nTherewas blood in the room. “Lemmer, he was a funny dude at the time.”\n\nHesmiled. “But you know that. This way a happy friend would have been if I had him back then. I didn’t care for him. …”

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: - Is it correct that the number of parameters used for generation is approximately doubled? While generation time is important, it would also be useful to know the VRAM requirements for generation.
- Is it possible to apply the importance sampling steps at a different interval, i.e. at the middle of the generation rather than at the start? This adjustment might yield improved performance.
soundness: 2
strengths: - The paper is clearly written and easy to follow.
- The method is intuitive and based on a solid mathematical foundation.
- The experiments follow the standard setup for evaluating Diffusion Language Models, making it easy to compare with other methods.
summary: This paper presents a new Discrete Diffusion Model that models an energy function $E_\phi$ to improve sampling procedure from an existing Diffusion Model. The authors propose a novel sampling procedure that incorporates importance sampling steps during generation.
weaknesses: - This method requires a pre-trained discrete diffusion model, which increases the overall computational requirements. Thus, it may be unfair to compare it directly with simpler methods like MLDM.
- While the proposed method reduces the Gen PPL metric, it also decreases the entropy of generated texts. One could even argue that it produces similar results to MLDM in terms of Gen PPL and entropy.

**Recommended Experiments**:

It would be interesting to see a more detailed trade-off between entropy and Gen PPL compared to the baselines and the proposed model. Although Figure 1 provides some insight, a separate figure with Entropy/Gen PPL axis would address concerns about the notably reduced entropy.

As mentioned, diffusion models are not as computationally efficient as simpler AR models. However, it may be possible to improve efficiency by allocating different computational budgets for training the base $p_\theta$ and $E_\phi$ models. Given the potential compute intensity of this experiment, a simpler approach might involve training the existing MLDM for the additional steps used to train the energy function  $E_\phi$ , then reporting the evaluation results.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 4
questions: 1. For EDLM-AR and EDLM-coAR, which pretrained AR model is used and what is the parameter size of the model?
2. For all EDLM variants, a separate transformer model is used to calculate the energy, which effectively doubles the parameter count and might contribute significantly to performance gains. While the reported performance and efficiency improvements likely go beyond what would be expected from simply doubling the model size, it would strengthen the validity of these findings to include an empirical comparison or an in-depth discussion with a baseline with an equivalent parameter count.
soundness: 4
strengths: - This paper is well-written and well-organized, presenting a simple yet elegant combination of energy-based and discrete diffusion models. EDLM demonstrates strong empirical results that match auto-regressive models with great improvements in sampling speed.
- The adaptation of auto-regressive models to a joint denoising distribution with masked inputs is innovative.
- The application and detailed analysis of importance sampling windows effectively improve the early sampling phases in discrete diffusion models, making a nice empirical contribution to the advancements of discrete diffusion language models.
summary: This work presents a new discrete diffusion language model that uses an energy-based model (EBM) to introduce token dependencies within the text sequence, instead of factorizing the denoising distribution over individual tokens. In practice, this EBM is parameterized in a residual form, with the energy function implemented by either a pretrained auto-regressive Transformer or a fine-tuned bidirectional Transformer. To facilitate efficient sampling, the paper employs self-normalizing importance sampling and draws several samples followed by resampling based on energies to complete each denoising step. The resulting framework, called Energy-based Diffusion Language Models (EDLMs), demonstrates significant advancements over previous discrete diffusion baselines and achieves comparable or better performance to traditional auto-regressive language models on text generation.
weaknesses: - While the study aims to address the independence assumptions in discrete diffusion models through EBMs, there is insufficient examination of relevant prior research also involving EBMs for language modeling (e.g., [1]). Given the similarity between EDLMs and [1], a more detailed discussion and comparison would clarify the position and relevance of this study.
- Another concern lies in the significance of applying EBMs to discrete diffusion models. Although vanilla discrete diffusion processes do factorize the denoising distribution over tokens, the bidirectional attention mechanism in transformers can already capture dependencies among tokens within a sequence. Therefore, to account for errors due to parallel decoding, feeding the decoded sequence back into the transformer (i.e., the next denoising step) could potentially identify erroneous tokens and assign low likelihoods to those positions. There is extensive literature on filtering and remasking tokens at each denoising step [2, 3, 4]. An in-depth discussion highlighting the advantages of using an EBM to capture token dependencies would strengthen the paper’s argument.

[1] Deng, Yuntian, et al. "Residual energy-based models for text generation." *arXiv preprint arXiv:2004.11714* (2020).

[2] Ghazvininejad, Marjan, et al. "Mask-predict: Parallel decoding of conditional masked language models." *arXiv preprint arXiv:1904.09324* (2019).

[3] Savinov, Nikolay, et al. "Step-unrolled denoising autoencoders for text generation." *arXiv preprint arXiv:2112.06749* (2021).

[4] Zheng, Lin, et al. "A reparameterized discrete diffusion model for text generation." *arXiv preprint arXiv:2302.05737* (2023).

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: Section 4.3 is not very clear to me. Could you provide a detailed explanation of how you estimated PPL?
soundness: 3
strengths: 1. Using an AR to help the sampling of Masked Discrete Diffusion Language Model is natural and straightforward.
2. The experiments demonstrate good improvements compared with MDLM baseline.
summary: This paper proposes using an autoregressive language model to help the sampling process of a Masked Discrete Diffusion Language Model (which is technically equivalent to a Masked Language Model) by rescoring a selected subset of the generated outputs. The authors further propose to train this model with noise contrastive estimation and find a method to estimate its PPL.
weaknesses: **My main concern is that the core technique introduced in this work has been present in the literature for a long time.**

1. First, the framework of the absorbing discrete diffusion model is essentially the same as the BERT-like masked language model (MLM). The forward process corresponds to masking tokens in the input, while the backward process corresponds to predicting and remasking tokens during iterative generation in MLMs. Therefore, this work primarily explores how to apply an energy-based framework to assist MLMs in generating text.

2. Importantly, the core part of the proposed method is using an autoregressive model as an energy function to guide the sampling process of an MLM (referred to as the denoising model in this paper). **The introduced energy function ($−logp_{AR}(x_0) + logp_θ(x_0|x_t)$) and the Importance Sampling method (Algorithm 1) are identical to reranking (or noisy parallel decoding) methods used in MLM generation/non-autoregressive generation, which have long been established in prior literature.** Examples can be found in [1,2,3]. Actually, the practice of using an AR model to rescore multiple outputs from an MLM has long been a standard baseline in non-autoregressive generation tasks.

[1] Non-autoregressive Neural Machine Translation, 2017  
[2] Mask-Predict: Parallel Decoding of Conditional Masked Language Models, 2019  
[3] Fully Non-autoregressive Neural Machine Translation: Tricks of the Trade, 2020

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: - What are the used window size and sampling size in Figure 1(a). I am somewhat confused about the 1.3x sampling speedup as EDLM-AR/NCE requires more computation on the intermediate steps. Could you please elaborate on the rationale behind the increased efficiency of EDLM-AR/NCE compared to EDLM?
- In the case of EDLM-NCE, it seems that when the value of t is large, the quality of multiple sampled $x_0$ might all degrade, leading to general closer low scores assigned by the classifier. Given this, why does adding the reranker in the early stages yield better results instead?
- In the context of EDLM-NCE, do you employ separate models or do you utilize a shared rerank and generator approach?
- I am somewhat confused about Figure 3. In my view, as the window size increases, one would expect a higher time cost. However, the graph appears to show the opposite trend. Could you provide insight into this discrepancy?
soundness: 3
strengths: - The writing is clear, making the paper easily comprehensible.
- The idea to inject sequence-level correlation in the intermediate denoising steps is intuitive.
summary: This paper points out that the existing discrete diffusion models aim to predict all missing tokens in parallel at each intermediate diffusion step, but the denoising joint distribution is simply parameterized as a product of token-wise independent distributions.
The authors propose Energy-based Diffusion Language Model (EDLM) by using an energy-based model to act as a reranker in the intermediate denoising steps to inject sequence-level correlation. Experimental results demonstrate that EDLM outperforms existing diffusion models and achieves sampling speedup without sacrificing generation quality.
weaknesses: - The primary contribution lies in the energy-based reranker, which mostly follows the approach presented in https://arxiv.org/abs/2004.11714. In my view, I tend to have concerns about the novelty given previous work.
- The description of the motivation and the independent factorization in Equation 6 does not seem entirely accurate. For line 156 "$\mu_\theta$ predicts each token in x_0 independently", it may be more precise to state that the sampling operation over the output distribution is done independently, as during training, predictions for all masked token positions are made simultaneously, leading to their hidden representations being interrelated. This implies that predicting a particular token will implicitly involve using representations of other tokens, which makes the the independent factorization in Equation 6 not entirely accurate.
- This paper only evaluates perplexity, and it remains uncertain whether such energy-based reranking translates effectively to downstream tasks.
- I tend to have concerns about the inference cost linked to repetitive sampling and energy function computations in each denoising steps.

Meta-review: This paper introduces Energy-based Diffusion Language Models (EDLM), which addresses token dependencies in discrete diffusion models by incorporating an energy-based model for reranking during the denoising process. The paper's main strengths are its clear presentation, strong empirical results showing improved performance over baseline discrete diffusion models, and the introduction of an innovative sampling procedure that combines denoising with importance sampling. The key weaknesses include concerns about novelty given similarities to previous work on energy-based models for text generation and non-autoregressive translation, potential computational overhead from repeated sampling, and initially limited evaluation on downstream tasks. The paper is recommended for acceptance primarily due to its technical soundness, clear improvements over baselines, and practical significance in advancing discrete diffusion models for text generation, even though some reviewers initially had concerns about novelty.

Decision: Accept (Poster)