PaperlayerView exact revision on arXiv

Source-linked AI summary

Beyond Straightness: Non-Crossing Flow Matching via Quantile AlignTree Coupling

Junyi Lin, Mengyu Li, Jingxuan Hu, Kejun He, Cheng Meng

arXiv:2609.03443v1cs.LG

TL;DR

Flow Matching depends on source–target coupling, but crossing paths can create local velocity ambiguity and OT-based coupling can be costly. QAT-FM uses tree-structured quantile alignment; it achieves competitive generation quality at lower construction complexity than OT-based methods, with non-crossing paths and greater intermediate separation than independent coupling.

  • Problem

    Flow Matching performance depends on source–target coupling, since crossing or closely passing paths with different target velocities create local velocity ambiguity.

    Source: §1 · pp. 1–2

  • Method

    QAT uses binary tree-structured quantile alignment to construct a coupling between an empirical data distribution and a continuous Gaussian prior.

    Source: §1 · p. 2

  • Results

    Experiments report competitive generation quality with substantially lower coupling construction complexity than OT-based methods.

    Source: §1 · p. 3

  • Takeaways & Limitations

    QAT-FM provides an efficient structured coupling mechanism for Flow Matching, with non-crossing paths and improved path separation that extends naturally to conditional generation.

    Source: §1 · p. 3, §5 · p. 10

Abstract

from arXiv · show

The performance of Flow Matching largely depends on the quality of the coupling between the source and target distributions. However, independent coupling often leads to path crossings and local velocity ambiguity, while OT-based couplings typically incur high construction costs. To address this challenge, we propose Quantile AlignTree Flow Matching (QAT-FM), an efficient structured coupling strategy that constructs a hierarchical coupling between a Gaussian prior and the target data distribution via a quantile-aligned tree structure. QAT-FM constructs the coupling in $\mathcal{O}(Nd\log N)$ time and supports per-pair source sampling with $\mathcal{O}(d)$ complexity, enabling scalable training for large-scale high-dimensional generative tasks. Theoretically, we prove that the QAT coupling satisfies marginal consistency, induces non-crossing linear interpolation paths, and consistently improves path separation at intermediate times compared with independent coupling, thereby alleviating local velocity ambiguity. QAT-FM further extends naturally to conditional generation, enabling structured conditional coupling while preserving global Gaussian alignment. Experiments across diverse benchmark datasets demonstrate that QAT-FM achieves competitive generative performance while substantially reducing coupling construction cost.

1 INTRODUCTION

Flow Matching depends on the source–target coupling: crossing or closely passing paths can create local velocity ambiguity, while existing approaches trade off coupling cost and favorable path geometry. QAT-FM uses tree-structured quantile alignment to provide an efficient coupling with favorable path structure, and experiments report competitive generation quality at lower construction complexity than OT-based methods.

1 INTRODUCTION

  • Crossing paths or paths passing closely through the same intermediate region with substantially different target velocities create local velocity ambiguity in Flow Matching.Under mean-squared-error training, inconsistent velocity directions are averaged, increasing the difficulty of fitting the velocity field.Source: §1 · pp. 1–2
  • Existing approaches trade off coupling quality and cost: reflow adds sampling and retraining, exact OT is expensive at scale, and mini-batch OT does not preserve full-dataset non-crossing structure.Condition-aware constraints or condition-wise couplings can further increase OT complexity in conditional generation.Source: §1 · p. 2
  • QAT constructs an exact coupling between an empirical data distribution and a continuous Gaussian prior in O(Nd log N) time, with O(d) per-pair sampling.The tree-structured quantile alignment also supports multiple split strategies and can incorporate conditioning information into tree construction and coupling.Source: §1 · p. 2
  • Experiments show competitive generation quality with substantially lower coupling construction complexity than OT-based methods.The experiments span synthetic data, unconditional and conditional generation, and large-scale image synthesis.Source: §1 · p. 3
  • QAT paths are non-crossing, and their intermediate path separation is uniformly greater than under independent coupling, alleviating local velocity ambiguity.In conditional generation, coupling strategies must also preserve alignment between samples and conditions.Source: §1 · pp. 2–3

2 PRELIMINARIES

The preliminaries introduce flow-based generative modeling and Flow Matching, then motivate the coupling problem and tree-based construction used by the paper. The choice of endpoint coupling governs trajectory geometry, while exact global OT is costly at large scale.

2 PRELIMINARIES

  • A flow map is the solution to an ODE driven by a time-dependent velocity field, and it transports the source distribution to the target distribution.Flow-based generative modeling learns a parameterized velocity field whose induced flow map sends the source distribution to the target.Source: §2 · p. 3
  • Flow Matching trains a parameterized velocity field by regression along a specified probability path, avoiding backpropagation through numerical ODE solvers during training.Conditional Flow Matching samples paired endpoints from a coupling and uses a conditional path and compatible target velocity; linear interpolation is a canonical choice.Source: §2 · p. 3
  • The endpoint coupling governs particle-trajectory geometry, but exact global OT is expensive at large scale, motivating binary tree-based space partitioning.Mini-batch OT reduces computation at the expense of coupling exactness and cross-batch consistency.Source: §2 · pp. 3–4

3 QUANTILE ALIGNTREE FLOW MATCHING

QAT transfers a binary data-tree topology to the source by matching split mass ratios, creating a structured coupling. QAT-FM retains linear interpolation and velocity regression while this coupling yields non-crossing paths and improved path separation; a joint-tree version supports conditional generation.

3.1 QUANTILE ALIGNTREE COUPLING

  • QAT transfers the data tree’s topology to the source by matching split mass ratios, giving corresponding data and source leaves equal mass.Each training pair draws a source point and target point from the conditionals associated with the same leaf.Source: §3.1 · pp. 4–5
  • With a standard Gaussian source, each axis-aligned leaf region supports parallel, coordinate-wise truncated-normal sampling.The region bounds come from split constraints along the root-to-leaf path.Source: Figure 1 · p. 5, §3.1 · p. 5

3.2 QAT FLOW MATCHING

  • Paths from distinct target leaves do not cross at any interpolation time, almost surely, which alleviates local velocity ambiguity.QAT-FM changes the sample pairing while retaining the linear path and velocity regression target.Source: §3.2 · pp. 5–6
  • Under non-degenerate splits, QAT strictly improves expected inter-path separation over independent coupling when the source is standard Gaussian.The gain is denoted ΔQAT and is induced by the QAT tree.Source: §3.2 · p. 6

3.3 CONDITIONAL QAT STRATEGY

  • Conditional QAT builds a tree over joint data-condition samples; condition-coordinate splits refine target groups without partitioning the source.It avoids explicitly solving product-space OT while retaining QAT’s efficient sampling form.Source: §3.3 · p. 6
  • Conditional QAT inherits non-crossing and path-separation results when distinct leaves first diverge at a data-coordinate split and the stated non-degeneracy conditions hold.The non-crossing guarantee is for paths from the corresponding leaf couplings.Source: §3.3 · p. 7

3.4 EFFICIENT QAT CONSTRUCTION AND SAMPLING

  • The adopted split rule selects the highest-variance coordinate and splits at its empirical mean; conditional terminal splitting prioritizes data coordinates.For discrete empirical data, recursion continues until each leaf contains at most one sample.Source: §3.4 · p. 7
  • For high-dimensional data, QAT operates on a rotated top-k subspace while preserving the global Gaussian source marginal.The rotation combines a Patch-Hadamard transform and PCA at the mean-pooling scale.Source: §3.4 · p. 7

4 EXPERIMENTS

Across synthetic, unconditional, and conditional generation benchmarks, QAT-FM shows competitive generation quality across solver regimes. The experiments also report that its gains are relatively insensitive to tree-construction hyperparameters.

4.1 2D SYNTHETIC EXPERIMENTS

  • On 2D checkerboard data, QAT-FM yields better local separation and more organized intermediate transport, recovering a clear checkerboard at small NFE.OT-FM and SD-FM capture the target's global structure earlier than I-FM.Source: §4.1 · p. 8, Figure 2 · p. 8

4.2 UNCONDITIONAL GENERATION ON CIFAR-10

  • On CIFAR-10, QAT-FM improves over I-FM and SD-FM with few-step Euler sampling and achieves the lowest FID with Dopri5.Figure 3 also shows strong performance during late-stage training.Source: §4.2 · p. 8, Table 2 · p. 8, Figure 3 · p. 8
  • Across the tested PCA dimensions, patch sizes, and tree-splitting strategies, QAT-FM is insensitive to these choices, with only slight FID and NFE variation.The authors attribute the gains primarily to global tree-structured quantile alignment rather than carefully tuned tree hyperparameters.Source: §4.2 · p. 8, Table 3 · p. 9

4.3 CONDITIONAL GENERATION EXPERIMENTS

  • On class-conditional ImageNet-32, QAT-FM remains competitive with few-step sampling and performs best under larger-step Euler solvers and Dopri5.C2OT performs best at small Euler step counts, but its advantage reverses with more accurate solvers.Source: §4.3 · p. 9, Table 4 · p. 9
  • On text-conditional CelebA-64, QAT-FM outperforms I-FM, OT-FM, and C2OT in FID across solvers while requiring the lowest average NFE.It matches I-FM and C2OT in CLIP score, exceeds OT-FM, and produces more coherent Euler-2 faces and sharper Euler-100 details.Source: Figure 4 · p. 9, §4.3 · p. 10
  • On ImageNet-256 latent generation, QAT-FM improves FID and Inception Score over I-FM across solver step counts while maintaining competitive sFID, precision, and recall.The reported improvements cover both few-step and high-accuracy sampling regimes.Source: §4.3 · p. 10, Table 5 · p. 10

5 CONCLUSION

The paper proposes QAT-FM, a tree-structured quantile-alignment coupling for Flow Matching. It reports favorable path geometry, lower coupling cost than OT-based methods, and results across synthetic, image, and latent-space generation.

5 CONCLUSION

  • QAT-FM is a Flow Matching coupling strategy based on tree-structured quantile alignment.The conclusion frames it as an efficient, structured coupling mechanism for improving Flow Matching training.Source: §5 · p. 10
  • QAT-FM avoids the high cost of OT-based coupling while preserving non-crossing paths and improved path separation.Source: §5 · p. 10
  • The method extends to conditional generation through condition-aware tree construction, and experiments cover synthetic data, image generation, and large-scale latent-space generation.Source: §5 · p. 10

A.1 PROOF OF THEOREM 1

The proof compares source-side regions aligned to the binary data tree with data leaves whose root-to-leaf paths first diverge at an internal node. It uses the split coordinate and aligned thresholds to establish that interpolations from distinct leaves do not coincide.

A.1 PROOF OF THEOREM 1

  • For any two distinct data leaves, the proof identifies their unique first divergence at an internal tree node.Source: §A.1 · p. 14
  • At that node, the data tree splits along a coordinate and threshold, while QAT reuses the coordinate with an aligned source-side threshold.Source: §A.1 · p. 14
  • Matching each data leaf to its source region yields strictly non-coincident linear interpolation paths for distinct leaves, almost surely.The proof uses positive coordinate differences to establish the non-coincidence result.Source: §A.1 · p. 14

A.2 PROOF OF THEOREM 2

The proof isolates the coupling-dependent cross term in expected pairwise path separation, then analyzes it recursively over the QAT tree. Under non-degenerate splits, the resulting separation increment is nonnegative and is positive if at least one effective split contributes positively.

A.2 PROOF OF THEOREM 2

  • The marginal-dependent terms in the interpolation expansion are identical for QAT and independent coupling, so their difference comes entirely from the cross term.Under independent coupling, the source and target differences are independent and zero-mean, making that cross term zero.Source: §A.2 · p. 14
  • At a leaf, the QAT coupling is a product coupling, so the cross term vanishes; internal nodes are analyzed by decomposing pairs across their children.The recursion defines a normalized coupling within each node and evaluates contributions from within-child and cross-child pairings.Source: §A.2 · pp. 14–15
  • Under non-degenerate splits, each effective node contributes Δη ≥ 0, and ΔQAT > 0 if at least one node contributes positively.The proof obtains the result by telescoping node contributions from the root, with zero contribution at leaves.Source: §A.2 · pp. 15–16

A.3 PROPERTIES OF CQAT

CQAT extends the tree coupling to joint data-condition samples while retaining Gaussian source alignment. Its properties depend on data-coordinate splits, and its construction and sampling have stated complexity bounds.

A.3 PROPERTIES OF CQAT

  • For two CQAT leaf paths whose first divergence is a data-coordinate split, the linear interpolation paths do not intersect.Condition-coordinate splits do not partition the source-side region and contribute zero; under non-degenerate splits, a positive data-coordinate contribution yields ΔCQAT > 0.Source: §A.3 · pp. 16–17

B.1 QAT COUPLING CONSTRUCTION AND SAMPLING

  • QAT construction costs O(Nd log N), and sampling one training pair costs O(d) under an approximately balanced tree.Sampling uses d independent truncated-Gaussian variates; epoch-level sampling costs O(Nd).Source: §B.1 · p. 17

B.2 CQAT COUPLING CONSTRUCTION AND SAMPLING

  • CQAT builds a tree over data and condition coordinates for O(N(d + m) log N) construction, while per-sample source generation remains O(d).When node size falls below ndata, candidate split coordinates are restricted to data dimensions; class labels can instead be handled by separate class-specific QAT trees.Source: §B.2 · pp. 18–19

B.3 HIGH-DIMENSIONAL IMAGE DATA PROCESSING

For image data, the method constructs QAT in a lower-dimensional representation while preserving the standard Gaussian source marginal. The implementation section also specifies that experiments vary the coupling strategy while sharing other training choices.

B.3 HIGH-DIMENSIONAL IMAGE DATA PROCESSING

  • A pure orthogonal rotation is used for image processing, preserving a standard Gaussian source without mean subtraction.Patch-wise Hadamard processing yields pooled patch representations, and PCA is estimated on those representations before retaining leading coordinates for tree construction.Source: §B.3 · p. 19
  • The low-dimensional QAT source is lifted to pixel space by filling residual PCA and high-frequency coordinates with independent standard Gaussian noise, yielding marginal N(0, Id).The stated per-sample cost is O(d + d′^2 + d log r), with batch sampling fully vectorizable.Source: §B.3 · pp. 19–20

C IMPLEMENTATION DETAILS

  • Across experiments, methods share the network architecture, training configuration, and sampling protocol, differing only in the coupling strategy used to form training pairs.The compared strategies are independent coupling, mini-batch OT coupling, and QAT or Conditional QAT coupling.Source: C IMPLEMENTATION DETAILS · p. 20

C.1 TWO-DIMENSIONAL EXPERIMENT

The experimental section specifies settings across synthetic, image, and conditional-generation tasks, along with supplementary runtime and sample visualizations. The supplied passages describe protocols and configurations, not comparative outcome values.

C.1 TWO-DIMENSIONAL EXPERIMENT

  • The two-dimensional experiment compares I-FM, OT-FM, SD-FM, and QAT-FM on transport from a Gaussian source to a checkerboard target under shared network and optimization settings.The setup uses 20K iterations with batch size 256; OT is solved within each minibatch, while SD-FM uses a regularized discrete reference measure.Source: §C.1 · p. 20

C.3 IMAGENET-32

  • ImageNet-1k class-conditional generation at 32 × 32 uses separate Conditional QAT couplings within each class subset.The setup includes horizontal-flip augmentations and evaluates generated samples against reference statistics from approximately 1.28 million training images.Source: §C.3 · p. 21

C.4 CELEBA-64

  • CelebA-64 text-conditional generation uses CelebA-Dialog descriptions and a joint Conditional QAT tree over low-dimensional image representations and text embeddings.Evaluation covers image fidelity by FID and text-image alignment by CLIP score.Source: §C.4 · pp. 21–22

C.5 IMAGENET-256

  • ImageNet-256 class-conditional generation trains LightningDiT in VA-VAE latent space using the recommended lightweight 64-epoch setting.Conditional QAT uses the leading 256 PCA coordinates; evaluation reports FID, IS, sFID, Precision, and Recall.Source: §C.5 · p. 22

D EXTENDED EXPERIMENTAL RESULTS

  • Supplementary experiments benchmark coupling construction and sampling wall-clock time and provide additional generated samples across four image datasets.The listed sample sets cover CIFAR-10, ImageNet-32, CelebA-64, and ImageNet-256.Source: D EXTENDED EXPERIMENTAL RESULTS · p. 22

D.1 RUNTIME COMPARISON

The section compares coupling construction and sampling costs over 100 training epochs and presents additional generated samples across datasets and sampling configurations.

D.1 RUNTIME COMPARISON

  • Full OT becomes rapidly infeasible as dataset size grows, while batch OT recomputes pairings each epoch and SD-FM's semi-discrete optimization is expensive.The comparison fixes feature dimension at d = 32 and reports cumulative times over 100 training epochs.Source: §D.1 · pp. 22–23
  • QAT achieves substantially lower construction time at all scales, with per-sample sampling cost close to independent random sampling.QAT builds offline; its sampling draws from truncated Gaussians using leaf bounds.Source: §D.1 · pp. 22–23, Figure 6 · p. 23

D.2 CIFAR-10

  • Figure 7 compares randomly generated unconditional CIFAR-10 samples from I-FM, OT-FM, SD-FM, and QAT-FM under Euler-2, Euler-100, and dopri5.Source: Figure 7 · p. 24

D.3 IMAGENET-32

  • Figure 8 shows class-conditional ImageNet-32 samples across four methods, three classes, and eight Euler or Dopri5 sampling configurations.The conditions are macaw, sports car, and pizza.Source: Figure 8 · p. 25

D.4 CELEBA-64

  • Figure 9 compares four methods and three sampling configurations for text-conditional CelebA-64 samples across CelebA-Dialog text conditions.The configurations are Euler-2, Euler-100, and Dopri5.Source: Figure 9 · p. 26

D.5 IMAGENET-256

  • Figure 10 presents class-conditional ImageNet-256 samples from I-FM and QAT-FM under the same random seed and class conditions.The samples are randomly generated without cherry-picking.Source: Figure 10 · p. 27

Related papers

Loading 2609.03443v1…