Energy-Based Diffusion Language Model

• Energy-Based Diffusion Language Models (EDLMs) are generative models that combine energy-based techniques with diffusion processes to produce stable and interpretable text outputs.
• They employ both continuous latent space formulation and discrete sequence corrections, addressing challenges like sampling instability and independence approximations.
• Empirical results demonstrate that EDLMs achieve competitive perplexity and improved generation speed, making them effective for interpretable and parallel text generation.

Energy-Based Diffusion LLMs (EDLMs) are a family of generative LLMs that integrate the structural flexibility and expressive capacity of energy-based models (EBMs) with the effective sampling and denoising capabilities of diffusion models. EDLMs have been instantiated in both continuous latent space form for interpretable text modeling (Yu et al., 2022) and discrete sequence form for parallel, non-autoregressive text generation (Xu et al., 28 Oct 2024). They address core challenges in both model classes: the instability and degeneration of latent EBMs under traditional MCMC, and the intrinsic approximation gap in discrete diffusion models that arises from neglecting inter-token sequence dependencies during denoising.

1. Foundations: Diffusion Models and Energy-Based Modeling

Diffusion models operate by gradually corrupting data through a Markovian noising process and learning to iteratively invert this corruption via a denoising process. In continuous latent EDLMs, the diffusion operates on a continuous latent variable $z \in \mathbb{R}^d$, while in discrete sequence EDLMs, the process corrupts token sequences by stochastically masking or replacing tokens.

Standard energy-based models define a probability distribution via an unnormalized score, or "energy," with difficult-to-compute partition functions. In text generation, latent- or sequence-level EBMs have been limited by intractable sampling or training instabilities. EDLMs resolve these through diffusion-based denoising, which confines EBMs to near-unimodal, tractable distributions at each step.

2. Latent-Space EDLMs for Interpretable Text Modeling

Latent EDLMs (Yu et al., 2022) assume a continuous latent $z$ and (optionally) a discrete symbol $y$ with an EBM prior: $$p_\phi(y, z) = \frac{1}{Z_\phi} \exp(\langle y, f_\phi(z) \rangle) p_0(z)$$ where $f_\phi: \mathbb{R}^d \to \mathbb{R}^K$ is an MLP, $p_0(z)$ is typically $\mathcal{N}(0, I)$, and $Z_\phi$ is the partition function. The marginal prior over $z$ is also energy-based, and the decoder $p_\theta(x \mid z)$ is an autoregressive LM.

Variational inference is performed with amortized encoder $q_\psi(z \mid x)$, optimizing the ELBO: $$\text{ELBO} = \mathbb{E}_{q_\psi(z \mid x)}[\log p_\theta(x \mid z)] - \text{KL}(q_\psi(z \mid x) \| p_\phi(z))$$ To remedy MCMC degradation, EDLMs inject a diffusion recovery chain on the latent space. The forward process applies Gaussian noise, and the reverse denoising is learned via conditional EBMs at each diffusion step: $$p_\phi(\tilde z_t \mid z_{t+1}) \propto \exp(F_\phi(\tilde z_t, t) - \frac{\|\tilde z_t - z_{t+1}\|^2}{2 \sigma_{t+1}^2})$$ Because each EBM operates on a nearly unimodal distribution, few Langevin steps suffice for efficient sampling.

3. Sequence-Level EDLMs for Parallel Generation

In the discrete sequence formulation (Xu et al., 28 Oct 2024), EDLMs address the “factorization mismatch” in traditional discrete diffusion LMs that predict tokens independently: $$p_\theta(x_0 \mid x_t) = \prod_i p_\theta(x_0^{(i)} \mid x_t)$$. This neglects inter-token dependencies and leads to cumulative decoding error as the number of denoising steps is reduced.

EDLMs correct this by defining a residual EBM at each denoising step: $$p_{\theta, \phi}(x_0 \mid x_t) \propto r_\theta(x_0 \mid x_t) \exp(-E_\phi(x_0, x_t, t))$$ where $r_\theta$ is the diffusion model’s learned denoiser, and $E_\phi$ injects global (sequence-level) energy-based corrections.

Parameterizations include:

• EDLM-AR: Energy derived from pretrained autoregressive LMs, via $$E_\phi(x_0, x_t) = -\log p_{\text{AR}}(x_0) + \log r_\theta(x_0 \mid x_t)$$.
• EDLM-NCE: An explicit Transformer-based EBM trained via noise-contrastive estimation (NCE) against $r_\theta$.

By incorporating these energies, sequence-level dependencies are modeled, leading to significantly improved sample quality and convergence.

4. Training Objectives and Algorithmic Advances

Latent EDLMs optimize an augmented ELBO involving both the forward diffusion-reconstruction likelihood and the reverse EBM chain: $$\text{ELBO}_{\text{Diff}} = \mathbb{E}_{q_\psi(z_0 \mid x)}[\log p_\theta(x \mid z_0) - \log q_\psi(z_0 \mid x)] + \mathbb{E}_{q_\psi(z_0, z_{1:T} \mid x)}\left[\log \frac{p_\phi(z_{0:T})}{q(z_{1:T} \mid z_0)}\right]$$ Additional regularizers include the information bottleneck—promoting disentanglement of $z_0$ from $x$—and geometric clustering—which sharpens latent modes to prevent collapse. The total loss is: $$\mathcal{L} = -\text{ELBO}_{\text{Diff}} + \beta_{\text{IB}} \mathcal{L}_{\text{IB}} + \beta_{\text{GC}} \mathcal{L}_{\text{GC}}$$

In sequence-level EDLMs, NCE is used for energy function estimation: $$\mathcal{L}_{\text{NCE}}(\phi) = -\mathbb{E}_{x_0, t} \left\{ \mathbb{E}_{x_+} [\log \sigma(-E_\phi(x_+, x_t, t))] + \mathbb{E}_{x_-} [\log \sigma(E_\phi(x_-, x_t, t))] \right\}$$ Perplexity evaluation exploits discrete diffusion variational bounds and importance-weighted estimates of the partition function.

A key algorithmic innovation is efficient parallel importance sampling during generation. Proposals are drawn from $r_\theta$, rescored via the EBM, and resampled—yielding wall-time speedups exceeding $1.3\times$ over standard diffusion sampling with negligible quality loss.

5. Inference and Sampling Procedures

Latent EDLM text generation samples $z_T \sim \mathcal{N}(0, I)$, reverses diffusion via Langevin-corrected denoising steps to $z_0$, then samples $x \sim p_\theta(x \mid z_0)$ with an autoregressive decoder. In discrete sequence EDLMs, the generation comprises the following for each diffusion step:

• Compute $r_\theta(\cdot \mid x_{\tau_n})$, sample $k$ proposals in parallel.
• Evaluate $E_\phi$ for each proposal.
• Resample using importance weights if in early (noisy) steps (controlled by a window parameter $\lambda$).
• Otherwise, revert to pure $r_\theta$-based sampling. This approach enables O(T) parallelization, and the computational overhead is typically dominated by the size of batched parallel proposals ($k \leq 16, \lambda \approx 0.2$).

6. Empirical Results

Latent EDLMs exhibit improved generation quality and interpretability across diverse benchmarks:

• On PTB, EDLM with geometric clustering achieves rPPL $164.6$ (best), BLEU $11.16$, and NLL $82.38$, outperforming both EBM-only and standard diffusion priors.
• For unsupervised clustering (DailyDialog), EDLM achieves higher mutual information (MI $3.94$ vs. EBM $2.42$), superior act- and emotion-homogeneity, and BLEU $28.75$ for reconstructions.
• On the Yelp sentiment-controlled dataset, 99.0% accuracy is realized with clear mode separation in latent scatterplots.
• Semi-supervised document classification (AGNews, $n=200$ labels): 87.4% (EDLM) vs. 86.4% (EBM-only).

Sequence-level EDLMs narrow the performance gap with autoregressive LMs:

• On Text8, EDLM achieves bits-per-character $\leq 1.24$, matching the AR baseline.
• On OpenWebText, EDLM-coAR reaches perplexity $17.58$ (AR: $17.56$), with increased robustness on out-of-domain splits.
• Sampling speedup is achieved: comparable generative perplexity is reached in $\sim 13$ s (EDLM-AR) vs. $\sim 17$ s (AR), a $1.3\times$ improvement.

The following table summarizes perplexity across key settings (Xu et al., 28 Oct 2024):

Model	OWT	PTB	Wiki	LM1B
AR	17.56	82.05	25.75	51.25
SEDD	24.56	100.09	34.28	68.20
MDLM	23.83	95.26	32.83	67.01
EDLM-NCE	21.52	93.21	30.77	63.19
EDLM-AR	20.49	89.67	29.24	60.80
EDLM-coAR	17.58	89.73	28.31	60.23

7. Implications, Limitations, and Future Directions

EDLMs provide a unified framework to combine the modeling strength and interpretability of EBMs with the efficient, stable sampling of diffusion approaches. They address major challenges in both hybrid and discrete autoregressive-free generative modeling, such as:

• Resolving sampling degeneration and instability in latent EBMs via diffusion recovery.
• Correcting the independence approximation in discrete diffusion LMs by sequence-level energetic correction.
• Achieving interpretable, controllable, and highly parallel text generation without significant sacrifice in perplexity or sample fidelity.

Limitations include the need for partition function estimation in non-AR EBMs, importance sampling overhead dependent on $k$ and $\lambda$, and increased memory requirements for large batches. Potential future extensions highlighted include alternative energy estimation (score-matching, adversarial), adaptive scheduling strategies, and generalization beyond language to other discrete modalities such as code or music (Yu et al., 2022, Xu et al., 28 Oct 2024).

A plausible implication is that EDLMs may enable non-autoregressive text generation to reach and even surpass the sample quality of left-to-right LMs, while expanding possibilities for order-agnostic, interpretable, and conditional text generation across settings previously dominated by autoregressive protocols.