One-Sided Quantile Coupling for Flow Matching

Technical Summary: One-Sided Quantile Coupling for Flow Matching

Problem Statement
Flow Matching (FM) trains continuous-time generative models by regressing the velocity field of a probability path connecting a simple source distribution (typically a standard Gaussian) to a target data distribution. The quality of the coupling—the pairing of source and target samples—critically influences the geometry of training trajectories. Standard FM uses independent random pairing, which often results in long, conflicting paths that increase the variance of the target velocity field and hinder training efficiency.

While structured couplings inspired by Optimal Transport (OT) can straighten these paths and improve generation quality, existing methods (e.g., OT-CFM) rely on solving a mini-batch assignment problem. These approaches require constructing pairwise cost matrices and solving assignment problems (e.g., via the Hungarian algorithm or Sinkhorn iterations), leading to computational costs that scale at least quadratically (or cubically for exact solves) with batch size. This limits scalability and can introduce instability. The paper addresses the question of whether the geometric benefits of structured coupling can be achieved without matching two pre-sampled batches, thereby avoiding the high cost of batchwise transport.

Methodology: Quantile Coupling Flow Matching (QC-FM)
The authors propose Quantile Coupling Flow Matching (QC-FM), a lightweight, one-sided coupling scheme. Unlike traditional methods that sample both noise and data batches and then re-pair them, QC-FM samples only the data batch and constructs the paired source samples directly.

1. One-Sided Construction:

   • Given a data batch $\mathcal{B}_x$, the method projects data points onto $k$ random orthogonal directions (slices) defined by a matrix $U$.
   • For each direction, the data points are ranked within the batch. These empirical ranks are mapped to Gaussian quantiles using the inverse standard normal CDF.
   • The source latent code $e^{(i)}$ is constructed to satisfy these projected quantile constraints. Specifically, the projection of the source onto the slice directions matches the Gaussian quantiles derived from the data ranks.
   • The remaining dimensions (the orthogonal complement of the slice subspace) are filled by sampling from a conditional Gaussian distribution. This ensures the source remains stochastic in unconstrained directions.

2. Gaussian-Remainder Hybrid Extensions:
   Since the full QC coupling (where all source points are constructed via ranks) deviates from the standard i.i.d. Gaussian prior for $k \ge 2$ due to dependencies among slice codes, the authors introduce two hybrid strategies to retain the benefits of QC while preserving the prior:

   • QC-FM-Mixture: Applies the QC construction to an "anchor" subset of the batch (e.g., 20%) and pairs the remaining "remainder" source samples with the remaining data samples using independent random pairing (the Baseline approach). This creates a mixture of structured and unstructured couplings.
   • QC-FM-Adjacency: Applies QC to the anchor subset but pairs the remainder source samples based on the local geometry induced by the anchors. Remainder data points are assigned to their nearest anchor (in the structured subspace), and the remainder source samples are paired within these groups. This propagates the structured bias without duplicating targets.

Key Contributions

• QC-FM Algorithm: A novel, batch-local structured coupling that imposes monotone one-dimensional transport along random slices. It eliminates the need for pairwise cost matrices and assignment solvers, reducing the dominant computational cost to $O(kB \log B)$ (sorting) rather than $O(B^2)$ or $O(B^3)$.
• Theoretical Analysis:
  • Variance Reduction: The authors prove that for a fixed frame, the QC coupling eliminates the irreducible regression variance along every selected slice, making the ideal flow exactly straight in those directions.
  • Prior Preservation: They quantify that for $k=1$, the source marginal is exactly Gaussian. For $k \ge 2$, the deviation from the Gaussian prior is confined to the copula of the slice codes, and the transport cost of this deviation is bounded.
• Hybrid Training Schemes: The introduction of Mixture and Adjacency variants that balance the injection of geometric bias with the preservation of the explicit Gaussian signal from the Baseline.

Experimental Results
The method was evaluated on CIFAR-10, CelebA-64, FFHQ-64, and ImageNet-64 under matched training budgets and architectures.

• Performance: QC-FM-Mixture consistently outperformed the independent-coupling Baseline and the exact Hungarian OT-CFM across all four datasets. Improvements over the Baseline ranged from 6.3% (CIFAR-10) to 12.9% (FFHQ) in terms of FID-50K.
• Efficiency: In terms of computational cost, QC-FM is significantly faster than OT-based methods. At a batch size of 2048, QC-FM-Mixture was over 800 times faster than the Hungarian OT solver and 35 times faster than the Sinkhorn solver.
• Toy Diagnostics: On synthetic datasets (Checkerboard and 8-Gaussians), QC-FM variants produced nearly straight learned trajectories (Path-Length Ratio $\approx 1$) and significantly lower target velocity variance compared to the Baseline, confirming the geometric benefits of the coupling.

Significance and Claims
The paper claims that preserving projected rank structure is a simple, scalable, and effective way to inject useful geometric bias into Flow Matching couplings without solving a mini-batch transport problem.

• Scalability: QC-FM offers an attractive middle ground between the unstructured Baseline and expensive batchwise transport methods. It is computationally lightweight and controllable via the number of projections ($k$).
• Practicality: The method requires no learned components for inference; generation still starts from the standard Gaussian. The coupling is constructed on-the-fly during training.
• Limitations: The authors modestly note that QC-FM is a batch-local surrogate for structured coupling, not an approximation of global Optimal Transport. The quantiles are estimated within mini-batches, and the theory controls per-slice variance rather than guaranteeing global trajectory straightness or specific FID improvements. The selection of slice count ($k$) and anchor ratio ($p$) currently relies on heuristics and dataset-specific tuning.

In conclusion, the work demonstrates that one-sided quantile coupling can improve sample quality and training efficiency in Flow Matching by structuring the noise-data pairing through rank preservation, offering a practical alternative to computationally prohibitive optimal transport methods.