Encoder-Decoder Diffusion Language Models for Efficient Training and Inference

TL;DR

This paper introduces an encoder-decoder architecture for discrete diffusion language models, significantly improving inference speed and training efficiency while maintaining high-quality output across various NLP tasks.

Contribution

It proposes a novel encoder-decoder framework for discrete diffusion models, enabling faster inference and training through specialized modules and algorithms.

Findings

Achieves superior quality-throughput trade-offs in summarization, translation, and reasoning.

Enables faster training of block diffusion models with better sequence partitioning.

Provides open-source code and models for practical adoption.

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