TLDR: A new framework called E2D2 introduces an encoder-decoder architecture for discrete diffusion language models, significantly improving both training and inference efficiency while maintaining or enhancing generation quality across various NLP tasks like summarization, translation, and mathematical reasoning. It achieves this by separating the computation for representing clean tokens and denoising corrupted tokens, allowing for a lightweight decoder and periodic encoder calls.
Researchers from Cornell University have introduced a novel approach to discrete diffusion language models, aiming to significantly boost their efficiency in both training and inference. The new framework, dubbed Efficient Encoder-Decoder Diffusion (E2D2), addresses a key limitation of previous diffusion models that relied on a decoder-only architecture, which often led to high computational costs.
Traditional discrete diffusion models are known for their ability to sample tokens in parallel, offering a speed advantage over autoregressive methods. However, their decoder-only design meant that the entire network had to be invoked at every denoising step, making them computationally intensive. The core insight behind E2D2 is that diffusion models perform two distinct types of computation: representing clean tokens and denoising corrupted ones. This distinction allows for the use of separate, specialized modules for each task.
A New Architecture for Faster AI
E2D2 proposes an encoder-decoder transformer architecture. A robust encoder is responsible for creating rich representations of clean tokens, while a lightweight decoder iteratively refines a noisy sequence, conditioned on the encoder’s output. This separation is crucial for efficiency. During inference, the lightweight decoder can be called multiple times to denoise tokens without needing to invoke the larger encoder repeatedly. The encoder is only called periodically to update its representations with newly generated tokens, leading to significantly faster generation.
Beyond faster inference, this encoder-decoder design also accelerates the training of block diffusion models. These models partition sequences into blocks to improve generation quality and support efficient Key-Value (KV) caching. Prior decoder-only block diffusion models incurred high training costs, often doubling the forward pass expense compared to standard diffusion. E2D2’s efficient training algorithm for block diffusion uses the encoder for the clean sequence and the decoder for the noised sequence, effectively halving training costs for models of comparable size.
The E2D2 framework includes not only this innovative architecture but also specialized training and sampling algorithms that enable faster inference and support KV caching. This comprehensive approach allows E2D2 to achieve superior trade-offs between generation quality and inference throughput across various tasks.
Demonstrated Performance Across Tasks
To showcase its effectiveness, the researchers trained domain-specific E2D2 models for summarization, translation, and mathematical reasoning. The results consistently showed that E2D2 outperformed decoder-only diffusion models in both generation speed and quality. For instance, in summarization, E2D2 matched or surpassed diffusion baselines while achieving approximately 75% higher throughput. In machine translation, it delivered higher throughput and better BLEU scores compared to existing block diffusion models.
On the GSM8K mathematical reasoning benchmark, E2D2 again demonstrated improved downstream performance and decoding throughput. By varying the decoder size, the researchers were able to map the Pareto frontier of efficiency and performance, illustrating that E2D2 consistently offered higher quality at various throughput levels, extending the boundaries of what’s possible in terms of speed and accuracy.
In language modeling, E2D2 models trained on OpenWebText achieved comparable perplexities to leading block diffusion models but were 40% faster to train, highlighting its efficiency gains in the training phase as well.
Also Read:
- Variational Masked Diffusion: A New Approach to Capturing Token Dependencies in Generative AI
- Improving Language Model Uncertainty Estimates Through Diverse Sampling
Design Choices and Future Directions
The paper also delves into architectural design decisions, comparing a “last hidden state” variant, where the decoder attends to the final encoder output, and a “shared KV cache” variant, which reuses the encoder’s intermediate Key-Value representations. The choice between these depends on the training scenario, with “last hidden state” performing better for training from scratch on large datasets, and “shared KV cache” proving more stable for fine-tuning from pretrained models on smaller datasets.
E2D2 represents a significant step forward in making discrete diffusion models more practical and efficient for real-world applications. While further innovations are needed to close the performance and throughput gaps entirely with autoregressive models, this work provides a promising direction for the future of language modeling. You can explore the full research paper for more details on this groundbreaking work. Read the full paper here.
