Discrete Wasserstein Flows for One-Step Generative Modeling
This paper extends one-step generative modeling to finite discrete spaces by replacing Euclidean drift with a discrete Wasserstein gradient flow. Using Maas geometry on a reversible Markov kernel, it defines a relative-KL flow that is realized at the particle level by Markov jumps and amortized into a neural generator. The result is a theoretically grounded analogue of continuous drifting that preserves one-step inference. The work is validated only on a small controlled torus, leaving real-world scalability as open future work.
The paper delivers a rigorous theoretical construction and a clean, controlled empirical validation, but its scope is deliberately narrow. The authors correctly identify that Euclidean drifting fails on finite spaces because particles cannot move infinitesimally, and they build a sound alternative using Maas discrete Wasserstein geometry with logarithmic-mean mobility. Their synthetic experiments on a 32x32 torus exactly verify KL dissipation, numerical scaling, and neural amortization feasibility. Nevertheless, the work stops short of demonstrating scalability or performance on real discrete data, and the conclusion explicitly flags this as future work.
The strongest aspects are the rigorous geometric formulation and the clarity of the empirical decomposition. The authors derive an exact normalization-free logit drift and a minimal positive-part particle realization that matches the law-level Euler step to first order. The experimental design is exemplary: by first running an exact tabular model on an enumerable $|\mathcal{X}|=1024$ torus, they isolate transport correctness from neural amortization error. The neural generator tracks the exact conditional oracle closely as the latent batch size grows, indicating that residual error stems from fitting rather than the underlying transport.
The primary limitation is the absence of any real structured discrete data. The entire empirical evaluation lives on a 32x32 periodic grid with 1,024 enumerable states, where densities, currents, and teachers can be computed exactly. This leaves scalable local ratio estimation and generalization to combinatorial state spaces completely unaddressed. There is also no empirical comparison to existing discrete diffusion or other one-step baselines, so it remains unclear whether the proposed transport offers practical advantages in wall-clock time or sample quality.
The evidence fully supports the claims within the paper's restricted scope. The tabular experiment confirms that the discrete transport step dissipates KL and obeys the predicted $\mathcal{O}(\eta^{2})$ local consistency, while the particle realization matches the law-level Euler step to machine precision. However, comparisons to related work remain conceptual. The authors correctly contrast their method with discrete diffusion by noting that inference avoids simulating a Markov process, but they do not run SEDD or any other baseline, and the citation to Deng et al.'s drifting models is used only as a continuous analogue.
Reproducibility is excellent in principle for the synthetic benchmark but incomplete for broader adoption. Appendices B and C provide exhaustive hyperparameters, architectures, optimizer settings, random seeds, and analytic target formulas, which allow exact replication of the torus experiments. Yet no code repository is mentioned or linked, and because the method has not been exercised on standard benchmarks, there is no way to verify whether the logarithmic-mean mobility or nearest-neighbor kernel transfer to non-synthetic domains. Independent reproduction on real discrete data is therefore blocked by both missing code and untested scalable estimation procedures.
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.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.