Hierarchical Rectified Flow Matching with Mini-Batch Couplings

## Abstract

Abstract Flow matching has emerged as a compelling generative modeling approach that is widely used across domains.
To generate data via a flow matching model, an ordinary differential equation (ODE) is numerically solved via forward integration of the modeled velocity field.
To better capture the multi-modality that is inherent in typical velocity fields, hierarchical flow matching was recently introduced. It uses a hierarchy of ODEs that are numerically integrated when generating data.
This hierarchy of ODEs captures the multi-modal velocity distribution just like vanilla flow matching is capable of modeling a multi-modal data distribution.
While this hierarchy enables to model multi-modal velocity distributions, the complexity of the modeled distribution remains identical across levels of the hierarchy.
In this paper, we study how to gradually adjust the complexity of the distributions across different levels of the hierarchy via mini-batch couplings.
We show the benefits of mini-batch couplings in hierarchical rectified flow matching via compelling results on synthetic and imaging data.
Code is available at https://riccizz.github.io/HRF_coupling .

## 1 Introduction

Flow matching (Lipman et al., 2023; Liu et al., 2023a; Albergo and Vanden-Eijnden, 2023) has gained significant attention across computer vision (Esser et al., 2024; Liu et al., 2023b), robotics (Zhang and Gienger, 2024), computational biology (Yim et al., 2023; Jing et al., 2024), and time series analysis (Chen et al., 2024; Zhang et al., 2024). This is largely due to its ability to generate high-quality data and due to its simple simulation-free learning of a data distribution. For this, it uses 1) an intermediate state, which is computed by (linearly) interpolating between a sample from a known source distribution and a randomly drawn data point, and 2) the velocity at this intermediate state. This velocity controls a neural ordinary differential equation (ODE), which governs the transformation of the samples from the source distribution to the target data distribution. Note, the distribution of velocities at an intermediate state is multi-modal (Zhang et al., 2025).

The velocity field is obtained by matching velocities from the interpolated samples with a parametric deep net using a mean squared error (MSE) objective. It is known that the MSE objective used in classic rectified flow does not permit to capture this multi-modal velocity distribution. Instead, classic rectified flow leads to a model that aims to capture the mean of the velocity distribution. Capturing the mean of the velocity distribution is sufficient for characterizing a multi-modal data distribution (Liu et al., 2023a). However, it inevitably results in curved forward integration paths, making the sampling process inefficient. Recently, hierarchical rectified flow (Zhang et al., 2025) was suggested as an approach to model the multi-modal velocity field via coupled ODEs.

To model the multi-modal velocity field, hierarchical rectified flow essentially applies a rectified flow formulation in the velocity space by matching ‘acceleration.’ It was also suggested to expand the idea further towards an arbitrary hierarchy level.
While this enables to model multi-modal velocity distributions, the complexity of the modeled distribution remains identical across all levels of the hierarchy. Said differently, the velocity distribution that a hierarchical rectified flow models across levels of its hierarchy is no easier than the original data distribution, potentially limiting benefits.

We hence wonder: can we gradually simplify the complexity of the “ground-truth” velocity distribution? For simplicity, in this paper, we focus on two hierarchy levels. Interestingly, we find mini-batch couplings to provide a compelling way to control the “ground-truth” velocity distribution.
Instead of computing intermediate states by interpolating between samples independently drawn from both the known source distribution and the dataset, we draw a mini-batch of samples from both the source distribution and the dataset, and subsequently couple them, e.g., via a procedure like optimal transport. Intuitively, considering as an extreme situation a mini-batch containing the entire dataset leads to a velocity distribution that is uni-modal.

Empirically, we find that hierarchical flow matching with mini-batch coupling in the data space consistently improves the generation quality of vanilla hierarchical rectified flow matching and vanilla flow matching with optimal transport coupling. Jointly coupling mini-batch samples in data and velocity space leads to further benefits if the number of neural function evaluations (NFEs) is low.

## 2 Preliminaries

Rectified Flow (RF).
A rectified flow models an unknown target data distribution $\rho_{1}$ given a dataset ${\cal D}=\{x_{1}\}$, where we assume data points $x_{1}\sim\rho_{1}$. Given a known source distribution $\rho_{0}$ (e.g., standard Gaussian), at inference time, source samples $x_{0}\sim\rho_{0}$ evolve from time $t=0$ to time $t=1$ following the ordinary differential equation (ODE)

$$ $dz_{t}=v(z_{t},t)dt,\,\text{with }z_{0}\sim\rho_{0},\quad t\in[0,1].$ (1) $$

Here, $v(z_{t},t)$ is a velocity field that depends on time $t$ and the current intermediate state $z_{t}$. This ODE-based sampling enables
to capture multi-modal data distributions.

At training time, flow matching learns the velocity field $v(z_{t},t)$ by minimizing the $\ell_{2}$-loss between the predicted velocity $v(x_{t},t)$ and a ground-truth velocity $v_{\text{gt}}(x_{t},t)$. To obtain the ground-truth velocity we first define an
intermediate state $x_{t}$ which, in a rectified flow formulation, is obtained by linearly interpolating between a randomly drawn source sample $x_{0}$ and a randomly drawn data point $x_{1}$, i.e.,

$$ $x_{t}=(1-t)x_{0}+tx_{1},\quad\text{where}\,x_{0}\sim\rho_{0},\,x_{1}\sim{\cal D}.$ (2) $$

Interpreting the intermediate state $x_{t}$ as a location, we obtain the ground-truth velocity $v_{\text{gt}}(x_{t},t)=\partial x_{t}/\partial t=x_{1}-x_{0}$. Combined,
the training objective is then:

$$ $\inf_{v}\mathbb{E}_{x_{0}\sim\rho_{0},x_{1}\sim{\cal D},t\sim U[0,1]}\left[\|x _{1}-x_{0}-v(x_{t},t)\|_{2}^{2}\right],$ (3) $$

where the optimization is over all measurable velocity fields. In practice, $v(x_{t},t)$ is parameterized by a neural network with trainable parameters $\theta$, i.e., $v(x_{t},t)\approx v_{\theta}(x_{t},t)$, and the optimization minimizes over $\theta$.

However, for a given $t$ and $x_{t}$, different pairs $(x_{0},x_{1})$ will yield different ground-truth velocities. The ground-truth velocity distribution at a given time $t$ and intermediate state $x_{t}$ is hence multi-modal. However, the $\ell_{2}$-loss averages these velocities, resulting in the ‘optimal’ velocity field:
$v^{*}(x_{t},t)=\mathbb{E}_{\{(x_{0},x_{1},t):(1-t)x_{0}+tx_{1}=x_{t}\}}[v_{
\mathrm{gt}}(x_{t},t)].$
According to Theorem 3.3 by Liu et al. (2023a), using $v^{*}$ in [Equation 1](https://arxiv.org/html/2507.13350v1#S2.E1) ensures that the stochastic process has marginal distributions consistent with the linear interpolation in [Equation 2](https://arxiv.org/html/2507.13350v1#S2.E2).

To capture
multi-modal velocity distributions, hierarchical rectified flow (Zhang et al., 2025) was introduced. It explicitly models the multi-modal velocity distributions at each time $t$ and intermediate state $x_{t}$, enabling a more expressive generative framework.

Hierarchical Rectified Flow (HRF).
To model the “ground-truth” velocity distribution more accurately, hierarchical rectified flow extends the classic rectified flow framework by focusing on velocities rather than locations. This approach effectively involves learning acceleration. In a classic rectified flow, the time-dependent location $x_{t}$ is computed from pairs $(x_{0},x_{1})$, and the ground-truth velocity $v_{\text{gt}}(x_{t},t)=\partial x_{t}/\partial t$ is used to train a velocity model $v_{\theta}(x_{t},t)$.

In hierarchical rectified flow, a source velocity sample $v_{0}\sim\pi_{0}$ is drawn from a known velocity distribution $\pi_{0}$, while a target velocity sample $v_{1}(x_{t},t)\sim\pi_{1}(v;x_{t},t)$ is defined at each time $t$ and location $x_{t}=(1-t)x_{0}+tx_{1}$. For rectified flow, $v_{1}(x_{t},t)$ is computed via $x_{1}-x_{0}$, and these samples follow the ground-truth velocity distribution $\pi_{1}(v;x_{t},t)$.

To learn acceleration, a new time axis $\tau\in[0,1]$ is introduced, and a time-dependent velocity $v_{\tau}(x_{t},t)=(1-\tau)v_{0}+\tau v_{1}(x_{t},t)$ is constructed. The ground-truth acceleration is then obtained as $a(x_{t},t,v_{\tau},\tau)=\partial v_{\tau}/\partial\tau=v_{1}(x_{t},t)-v_{0}=x
_{1}-x_{0}-v_{0}$. For a fixed $(x_{t},t)$, this leads to the ODE:

$$ $du_{\tau}(x_{t},t)=a(x_{t},t,u_{\tau},\tau)d\tau,\quad\text{with }u_{0}\sim\pi _{0}.$ (4) $$

Here, $a(x_{t},t,u_{\tau},\tau)$ is the expected acceleration vector field: $a(x_{t},t,u_{\tau},\tau)=\mathbb{E}_{\pi_{0},\pi_{1}(v;x_{t},t)}[V_{1}-V_{0}|V
_{\tau}=u]$. The acceleration vector field is learned by minimizing the following objective:

$$ $\inf_{a}\mathbb{E}_{x_{0}\sim\rho_{0},x_{1}\sim\mathcal{D},t\sim U[0,1],v_{0} \sim\pi_{0},\tau\sim U[0,1]}\left[\|(x_{1}-x_{0}-v_{0})-a(x_{t},t,v_{\tau}, \tau)\|_{2}^{2}\right].$ (5) $$

In practice, the acceleration is parameterized via a deep net $a_{\theta}(x_{t},t,v_{\tau},\tau)$, and the model is trained by minimizing this objective over the parameters $\theta$.

During sampling, coupled ODEs are used:

$$ $\begin{cases}du_{\tau}(z_{t},t)=a(z_{t},t,u_{\tau},\tau)d\tau,&u_{0}\sim\pi_{0 },\ \tau\in[0,1],\\ dz_{t}=u_{1}(z_{t},t)dt,&z_{0}\sim\rho_{0},\ t\in[0,1].\end{cases}$ (6) $$

These ODEs map $z_{0}\sim\rho_{0}$ to $z_{1}\sim\rho_{1}$. Sampling involves drawing $v_{0}\sim\pi_{0}$ and $x_{0}\sim\rho_{0}$, integrating forward to obtain $v_{1}(x_{0},0)$, and then performing location updates iteratively until reaching $x_{1}$. This procedure can be implemented using the vanilla Euler method and the trained $a_{\theta}$.

Considering the training objective for acceleration matching ([Equation 5](https://arxiv.org/html/2507.13350v1#S2.E5)) and the coupled ODEs for sampling ([Equation 6](https://arxiv.org/html/2507.13350v1#S2.E6)), both can be naturally extended to any depth. In this paper, we focus solely on depth-two HRF (HRF2) models.

Minibatch Optimal Transport.
Optimal Transport (OT) seeks to find an optimal coupling of two distributions that minimizes an expected transport cost (Villani, 2009). Suppose $\alpha$ and $\beta$ are two distributions in $\mathbb{R}^{d}$, and $c:\mathbb{R}^{d}\times\mathbb{R}^{d}\rightarrow\mathbb{R}$ is some distance. Then OT aims to find the solution of the following optimization problem:

$$ $\inf_{\gamma\in\Gamma}\int_{\mathbf{R}^{d}\times\mathbf{R}^{d}}c^{2}(x,y)d \gamma(x,y),$ (7) $$

where $\Gamma$ is the set of all joint distributions with marginals $\alpha$ and $\beta$. When $\alpha$ and $\beta$ are both empirical distributions, OT reduces to linear programming, which is computationally expensive when the data size is large (Peyré et al., 2019). While OT is computationally expensive for large datasets, mini-batch OT (Fatras et al., 2020, 2021) was introduced as an alternative: a small batch of the data is used to calculate the coupling, obtaining an unbiased estimator of the underlying truth (Fatras et al., 2020). Although mini-batch OT incurs an error compared to the exact OT, it has found use in practice (Deshpande et al., 2018, 2019; Pooladian et al., 2023; Tong et al., 2024; Cheng and Schwing, 2025). Tong et al. (2024); Pooladian et al. (2023) showed that training and inference are more efficient with mini-batch OT.

## 3 Approach

In [Section 3.1](https://arxiv.org/html/2507.13350v1#S3.SS1), we use a simple 1D example to illustrate how the mini-batch couplings in data space and velocity space affect the velocity distribution and the generation of velocity samples. This motivates the development of HRF with mini-batch coupling. In [Section 3.2](https://arxiv.org/html/2507.13350v1#S3.SS2), we introduce the training of HRF2 with mini-batch coupled data points. In [Section 3.3](https://arxiv.org/html/2507.13350v1#S3.SS3), we explain how mini-batch coupling for velocity is achieved by leveraging a pre-trained model. In [Section 3.4](https://arxiv.org/html/2507.13350v1#S3.SS4), we introduce a two-stage approach that combines mini-batch data coupling and velocity coupling.

Figure: Figure 1: The generated velocity distributions at $(x_{t},t)=(-1,0)$ for the dataset 1D $\mathcal{N}\to 2\mathcal{N}$ using HRF2, (a) without couplings, (b)-(c) with data coupling (batch sizes: 5 and 100), (d) with velocity coupling (batch size: 100), (e)-(f) with velocity coupling (batch size: 100) and data coupling (batch sizes: 5 and 100). Data coupling simplifies the velocity distribution (cf. (a)-(c)), while velocity coupling reduces the number of sampling steps.
Refer to caption: x1.png

### 3.1 Velocity Distribution

HRF was designed with independently sampled $x_{0}$ and $x_{1}$. However, in this paper, we first show that the underlying theory can be generalized to an arbitrary joint distribution over $x_{0}$ and $x_{1}$, i.e., $\gamma(x_{0},x_{1})$, which has the correct marginal distributions, i.e.,

$$ $\int\gamma(x_{0},x_{1})dx_{1}=\rho_{0}(x_{0})\,\text{and}\,\int\gamma(x_{0},x_ {1})dx_{0}=\rho_{1}(x_{1}).$ (8) $$

The following theorem characterizes the distribution of the velocity at a specific space-time location $(x_{t},t)$ if an arbitrary joint distribution $\gamma$ is used instead of a product of two independent distributions.

###### Theorem 3.1 .

The velocity distribution $\pi_{1}(v;x_{t},t)$ at the space-time location $(x_{t},t)$ induced by the linear interpolation in [Equation 2](https://arxiv.org/html/2507.13350v1#S2.E2) for $(x_{0},x_{1})\sim\gamma(x_{0},x_{1})$ is

$$ $\displaystyle\pi_{1}(v;x_{t},t)=\frac{\gamma(x_{t}-tv,x_{t}+(1-t)v)}{\rho_{t}( x_{t})},$ (9) $$

where

$$ $\rho_{t}(x_{t})=\int\gamma\left(x_{0},\frac{x_{t}-(1-t)x_{0}}{t}\right)dx_{0}= \int\gamma\left(\frac{x_{t}-tx_{1}}{1-t},x_{1}\right)dx_{1},$ (10) $$

and $\rho_{t}(x_{t})\neq 0$.
The distribution $\pi_{1}(v;x_{t},t)$ is undefined if $\rho_{t}(x_{t})=0$.

The proof of [Theorem 3.1](https://arxiv.org/html/2507.13350v1#S3.Thmtheorem1) is deferred to [Appendix A](https://arxiv.org/html/2507.13350v1#A1).

For vanilla HRF2, the source and target distributions are independent, meaning $\gamma(x_{0},x_{1})=\rho_{0}(x_{0})\rho_{1}(x_{1})$. Consequently, as derived by Zhang et al. (2025) for a special case of [Theorem 3.1](https://arxiv.org/html/2507.13350v1#S3.Thmtheorem1), at time $t=0$, the velocity distribution becomes $\pi_{1}(v;x_{t},t)=\rho_{1}(x_{t}+v)$, making it a shifted version of the data distribution. Hence, learning this distribution is as challenging as directly modeling the data distribution.

To progressively simplify the velocity distribution across the levels of the HRF hierarchy, we introduce data coupling and velocity coupling. This makes the distributions easier to learn and improves overall model performance. Notably, coupling at one hierarchy level simplifies the distributions at all lower levels, thereby facilitating the matching process at the current level. Importantly, the complexity of the learned distribution can be controlled by adjusting the batch size in the mini-batch coupling process.

To illustrate the aforementioned distribution simplification we provide an example with 1D data. In this example, the source distribution is a standard Gaussian, while the target distribution is a mixture of two Gaussians with means located at $-1$ and $1$. As shown in [Figure 1](https://arxiv.org/html/2507.13350v1#S3.F1)(a-c), after applying data coupling (depth 1), the velocity distribution (depth 2) collapses into a single-mode Gaussian as the coupling batch size (bs) increases, effectively simplifying the velocity layer’s distribution. The number given in the legend refers to the number of used velocity ODE integration steps.

From [Figure 1](https://arxiv.org/html/2507.13350v1#S3.F1)(d), we observe that velocity coupling on its own does not simplify the velocity distribution. Instead, it simplifies the distribution at the next level (acceleration) and straightens the paths for velocity samples, reducing the number of integration steps needed to model the velocity distribution. [Figure 1](https://arxiv.org/html/2507.13350v1#S3.F1)(e,f) demonstrates that data coupling and velocity coupling are not mutually exclusive. They can be applied simultaneously to complement each other.

Finally, comparing the columns in [Figure 1](https://arxiv.org/html/2507.13350v1#S3.F1) shows that a larger batch size for mini-batch coupling results in a uni-modal distribution. Therefore, in practice, it is important to choose an appropriate batch size to control the multi-modality at each level. Next, we detail how data coupling and velocity coupling can be achieved.

###### Theorem 3.1 .

### 3.2 HRF2 with Data Coupling

To simplify the velocity distribution by reducing its multi-modality, it is crucial to first understand its root cause. During training, if the source data $x_{0}$ and target data $x_{1}$ are sampled independently, the multi-modality inherent in the data is preserved in the velocity distribution at $t=0$. We discussed this in the sequel of [Theorem 3.1](https://arxiv.org/html/2507.13350v1#S3.Thmtheorem1). Therefore, breaking this independence is key to simplifying the velocity distribution.
We find that data couplings that restrict flexibility, e.g., mini-batch OT, provide an opportunity to do this. Intuitively, using mini-batch OT results in a coupling of source and target data that isn’t entirely arbitrary, which inherently simplifies the velocity distribution.

Figure: Algorithm 1 HRF2 with Data Coupling

Following Tong et al. (2024); Pooladian et al. (2023), we apply mini-batch OT on the data used for HRF2 training. Let $\{x_{0}^{(i)}\}_{i=1}^{B}\sim\rho_{0}$ and $\{x_{1}^{(i)}\}_{i=1}^{B}\sim\mathcal{D}$. The OT problem in [Equation 7](https://arxiv.org/html/2507.13350v1#S2.E7) can be solved exactly and efficiently on a small batch size using standard solvers, e.g., POT (Flamary et al., 2021). The resulting coupling from the algorithm gives us a permutation matrix that pairs $x_{0}^{(i)}$ with $x_{1}^{(\sigma(i))}$ for $i\in\{1,\dots,B\}$.
Instead of sampling $x_{0}$ and $x_{1}$ independently from $\rho_{0}$ and the dataset $\mathcal{D}$ during training, we jointly sample $(x_{0},x_{1})$ from the joint distribution $\gamma(x_{0},x_{1})$ characterized by the mini-batch OT result. Using these samples, the training objective reads as follows:

$$ $\min_{\theta}\mathbb{E}_{(x_{0},x_{1})\sim\gamma,t\sim U[0,1],v_{0}\sim\pi_{0} ,\tau\sim U[0,1]}\left[\|(x_{1}-x_{0}-v_{0})-a_{\theta}(x_{t},t,v_{\tau},\tau) \|_{2}^{2}\right].$ (11) $$

The optimization procedure is detailed in [Algorithm 1](https://arxiv.org/html/2507.13350v1#algorithm1).

### 3.3 HRF2 with Velocity Coupling

Similar to data coupling, velocity coupling also aims to eliminate the independence between $v_{0}$ and $v_{1}(x_{t},t)$. With mini-batch coupled velocity samples that are drawn from an underlying joint distribution $\kappa_{x_{t},t}(v_{0},v_{1}(x_{t},t))$, the corresponding objective function is defined as follows:

$$ $\min_{\theta}\mathbb{E}_{x_{0}\sim\rho_{0},x_{1}\sim\mathcal{D},(v_{0},v_{1}) \sim\kappa_{x_{t},t},t\sim U[0,1],\tau\sim U[0,1]}\left[\|(v_{1}(x_{t},t)-v_{0 })-a_{\theta}(x_{t},t,v_{\tau},\tau)\|_{2}^{2}\right].$ (12) $$

[Algorithm 2](https://arxiv.org/html/2507.13350v1#algorithm2) summarizes the optimization procedure,
for which we study the following coupling.

Figure: Algorithm 3 Velocity Coupling via Mini-Batch OT

Velocity Coupling via Batch OT.
Different from the data coupling, where the target data samples are readily available, obtaining velocity samples for a fixed $(x_{t},t)$ requires simulation. Note that the target velocity samples $v_{1}(x_{t},t)\sim\pi_{1}(v;x_{t},t)$. To correctly couple $v_{0}$ and $v_{1}(x_{t},t)$, it is essential to fix the space-time location $(x_{t},t)$, ensuring that the samples $v_{1}(x_{t},t)$ are drawn from the same velocity distribution. Thus, the first step of velocity coupling is to obtain a batch of $v_{1}(x_{t},t)$ at a fixed $(x_{t},t)$. We achieve this by using a pre-trained HRF2 model $a_{\theta}(x_{t},t,v_{\tau},\tau)$, which transport $v_{0}$ to $v_{1}(x_{t},t)$ according to [Equation 4](https://arxiv.org/html/2507.13350v1#S2.E4). We then use the 2-Wasserstein optimal transport to sample coupled pairs $(v_{0},v_{1}(x_{t},t))$. This is detailed in [Algorithm 3](https://arxiv.org/html/2507.13350v1#algorithm3).

### 3.4 HRF2 with Hierarchical Data & Velocity Couplings

As shown before, data and velocity coupling complement each other.
To use both couplings we need $(x_{0},x_{1})\sim\gamma$ and $(v_{0},v_{1}(x_{t},t))\sim\kappa_{x_{t},t}$. To achieve this, we apply a two-stage training. First, we use data coupling to train the model $a_{\theta}$ according to [Algorithm 1](https://arxiv.org/html/2507.13350v1#algorithm1). In the second stage, using this pre-trained model, we generate paired samples $(v_{0},v_{1}(x_{t},t))$ according to [Algorithm 3](https://arxiv.org/html/2507.13350v1#algorithm3). We then train our network using the following objective:

$$ $\min_{\theta}\mathbb{E}_{(x_{0},x_{1})\sim\gamma,(v_{0},v_{1})\sim\kappa_{x_{t },t},t\sim U[0,1],\tau\sim U[0,1]}\left[\|(v_{1}(x_{t},t)-v_{0})-a_{\theta}(x_ {t},t,v_{\tau},\tau)\|_{2}^{2}\right].$ (13) $$

Importantly, note that the coupling of the velocities depends on the coupling of the data.

### 3.5 Marginal Preserving Property

The consistency of the velocity distribution with mini-batch velocity coupling directly follows prior works that use mini-batch coupling and reflow for data generation (Tong et al., 2024; Pooladian et al., 2023).

In addition, we can prove that the generation process according to [Equation 6](https://arxiv.org/html/2507.13350v1#S2.E6) with trained $a_{\theta}$ using mini-batch data coupling preserves the target data distribution and leads to correct marginals for all times $t\in[0,1]$.

###### Theorem 3.2 .

The time-differentiable stochastic process $\bm{Z}=\{Z_{t}:t\in[0,1]\}$ generated by [Equation 6](https://arxiv.org/html/2507.13350v1#S2.E6) has the same marginal distribution as the time-differentiable stochastic process $\bm{X}=\{X_{t}:t\in[0,1]\}$ generated by the linear interpolation in [Equation 2](https://arxiv.org/html/2507.13350v1#S2.E2) with the joint distribution $\gamma$ induced by mini-batch coupling.

The proof of [Theorem 3.2](https://arxiv.org/html/2507.13350v1#S3.Thmtheorem2) is deferred to [Appendix B](https://arxiv.org/html/2507.13350v1#A2).

###### Theorem 3.2 .

## 4 Experiments

In this section, we explore how data coupling and velocity coupling influence the velocity distribution and assess whether simplifying it leads to performance improvements. For all experiments, we report the total neural function evaluations (NFEs), which represents the product of the number of integration steps across all HRF levels.

### 4.1 Synthetic Data

We conduct experiments on four synthetic datasets used by Zhang et al. (2025) to ensure a fair comparison with HRF2. These datasets include two 1D cases ($\mathcal{N}\to 2\mathcal{N}$, $\mathcal{N}\to 5\mathcal{N}$), and two 2D cases ($\mathcal{N}\to 6\mathcal{N}$ and $8\mathcal{N}\to$ moon-shaped data). We use the Wasserstein and sliced 2-Wasserstein distances to evaluate 1D and 2D experiments, respectively. A complete description of the model architecture, parameter settings, and training details is provided in [Section D.1](https://arxiv.org/html/2507.13350v1#A4.SS1).

Recall that in [Figure 1](https://arxiv.org/html/2507.13350v1#S3.F1), we have observed that for the 1D $\mathcal{N}\to 2\mathcal{N}$ dataset, data coupling simplifies the velocity distribution. Now, we extend this analysis to 2D datasets. We denote HRF2 with data coupling as HRF2-D and HRF2 with joint data and velocity coupling as HRF2-D&V. As shown in [Figure 2](https://arxiv.org/html/2507.13350v1#S4.F2), the 2D results corroborate the 1D findings. For the original HRF2, the velocity distribution at $t=0$ is simply a shifted version of the data distribution. After applying data coupling, the velocity distribution at a given space-time location $(x_{t},t)$ becomes more unimodal, effectively aligning with a portion of the target distribution. For example, in [Figure 2](https://arxiv.org/html/2507.13350v1#S4.F2)(d), we observe that the velocity distribution primarily consists of the region of the target distribution closest to $x_{t}$.

Figure: Figure 2: Results on $8\mathcal{N}\to$ moon dataset. Three rows are HRF2, HRF2 with data coupling, HRF2 with data & velocity coupling. (a) and (c) are trajectories (green) of sample particles flowing from source distribution (grey) to target distribution (blue) with total NFEs 1 and 100. (b) and (d) are velocity distributions at the center of the bottom left Gaussian mode at $t=0$. Data coupling simplifies the velocity distribution and velocity coupling reduces the required sampling steps.
Refer to caption: x7.png

Figure: Figure 3: Results on synthetic datasets: (a) 1D $\mathcal{N}\to 2\mathcal{N}$ (b) 1D $\mathcal{N}\to 5\mathcal{N}$ (c) 2D $\mathcal{N}\to 6\mathcal{N}$ (d) 2D $8\mathcal{N}\to$ moon. Top row: HRF2-D&V generated data distributions. Bottom row: performance vs. total NFEs. We use Wasserstein and sliced 2-Wasserstein distances for 1D and 2D data, respectively.
Refer to caption: x19.png

Figure: Figure 4: Results on MNIST, CIFAR-10 and CelebA-HQ 256 datasets. (a) HRF2-D with total NFE=1. (b) HRF2-D with total NFE=500. (c) HRF2-D&V with total NFE=1. Here we report the results with HRF2-D&V-OT. (d) FID scores with respect to total NFEs. With joint data coupling and velocity coupling, HRF2-D&V can generate reasonably good results with only 1 step.
Refer to caption: x27.png

In contrast, velocity coupling does not modify the velocity distribution itself but significantly reduces the number of required sampling steps. As shown in [Figure 2](https://arxiv.org/html/2507.13350v1#S4.F2)(b), a single integration step already produces a reasonable velocity distribution with joint data and velocity couplings. The results in [Figure 3](https://arxiv.org/html/2507.13350v1#S4.F3) demonstrate that joint data and velocity couplings effectively enhance the model performance, particularly when NFEs are low. Additional results on synthetic data are provided in [Section C.1](https://arxiv.org/html/2507.13350v1#A3.SS1).

### 4.2 Image Data

For high-dimensional image data, we conduct experiments on MNIST (LeCun et al., 1998), CIFAR-10 (Krizhevsky, 2009), and CelebA-HQ 256 (Karras et al., 2018), using Fréchet Inception Distance (FID) as the evaluation metric. For MNIST and CIFAR-10, we directly operate in the pixel space, with input dimensions of $1\times 28\times 28$ and $3\times 32\times 32$, respectively. For CelebA-HQ, we first encode the original $3\times 256\times 256$ images into a $4\times 32\times 32$ latent space using a pretrained VAE encoder, and conduct training and inference in the latent space. The experiments on CelebA-HQ demonstrate that our methods scale well on higher-dimensional data.

We compare our method to RF (Liu, 2022), HRF2 (Zhang et al., 2025) and OT-CFM (Tong et al., 2024) baselines. HRF2 is the base pre-trained model for our HRF2-D and HRF2-D&V. OT-CFM is essentially equivalent to HRF1-D. As shown in [Figure 4](https://arxiv.org/html/2507.13350v1#S4.F4)(d), data coupling significantly improves performance across both low and high total NFEs. However, applying velocity coupling on top of data coupling only yields substantial improvements in the low-NFE regime. From [Figure 4](https://arxiv.org/html/2507.13350v1#S4.F4)(a)-(c), we observe that data coupling alone enhances performance at low NFEs, but HRF2-D still struggles in extreme cases. Notably, incorporating velocity coupling enables the model to generate compelling results even under the extreme condition of total NFE = 1. More results are presented in [Section C.2](https://arxiv.org/html/2507.13350v1#A3.SS2).

The details of the model architectures are deferred to [Section D.2](https://arxiv.org/html/2507.13350v1#A4.SS2).
Since the model architecture remains unchanged, HRF2-D and HRF2-D&V have the same memory usage and inference time as HRF2. For training, we apply data coupling and velocity coupling following [Algorithms 1](https://arxiv.org/html/2507.13350v1#algorithm1) and [2](https://arxiv.org/html/2507.13350v1#algorithm2). To obtain accurate velocity pairs $(v_{0},v_{1})$ for velocity coupling, more integration steps are required here than for synthetic data.
More details are provided in [Section D.2](https://arxiv.org/html/2507.13350v1#A4.SS2).

## 5 Related Work

Flow Matching:
Concurrently, Liu et al. (2023a); Lipman et al. (2023); Albergo and Vanden-Eijnden (2023) presented learning of the ODE velocity that governs the generation of new data through a time-differentiable stochastic process defined by interpolating between samples from the source and data distributions. This provides flexibility by enabling precise connections between any two densities over finite time intervals. Liu et al. (2023a) concentrated on a linear interpolation, which provides straight paths connecting points from the source and the target distributions. Lipman et al. (2023) introduced the interpolation through the lens of conditional probability paths leading to a Gaussian. Albergo and Vanden-Eijnden (2023); Albergo et al. (2023) introduced stochastic interpolants with more general forms. They all learn the expected velocity field, which leads to curved sampling paths for data generation. Flow matching has been extended to handle discrete data (Gat et al., 2024; Stark et al., 2024) and manifold data (Chen and Lipman, 2024).

Straightening Flows:
Liu et al. (2023a) proposed an iterative method called reflow, which connects points from the source and target distributions using a trained rectified flow model to smooth the transport path. They demonstrated that repeating this process results in an optimal transport map. However, in practice, errors in the learned velocity field can introduce bias. Other related studies address this by adjusting how noise and data are sampled during training, rather than using iterations. For instance, Pooladian et al. (2023); Tong et al. (2024) computed mini-batch optimal transport couplings between the source and data distributions to reduce transport costs and flow variance. Park et al. (2024) address the curved paths in flow matching by learning both initial velocity and acceleration, such that the sampling paths can cross. However, it requires a pre-trained diffusion model to acquire noise-data pairs. Cheng and Schwing (2025) study conditional data.

Distribution of flow fields:  Zhang et al. (2025) capture the distribution of the random flow fields induced by the linear interpolation of source and target data. The sampling process is governed by coupled ODEs, which allows sampling paths to cross. Guo and Schwing (2025) model the flow field distribution using a variational autoencoder.

Building upon work by Zhang et al. (2025), we show that modeling the velocity distributions after the mini-batch coupling in data space improves the performance of HRF2 and OT-CFM. Hierarchically coupling the data and velocity leads to significantly improved results at low NFEs.

## 6 Conclusion

We study ways to control the complexity of the multi-modal velocity distribution and their impact on capturing this distribution with hierarchical flow matching. We find hierarchical flow matching with mini-batch coupling in the data space consistently improves the generation quality compared to vanilla hierarchical rectified flow matching and vanilla flow matching with mini-batch optimal transport. Joint coupling in the data space and the velocity space leads to further improvements if few function evaluations are used.

Limitations and broader impacts: Our proposed methods offer faster and more accurate data generation. It can help advance scientific modeling and simulations, contributing to advances in areas like physics, healthcare, and drug discovery. For the limitations: the current velocity coupling approach requires simulated target velocity samples during training. Simulation-free velocity coupling is an interesting direction for future research.

Acknowledgments: Work supported in part by NSF grants 1934757, 2008387, 2045586, 2106825, MRI 1725729, NIFA award 2020-67021-32799, and the Alfred P. Sloan Foundation.

## References

- Albergo and Vanden-Eijnden [2023]
M. Albergo and E. Vanden-Eijnden.
Building normalizing flows with stochastic interpolants.
In *Proc. ICLR*, 2023.
- Albergo et al. [2023]
M. Albergo, N. Boffi, and E. Vanden-Eijnden.
Stochastic interpolants: A unifying framework for flows and diffusions.
*arXiv preprint arXiv:2303.08797*, 2023.
- Chen and Lipman [2024]
R. T. Chen and Y. Lipman.
Flow matching on general geometries.
In *ICLR*, 2024.
- Chen et al. [2024]
Y. Chen, M. Goldstein, M. Hua, M. S. Albergo, N. M. Boffi, and E. Vanden-Eijnden.
Probabilistic forecasting with stochastic interpolants and f$\backslash$" ollmer processes.
In *ICML*, 2024.
- 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 *Proc. ICCV*, 2025.
- Dao et al. [2023]
Q. Dao, H. Phung, B. Nguyen, and A. Tran.
Flow matching in latent space, 2023.
URL [https://arxiv.org/abs/2307.08698](https://arxiv.org/abs/2307.08698).
- Deshpande et al. [2018]
I. Deshpande, Z. Zhang, and A. G. Schwing.
Generative Modeling using the Sliced Wasserstein Distance.
In *Proc. CVPR*, 2018.
- Deshpande et al. [2019]
I. Deshpande, Y.-T. Hu, R. Sun, A. Pyrros, N. Siddiqui, S. Koyejo, Z. Zhao, D. Forsyth, and A. G. Schwing.
Max-Sliced Wasserstein Distance and its use for GANs.
In *Proc. CVPR*, 2019.
- 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, et al.
Scaling rectified flow transformers for high-resolution image synthesis.
In *Forty-first International Conference on Machine Learning*, 2024.
- Fatras et al. [2020]
K. Fatras, Y. Zine, R. Flamary, R. Gribonval, and N. Courty.
Learning with minibatch wasserstein: asymptotic and gradient properties.
In *AISTATS 2020-23nd International Conference on Artificial Intelligence and Statistics*, volume 108, pages 1–20, 2020.
- Fatras et al. [2021]
K. Fatras, Y. Zine, S. Majewski, R. Flamary, R. Gribonval, and N. Courty.
Minibatch optimal transport distances; analysis and applications.
*arXiv preprint arXiv:2101.01792*, 2021.
- 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, et al.
POT: Python optimal transport.
*Journal of Machine Learning Research*, 22(78):1–8, 2021.
- Gat et al. [2024]
I. Gat, T. Remez, N. Shaul, F. Kreuk, R. T. Chen, G. Synnaeve, Y. Adi, and Y. Lipman.
Discrete flow matching.
In *NeurIPS*, 2024.
- Guo and Schwing [2025]
P. Guo and A. G. Schwing.
Variational rectified flow matching.
*arXiv preprint arXiv:2502.09616*, 2025.
- Jing et al. [2024]
B. Jing, B. Berger, and T. Jaakkola.
Alphafold meets flow matching for generating protein ensembles.
*arXiv preprint arXiv:2402.04845*, 2024.
- Karras et al. [2018]
T. Karras, T. Aila, S. Laine, and J. Lehtinen.
Progressive growing of GANs for improved quality, stability, and variation.
In *International Conference on Learning Representations*, 2018.
URL [https://openreview.net/forum?id=Hk99zCeAb](https://openreview.net/forum?id=Hk99zCeAb).
- Krizhevsky [2009]
A. Krizhevsky.
Learning multiple layers of features from tiny images.
Technical report, University of Toronto, 2009.
- LeCun et al. [1998]
Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner.
Gradient-based learning applied to document recognition.
*Proceedings of the IEEE*, 1998.
- Lipman et al. [2023]
Y. Lipman, R. Chen, H. Ben-Hamu, M. Nickel, and M. Le.
Flow Matching for Generative Modeling.
In *Proc. ICLR*, 2023.
- Liu [2022]
Q. Liu.
Rectified flow: A marginal preserving approach to optimal transport.
*arXiv preprint arXiv:2209.14577*, 2022.
- Liu et al. [2023a]
X. Liu, C. Gong, and Q. Liu.
Flow straight and fast: Learning to generate and transfer data with rectified flow.
In *Proc. ICLR*, 2023a.
- Liu et al. [2023b]
X. Liu, X. Zhang, J. Ma, and J. Peng.
Instaflow: One step is enough for high-quality diffusion-based text-to-image generation.
In *ICLR*, 2023b.
- Park et al. [2024]
D. Park, S. Lee, S. Kim, T. Lee, Y. Hong, and H. J. Kim.
Constant acceleration flow.
In *NeurIPS*, 2024.
- Peebles and Xie [2023]
W. Peebles and S. Xie.
Scalable diffusion models with transformers, 2023.
URL [https://arxiv.org/abs/2212.09748](https://arxiv.org/abs/2212.09748).
- Peyré et al. [2019]
G. Peyré, M. Cuturi, et al.
Computational optimal transport: With applications to data science.
*Foundations and Trends® in Machine Learning*, 11(5-6):355–607, 2019.
- Pooladian et al. [2023]
A.-A. Pooladian, H. Ben-Hamu, C. Domingo-Enrich, B. Amos, Y. Lipman, and R. Chen.
Multisample flow matching: Straightening flows with minibatch couplings.
In *Proc. ICML*, 2023.
- Rombach et al. [2022]
R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer.
High-resolution image synthesis with latent diffusion models.
In *Proc. CVPR*, 2022.
- Stark et al. [2024]
H. Stark, B. Jing, C. Wang, G. Corso, B. Berger, R. Barzilay, and T. Jaakkola.
Dirichlet flow matching with applications to dna sequence design.
In *ICML*, 2024.
- Tong et al. [2024]
A. Tong, N. Malkin, G. Huguet, Y. Zhang, J. Rector-Brooks, K. Fatras, G. Wolf, and Y. Bengio.
Improving and generalizing flow-based generative models with minibatch optimal transport.
*TMLR*, 2024.
- Villani [2009]
C. Villani.
*Optimal transport: old and new*, volume 338.
Springer, 2009.
- Yim et al. [2023]
J. Yim, A. Campbell, A. Y. Foong, M. Gastegger, J. Jiménez-Luna, S. Lewis, V. G. Satorras, B. S. Veeling, R. Barzilay, T. Jaakkola, and F. Noe.
Fast protein backbone generation with SE(3) flow matching.
In *NeurIPS Workshop: Machine Learning in Structural Biology*, 2023.
- Zhang and Gienger [2024]
F. Zhang and M. Gienger.
Affordance-based robot manipulation with flow matching.
*arXiv preprint arXiv:2409.01083*, 2024.
- Zhang et al. [2024]
X. Zhang, Y. Pu, Y. Kawamura, A. Loza, Y. Bengio, D. L. Shung, and A. Tong.
Trajectory flow matching with applications to clinical time series modeling.
In *NeurIPS*, 2024.
- Zhang et al. [2025]
Y. Zhang, Y. Yan, A. Schwing, and Z. Zhao.
Towards hierarchical rectified flow.
In *The Thirteenth International Conference on Learning Representations*, 2025.

## Appendix: Hierarchical Rectified Flow Matching with Mini-Batch Couplings

This appendix is structured as follows:
in [Appendix A](https://arxiv.org/html/2507.13350v1#A1) we provide the proof of [Theorem 3.1](https://arxiv.org/html/2507.13350v1#S3.Thmtheorem1); in [Appendix B](https://arxiv.org/html/2507.13350v1#A2) we provide the proof of [Theorem 3.2](https://arxiv.org/html/2507.13350v1#S3.Thmtheorem2);
in [Appendix C](https://arxiv.org/html/2507.13350v1#A3) we provide additional results for both synthetic and image data; and in [Appendix D](https://arxiv.org/html/2507.13350v1#A4) we discuss implementation details.

## Appendix A Proof of Theorem 3.1

Proof of [Theorem 3.1](https://arxiv.org/html/2507.13350v1#S3.Thmtheorem1):
For simplicity, we show the proof for 1D random variables $x_{0}$ and $x_{1}$ drawn from a joint distribution $\gamma(x_{0},x_{1})$. The joint distribution of $v(x_{t},t)$ and $x_{t}$ is $\pi_{1}(v;x_{t},t)\rho_{t}(x_{t})$, since $\pi_{1}(v;x_{t},t)$ corresponds to the conditional distribution of the velocity given location $x_{t}$. According to the linear interpolation in [Equation 2](https://arxiv.org/html/2507.13350v1#S2.E2), we have,

$$ $\displaystyle\begin{bmatrix}v(x_{t},t)\\ x_{t}\end{bmatrix}=\begin{bmatrix}1&-1\\ t&(1-t)\end{bmatrix}\begin{bmatrix}x_{1}\\ x_{0}\end{bmatrix}=A\begin{bmatrix}x_{1}\\ x_{0}\end{bmatrix},$ (14) $$

where the matrix $A$ has determinant 1. Since $[v(x_{t},t),x_{t}]^{T}$ is a linear transformation of $[x_{1},x_{0}]^{T}$, we have the following expression for the joint distribution of $v(x_{t},t)$ and $x_{t}$,

$$ $\pi_{1}(v;x_{t},t)\rho_{t}(x_{t})=\frac{1}{\det(A)}\gamma\left(A^{-1}\begin{ bmatrix}v(x_{t},t)\\ x_{t}\end{bmatrix}\right)=\gamma(x_{t}-tv,x_{t}+(1-t)v).$ (15) $$

After rearranging we get $\pi_{1}(v;x_{t},t)=\frac{\gamma(x_{t}-tv,x_{t}+(1-t)v)}{\rho_{t}(x_{t})}$. For the higher dimensional case, the relation in [Equation 15](https://arxiv.org/html/2507.13350v1#A1.E15) still holds. This completes the proof.
$\blacksquare$

## Appendix B Proof of Theorem 3.2

Proof of [Theorem 3.2](https://arxiv.org/html/2507.13350v1#S3.Thmtheorem2):
We consider the characteristic function of $Z_{t+\Delta t}=Z_{t}+V\Delta t$ for $t\in[0,1]$ and $\Delta t\in[0,1-t]$, assuming that $Z_{t}$ has the same distribution as $X_{t}$. If the characteristic functions of $Z_{t+\Delta t}$ and $X_{t+\Delta t}$ agree, then $Z_{t+\Delta t}$ and $X_{t+\Delta t}$ have the same distribution.

To show this, we evaluate the characteristic function of $Z_{t+\Delta t}$,

$$ $\displaystyle\mathbb{E}\left[e^{\imath\langle k,Z_{t+\Delta t}\rangle}\right]$ $\displaystyle=\mathbb{E}_{\rho_{t},\pi_{1}}\left[e^{\imath\langle k,X_{t}+V \Delta t\rangle}\right]$ $\displaystyle=\int\int e^{\imath\langle k,x_{t}+v\Delta t\rangle}\pi_{1}(v;x_{ t},t)\rho_{t}(x_{t})dvdx_{t}$ $\displaystyle\stackrel{{\scriptstyle a}}{{=}}\int\int e^{\imath\langle k,x_{t} +v\Delta t\rangle}\frac{\gamma(x_{t}-vt,x_{t}+(1-t)v)}{\rho_{t}(x_{t})}\rho_{t }(x_{t})dvdx_{t}$ $\displaystyle=\int\int e^{\imath\langle k,(x_{t}+v\Delta t)\rangle}\gamma(x_{t }-tv,x_{t}+(1-t)v)dvdx_{t}$ $\displaystyle\stackrel{{\scriptstyle b}}{{=}}\int\int e^{\imath\langle k,(1-t- \Delta t)x_{0}+(t+\Delta t)x_{1}\rangle}\gamma(x_{0},x_{1})dx_{0}dx_{1}$ $\displaystyle=\mathbb{E}_{\rho_{t+\Delta t}}\left[e^{\imath\langle k,X_{t+ \Delta t}\rangle}\right].$ (16) $$

We use the notation $\langle\cdot,\cdot\rangle$ to denote the inner product. The equality $a$ is valid due to [Theorem 3.2](https://arxiv.org/html/2507.13350v1#S3.Thmtheorem2). The equality $b$ holds because $x_{0}=x_{t}-tv$ and $x_{1}=x_{t}+(1-t)v$ with the linear interpolation. Therefore, we find that $Z_{t+\Delta t}$ and $X_{t+\Delta t}$ follow the same distribution. In addition, since $Z_{0}$ and $X_{0}$ follow the same distribution $\rho_{0}$, we can conclude that $Z_{t}$ and $X_{t}$ follow the same marginal distribution at $t$ for $t\in[0,1]$. This completes the proof. $\blacksquare$

## Appendix C Additional Experimental Results

### C.1 Synthetic Data Results

We present more results on synthetic data: [Figure 5](https://arxiv.org/html/2507.13350v1#A3.F5) for 1D $\mathcal{N}\to 2\mathcal{N}$ data, [Figure 6](https://arxiv.org/html/2507.13350v1#A3.F6) for 1D $\mathcal{N}\to 5\mathcal{N}$ data, and [Figure 7](https://arxiv.org/html/2507.13350v1#A3.F7) for 2D $\mathcal{N}\to 6\mathcal{N}$ data. Across all these experiments, we consistently observe that data coupling simplifies the velocity distribution, while velocity coupling significantly reduces the required sampling steps.

Figure: Figure 5: Results on 1D $\mathcal{N}\to 2\mathcal{N}$ data. The three rows correspond to HRF2, HRF2 with data coupling, HRF2 with data & velocity coupling. (a) and (b) are generated data distribution with total NFEs 1 and 100. (c) is velocity distribution at $(x_{t},t)=(-1,0)$.
Refer to caption: x39.png

Figure: Figure 6: Results on 1D $\mathcal{N}\to 5\mathcal{N}$ data. The three rows correspond to HRF2, HRF2 with data coupling, HRF2 with data & velocity coupling. (a) and (b) are generated data distribution with total NFEs 1 and 100. (c) is velocity distribution at $(x_{t},t)=(0,0)$.
Refer to caption: x48.png

Figure: Figure 7: Results on 2D $\mathcal{N}\to 6\mathcal{N}$ data. The three rows correspond to HRF2, HRF2 with data coupling, HRF2 with data & velocity coupling. (a) and (c) are trajectories (green) of sample particles flowing from source distribution (grey) to target distribution (blue) with total NFEs 1 and 100. (b) and (d) are velocity distributions at $(0,1)$ at $t=0$.
Refer to caption: x57.png

### C.2 Image Data Results

We present more results on image data: [Table 1](https://arxiv.org/html/2507.13350v1#A3.T1) for MNIST, [Table 2](https://arxiv.org/html/2507.13350v1#A3.T2) for CIFAR-10, and [Tables 3](https://arxiv.org/html/2507.13350v1#A3.T3) and [8](https://arxiv.org/html/2507.13350v1#A3.F8) for CelebA-HQ 256. Again, we consistently observe that data coupling enhances sampling quality for both low and high total NFEs, but collapses when total NFE is reduced to 1, while velocity coupling produces high-quality samples even under this extreme case.

**Table 1: FID performance on MNIST under different total NFE settings. Bold for the best.**
| Total NFEs | RF (1.08M) | OT-CFM (1.08M) | HRF2 (1.07M) | HRF2-D (1.07M) | HRF2-D&V (1.07M) |
| --- | --- | --- | --- | --- | --- |
| 5 | 19.187 ± 0.188 | 13.977 ± 0.166 | 15.798 ± 0.151 | 10.167 ± 0.136 | 5.519 ± 0.112 |
| 10 | 7.974 ± 0.119 | 4.477 ± 0.099 | 6.644 ± 0.076 | 3.823 ± 0.038 | 3.861 ± 0.089 |
| 20 | 6.151 ± 0.090 | 2.763 ± 0.036 | 3.408 ± 0.076 | 2.318 ± 0.053 | 3.720 ± 0.045 |
| 50 | 5.605 ± 0.057 | 2.321 ± 0.038 | 2.664 ± 0.058 | 1.929 ± 0.031 | 3.604 ± 0.016 |
| 100 | 5.563 ± 0.049 | 2.346 ± 0.023 | 2.588 ± 0.075 | 1.847 ± 0.011 | 3.423 ± 0.003 |
| 500 | 5.453 ± 0.047 | 2.296 ± 0.007 | 2.574 ± 0.121 | 1.913 ± 0.043 | 3.546 ± 0.107 |

**Table 2: FID performance on CIFAR-10 under different total NFE settings. Bold for the best.**
| Total NFEs | RF (35.75M) | OT-CFM (35.75M) | HRF2 (44.81M) | HRF2-D (44.81M) | HRF2-D&V (44.81M) |
| --- | --- | --- | --- | --- | --- |
| 5 | 36.209 ± 0.142 | 23.111 ± 0.010 | 30.884 ± 0.104 | 22.817 ± 0.072 | 6.315 ± 0.057 |
| 10 | 14.113 ± 0.092 | 12.564 ± 0.016 | 12.065 ± 0.024 | 10.969 ± 0.025 | 5.739 ± 0.017 |
| 20 | 8.355 ± 0.065 | 8.553 ± 0.002 | 7.129 ± 0.027 | 6.860 ± 0.022 | 5.332 ± 0.009 |
| 50 | 5.514 ± 0.034 | 5.911 ± 0.005 | 4.847 ± 0.028 | 4.739 ± 0.006 | 5.142 ± 0.024 |
| 100 | 4.588 ± 0.013 | 4.952 ± 0.012 | 4.334 ± 0.054 | 4.301 ± 0.022 | 5.078 ± 0.044 |
| 500 | 3.887 ± 0.035 | 4.184 ± 0.086 | 3.706 ± 0.043 | 3.578 ± 0.028 | 5.095 ± 0.032 |

**Table 3: FID performance on CelebA-HQ 256 under different total NFE settings. Bold for the best.**
| Total NFEs | RF (457.06M) | OT-CFM (457.06M) | HRF2 (616.20M) | HRF2-D (616.20M) | HRF2-D&V (616.20M) |
| --- | --- | --- | --- | --- | --- |
| 5 | 38.641 ± 0.126 | 29.646 ± 0.093 | 34.246 ± 0.107 | 32.918 ± 0.085 | 7.612 ± 0.015 |
| 10 | 16.876 ± 0.088 | 12.879 ± 0.083 | 15.391 ± 0.074 | 13.424 ± 0.022 | 6.931 ± 0.038 |
| 20 | 10.027 ± 0.060 | 7.426 ± 0.042 | 9.291 ± 0.042 | 7.048 ± 0.033 | 6.560 ± 0.039 |
| 50 | 7.395 ± 0.021 | 5.545 ± 0.023 | 6.927 ± 0.041 | 5.529 ± 0.021 | 6.330 ± 0.022 |
| 100 | 6.850 ± 0.064 | 5.236 ± 0.034 | 6.450 ± 0.062 | 4.961 ± 0.019 | 6.248 ± 0.023 |
| 500 | 6.418 ± 0.026 | 5.094 ± 0.019 | 6.188 ± 0.056 | 4.624 ± 0.029 | 6.225 ± 0.015 |

Figure: Figure 8: Examples of the generated images for CelebA-HQ starting from the same noise for all models. The 5 rows from top to bottom correspond to total NFEs = 1, 5, 10, 50, 500.
Refer to caption: x69.png

## Appendix D Implementation Details

### D.1 Synthetic Data

For synthetic data experiments, we employ a neural network architecture with two distinct stages. The first stage separately encodes spatial and temporal inputs with linear layers and Sinusoidal Positional Embeddings. The second stage concatenates the processed features and refines them through multiple linear layers to produce the final output. The model consists of 304,513 parameters, totaling 0.30M in size.

For data coupling, we train the model from scratch following [Algorithm 1](https://arxiv.org/html/2507.13350v1#algorithm1) strictly. A key consideration is the choice of batch sizes, as two different batch sizes are involved – one for batch OT and another for training. In 1D and 2D experiments, a large batch size is necessary for stable training, but using an excessively large batch size for batch OT is computationally inefficient. To address this, we set the batch size for batch OT to 100 while using a batch size of 1,000 for gradient computation. This means that in each training iteration, we perform batch OT on 100 data points 10 times to accumulate a full batch for gradient updates.

For velocity coupling, we use the HRF2-D model from the previous step as the base model to generate $(v_{0},v_{1})$ pairs at a fixed space-time location $(x_{t},t)$, following [Algorithm 2](https://arxiv.org/html/2507.13350v1#algorithm2). During training, we observed that the performance depends on the quality of the base model. To mitigate this, we save multiple checkpoints of HRF2-D and select the best-performing checkpoint via a validation dataset as the base model for velocity coupling.

Computational requirements during training are shown in [Table 4](https://arxiv.org/html/2507.13350v1#A4.T4). In the low-dimensional setting, batch OT becomes more time-consuming than the training itself. As a result, HRF2-D trains significantly slower than HRF2. In contrast, HRF2-D&V uses precomputed velocity pairs and therefore does not require batch OT during training. Moreover, it operates with a smaller batch size (1000) than HRF2 (5000), resulting in lower memory usage and faster training.

During the evaluation, we select the best checkpoint from a validation set for all models (HRF2, HRF2-D, HRF2-D&V). For each seed, we conduct the experiment three times, yielding three best models per seed. Each model is then evaluated three times, resulting in nine experimental results per seed. Finally, we report the mean and standard deviation over three different seeds, totaling 27 experimental results.

**Table 4: Computational requirements during training on synthetic datasets.**
| Training | 1D data | 2D data |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | HRF2 (0.30M) | HRF2-D (0.30M) | HRF2-D&V (0.30M) | HRF2 (0.32M) | HRF2-D (0.32M) | HRF2-D&V (0.32M) |
| Time (s/iter) | 0.0028 | 0.0581 | 0.0025 | 0.0029 | 0.0588 | 0.0027 |
| Memory (MB) | 658 | 658 | 566 | 660 | 660 | 568 |
| Param. Counts | 304,513 | 304,513 | 304,513 | 321,154 | 321,154 | 321,154 |

### D.2 Image Data

We adapt and modify the model architectures based on Zhang et al. [2025] for MNIST and CIFAR-10 datasets and Dao et al. [2023] for the CelebA-HQ dataset.

MNIST. For MNIST, we use the standard UNet. The ResNet blocks in the UNet function similarly to the model used for synthetic data. They process spatial and temporal inputs separately using convolutional and linear layers, respectively. The processed features are then concatenated and passed through a series of linear layers to capture space-time dependencies. The model consists of 1.07M parameters.

CIFAR-10. For CIFAR-10, the model consists of two UNets: a large UNet for processing $v_{\tau}$ and $\tau$, and a smaller UNet (one-fourth the size) for processing $x_{t}$ and $t$. The outputs of each ResNet block in the smaller UNet are input to the corresponding ResNet blocks in the larger UNet, facilitating information exchange between different scales. The model consists of 44.81M parameters.

CelebA-HQ. For CelebA-HQ, we first encode images into a latent space using the pretrained VAE encoder from Stable Diffusion [Rombach et al., 2022]. We then use DiT [Peebles and Xie, 2023] as the backbone to process $v_{\tau}$ in this latent space. To condition the velocity prediction on $x_{t}$, we inject $x_{t}$ into each DiT block via cross-attention layers, while keeping the main DiT architecture unchanged. The time embedding is also modified by replacing $\text{embedding}(t)$ with $\text{embedding}(t)+\text{embedding}(\tau)$ to incorporate time information in both time axes.

For training RF and OT-CFM on MNIST and CIFAR-10, we follow the procedures and hyperparameter settings from Tong et al. [2024] and Lipman et al. [2023]. For HRF2 on the same datasets, we adopt the training setup from Zhang et al. [2025]. For all models on CelebA-HQ, we follow the procedures and hyperparameters from Dao et al. [2023].

For data coupling, we train the model from scratch following [Algorithm 1](https://arxiv.org/html/2507.13350v1#algorithm1). Both the batch OT and training batch sizes are set to 128 for MNIST and CIFAR-10 and 256 for CelebA-HQ.

For velocity coupling, we start from the HRF2-D model obtained in the previous step. Following the synthetic data experiments, we select the best-performing model on the validation dataset to ensure training quality. The training speed for velocity coupling is primarily limited by the velocity sample generation. Therefore, we generate velocity pairs before training and perform the training offline.

We train the UNet for MNIST and CIFAR-10 on 1 NVIDIA L40S GPU and the DiT for CelebA-HQ on 8 NVIDIA L40S GPUs. Computational requirements, including training time and memory usage, are shown in [Tables 5](https://arxiv.org/html/2507.13350v1#A4.T5), [6](https://arxiv.org/html/2507.13350v1#A4.T6) and [7](https://arxiv.org/html/2507.13350v1#A4.T7).

**Table 5: Computational cost and model size for different methods on MNIST.**
| MNIST | RF (1.08M) | OT-CFM (1.08M) | HRF2 (1.07M) | HRF2-D (1.07M) | HRF2-D&V (1.07M) |
| --- | --- | --- | --- | --- | --- |
| Time (s/iter) | 0.045 | 0.046 | 0.046 | 0.046 | 0.046 |
| Memory (MB) | 2546 | 2546 | 2546 | 2546 | 2546 |
| Param. Counts | 1,075,361 | 1,075,361 | 1,065,698 | 1,065,698 | 1,065,698 |

**Table 6: Computational cost and model size for different methods on CIFAR-10.**
| CIFAR-10 | RF (35.75M) | OT-CFM (35.75M) | HRF2 (44.81M) | HRF2-D (44.81M) | HRF2-D&V (44.81M) |
| --- | --- | --- | --- | --- | --- |
| Time (s/iter) | 0.166 | 0.169 | 0.196 | 0.202 | 0.200 |
| Memory (MB) | 7480 | 7480 | 9220 | 9220 | 9220 |
| Param. Counts | 35,746,307 | 35,746,307 | 44,807,843 | 44,807,843 | 44,807,843 |

**Table 7: Computational cost and model size for different methods on CelebA-HQ.**
| CelebA-HQ 256 | RF (457.06M) | OT-CFM (457.06M) | HRF2 (616.20M) | HRF2-D (616.20M) | HRF2-D&V (616.20M) |
| --- | --- | --- | --- | --- | --- |
| Time (s/iter) | 0.418 | 0.420 | 0.688 | 0.688 | 0.688 |
| Memory per GPU (GB) | 26.65 | 26.65 | 39.89 | 39.89 | 36.79 |
| Param. Counts | 457,062,416 | 457,062,416 | 616,197,186 | 616,197,186 | 616,197,186 |

For each model, we conduct five evaluation runs, and report the means and standard deviations.

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 2
presentation: 2
questions: 1. In Fig. 1(a), after 100 steps of ODE steps, why does the velocity distribution not reach a univariate Gaussian (as shown in Fig. 2 in the original HRF paper, Zhang et al. 2025)?

2. Please enlarge the legends in the 2d sample plots, Fig. 3, top panel (two plots on the right).

3. Please report the overall performance of all the competing approaches at the final Euler step for the image benchmarks.

4. Can the author comment on the use of adaptive solvers instead of fixed-step Euler?
soundness: 3
strengths: I like the presentation of the approach. The use of mini-batch OT in the two levels of hierarchical rectified flow is intuitive and well-motivated from experiments on bi-modal Gaussian data. The theoretical results are interesting and easy to follow. The experiments do show improvement in lower FID compared to baselines with low NFEs.
summary: The paper extends the hierarchical rectified flow matching models by incorporating mini-batch optimal transport couplings between the source and target samples for the data and velocity. It is shown through theory and empirical evidence that the mini-batch coupling (for a large batch size) simplifies the distribution of velocity close to t=0 (closer to the source distribution) for rectified flow models. This observation motivated the authors to apply mini-batch OT in the two-level hierarchical rectified flow model (HRF2). The paper shows improvement in FID with a few NFEs for popular image benchmarks and synthetic datasets.
weaknesses: Major:

Novelty - The paper's main contribution is to use mini-batch OT for training the HRF2 model. The mini-batch OT has been previously used for CFM models (OT-CFM). I understand that this paper extends this to HRF models, but in my view, this is a very limited novelty.

Training cost - Flow matching training with the mini-batch OT is expensive. This method adds computing another OT map for the velocity distribution. In addition, HRF2 with velocity coupling requires a simulation from ODE flow for the velocity distribution. I would request the authors to add the training cost analysis of their method.

Other:

Suboptimal baselines - The OT-CFM results on CIFAR10 seem to have FID > 5 for NFE 100. However, Tong et al. (OT-CFM), Guo and Schwing (VRFM), and Samaddar et al. (Latent-CFM) all reported FID < 5 for NFE 100. This suggests that the baseline training or inference was suboptimal.

---

confidence: 4
contribution: 1
initial_rating: 4
presentation: 3
questions: •	As shown in Figure 3(d), HRF2-D and HRF2-D&V performs nearly the same, with HRF2-D&V slightly better. In contrast, HRF2-D&V performs worse than HRF2-D when NFE is large on the real image datasets. How to choose between “data coupling” and “joint data and velocity coupling” given a certain task and a certain NFE? 

•	Could the authors provide a theoretical validation of “minibatch OT coupling inherently simplifies the velocity distribution on the next hierarchy” on a simple distribution like Gaussian distribution?
soundness: 2
strengths: •	The paper is well-written and clearly-organized.

•	The empirical discovery that using mini-batch OT coupling in the current hierarchy level simplifies the distribution at the next level is interesting and could be considered in other scenarios of generative modeling like diffusion models.
summary: This paper mainly focuses on gradually simplifying the complexity of the distributions across hierarchy in hierarchical flow matching using mini-batch optimal transport (OT) coupling. The paper empirically compares different schemes of coupling, including data coupling and joint data-velocity coupling. The paper experiments on synthetic and image datasets to show the effectiveness of the proposed method.
weaknesses: •	The novelty of the paper is limited and incremental, the usage of minibatch OT coupling in flow matching is already proposed in [1]. The main point of the paper that using mini-batch OT inherently simplifies the velocity distribution is only intuitively explained without further theoretical mathematical evaluation.

•	The paper lacks further theoretical discussion and comparison between “data coupling” and “joint data and velocity coupling”.

•	The method is not simulation-free, which increases the training cost.

•	[1] Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky TQ Chen. Multisample flow matching: Straightening flows with minibatch couplings. In Proc. ICML, 2023

---

confidence: 3
contribution: 2
initial_rating: 4
presentation: 3
questions: Please see weakness part.
soundness: 2
strengths: The paper extends the exising HRF framework by introducing coupling, showing that the velocity distribution learnt from coupling data can recover the data distrbution as well. Through some 1D experiments, the authors show that coupling can reduce the multi-modality of velocity distribution, which is convincing. Results of large scale experiments also verify the effectiveness of proposed methods.
summary: The paper proposes Hierarchical Flow Matchining with mini-batch coupling, extending the existing HRF method. The coupling can be conducted in two levels: data distribution and velocity distribution, which is shown to gradually adjust the complexity of the distributions across
different levels of the hierarchy via mini-batch couplings.
Extensive numerical experiments are conducted to verify the efficiency.
weaknesses: 1. The proposed method is a naive combination of two existing methods: HRF and mini-batch coupling, and provides little new insights. It is a well known result that mini-batch coupling can help straighten the velocity and thus address multi-modality in some sense. The authors didn't provide further useful insights for why mini-batch coupling could benefit HRF other than some 1D numerical experiments. Can the authors develop some theories (even toy example is fine) to demonstrate why mini-batching coupling is good for HRF?

2. If I understand correctly, the velocity coupling requires a pretrained velocity model. This limits the practical application of the proposed method, which can only do distillation instead of training from scratch.

3. The authors only compared the proposed algorithm with FM, HRF and FM with coupling, claiming benefits in low NFE regime. The gains in high NFE regime seem quite marginal according to Figure 4(d). In terms of reducing NFE, the proposed method should also compare with other distillation/one-step FM algorithms such as [1,2]. 

[1] Geng, Zhengyang, et al. "Mean flows for one-step generative modeling." arXiv preprint arXiv:2505.13447 (2025).

[2] Frans, Kevin, et al. "One step diffusion via shortcut models." arXiv preprint arXiv:2410.12557 (2024).


Decision: Reject