Discrete Wasserstein Flows for One-Step Generative Modeling

2026-10-01

AI Summary

This study addresses the limitation that continuous-domain drift methods cannot be directly applied to discrete state spaces, which hinders the development of one-step generative modeling. To overcome this, the paper presents the first extension of the drift framework to finite state spaces by defining a KL gradient flow based on discrete Wasserstein geometry and reversible Markov kernels, where probability evolution is achieved through particle-level Markov jumps. The approach employs iterative optimization during training while enabling one-step generation at inference, further incorporating a latent conditional neural generator to enhance modeling capacity. Experimental results validate the KL dissipation properties and numerical scaling laws, demonstrating that finite-capacity neural networks can effectively track exact transport targets while preserving efficient one-step generation. This work thereby overcomes key theoretical bottlenecks in discrete data generation.

Abstract

We introduce a new framework for one-step generative modelling on finite state spaces. To extend drifting beyond continuous domains, we use discrete Wasserstein geometry to define a target-relative KL gradient flow over the transitions of a reversible Markov kernel. We realize this probability flow at the particle level through Markov jumps and amortize the resulting transport updates into a latent-conditioned generator, so that the iterative dynamics are required only during training while inference remains one-step. In a controlled setting where the underlying distributions and transport dynamics can be computed exactly, we verify KL dissipation, consistency between the particle dynamics and the probability flow, and the predicted numerical scaling. We further show that a finite-capacity neural generator can track these exact transport targets while retaining one-step generation. These results validate the basic construction and provide a foundation for scaling Discrete Drifting to structured discrete data.

Problem

• discrete generative modeling
• one-step generation
• Wasserstein flows
• finite state spaces
• Markov kernel

Innovation

• Discrete Wasserstein Flows
• One-Step Generative Modeling
• KL Gradient Flow
• Markov Jumps
• Amortized Transport