Energy-Based Diffusion Language Models for Text Generation
Authors: Minkai Xu
Tomas Geffner
Karsten Kreis
Weili Nie
Yilun Xu
Jure Leskovec
Stefano Ermon
Arash Vahdat
Venue: ICLR 2025 Poster (2025)
Keywords: Language Models
Discrete Diffusion Models
Energy-based Models
Primary area: generative models

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.3x sampling speedup over existing diffusion models. Reproduced code is available at https://github.com/MinkaiXu/Energy-Diffusion-LLM.

Published as a conference paper at ICLR 2025
ENERGY-BASED DIFFUSION LANGUAGE MODELS FOR TEXT GENERATION
Minkai Xu Tomas Geffner", Karsten Kreis", Weili Nie², Yilun Xu², Jure Leskovec', Stefano Ermon', Arash Vahdat Stanford University PNVIDIA
ABSTRACT
Despite remarkable progress in autoregressive language models, alternative gen- erative paradigms beyond left-to-right generation still being actively explored. Discrete diffusion models, with the capacity for parallel generation, have recently emerged promising alternative. Unfortunately, these models still underper- form the autoregressive counterparts, with the performance gap increasing when reducing the number of sampling steps. Our analysis reveals that this degrada- tion is a consequence of an imperfect approximation used by diffusion models. In this work, we propose Energy- Diffusion Language Model (EDLM), an energy 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 obtained by leveraging a pretrained autoregressive model or by finetuning 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 model can consistently outperform state-of-the-ar diffusion models by a signifi- cant margin, and approaches autoregressive models' perplexity. We further show that, without any generation performance drop, our framework offers a 1.3x sam- pling speedup over existing diffusion models. Reproduced code is available https: http//ggtt// com woom/iiiiee /t////iiiiaiiii wawnnaaaii/e/rgyaaaaaae oaffffuiion LLM.
INTRODUCTION
In recent years, autoregressive (AR) models have achieved remarkable advances in modern large language modeling (Vaswani et al., 2017: Wolf et al., 2020; Brown et al., 2020; Chowdhery et al., 2023; Touvron et al., 2023), showing impressive results in tasks such as expert-level coding (Roziere et al., 2023) and reasoning (Wei et al., 2022; Trinh et al., 2024). Despite the considerable progress achieved, the left-to-right generative paradigm exhibits several long-standing drawbacks. For ex- ample, training and sampling of autoregressive models require fixed ordering of the sequence, constraining their generation flexibility (Hoogeboom et al., 2022: Shih et al., 2022). Furthermore, at test time, the generation is heavily conditioned previous generations, which leads to accumulated error a.k.a exposure bias (Bengio et al., 2015). As a result, alternative generative paradigms beyond left-to-right generation actively being explored.
Discrete diffusion models (DMs) have recently emerged as a promising competitor for discrete se- quence generation (Austin al., 2021; Dieleman et al., 2022: Gat al., 2024). Unlike AR models, discrete DMs conduct generation by progressively decoding the full sequence in parallel starting from fully masked sequence, offering great potential bidirectional controllable generation and sampling acceleration. Most recently, great efforts have been invested in improving their perfor- mance by extending discrete diffusions to operate in continuous time (Campbell et al., 2022: 2024; Sahoo et al., 2024; Shi et al., 2024). Unfortunately, despite their promise, these models still under- perform AR models by a large margin, resulting limited practical usage. Developing principled and powerful discrete DMs remains an open problem.
In this work, we take closer look at the existing discrete diffusion models and reveal vital mis- match issue between the training and sampling distribution in its current form. Specifically, the
Work done during internship NVIDIA. Corresponding author: mi nnk Ees. staan edd..

---

Published as a conference paper at ICLR 2025
existing discrete diffusion models aim to prediet all missing tokens in parallel each intermediate diffusion step, but the denoising joint distribution is simply parameterized as a product of token-wing independent distributions. As result, intermediate denoising steps ignore sequence-level correla- tions, 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 Lan- guage Model (EDLM), an unnormalized energy-based model (EBM) that learns to jointly denoise the full sequence each diffusion step. Our key innovation is to learn an EBM for each denoising distribution p(x(-1(x;), 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 a novel residual form
(X = Paiffussion (xptiassionnx41--000 exp(
(1)
applied over pretrained diffusion models Paiffusion- Such formulation enjoys several distinct advan- tages. 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 leverag- ing pretraiend AR as the energy function, approach can be interpreted as parallel sampling from pretrained language models using diffusion models as the proposal distribution.
EDLM can be viewed new family of discrete generative models that marries energy-based and diffusion-bass models. The novel combination addresses the training and sampling distribution mismatch problem, enjoys better generation quality with less accumulated error, and improves sam- pling efficiency by reducing the number of sampling steps. We further provide a formal estimator to calculate the perplexity of EDLM, allowing our model be compared against other models in a standard way. We conduct comprehensive experiments on two common language modeling benchmarks evaluate the performance of our proposed method. Results show that on the per- plexity metric, EDLM can consistently achieve state-of-thh-arr performance among diffusiin-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.3x sampling speedup when keeping the same sampling performance.
2 RELATED WORK
Diffusion models are powerful models with surprising results in generating high-quality im- ages (Sohl-Dickstein et al., 2015; Song & Ermon, 2019; Ho et al., 2020; Dhariwal & Nichol, 2021). These models were originally designed for continuous data, with both forward and backward re- verse 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., 2021; Hoogeboom et al., 2022), with Campbell et al. (2022); Zhao et al. (2024) further extending the framework to model discrete data in continuous time. Most recently, several concurrent works (Lou et al., 2024; Sahoo et al., 2024; Shi et al., 2024) archived impressive progress in large-scale language modeling by scaling the model and simplifying the training and sampling processes. Additionally, Campbell et al. (2024); Gat et al. (2024) developed flow matching methods for discrete data, which rely on formulation similar 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 mis- match 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 small number of denois- ing timesteps. Deng et al. (2020) studies similar accumulated error problem (a.k.a exposure bias) in the autoregressive model setting and proposes energy-based models for modeling global con- text, but the methodology focuses on improving autoregressive models and therefore fundamentally different from ours. Other work (Gu et al., 2018; Ghazvininejad et al., 2019; Gu & Kong. 2021; Savinov et al., 2022; Zheng et al., 2023) study similar non non-autoregressive text generation but com- bined with reranking and/or remasking methods. However, these methods typically result in biased sampling, which can generate high-quality samples but fail to model the distribution and suffer from
2

---

Published as a conference paper at ICLR 2025
low sampling diversity. Therefore, these methods mainly focus different tasks such machine translation. Notably, Lezama et al. (2023) notices a similar issue (focusing on diffusion models in the image generation domain) and proposes to learn corrector network to correct the independent decoding error. However, this work concentrates on the image generation domain, and the correction process needs to be sequentially, which takes additional inference time. In contrast, propose an energy-based denoising process that directly captures the correlation at each denoising step and enables efficient parallel decoding via importance sampling scheme.
DISCRETE DIFFUSION MODELS
General Discrete Diffusion. Assume data lives in finite discrete space of size m 1. In this paper, we augment the categorical space with an additional mask state with index m. In general discrete diffusion framework (Austin al., 2021), the diffusion process defined Markov chain ((x,\x{-1) - Cat(x;; Qix;-1), which repeatedly multiplies x with matrices Q, over T' discrete time steps. Given these transitions, the marginal distributions each timestep can be written in closed-form as g(x, x) = Cat(x;; Qxx) = Cat(x;; Qi Q1x). Such forward process can also be viewed as an interpolation between a clean data sample x and a reference distribution Cat(:; x) induced by Qr:
(x, x) Cat(x;; oxXu + (1 (4)
where E [0, 1] is strictly decreasing function w.n.f t, with ag ~ and a1 R 0.
In the continuous time limit, for two arbitrary times 0 < s S t S I, the transition distributions can be written as q(x,/x,) = Cat(x;; X + (1 )T). where Otzj at/a, (Zhao et al., 2024; Shi et al., 2024). This implies that during each diffusion step s the token will jump to a sample from the prior distribution и with a probability of (1 on The forward process allows us to compute many distributions closed form. One particular distribution interest is the reversal of the forward process given Xo. that is, the posterior distribution given by
(xx, xb) = Cat Xx: [artaxe + (1 (t 1. T @ [a,x0 (1 as)x)
(3)
0x X (1 Oz )x,
Masked Diffusion. In this paper, we focus on masked (i.e., absorbing state) diffusion models, where the target distribution set as æ = m. each diffusion step each token transitions to the 'masked' state with some probability. Under such masking framework, the forward marginals (Eq. (2)) are given by ((x)(x)) = ax + (1 - 0g )m, and the posterior (Eq. (3)) can be simplified as
Cat(xx;;x))
X + m,
q(x Xx, xo)
=
Cat Xxi (1-a,,m+
(4)
1-ac
X m.
Diffusion models aim to learn a backward model (o(xxlz;) to approximate the reversal of the for- ward process. Leveraging Eq. (4) we can parameterize the model pe (x, (x:) as
(Ctt(x;;x:),
X4 m.
Po (x,(xxx) = q(xx, Xx, Xo = (x,, t))
Cat X3; (I-a))++aa (a, (a
(5)
1-a;
X
where ps = Po (x) is known as Xy predictor, since it predicts the mean of the distribution over Xo given xt. However, explained next, the practical implementation of this x. predictor results in a mismatch between training and sampling distributions.
Problem Statement. Existing discrete DMs learn denoising distribution μn px (xx) match the true reversal g(x,|xx) Let x denote the tokens a full sequence in this section. In practice, the model is parameterized factorized denoising model:
Po (xo x,) (x Tx,,( t), where no (x0,
m, X₁ A m,
xx
(6)
and the predictor μ4 factorizes each token in X₂ independently. This factorization enables us to conduct efficient denoising step po(x,|x,) |x. by first sampling all Xoje tokens from po (x (x,) in
3

---

Published as a conference paper at ICLR 2025
parallel and then masking certain tokens according the forward (xx,]/x,,
Xo However, this parameterization ignores dependencies between tokens in the sequence, a fundamental limitation which implies that (xooo/x,, (x,) can never match the exact backward (x,) As result, this parallel sampling introduces accumulated errors, as the factorized denoising step Po does not match the original generative model po for the joint distribution of the elements of Xp.
4
METHOD
In this section, we formally introduce Energy-based Diffusion Language Model (EDLM). Our work aims to design a new family of energy--aasse based discrete generative models that address the fundamental mismatch problem between pa p*(xo (xo]x{) and g(xo|xx (x;) in existing models. We first describe the general EBM formulation in Section 4. and then elaborate on how to obtain the energy function by either leveraging pretrained AR models or finetuning via noise contrastive estimation Section 4.2. Then 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 Xo predictor for each denoising step For example, original discrete DMs learn (xxx,, t) to directly predict independent distributions for each token in Xo- our framework, given diffused data X₁ at timestep t, we introduce the generative denoising kernel as unnormalized density:
(x. = (xo) exp( Es Ex x,, t)) (x(x))
where pa is the pretrained diffusion model, Eg the energy introduced to capture the correlation in the No sequence, and (xx) is normalizing factor known as the partition function. In the following text, we use Pa,ş to denote the joint model, Eg the (residual) energy function, and keep the pretrained p4 fixed. Computing the partition function is intractable since it requires summing over the whole x space, which exponential in the sequence length. our language modeling experiments, the vocabulary size is around 50k and the generation length is 1024 tokens, resulting space size around 50, 2004024, 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 g(xo (x.), while avoiding computing the partition function.
4.2 IMPLEMENTATION OF ENERGY FUNCTION
Training energy- models with the intractable partition function long-standing challenge in machine learning (Hinton, 2002; Carreira-Perpinat & Hinton, 2005; LeCun et al., 2006). Typ- ical 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 's parameters for the xo predictor (x.+(xo|xx)) xx.) In the following, 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.
describe two methods to do this. One involves taking
4.2.1 LEVERAGE PRE-TRAINED AUTOREGRESSIVE MODELS
Let PAR(Xo) be autoregressive model trained over clean samples Xp, Lee, PXR(Xo) IRPAR (x|1x, ")). Since AR models are only defined over clean samples Xo. it is not directly clear how to leverage them in denoising tasks, since the models are never trained on diffused data X,. Our key insight is that, in absorbing discrete diffusion. the diffused x, is the same as Xo with certain tokens masked according to the forward process. More specifically, letting Xo = X2 [x, # m] and Xo/ Xo denote the set of Xo tokens corresponding to unmasked and masked components in X
4

---

Published as a conference paper at ICLR 2025
observe that can induce the denoising transitions using
X+) = PAR Xo) PAR(Xo, Xo. xo)
PAR (X0)
PAR PAR(Xo
(8)
PAR (Xo)
PAR (xo)
where the normalizing constant is given by PAR (xo) = Exe/ / PAR (xo/ X0, Xo). This partition function also involves a sum over the X space and is intractable to compute, However, noting that in the reversal of masked diffusion (Eq. (5)) unmasked tokens X, are fixed, we have that PAR xx) x
PAR (Xo). which can be computed efficienctly.
Unlike the original diffusion models (Eq. (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, PAR (xo x,) ((xo xx). The key challenge is that sampling the autoregressive model PAR [xo Xt) at each denoising step is com- putationally expensive. However, it is efficient to evaluate PAR(Xo), which is proportional to the posterior (xR(X)) xx), and thus can be used in our residual energy-based based formulation.
exp ( (xo.x.)) to approximate
Recall that our residual energy model (Eq. (7)) uses pa xx) ((xo xz). By taking PAR(xo x₂) as the proxy of g(xo x,), we approximate the optimal resid-
(x(
ual energy as
E (Xo: xt) = - -lo x.) + po xx) log 7
log xx) log (xo xx) - log Z
q(xo Xx) pAR(xo|xx)
- -ooo Xx) lx xx) + PAR PAR (Xo) - log Z
See Eq. (8).
New Partition Function
(9)
For common EBM model inference, such as MCMC, we only require energies up to a constant. Therefore, as shown in the equation above, with pretrained AR model we can simply define the energy function as log PAR(X,) + log (x xx). 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 pa (xo xx) 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 the proposal distribution, providing 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 X; to Xg, i.e., ((x,)x,) = Cat(;;,) for x, * 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 fsf network to simply copy the unmasked tokens of Xs. Similarly, in our AR EBM model, we can also carry over the unmasked tokens in each denoising step, by directly setting PAR (ax(xx/x x55 = x for x, * m. This "carry-over autoregressive" (coAR) EBM implementation allows simply computing POOAR (xo xx)= = PeaaR
POWAR (xa) by just PoOAR (Xo), as we always carry over those tokens from x, and therefore PooAR (Xo) 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, 2010), specifically its conditional version (Ma Collins, 2018). First, NCE training relies contrastive samples from the data distribution and 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 Xo and diffused data Xt, we set true posterior q(Eo|x xo)² as positive
This property does not require the autoregressive structure PAR, but only relies on the fact that during the diffusion reversal process unmasked tokens do not change. In this subsection, we notation &o instead Xo to distinguish the generated data and clean data Xo-
5

---

Published as a conference paper at ICLR 2025
distribution and the denoising distribution po (x) the negative distribution. Thanks to the
residual energy formulation from Eq. (7), the log-odds reduce to Po. Po.s log pe = -- and the objective is simplified into the binary classification objective
ENCE (d: 8) -Exo Pux - ((x(x
1
1
Ex xx) log
+ Ex (xoooxx)))) log
+ exppex((x))) (x4, Xt,())
1 1 expl
Ex (x_ Xx,
1)
(10)
where x is positive data from the true posterior and is negative data from diffusion model, given the diffused sequence X,. Here, the true posterior g(xo|xe, Xt, Xn) := Xo is defined as simply recovering the true data. Training the energy function can be viewed as training 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 X generations, assigning negative energy to real data and positive energy to data produced by the denoiser network. As a result, the joint model pa.á 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 AOO BLACKWELLIZED LIKELIHOOD BOUNDS
A common protocol for evaluating discrete generative models on language modeling is perplexity (PPL), which relies computing oog-likelhhoods for hold-out test dataset. In this section, explain how to estimate the log-likeling lieeehood with EDLM. First, with the energy-based Xy predictor pa.s (xo xx), the step denoising model (x, is given by
po (x.o(xx(xx) = Epo, (xolx,) (x,,) (x,//xx xo)
0g Cat (X,; m) 0g a4 Pa. (xo(x, + m] x, x,(x, * ml|xx)
(1I)
at
For the simplified reverse process in Eq. (5), the log-likelihood can be simply computed form of Rao-Blas Blackwellization. The BloBBaaewellized likelihood bound can be estimated by the discrete-
time diffusion loss of finite T: Ldiffusion EDLM, the loss be written as:
ETT [KI. (g(x., X x)]lpo.d (x.,xx)))) Under
T
Ediffusion Σ Eg
at as, (log Po
(x0: X Xx: tx) log Z Zo(xx:))
(12)
att
where po (xo(xx)) = (X,, x.) is simple cross-entropy between μs predicted logits and the clean data. We can extend the objective Eq. (12) to the continuous limit by taking T - oc, which induces the following negative evidence lower bound (NELBO) Loo:
df 1-
C.
EEg
(logPa(xol - Ex(xo, Xx, t) log Zo (x:)) dt.
(13)
where af denotes the derivative of ar w.rl I. Such ELBO is tighter than the discrete-time version, and invariant the noise schedule. However, this bound requires the partition function log (x,)
log Ex, po (x(o)(X)) exp(- (x0, (x)x = log Ex.<po exp( - (Xo. x²)) and thus is intractable to compute. We use two estimators based on the work (Nowozin, 2018; Deng et al., 2020) estimate the partition function log Z
Theorem Given diffused data X, at timestep I, log 2n denote the empirical estimation of log (xx) = log Exo~po exp( Exp(-EEgo exp(-E,(Xx,) (xxxxx)x))) x:)) with samples (i) ~ poli = -I,...,, x)):);
,
log 2, log EL exp(- E(x) (i)
x)) Then Ve O. EN such that Vn. N we have
log Zz < Eloog 2n) < log Zz < E[ 1) log 2, 2(n 1) log 2n-1) < log Z +. (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 is sufficiently large. In practice, we follow Nowozin (2018) to use the leave-one-out strategy to estimate 2,-1, which is proven to yield an estimator with lower variance. In addition, want recall that in autoregressive energy function instance (section 1), we introduce the carry-over parameterization, where the joint energy model is self- normalized, which enables us exactly compute the ELOO without estimating 7g-

---

Published as a conference paper at ICLR 2025
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, 2000; Hinton, 2002) were originally applied to bi- nary 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, 2013; Hammersley, 2013). With our joint model as the product of the diffusion model po and residual energy function Eg, and given intermediate diffusion data X₃ at timestep t, we can conduct effi cient parallel sampling by: 1) sampling multiple Xa predictions {xo from the diffusion denoiser po(xo]Xx x)); 2) feeding samples into energy function in parallel compute energies; and 3) resam- pling single Xo from the pool (x6) according to the energy values. The sampled Xp is then fed into the posterior formulation ((x, Xo) perform one reverse step (i.e. one-step denoising).
Importance Sampling Window. While im- portance sampling introduces additional com- putation, it yields significant reduction in the parallel decoding error. Therefore, the
Algorithm Denoising via Importance Sampling
1: Input: discrete diffusion model 0, energy- based model ø, sequence of timesteps Ti > T2 > > TN-1, number of sampling size
method allows us to conduct diffusion sampling with fewer denoising steps, reducing the over- all sampling wall-clock time. To further ac- celerate the sampling speed, we introduce the concept of importance sampling window length
k importance sampling window w
2: X, f i
3:
for = to N - 1
4:
po (Xu x,,)
(
w E [0, 1], which sets the timestep for stop- ping importance sampling, i.e., only con- duct importance sampling in the time window
5:
if 2 1 w then
6:
{xn: xx¹ ~ pa XT.)
7:
Compute energies e¹ Eg (x)= XTn
I [1 w, 1] Interestingly, in our experi- ment, we notice that the energy-based impor- tance sampling during the early stage of de- noising sampling contributes more to the gen- eration 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 xo prediction, since there is little information on the full sequence. This encourages plore sampling with a short time window w for higher efficiency, which we discuss in detail in section 5.3.
for x ~ [x) x
8:
Sample xo ~ (x) ... x(i
with probability
exp( -ce)
Σ5 -ep)
9:
else
10: Sample Xo ~ (xo
11:
end if
12: X ~ g(Xrata 1 X0)
end for
to ex-
14: Output: x = XTN
Detailed pseudo code for the sampling procedure is provided in Algorithm 1. In practice, we ob- serve that relatively small importance sampling size 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 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. 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, 2006), a relatively small-scale, character- level text modeling benchmark extracted from English Wikipedia, and 2) OpenWepTee source replica of the unreleased WebText (Gokaslan & Cohen, 2019) dataset used to train GPT-2.
open-
Baselines. We compare EDLM against ofate-of-the-art discrete diffusion models and transformer AR modelS (Vaswani et al., 2017). Discrete diffusion baselines include Discrete Diffusson Model (D3PM) (Austin al., 2021), Score Entropy Discrete Diffusion (SEDI) (Lou et al., 2024), Masked
7

---

Published as a conference paper at ICLR 2025
Diffusion Language Model (MDLM) (Sahoo al., 2024) and MD4 (Shi al., 2024). We no- tice other recent works such as discrete flow matching (Campbell et al., 2024: Gat et al., 2024), but they are either not open-sourced focus non-language settings. Therefore, we take their concurrent work with similar formulation MDLM (Sahoo et al., 2024) and MD4 (Shi et al., 2024) as the representative for comparison. On the small-scale Text8 benchmark, we additionally eval- uate other discrete generative models including Plaid (Gulrajani & Hashimoto, 2024), Bayesian Flow Network (Graves et al., 2023), Any-order Autoregressive Models ARDM (Hoogeboom al., 2022) and MAC (Shih et al., 2022), and flow-based methods IAF/SCF (Ziegler & Rush, 2019), AR Argmax Flow (Hoogeboom et al., 2021). Discrete Flow (Tran et al., 2019), and Multinomial Diffu- sion (Hoogeboom et al., 2021). 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.
Metrics. We follow the convention (Shi et al., 2024) to use the common protocols Bits Per Character (BPC), Perplexity (PPL), and Generative Perplexity (Gen PPL) to evaluate gen- crative sequence models. For a sequence of length L, BPC metric is given by - Et=n loga p(x;) and PPL. is given by exp E log p(x,), which can be viewed as the aver- age 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 PPI. instead consists of likelihoods calcu- lated by another large oracle model on data generated by the evaluated models. Intuitively, PPL and BPC similarly eval- uate the likelihood modeling capacity, while Gen PPL evalu- ates generation quality and consistency. To compute Gen PPL we generate 2048 samples of 1024.
Table 1: Bits Per Character (BPC)
on Text8 test set.
Method
BPC (4)
Autoregressive
Transformer AR IAF/SCF
1,23
1.88
AR Argmax Flow
1.39
AR Discrete Flow
1.23
Any-order Autoregressive ARDM
1.43
MAC
1.40
Continuous Diffusiion
Plaid
1.48
BFN
1.41
Implementation Details. For all models including our meth- ods and baselines, we follow the common practice of us- ing standard 12-layer transformers similar GP'2--maal scale (Radford et al., 2019; Shi et al., 2024). Our proposed EDLM combines two models, the diffusion model pe and the energy function Ez. For all experiments, we use pre- trained MDLM (Sahoo et al., 2024) as the diffusion model pe- For AR-based EBM (see section 4.2.1. named E.L.--ARA or E.D.M-COAR when using cvery- we directly leverage the pretrained AR model as the energy function. For NCE fine-tuned EBM (see section 4.2.2, named
Discrete Diffusion
Mult. Diffusion
S 1.72
D3PM Uniform
1.61
D3PM Absorb
1.45
SEDD Absorb
1.41
MDLM
1.40
MD4
S 1.37
EDLM (Ours)
S 1.24
EDLM-NCE). we finetune the pretrained MDLM, with the energy function computed by project- ing the oncan-pooled last token embeddings down to 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., 2021; Lou et al., 2024), 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 Open WeeenWebexxt experiments next.
Teb We report the perplexity results of all models trained on the large-scale Open- WebText dataset in Table 2. We evaluate the model capacity for both the OpenWehText test set and
8

---

Published as a conference paper at ICLR 2025
Table 2: Test perplexities (U). 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
LMIB
Lambuda
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
Results of baseline methods unseen datasets are borrowed from (Sahoo al., 2024).
Table 3: Generative perplexities on unconditional text generation. LLAMA2. LLAMA3, and GPT2 (4) correspond to Generative perplexities evaluated by different oracle models.
Method
Timesteps
ILAMA22
LLAMARI
GPT2L
Entropy
7.0
14.7
Autoregressive
1024
22.9
40.3
35.7
8.1
SUNDAE
200
29.5
45.1
34.7
5.2
Ssd-LM
>+0000
99.2
D3PM Absorb
1024
692
754,9
842.3
7.6
SEDD
256/512/ 1024/ 2048
55.3/
33.7/36.2
618/5222/1.5/337
7.//7.7776/75
MDLM
256/512/
1024/
27.6//
668/524/1426/133
79/78/76675
EDLM-AR (Ours)
256/512/
1024 2048
34.7/26.87 19.6 14.6
62.2/444/
28.8/
20.8
2.1/420/25.5/73
7.9/76/72//69
EDLM-NCE (Ours)
256/512/
1024/ 2048
35.7 26.3/ 19.0/ 14.6
62.9 /4.1 / 20.7
61.7/42.5/555 24.52////55/77.
79//76/7.3699
Results of SUNDAF and Ssd-L.M are borrowed from (Gat et al., 2024).
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 al., 1993)). Wikitext (Merity et al.. 2016), LMIB, Lambada (Paperno et al., 2016), AG News (Zhang et al., 2015), and Scien- tific Papers (Pubmed and Arxiv) (Cohan et al., 2018). As shown in the results, we again observe that EDLM outperforms existing diffusion methods by 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 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 Open WebText, and results are presented in Table 3. We addi- tionally incorporate unrollnroleed Denoising Autoencoder (SUNDAE) (Savinov et al., 2022) and Ssd-LM (Han et al., 2022) 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 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 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.
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 emror. 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 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 accel- erate the sampling speed, in this experiment, set the importance sampling size k window

---

Published as a conference paper at ICLR 2025
Generative Perplexity
50
EDUM-NCE EDLM SEDO MOLM
POLM-NCE EDUMAI
2-
SEDD
Size
Entrapy
AR
40
7.8
35
16
30
1.0 0.0
0.6
0.2
Sampling Window Length
(a) Generative Perplexity vs. Time.
(b) Entropy Time.
(c) Importance Sampling Ablation.
Figure 1: Analysis and ablation study for EDLM. Figures Ia and 1b: we run AR baseline and diffusion-basee models with [512, 768, 1024] denoising steps, and plot the curve of corresponding metric vs. wall-clock time. For generative perplexity, the metric is evaluated by GPT-2, and a curve on the bottom-left indicates better sampling quality vs time trade-off. Figure le ablation study of EDLM under different importance sampling size and window length.
as w 0.2, which empirically yields the best generation quality with the same wall-clock sampling time budget. w 0.2 means we use importance sampling for t E [0.8, 1], and just use the original diffusion model jsg for denoising for t E [0. 0.8). We provide additional results with other win- dow size in Appendix 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 la and 1b. As shown in the figure, EDLM achieves better sampling quality vs time trade-off compared against MDLM. Specifically, when achieving a similar generative perplexity as the AR baseline (35.7 7), we ~+1 seconds while the MDLM baseline takes 1~ seconds, indicating a 11113 accelcration without any performance drop. This speedup further highlights the advantage of our energy-base denois- ing formulation and importance sampling scheme, where we can conduct sampling with reduced denoising steps and correct the decoding errors by parallel importance sampling.
from the figure that EDLM only takes
Effect of Importance Sampling. We further study the effect of importance sampling hyperpa- rameters in this section. We fix the number of denoising timesteps 1024, and investigate the generation performance with varying importance sampling size and window length. We refer read- ers Section 4.4 and Algorithm for details of these two hyperparameters, Again, found that EDLM-NCE and EDLM-AR exhibit similar trends in the study, and therefore show one study re- sult here. The ablation study results are summarized in Figure Ic, from which we highlight two key observations: 1) Row-wis wiww comparison indicates that the sampling quality is not sensitive to the im- portance sampling size. However, we emphasize that the conclusion is based on sizes smaller than 16, as 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 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 signif- icantly reduces the Gen PPL to 30 while longer window only marginally further improves the results. This phenomenon suggests that major decoding error happens during early-stage denoising sampling, and 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 gener- ation. By leveraging a residual energy-based approach, EDLM effectively reduces the mismatch between training and sampling distributions discrete diffusion models, resulting improved generation quality and efficiency. Through experiments both small and large language mod- cling 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 model- ing, setting the stage for further exploration of parallel generation techniques.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBLITY STATEMENT
We provide detailed pseudo codes of relevant algorithms in Algorithm 1 and Algorithm 2. We provide the formal proof for theorem 1 in Appendix A. We describe the details for used datasets, data processing, model architecture, training and evaluation implementations in Section 5.1 and ap- pendix C.I. We emphasize transparency and reproducibility for machine learning research.
ACKNOWLEDGMENT
MX thanks the generous support of Sequoia Capital Stanford Graduate Fellowship. SE gratefully acknowledge the support ARO (W911FF-21-1-0125), ONR (N00014-23-1-2159), and Chan Zuckerberg Biohub. JL gratefully acknowledge the support of NSF under Nos. OAC-1835598 (CINES), CCF-1918940 8940 (Expeditions)), DMS-2327709 (IHBEM), IIS-2403318 (III); Stanford Data Applications Initiative, Wu Tsai Neurosciences Institute, Stanford Institute for Human-Centered AI, Chan Zuckerberg Initiative, Amazon, Genentech, GSK, Hitachi, SAP, and UCB.
REFERENCES
Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Struc- tured denoising diffusion models in discrete state-spaces. In Advances in Neural Information Processing Systems, 2021.
Samy Bengio, Oriol Vinyals, Navdeep Jaitly, and Noam Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. Advances in neural information processing systems, 28, 2015.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Hemighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. Advances in Neural Information Processing Systems, volume 33, 2020.
Andrew Campbell, Joe Benton, Valentin De Bortoli, Tom Rainforth, George Deligiannidis, and A. Doucet. continuous time framework for discrete denoising models. Advances Neural Information Processing Systems, 2022.
Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Gen- erative flows on discrete state-spaces: Enabling multimodal flows with applications to pro- tein co-design. In Forty-first International Conference Machine Learning, 2024. URL ht pttts:://peerreee
.et/
Miguel A Parreira-Perpinan and Geoffrey Hinton. On contrastive divergence learning, Interna- tional workshop artificial intelligence and statistics, pp. 33 40 PMLR, 2005.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240): 1-113, 2023.
Arman Cohan, Franck Dernoncourt, Doo Soon Kim, Trung Bui, Seokhwan Kim, Walter Chang, and Nazli Goharian. A discourse-aware attention model for abstractive summarization of long documents. arXiv preprint arxiv: 2018
Yuntian Deng, Anton Bakhtin, Myle Ott, Arthur Szlam, and Mare' Aurclio Ranzato. Residual energy-based models for text generation. In International Conference Learning Represen- tations, 2020. URL htw.h:tpp:///c.....m..tt /ttp://////p//ei///e/ee n ftrr//orummp/
1d=B1 (atpummmd==i1ss@sa0 145gHHHH
Prafulla Dhariwal and Alex Nichol. Diffusson models beat gans on image synthesis. In Neural Information Processing Systems, 2021.
11

---

Published as a conference paper at ICLR 2025
Sander Dieleman, Laurent Sartran, Arman Roshannai, Nikolay Savinov, Yaroslav Ganin, Pierre H. Richemond, A. Doucet, Robin Strudel, Chris Dyer, Conor Durkan, Curtis Hawthorne, Rémi Leblond, Will Grathwohl, and Jonas Adler. Continuous diffusion for categorical data. ArXiv, abs/2211.1500 2022.
Itai Gat, Tal Remez, Neta Shaul, Felix Kreuk, Ricky TQ Chen, Gabriel Synnaeve, Yossi Adi, and Yaron Lipman. Discrete flow matching. arXiv preprint arXiv:2407.15595, 2024.
Alan E Gelfand. Gibbs sampling. Journal of the American statistical Association, 93(452):1300-
1304, 2000.
Marjan Ghazvininejad, Omer Levy, Yinhan Liu, and Luke Zettlemoyer. Mask- Parallel decoding of conditional masked language models. In Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP), pp. 61 6112-6121, 2019.
Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. httpppsstttttp/////ttta wwt.onn..o..........o..... io/
OpenWeebTextCorpuss 2019.
Alex Graves, Rupesh Kumar Srivastava, Timothy Atkinson, and Faustino Gomez. Bayesian flow networks. arXiv preprint arXin:2308.07037, 2023.
Jiatao Gu and Xiang Kong. Fully non-autoregressive neural machine translation: Tricks the trade. In Findings of the Association for Computational Linguisties: ACL-IJCNLP 2021, pp. 120-133, 2021.
Jiatao Gu, James Bradbury, Caiming Xiong, Victor O.K. Li, and Richard Socher. Non-autoregreessive
neural machine translation. In International Conference on Learning Representations, 2018. URL
https: wwwpeeeeevviveeeeeeeeeeeeeeeeeee ew net/ / forum? i d= 8Bt 1 Cb.
Ishaan Gulrajani and Tatsunori B Hashimoto. Likelihood based diffusion language models. Ad-
vances in Neural Information Processing Systems, 36, 2024.
Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: A new estimation principle for unnormalized statistical models. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pp. 297-304. JMLR Workshop and Conference Proceedings, 2010.
John Hammersley. Monte carlo methods. Springer Science & Business Media, 2013.
Xiaochuang Han, Sachin Kumar, and Yulia Tsvetkov. Ssd-Im: Semi-autoregressive simplex- based diffusion language model for text generation and modular control. arXiv preprint WrXiv:2210.17432. 2022.
Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8) 2002.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, 2020.
Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling, Argmax flows and multinomial diffusion: Learning categorical distributions. Advances in Neural Information Processing Systems, 34:12454-12468
2021.
Emiel Hoogeboom, Alexey A. Gritsenko, Jasmijn Bastings, Ben Poole, Rianne van den Berg, and Tim Salimans. Autoregressive diffusion models. In International Conference on Learning Rep- resentations, 2022. URL https: //tp///eeeeee/ee/////// net
IV.
Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, Fujie Huang, al. A tutorial energy- based learning. Predicting structured data, (0)) 2006.
Jose Lezama, Tim Salimans, Lu Jiang, Huiwen Chang, Jonathan Ho, and Irfan Essa Discrete predictor-comeecorr diffusion models for image synthesis. In The Eleventh International Confer- ence on Learning Representations, 2023. URL https: /ttp///weeeeeeeee////e// net fft.umririiieee VM@batVBivv.
12

---

Published as a conference paper at ICLR 2025
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion modeling by estimating the ratios of the data distribution. In Proceedings of the 41st International Conference on Machine Learning. pp.
PMLR, 2024.
Zhuang Ma and Michael Collins. Noise contrastive estimation and negative sampling for condi- tional models: Consistency and statistical efficiency. In Proceedings of the 2018 Conference Empirical Methods Natural Language Processing, pp. 3698-3707, 2018
Matt Mahoney. Large text compression benchmark. 2006. URL https: / http://wwwwttttt/////atta// www.:matttaattt mmthomey.... net/dc/text html.
Mitch Marcus, Beatrice Santorini, and Mary Ann Marcinkiewicz. Building a large annotated corpus of english: The penn treebank. Computational linguistics, 19(2):11-3330 1993.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. arXiv preprint arXivv1609.07843 2016.
Scbastian Nowozin. Debiasing evidence approximations: On importance-weighted autoencoders and jackknife variational inference. In International Conference on Learning Representations, 2018. URL https::/ net////wweeeeeeeeeeeeeeeeeeeeeee Hot.:/rrrorm---HH211o
Art B. Owen. Monte Carlo theory, methods and examples. https: //ttt:://artowen.. su. www.ama...dd..oooom.mmmoooo.. www.ddmiiinn. ins/ mc/, 2013.
Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring a broad discourse context. arXiv preprint arviv 1606.06666603) 2016.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Baptiste Roziere, Jonas Gehring, Fabian Gloeckle, Sten Sootla, Itai Gat, Xiaoqing Ellen Tan, Yossi Adi, Jingyu Liu, Romain Sauvestre, Tal Remez, et al. Code Ilama: Open foundation models for code. arXiv preprint arXiv:2308. 2tt.::33...190... 2023.
Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin T Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models, 2024.
Nikolay Savinov, Junyoung Chung, Mikolaj Binkowski, Erich Elsen, and Aaron van den Oord. Step- unrolled denoising autoencoders for text generation. In International Conference Learning Representations, 2022. URL https: : //wp//eeeee...../////// net/ Eorum? Ioaia==TTG1p1111 g6.
Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis K Titsias. Simplified and gener- alized masked diffusion for discrete data. arXiv preprint arXiv:2406.04 044.22906.04399, 2024.
Andy Shih, Dorsa Sadigh, and Stefano Ermon. Training and inference on any-order autoregressive models the right way. Advances in Neural Information Processing Systems, 35:2762-2775, 2022.
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In Proceedings of the 32nd International Con- ference on Machine Learning, 2015.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Neural Information Processing Systems, 2019.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint 23X2::2302.13977 2023.
Dustin Tran, Keyon Vafa, Kumar Agrawal, Laurent Dinh, and Ben Poole. Discrete flows: Invertible generative models of discrete data. Advances in Neural Information Processing Systems, 32, 2019.
13

---

Published as a conference paper at ICLR 2025
Trieu H. Trinh, Yuhuai Wu, Quoe V Le, He He, and Thang Luong. Solving olympiad geometry without human demonstrations. Nature, 625(7995):476-455 482, 2024.
Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In NIPS, 2017.
Jason Wei, Xuezhi Wang. Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoe V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824 24837, 2022.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Remi Louf, Morgan Funtowicz, Joe Davison, Sam Shleifer, Patrick Von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander Rush. Transformers: State-of-the-an natural language processing. In Qun Liu and David Schlangen (eds.), Proceedings of the 2020 Confer- ence on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 38- 45, Online, October 2020. Association for Computational Linguistics. doi: 10.18653/v1/2020. emnlp-demoo URL https: ac anthology, org/ 2020 emnlp-demos 6.
Xiang Zhang, Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text clas- sification. Advances in neural information processing systems, 28, 2015.
Lingxiao Zhao, Xueying Ding, Lijun Yu, and Leman Akoglu. Improving and unifying discretee&cootinuuus-tiiss discrete denoising diffusion. arXiv preprint arXiv:2402.03701, 2024.
Lin Zheng, Jianbo Yuan, Lei Yu, and Lingpeng Kong. A reparameterized discrete diffusion model for text generation. arXiv preprint arXiv:2302.05737, 2023.
Zachary Ziegler and Alexander Rush. Latent normalizing flows for discrete sequences. Interna- tional Conference on Machine Learning, pp. 7673-7682, PMLR, 2019.
14

---

Published as a conference paper at ICLR 2025
A
PROOF
To make the paper self-contained, we incorporate the relevant theoretical proof for Theorem from (Nowozin, 2018; Deng al., 2020). Our key difference is correcting minor misargument from estimation on Zg to estimation on log Z
Theorem Given diffused data xx at timestep t, let log 2n denote the empirical estimation of log Z (x+) = log Exo-po exp( exp(-Exg (xo, xt)) with n samples x( ~ poli = 1,,.. , Xt): ,3(x));
(8) x()) Then Ve EN such that Vn - we have log Zn log Ei= exp(-E((x))
log Zz - < Eloog 2n] < log Z E[(2n - 1) log nn 2(n - 1) < log Zo + (15)
Proof. Following Eq. 35 Nowozin (2018), we have Ellog 2n]
Ellog& ] = log Zo
27 n 373
2/2 +
1
n²
73
4n1
n²
12 n
(74 -3
5/5
1 10
n
7372
n (1s - (20011
1
(16)
where 7 = Eloo 2n). 7 = Bilogg 2n 7
By omitting the higher-order expansion, we equivalently have Eloog 2n] = log Zo which indicates that lim,. Bllog _m] log Zo. Then, we have
22
+
242 n
VE > 0, 3N1 > such that for n > N1. Bllog Z,.] > log Zo - 6.
Since lim, n(log Zo - Ellog 2n.) lim such that for n > N2, log Ze > logg 2n).
27
72 + o(1) =
72 272
> 0, 80
NN₂
> 0
In summary, we have proved that log Ze- - e Ellog Zn] log Ze.
Again, by omitting higher-order expansions using Eq. (16), we can have another equivalent form. Ellog 2, log Zo
Ta 1 + 5. + +onn where e a constant, and therefore E 1) log 2,, 27 n
2(n 1) log -n] (2n [log 2n] 2(n E)E 2n-1] log Zo + 24 + onn-1). This
indicates lim,, [2[ 1) log Zn 2(n 1) log Zn 1] log Zg. Therefore, we have
VE > 0, 3N3 > such that Vn > N₂ EE (2n - 1) log nn, 2( 1) log 2n-1] < log Za +
Since lim n (B[[(nn1)) log Z,, - 2 log 2n, -log Zu) = lim, +00 277
22 +0 >
0, so 3N4 > 0 such that when n > N4. [2n 1) log Z, 2(n 1) log Zn-1] > log Ze.
Given all the results above, we have that Ve > 0, for Vn > N that N max (N1, N2, N2. NaB log Zo < Ellog 2n] < log Zo < E[ 1) log Z,, 2(n 1) log 2,, 1| log Zo +
B ALGORITHM
We provide the detailed pseudo code for noise contrastive training of the HBM in Algorithm 2.
Algorithm 2 Noise Contrastive Estimation Training
1: Input: dataset 'D discrete diffusion model 0, energy-based model Φ learning rate !
2: repeat
3: Sample Xo ~ D and t ~ U(0, 1]
4: X, ~ q(xx{1(x))) (xo)
5: Sample positive data x ~ g(-|x,, xp)
6: Sample negative data x_ (-
7:
L(d; 0) +- (X in Eq. (10)
8: Φ b Φ - Vobb;
9:
until convergence
15

---

Published as a conference paper at ICLR 2025
C ADDITIONAL EXPERIMENT DETAILS AND RESULTS
C.I ADDITIONAL EXPERIMENT DETAILS
We provide additional experiment setup details in this section.
Text8. We follow all the common practices in Austin et al. (2021); Campbell et al. (2024) to conduct Text8 experiments, which has a dictionary size of 28 with 26 lowercase letters, a white-space token, and mask token. We follow the standard dataset split and train MDLM using standard 12-layer transformer architecture. the transformer also has the same number of heads (12) and hidden dimen- sion (784) as Austin al. (2021). The model is trained text chunks of length 256 for 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 con- verge 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 linear warm-up of 2000 steps. We set the channel-wise dropout rate as 0.05 and conducted optimization with AdamW and 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.
Wext We follow the standard data split in (Sahoo et al., 2024) to leave a validation split with the last 100k does as the validation set. We tokenize Open WebText with the GPT2 tokenizer, with a vocabulary size of around 50K. All models are trained with sequences wrapped to a length of 1, 020 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 timestep embedding of 128 when applicable. Word embeddings are not tied between the input and output. Other training details are also kept the same, ie., 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 func- tion across different denoising steps. Specifically, for each timestep (xx (x0), and then sample positive data x4 ~ q(-|xx, xo) and negative data x_ ~ pa (x),, similar to Algorithm 2. Then we study the energy value for (x(,,, and Ex(x),x.)
draw diffused data from
We summarize all plots for the energy w.r.I 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 e' denote the energy for i-th negative sample, we first normalize the energies to è' Eie oi and then ESS is given by
E 1=1 0²)² Σ (00)
ESS
(17)
which intuitively helps quantify the degeneracy of the sample weights in importance sampling. A low ESS indicates that 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, 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 did not observe that EDLM-NCE is clearly better than EDLM-AR, and we conclude that this is mainly because ESS analysis is conducted with 16 samples instead of extremely large
note that
16

---

Published as a conference paper at ICLR 2025
2000
1000
Energy
PosRive Energy Negattre Energy
Enwrgy
Enengy
1000
-5
-10
2000
2200
-3000
Positive Energy
4000
Positive Energy
Negalive Energy
Negalive Enwrgy
5000
0.0
0.2
0.4
0.6
0.8
1.0
3-4
0.6
2.6
0.5
0.2
2.4
Time
Time
Time
15
2000
20
Ererpy
Enwigs
2500
Enarge
3,500
5
1000
2000
-5
Negative Enargy Min Negalive Energy
1500
Max Negatitve Energy
Min Negathve Enengy
500
Max Negative Energy
Min Negative Energy
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Time
Time
Time
166
ESS
14
ESS
12
10
10
2
0.2
1.0 0.0
0.6
0.2
0.4
1.0
Time
Time
Time
(a) EDLM-NCE
(b) EDLM-AR
(c) EDLM-coAR
Figure 2: Behavior the energy function under different parameterization. The first row plots the energy of positive samples and the average energy of negative samples; the second plots the maximum and minimum energy values of the 16 negative samples; the third row plots the effective sampling size (ESS) for energies of the 16 negative samples. Different columns correspond to results for different EDLM parameterization.
50
EDLM-NCE w= 0.20
8.0
EDLM-NCE w= 0.20
Generative Perplexity
EDLM-NCE 0.40
EDLM-NCE w= 0.40
45
EDLM-NCE 0.80
EDLM-NCE w= 0.60
7.8
EDLM-NC NCE w= 0.80
EDLM-NCE 0.60
Entropy
40
EDLM-AR w= 0.20
EDLM-AR 0.20
EDLM-AR 0.40
7.6
EDLM-AR 0.40
35
EDLM-AR w 0.60
EDLM-AR w- 0.60
EDLM-AR w= 0.80
EDLM-AR 0.80
7.4
30
SEDD
SEDD
MDLM
MDLM
25
10
15
20
25
30
35
40
10
15
20
25
30
35
40
Time (s)
Time (s)
(a) Generative Perplexity vs. Time.
(b) Entropy vs. Time.
Figure 3: Additional EDLM sampling results with varying importance sampling window w. Similar to section 5.3, we run each setting with [512, 768, 1024] denoising steps, and plot the curve of the corresponding metric wall-clock time. For generative perplexity, the metric is evaluated by GPT- 2, and a curve on the bottom-left indicates a better sampling quality vs time trade-off.
number. We recall here that GPU memory bound the sampling batch size to 16, and therefore large ESS analysis in practice is not useful.
17

---

Published as a conference paper at ICLR 2025
Additional results with different sampling window length w. We provide additional sampling results with different importance sampling window sizes w. All results are provided in Figure 3. see that generally
As shown in the figure, especially the reference line of AR in Figure 3a, w I 0.6 can all lead to better efficiency, among which w = 0.2 achieve the best speedup.
Additional visualization of Gen PPL and entropy results.
We provide a more detailed trade-off between entropy and Gen PPL compared the baselines and the proposed model.
Generative Perplexity
50
Specifically, we rearrange Figures la and lb to a new figure with Entropy and Gen PPL axis, and provide the figure in Fig- ure 3b. As shown in the figure, EDLM shows a trade-off sim- ilar to MDLM (< 0.1 difference of entropy in the overlapping part). However, EDLM aims to reduce the sampling errors and therefore has a similar performance as MDLM but with much fewer sampling timesteps
45
40
EDLM-NCE 0,20
EDLM-AR w- 0.20
35
SEDD
MOLM
30
7.6
0.0
Entropy
Figure 4: Gen. PPL vs. Entropy.
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 human. And said, YYou didn't wanna get into that." Yeah,, that's s just physically [grown-upl, I kind of weighed down with that. An \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.An Antn you had to learn something else, more of appreciation of it, watching guys play [inside out] as you grew. Everything they showed up you knew as your parents wanted you. AnAnAnd [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. \nlnie had to find some way out. It got borderline ridiculous, but we found a way.An AnAnHe and I speak and we talk a lot. That's s really the thing about us. Just the two dogs the right time. An funny man. Just like we are, baby. We're re not stupid."Inin when you're re getting pumps?"\n\n\We" of the eye. I wish I could. I can''t see it on my face. Yeah, I don''t really. picked it up and it's not the right way. But the way I got
always in the way they wanted, as just
happened like I can''t see it from that side
Cousins\n\nI think I'm not confused. He was so close to me. But I'm never confused either. AnAn\Are you playing the same type of game now?"\n\n\Alwen 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 good season. He was defensive superstar. 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 s part of this life. Like my dad and had little bit, and that got you sometimes it didn't work to the liking. [You] would physically get him. \n\nWe did, but he always demanded that I catch up to him. And when I tell him, \Catch me you say something stupid, you're re not stupid." " If he does say stupid something, you know, I'm going to respond to it. Annnnii like \What'ss going to do catch \'n' it? If he's s gonna catch it, thatV's my one. "A "innnnnythhin like that?\n\n\ItV
the same way. You
we would go through the type
funny because 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 quy and it affects the game. "An/n\That's a good one. \n\N\His skills that he play, can he tell you?"An He's an unvalievable
18

---

Published as a conference paper at ICLR 2025
dude. \n\nYeah, he's s not the quy that everybody goes to play with, just out on the team. He plays stuff because he takes the game seriously. But you see player that goes and develop that's s easy to feed off of, and that's s good. They win games, too, by being a that's all. A.I.AAnnHoonstly.
teammate, in that rnvironmmet..An hasn't t been really easy for me. I played 35 percent of the time and these guys ask a lot about me [laughs]. guys. I know that when I was young I would not have made it. I was just born that way when was there. So don't call me the only man that I really liked or the backbone of the team, ever. That guy is the person in the your team, you don't have to make credit for what you do. An definitely miss Tom or Cole. Everybody does that to me. I love Chad Mac. It's the Strict style. He''s always teaching
it
told you guys those
An/nhen you' re all of the guys in
being an open guy. An\nThat kind 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 manager will have a tremendous benefit as player, and that's very important, that is, why hockey. That's what you want 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 Totally. I just
player. AnAnRNS: And you can do that?\n\ndRP think you can. So, you' ve gotta know, if you' re going to play that role of manager, you'll be the one that knows where to go. And that's s how great it is, every year. AnAnRNS: If it's difficult as manager in your last one, did you find yourself right about that decision? Did you feel better?ining \elll, I sat there the other day thinking, 'OK, have to do this job to get career in that position.' So when I do that: one-thousand up, no down. could not go on. didn't do what I need to say about that and I can't 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." AI can't t do that' 'I let the people out of my decisions,'' then I think you are one of that. \nAnBut you 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, decide which is different. So we' re the same positions but we talk about how to do An'Ann just, more and more ve've said: that's how you' re gonna get there. If you want that manager role, you should do this. Because you've 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
when you
hockey player, you
19

---

Published as a conference paper at ICLR 2025
and positions?\nlnJPP Yeah, you never know. And you know how they' re actually putting these things. So when you start a coach, you're 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 always tell people after you' re done doing the things, you trust me that it's the most important to take place next. AnhnSo, when I started coaching, do not think that was great leap because that was my first time in the NHL. But I would like to tell the coaches, Youu' 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' 11 be the one and that season you'll be the one. And that's what you gotta think about is that if you're re more involved and in unique ways, if you're 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 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...
best player you can be from your experiences.
D.3 GENERATED SAMPLES FROM EDLM-COAR WITH 1024 TIMESTEPS.
.o.to explain that to someone who has not shared information with Google itself and not the powers that be. And Kim didn't. t. never wanted to do this because before it, dian't know anything. Finally, you feel as if you "nnni laughed. An\n didn't feel right. Butt yeah, we stayed together for a while, Kim said, \and when we had time for us to, almost in silence. \I'm m not able to act out that way, she said, Aso 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 "inny."nnnn" to other girls friends and friends and I had my first year-long meeting with social media guys. As a women's boss. "AnAnAI feel like different company here," she said. \I mean, it's basically an company. "\nAn\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. f'iends.."n'nhhen "\n\nWhen she was speaking to North, she felt in touch with friends and family in Marlboro where 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, 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. Sone conversations on the phone and for the other AnAnKim Airls.AnAnKim leaned in very talkily and he said to her, \We don't like you much," but he wasn't trying to make California's jealous. Youu' 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 t go out to anything that's not good for you." And that's s all too true. AnAnI first 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 quy that is always doing something great," West had said. AAnd then I'll come up short, do well myself, and I go, 'I really want someone on this,'' she said, An'd) sad, " Jörg smiled, \this kind of guy who's not that smart and not do well. I
were talking. "AnAnI was
day, we were talking
20

---

Published as a conference paper at ICLR 2025
mean, you looked the fuck you down and you figured out that's s more important to me than you thought you were, basically. You wanted to do "Mnethiing."nllVee years trying to figure out about you so that I could link you and spend time with somebody who's gonna roll.")ninly could just read his shit and die, " she said, gesturing to West's letters. AAnd that's s how was left behind. It's always struck me that didn't know how I wanted to spend with this group or another. "Anin you dated her the past?" she asked, asking if he had more of problem with her? \No, not really, he said. \I'm nice guy and I would talk back to anyone and that is why I actually talk to people. "Thtn)That would be really creepy place to be in," " he told her, adding, Vit was like... friend. "\n\nThere was blood in the room. \Lemme, he was funny dude at the time. \nnHe smiled. \Btt you know that. This way a happy friend would have been if I had him back then. I didn't t care for him. ..."
sad, " he said. \That was all had my
first time I've lost a
21


Reviews:
[{"review_id": "V3y1N8sPYQ", "reviewer": "Reviewer_LAtw", "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.", "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.\n- 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.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The paper is clearly written and easy to follow.\n- The method is intuitive and based on a solid mathematical foundation.\n- The experiments follow the standard setup for evaluating Diffusion Language Models, making it easy to compare with other methods.", "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.\n- 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.\n\n**Recommended Experiments**:\n\nIt 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.\n\nAs 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "9BEgDhVhVX", "reviewer": "Reviewer_r3kY", "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.", "questions": "1. For EDLM-AR and EDLM-coAR, which pretrained AR model is used and what is the parameter size of the model?\n2. 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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "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.\n- The adaptation of auto-regressive models to a joint denoising distribution with masked inputs is innovative.\n- 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.", "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.\n- 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.\n\n[1] Deng, Yuntian, et al. \"Residual energy-based models for text generation.\" *arXiv preprint arXiv:2004.11714* (2020).\n\n[2] Ghazvininejad, Marjan, et al. \"Mask-predict: Parallel decoding of conditional masked language models.\" *arXiv preprint arXiv:1904.09324* (2019).\n\n[3] Savinov, Nikolay, et al. \"Step-unrolled denoising autoencoders for text generation.\" *arXiv preprint arXiv:2112.06749* (2021).\n\n[4] Zheng, Lin, et al. \"A reparameterized discrete diffusion model for text generation.\" *arXiv preprint arXiv:2302.05737* (2023).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bhgc0EzaWI", "reviewer": "Reviewer_9hNB", "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.", "questions": "Section 4.3 is not very clear to me. Could you provide a detailed explanation of how you estimated PPL?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Using an AR to help the sampling of Masked Discrete Diffusion Language Model is natural and straightforward.\n2. The experiments demonstrate good improvements compared with MDLM baseline.", "weaknesses": "**My main concern is that the core technique introduced in this work has been present in the literature for a long time.**\n\n1. 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.\n\n2. 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.\n\n[1] Non-autoregressive Neural Machine Translation, 2017  \n[2] Mask-Predict: Parallel Decoding of Conditional Masked Language Models, 2019  \n[3] Fully Non-autoregressive Neural Machine Translation: Tricks of the Trade, 2020", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jrJjpMO0jW", "reviewer": "Reviewer_wiWZ", "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.\nThe 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.", "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?\n- 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?\n- In the context of EDLM-NCE, do you employ separate models or do you utilize a shared rerank and generator approach?\n- 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?", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The writing is clear, making the paper easily comprehensible.\n- The idea to inject sequence-level correlation in the intermediate denoising steps is intuitive.", "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.\n- 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.\n- This paper only evaluates perplexity, and it remains uncertain whether such energy-based reranking translates effectively to downstream tasks.\n- I tend to have concerns about the inference cost linked to repetitive sampling and energy function computations in each denoising steps.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)