Encoder-Decoder Diffusion Language Models for Efficient Training and Inference

TL;DR

The paper tackles the high computation cost of decoder-only discrete diffusion language models by introducing Efficient Encoder-Decoder Diffusion (E2D2), which decouples the work of representing clean tokens (encoder) from denoising corrupted tokens (decoder). By enabling multiple lightweight decoding passes per generation step and periodic encoder updates, E2D2 achieves faster inference and training, particularly benefiting block diffusion with KV caching. The authors provide efficient sampling and training algorithms, and demonstrate that E2D2 improves throughput and quality on summarization, translation, and mathematical reasoning tasks, mapping a favorable Pareto frontier against decoder-only baselines. This approach holds practical impact for scalable diffusion-based NLP with improved efficiency and versatility across tasks.

Abstract

Discrete diffusion models enable parallel token sampling for faster inference than autoregressive approaches. However, prior diffusion models use a decoder-only architecture, which requires sampling algorithms that invoke the full network at every denoising step and incur high computational cost. Our key insight is that discrete diffusion models perform two types of computation: 1) representing clean tokens and 2) denoising corrupted tokens, which enables us to use separate modules for each task. We propose an encoder-decoder architecture to accelerate discrete diffusion inference, which relies on an encoder to represent clean tokens and a lightweight decoder to iteratively refine a noised sequence. We also show that this architecture enables faster training of block diffusion models, which partition sequences into blocks for better quality and are commonly used in diffusion language model inference. We introduce a framework for Efficient Encoder-Decoder Diffusion (E2D2), consisting of an architecture with specialized training and sampling algorithms, and we show that E2D2 achieves superior trade-offs between generation quality and inference throughput on summarization, translation, and mathematical reasoning tasks. We provide the code, model weights, and blog post on the project page: https://m-arriola.com/e2d2

Figures (8)

• Figure 1: Efficient Encoder-Decoder Diffusion (E2D2) enables faster generation than decoder-only architectures. We accelerate inference by using a lightweight decoder to iteratively denoise for a fixed number of sampling steps, without invoking the encoder. The encoder processes the newly generated tokens periodically to update its representations.
• Figure 1: Forward-pass training FLOPs comparison. $N$ denotes the number of layers, $L$ the sequence length, $S$ the block size, and $D$ the hidden dimension. The expression highlighted in blue denotes the number of attention operations. E2D2 refers to block diffusion parameterization. For E2D2, we separate $N$ into $N_{\text{Enc}}$ encoder and $N_{\text{Dec}}$ decoder layers. We provide the derivation in Appendix \ref{['appsec:training-flops']}.
• Figure 2: Example attention masks used in block diffusion training with an encoder-decoder architecture, for $L=6$ tokens with blocks of size $S=2$. Left: The encoder mask $\mathcal{M}_\text{Enc} \in \{0, 1\}^{L \times L}$ enables clean tokens to attend within their respective block and to previous blocks. Right: The decoder mask $\mathcal{M}_\text{Dec} \in \{ 0,1 \}^{L \times 2L}$ uses self-attention within noised blocks and cross-attention to previous clean blocks using the encoder's output representation ${\mathbf h}^{1:L}$.
• Figure 3: CNN/DailyMail test set ROUGE scores. Best values for our trained models are bolded.$N$ refers to number of transformer layers ($N_\text{Enc}$/$N_\text{Dec}$ for E2D2). Decoding throughput (Tput) is measured in tokens / sec on 1 H100 80GB machine. For all models, we use $T=L$ sampling steps, so the throughput can be higher for diffusion when $T<L$. We report mean $\pm$ standard deviation for 100 samples.
• Figure 4: Example attention masks used in training full-sequence diffusion with an encoder-decoder architecture, for $L=6$ tokens. We denote clean tokens in ${\mathbf z}_t^{1:L}$ as ${\mathbf x}^{C_t}$, where $C_t$ corresponds to their token indices. In this example, $|C_t|=3$. Left: The encoder mask $\mathcal{M}_\text{Enc} \in \{0, 1\}^{|C_t| \times |C_t|}$ enables clean tokens to attend to other clean tokens in the sequence. Right: The decoder mask $\mathcal{M}_\text{Dec} \in \{ 0,1 \}^{L \times (L + |C_t|)}$ uses bidirectional attention across the noised sequence ${\mathbf z}_t^{1:L}$ and cross-attention to the encoder's output representation ${\mathbf h}^{C_t}$.