DiffuMamba: High-Throughput Diffusion LMs with Mamba Backbone
Abstract
Diffusion language models (DLMs) have emerged as a promising alternative to autoregressive (AR) generation, yet their reliance on Transformer backbones limits inference efficiency due to quadratic attention or KV-cache overhead. We introduce DiffuMamba, a masked diffusion language model built on a bidirectional Mamba backbone that combines the diffusion objective with linear-time sequence modeling, and DiffuMamba-H, a hybrid variant with interleaved attention. Across scales up to 1.3B parameters, our models match Transformer-based diffusion in downstream performance while achieving up to 8.2× and 4.3× higher inference throughput, respectively, on long sequences. We further present a systematic analysis of inference efficiency across modern DLM variants, combining asymptotic complexity with empirical measurements. Notably, cache-efficient block diffusion with Mamba mixers emerges as the only strategy that scales linearly with sequence length and achieves the strongest performance across all baselines, suggesting a promising direction for future diffusion-based generation systems.
Lay Summary
Today's AI chatbots write one word at a time, like reading a book aloud, fast for short notes but slow for long ones. Diffusion language models instead generate whole paragraphs in parallel by polishing gibberish into text, like a sculptor at work. But each polishing pass uses attention, which compares every word to every other word, making the cost explode at long lengths. We replaced attention with Mamba, an engine that reads text in a single sweep instead of cross-comparing every pair of words. Our two models: DiffuMamba and the hybrid DiffuMamba-H plug this engine into the same diffusion recipe, trained up to 1.3B parameters. On standard benchmarks, they match or beat the attention baseline, generate text up to 8.2× faster on long inputs, and stay stable at lengths where the attention baseline collapses.