Minibatch Optimal Transport and Perplexity Bound Estimation in Discrete Flow Matching

TL;DR

This paper introduces a novel discrete flow matching framework using minibatch optimal transport, reducing transitions significantly and proposing bounds for perplexity estimation, with Multimask Flows outperforming previous methods.

Contribution

It develops a dynamic optimal transport-based minimization for discrete flows, derives its Kantorovich formulation, and introduces Multimask Flows for improved generative performance.

Findings

Reduced number of state transitions by up to 32 times.

Proposed perplexity bounds enable better model evaluation.

Multimask Flows outperform masked flows in perplexity without losing diversity.

Abstract

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 similarity 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.