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

TL;DR

PairFlow is a lightweight preprocessing method that enables few-step sampling in discrete flow models without requiring a pretrained teacher, significantly reducing training overhead while maintaining or improving performance.

Contribution

We propose PairFlow, a novel closed-form source-target coupling technique for training discrete flow models efficiently without finetuning or pretrained teachers.

Findings

Achieves up to 98.3% reduction in compute compared to full training.

Matches or surpasses performance of two-stage finetuning methods.

Effective on molecular, binary, and RGB image data.

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.