Latent-Kernel Discrete Flow Maps for Few-Step Generation

2026-07-29

AI Summary

This work addresses the challenge that few-step discrete diffusion and flow matching models struggle to capture inter-token dependencies—such as subject-verb agreement—in text generation. To this end, the authors propose Latent Kernel Discrete Flow (LKF), the first framework enabling native few-step generation with correlated updates without relying on teacher distillation. LKF efficiently models cross-position dependencies within a single step by coupling factorized components through shared latent variables, and enhances expressiveness while preserving linear sampling complexity via a latent-variable mixture kernel, closed-form marginalization, and a reuse mechanism. Notably, the authors show that MDLM is a special case of LKF with M=1. Experiments on LM1B and WikiText-103 demonstrate that LKF with M=8 improves generation perplexity by 2.1–3.3× over likelihood-based baselines without sacrificing diversity, outperforming existing distillation and correction approaches.

Abstract

Discrete diffusion and flow-matching models denoise a sequence over many steps, but to keep each step cheap, they factorize the transition across positions and decide every token independently. This makes few-step generation challenging for text when the target couples two positions, such as a subject and a verb that must agree. An independent update commits to them separately, and many function evaluations are spent repairing the mismatch. Existing few-step methods buy back the lost correlation by distilling or rectifying a slow teacher, and so inherit the teacher's quality ceiling. We ask instead whether a model can express correlated steps natively, and answer with Latent-Kernel Discrete Flow Maps (LKF), a from-scratch flow-map kernel that is a mixture of M factorized components tied by a single shared latent. Conditioned on the latent, each component is cheap, and the mixture is summed over the latent in closed form for small M. We show that a single step places mass on correlated completions with the same sampling time complexity as a factorized model, since one latent is drawn per sequence and reused across the entire denoising trajectory. We also show that the Masked Diffusion Language Model (MDLM) is a special case of our LKF model at M=1. The experiments for unconditional text generation on the One-Billion-Word (LM1B) and WikiText-103 benchmarks show that our LKF model learns strongly heterogeneous components and improves generative perplexity by 2.1x to 3.3x over the likelihood baselines without losing diversity. The gain grows with M, and at M=8, it surpasses distilled and rectified few-step samplers. The source code is available at: https://github.com/mansoor181/lkf.git

Problem

• few-step generation
• discrete diffusion
• position correlation
• text generation
• flow matching

Innovation

• Latent-Kernel
• Discrete Flow Matching
• Few-Step Generation
• Correlated Token Modeling
• Mixture Flow Maps