FS-DFM (Few-Step Discrete Flow-Matching)
FS-DFM: Fast and Accurate Long Text Generation with Few-Step Diffusion Language Model
FS-DFM is a token-space diffusion / flow-matching language model designed for fast long-text generation by explicitly training for a user-specified step budget (e.g., 1–8 steps), while preserving a CTMC-based discrete flow formulation.
What’s in this repo
Checkpoint files
FS_DFM_checkpoint.pthFS-DFM 1.3B , uniform source, RK4 teacher distilled DFM_checkpoint.pthDFM 1.3B , uniform source, DFM pretrained initialization
Model summary
Core idea (high level):
Condition the model on a target inference step size/budget and train it so that one big step matches many small steps .
Use a cumulative scalar update to make large steps stable on the probability simplex.
Use student–teacher distillation (Runge–Kutta shortcut teachers, EMA stabilization) to improve few-step fidelity.
Formulation: discrete flow-matching over a CTMC on token sequences; sampling uses custom solvers (e.g., mixture_euler_with_cumulative_scalar).
Comparison of Methods
ARM DFM FS-DFM (Ours) ARM DFM FS-DFM
Architecture
From the paper’s implementation details:
Backbone is a DiT-style transformer with rotary attention
Adaptive LayerNorm conditioning in each blockConditioning includes continuous time embedding + step-size embedding
Final linear head produces logits; conversion from logits to a CTMC generator + stepping happens in the solver
Tokenizer: GPT-2 tokenizer 1024-token blocks (EOS appended, then packed/concatenated).
Training data & evaluation data
Training: FineWeb-Edu
Evaluation: WikiText-103
(See the paper for details and the exact preprocessing pipeline.)
How to use
FS-DFM uses custom discrete solvers and is not a drop-in transformers model. The intended usage is via the official training/evaluation scripts.
1) Install the official code
1 git clone https://github.com/apple/ml-fs-dfm
2 cd ml-fs-dfm
3
4 conda env create -f fsdfm_environment.yml
5 conda activate FSDFM
6
7 pip install -e .
