Minibatch Optimal Transport and Perplexity Bound Estimation in Discrete Flow Matching

Discrete flow matching, a recent framework for modeling categorical data, has shown competitive performance with autoregressive models. However, unlike continuous flow matching, the rectification strategy cannot be applied due to the stochasticity of discrete paths, necessitating alternative methods to minimize state transitions. We propose a dynamic-optimal-transport-like minimization objective and derive its Kantorovich formulation for discrete flows with convex interpolants, where transport cost depends solely on inter-state dissimilarity and can be optimized via minibatch strategies. We show that such methods can reduce the number of transitions up to 32 times (1024 to 32) to reach the same generative perplexity without compromising diversity. Additionally, path nondeterminism in discrete flows precludes an instantaneous change-of-variables analogue, preventing precise probability estimation available to continuous flows. We therefore propose two upper bounds on perplexity, enabling principled training, evaluation and model comparison. Finally, we introduce Multimask Flows which outperform masked flows in generative perplexity without compromising diversity, particularly when utilizing minibatch Optimal Transport.

Some modern AI systems generate text not word-by-word, but by starting from a rough draft and repeatedly editing it, taking hundreds or even thousands of tiny steps before arriving at a finished passage. Discrete flow matching is a recent and promising method of this kind for text generation, but it has lacked two practical tools: a way to make generation faster, and a reliable way to measure how good the resulting text is. We tackled the speed problem by borrowing a classic idea from mathematics known as optimal transport, which studies the most efficient way to move things from one arrangement to another. By cleverly pairing the model's rough starting points with its targets, we let it reach the same destination while changing fewer words along the way. We also derived formulas that put a trustworthy ceiling on a model's error, so its quality can be checked during training and compared against other systems. These advances let our models generate text of the same quality using up to 32 times fewer editing steps, without making the output any less varied or interesting. This makes this style of text generation more practical and easier to evaluate.