FS-DFM: Fast and Accurate Long Text Generation with Few-Step Diffusion Language Models

Autoregressive language models (ARMs) deliver strong likelihoods, but are inherently serial: they generate one token per forward pass, which limits throughput and inflates latency for long sequences. Diffusion Language Models (DLMs) parallelize across positions and thus appear promising for language generation, yet standard discrete diffusion typically needs hundreds to thousands of model evaluations to reach high quality, trading serial depth for iterative breadth. We introduce FS-DFM, Few-Step Discrete Flow-Matching. A discrete flow-matching model designed for speed without sacrificing quality. The core idea is simple: make the number of sampling steps an explicit parameter and train the model to be consistent across step budgets, so one big move lands where many small moves would. We pair this with a reliable update rule that moves probability in the right direction without overshooting, and with strong teacher guidance distilled from long-run trajectories. Together, these choices make few-step sampling stable, accurate, and easy to control. On language modeling benchmarks, FS-DFM with 8 sampling steps achieves perplexity parity with a 1,024-step discrete-flow baseline for generating 1,024 tokens using a similar-size model, delivering up to 128 times faster sampling and corresponding latency/throughput gains. Code & pretrained checkpoints: https://github.com/apple/ml-fs-dfm

💡 Research Summary

FS‑DFM (Few‑Step Discrete Flow‑Matching) tackles the core inefficiency of diffusion‑based language models: the need for hundreds or thousands of refinement steps to achieve high‑quality text. Building on the discrete flow‑matching (DFM) framework, the authors introduce two complementary mechanisms that enable accurate generation with only a handful of steps. First, they make the sampling budget (the number of function evaluations, NFE) an explicit conditioning variable for the model. During training the network receives the desired step count h and is forced, via a self‑consistency loss, to produce the same final distribution whether it takes many small steps or a few large ones. This “step‑aware” design teaches the model to learn a mapping that is invariant to the granularity of the diffusion trajectory. Second, they employ a shortcut teacher based on Runge‑Kutta ODE solvers (Heun’s RK‑2 and classical RK‑4). The teacher runs a high‑resolution (1 024‑step) flow and provides precise conditional token distributions p₁|ₜ(xᵢ|z) for every intermediate time t. These distributions are distilled into the student model, which learns a cumulative scalar S(h) that scales the infinitesimal generator uₜ appropriately for the chosen h. By combining the scalar with the RK‑based update rule, FS‑DFM avoids overshooting and maintains numerical stability even when h is large. The overall training objective combines the original Bregman‑divergence loss for flow‑matching, a self‑consistency term across different h values, and a teacher‑distillation KL term, weighted by hyper‑parameters tuned on a validation set. Experiments span three model sizes (0.17 B, 7 B, and 8 B parameters) and standard text corpora (OpenWebText, C4). Across all scales, FS‑DFM with only 8 NFEs matches the perplexity (≈7.5) and token‑level accuracy (≈0.92) of a baseline DFM that requires 1 024 steps. This translates into up to a 128× speed‑up in wall‑clock inference time, reducing per‑token latency to a few milliseconds on modern GPUs/TPUs—fast enough for real‑time applications. Qualitative analysis on unconditional 1 024‑token generation shows that FS‑DFM produces coherent, punctuation‑correct text with minimal artifacts, whereas competing few‑step attempts (e.g., LLaDA‑8B‑Instruct, Dream‑7B‑Instruct) suffer from repeated commas, blank tokens, and abrupt truncation. Ablation studies reveal that the RK‑2 (Heun) teacher offers the best trade‑off between accuracy and computational cost, while RK‑4 yields marginal gains at higher expense. The step‑aware self‑consistency loss is crucial: removing it degrades performance dramatically for h ≠ 1 024, confirming that explicit budgeting is essential for few‑step reliability. Limitations include a focus on unconditional generation; conditional (prompt‑guided) scenarios remain to be explored, and extreme budgets (1 step or > 512 steps) show slight degradation, suggesting future work on adaptive step scheduling. Overall, FS‑DFM demonstrates that diffusion language models can be re‑engineered to operate efficiently in the few‑step regime without sacrificing the intrinsic advantages of diffusion—parallel token updates, bidirectional context, and flexible controllability. This breakthrough opens the door to deploying diffusion‑based text generators in latency‑sensitive settings such as interactive chatbots, on‑device assistants, and large‑scale content creation pipelines.