Source-linked AI summary
PairFlow: Closed-Form Source-Target Coupling for Few-Step Generation in Discrete Flow Models
Mingue Park, Jisung Hwang, Seungwoo Yoo, Kyeongmin Yeo, Minhyuk Sung
TL;DR
DFMs provide discrete generation but sample slowly because their iterative, factorized dynamics struggle in few-step settings, while existing acceleration relies on costly finetuning. PairFlow uses closed-form inversion and source–target coupling as lightweight preprocessing without a pretrained teacher. It achieves comparable or better few-step performance across molecules and images, while using at most 1.7% of full training compute and improving subsequent distillation.
Problem
DFMs sample slowly because generation is iterative, and existing acceleration methods impose substantial finetuning overhead.
Method
PairFlow constructs source–target training pairs through closed-form velocities and backward inversion, without requiring a pretrained teacher.
Results
PairFlow matches or surpasses distilled models across molecular and image benchmarks while preprocessing requires at most 1.7% of full model-training compute.
Takeaways & Limitations
PairFlow provides a low-cost route to few-step DFM sampling and produces stronger base models for subsequent distillation.
Takeaways & Limitations
Forward closed-form pairing can fail to cover all target data because multiple source samples may map to the same target, requiring impractically many source samples.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
DFMs enable discrete generation but are slow because sampling is iterative, while existing acceleration methods add substantial finetuning cost. PairFlow constructs source–target pairs through closed-form velocities, achieving few-step sampling with lightweight preprocessing and strong results across molecules and images.
- Motivation: Distillation-based acceleration can require 10–20% of base-model training time, motivating a preprocessing approach requiring only tens of GPU minutes.
- Contribution: PairFlow constructs paired source–target samples using closed-form velocities, avoiding a pretrained teacher and enabling few-step DFM sampling without finetuning.The pairing algorithm is fully parallelizable and requires at most 1.7% of full model-training compute.
- Method: Closed-form forward velocity is determined by Hamming distance in uniform-state DFMs, mapping source samples toward targets while revealing a challenge from many-to-one mappings.The resulting coupling must address insufficient target coverage when multiple source samples map to one target.
- Results: PairFlow matches or surpasses distilled models while using up to 143× less finetuning computation on molecular benchmarks and achieves comparable MNIST-Binary performance up to 35× faster.
- Results: Experiments cover molecular datasets and 2D images, including QM9, ZINC-250k, MNIST-Binary, and CIFAR-10.
2 RELATED WORK
Related work extends flow-matching and distillation ideas to accelerate discrete generation, but approaches trade optimization simplicity or training cost for faster sampling. PairFlow is positioned against timestep optimization, mixture modeling, consistency distillation, and iterative coupling refinement.
- Discrete Flow Models: Discrete flow models extend flow matching to high-dimensional and structured data, including scientific and generative applications.
- Discrete Flow Models: Uniform-state DFMs have self-correcting properties but degrade markedly in few-step settings, limiting efficient generation under tight compute budgets.
- Acceleration Methods: Prior acceleration methods optimize sampling timesteps, add mixture-model loss terms, or adapt consistency distillation to reduce decoding errors and improve parallelism.
- Coupling-Based Acceleration: ReDi iteratively optimizes data–noise pairs, drawing on ReFlow’s straight-flow idea to reduce conditional total correlation for few-step generation.
3 PRELIMINARIES
Discrete Flow Matching learns a probability path from an easy-to-sample source distribution to an unknown data distribution and transports source samples along it. Few-step generation is difficult because factorized token updates mismatch correlated joint transitions, motivating path straightening and coupling refinement.
- 3.1 DISCRETE FLOW MATCHING: DFM learns a path p_t connecting source p to target q, then transports samples drawn from p to generate data from q.
- 3.1 DISCRETE FLOW MATCHING: The model approximates the marginal velocity field with a network whose outputs are categorical distributions over token values.
- 3.1 DISCRETE FLOW MATCHING: Conditional paths interpolate source samples x_0 and data samples x_1 through token-wise mixture paths governed by an increasing scheduler κ_t.
- 3.1 DISCRETE FLOW MATCHING: Sampling integrates the learned velocity field over [0, 1] using discrete update steps controlled by step size h.
- 3.2 STRAIGHTENING PROBABILITY PATHS FOR ACCELERATED SAMPLING: ReFlow addresses curved paths by training on teacher-generated source–target pairs, while discrete analogues iteratively refine couplings and optimize the DFM objective.
- 3.2 STRAIGHTENING PROBABILITY PATHS FOR ACCELERATED SAMPLING: Few-step generation suffers when factorized per-token updates miss correlated joint transitions, so distillation methods reduce conditional total correlation to straighten trajectories.
4 PAIRFLOW
PAIRFLOW constructs well-aligned source–target pairs for discrete flow models using closed-form velocity fields and a lightweight preprocessing procedure, without pretrained models. Backward inversion ensures dataset coverage while producing straighter paths, and the preprocessing requires only a small fraction of full-training compute.
- PAIRFLOW: PAIRFLOW discovers well-aligned source–target pairs directly from data using closed-form velocity fields, without relying on pretrained models.The framework assumes a uniform source distribution and extends closed-form velocity construction to discrete domains.
- Finding pairs via closed-form backward velocity fields: The pair-construction procedure is summarized by computing the closed-form backward quantities and applying the backward update for a fixed number of iterations.The overall procedure is described as fully algorithmic preprocessing before model training.
- Finding pairs via closed-form forward velocity fields: The resulting inversion lowers average Hamming distance to 6.47 versus 9.0 for UDLM, promoting straighter training paths.Figure 1 contrasts PAIRFLOW inversion with UDLM’s standard corruption process.
- Finding pairs via closed-form forward velocity fields: The forward construction weights dataset sequences by Hamming-distance proximity, but generated samples may fail to cover the empirical target distribution efficiently.This coverage limitation motivates the backward construction as an alternative.
- Finding pairs via closed-form backward velocity fields: PAIRFLOW backtraces data samples toward the uniform source, guaranteeing inclusion of all data points in the resulting pairs and keeping sources closer to the original data.The backward velocity is obtained through a closed-form noise predictor and then applied iteratively to construct source–target pairs.
- PAIRFLOW: Table 1 reports dataset and training statistics, including token count, vocabulary size, dataset size, and wall-clock runtimes relative to base training.The table uses RTX A6000 wall-clock measurements and reports proportions relative to TBase.
5 EXPERIMENTAL RESULTS
Experiments across molecular and image benchmarks show that PAIRFLOW improves few-step generation over UDLM and is competitive with, or better than, distillation-based methods while supporting stronger subsequent distillation.
- 5 EXPERIMENTAL RESULTS: PAIRFLOW improves UDLM across molecular and image benchmarks, matching or surpassing distillation-based methods in few-step generation.The evaluation covers QM9, ZINC-250k, MNIST-Binary, and CIFAR-10, with molecular validity, uniqueness, novelty, FID, and IS metrics.
- 5.2 MOLECULE GENERATION: On QM9, 1-step validity rises from 17.5 to 223.4, a 12.8× improvement, while 2-step and 4-step validity improve by 231% and 47.6%.PAIRFLOW’s 2-step and 4-step validities are comparable to UDLM’s 4-step and 8-step results, respectively.
- 5.2 MOLECULE GENERATION: PAIRFLOW achieves comparable or better molecular performance than DCD and ReDi with less than 2% training overhead and no pretrained model.On QM9, 2-step validity is 416.0 for PAIRFLOW versus 530.8 for UDLM + DCD; on ZINC-250k, it is 146.3 versus 75.9 for UDLM + ReDi.
- 5.3 IMAGE GENERATION: On MNIST-Binary, PAIRFLOW reduces 1-step FID by 68.9% to 40.59 and reduces 2-step FID by 63.3% to 15.61 versus UDLM.On CIFAR-10, PAIRFLOW outperforms UDLM, while DCD and ReDi degrade performance according to FID and IS.
- 5.3 IMAGE GENERATION: PAIRFLOW preprocessing takes 1.4 minutes, yielding 28.6× and 35× speedups over DCD and ReDi preprocessing, respectively.The comparison uses 40 minutes for DCD and 49 minutes for ReDi, versus 1.4 minutes for PAIRFLOW.
- 5.4 DISTILLING MODELS TRAINED WITH ALIGNED PAIRS: Distilling PAIRFLOW further improves results, with QM9 1-step validity of 453.8 versus 323 for UDLM+DCD and ZINC-250k 2-step validity of 221.5 versus 75.9 for UDLM+ReDi.These gains require only 0.77% additional preprocessing cost on QM9 and 6.42% on ZINC-250k relative to distillation runtime.
6 CONCLUSION
PAIRFLOW accelerates Discrete Flow Models through lightweight preprocessing that couples source and target samples using closed-form inversion, requiring no pretrained teacher. It costs up to 1.7% of full training while achieving comparable or superior performance to finetuned alternatives.
- PAIRFLOW requires up to 1.7% of base-model training compute, uses no pretrained teacher, and achieves comparable or superior performance to finetuned alternatives.Its closed-form inversion enables source–target coupling before training rather than acceleration through finetuning.
- Closed-form forward and backward velocities are derived under a uniform source-distribution assumption.The appendix provides proofs for both velocity formulations and notes that the backward expression at t = 1 is interpreted through its limiting value.
- PAIRFLOW improves UDLM on CIFAR-10 while remaining competitive with distilled models across sampling settings.The comparison reports FID and IS using 1,000 steps and 50K samples.
C.1 COVERAGE OF TRAINING DATASET BY SAMPLING WITH FORWARD VELOCITY
Sampling source points with the closed-form forward velocity does not fully cover the training dataset when the number of draws equals the dataset size. This motivates the backward-velocity construction used by PAIRFLOW.
- Forward-velocity pairing has substantially higher cost than backward-velocity pairing when full training-set coverage is desired.The forward procedure samples source points and transports them toward dataset points, but duplicate recovery limits coverage.
- Coverage is defined as the ratio of unique recovered samples to the training-set size.The definition is COV = k̄/|X1|.
- Sampling |X1| source points recovers only about 63% of the training distribution in practice.The section compares empirical and theoretical coverage across datasets using unique generated samples.
D DETAILED EXPERIMENTAL RESULTS
Detailed experiments evaluate molecular validity, uniqueness, and novelty alongside image FID and IS across multiple sampling budgets. The reported tables show PAIRFLOW results across QM9, ZINC-250k, MNIST-Binary, and CIFAR-10.
- Detailed evaluation covers 1,024 molecular samples over 10 trials and reports FID or FID plus IS for image datasets.QM9 and ZINC-250k use validity, uniqueness, and novelty; MNIST-Binary uses FID, while CIFAR-10 uses FID and IS.
- The ZINC-250k and image results are organized across validity, uniqueness, novelty, FID, and IS tables for different sampling steps.Tables 9–13 summarize these molecular and image metrics, with best values highlighted by column.
E EXPERIMENT ON CONTINUOUS FLOW MATCHING
Continuous experiments adapt closed-form source–target pairing to rectified flows and evaluate both MNIST and dimension-varying two-moons data. Closed-form pairing improves few-step MNIST performance, while its gains diminish as synthetic dimension increases.
- E. EXPERIMENT ON CONTINUOUS FLOW MATCHING: The continuous extension trains rectified-flow models on source–target pairs constructed with the continuous variant of PAIRFLOW.The experiments include MNIST and dimensions from 2 through 256 for the product two-moons dataset.
- E.1 CLOSED-FORM VELOCITY IN CONTINUOUS FLOW MATCHING: The continuous velocity is the conditional drift for a linear probability path and has a closed form computable from the source and target distributions.The construction uses independent source and target variables and an empirical target distribution when data samples are available.
- E.1 CLOSED-FORM VELOCITY IN CONTINUOUS FLOW MATCHING: Closed-form pairing yields better MNIST FID than independent pairing in few-step settings and preserves performance at many steps.After distillation, PAIRFLOW+RF achieves lower FID than ReFlow applied to the independently paired baseline at every sampling budget.
- E.3 CONTINUOUS RECTIFIED FLOWS ON DIMENSION-VARYING SYNTHETIC DATA: As data dimension increases in the N-fold two-moons experiment, the improvement from closed-form pairing decreases.This trend suggests a practical limitation for high-dimensional continuous data; Chamfer distance is measured using the first two coordinates.
F ADDITIONAL QUALITATIVE RESULTS
Additional experiments show that PAIRFLOW maintains strong few-step generation across molecular, image, higher-dimensional, and large-vocabulary settings, while subset pairing reduces preprocessing cost without sacrificing comparable performance.
- QM9 re-flow iterations: PAIRFLOW outperforms UDLM on QM9 across re-flow settings, including comparisons against UDLM with multiple additional iterations.Validity, uniqueness, and novelty are averaged over 10 independent runs across timesteps and re-flow iterations.
- ZINC-250k subset pairing: Subset pairing on ZINC-250k reduces pairing computation time while maintaining performance comparable to the full-set baseline.The evaluation covers validity, uniqueness, novelty, and pairing-time cost across subset partitions and NFEs.
- Additional domains: Additional evaluations extend the method to FFHQ at 64 × 64 resolution and LM1B across sequence lengths of 16, 32, 64, and 128.FFHQ is evaluated with FID, while LM1B uses generative perplexity and entropy across NFE settings.
J ANALYSIS FOR THE OVERFITTING IN IMAGE DOMAINS
Image-domain evaluations test whether PAIRFLOW’s improved generation quality reflects memorization. Across CIFAR-10 and MNIST-Binary, quality improves over baselines while nearest-neighbor distances remain comparable or slightly larger.
- Generation quality: PAIRFLOW delivers improved generation quality over baseline models across CIFAR-10 and MNIST-Binary evaluation metrics.CIFAR-10 reports FID and FID-Dino, while MNIST-Binary reports FID across multiple NFE settings.
- Nearest-neighbor analysis: Nearest-neighbor distances for PAIRFLOW are comparable to or slightly larger than those of reference models across the evaluated image settings.The analysis uses pixel-space ℓ2 distances for MNIST-Binary and ℓ2 plus DINOv2 cosine distances for CIFAR-10.
K LIMITATIONS AND FUTURE WORK
The paper identifies reduced benefits in higher-dimensional settings and frames continuous-flow extension as future work, while emphasizing PAIRFLOW’s fit for low-dimensional discrete scientific data.
- Future impact: The paper connects reducing training compute with potential reductions in training energy consumption and broader access to foundation-model development.These implications are presented as motivation for reducing compute while retaining fast generation.
- Continuous-flow extension: Initial continuous Flow Matching experiments indicate effectiveness on relatively low-dimensional synthetic data, but the paper expects more source–target pairs may be needed in continuous settings.The authors identify further investigation of continuous data as future work.
- Dimensionality limitation: PAIRFLOW’s advantage diminishes for higher-dimensional data, as shown by the shrinking margin over standard CondOT with increasing dimension.The two-moons analysis reports Chamfer distance against sampling steps and finds the strongest gains at few sampling steps.
- Supported scope: The authors emphasize that PAIRFLOW is particularly suited to low-dimensional discrete data, including molecular and protein structures.This scope is presented as a current strength while broader continuous-data effects remain under investigation.