Encoder-Decoder Diffusion Language Models for Efficient Training and Inference

AI summary

Overview

Research area: Efficient generative modeling with discrete diffusion language models. Technical level: Intermediate (readers should be comfortable with transformers, attention, and the basic idea of diffusion-based text generation). Scope: The paper proposes and evaluates E2D2, an encoder-decoder diffusion framework that reduces the cost of both training and inference for diffusion language models by separating clean-token representation from noisy-token denoising.

What This Paper Is About

Diffusion language models generate text by starting from a fully masked sequence and iteratively "denoising" it in parallel, which is intended to be faster than the token-by-token decoding used by autoregressive (AR) models. The catch is that existing diffusion models are decoder-only, so every denoising step calls the entire network, making inference expensive. This paper shows that the two internal jobs of a diffusion model — representing clean tokens and denoising corrupted tokens — can be assigned to separate modules, letting a large encoder run occasionally and a small decoder run repeatedly.

Key Contributions

1. An encoder-decoder architecture for discrete diffusion. A transformer encoder produces representations of clean tokens (prompt plus already-generated text), while a lightweight transformer decoder iteratively denoises the noised block conditioned on those representations through cross-attention. Encoder and decoder parameters can be reallocated independently, unlike a single decoder-only stack.

2. Efficient sampling and training algorithms. The sampling procedure invokes the small decoder for many denoising steps and the encoder only periodically, with KV caching throughout. A vectorized block-diffusion training algorithm uses custom attention masks so all blocks are processed in one encoder pass and one decoder pass, halving training FLOPs relative to a decoder-only block diffusion model of equal size.

3. Two decoder-conditioning variants. A "last hidden state" design (T5-style, where each decoder layer sees the final encoder output) for training from scratch, and a "shared KV cache" design (where decoder layers reuse KV pairs from corresponding encoder layers) that makes it practical to convert a pretrained decoder-only LLM into an encoder-decoder by copying weights and truncating the decoder.

4. Empirical validation across four settings. Domain-specific models for summarization, translation, mathematical reasoning, and open-domain pretraining, with a Pareto analysis of quality versus throughput and public release of code, weights, and a project page.

Main Findings

• Summarization (CNN/DailyMail): E2D2 with 20 encoder and 8 decoder layers reaches ROUGE-1 of 36.0 at 155.8 tokens/sec, beating the decoder-only BD3LM baseline (35.8 at 135.1 tokens/sec), MDLM (30.6 at 49.3 tokens/sec), and even a 28-layer AR model (31.7 at 89.1 tokens/sec), roughly a 75% throughput gain over AR.
• Machine translation (WMT14 de-en): E2D2 at 28 encoder/4 decoder layers scores BLEU 24.8 at 162.0 tokens/sec, versus 24.0 at 102.4 tokens/sec for a 16-layer BD3LM and 18.4 at 60.4 tokens/sec for MDLM.
• Mathematical reasoning (GSM8K): Fine-tuned E2D2 achieves 47.9% pass@1 at 102.8 tokens/sec, versus 33.2% at 86.6 tokens/sec for BD3LM and 14.0% for MDLM. AR remains stronger on raw accuracy (66.6%) but is slower to decode.
• Language modeling (OpenWebText): E2D2 attains validation perplexity 21.73, better than MDLM (22.98) and SEDD (24.10) and close to BD3LM (20.73), while training about 40% faster than BD3LM (8.4e4 vs 5.9e4 tokens/sec).
• Pareto frontier: By sweeping decoder depth, the authors show that at matched decoding throughput, E2D2 consistently produces higher-quality models than BD3LM, shifting the quality-speed trade-off curve outward.
• Block size ablation: Smaller blocks improve quality (pass@1 50.1 at block size 2) but reduce throughput (34.8 tokens/sec); larger blocks invert the trade-off (20.9 pass@1 at 62.2 tokens/sec for block size 32). Encoder invocations are the throughput bottleneck.
• Diffusion-step ablation: E2D2 outperforms BD3LM in quality at every diffusion-step count tested; the throughput advantage shrinks as the number of steps drops, since fewer decoder calls mean the encoder's relative cost grows.
• FLOP accounting: For equal total layers, E2D2's training forward pass uses about half the FLOPs of decoder-only block diffusion, and its inference FLOPs scale with a cheap decoder call repeated across steps plus an occasional expensive encoder call.

Methodology in Plain English

The authors start from the observation that a diffusion language model does two conceptually different things at each step: it builds an internal understanding of the text that is already known (the prompt and previously decoded tokens), and it predicts replacements for the currently masked tokens. In a decoder-only model these are entangled, so both happen at full network cost every step.

E2D2 disentangles them. A transformer encoder reads the clean tokens and produces rich representations, cached as keys and values. A small transformer decoder takes the currently noised block and cross-attends to those cached representations to produce token predictions. Because the encoder output does not change while a block is being denoised, the decoder can be run many times without touching the encoder. Only when a block is finished does the encoder absorb the new clean tokens and refresh its cache.

Two wiring choices are explored. The last hidden state variant concatenates the encoder's final-layer output onto the decoder's input embeddings, similar to T5. The shared KV cache variant instead feeds encoder key-value pairs from specific layers into specific decoder layers, chosen so that they line up with what a pretrained decoder-only model expects — this enables fine-tuning a pretrained LLM (Qwen3) into an E2D2 model rather than training from scratch.

For training, block diffusion normally requires processing both the full clean sequence and the full noised sequence, roughly doubling cost. E2D2 sends the clean sequence through the encoder once with a block-causal mask and the noised sequence through the decoder once with a custom mask, so all blocks are trained in parallel with half the FLOPs.

Evaluation covers four tasks: CNN/DailyMail summarization with ROUGE, WMT14 German-English translation with BLEU, GSM8K math reasoning with pass@1 accuracy, and OpenWebText pretraining with perplexity on held-out corpora. Throughput is measured in tokens per second on a single H100 80GB (and A100 80GB for ablations), and decoder depth is varied to trace the quality-throughput trade-off.

Why This Matters

Impact on research. The paper reframes an architectural assumption that has gone largely unchallenged in diffusion language modeling: that a single decoder-only stack must handle both representation and denoising. It provides a concrete, FLOP-accounted argument that encoder-decoder separation is strictly more efficient under equal layer budgets, plus attention-mask designs that make the training practical. This creates a new axis for the diffusion-versus-autoregressive debate — not just quality, but cost per generated token — and offers a recipe for recycling pretrained AR weights into diffusion models.

Real-world applications:

• Latency-sensitive summarization and document processing, where generating many tokens in parallel at high throughput directly reduces serving cost.
• Machine translation services that need high-quality output with predictable, low per-token latency.
• Reasoning assistants and math tutoring systems that must produce multi-step chains of thought quickly.
• On-device or edge deployment, where the ability to use a small decoder for most steps and a larger encoder only occasionally fits constrained memory and compute budgets.

Industry relevance. Inference cost dominates the economics of serving language models. A framework that demonstrably improves the quality-throughput frontier, and that can be initialized from existing pretrained checkpoints via the shared-KV-cache trick, is directly applicable to production systems looking to cut GPU spend without retraining from scratch.

Future Directions

• Scaling to frontier sizes. Results are at 170M to roughly 1.7B parameters; whether the encoder-decoder advantage holds and compounds at the 7-8B scale used by LLaDA, Seed Diffusion, and MMaDA is untested.
• Reconciling small block sizes with throughput. The best quality comes from block size 2, but that maximizes encoder invocations; better scheduling or partial encoder refreshes could break this coupling.
• Unifying with approximate KV caching. Models trained with full-sequence diffusion rely on approximate caching at inference; combining E2D2's exact encoder-decoder caching with those methods could yield further gains.
• Extending beyond text. The same encoder-decoder decomposition could apply to discrete diffusion in protein design, music, and other sequence domains where clean-context representation and denoising have distinct costs.
• Better understanding of the design space. The paper compares two encoder-to-decoder conditioning schemes; systematic study of layer pairings, parameter sharing, and decoder depth allocation remains open.

Target Audience

Machine learning researchers working on diffusion models and efficient generative architectures; NLP engineers concerned with inference throughput and serving costs; practitioners interested in converting pretrained autoregressive LLMs into diffusion models; and graduate students studying the trade-offs between autoregressive and parallel decoding paradigms. Readers with a working knowledge of transformer attention and the basics of discrete diffusion will get the most out of it, though the high-level argument is accessible to anyone familiar with modern language model architectures.

Authors’ 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