Latent-Kernel Discrete Flow Maps for Few-Step Generation

Discrete diffusion and flow-matching models typically generate text by factorizing transitions across positions, treating each token independently to keep computation cheap. This independence creates a bottleneck for few-step generation, as correlated tokens—such as subject-verb agreement—cannot be resolved in a single step, leading to inconsistent outputs that require many additional iterations to repair. The authors ask whether a model can natively express these correlated steps without relying on distillation or rectification from a slow teacher.

To solve this, the authors propose Latent-Kernel Discrete Flow Maps (LKF). LKF replaces the standard factorized transition with a mixture of M factorized components, tied together by a single shared latent variable. During each denoising step, the model draws one latent variable per sequence, which then conditions the independent sampling of all tokens. Because the latent is shared across the entire sequence, it allows the model to capture global dependencies (like grammatical agreement) while maintaining the computational efficiency of factorized models. The model is trained from scratch using an exact mixture log-likelihood objective, and at inference, it uses a best-of-M decoding strategy to select the most coherent sequence.

Experiments on the One-Billion-Word (LM1B) and WikiText-103 benchmarks demonstrate that LKF learns strongly heterogeneous components. The model achieves a 2.1x to 3.3x improvement in generative perplexity over standard likelihood baselines. Notably, as the number of components (M) increases, LKF surpasses existing distilled and rectified few-step samplers without inheriting their teacher-imposed quality ceilings. The authors also provide theoretical bounds showing that the model's ability to capture total correlation grows with the number of latent components.