Discrete Wasserstein flows enable one-step discrete data generation
Discrete Wasserstein Flows for One-Step Generative Modeling
Machine LearningArtificial Intelligence
Summary
Creating new examples of data that come in fixed categories, like text or graphs, is tricky because these spaces are not continuous. The authors developed a method that carefully moves probabilities around these finite categories using ideas from Wasserstein geometry, a way of measuring distances between distributions. This method transforms data in one step during use but trains through many small adjustments. Their tests confirm the method works as expected and can be tracked by a neural network that generates new discrete data quickly.
What this means in practice
- •For machine learning engineers: Generate discrete structured data such as categorical sequences or graphs in a single step after training without iterative sampling.
- •For data compression teams: Implement fast generative models that produce compressed representations of discrete data with controlled probability flows.
Authors
Alessandro Micheli, Andrea Zerio, Samir Bhatt
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.