Table of Contents
Fetching ...

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

Amin Karimi Monsefi, Nikhil Bhendawade, Manuel Rafael Ciosici, Dominic Culver, Yizhe Zhang, Irina Belousova

TL;DR

FS-DFM, Few-Step Discrete Flow-Matching, a discrete flow-matching model designed for speed without sacrificing quality is introduced, delivering up to 128 times faster sampling and corresponding latency/throughput gains.

Abstract

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.

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

TL;DR

FS-DFM, Few-Step Discrete Flow-Matching, a discrete flow-matching model designed for speed without sacrificing quality is introduced, delivering up to 128 times faster sampling and corresponding latency/throughput gains.

Abstract

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.

Paper Structure

This paper contains 37 sections, 1 theorem, 42 equations, 9 figures, 6 tables, 2 algorithms.

Key Result

Proposition A.1

Let $u_t$ be the infinitesimal generator of a CTMC, then $u_t$ satisfies the following conditions.

Figures (9)

  • Figure 1: Generation quality across model sizes (perplexity and accuracy vs. NFE). FS-DFM reaches the strong-quality regime in few steps across all sizes, while DFM needs far more evaluations. Gold stars (NFE=8) highlight FS-DFM in a few-step regime, with accuracy quickly saturating and entropy converging to similar ranges as steps increase. The average value of entropy for all the models is $7.41$ to $8.07$.
  • Figure 2: Eight-step long-horizon generation: 1 024-token unconditional generation in 8 sampling steps. FS-DFM (0.17B) successfully produces 1 024 tokens under the 8-step constraint. Despite having 40x more parameters, LLaDA-8B-Instruct and Dream-7B-Instruct's 8-steps generations exhibit trailing blanks and punctuation artifacts (e.g., repeated commas). Generations are truncated. Complete output in \ref{['app:sample_outputs']}.
  • Figure 3: RK-2 vs. RK-4 across NFE. Top panels show entropy (linear y) and perplexity (log–log), with ribbons and vertical connectors highlighting pointwise gaps. The bottom shows the deltas ($\Delta$ entropy = RK-4 $-$ RK-2; %$\Delta$ perplexity = RK-4/RK-2 $-1$). RK-4 has consistently lower perplexity ratio $\simeq$$0.88 \times$ with a small entropy trade-off (median 7.63 vs. 7.5), making it the stronger choice for generation over most NFE settings.
  • Figure 4: Discrete flow-matching on a $128\times128$ checkerboard under the instantaneous scale $g(t)$. (a) With an all-[MASK] source, early frames exhibit almost no jumps (stalling near $t\approx0$). (b) A uniform source jumps earlier but still under-updates at tiny $t$. Both patterns motivate the Cumulative Scalar$\bar{g}_{t,h}$ (\ref{['eq:sample:average:velocity']}) to supply the correct probability flow over a finite step.
  • Figure 5: 1 024-token unconditional generation in 8 sampling steps. FS-DFM (0.17B) successfully produces 1 024 tokens under the 8-step constraint. Despite having 40x more parameters, LLaDA-8B-Instruct and Dream-7B-Instruct's 8-steps generations exhibit trailing blanks and punctuation artifacts (e.g., repeated commas). Generations are truncated. Complete version of \ref{['fig:demo']}.
  • ...and 4 more figures

Theorems & Definitions (1)

  • Proposition A.1