Energy-Based Diffusion Language Models for Text Generation

TL;DR

This paper introduces Energy-based Diffusion Language Models (EDLM), which improve diffusion-based text generation by using an energy-based approach at the sequence level, achieving better performance and faster sampling compared to existing diffusion models.

Contribution

The paper proposes EDLM, an energy-based diffusion model for text generation that enhances approximation quality and sampling efficiency, bridging the gap with autoregressive models.

Findings

EDLM outperforms state-of-the-art diffusion models on benchmarks.

EDLM approaches autoregressive models' perplexity.

EDLM offers a 1.3× speedup in sampling without performance loss.

Abstract

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