PairFlow: Closed-Form Source-Target Coupling for Few-Step Generation in Discrete Flow Models

TL;DR

This work addresses slow sampling in discrete flow models by introducing PairFlow, a lightweight preprocessing step that constructs well-aligned source–target pairs without requiring a pretrained teacher. It introduces closed-form forward and backward velocity fields in discrete spaces to build paired samples directly from data, enabling few-step generation with at most $1.7\%$ of base training compute. Across molecular datasets (QM9, ZINC-250k) and image benchmarks (MNIST-Binary, CIFAR-10), PairFlow matches or surpasses distillation-based acceleration while dramatically reducing preprocessing costs; when followed by distillation, PairFlow-trained bases further improve downstream performance. The method broadens the practicality of discrete flow models for real-world discrete data, with strong implications for efficiency in molecular generation and high-dimensional discrete vision tasks. It also establishes a principled connection between velocity-based inversions and source–target coupling quality in discrete flows, suggesting avenues for future integration with continuous-flow techniques.

Abstract

We introduce $\texttt{PairFlow}$, a lightweight preprocessing step for training Discrete Flow Models (DFMs) to achieve few-step sampling without requiring a pretrained teacher. DFMs have recently emerged as a new class of generative models for discrete data, offering strong performance. However, they suffer from slow sampling due to their iterative nature. Existing acceleration methods largely depend on finetuning, which introduces substantial additional training overhead. $\texttt{PairFlow}$ addresses this issue with a lightweight preprocessing step. Inspired by ReFlow and its extension to DFMs, we train DFMs from coupled samples of source and target distributions, without requiring any pretrained teacher. At the core of our approach is a closed-form inversion for DFMs, which allows efficient construction of paired source-target samples. Despite its extremely low cost, taking only up to 1.7% of the compute needed for full model training, $\texttt{PairFlow}$ matches or even surpasses the performance of two-stage training involving finetuning. Furthermore, models trained with our framework provide stronger base models for subsequent distillation, yielding further acceleration after finetuning. Experiments on molecular data as well as binary and RGB images demonstrate the broad applicability and effectiveness of our approach.

Figures (10)

• Figure 1: Illustrations of data inversion in PairFlow and the standard corruption process in UDLM. PairFlow achieves a lower average Hamming distance (6.47 vs. 9.0), promoting straighter paths during training.
• Figure 2: Step-wise performance analysis on the QM9 dataset ramakrishnan2014quantum. Each plot reports the number of valid (left), unique (middle), and novel (right) SMILES strings weininger1988smiles out of 1,024 generated samples. Best viewed when zoomed in.
• Figure 3: Step-wise performance analysis on the ZINC-250k dataset irwin2012zinc. Each plot reports the number of valid (left), unique (middle), and novel (right) SMILES strings weininger1988smiles out of 1,024 generated samples. Best viewed when zoomed in.
• Figure 4: Step-wise performance analysis on discretized image datasets. From left to right: FID on MNIST-Binary lecun2002gradient, FID on CIFAR-10 krizhevsky2009learning, and Inception Scores (IS) on CIFAR-10. Best viewed when zoomed in.
• Figure 5: Qualitative results of 1-step generation on MNIST-Binary ($28 \times 28$; left) and 64-step (top right) and 256-step (bottom right) generation on CIFAR-10 ($32 \times 32$).