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

Junyi Lin Affiliation: Institute of Statistics and Big Data, Renmin University of China, Beijing, China Email: junyilin@ruc.edu.cn Mengyu Li Affiliation: Department of Statistics and Data Science, Tsinghua University, Beijing, China Email: 2025104242@ruc.edu.cn Jingxuan Hu Affiliation: Institute of Statistics and Big Data, Renmin University of China, Beijing, China Email: chengmeng@ruc.edu.cn Kejun He ^†^†thanks: Corresponding author Affiliation: Institute of Statistics and Big Data, Renmin University of China, Beijing, China Email: kejunhe@ruc.edu.cn Cheng Meng ^†^†thanks: Corresponding author Affiliation: Institute of Statistics and Big Data, Renmin University of China, Beijing, China Email: mengyuli@tsinghua.edu.cn

Abstract

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

In recent years, continuous-time flow-based generative models have become increasingly prominent in generative modeling (Chen et al., 2018; Grathwohl et al., 2019; Lipman et al., 2023; Liu et al., 2023). A representative class is continuous normalizing flows (CNFs), which define continuous transformations as the solution flows of ordinary differential equations (ODEs) whose time-dependent velocity fields are parameterized by neural networks. During sampling, new data are generated by numerically integrating the learned ODE dynamics. To improve training efficiency, Flow Matching (Lipman et al., 2023) directly supervises the velocity field with targets derived from pre-specified conditional probability paths, avoiding the repeated numerical ODE integration required by conventional CNF training, enabling scalable training of generative models. Flow Matching has achieved competitive performance across diverse domains, including image generation (Ma et al., 2024; Esser et al., 2024), speech synthesis (Le et al., 2023; Liu et al., 2024), and text generation (Hu et al., 2024; Hu et al., 2026).

Although Flow Matching improves the scalability of training continuous-time generative models, its performance remains sensitive to the conditional probability paths and source–target couplings constructed during training (Lipman et al., 2023; Tong et al., 2024; Pooladian et al., 2023). In common linear-path formulations, Flow Matching trains the velocity field by regressing target velocities along linear interpolation paths between noise and data samples; different coupling choices therefore directly alter the structure of the supervision signal. When training paths cross, or pass too closely through the same intermediate region with substantially different target velocities, velocity ambiguity arises locally; under mean-squared-error training, these inconsistent velocity directions are averaged, increasing the difficulty of fitting the velocity field (Guo and Schwing, 2025). By contrast, training paths with the non-crossing property or good path separation reduce this local ambiguity, yielding a more deterministic mapping from intermediate states to target velocities and thereby improving training stability. In conditional generation, coupling strategies must also preserve the alignment between samples and conditions, since couplings that ignore the conditioning information may produce geometrically favorable but conditionally inconsistent supervision signals (Chemseddine et al., 2025; Cheng and Schwing, 2025).

Several existing approaches seek to alleviate path crossing and velocity ambiguity by improving coupling quality. The reflow procedure in Rectified Flow (Liu et al., 2023) reconstructs pairings between noise and data using a learned model, yielding more deterministic and straighter generation trajectories; however, this requires additional sampling and re-training steps, imposing substantial computational overhead. OT-based methods (Villani and others, 2009) construct geometrically optimal pairings by minimizing transport cost; under a quadratic cost, the induced displacement interpolation exhibits a favorable non-crossing structure that alleviates velocity ambiguity. Nonetheless, exact empirical OT is computationally expensive and difficult to scale to large datasets; in conditional generation, maintaining conditioning consistency typically requires condition-aware constraints or condition-wise couplings, further increasing algorithmic complexity. Mini-batch OT (Pooladian et al., 2023; Tong et al., 2024) substantially reduces this cost by solving OT within each mini-batch, but the resulting coupling is primarily local and cannot strictly preserve the non-crossing structure at the full-dataset level. Most of the above methods construct couplings only over finite Gaussian samples; SD-FM (Mousavi-Hosseini et al., 2026) further improves the alignment of continuous Gaussian information via gradient-based optimization, but its iterative optimization incurs substantial sampling or optimization cost and does not yield an explicit low-cost global coupling construction.

Table 1: Comparison of Flow Matching coupling strategies.

Method	Build Cost¹	Sample Cost	Non-Cross	Path Sep.	Gaussian Align.
FM	—	$\mathcal{O}(Nd)$	✗	✗	✓✓
RecFlow	$\mathcal{O}(NE\Theta)$	$\mathcal{O}(N\Theta)$	✓	✗	✗
OT-FM ${}_{\text{(Full)}}$	$\mathcal{O}(N^{3}+N^{2}d)$	$\mathcal{O}(N^{2}+Nd)$	✓	✓	✗
OT-FM ${}_{\text{(Minibatch)}}$	$\mathcal{O}(NB+NBd)$	$\mathcal{O}(NB+Nd)$	✗	✗	✗
SD-FM	$\mathcal{O}(NKd)$	$\mathcal{O}(N^{2}d)$	✓	✓	✓
QAT-FM ${}_{\textbf{(Ours)}}$	$\mathcal{O}(Nd\log N)$	$\mathcal{O}(Nd)$	✓	✓✓²	✓✓

1

Notation. $N$ : dataset size; $B$ : mini-batch size; $d$ : data dimension; $K$ : SD optimization steps; $E$ : reflow rounds; $\Theta$ : one network forward pass. Complexities are reported up to constant factors.
2

✓✓indicates that the property is also preserved in conditional generation.

Table 1 compares representative Flow Matching coupling strategies along five axes: coupling construction cost (Build Cost), sampling cost (Sample Cost), whether the coupling induces non-crossing paths at the full-dataset level (Non-Cross), whether it improves path separation (Path Sep.), and whether the source marginal is aligned with a continuous Gaussian prior (Gaussian Align.). As shown, prior methods generally trade off computational efficiency, path geometry, and Gaussian-prior alignment. To address these limitations, we propose Quantile AlignTree Flow Matching (QAT-FM), a low-cost structured coupling strategy that achieves non-crossing paths, improved path separation, and continuous Gaussian-prior alignment, with several of these properties naturally extending to conditional generation.

Contributions. Our main contributions are as follows:

1)

We propose Quantile AlignTree (QAT), an efficient semi-discrete transport construction based on binary tree-structured quantile alignment. QAT induces an exact coupling between an empirical data distribution and a continuous Gaussian prior. With construction complexity $\mathcal{O}(Nd\log N)$ and per-pair sampling complexity $\mathcal{O}(d)$ , QAT scales to large-scale, high-dimensional generative modeling tasks.
2)

The QAT coupling supports multiple tree-splitting strategies and can naturally incorporate conditioning information into both tree construction and coupling.
3)

We prove that QAT induces favorable path structure: the resulting linear interpolation paths satisfy the non-crossing property, and path separation at intermediate times is uniformly greater than that of standard independent-coupling flow matching, alleviating local velocity ambiguity.
4)

We conduct systematic experiments on low-dimensional synthetic data, unconditional generation, discrete class-conditional generation, continuous text-conditional generation, and large-scale image synthesis. Results show that QAT-FM achieves competitive generation quality with substantially lower coupling construction complexity than OT-based methods.

2 Preliminaries

Notation.

Let $\mathcal{P}(\mathbb{R}^{d})$ denote the set of Borel probability measures on $\mathbb{R}^{d}$ . For $\alpha\in\mathcal{P}(\mathbb{R}^{d})$ and $\beta\in\mathcal{P}(\mathbb{R}^{d^{\prime}})$ , a joint distribution $\pi\in\mathcal{P}(\mathbb{R}^{d}\times\mathbb{R}^{d^{\prime}})$ is called a coupling of $\alpha$ and $\beta$ if its marginals are $\alpha$ and $\beta$ ; we write $\Pi(\alpha,\beta)$ for the collection of all such couplings. For a vector $\bm{x}\in\mathbb{R}^{d}$ , let $\bm{x}^{(j)}$ denote its $j$ -th coordinate. For $\rho\in\mathcal{P}(\mathbb{R}^{d})$ , let $\rho^{(j)}$ denote its $j$ -th marginal distribution, $F_{\rho}^{(j)}$ the cumulative distribution function (CDF) of $\rho^{(j)}$ , and $\bigl(F_{\rho}^{(j)}\bigr)^{-1}$ its generalized inverse. For a measurable map $T:\mathbb{R}^{d}\to\mathbb{R}^{k}$ , we write $T_{\#}\rho$ for the pushforward of $\rho$ under $T$ .

Flow-based Generative Modeling.

Given a time-dependent velocity field $v:[0,1]\times\mathbb{R}^{d}\to\mathbb{R}^{d}$ , write $v_{t}(x):=v(t,x)$ . The flow map $\phi_{t}:\mathbb{R}^{d}\to\mathbb{R}^{d}$ induced by $v_{t}$ is defined as the solution to the ODE

\frac{d}{dt}\phi_{t}(x)=v_{t}(\phi_{t}(x)),\qquad\phi_{0}(x)=x.

Setting $\rho_{t}=(\phi_{t})_{\#}\gamma$ , where $\gamma$ denotes the source distribution, the pair $(\rho_{t},v_{t})$ satisfies the continuity equation

\partial_{t}\rho_{t}+\nabla\cdot(\rho_{t}v_{t})=0,

(1)

with $\rho_{0}=\gamma$ . When $(\phi_{1})_{\#}\gamma=\mu$ , the flow transports the source $\gamma$ to the target distribution $\mu$ . The goal of flow-based generative modeling is therefore to learn a parameterized velocity field $v_{\theta}$ whose induced flow map satisfies $(\phi_{1})_{\#}\gamma=\mu$ .

Flow Matching (Lipman et al., 2023) provides a tractable training procedure for this velocity field. It specifies a probability path $(\rho_{t})_{t\in[0,1]}$ connecting $\gamma$ to $\mu$ together with a compatible target velocity field $u_{t}$ . The parameterized velocity field is then learned via the regression objective

\mathcal{L}_{\mathrm{FM}}(\theta)=\mathbb{E}_{\begin{subarray}{c}t\sim\mathrm{Unif}[0,1]\\ x_{t}\sim\rho_{t}\end{subarray}}\left\|v_{\theta}(t,x_{t})-u_{t}(x_{t})\right\|^{2}.

(2)

This formulation avoids backpropagating through numerical ODE solvers during training.

Conditional Flow Matching.

In practice, the marginal velocity field $u_{t}$ is often intractable to construct directly. Conditional Flow Matching¹¹ 1 Here, “conditional” refers to conditioning on paired endpoints. To avoid ambiguity with conditional generation, we refer to this formulation simply as Flow Matching (FM) throughout the remainder of the paper unless otherwise specified. (CFM) (Lipman et al., 2023; Tong et al., 2024) addresses this by conditioning on paired endpoints: given $\pi\in\Pi(\gamma,\mu)$ , draw $(x_{0},x_{1})\sim\pi$ , specify a conditional path $p_{t}(\cdot\mid x_{0},x_{1})$ and its compatible velocity $u_{t}(\cdot\mid x_{0},x_{1})$ , and train $v_{\theta}(t,x_{t})$ to match $u_{t}(x_{t}\mid x_{0},x_{1})$ , where $x_{t}\sim p_{t}(\cdot\mid x_{0},x_{1})$ . A canonical and widely used choice is the linear interpolation path

x_{t}=(1-t)x_{0}+tx_{1},\qquad u_{t}(x_{t}\mid x_{0},x_{1})=x_{1}-x_{0},

under which the training objective reduces to

\mathcal{L}_{\mathrm{CFM}}(\theta)=\mathbb{E}_{\begin{subarray}{c}t\sim\mathrm{Unif}[0,1]\\ (x_{0},x_{1})\sim\pi\end{subarray}}\left\|v_{\theta}\!\left(t,(1-t)x_{0}+tx_{1}\right)-(x_{1}-x_{0})\right\|^{2}.

(3)

This reveals that the endpoint pairing is governed by the coupling $\pi\in\Pi(\gamma,\mu)$ . The simplest choice is the independent coupling $\pi=\gamma\otimes\mu$ . A more principled alternative is the OT coupling

\pi_{\mathrm{OT}}\in\arg\min_{\pi\in\Pi(\gamma,\mu)}\int_{\mathbb{R}^{d}\times\mathbb{R}^{d}}\|x_{0}-x_{1}\|^{2}\,d\pi(x_{0},x_{1}),

(4)

which under the linear path tends to produce shorter, geometrically natural paired trajectories (Pooladian et al., 2023; Tong et al., 2024). However, computing the exact global OT coupling is expensive at large scale; in practice, mini-batch OT is used as an approximation, reducing computational cost at the expense of coupling exactness and cross-batch consistency. For the special case $\gamma=\mathcal{N}(0,I_{d})$ , Mousavi-Hosseini et al. (2026) proposes solving a semi-discrete OT problem in (4) to better align the continuous Gaussian source with the empirical target distribution.

Data Partitioning by Trees.

The preceding discussion shows that the choice of coupling $\pi$ directly governs the geometry of particle trajectories in FM, while exact global OT coupling is prohibitively expensive at large scale (Zhang et al., 2025; Mousavi-Hosseini et al., 2026). This motivates binary tree-based space partitioning, which provides a simple way to organize data hierarchically and will serve as the foundation of our coupling construction. Given a data space $\mathcal{X}\subseteq\mathbb{R}^{d}$ and a dataset $\mathcal{D}=\{x_{i}\}_{i=1}^{N}\subseteq\mathcal{X}$ , let $\widehat{\mu}_{\mathcal{D}}=N^{-1}\sum_{i=1}^{N}\delta_{x_{i}}$ denote its empirical distribution. A binary partition tree $T$ consists of an internal node set $\mathcal{I}$ and a leaf node set $\mathcal{L}$ . Each node $\eta$ is associated with a region $R_{\eta}\subseteq\mathcal{X}$ and data subset $\mathcal{D}_{\eta}=\mathcal{D}\cap R_{\eta}$ , where the root $r$ satisfies $R_{r}=\mathcal{X}$ .

For each internal node $\eta\in\mathcal{I}$ , let $\widehat{\mu}_{\eta}=|\mathcal{D}_{\eta}|^{-1}\sum_{x\in\mathcal{D}_{\eta}}\delta_{x}$ denote the conditional empirical measure on $\mathcal{D}_{\eta}$ . A split strategy selects a split coordinate $\kappa_{\eta}\in[d]$ and threshold $s_{\eta}\in\mathbb{R}$ based on $\widehat{\mu}_{\eta}$ , partitioning $R_{\eta}$ into

R_{\eta}^{-}=\{x\in R_{\eta}:x^{(\kappa_{\eta})}\leq s_{\eta}\},\qquad R_{\eta}^{+}=\{x\in R_{\eta}:x^{(\kappa_{\eta})}>s_{\eta}\},

and correspondingly partitioning $\mathcal{D}_{\eta}$ into $\mathcal{D}_{\eta}^{-}=\mathcal{D}_{\eta}\cap R_{\eta}^{-}$ and $\mathcal{D}_{\eta}^{+}=\mathcal{D}_{\eta}\cap R_{\eta}^{+}$ . After recursive splitting, the leaf regions form a partition of $\mathcal{X}$ , denoted $\mathcal{R}(T):=\{R_{\ell}:\ell\in\mathcal{L}\}$ .

3 Quantile AlignTree Flow Matching

3.1 Quantile AlignTree Coupling

The preceding analysis shows that in the linear FM objective, the coupling $\pi$ governs how source samples are paired with target samples, directly shaping the geometry of flow paths. The key idea of QAT is to build a binary partition tree over the empirical data distribution and transfer its topology to the source distribution by matching the split mass ratios at each internal node. Figure 1 summarizes the resulting coupling and sampling procedure.

Let $\widehat{\mu}_{\mathcal{D}}$ denote the empirical measure of dataset $\mathcal{D}=\{x_{i}\}_{i=1}^{N}\subseteq\mathbb{R}^{d}$ , and let $\gamma\in\mathcal{P}(\mathbb{R}^{d})$ be the source distribution, e.g., $\gamma=\mathcal{N}(0,I_{d})$ . Following the binary partition tree construction of Section 2, we build a data tree $T(\widehat{\mu}_{\mathcal{D}})$ on $\widehat{\mu}_{\mathcal{D}}$ . For each internal node $\eta\in\mathcal{I}$ , let $(\kappa_{\eta},s_{\eta})$ denote its split coordinate and threshold, and let $\widehat{\mu}_{\eta}$ be the conditional empirical measure at that node. Based on the data tree, we define the Quantile AlignTree as follows.

Definition 1 (Quantile AlignTree (QAT))

Given the data tree $T(\widehat{\mu}_{\mathcal{D}})$ and source distribution $\gamma$ , the Quantile AlignTree of $\gamma$ with respect to the data tree, denoted $\widetilde{T}_{\gamma}(T)$ , is constructed recursively with the same topology as $T(\widehat{\mu}_{\mathcal{D}})$ . At each internal node $\eta$ , the split coordinate $\kappa_{\eta}$ is inherited from the data tree, but the data-side threshold $s_{\eta}$ is replaced by the source-side aligned threshold

\widetilde{s}_{\eta}=\left(F_{\gamma_{\eta}}^{(\kappa_{\eta})}\right)^{-1}\!\left(F_{\widehat{\mu}_{\eta}}^{(\kappa_{\eta})}(s_{\eta})\right),

(5)

where $\gamma_{\eta}$ denotes the conditional distribution of $\gamma$ on the current source-side region $\widetilde{R}_{\eta}$ . The source-side region is then partitioned as

\widetilde{R}_{\eta}^{-}=\{z\in\widetilde{R}_{\eta}:z^{(\kappa_{\eta})}\leq\widetilde{s}_{\eta}\},\qquad\widetilde{R}_{\eta}^{+}=\{z\in\widetilde{R}_{\eta}:z^{(\kappa_{\eta})}>\widetilde{s}_{\eta}\},

and the left and right child nodes are constructed recursively.

The quantile-based aligned threshold (5) guarantees that the source side and the data side carry equal mass ratios at every split. Specifically, when the source-side conditional marginal is continuous,

\gamma_{\eta}(\widetilde{R}_{\eta}^{-})=\widehat{\mu}_{\eta}(R_{\eta}^{-}),\qquad\gamma_{\eta}(\widetilde{R}_{\eta}^{+})=\widehat{\mu}_{\eta}(R_{\eta}^{+}).

Consequently, the data tree and the aligned source tree carry equal mass at corresponding leaf nodes, naturally inducing a coupling from $\gamma$ to $\widehat{\mu}_{\mathcal{D}}$ .

Refer to caption — Figure 1: Overview of the QAT coupling and sampling procedure. (*Left*) QAT builds a binary partition tree over the data using feature-based or condition-aware splits, and transfers the same tree topology to the Gaussian source by quantile alignment, yielding the structured coupling $\pi_{\mathrm{QAT}}$ . (*Right*) each data leaf is paired with an aligned Gaussian source region. Since these source regions are axis-aligned hyperrectangles, sampling from the corresponding QAT source conditional reduces to parallel per-coordinate truncated normal sampling.

Definition 2 (QAT Coupling)

Let $\mathcal{L}$ be the nonempty leaf node set of data tree $T(\widehat{\mu}_{\mathcal{D}})$ . For each leaf $\ell\in\mathcal{L}$ , let $\mathcal{D}_{\ell}=\mathcal{D}\cap R_{\ell}$ and $\alpha_{\ell}:=\tfrac{1}{N}|\mathcal{D}_{\ell}|$ . Let $\widehat{\mu}_{\ell}=|\mathcal{D}_{\ell}|^{-1}\sum_{x_{i}\in\mathcal{D}_{\ell}}\delta_{x_{i}}$ denote the conditional empirical measure on $\mathcal{D}_{\ell}$ , and let $\gamma_{\ell}$ denote the conditional source measure associated with leaf $\ell$ in the Quantile AlignTree. Define the QAT coupling as

\pi_{\mathrm{QAT}}:=\sum_{\ell\in\mathcal{L}}\alpha_{\ell}\,\gamma_{\ell}\otimes\widehat{\mu}_{\ell}.

(6)

Then $\pi_{\mathrm{QAT}}\in\Pi(\gamma,\widehat{\mu}_{\mathcal{D}})$ .

Intuitively, the QAT coupling matches each data leaf $\mathcal{D}_{\ell}$ with an aligned source region. To draw a training pair $(x_{0},x_{1})\sim\pi_{\mathrm{QAT}}$ , one samples a leaf $\ell$ with probability $\alpha_{\ell}$ , draws $x_{0}\sim\gamma_{\ell}$ from the corresponding conditional source distribution, and draws $x_{1}\sim\widehat{\mu}_{\ell}$ from the conditional empirical distribution on that leaf.

Remark 1 (Gaussian Sampling)

When $\gamma=\mathcal{N}(0,I_{d})$ , sampling from $\gamma_{\ell}$ is especially convenient. Since all splits are axis-aligned, every leaf source region is an axis-aligned hyperrectangle. If leaf $\ell$ corresponds to source region $\widetilde{R}_{\ell}=\prod_{j=1}^{d}[a_{j}^{(\ell)},b_{j}^{(\ell)}]$ , sampling from $\gamma_{\ell}$ reduces to drawing each coordinate independently from a truncated standard normal:

x_{0}^{(j)}\sim\mathcal{N}(0,1)\big|_{[a_{j}^{(\ell)},b_{j}^{(\ell)}]},\qquad j=1,\ldots,d.

The interval endpoints are determined by all split constraints along the root-to-leaf path, and the sampling is fully parallelizable across coordinates.

3.2 QAT Flow Matching

Substituting the QAT coupling into the linear FM objective (3) yields the QAT-FM training objective. Let $\widehat{\mu}_{\mathcal{D}}$ be the empirical data distribution, $\gamma$ the source distribution, and $\pi_{\mathrm{QAT}}\in\Pi(\gamma,\widehat{\mu}_{\mathcal{D}})$ the QAT coupling of Definition 2. QAT-FM learns the velocity field via

\mathcal{L}_{\mathrm{QAT\text{-}FM}}(\theta):=\mathbb{E}_{\begin{subarray}{c}t\sim\mathrm{Unif}[0,1]\\ (x_{0},x_{1})\sim\pi_{\mathrm{QAT}}\end{subarray}}\left\|v_{\theta}\bigl(t,(1-t)x_{0}+tx_{1}\bigr)-(x_{1}-x_{0})\right\|^{2}.

(7)

Compared to standard FM, QAT-FM differs only in the sample pairing: $\pi_{\mathrm{QAT}}$ replaces the independent coupling, while the linear interpolation path and velocity regression target remain unchanged. This structured pairing yields the following geometric properties.

Theorem 1 (Non-Crossing Property)

Let $\pi_{\mathrm{QAT}}\in\Pi(\gamma,\widehat{\mu}_{\mathcal{D}})$ be the QAT coupling. Consider two paired samples $(x_{0},x_{1}),(x_{0}^{\prime},x_{1}^{\prime})\in\mathrm{supp}(\pi_{\mathrm{QAT}})$ . If $x_{1}$ and $x_{1}^{\prime}$ belong to two distinct leaf nodes $\ell\neq\ell^{\prime}$ of the data tree, then the corresponding linear interpolation paths

x_{t}=(1-t)x_{0}+tx_{1},\qquad x_{t}^{\prime}=(1-t)x_{0}^{\prime}+tx_{1}^{\prime}

do not intersect, i.e., $x_{t}\neq x_{t}^{\prime}$ for all $t\in[0,1]$ , almost surely.

Theorem 1 establishes that paths associated with distinct target leaves are strictly non-crossing. This property alleviates the velocity ambiguity that can arise from overlapping training paths: since the FM objective fits $v_{\theta}$ pointwise at each $t\in[0,1]$ , two training paths satisfying $x_{t}=x_{t}^{\prime}$ but $x_{1}-x_{0}\neq x_{1}^{\prime}-x_{0}^{\prime}$ impose conflicting supervision signals at the same spatial location, increasing the difficulty of fitting $v_{\theta}$ . Paths that remain well-separated at intermediate times are therefore desirable. Beyond preventing cross-leaf path crossings, the QAT coupling also consistently improves the expected inter-path distance, as the following theorem quantifies.

Theorem 2 (Improved Path Separation)

Let $\gamma=\mathcal{N}(0,I_{d})$ , let $\pi_{\mathrm{QAT}}\in\Pi(\gamma,\widehat{\mu}_{\mathcal{D}})$ be the QAT coupling, and let $\pi_{\mathrm{ind}}:=\gamma\otimes\widehat{\mu}_{\mathcal{D}}$ be the independent coupling. For any coupling $\pi$ , let $\mathbb{E}_{\pi^{\otimes 2}}$ denote expectation over two independent pairs $(x_{0},x_{1}),(x_{0}^{\prime},x_{1}^{\prime})\sim\pi$ , and define $x_{t}=(1-t)x_{0}+tx_{1}$ and $x_{t}^{\prime}=(1-t)x_{0}^{\prime}+tx_{1}^{\prime}$ . Then for every $t\in(0,1)$ ,

\mathbb{E}_{\pi_{\mathrm{QAT}}^{\otimes 2}}\!\left[\|x_{t}-x_{t}^{\prime}\|_{2}^{2}\right]-\mathbb{E}_{\pi_{\mathrm{ind}}^{\otimes 2}}\!\left[\|x_{t}-x_{t}^{\prime}\|_{2}^{2}\right]=4t(1-t)\,\Delta_{\mathrm{QAT}},

where $\Delta_{\mathrm{QAT}}$ is the node-level path-separation gain induced by the QAT tree:

\Delta_{\mathrm{QAT}}:=\sum_{\eta\in\mathcal{I}}\alpha_{\eta}q_{\eta}(1-q_{\eta})\Bigl(\mathbb{E}_{\gamma_{\eta^{+}}}\!\bigl[X^{(\kappa_{\eta})}\bigr]-\mathbb{E}_{\gamma_{\eta^{-}}}\!\bigl[X^{(\kappa_{\eta})}\bigr]\Bigr)\Bigl(\mathbb{E}_{\widehat{\mu}_{\eta^{+}}}\!\bigl[X^{(\kappa_{\eta})}\bigr]-\mathbb{E}_{\widehat{\mu}_{\eta^{-}}}\!\bigl[X^{(\kappa_{\eta})}\bigr]\Bigr).

(8)

Here $\alpha_{\eta}:=\widehat{\mu}_{\mathcal{D}}(R_{\eta})$ is the data mass at node $\eta$ , $q_{\eta}:=F_{\widehat{\mu}_{\eta}}^{(\kappa_{\eta})}(s_{\eta})$ is the conditional left-child mass fraction at split coordinate $\kappa_{\eta}$ , and $\gamma_{\eta^{\pm}}$ , $\widehat{\mu}_{\eta^{\pm}}$ denote the conditional source and empirical distributions on the left and right child regions of $\eta$ . Under non-degenerate splits, $\Delta_{\mathrm{QAT}}>0$ , so the QAT coupling strictly dominates the independent coupling in expected inter-path separation.

3.3 Conditional QAT Strategy

In conditional generation tasks, data samples carry conditioning variables. Let $\{(x_{i},c_{i})\}_{i=1}^{N}\subseteq\mathbb{R}^{d}\times\mathbb{R}^{m}$ be the joint dataset, $\widehat{\mu}_{\mathcal{D}_{C}}:=N^{-1}\sum_{i=1}^{N}\delta_{(x_{i},c_{i})}$ the joint empirical distribution, and $\widehat{\mu}_{C}:=N^{-1}\sum_{i=1}^{N}\delta_{c_{i}}$ the marginal over condition space. A standard OT coupling based only on data-space distances can ignore the alignment between samples and their conditioning variables. Alternatively, solving OT directly between the joint source $\gamma\otimes\widehat{\mu}_{C}$ and the joint empirical distribution $\widehat{\mu}_{\mathcal{D}_{C}}$ increases the optimization dimension and computational cost (Cheng and Schwing, 2025; Mousavi-Hosseini et al., 2026). We therefore propose the following Conditional QAT strategy.

Definition 3 (Conditional QAT)

Let $\{c_{i}\}_{i=1}^{N}\subseteq\mathbb{R}^{m}$ be conditioning variables. Conditional QAT constructs a data tree on the joint samples $\{(x_{i},c_{i})\}_{i=1}^{N}\subseteq\mathbb{R}^{d+m}$ . For each internal node $\eta$ , let $\kappa_{\eta}\in[d+m]$ denote its split coordinate. When building the source-side AlignTree, if $\kappa_{\eta}\leq d$ (a data coordinate), the source side is partitioned using the standard QAT aligned threshold in (5); if $\kappa_{\eta}>d$ (a condition coordinate), no split is applied to the source and both child nodes inherit the parent source region, the split is used only to refine the target-side grouping: $\widetilde{R}_{\eta^{-}}=\widetilde{R}_{\eta^{+}}=\widetilde{R}_{\eta}$ .

The Conditional QAT coupling takes the same form as Definition 2, with the target empirical distribution replaced by the joint distribution $\widehat{\mu}_{\mathcal{D}_{C}}$ :

\pi_{\mathrm{CQAT}}=\sum_{\ell\in\mathcal{L}_{C}}\alpha_{\ell}\,\gamma_{\ell}\otimes\widehat{\mu}_{\mathcal{D}_{C},\ell}\in\Pi(\gamma,\widehat{\mu}_{\mathcal{D}_{C}}),

where $\mathcal{L}_{C}$ is the set of nonempty leaves of the joint data tree, and $\widehat{\mu}_{\mathcal{D}_{C},\ell}$ is the conditional joint empirical distribution at leaf $\ell$ . A sample from $\pi_{\mathrm{CQAT}}$ can be written as $(x_{0},(x_{1},c))$ , where the linear FM path is applied only in the data space between $x_{0}$ and $x_{1}$ , while $c$ is provided as conditioning input to the velocity model. By encoding conditional structure through leaf-level matching in the joint tree, this construction avoids explicitly solving the product-space OT problem between $\gamma\otimes\widehat{\mu}_{C}$ and $\widehat{\mu}_{\mathcal{D}_{C}}$ , while retaining the efficient sampling form of unconditional QAT.

Remark 2 (Theoretical extensions)

Conditional QAT inherits the theoretical properties of QAT under appropriate conditions. Specifically, if the root-to- $\ell$ and root-to- $\ell^{\prime}$ paths first diverge at a data-coordinate split node (i.e., $\kappa_{\eta}\leq d$ ), then the linear interpolation paths drawn from the corresponding leaf couplings satisfy the non-crossing property of Theorem 1. Under the same non-degeneracy conditions, Conditional QAT also inherits the path-separation bound of Theorem 2. Proofs and further details are given in Appendix A.3.

3.4 Efficient QAT Construction and Sampling

When $\gamma=\mathcal{N}(0,I_{d})$ , each source-side leaf region $\widetilde{R}_{\ell}$ is fully characterized by its coordinate-wise bounds $(a_{\ell},b_{\ell})$ (Remark 1). Regardless of whether the data tree is built from unconditional or conditional samples, each leaf $\ell$ stores only the data index set $\mathcal{S}_{\ell}$ and the source-side bounds $(a_{\ell},b_{\ell})$ , which fully specify the leaf-wise product coupling $\gamma_{\ell}\otimes\widehat{\mu}_{\ell}$ . Training pairs are drawn by sampling a leaf $\ell$ proportionally to its mass, then independently drawing a source point from the corresponding truncated Gaussian region and a target point from $\mathcal{S}_{\ell}$ . The full construction, sampling algorithm, and complexity analysis are given in Appendix B.

Split Strategy.

QAT is compatible with a variety of split strategies. Throughout this work we adopt a maximum-variance rule: for each node $\eta$ , the split coordinate is the dimension of highest empirical variance and the split threshold is the empirical mean along that coordinate,

\kappa_{\eta}=\arg\max_{j}\operatorname{Var}_{\widehat{\mu}_{\eta}}\!\big(Y^{(j)}\big),\qquad s_{\eta}=\mathbb{E}_{\widehat{\mu}_{\eta}}\!\big[Y^{(\kappa_{\eta})}\big],

where $Y=X$ in the unconditional case and $Y=(X,C)$ in the conditional case. For discrete empirical data, we recurse until each leaf contains at most one sample; in the conditional setting, the terminal splitting phase prioritizes data coordinates over condition coordinates, maximizing the number of source-side splits and thereby increasing the path-separation gain of Theorem 2.

High-Dimensional Data Handling.

For high-dimensional data, axis-aligned partitioning in the raw data space may fail to capture the principal directions of variation. We therefore first rotate the data via an orthogonal transform $P\in\mathbb{R}^{d\times d}$ , setting $z_{i}=Px_{i}$ and constructing the QAT coupling on the top- $k$ subspace $z_{i}^{(1:k)}$ . By the rotational invariance of the standard Gaussian, $P_{\#}\mathcal{N}(0,I_{d})=\mathcal{N}(0,I_{d})$ for any orthogonal $P$ . Training then proceeds by sampling $(\epsilon^{(1:k)},z_{i}^{(1:k)})\sim\pi_{\mathrm{QAT}}$ , independently drawing $\epsilon^{(-k)}\sim\mathcal{N}(0,I_{d-k})$ , and setting

x_{0}=P^{\top}\!\big(\epsilon^{(1:k)},\epsilon^{(-k)}\big).

The marginal distribution of $x_{0}$ remains $\mathcal{N}(0,I_{d})$ , while the QAT tree operates entirely within the low-dimensional subspace. In practice, $P$ is composed of a Patch-Hadamard transform and a PCA rotation at the mean-pooling scale, concentrating the principal data variation into a small number of coordinates. Details are given in Appendix B.3.

4 Experiments

We evaluate QAT-FM through a series of experiments: we begin with path visualizations on two-dimensional synthetic data, then proceed to high-dimensional image generation benchmarks covering unconditional generation (CIFAR-10), class-conditional generation (ImageNet-32), text-conditional generation (CelebA-64), and large-scale latent-space generation (ImageNet-256). Additional implementation details are provided in Appendix C; extended experimental results appear in Appendix D.

4.1 2D Synthetic Experiments

Following Lipman et al. (2023); Pooladian et al. (2023), we compare different coupling methods on two-dimensional checkerboard data. Four methods are evaluated: I-FM, OT-FM (Tong et al., 2024), SD-FM (Mousavi-Hosseini et al., 2026), and the proposed QAT-FM. All methods share the same network architecture and training configuration.

Figure 2 (left) shows sample evolution from $t=0$ to $t=1$ . OT-FM and SD-FM capture the global structure of the target distribution earlier than I-FM, while QAT-FM yields better local separation and more organized intermediate transport. Figure 2 (right) shows generated samples with varying NFE. QAT-FM recovers a clear checkerboard structure at small NFE, indicating that its structured coupling improves few-step sample quality.

4.2 Unconditional Generation on CIFAR-10

We evaluate QAT-FM on unconditional generation using CIFAR-10 (Krizhevsky et al., 2009), which contains $50{,}000$ RGB training images of size $32\times 32$ across $10$ categories. The velocity-field network follows the architecture of Tong et al. (2024); training details are provided in Appendix C. We generate $50{,}000$ samples using several numerical solvers and compute FID (Heusel et al., 2017). Results are averaged over $5$ random seeds and reported with standard deviations.

Method	Euler-2	Euler-100	Dopri5	NFE $\downarrow$
I-FM	164.18 ${}_{\pm\text{0.13}}$	3.79 ${}_{\pm\textbf{0.04}}$	3.38 ${}_{\pm\text{0.03}}$	137.30 ${}_{\pm\text{0.30}}$
OT-FM	77.26 ${}_{\pm\textbf{0.12}}$	3.96 ${}_{\pm\text{0.04}}$	3.49 ${}_{\pm\text{0.03}}$	129.52 ${}_{\pm\textbf{0.86}}$
SD-FM	162.66 ${}_{\pm\text{0.23}}$	4.03 ${}_{\pm\text{0.02}}$	3.61 ${}_{\pm\text{0.01}}$	135.78 ${}_{\pm\text{0.36}}$
QAT-FM	98.68 ${}_{\pm\text{0.17}}$	3.91 ${}_{\pm\text{0.04}}$	3.26 ${}_{\pm\textbf{0.02}}$	145.55 ${}_{\pm\text{1.05}}$

Table 2: FID results on unconditional CIFAR-10 generation at

100

k training steps.

Figure 3: FID (Dopri5) vs. steps.

Table 2 reports CIFAR-10 results at $100$ k training steps. Under few-step Euler sampling, OT-FM performs best, while QAT-FM substantially improves over I-FM and SD-FM, indicating that structured coupling can improve generation quality under coarse numerical integration. Under the adaptive Dopri5 solver, QAT-FM achieves the lowest FID, showing competitive final generation quality in the unconditional setting. Figure 3 further shows that QAT-FM maintains strong performance during late-stage training.

Ablation Studies.

Our main experiments use patch size $p=4$ , PCA dimension $k=32$ , and mean-based tree splitting. To assess the sensitivity of these choices, we conduct three ablation studies: (1) varying the PCA dimension $k\in\{32,64,128\}$ ; (2) varying the patch size $p\in\{1,2,4\}$ ; and (3) comparing three split strategies—mean split, median split, and random split with ratio $\sim\mathcal{U}(0.25,0.75)$ . Results are summarized in Table 3.

Table 3: Ablation studies on PCA dimension, patch size, and tree splitting strategy.

PCA Dim.	FID $\downarrow$	NFE $\downarrow$
$k=32$	3.26	145.48
$k=64$	3.27	143.50
$k=128$	3.24	143.76

Patch Size	FID $\downarrow$	NFE $\downarrow$
$p=1$	3.24	145.36
$p=2$	3.24	145.55
$p=4$	3.26	145.48

Split Rule	FID $\downarrow$	NFE $\downarrow$
Mean	3.26	145.48
Median	3.22	143.89
Random	3.24	144.87

As shown in Table 3, QAT-FM is insensitive to the choice of PCA dimension, patch size, and tree splitting strategy; FID and NFE vary only slightly across all configurations. These results suggest that the performance gains of QAT-FM do not rely on carefully tuned tree construction hyperparameters, but stem primarily from the global tree-structured quantile alignment mechanism.

4.3 Conditional Generation Experiments

Class-conditional generation.

We evaluate QAT-FM on class-conditional generation using ImageNet-1k (Deng et al., 2009), which contains $1{,}000$ categories and over $1$ M training images. All images are resized to $32\times 32$ , with random horizontal flipping used for augmentation. We build the Conditional QAT coupling using both original and horizontally flipped training images. At evaluation, we generate $50$ K samples and compute FID against the training set. Results are averaged over $5$ random seeds and reported with standard deviations.

Table 4: FID results on class-conditional ImageNet-

32\times 32

generation.

Method	Euler-2	Euler-5	Euler-10	Euler-20	Euler-50	Euler-100	Dopri5	NFE $\downarrow$
I-FM	122.40 ${}_{\pm\text{0.08}}$	23.04 ${}_{\pm\text{0.07}}$	10.12 ${}_{\pm\text{0.07}}$	6.48 ${}_{\pm\text{0.07}}$	4.77 ${}_{\pm\text{0.07}}$	4.25 ${}_{\pm\text{0.06}}$	3.82 ${}_{\pm\text{0.04}}$	140.27 ${}_{\pm\textbf{1.99}}$
OT-FM	115.63 ${}_{\pm\text{0.42}}$	31.82 ${}_{\pm\text{0.11}}$	15.79 ${}_{\pm\text{0.02}}$	10.59 ${}_{\pm\text{0.00}}$	8.14 ${}_{\pm\text{0.02}}$	7.37 ${}_{\pm\text{0.00}}$	6.65 ${}_{\pm\text{0.02}}$	142.51 ${}_{\pm\text{0.45}}$
C2OT	100.41 ${}_{\pm\textbf{0.05}}$	20.94 ${}_{\pm\textbf{0.08}}$	10.10 ${}_{\pm\textbf{0.06}}$	6.80 ${}_{\pm\text{0.04}}$	5.15 ${}_{\pm\text{0.04}}$	4.62 ${}_{\pm\text{0.04}}$	4.16 ${}_{\pm\text{0.02}}$	143.97 ${}_{\pm\text{0.75}}$
QAT-FM	105.16 ${}_{\pm\text{0.30}}$	22.88 ${}_{\pm\textbf{0.05}}$	10.30 ${}_{\pm\text{0.08}}$	6.40 ${}_{\pm\textbf{0.07}}$	4.54 ${}_{\pm\textbf{0.07}}$	3.99 ${}_{\pm\textbf{0.07}}$	3.56 ${}_{\pm\textbf{0.06}}$	143.35 ${}_{\pm\text{2.22}}$

Table 4 reports FID on class-conditional ImageNet- $32\times 32$ generation. Both C2OT and QAT-FM improve few-step generation over I-FM, demonstrating the effectiveness of condition-aware coupling. While C2OT performs best under small Euler step counts, its advantage reverses as the solver becomes more accurate. By contrast, QAT-FM remains competitive in few-step sampling and achieves the best performance under larger-step Euler solvers and Dopri5. This indicates that Conditional QAT provides more stable supervision for high-quality conditional FM training.

Text-conditional generation.

We further evaluate QAT-FM on text-conditional CelebA- $64\times 64$ generation. CelebA (Liu et al., 2015) contains over $200$ K celebrity face images with attribute annotations, and CelebA-Dialog (Jiang et al., 2021) provides fine-grained natural-language descriptions for these images. We use the text captions as conditioning inputs. Following C2OT (Cheng and Schwing, 2025), each caption is encoded by a CLIP (Radford et al., 2021)-like DFN text encoder (Fang et al., 2024), and the resulting embedding is used as the condition variable. We compare QAT-FM with I-FM, OT-FM, and C2OT under the same text-conditioning setting. At evaluation, we generate $50$ K samples and report FID for image fidelity and SigLIP-2 (Tschannen et al., 2025) CLIP score for text-image alignment.

	Method	Euler-2	Euler-100	Dopri5	NFE $\downarrow$
FID $\downarrow$	I-FM	82.40	3.60	3.13	149.97
OT-FM	70.24	3.42	2.98	146.53
C2OT	81.39	3.50	3.08	148.75
QAT-FM	61.84	3.17	2.68	144.08
CLIP $\uparrow$	I-FM	0.11	0.10	0.10	149.97
OT-FM	0.09	0.08	0.08	146.53
C2OT	0.11	0.10	0.10	148.75
QAT-FM	0.11	0.10	0.10	144.08

Figure 4 shows that QAT-FM consistently outperforms I-FM, OT-FM, and C2OT in FID across different solvers, while also requiring the lowest average NFE. Meanwhile, QAT-FM attains the same CLIP scores as I-FM and C2OT, and higher CLIP scores than OT-FM, indicating that the improvement in image fidelity does not come at the cost of weaker text-image alignment. The qualitative samples further show that QAT-FM produces more coherent faces under Euler-2 and sharper details under Euler-100, consistent with the quantitative results.

Table 5: Class-conditional generation results on ImageNet-

256\times 256

latent generation.

Method	Steps	IS $\uparrow$	FID $\downarrow$	sFID $\downarrow$	Pre. $\uparrow$	Rec. $\uparrow$
I-FM	10	152.19	11.83	20.73	0.64	0.49
QAT-FM	10	162.15	10.67	20.74	0.66	0.48
I-FM	50	242.65	2.57	5.62	0.80	0.57
QAT-FM	50	250.36	2.46	5.69	0.80	0.58
I-FM	250	248.94	2.08	4.29	0.81	0.59
QAT-FM	250	256.96	1.99	4.25	0.80	0.60

Large-scale latent generation.

To assess scalability, we further evaluate QAT-FM on class-conditional ImageNet- $256\times 256$ latent generation (Rombach et al., 2022; Ma et al., 2024). Following Yao et al. (2025), images are encoded into VA-VAE latents, and the flow model is trained in the latent space. We construct the Conditional QAT directly over the VAE latents, while following the network architecture, training hyperparameters, and evaluation protocol of Yao et al. (2025). Due to computational constraints, we train for $64$ epochs using the official recommended setting. The implementation and evaluation details are provided in Appendix C.

Table 5 shows that QAT-FM consistently improves FID and Inception Score over I-FM across all solver step counts, while maintaining competitive sFID, precision, and recall. The improvements hold in both few-step and high-accuracy sampling regimes, indicating that the QAT coupling benefits latent-space flow training beyond a particular solver choice. Figure 5 provides qualitative examples under Euler-250 sampling. These results suggest that QAT-FM can be integrated into modern latent-space flow backbones with minimal architectural changes while improving generation quality.

5 Conclusion

We study the prior-data coupling problem in Flow Matching and propose QAT-FM, a simple yet effective coupling strategy based on tree-structured quantile alignment. QAT-FM avoids the high cost of OT-based coupling while preserving favorable path geometry, including non-crossing paths and improved path separation. It also extends naturally to conditional generation through condition-aware tree construction. Extensive experiments verify the effectiveness and scalability of our method across synthetic data, image generation, and large-scale latent-space generation. Overall, QAT-FM provides an efficient and structured coupling mechanism for improving Flow Matching training.

References

Chemseddine et al. (2025) J. Chemseddine, P. Hagemann, G. Steidl, and C. Wald Conditional wasserstein distances with applications in bayesian ot flow matching. Journal of Machine Learning Research 26 (141), pp. 1–47. External Links: Link Cited by: §1.
Chen et al. (2018) R. T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud Neural ordinary differential equations. In Advances in Neural Information Processing Systems, Vol. 31, pp. . External Links: Link Cited by: §1.
Cheng and Schwing (2025) H. K. Cheng and A. Schwing The curse of conditions: analyzing and improving optimal transport for conditional flow-based generation. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 15875–15884. Cited by: §B.2, §C.3, §C.4, §C.4, §1, §3.3, §4.3.
Deng et al. (2009) J. Deng, W. Dong, R. Socher, L. Li, K. Li, and L. Fei-Fei ImageNet: a large-scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, Vol. , pp. 248–255. External Links: Document Cited by: §C.3, §C.5, §4.3.
Dhariwal and Nichol (2021) P. Dhariwal and A. Nichol Diffusion models beat gans on image synthesis. In Advances in Neural Information Processing Systems, Vol. 34, pp. 8780–8794. External Links: Link Cited by: §C.3, §C.5.
Esser et al. (2024) P. Esser, S. Kulal, A. Blattmann, R. Entezari, J. Müller, H. Saini, Y. Levi, D. Lorenz, A. Sauer, F. Boesel, D. Podell, T. Dockhorn, Z. English, and R. Rombach Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first International Conference on Machine Learning, External Links: Link Cited by: §1.
Fang et al. (2024) A. Fang, A. M. Jose, A. Jain, L. Schmidt, A. T. Toshev, and V. Shankar Data filtering networks. In The Twelfth International Conference on Learning Representations, External Links: Link Cited by: §C.4, §4.3.
Flamary et al. (2021) R. Flamary, N. Courty, A. Gramfort, M. Z. Alaya, A. Boisbunon, S. Chambon, L. Chapel, A. Corenflos, K. Fatras, N. Fournier, L. Gautheron, N. T.H. Gayraud, H. Janati, A. Rakotomamonjy, I. Redko, A. Rolet, A. Schutz, V. Seguy, D. J. Sutherland, R. Tavenard, A. Tong, and T. Vayer POT: python optimal transport. Journal of Machine Learning Research 22 (78), pp. 1–8. External Links: Link Cited by: §C.1, §D.1.
Genevay et al. (2016) A. Genevay, M. Cuturi, G. Peyré, and F. Bach Stochastic optimization for large-scale optimal transport. In Advances in Neural Information Processing Systems, Vol. 29, pp. . External Links: Link Cited by: §C.1.
Grathwohl et al. (2019) W. Grathwohl, R. T. Q. Chen, J. Bettencourt, and D. Duvenaud FFJORD: free-form continuous dynamics for scalable reversible generative models. In International Conference on Learning Representations, External Links: Link Cited by: §1.
Guo and Schwing (2025) P. Guo and A. Schwing Variational rectified flow matching. In Forty-second International Conference on Machine Learning, External Links: Link Cited by: §1.
Heusel et al. (2017) M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter GANs trained by a two time-scale update rule converge to a local nash equilibrium. In Advances in Neural Information Processing Systems, Vol. 30, pp. . External Links: Link Cited by: §4.2.
Horadam (2012) K. Horadam Hadamard matrices and their applications. Princeton university press. Cited by: §B.3, §B.3.
Hu et al. (2026) K. Hu, L. Qiu, Y. Lu, H. Zhao, T. Li, Y. Kim, J. Andreas, and K. He ELF: embedded language flows. External Links: 2605.10938, Link Cited by: §1.
Hu et al. (2024) V. Hu, D. Wu, Y. Asano, P. Mettes, B. Fernando, B. Ommer, and C. Snoek Flow matching for conditional text generation in a few sampling steps. In Proceedings of the 18th Conference of the European Chapter of the Association for Computational Linguistics (Volume 2: Short Papers), St. Julian’s, Malta, pp. 380–392. External Links: Link, Document Cited by: §1.
Jiang et al. (2021) Y. Jiang, Z. Huang, X. Pan, C. C. Loy, and Z. Liu Talk-to-edit: fine-grained facial editing via dialog. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 13799–13808. Cited by: §C.4, §4.3.
Kingma and Ba (2015) D. P. Kingma and J. Ba Adam: a method for stochastic optimization. In International Conference on Learning Representations, Cited by: §C.1.
Krizhevsky et al. (2009) A. Krizhevsky G. Hinton et al. Learning multiple layers of features from tiny images. Cited by: §C.2, §4.2.
Le et al. (2023) M. Le, A. Vyas, B. Shi, B. Karrer, L. Sari, R. Moritz, M. Williamson, V. Manohar, Y. Adi, J. Mahadeokar, and W. Hsu Voicebox: text-guided multilingual universal speech generation at scale. In Thirty-seventh Conference on Neural Information Processing Systems, External Links: Link Cited by: §1.
Lipman et al. (2023) Y. Lipman, R. T. Q. Chen, H. Ben-Hamu, M. Nickel, and M. Le Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, External Links: Link Cited by: §C.1, §1, §1, §2, §2, §4.1.
Lipman et al. (2024) Y. Lipman, M. Havasi, P. Holderrieth, N. Shaul, M. Le, B. Karrer, R. T. Q. Chen, D. Lopez-Paz, H. Ben-Hamu, and I. Gat Flow matching guide and code. External Links: 2412.06264, Link Cited by: §C.1.
Liu et al. (2024) A. H. Liu, M. Le, A. Vyas, B. Shi, A. Tjandra, and W. Hsu Generative pre-training for speech with flow matching. In The Twelfth International Conference on Learning Representations, External Links: Link Cited by: §1.
Liu et al. (2023) X. Liu, C. Gong, and qiang liu Flow straight and fast: learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, External Links: Link Cited by: §1, §1.
Liu et al. (2015) Z. Liu, P. Luo, X. Wang, and X. Tang Deep learning face attributes in the wild. In Proceedings of the IEEE International Conference on Computer Vision (ICCV), Cited by: §C.4, §4.3.
Ma et al. (2024) N. Ma, M. Goldstein, M. S. Albergo, N. M. Boffi, E. Vanden-Eijnden, and S. Xie SiT: exploring flow and diffusion-based generative models with scalable interpolant transformers. In European Conference on Computer Vision, Cham, pp. 23–40. External Links: ISBN 978-3-031-72980-5 Cited by: §C.3, §1, §4.3.
Mousavi-Hosseini et al. (2026) A. Mousavi-Hosseini, S. Y. Zhang, M. Klein, and marco cuturi Flow matching with semidiscrete couplings. In The Fourteenth International Conference on Learning Representations, External Links: Link Cited by: §C.1, §D.1, §1, §2, §2, §3.3, §4.1.
Ning et al. (2023) M. Ning, E. Sangineto, A. Porrello, S. Calderara, and R. Cucchiara Input perturbation reduces exposure bias in diffusion models. In Proceedings of the 40th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 202, pp. 26245–26265. External Links: Link Cited by: §C.4.
Parmar et al. (2022) G. Parmar, R. Zhang, and J. Zhu On aliased resizing and surprising subtleties in gan evaluation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 11410–11420. Cited by: §C.2.
Pooladian et al. (2023) A. Pooladian, H. Ben-Hamu, C. Domingo-Enrich, B. Amos, Y. Lipman, and R. T. Q. Chen Multisample flow matching: straightening flows with minibatch couplings. In Proceedings of the 40th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 202, pp. 28100–28127. External Links: Link Cited by: §C.1, §1, §1, §2, §4.1.
Radford et al. (2021) A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, P. Mishkin, J. Clark, G. Krueger, and I. Sutskever Learning transferable visual models from natural language supervision. In Proceedings of the 38th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 139, pp. 8748–8763. External Links: Link Cited by: §4.3.
Ramachandran et al. (2017) P. Ramachandran, B. Zoph, and Q. V. Le Searching for activation functions. External Links: 1710.05941, Link Cited by: §C.1.
Rombach et al. (2022) R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 10684–10695. Cited by: §4.3.
Tong et al. (2024) A. Tong, K. FATRAS, N. Malkin, G. Huguet, Y. Zhang, J. Rector-Brooks, G. Wolf, and Y. Bengio Improving and generalizing flow-based generative models with minibatch optimal transport. Transactions on Machine Learning Research. Note: Expert Certification External Links: ISSN 2835-8856, Link Cited by: §C.2, §C.2, §C.3, §1, §1, §2, §2, §4.1, §4.2.
Tschannen et al. (2025) M. Tschannen, A. Gritsenko, X. Wang, M. F. Naeem, I. Alabdulmohsin, N. Parthasarathy, T. Evans, L. Beyer, Y. Xia, B. Mustafa, O. Hénaff, J. Harmsen, A. Steiner, and X. Zhai SigLIP 2: multilingual vision-language encoders with improved semantic understanding, localization, and dense features. External Links: 2502.14786, Link Cited by: §C.4, §4.3.
Villani et al. (2009) C. Villani et al. Optimal transport: old and new. Vol. 338, Springer. Cited by: §1.
Wang et al. (2023) Z. Wang, Y. Jiang, H. Zheng, P. Wang, P. He, Z. Wang, W. Chen, and M. Zhou Patch diffusion: faster and more data-efficient training of diffusion models. In Thirty-seventh Conference on Neural Information Processing Systems, External Links: Link Cited by: §C.4.
Yao et al. (2025) J. Yao, B. Yang, and X. Wang Reconstruction vs. generation: taming optimization dilemma in latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 15703–15712. Cited by: §C.5, §C.5, §4.3.
Zhang et al. (2025) S. Zhang, A. Mousavi-Hosseini, M. Klein, and M. Cuturi On fitting flow models with large sinkhorn couplings. External Links: 2506.05526, Link Cited by: §2.

Appendix A Proofs

A.1 Proof of Theorem 1

Proof. By the binary tree structure, for any two distinct data leaves $\ell\neq\ell^{\prime}$ , there exists a unique first internal node $\eta$ at which the root-to- $\ell$ and root-to- $\ell^{\prime}$ paths diverge. Without loss of generality, assume $\ell$ lies in the left subtree of $\eta$ and $\ell^{\prime}$ lies in the right subtree.

By the construction of the data tree, node $\eta$ is split along coordinate $\kappa_{\eta}$ with threshold $s_{\eta}$ , so its left and right child regions are

R_{\eta^{-}}=\{x\in R_{\eta}:x^{(\kappa_{\eta})}\leq s_{\eta}\},\qquad R_{\eta^{+}}=\{x\in R_{\eta}:x^{(\kappa_{\eta})}>s_{\eta}\}.

Since $x_{1}\in R_{\eta^{-}}$ and $x_{1}^{\prime}\in R_{\eta^{+}}$ ,

x_{1}^{(\kappa_{\eta})}\leq s_{\eta}<(x_{1}^{\prime})^{(\kappa_{\eta})}.

(9)

On the source side, QAT reuses the same split coordinate $\kappa_{\eta}$ and replaces the data-side threshold $s_{\eta}$ by the aligned threshold $\widetilde{s}_{\eta}$ . The corresponding source-side child regions are

\widetilde{R}_{\eta^{-}}=\{z\in\widetilde{R}_{\eta}:z^{(\kappa_{\eta})}\leq\widetilde{s}_{\eta}\},\qquad\widetilde{R}_{\eta^{+}}=\{z\in\widetilde{R}_{\eta}:z^{(\kappa_{\eta})}>\widetilde{s}_{\eta}\}.

Since QAT matches each data leaf with the corresponding source leaf region, the source samples paired with leaves $\ell$ and $\ell^{\prime}$ satisfy $x_{0}\in\widetilde{R}_{\eta^{-}}$ and $x_{0}^{\prime}\in\widetilde{R}_{\eta^{+}}$ . Hence

x_{0}^{(\kappa_{\eta})}\leq\widetilde{s}_{\eta}<(x_{0}^{\prime})^{(\kappa_{\eta})}.

(10)

For any $t\in[0,1]$ , the two linear interpolation paths are

x_{t}=(1-t)x_{0}+tx_{1},\qquad x_{t}^{\prime}=(1-t)x_{0}^{\prime}+tx_{1}^{\prime}.

Examining the $\kappa_{\eta}$ -th coordinate,

(x_{t}^{\prime})^{(\kappa_{\eta})}-x_{t}^{(\kappa_{\eta})}=(1-t)\bigl\{(x_{0}^{\prime})^{(\kappa_{\eta})}-x_{0}^{(\kappa_{\eta})}\bigr\}+t\bigl\{(x_{1}^{\prime})^{(\kappa_{\eta})}-x_{1}^{(\kappa_{\eta})}\bigr\}.

By (9) and (10), both differences on the right-hand side are strictly positive. Therefore,

(x_{t}^{\prime})^{(\kappa_{\eta})}-x_{t}^{(\kappa_{\eta})}>0,\qquad\forall\,t\in[0,1],

so $x_{t}^{\prime}\neq x_{t}$ for all $t\in[0,1]$ . This proves that linear interpolation paths from two distinct QAT leaves are strictly non-coincident almost surely. $\square$

A.2 Proof of Theorem 2

Proof. For any coupling $\pi\in\Pi(\gamma,\widehat{\mu}_{\mathcal{D}})$ , let $(X_{0},X_{1})$ and $(X_{0}^{\prime},X_{1}^{\prime})$ be two independent pairs drawn from $\pi$ . By the linear interpolation $X_{t}-X_{t}^{\prime}=(1-t)(X_{0}-X_{0}^{\prime})+t(X_{1}-X_{1}^{\prime})$ ,

	$\displaystyle\mathbb{E}_{\pi^{\otimes 2}}\\|X_{t}-X_{t}^{\prime}\\|_{2}^{2}=$	$\displaystyle(1-t)^{2}\mathbb{E}\\|X_{0}-X_{0}^{\prime}\\|_{2}^{2}+t^{2}\mathbb{E}\\|X_{1}-X_{1}^{\prime}\\|_{2}^{2}$		(11)
		$\displaystyle+2t(1-t)\mathbb{E}_{\pi^{\otimes 2}}\!\left[(X_{0}-X_{0}^{\prime})^{\top}(X_{1}-X_{1}^{\prime})\right].$		(11)

The first two terms depend only on the marginal distributions $\gamma$ and $\widehat{\mu}_{\mathcal{D}}$ , and are therefore identical under $\pi_{\mathrm{QAT}}$ and $\pi_{\mathrm{ind}}$ . The difference between the two couplings arises entirely from the cross term. Under the independent coupling $\pi_{\mathrm{ind}}=\gamma\otimes\widehat{\mu}_{\mathcal{D}}$ , $X_{0}-X_{0}^{\prime}$ and $X_{1}-X_{1}^{\prime}$ are mutually independent with zero mean, giving

\mathbb{E}_{\pi_{\mathrm{ind}}^{\otimes 2}}\!\left[(X_{0}-X_{0}^{\prime})^{\top}(X_{1}-X_{1}^{\prime})\right]=0.

(12)

Recursive analysis of the QAT cross term. For each node $\eta$ , let $\pi_{\eta}$ denote the normalized QAT coupling conditioned on node $\eta$ ’s region, and let $\mathbb{E}_{\eta}$ denote expectation when both pairs are drawn independently from $\pi_{\eta}$ . Define

\Gamma_{\eta}:=\mathbb{E}_{\eta}\!\left[(X_{0}-X_{0}^{\prime})^{\top}(X_{1}-X_{1}^{\prime})\right].

If $\eta$ is a leaf, the QAT coupling within that leaf is $\gamma_{\eta}\otimes\widehat{\mu}_{\eta}$ , so $X_{0}-X_{0}^{\prime}$ and $X_{1}-X_{1}^{\prime}$ are independent with zero mean, giving $\Gamma_{\eta}=0$ .

For any internal node $\eta\in\mathcal{I}$ with left and right children $\eta^{-}$ and $\eta^{+}$ , let

Z_{\eta}:=\mathbf{1}\{X_{0}\in\widetilde{R}_{\eta^{-}},\,X_{1}\in R_{\eta^{-}}\},\qquad Z_{\eta}^{\prime}:=\mathbf{1}\{X_{0}^{\prime}\in\widetilde{R}_{\eta^{-}},\,X_{1}^{\prime}\in R_{\eta^{-}}\}

indicate whether the first and second pairs fall into the left child. Write $p_{\eta}^{-}:=\widehat{\mu}_{\eta}(R_{\eta^{-}})$ and $p_{\eta}^{+}:=\widehat{\mu}_{\eta}(R_{\eta^{+}})$ . By the law of total expectation,

$\displaystyle\Gamma_{\eta}=$	$\displaystyle\mathbb{E}\!\left[\mathbb{E}_{\eta}\!\left[(X_{0}-X_{0}^{\prime})^{\top}(X_{1}-X_{1}^{\prime})\,\middle\|\,Z_{\eta},Z_{\eta}^{\prime}\right]\right]$	(13)
$\displaystyle=$	$\displaystyle(p_{\eta}^{-})^{2}\Gamma_{\eta^{-}}+(p_{\eta}^{+})^{2}\Gamma_{\eta^{+}}$
	$\displaystyle+p_{\eta}^{-}p_{\eta}^{+}\,\mathbb{E}_{\eta}\!\left[(X_{0}-X_{0}^{\prime})^{\top}(X_{1}-X_{1}^{\prime})\,\middle\|\,Z_{\eta}=1,\,Z_{\eta}^{\prime}=0\right]$
	$\displaystyle+p_{\eta}^{+}p_{\eta}^{-}\,\mathbb{E}_{\eta}\!\left[(X_{0}-X_{0}^{\prime})^{\top}(X_{1}-X_{1}^{\prime})\,\middle\|\,Z_{\eta}=0,\,Z_{\eta}^{\prime}=1\right].$

When both pairs fall into the same child, they contribute $(p_{\eta}^{-})^{2}\Gamma_{\eta^{-}}$ and $(p_{\eta}^{+})^{2}\Gamma_{\eta^{+}}$ respectively. For the cross-subtree terms, write $m_{0,\eta}^{\pm}:=\mathbb{E}_{\gamma_{\eta^{\pm}}}[X]$ and $m_{1,\eta}^{\pm}:=\mathbb{E}_{\widehat{\mu}_{\eta^{\pm}}}[X]$ . Taking the case $Z_{\eta}=0$ , $Z_{\eta}^{\prime}=1$ (first pair in right child, second in left), the cross-subtree contribution is

		$\displaystyle p_{\eta}^{+}p_{\eta}^{-}\,\mathbb{E}_{\eta}\!\left[(X_{0}-X_{0}^{\prime})^{\top}(X_{1}-X_{1}^{\prime})\,\middle\|\,Z_{\eta}=0,\,Z_{\eta}^{\prime}=1\right]$
	$\displaystyle=$	$\displaystyle p_{\eta}^{+}p_{\eta}^{-}\bigg\{\mathbb{E}_{\eta^{+}}[X_{0}^{\top}X_{1}]-(m_{0,\eta}^{+})^{\top}m_{1,\eta}^{-}-(m_{0,\eta}^{-})^{\top}m_{1,\eta}^{+}+\mathbb{E}_{\eta^{-}}[X_{0}^{\top}X_{1}]\bigg\}.$

Since $\Gamma_{\eta^{\pm}}=2\bigl\{\mathbb{E}_{\eta^{\pm}}[X_{0}^{\top}X_{1}]-(m_{0,\eta}^{\pm})^{\top}m_{1,\eta}^{\pm}\bigr\}$ , the above simplifies to

p_{\eta}^{+}p_{\eta}^{-}\bigg\{\frac{1}{2}\Gamma_{\eta^{+}}+\frac{1}{2}\Gamma_{\eta^{-}}+\left(m_{0,\eta}^{+}-m_{0,\eta}^{-}\right)^{\top}\left(m_{1,\eta}^{+}-m_{1,\eta}^{-}\right)\bigg\}.

The case $Z_{\eta}=1$ , $Z_{\eta}^{\prime}=0$ contributes the same amount. Substituting into (13),

\Gamma_{\eta}=\,p_{\eta}^{-}\Gamma_{\eta^{-}}+p_{\eta}^{+}\Gamma_{\eta^{+}}+2p_{\eta}^{-}p_{\eta}^{+}\left(m_{0,\eta}^{+}-m_{0,\eta}^{-}\right)^{\top}\left(m_{1,\eta}^{+}-m_{1,\eta}^{-}\right).

(14)

Reduction to split-coordinate mean differences. Since the source distribution is $\mathcal{N}(0,I_{d})$ and all QAT splits are axis-aligned, the source regions $\widetilde{R}_{\eta^{-}}$ and $\widetilde{R}_{\eta^{+}}$ differ only in the $\kappa_{\eta}$ -th coordinate interval. By Gaussian coordinate independence, for any $j\neq\kappa_{\eta}$ ,

(\mathrm{proj}_{j})_{\#}\gamma_{\eta^{+}}=(\mathrm{proj}_{j})_{\#}\gamma_{\eta^{-}},

so $(m_{0,\eta}^{+})^{(j)}=(m_{0,\eta}^{-})^{(j)}$ for all $j\neq\kappa_{\eta}$ . Writing

	$\displaystyle\Delta_{\eta}$	$\displaystyle:=\left(m_{0,\eta}^{+}-m_{0,\eta}^{-}\right)^{\top}\left(m_{1,\eta}^{+}-m_{1,\eta}^{-}\right)$
		$\displaystyle=\,\left[(m_{0,\eta}^{+})^{(\kappa_{\eta})}-(m_{0,\eta}^{-})^{(\kappa_{\eta})}\right]\left[(m_{1,\eta}^{+})^{(\kappa_{\eta})}-(m_{1,\eta}^{-})^{(\kappa_{\eta})}\right],$

equation (14) becomes

\Gamma_{\eta}=p_{\eta}^{-}\Gamma_{\eta^{-}}+p_{\eta}^{+}\Gamma_{\eta^{+}}+2p_{\eta}^{-}p_{\eta}^{+}\Delta_{\eta}.

(15)

Telescoping from the root. Let $\alpha_{\eta}:=\widehat{\mu}_{\mathcal{D}}(R_{\eta})$ , $\alpha_{\eta^{-}}:=\widehat{\mu}_{\mathcal{D}}(R_{\eta^{-}})$ , $\alpha_{\eta^{+}}:=\widehat{\mu}_{\mathcal{D}}(R_{\eta^{+}})$ . Since $p_{\eta}^{-}=\alpha_{\eta^{-}}/\alpha_{\eta}$ and $p_{\eta}^{+}=\alpha_{\eta^{+}}/\alpha_{\eta}$ , multiplying (15) by $\alpha_{\eta}$ gives

\alpha_{\eta}\Gamma_{\eta}=\alpha_{\eta^{-}}\Gamma_{\eta^{-}}+\alpha_{\eta^{+}}\Gamma_{\eta^{+}}+2\frac{\alpha_{\eta^{-}}\alpha_{\eta^{+}}}{\alpha_{\eta}}\Delta_{\eta}.

(16)

Telescoping from the root $r$ (where $\alpha_{r}=1$ ) and using $\Gamma_{\ell}=0$ at all leaves,

\mathbb{E}_{\pi_{\mathrm{QAT}}^{\otimes 2}}\!\left[(X_{0}-X_{0}^{\prime})^{\top}(X_{1}-X_{1}^{\prime})\right]=\Gamma_{r}=2\sum_{\eta\in\mathcal{I}}\frac{\alpha_{\eta^{-}}\alpha_{\eta^{+}}}{\alpha_{\eta}}\Delta_{\eta}.

(17)

Setting $q_{\eta}:=\alpha_{\eta^{-}}/\alpha_{\eta}$ , so that $\frac{\alpha_{\eta^{-}}\alpha_{\eta^{+}}}{\alpha_{\eta}}=\alpha_{\eta}q_{\eta}(1-q_{\eta})$ ,

\mathbb{E}_{\pi_{\mathrm{QAT}}^{\otimes 2}}\!\left[(X_{0}-X_{0}^{\prime})^{\top}(X_{1}-X_{1}^{\prime})\right]=2\Delta_{\mathrm{QAT}}.

(18)

Combining (11), (12), and (18),

	$\displaystyle\mathbb{E}_{\pi_{\mathrm{QAT}}^{\otimes 2}}\!\left[\\|X_{t}-X_{t}^{\prime}\\|_{2}^{2}\right]-\mathbb{E}_{\pi_{\mathrm{ind}}^{\otimes 2}}\!\left[\\|X_{t}-X_{t}^{\prime}\\|_{2}^{2}\right]$	(19)
$\displaystyle=$	$\displaystyle 2t(1-t)\mathbb{E}_{\pi_{\mathrm{QAT}}^{\otimes 2}}\!\left[(X_{0}-X_{0}^{\prime})^{\top}(X_{1}-X_{1}^{\prime})\right]$
$\displaystyle=$	$\displaystyle 4t(1-t)\Delta_{\mathrm{QAT}}.$

Under non-degenerate splits, each effective split node satisfies $\Delta_{\eta}\geq 0$ ; if at least one node has $\Delta_{\eta}>0$ , then $\Delta_{\mathrm{QAT}}>0$ . This completes the proof. $\square$

A.3 Properties of CQAT

Let $\mathcal{D}_{c}=\{(x_{i},c_{i})\}_{i=1}^{N}\subseteq\mathbb{R}^{d}\times\mathbb{R}^{m}$ be the conditional dataset with joint empirical distribution $\widehat{\mu}_{\mathcal{D}_{c}}$ . Denote the Conditional QAT coupling by $\pi_{\mathrm{CQAT}}\in\Pi(\gamma,\widehat{\mu}_{\mathcal{D}_{c}})$ , where $\gamma=\mathcal{N}(0,I_{d})$ . Throughout, $X_{1}$ refers to the data-coordinate component of a joint sample $(X_{1},C)$ .

Proposition 1 (Non-crossing property of CQAT)

Let $\ell\neq\ell^{\prime}$ be two leaf nodes of the Conditional QAT tree, and let $\eta$ be the first internal node at which the root-to- $\ell$ and root-to- $\ell^{\prime}$ paths diverge. If $\eta$ is a data-coordinate split node, i.e., $\kappa_{\eta}\leq d$ , then the linear interpolation paths drawn from the two leaf couplings do not intersect.

Proof. Without loss of generality, assume $\ell$ lies in the left subtree $\eta^{-}$ and $\ell^{\prime}$ in the right subtree $\eta^{+}$ . Since $\kappa_{\eta}\leq d$ , the CQAT tree performs a data-coordinate split at $\eta$ . By the tree construction, for any $(x_{1},c_{1})\in R_{\ell}$ and $(x_{1}^{\prime},c_{1}^{\prime})\in R_{\ell^{\prime}}$ ,

x_{1}^{(\kappa_{\eta})}\leq s_{\eta}<(x_{1}^{\prime})^{(\kappa_{\eta})}.

Moreover, since $\eta$ is a data-coordinate split node, the source-side region is simultaneously partitioned: for any $x_{0}\in\widetilde{R}_{\ell}$ and $x_{0}^{\prime}\in\widetilde{R}_{\ell^{\prime}}$ ,

x_{0}^{(\kappa_{\eta})}\leq\widetilde{s}_{\eta}<(x_{0}^{\prime})^{(\kappa_{\eta})}.

Hence, for all $t\in[0,1]$ ,

(x_{t}^{\prime})^{(\kappa_{\eta})}-x_{t}^{(\kappa_{\eta})}=(1-t)\bigl[(x_{0}^{\prime})^{(\kappa_{\eta})}-x_{0}^{(\kappa_{\eta})}\bigr]+t\bigl[(x_{1}^{\prime})^{(\kappa_{\eta})}-x_{1}^{(\kappa_{\eta})}\bigr]>0,

which implies $x_{t}^{\prime}\neq x_{t}$ for all $t\in[0,1]$ . $\square$

Proposition 2 (Path separation property of CQAT)

Let $x_{t}=(1-t)x_{0}+tx_{1}$ and $x_{t}^{\prime}=(1-t)x_{0}^{\prime}+tx_{1}^{\prime}$ , where $x_{1},x_{1}^{\prime}$ denote the data-coordinate components of joint samples. For any $t\in(0,1)$ ,

\mathbb{E}_{\pi_{\mathrm{CQAT}}^{\otimes 2}}\!\left[\|x_{t}-x_{t}^{\prime}\|_{2}^{2}\right]-\mathbb{E}_{\pi_{\mathrm{ind}}^{\otimes 2}}\!\left[\|x_{t}-x_{t}^{\prime}\|_{2}^{2}\right]=4t(1-t)\Delta_{\mathrm{CQAT}},

where $\Delta_{\mathrm{CQAT}}$ is the node-level separation induced by the Conditional QAT tree structure:

	$\displaystyle\Delta_{\mathrm{CQAT}}$	$\displaystyle:=\sum_{\eta\in\mathcal{I}_{c}:\,\kappa_{\eta}\leq d}\alpha_{\eta}q_{\eta}(1-q_{\eta})\Bigl(\mathbb{E}_{\gamma_{\eta^{+}}}\!\bigl[X^{(\kappa_{\eta})}\bigr]-\mathbb{E}_{\gamma_{\eta^{-}}}\!\bigl[X^{(\kappa_{\eta})}\bigr]\Bigr)$		(20)
		$\displaystyle\times\Bigl(\mathbb{E}_{\widehat{\mu}_{\mathcal{D}_{c},\eta^{+}}}\!\bigl[X^{(\kappa_{\eta})}\bigr]-\mathbb{E}_{\widehat{\mu}_{\mathcal{D}_{c},\eta^{-}}}\!\bigl[X^{(\kappa_{\eta})}\bigr]\Bigr).$		(20)

Here $\mathcal{I}_{c}$ is the set of internal nodes of the Conditional QAT tree; $\alpha_{\eta}:=\widehat{\mu}_{\mathcal{D}_{c}}(R_{\eta})$ is the joint data mass at node $\eta$ ; and $q_{\eta}:=\widehat{\mu}_{\mathcal{D}_{c},\eta}(R_{\eta^{-}})$ is the conditional mass of the left child.

Proof. The telescoping recursion (17) in the proof of Theorem 2 relies only on the left-right child decomposition at each internal node and the product structure of the leaf couplings. Replacing the QAT tree with the Conditional QAT tree and $\widehat{\mu}_{\mathcal{D}}$ with the joint empirical distribution $\widehat{\mu}_{\mathcal{D}_{c}}$ , the same argument gives

\mathbb{E}_{\pi_{\mathrm{CQAT}}^{\otimes 2}}\!\left[(X_{0}-X_{0}^{\prime})^{\top}(X_{1}-X_{1}^{\prime})\right]=2\sum_{\eta\in\mathcal{I}_{c}}\frac{\alpha_{\eta^{-}}\alpha_{\eta^{+}}}{\alpha_{\eta}}\Delta_{\eta},

where $\alpha_{\eta}:=\widehat{\mu}_{\mathcal{D}_{c}}(R_{\eta})$ and $\Delta_{\eta}$ is the local cross-term contribution at node $\eta$ .

It remains to show that condition-coordinate split nodes contribute zero. If $\kappa_{\eta}>d$ , the CQAT tree does not partition the source-side region at $\eta$ , so both child nodes inherit the same source region: $\gamma_{\eta^{-}}=\gamma_{\eta^{+}}$ . Consequently,

\mathbb{E}_{\gamma_{\eta^{+}}}[X]=\mathbb{E}_{\gamma_{\eta^{-}}}[X],

and the source-side mean difference vanishes, giving $\Delta_{\eta}=0$ . Therefore, only data-coordinate split nodes ( $\kappa_{\eta}\leq d$ ) contribute, and the sum reduces to

\mathbb{E}_{\pi_{\mathrm{CQAT}}^{\otimes 2}}\!\left[(X_{0}-X_{0}^{\prime})^{\top}(X_{1}-X_{1}^{\prime})\right]=2\Delta_{\mathrm{CQAT}}.

Combining with the linear interpolation expansion (11) and the fact that the cross-term vanishes under the independent coupling $\pi_{\mathrm{ind}}$ , we obtain

\mathbb{E}_{\pi_{\mathrm{CQAT}}^{\otimes 2}}\!\left[\|X_{t}-X_{t}^{\prime}\|_{2}^{2}\right]-\mathbb{E}_{\pi_{\mathrm{ind}}^{\otimes 2}}\!\left[\|X_{t}-X_{t}^{\prime}\|_{2}^{2}\right]=4t(1-t)\Delta_{\mathrm{CQAT}}.

Under non-degenerate splits, if at least one data-coordinate split node contributes positively, then $\Delta_{\mathrm{CQAT}}>0$ . This completes the proof. $\square$

Appendix B Algorithms

B.1 QAT Coupling Construction and Sampling

This section presents pseudocode for QAT coupling construction and training-pair sampling. QAT uses $\gamma=\mathcal{N}(0,I_{d})$ as the source distribution and maintains source-side coordinate bounds $(a_{\ell},b_{\ell})\in(\mathbb{R}\cup\{\pm\infty\})^{d}\times(\mathbb{R}\cup\{\pm\infty\})^{d}$ at each leaf node. Throughout, $\Phi$ and $\Phi^{-1}$ denote the standard-normal CDF and its inverse, respectively.

Algorithm 1 QAT Coupling Construction

1: Dataset

\mathcal{D}=\{x_{i}\}_{i=1}^{N}\subseteq\mathbb{R}^{d}

\mathcal{S}_{r}\leftarrow\{1,\ldots,N\}

a_{r}\leftarrow(-\infty,\ldots,-\infty)

b_{r}\leftarrow(+\infty,\ldots,+\infty)

\mathcal{Q}\leftarrow\{r\}

\mathcal{L}\leftarrow\varnothing

4: while

\mathcal{Q}\neq\varnothing

5: Pop one node

\eta

from

\mathcal{Q}

6: if

|\mathcal{S}_{\eta}|\leq 1

then

7: Add

(\mathcal{S}_{\eta},\,a_{\eta},\,b_{\eta})

\mathcal{L}

8: else

\kappa_{\eta}\leftarrow\arg\max_{j\in[d]}\operatorname{Var}_{\{x_{i}:i\in\mathcal{S}_{\eta}\}}\!(x^{(j)})

s_{\eta}\leftarrow\frac{1}{|\mathcal{S}_{\eta}|}\sum_{i\in\mathcal{S}_{\eta}}x_{i}^{(\kappa_{\eta})}

10:

\mathcal{S}_{\eta^{-}}\leftarrow\{i\in\mathcal{S}_{\eta}:x_{i}^{(\kappa_{\eta})}\leq s_{\eta}\}

\mathcal{S}_{\eta^{+}}\leftarrow\mathcal{S}_{\eta}\setminus\mathcal{S}_{\eta^{-}}

q_{\eta}\leftarrow|\mathcal{S}_{\eta^{-}}|/|\mathcal{S}_{\eta}|

11:

\widetilde{s}_{\eta}\leftarrow\Phi^{-1}\!\bigl(\Phi(a_{\eta}^{(\kappa_{\eta})})+q_{\eta}\cdot\bigl(\Phi(b_{\eta}^{(\kappa_{\eta})})-\Phi(a_{\eta}^{(\kappa_{\eta})})\bigr)\bigr)

12:

(a_{\eta^{-}},b_{\eta^{-}})\leftarrow(a_{\eta},b_{\eta})

;

(a_{\eta^{+}},b_{\eta^{+}})\leftarrow(a_{\eta},b_{\eta})

13:

b_{\eta^{-}}^{(\kappa_{\eta})}\leftarrow\widetilde{s}_{\eta}

;

a_{\eta^{+}}^{(\kappa_{\eta})}\leftarrow\widetilde{s}_{\eta}

14: Push

\eta^{-}

and

\eta^{+}

into

\mathcal{Q}

15: end if

16: end while

17: return

\mathcal{L}=\{(\mathcal{S}_{\ell},\,a_{\ell},\,b_{\ell}):\ell\in\mathcal{L}\}

Because the construction terminates when $|\mathcal{S}_{\ell}|\leq 1$ , each non-empty leaf holds exactly one training sample. Every leaf therefore carries mass $1/N$ , and sampling a leaf is equivalent to drawing a data index uniformly at random.

Algorithm 2 QAT Coupling Training-Pair Sampling

1: Leaf list

\mathcal{L}=\{(\mathcal{S}_{\ell},\,a_{\ell},\,b_{\ell}):\ell\in\mathcal{L}\}

2: Sample

\ell\sim\operatorname{Unif}(\mathcal{L})

; let

\mathcal{S}_{\ell}=\{i\}

and set

x\leftarrow x_{i}

3: for

j=1,\ldots,d

4: Sample

u_{j}\sim\operatorname{Unif}\bigl(\Phi(a_{\ell}^{(j)}),\,\Phi(b_{\ell}^{(j)})\bigr)

; set

\epsilon^{(j)}\leftarrow\Phi^{-1}(u_{j})

5: end for

6: return

(\epsilon,\,x)

In flow matching training, the sampled pair $(\epsilon,x)$ serves as $(x_{0},x_{1})$ .

Complexity analysis.

Assuming the tree is approximately balanced, it has $\mathcal{O}(\log N)$ levels. At each level, computing coordinate means and variances over all $N$ samples costs $\mathcal{O}(Nd)$ , giving a total construction cost of $\mathcal{O}(Nd\log N)$ . Sampling a single training pair requires generating $d$ independent truncated-Gaussian variates via the inverse-CDF method, each in $O(1)$ , for a per-sample cost of $\mathcal{O}(d)$ . Epoch-level sampling therefore costs $\mathcal{O}(Nd)$ , matching standard Gaussian sampling in asymptotic complexity.

Vectorized implementation.

In practice, the source-side bounds for each sample’s corresponding leaf are precomputed and stacked into two matrices $\mathbf{A},\mathbf{B}\in(\mathbb{R}\cup\{\pm\infty\})^{N\times d}$ indexed by data order. Given a batch of indices $\mathcal{I}=\{i_{b}\}_{b=1}^{B}$ , one retrieves $\mathbf{A}_{\mathcal{I},:}$ and $\mathbf{B}_{\mathcal{I},:}$ directly and draws $U_{bj}\sim\operatorname{Unif}\bigl(\Phi(\mathbf{A}_{i_{b},j}),\Phi(\mathbf{B}_{i_{b},j})\bigr)$ in parallel, then sets $\epsilon_{bj}=\Phi^{-1}(U_{bj})$ . The entire batch sampling step is therefore fully vectorizable.

B.2 CQAT Coupling Construction and Sampling

In CQAT, the condition variable is embedded as a Euclidean vector and concatenated with the data feature to form $y_{i}=(x_{i},c_{i})\in\mathbb{R}^{d+m}$ ; the algorithm then builds a tree directly on the joint sample set $\mathcal{D}_{c}=\{y_{i}\}_{i=1}^{N}$ . To prevent late-stage splits driven purely by condition coordinates — which partition the data index set without updating the source-side bounds — we restrict the candidate split coordinates to $[d]$ whenever the node sample count falls below a threshold $n_{\mathrm{data}}$ . Algorithm 3 gives the full construction.

Algorithm 3 CQAT Coupling Construction

1: Joint dataset

\mathcal{D}_{c}=\{y_{i}\}_{i=1}^{N}\subseteq\mathbb{R}^{d+m}

, data-split threshold

n_{\mathrm{data}}

\mathcal{S}_{r}\leftarrow\{1,\ldots,N\}

a_{r}\leftarrow(-\infty,\ldots,-\infty)

b_{r}\leftarrow(+\infty,\ldots,+\infty)

\mathcal{Q}\leftarrow\{r\}

\mathcal{L}\leftarrow\varnothing

4: while

\mathcal{Q}\neq\varnothing

5: Pop one node

\eta

from

\mathcal{Q}

6: if

|\mathcal{S}_{\eta}|\leq 1

then

7: Add

(\mathcal{S}_{\eta},\,a_{\eta},\,b_{\eta})

\mathcal{L}

8: else

\mathcal{J}_{\eta}\leftarrow[d]

|\mathcal{S}_{\eta}|\leq n_{\mathrm{data}}

, else

\mathcal{J}_{\eta}\leftarrow[d+m]

10:

\kappa_{\eta}\leftarrow\arg\max_{j\in\mathcal{J}_{\eta}}\operatorname{Var}_{\{y_{i}:i\in\mathcal{S}_{\eta}\}}\!(y^{(j)})

s_{\eta}\leftarrow\frac{1}{|\mathcal{S}_{\eta}|}\sum_{i\in\mathcal{S}_{\eta}}y_{i}^{(\kappa_{\eta})}

11:

\mathcal{S}_{\eta^{-}}\leftarrow\{i\in\mathcal{S}_{\eta}:y_{i}^{(\kappa_{\eta})}\leq s_{\eta}\}

\mathcal{S}_{\eta^{+}}\leftarrow\mathcal{S}_{\eta}\setminus\mathcal{S}_{\eta^{-}}

12:

(a_{\eta^{-}},b_{\eta^{-}})\leftarrow(a_{\eta},b_{\eta})

;

(a_{\eta^{+}},b_{\eta^{+}})\leftarrow(a_{\eta},b_{\eta})

13: if

\kappa_{\eta}\leq d

then

14:

q_{\eta}\leftarrow|\mathcal{S}_{\eta^{-}}|/|\mathcal{S}_{\eta}|

15:

\widetilde{s}_{\eta}\leftarrow\Phi^{-1}\!\bigl(\Phi(a_{\eta}^{(\kappa_{\eta})})+q_{\eta}\cdot\bigl(\Phi(b_{\eta}^{(\kappa_{\eta})})-\Phi(a_{\eta}^{(\kappa_{\eta})})\bigr)\bigr)

16:

b_{\eta^{-}}^{(\kappa_{\eta})}\leftarrow\widetilde{s}_{\eta}

;

a_{\eta^{+}}^{(\kappa_{\eta})}\leftarrow\widetilde{s}_{\eta}

17: end if

18: Push

\eta^{-}

and

\eta^{+}

into

\mathcal{Q}

19: end if

20: end while

21: return

\mathcal{L}=\{(\mathcal{S}_{\ell},\,a_{\ell},\,b_{\ell}):\ell\in\mathcal{L}\}

Scale balancing.

Since CQAT builds on data and condition coordinates jointly, scale imbalance between $x_{i}$ and $c_{i}$ may bias split selection toward high-variance condition dimensions. We therefore build the tree on $(X,\,wC)$ . For discrete class labels, we set $w=\infty$ , which is formally equivalent to first partitioning all samples by class and then constructing an independent QAT within each class subtree — matching the approach of C2OT (Cheng and Schwing, 2025). For continuous conditions, we adopt the variance-normalized weight

w=\left(\frac{\sum_{j}\operatorname{Var}(X^{(j)})}{\sum_{k}\operatorname{Var}(C^{(k)})}\right)^{1/2},

(21)

so that the rescaled condition variable has aggregate variance comparable to that of the data features.

Complexity analysis.

Building the tree over $d+m$ coordinates gives a construction cost of $\mathcal{O}(N(d+m)\log N)$ under a balanced-tree assumption. Condition-coordinate splits ( $\kappa_{\eta}>d$ ) partition the data index set without touching the source-side bounds; data-coordinate splits ( $\kappa_{\eta}\leq d$ ) perform quantile alignment exactly as in QAT. At training time, Algorithm 2 applies directly: one returns the data sample $x_{i}$ together with its paired condition variable $c_{i}$ . Since source sampling remains confined to $d$ -dimensional Gaussian space, the per-sample cost is $\mathcal{O}(d)$ , unchanged from QAT.

B.3 High-Dimensional Image Data Processing

For high-dimensional image data, we do not build the QAT directly in raw pixel space. Instead, we first apply an orthogonal transform to concentrate the principal variation into the leading coordinates. Crucially, only a pure rotation is applied — no mean subtraction — so that if the source is $\mathcal{N}(0,I_{d})$ , it remains standard Gaussian under both the forward and inverse transforms.

Patch-Hadamard transform.

Let the input image have spatial resolution $H\times W$ with $C$ channels, flattened to $x_{i}\in\mathbb{R}^{d}$ ( $d=HWC$ ). For patch size $p$ , the image is divided into $d^{\prime}=(H/p)(W/p)C$ non-overlapping patches, each of spatial dimension $r=p^{2}$ . A normalized Hadamard matrix $H_{\mathrm{patch}}\in\mathbb{R}^{r\times r}$ (Horadam, 2012) is applied independently to each patch. Because the normalized Hadamard matrix is orthogonal with constant direction $r^{-1/2}\mathbf{1}_{r}$ , extracting the coefficient along this direction is equivalent to patch-level mean pooling. We denote the resulting mean-pooled representation as $\bar{x}_{i}\in\mathbb{R}^{d^{\prime}}$ .

Let $H_{p}\in\mathbb{R}^{d\times d}$ denote the global patch-wise Hadamard transform induced by patch size $p$ . Since each per-patch transform is orthogonal, $H_{p}$ is itself orthogonal. We order the output coordinates so that the first $d^{\prime}$ entries of $H_{p}x_{i}$ correspond to the constant-direction coefficients across all patches. In practice, if only the low-dimensional tree representation is needed, $\bar{x}_{i}$ can be obtained directly by patch-level mean pooling, with PCA directions estimated on this $d^{\prime}$ -dimensional representation.

PCA rotation and low-dimensional QAT.

Patch-Hadamard pooling avoids the cost of PCA in the full $d$ -dimensional pixel space. We estimate the PCA rotation matrix $\mathbf{P}\in\mathbb{R}^{d^{\prime}\times d^{\prime}}$ from $\{\bar{x}_{i}\}_{i=1}^{N}$ without mean shift. Setting $z_{i}=\mathbf{P}\bar{x}_{i}\in\mathbb{R}^{d^{\prime}}$ , we use only the leading $k$ coordinates $z_{i}^{(1:k)}$ to construct the QAT coupling. Algorithm 4 summarizes the full construction pipeline.

Algorithm 4 High-Dimensional QAT Construction

1: Image dataset

\mathcal{D}=\{x_{i}\}_{i=1}^{N}

, patch size

p

, reduced dimension

k

d^{\prime}\leftarrow(H/p)(W/p)C

\bar{x}_{i}\leftarrow(H_{p}x_{i})^{(1:d^{\prime})}

, for

i=1,\ldots,N

4: Estimate PCA rotation

\mathbf{P}\in\mathbb{R}^{d^{\prime}\times d^{\prime}}

from

\{\bar{x}_{i}\}_{i=1}^{N}

without mean shift

\widetilde{z}_{i}\leftarrow(\mathbf{P}\bar{x}_{i})^{(1:k)}

, for

i=1,\ldots,N

6: Run Algorithm 1 on

\{\widetilde{z}_{i}\}_{i=1}^{N}

7: return

\mathcal{L}

\mathbf{P}

H_{p}

Inverse-transform sampling.

The low-dimensional QAT constructs a coupling over the leading $k$ PCA-pooled coordinates, whose source marginal is $\mathcal{N}(0,I_{k})$ . To recover a source sample in the original $d$ -dimensional pixel space, we fill the remaining coordinates with independent standard Gaussian noise: first the $d^{\prime}-k$ residual PCA-pooled coordinates, then the high-frequency Hadamard coefficients for each patch (all per-patch dimensions except the constant direction). Since $\mathbf{P}$ and $H_{p}$ are both orthogonal and all unfilled coordinates are drawn as independent standard Gaussians, the resulting source sample has marginal exactly $\mathcal{N}(0,I_{d})$ .

Algorithm 5 High-Dimensional QAT Sampling

\mathcal{D}=\{x_{i}\}_{i=1}^{N}

, leaf list

\mathcal{L}=\{(\mathcal{S}_{\ell},\,a_{\ell},\,b_{\ell}):\ell\in\mathcal{L}\}

\mathbf{P}

H_{p}

d^{\prime}

k

2: Sample one leaf

\ell\sim\operatorname{Unif}(\mathcal{L})

; let

\mathcal{S}_{\ell}=\{i\}

and set

x\leftarrow x_{i}

3: for

j=1,\ldots,k

4: Sample

u_{j}\sim\operatorname{Unif}\bigl(\Phi(a_{\ell}^{(j)}),\,\Phi(b_{\ell}^{(j)})\bigr)

; set

\epsilon^{(j)}\leftarrow\Phi^{-1}(u_{j})

5: end for

6: Sample

\epsilon^{(k+1:d^{\prime})}\sim\mathcal{N}(0,I_{d^{\prime}-k})

\epsilon_{\mathrm{pool}}\leftarrow\bigl(\epsilon^{(1:k)},\,\epsilon^{(k+1:d^{\prime})}\bigr)

\bar{\epsilon}\leftarrow\mathbf{P}^{\top}\epsilon_{\mathrm{pool}}

9: Sample

\epsilon_{\mathrm{high}}\sim\mathcal{N}(0,I_{d-d^{\prime}})

10:

\epsilon_{H}\leftarrow\bigl(\bar{\epsilon},\,\epsilon_{\mathrm{high}}\bigr)

11:

x_{0}\leftarrow H_{p}^{\top}\epsilon_{H}

12: return

(x_{0},\,x)

Complexity analysis.

Let $d=HWC$ , $r=p^{2}$ , and $d^{\prime}=(H/p)(W/p)C$ . Computing patch-level mean pooling costs $\mathcal{O}(Nd)$ ; the full patch-wise Hadamard transform can be applied via the Fast Walsh-Hadamard Transform (Horadam, 2012) in $\mathcal{O}(Nd\log r)$ . PCA is estimated on the $d^{\prime}$ -dimensional pooled representation, and the low-dimensional QAT is built in $k$ -dimensional space, giving tree construction cost $\mathcal{O}(Nk\log N)$ . At sampling time, the inverse PCA rotation costs $\mathcal{O}(d^{\prime 2})$ per sample and the inverse patch-wise Hadamard transform costs $\mathcal{O}(d\log r)$ , giving a total per-sample cost of $\mathcal{O}(d+d^{\prime 2}+d\log r)$ ; batch sampling is fully vectorizable.

Appendix C Implementation Details

This section details the per-experiment training configurations. Unless stated otherwise, all methods share the same network architecture, training configuration, and sampling protocol, differing only in the coupling strategy used to form training pairs: I-FM uses independent coupling, OT-FM uses mini-batch OT coupling, and QAT-FM uses the QAT or Conditional QAT coupling proposed in this work.

C.1 Two-Dimensional Experiment

Data.

Following Lipman et al. (2023); Pooladian et al. (2023), we consider transport from a two-dimensional standard Gaussian source distribution to a checkerboard target distribution. Training samples are drawn independently from the source and target; the flow paths learned by I-FM, OT-FM, SD-FM, and QAT-FM are then compared under the same network architecture and optimization configuration.

Network and training.

Following Lipman et al. (2024), the velocity field is parameterized by a five-layer MLP with hidden width $512$ and Swish activations (Ramachandran et al., 2017). All methods are trained for $20$ K iterations with batch size $256$ and Adam (Kingma and Ba, 2015) at learning rate $10^{-3}$ . For QAT-FM, $50$ K samples are drawn from the checkerboard target to build the QAT tree, and training pairs are sampled from the induced coupling. For OT-FM, exact EMD is solved within each mini-batch using POT (Flamary et al., 2021). For SD-FM, a discrete reference measure is built from $50$ K target samples and the corresponding regularized problem is solved with entropy parameter $0.05$ ; we optimize with Adam, as the SGD-based strategy of Genevay et al. (2016); Mousavi-Hosseini et al. (2026) proved unstable and slow to converge in this setting.

Table 6: Hyperparameter settings for image-space experiments.

	CIFAR-10	ImageNet- $32$	CelebA- $64$
Channels	192	256	192
Depth	2	3	3
Channel multiplier	1,2,2,2	1,2,2,2	1,2,3,4
Heads	4	4	4
Head channels	64	64	64
Attention resolution	16	4	8
Dropout	0.0	0.0	0.0
Batch size / GPU	256	128	64
GPUs	1	4	4
Effective batch size	256	512	256
Iterations	$100$ k	$200$ k	$100$ k
Learning rate	2e-4	1e-4	1e-4
LR scheduler	Constant	Linear decay	Constant
Warmup steps	5k	20k	10k

C.2 CIFAR-10

We conduct unconditional image generation on CIFAR-10 (Krizhevsky et al., 2009), using the U-Net architecture of Tong et al. (2024). Main network and training hyperparameters are listed in Table 6. We train with Adam, a linear warmup to $2\times 10^{-4}$ over $5$ k steps (constant thereafter), gradient norm clipping at $1.0$ , and an EMA decay of $0.9999$ .

For QAT-FM, the QAT coupling is built from the $50{,}000$ CIFAR-10 training images together with their horizontally flipped augmentations, giving $n=100{,}000$ tree samples. Images are first processed by patch-level mean pooling with patch size $p=4$ ; PCA rotation is then estimated on the pooled representation, and the leading $k=32$ coordinates are retained for tree construction. The tree splits on the maximum-variance coordinate at the empirical mean threshold and recurses until each leaf contains at most one sample. At evaluation, $50$ K images are generated, and FID is computed using clean-fid (Parmar et al., 2022), following the protocol of Tong et al. (2024).

C.3 ImageNet-32

We evaluate class-conditional generation on ImageNet-1k (Deng et al., 2009) at $32\times 32$ resolution. All images are resized to $32\times 32$ with random horizontal flipping, and class labels are provided as condition inputs to the velocity-field network. We follow the class-conditional U-Net architecture of Tong et al. (2024); main hyperparameters are listed in Table 6. Optimizer, gradient clipping, and EMA settings match the CIFAR-10 setup. The learning rate warms up linearly to $1\times 10^{-4}$ over $20$ k steps and then decays linearly to $3\times 10^{-5}$ .

For QAT-FM, the Conditional QAT coupling is built from ImageNet-1k training images and their horizontal-flip augmentations, yielding over $2$ M tree samples. Patch-level mean pooling ( $p=4$ ) and PCA rotation are applied as in the CIFAR-10 setup, retaining the leading $k=32$ coordinates. Since the condition is a discrete class label, training samples are first partitioned by class and an independent QAT coupling is constructed within each class subset. For C2OT (Cheng and Schwing, 2025), we follow its oversampling strategy with ratio $50$ , yielding a local OT batch size of $128\times 50=6400$ per GPU. Following the C2OT setting for class-conditional generation, we use its conditional weighting scheme and set the condition weight to $10^{8}$ .

At evaluation, $50$ K image samples are generated to estimate Inception feature statistics of the generated distribution; reference statistics are computed from all ${\sim}1.28$ M ImageNet-1k training images, following Dhariwal and Nichol (2021); Ma et al. (2024).

C.4 CelebA-64

Data.

We evaluate text-conditional generation on CelebA (Liu et al., 2015) at $64\times 64$ resolution. Since CelebA does not provide text annotations by default, we use the natural-language descriptions from CelebA-Dialog (Jiang et al., 2021) as text conditions.

Network and Training.

We use the U-Net architecture with main hyperparameters in Table 6. Following Cheng and Schwing (2025), scale-shift normalization is applied in the U-Net residual blocks, and text embeddings are projected to the conditioning dimension via a linear layer before being fed to the network. Optimizer, gradient clipping, and EMA settings match the CIFAR-10 experiment; the learning rate warms up linearly to $1\times 10^{-4}$ over $10$ k steps and remains constant thereafter.

For QAT-FM, a Conditional QAT coupling is constructed in the joint space of low-dimensional image representations and text embeddings. The image-side representation is obtained via patch-level mean pooling ( $p=4$ ) and PCA rotation, retaining the leading $k=32$ coordinates. Text conditions are encoded by a DFN5B-CLIP-ViT-H-14 text encoder (Fang et al., 2024); the resulting embeddings are projected to $32$ dimensions via PCA and used as condition-side features in the joint tree. To prevent late-stage splits driven purely by condition coordinates, candidate split dimensions are restricted to data coordinates when a node’s sample count falls below $n_{\mathrm{data}}=128$ . For C2OT (Cheng and Schwing, 2025), we follow its oversampling strategy with ratio $50$ , a local OT batch size of $64\times 50=3200$ per GPU, which substantially increases the coupling cost; see Appendix D.1. For a fair conditional comparison, we use the Euclidean distance between text embeddings as the conditional distance and apply the weighting scheme in (21).

Evaluation.

At evaluation, $50$ K images are generated and assessed along two axes: image fidelity via FID against training-set reference statistics (Ning et al., 2023; Wang et al., 2023), and text-image alignment via CLIP score computed with the siglip2-base-patch16-224 checkpoint of SigLIP-2 (Tschannen et al., 2025).

C.5 ImageNet-256

We conduct class-conditional latent-space generation on ImageNet- $256\times 256$ (Deng et al., 2009). Images are encoded by VA-VAE (Yao et al., 2025) into latent representations, and a flow-based LightningDiT (Yao et al., 2025) is trained in the resulting latent space with class labels as condition inputs. All network architecture, optimizer, learning-rate schedule, EMA, precision, and sampling configurations follow the official implementation of Yao et al. (2025); due to computational constraints, we adopt the recommended lightweight training setting of $64$ epochs.

For QAT-FM, the Conditional QAT coupling is built directly over VA-VAE latent representations. PCA projection is applied, retaining the leading $k=256$ coordinates for tree construction; patch-level mean pooling is omitted since VA-VAE already maps images to a compact latent space. The tree splits on the maximum-variance coordinate at the empirical mean and recurses until each leaf contains at most one sample. Classifier-free guidance is applied during inference with a scale of $10$ .

At evaluation, we follow the protocol of Yao et al. (2025) and compute FID against the ImageNet- $256\times 256$ reference statistics from guided-diffusion (Dhariwal and Nichol, 2021). We additionally report IS, sFID, Precision, and Recall.

Appendix D Extended Experimental Results

This section provides supplementary results comprising a runtime comparison of coupling strategies and additional visualization of generated samples. Appendix D.1 benchmarks the wall-clock time of each coupling approach at both construction and sampling phases; the subsequent four subsections present additional generated samples from CIFAR-10, ImageNet- $32\times 32$ , CelebA- $64\times 64$ , and ImageNet- $256\times 256$ .

D.1 Runtime Comparison

We compare the wall-clock time of different coupling strategies at both construction and sampling phases. All experiments fix the feature dimension at $d=32$ and vary the dataset size $n$ . To reflect the repeated pairing demand of FM training, we report cumulative times over $100$ training epochs. For OT-based methods, source points are resampled from a standard Gaussian and the coupling is recomputed from scratch at every epoch; by contrast, the QAT and SD-FM couplings are constructed once offline and reused throughout training. Full OT is solved with the EMD solver from POT (Flamary et al., 2021); mini-batch OT (OTBatch) uses ot.dist_batch to compute batch pairwise distances, fully exploiting GPU parallelism. Entropy regularization with $\varepsilon=0.05$ is applied for both SD-FM and batch OT. SD-FM optimizes a semi-discrete coupling with SGD (Mousavi-Hosseini et al., 2026), running until convergence or $100$ optimization epochs.

Figure 6: Runtime comparison of different coupling strategies. Left: total coupling construction time over

100

epochs. Right: total sampling time over

100

epochs.

Figure 6 reports cumulative construction and sampling times over $100$ epochs. Full OT becomes rapidly infeasible as $n$ grows; batch OT reduces the per-epoch cost but must recompute the coupling every epoch, keeping its total time high. SD-FM requires only a single construction pass, but the semi-discrete optimization itself is expensive. By contrast, QAT constructs the tree and records the Gaussian bounds at each leaf in a single offline pass, achieving substantially lower construction time at all scales. At sampling time, QAT draws source points by querying leaf bounds and sampling from truncated Gaussians, placing its per-sample cost close to that of independent (random) sampling. OT-based methods must sample from the coupling’s joint distribution, and SD-FM requires a nearest-neighbor lookup over all training points at each draw—both become prohibitive at large $n$ . Overall, QAT achieves structured coupling at a sampling efficiency comparable to independent coupling.

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

Abstract

1 Introduction

2 Preliminaries

Notation.

Flow-based Generative Modeling.

Conditional Flow Matching.

Data Partitioning by Trees.

3 Quantile AlignTree Flow Matching

3.1 Quantile AlignTree Coupling

Definition 1 (Quantile AlignTree (QAT))

Definition 2 (QAT Coupling)

Remark 1 (Gaussian Sampling)

3.2 QAT Flow Matching

Theorem 1 (Non-Crossing Property)

Theorem 2 (Improved Path Separation)

3.3 Conditional QAT Strategy

Definition 3 (Conditional QAT)

Remark 2 (Theoretical extensions)

3.4 Efficient QAT Construction and Sampling

Split Strategy.

High-Dimensional Data Handling.

4 Experiments

4.1 2D Synthetic Experiments

4.2 Unconditional Generation on CIFAR-10

Ablation Studies.

4.3 Conditional Generation Experiments

Class-conditional generation.

Text-conditional generation.

Large-scale latent generation.

5 Conclusion

References

Appendix A Proofs

A.1 Proof of Theorem 1

A.2 Proof of Theorem 2

A.3 Properties of CQAT

Proposition 1 (Non-crossing property of CQAT)

Proposition 2 (Path separation property of CQAT)

Appendix B Algorithms

B.1 QAT Coupling Construction and Sampling

Complexity analysis.

Vectorized implementation.

B.2 CQAT Coupling Construction and Sampling

Scale balancing.

Complexity analysis.

B.3 High-Dimensional Image Data Processing

Patch-Hadamard transform.

PCA rotation and low-dimensional QAT.

Inverse-transform sampling.

Complexity analysis.

Appendix C Implementation Details

C.1 Two-Dimensional Experiment

Data.

Network and training.

C.2 CIFAR-10

C.3 ImageNet-32

C.4 CelebA-64

Data.

Network and Training.

Evaluation.

C.5 ImageNet-256

Appendix D Extended Experimental Results

D.1 Runtime Comparison

D.2 CIFAR-10

D.3 ImageNet-32

D.4 CelebA-64

D.5 ImageNet-256