Local Flow Matching Generative Models

## Abstract

Abstract Flow Matching (FM) is a simulation-free method for learning a continuous and invertible flow to interpolate between two distributions, and in particular to generate data from noise in generative modeling. In this paper, we introduce Local Flow Matching ( LFM ), which learns a sequence of FM sub-models and each matches a diffusion process up to the time of the step size in the data-to-noise direction. In each step, the two distributions to be interpolated by the sub-model are closer to each other than data vs. noise, and this enables the use of smaller models with faster training. The stepwise structure of LFM is natural to be distilled and different distillation techniques can be adopted to speed up generation. Theoretically, we prove a generation guarantee of the proposed flow model in terms of the χ 2 superscript 𝜒 2 \chi^{2} italic_χ start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT -divergence between the generated and true data distributions. In experiments, we demonstrate the improved training efficiency and competitive generative performance of LFM compared to FM on the unconditional generation of tabular data and image datasets, and also on the conditional generation of robotic manipulation policies.

## 1 Introduction

Generative modeling has revolutionized the field of machine learning, enabling the creation of realistic synthetic data across various domains.
Recently, diffusion models [, ] started to take over earlier models like Generative Adversarial Networks (GAN) [],
Variational Autoencoders (VAE) [] and Normalizing Flows [],
offering benefits in terms of stability, diversity, and scalability. Diffusion models have found diverse applications in audio synthesis [], text-to-image synthesis [], imitation learning for robotics [], among others.
A notable advantage of score-based diffusion models is their simulation-free training, meaning that the training objective is a “score matching” loss taking the form of an $L^{2}$ loss averaged over data samples.
Under a Stochastic Differential Equation (SDE) formulation [], this allows training of a continuous-time score function parametrized by a neural network from a large amount of data in high dimensions.

Compared to the SDE generation in diffusion models, the Ordinary Differential Equation (ODE) generation of a trained diffusion or flow model is deterministic and typically uses fewer time steps.
Using an ODE formulation, Flow Matching (FM) models [, , ] proposed a simulation-free training of flow models based on regressing vector fields using an $L^{2}$ “matching” loss
and have shown state-of-the-art generation performance on various tasks,
including text-to-image generation [],
humanoid robot control [],
audio generation [],
and application to discrete data in programming code generation [].
The simulation-free training of FM models significantly alleviates the computational issue of earlier Continuous Normalizing Flow (CNF) models using likelihood-based training [].
Going beyond the fast generation of ODE flow models, recent works on model distillation [, , ] have shown that the generation of large (diffusion and flow) generative models can be significantly accelerated and in some cases computed in an extremely small number of steps.

In this work, we propose a simulation-free flow model called “Local Flow Matching” (LFM), which can be seen as a stepwise version of FM, and we call previous FMs the global ones.
The proposed LFM trains a sequence of small (sub-)flow models, which, after concatenation, transport invertibly from data to noise and back.
In each sub-flow, any FM model can be plugged in.
In other words, the global FM tries to interpolate directly between noise and data distributions which may differ a lot,
while our local FM breaks down the job into smaller steps and interpolates between a pair of distributions which are closer to each other (namely “local”) in each step, see Figure 1.
Because the source and target distributions are not too far away, we expect to train a local FM model with potentially smaller model and faster convergence in training. This would allow reduced memory load and computational cost without sacrificing model quality.
In addition to the benefits when training from scratch, our framework is compatible with various distillation techniques, and we empirically found that our model can have advantage over global FMs after distillation.

Specifically, in each step of LFM, we train a sub-flow model to interpolate between $(p_{n-1},p_{n}^{*})$, where $p_{n}^{*}$ evolves from $p_{n-1}$ along the diffusion process for time up to the step size. The forward process (data-to-noise) starts from $p_{0}$ being the data distribution and ends at some $p_{N}$ which is close to normal distribution. The reverse process (noise-to-data) uses the invertibility of each sub-flow model to generate data from noise by ODE integration.
The construction of using (the marginal distribution of) a diffusion process as a “target” in each step allows us to theoretically prove the generation guarantee of LFM by connecting to the diffusion theory.

The stepwise approach in our work is inspired by similar constructions in the literature,
including block-wise training of ResNet under the GAN framework []
and
flow-based generative models implementing a discrete-time gradient descent in Wasserstein space following the Jordan-Kinderleherer-Otto (JKO) scheme [, , , ].
Unlike previous models, the proposed LFM is simulation-free, which allows application to high-dimensional data like large images and robotics data as shown in our experiments.

In summary, the contributions of the work are:

- •
To generate data from the normal distribution, we propose to train a series of FM sub-models that approximately follow the diffusion process in the data-to-noise direction.
Such breakdown of the global flow into small pieces makes the pair of distributions closer to each other in each step, and then the FM sub-model can be obtained with reduced model size and faster training convergence.
Concatenation of all the sub-models provides an invertible flow that can go from noise to data in the reverse direction.
- •
Theoretically, we prove the generation guarantee, namely how the generated distribution by LFM approximates the data distribution $P$, under $\chi^{2}$ divergence which implies Kullback–Leibler (KL) and Total Variation (TV) guarantees.
We prove an $O(\varepsilon^{1/2})$-$\chi^{2}$ guarantee, where $\varepsilon$ is the $L^{2}$ error of FM training objective,
and the other technical assumptions are motivated by our stepwise FM to the OU process.
Our theory applies when data density is regular
and also covers cases when $P$ merely has finite second moments (where we use a short-time initial diffusion to smooth $P$),
e.g., when $P$ is compactly supported.
- •
Our framework allows readily plugging in different designs of the FM in each local sub-flow model.
In addition, the stepwise structure of LFM renders it natural to be distilled, and our approach is compatible with different distillation techniques.
Empirically, the proposed LFM shows
improved training efficiency and competitive generative performance against existing FM methods on likelihood estimation, image generation, and robotic manipulation tasks.
The model gives strong performance when trained from scratch and in the distilled setting.

#### Notations.

We use the same notation for the distribution and its density (with respect to the Lebesgue measure on $\mathbb{R}^{d}$) when there is no confusion.
For a distribution $P$, $M_{2}(P):=\int_{\mathbb{R}^{d}}\|x\|^{2}dP(x)$.
Let $\mathcal{P}_{2}=\{P\text{ on $\mathbb{R}^{d}$},\,s.t.,M_{2}(P)<\infty\}$.
The Wasserstein-2 distance, denoted by $\mathcal{W}_{2}(P,Q)$, gives a metric on $\mathcal{P}_{2}$.
For $T:\mathbb{R}^{d}\to\mathbb{R}^{d}$,
$T_{\#}P$ denotes the pushforward of $P$, i.e.,
$T_{\#}P(A)=P(T^{-1}(A))$ for a measurable set $A$.
We also write $T_{\#}p$ for the pushforwarded density.

### 1.1 Related works

#### Continuous normalizing flow (CNF).

CNF trains a neural ODE model [] by maximizing the model likelihood, which the ODE parametrization can compute,
on observed data samples [].
To facilitate training and inference, subsequent works have proposed advanced techniques such as trajectory regularization [, ] and block-wise training [, ]. These techniques help stabilize the training process and improve the model’s performance.
Despite successful applications in time-series analyses [], uncertainty estimation [], optimal transport [], and astrophysics [],
a main drawback of CNF is its computational cost, since backpropagating the neural ODE in likelihood-based training is expensive (non-simulation free) and not scalable to high dimensions.

#### Simulation-free flow models.

Flow Matching (FM) models [, , ] are simulation-free and a leading class of generative models. We review the technical details of FM in Section 3.1.
FM methods are compatible with different choices to interpolate two random end-points drawn from the source and target distributions, e.g. straight lines (called “Optimal Transport (OT) path”) or motivated by the diffusion process (called “diffusion path”) [].
Later works also considered pre-computed OT interpolation [], and stochastic interpolation paths [].
All previous works train a global flow model to match between the two distributions, which could require a large model that takes a longer time to train.
In this work, we propose to train multiple smaller flow models. Our approach is compatible with any existing FM method to train these so-called local flows, making it a flexible and extensible framework.

#### Accelerated generation and model distillation.

Model compression and distillation have been intensively developed to accelerate the generation of large generative models.
[] proposed learning a compressed student normalizing flow model by minimizing the reconstruction loss from a teacher model.
For diffusion models, progressive distillation was developed in [],
and Consistency Models [] demonstrated high quality sample generation by directly mapping noise to data.
For FM models, [] proposed to distill the ODE trajectory into a single mapping parametrized by a network, which can reduce the number of function evaluations (NFE) to be one.
The approach was later effectively applied to large text-to-image generation [].
More recent techniques to distill FM models include dynamic programming to optimize stepsize given a budget of NFE [].
In our work,
each local flow model can be distilled into a single-step mapping following [], and the model can be further compressed if needed. Our framework is compatible with different distillation techniques.

#### Theoretical guarantees of generative models.

Guarantees of diffusion models, where the generation process utilizes an SDE (random) [, , , ]
or ODE (deterministic) sampler [, , , , ],
have been recently intensively developed.
In comparison, there are fewer theoretical findings for ODE flow models
both in training (forward process) and generation (reverse processes)
like CNF or FM.
For FM models, $\mathcal{W}_{2}$-guarantee was proved in []
and in [] with sample complexity analysis.
The Wasserstein bound does not imply KL or TV bound, which are more relevant for information-theoretical and statistical interpretation.
For CNF trained by maximizing likelihood, non-parametric statistical convergence rates were proved in [].
KL guarantee of a step-wisely trained CNF model motivated by JKO scheme [] was proved in [], yet the approach is likelihood-based and not simulation-free.
Recently, [] proved KL guarantee for an SDE interpolation version of the FM model introduced in [].
Our work analyzes a stepwise ODE FM model, which is simulation-free, and the guarantee is in $\chi^{2}$ divergence which implies KL (and TV) guarantees.

## 2 Preliminaries

#### Flow models and neural ODE.

In the context of generative models, the goal is to generate data distribution $P$, which is usually only accessible from finite training samples.
When $P$ has density we denote it by $p$.
A continuous-time flow model trains a neural ODE [] to generate data distribution $p$ from a standard distribution $q$, which is typically $\mathcal{N}(0,I)$ and called “noise”.

Specifically, a neural ODE model provides a velocity field $v(x,t;\theta)$ on $\mathbb{R}^{d}\times[0,T]$ parametrized by a neural network, and $\theta$ consists of trainable parameters. In a flow model, the solution of the ODE

$$ $\dot{x}(t)=v(x(t),t;\theta),\quad t\in[0,T],\quad x(0)\sim\rho_{0},$ (1) $$

where $\rho_{0}$ is a distribution, and we denote the law of $x(t)$ as $\rho_{t}$.
Whenever the ODE is well-posed, it gives a continuous invertible mapping from the initial value $x(0)$ to the terminal value $x(T)$, and the inverse map from $x(T)$ to $x(0)$ can be computed by integrating the ODE (1) reverse in time.
Thus, we can either set $\rho_{0}$ to be noise and $\rho_{T}$ to be data or the other way round.
Earlier flow models like (continuous-time) CNFs [] train the flow $v(x,t;\theta)$ based on likelihood, i.e. letting $\rho_{0}$ be noise and maximizing the likelihood of $\rho_{T}$ on data samples. Such approaches are not simulation-free and can be difficult to scale to large-size computations.

#### Ornstein-Uhlenbeck (OU) process.

The OU process in $\mathbb{R}^{d}$ observes the following SDE

$$ $dX_{t}=-X_{t}dt+\sqrt{2}dW_{t},$ (2) $$

where the equilibrium density is $q\propto e^{-V}$ with $V(x)=\|x\|^{2}/2$. In other words, $q$ is the standard normal density $\mathcal{N}(0,I)$.
Suppose $X_{0}\sim\rho_{0}$ which is some initial density at time zero (the initial distribution may not have density),
and denote by $\rho_{t}$ the marginal density of $X_{t}$ for $t>0$.
The time evolution of $\rho_{t}$ is described by the Fokker–Planck Equation (FPE)
$\partial_{t}\rho_{t}=\nabla\cdot(\rho_{t}\nabla V+\nabla\rho_{t})$, $V(x)=\|x\|^{2}/2$,
which determines $\rho_{t}$ given the initial value $\rho_{0}$.
We also introduce the operator $({\rm OU})_{0}^{t}$ and write

$$ $\rho_{t}=({\rm OU})_{0}^{t}\rho_{0}.$ (3) $$

Equivalently, $\rho_{t}$ is the probability density of the random vector
$Z_{t}:=e^{-t}X_{0}+\sigma_{t}Z$, where $\sigma_{t}^{2}:=1-e^{-2t}$,
$Z\sim\mathcal{N}(0,I_{d})$ and is independent from $X_{0}$.

Figure: Figure 1: Illustration of the proposed LFM model. In the $n$-step step, a local FM model is trained to interpolate from $p_{n-1}$ to $p_{n}^{*}$ and it gives the learned (invertible) transport map $T_{n}$ which pushforwards $p_{n-1}$ to $p_{n}$. The concatenation of the $N$ sub-models gives a flow between data and noise.
Refer to caption: extracted/5897716/diag_lfm.png

## 3 Local Flow Matching ( LFM )

The essence of the proposed LFM model is to break down a single flow from data to noise (and back) into several pieces, and apply FM on each piece sequentially.
We introduce the method in this section, and further algorithmic details are provided in Section 5.

### 3.1 Review of Flow Matching

Flow Matching [],
also proposed as “stochastic interpolants” [] and “rectified flow” [],
trains the flow $v(x,t;\theta)$ by minimizing an $L^{2}$ loss and is simulation-free; our stepwise approach in this work will be based on FM.
Rescale the time interval to be $[0,1]$.
FM utilizes a pre-specified interpolation function $I_{t}$ given two endpoints $x_{l}$ and $x_{r}$ (l for ‘left’ and r for ‘right’) as

$$ $\phi(t):=I_{t}(x_{l},x_{r}),\quad t\in[0,1],$ (4) $$

where $x_{l}\sim p$, $x_{r}\sim q$, and $I_{t}$ can be analytically designed.
The (population) training objective is

$$ $\min_{\theta}\int_{0}^{1}\mathbb{E}_{x_{l},x_{r}}\|v(\phi(t),t;\theta)-\frac{d }{dt}\phi(t)\|_{2}^{2}dt,$ (5) $$

and, as has been proven in the literature,
the velocity field $v$ that minimizes (5) has the expression $v(x,t)=\mathbb{E}_{x_{l},x_{r}}[\frac{d}{dt}I_{t}(x_{l},x_{r})|I_{t}(x_{l},x_{
r})=x]$
and it induces a flow that transports from $p$ to $q$.
We call this $v$ the target velocity field, and say that the FM model interpolates the pair of distributions $(p,q)$.
See Section 5.1 for more algorithmic details of FM, such as the choice of $I_{t}$.

### 3.2 Stepwise training of N 𝑁 N italic_N FM sub-models

We propose to train a sequence of $N$ FM sub-models (called sub-flows) on time $[0,T]$,
each by its own training objective to interpolate a pair of distributions,
and collectively the $N$ sub-flows fulfill the transport from data $p$ to noise $q$ (and back by the reverse flow).
Our method will train the $N$ sub-flows sequentially,
where in each step the flow tries to match the terminal density evolved by the OU process of time up to the step size.
Because the step size is not large and the pair of end-point distributions to interpolate are not far away, we call our method local FM.

#### Training.

On the time interval $[0,T]$,
we first specify a time step schedule $\{\gamma_{n}\}_{n=1}^{N}$, $\sum_{n}\gamma_{n}=T$,
and for the simplest case $\gamma_{n}=\gamma=T/N$.
Suppose the data distribution $P$ has a regular density $p_{0}$ –
if not, we can apply a short-time diffusion to $P$, and take the resulting density as $p_{0}$, see more in Section 4.3.
Starting from $p_{0}$ (when $n=1$),
we recursively construct target density $p_{n}^{*}$ by

$$ $p_{n}^{*}=(\text{OU})_{0}^{\gamma_{n}}p_{n-1},$ $$

where the operator $(\text{OU})_{0}^{t}$ is as defined in (3).
In other words, let $x_{l}\sim p_{n-1}$, and

$$ $x_{r}:=e^{-\gamma_{n}}x_{l}^{\prime}+\sqrt{1-e^{-2\gamma_{n}}}g,\quad x_{l}^{ \prime}\sim p_{n-1},\quad g\sim\mathcal{N}(0,I_{d}),\quad g\perp x_{l}^{\prime},$ (6) $$

then the marginal distribution of $x_{r}$ is $p_{n}^{*}$.
We now train the $n$-th sub-flow, denoted by $\hat{v}_{n}(x,t;\theta)$,
by FM to interpolate the pair $(p_{n-1},p_{n}^{*})$ using (4)(5) (the time interval $[0,\gamma_{n}]$ is rescaled to be $[0,1]$ in FM training).
The velocity field $\hat{v}_{n}$ can have its own parametrization $\theta_{n}$, which then can be trained independently from previous sub-flows.

After the sub-flow $\hat{v}_{n}$ is trained, it provides a transport map $T_{n}$ that maps $x_{n-1}\sim p_{n-1}$ to $x_{n}$ by

$$ $x_{n}:=T_{n}(x_{n-1})=x(\gamma_{n})=x_{n-1}+\int_{0}^{\gamma_{n}}\hat{v}_{n}(x (t),t;\theta)dt,$ (7) $$

where $x(t)$ solves the ODE $\dot{x}(t)=\hat{v}_{n}(x(t),t;\theta)$ from $x(0)=x_{n-1}$.
The mapping $T_{n}$ is invertible and $T_{n}^{-1}$ is by integrating the reverse-time ODE.
We denote the distribution of $x_{n}$ as $p_{n}$, that is,

$$ $p_{n}=(T_{n})_{\#}p_{n-1}.$ $$

From this $p_{n}$, we can train the next sub-flow by interpolate $(p_{n},p_{n+1}^{*})$.
This recursive scheme is illustrated in Figure 1.

If the flow matching in $n$-th step is successful, then we expect the trained $\hat{v}_{n}$ makes $p_{n}\approx p_{n}^{*}$.
Define the time stamps $\{t_{n}\}_{n=0}^{N}$, $t_{0}=0$, $t_{n}-t_{n-1}=\gamma_{n}$, and then $t_{N}=T$.
If $p_{n}$ exactly equals $p_{n}^{*}=(\text{OU})_{0}^{\gamma_{n}}p_{n-1}$, then over $N$ steps we have $p_{N}=(\text{OU})_{0}^{T}p_{0}$, which approximates the equilibrium $q$ exponentially fast as $T$ increases.
In practice, the trained flow has some finite flow matching error and $p_{n}$ has some discrepancy from $p_{n}^{*}$,
yet if the error is small, we still expect $p_{N}\approx q$, assuming $T=\sum_{n}\gamma_{n}$ is sufficiently large.
This will be theoretically analyzed in Section 4.

In practice, when training the $n$-step sub-flow, the finite samples of $p_{n-1}$ are obtained by transporting samples from $p_{0}$ through the previous sub-flows.
This pushforward from $p_{n-1}$ to $p_{n}$ can be done for all training samples once the $n$-th sub-flow is trained, see Algorithm 1.
The proposed LFM model can be trained from scratch,
and we also distill the model to improve generation efficiency.
See Section 5 for details.

#### Generation.

Once the $N$ sub-flows are trained, we go backward from the $N$-th to the 1st sub-flows to generate data from noise.
Specifically, we sample $y_{N}\sim q$, and let $y_{n-1}=T_{n}^{-1}(y_{n})$ for $n=N,\cdots,1$, where $T_{n}^{-1}$ is computed by integrating the ODE with velocity field $\hat{v}_{n}$ reverse in time.
We then use $y_{0}$ as the generated data samples.
The closeness of the distribution of $y_{0}$ to the data distribution $P$ will be theoretically shown in Section 4.

## 4 Theoretical guarantee of generation

In this section, we prove the generation guarantee, namely how the generated density by the trained LFM model is close to the true data distribution. All proofs are left to Appendix A.

### 4.1 Summary of forward and reverse processes

Recall that $P$ is the distribution of data in $\mathbb{R}^{d}$, and $q$ the density of $\mathcal{N}(0,I_{d})$.
The procedures of training and generation in Section 3 can be summarized in the following forward (data-to-noise training) and reverse (noise-to-data generation) processes respectively:

$$ $\begin{split}\text{(forward)}\quad&p=p_{0}\xrightarrow{T_{1}}{p_{1}} \xrightarrow{T_{2}}{}\cdots\xrightarrow{T_{N}}{p_{N}}\approx q,\\ \text{(reverse)}\quad&p\approx q_{0}\xleftarrow{T_{1}^{-1}}{q_{1}}\xleftarrow{ T_{2}^{-1}}{}\cdots\xleftarrow{T_{N}^{-1}}{q_{N}}=q,\end{split}$ (8) $$

where $T_{n}$ is by the learned $n$-th step sub-flow as defined in (7).
When $P$ has a regular density $p$ we set it as $p_{0}$; otherwise, $p_{0}$ will be a smoothed version of $P$, see more below.
The reverse process gives the final output samples which have density $q_{0}$.
The goal of our analysis is to show that $q_{0}\approx p$, namely the generated density is close to the data density.

To keep exhibition simple, we consider when $\gamma_{n}=\gamma>0$ for all $n$.
Using the time stamps $t_{n}$, the $n$-th step sub-flow is on the time interval $[t_{n-1},t_{n}]$, which can be shifted to be $[0,\gamma]$.
Our analysis will be based on comparing the true or target flow (that transports to terminal density $p_{n}^{*}$)
with the learned flow (that transports to terminal density $p_{n}$),
and we introduce the notations for the corresponding transport equations.

For fixed $n$, on the shifted time interval $[0,\gamma]$, the target flow and the learned flow are induced by the velocity field $v(x,t)$ and $\hat{v}(x,t)=\hat{v}_{n}(x,t;\theta)$ respectively. We omit n in the notation.
As was shown in Section 3.1,
the target $v$ depends on the choice of $I_{t}$ yet it always transports from $p_{n-1}$ to $p_{n}^{*}$.
Let $\rho_{t}(x,t)$ be the law of $x(t)$ that solves the ODE with velocity field $v$, where $x(0)\sim p_{n-1}$,
and $\hat{\rho}_{t}(x,t)$ be the law of $x(t)$ that solves the ODE with the learned $\hat{v}$.
($\rho_{t}$ is not necessarily the density of an OU process $X_{t}$, though $\rho_{\gamma}=p_{n}^{*}=(\text{OU})_{0}^{\gamma_{n}}p_{n-1}$.)
We also denote $\rho(x,t)$ as $\rho_{t}$ (omitting the variable $x$)
or $\rho$ (omitting both $x$ and $t$), depending on the context, and similarly for $\hat{\rho}(x,t)$.
The target and learned flows have the transport equations as

$$ $\begin{split}\partial_{t}\rho+\nabla\cdot(\rho v)&=0,\quad\rho_{0}=p_{n-1}, \quad\rho_{\gamma}=p_{n}^{*}.\\ \partial_{t}\hat{\rho}+\nabla\cdot(\hat{\rho}\hat{v})&=0,\quad\hat{\rho}_{0}=p _{n-1},\quad\hat{\rho}_{\gamma}=p_{n}.\end{split}$ (9) $$

### 4.2 Exponential convergence of the forward process in χ 2 superscript 𝜒 2 \chi^{2} italic_χ start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT

We introduce the following assumption about the learned flow $\hat{v}_{n}$:

###### Assumption 1 .

(A0)
For all $n$,
the learned $\hat{v}_{n}$ ensures that $T_{n}$ and $T_{n}^{-1}$ are non-degenerate (Definition A.1),
and, on the time interval $[t_{n-1},t_{n}]$ shifted to be $[0,\gamma]$,
$\int_{0}^{\gamma}\int_{\mathbb{R}^{d}}\|v-\hat{v}\|^{2}{\rho}dxdt\leq
\varepsilon^{2}$.
Without loss of generality, assume $\varepsilon<1$.

The training objective of FM is equivalent to minimizing the $L^{2}$ loss $\int_{0}^{\gamma}\int_{\mathbb{R}^{d}}\|v-\hat{v}\|^{2}\rho dxdt$ [],
thus our $\varepsilon$ is the learning error of FM (in each step and uniform for all $n$).

###### Assumption 2 .

There are positive constants $C_{1},C_{2},L$ such that, for all $n$, on the time interval $[t_{n-1},t_{n}]$
shifted to be $[0,\gamma]$,

(A1)
$\rho_{t}$, $\hat{\rho}_{t}$ for any $t\in[0,\gamma]$ are positive on $\mathbb{R}^{d}$ and
$\rho_{t}(x),\,\hat{\rho}_{t}(x)\leq C_{1}e^{-\|x\|^{2}/2}$;

(A2)
$\forall t\in[0,\gamma]$, $\rho_{t},\,\hat{\rho}_{t}$ are $C^{1}$ on $\mathbb{R}^{d}$
and $\|\nabla\log\rho_{t}(x)\|$, $\|\nabla\log\hat{\rho}_{t}(x)\|\leq L(1+\|x\|)$, $\forall x\in\mathbb{R}^{d}$;

(A3)
$\forall t\in[0,\gamma]$, $\int_{\mathbb{R}^{d}}(1+\|x\|)^{2}({\rho_{t}^{3}}/{\hat{\rho}_{t}^{2}})(x)dx
\leq C_{2}$.

At $t=0$, $\rho_{0}=\hat{\rho}_{0}=p_{n-1}$, thus Assumption 2 requires that $f=p_{n}$ for $n=0,1,\cdots$ satisfies

$$ $f(x)\leq C_{1}e^{-\|x\|^{2}/2},\quad\|\nabla\log f(x)\|\leq L(1+\|x\|),\quad \int_{\mathbb{R}^{d}}(1+\|x\|)^{2}f(x)dx\leq C_{2},$ (10) $$

by (A1)(A2)(A3) respectively.
The first two inequalities require $p_{n}$ to have a Gaussian decay envelope
and the score of $p_{n}$ has linear growth (can be induced by Lipschitz regularity),
and the third inequality can be implied by the first one.
The condition (10) poses regularity conditions on $p_{0}$ which can be satisfied by many data densities in applications, and, in particular, if $P$ has finite support then these hold after $P$ is smoothed by an initial short-time diffusion (Lemma A.6).

Technically, Assumption 2 poses the Gaussian envelope and regularity requirements on all $p_{n}$ and also $\rho_{t}$ and $\hat{\rho}_{t}$ for all time.
This can be expected to hold at least when FM is well-trained:
suppose $\rho_{t}$ satisfies (A1)(A2) for all $t$
due to the regularity of $v$ (by the analytic $I_{t}$ and the regularity of the pair of densities $(p_{n-1},p_{n}^{*})$),
when the true and learned flows match each other, we have $\hat{\rho}\approx\rho$,
thus we also expect (A1)(A2) to hold for $\hat{\rho}_{t}$;
Meanwhile, the ratio $\rho_{t}/\hat{\rho}_{t}$ is close to 1 and if can be assumed to be uniformly bounded,
then (A3) can be implied by the boundedness of $\int_{\mathbb{R}^{d}}(1+\|x\|)^{2}{\rho_{t}}(x)dx$
which can be implied by the Gaussian envelope (A1) of $\rho_{t}$.

###### Proposition 4.1 (Exponential convergence of the forward process) .

Under Assumptions 1-2,

$$ $\chi^{2}(p_{n}\|q)\leq e^{-2\gamma n}\chi^{2}(p_{0}\|q)+\frac{C_{4}}{1-e^{-2 \gamma}}\varepsilon^{1/2}$ (11) $$

for $n=1,2,\cdots$, where the constant $C_{4}$ defined in (20) is determined by $C_{1}$, $C_{2}$, $L$, $\gamma$ and $d$.

###### Assumption 1 .

###### Assumption 2 .

###### Proposition 4.1 (Exponential convergence of the forward process) .

### 4.3 Generation guarantee of the backward process

#### P 𝑃 P italic_P with regular density.

Because the composed transform $T_{N}\circ\cdots\circ T_{1}$ from $p_{0}$ to $p_{N}$ is invertible,
and the inverse map transforms from $q=q_{N}$ to $q_{0}$,
the smallness of $\chi^{2}(p_{N}\|q_{N})$ implies the smallness of $\chi^{2}(p_{0}\|q_{0})$
due to a bi-directional version of data processing inequality (DPI), see Lemma A.4,
As a result, the exponential convergence of the forward process in Proposition 4.1 directly gives the $\chi^{2}$-guarantee $q_{0}\approx p$.

###### Theorem 4.2 (Generation guarantee of regular data density) .

Suppose $P\in\mathcal{P}_{2}$ has density $p$,
let $p_{0}=p$ and $p_{0}$ satisfies (10).
Under Assumptions 1-2,
if we use the number of steps
$N\geq\frac{1}{2\gamma}\left(\log\chi^{2}(p_{0}\|q)+\frac{1}{2}\log(\frac{1}{
\varepsilon})\right)\sim\log(1/\varepsilon)$,
we have
$\chi^{2}(p\|q_{0})\leq C\varepsilon^{1/2}$, $C:=(1+\frac{C_{4}}{1-e^{-2\gamma}})$.

By that ${\rm KL}(p\|q)\leq\chi^{2}(p\|q)$ (Lemma A.5),
the $\chi^{2}$-guarantee of $O(\varepsilon^{1/2})$ in Theorem 4.2
implies that
${\rm KL}(p\|q_{0})=O(\varepsilon^{1/2})$,
and consequently ${\rm TV}(p,q_{0})=O(\varepsilon^{1/4})$ by Pinsker’s inequality.

###### Theorem 4.2 (Generation guarantee of regular data density) .

#### P 𝑃 P italic_P up to initial diffusion.

For data distribution $P$ that may not have a density or the density does not satisfy the regularity conditions,
we introduce a short-time diffusion to obtain a smooth density $\rho_{\delta}$ from $P$ and use it as $p_{0}$.
This construction can ensure that $\rho_{\delta}$ is close to $P$, e.g. in $\mathcal{W}_{2}$ distance,
and is commonly used in diffusion models []
and also the theoretical analysis (known as “early stopping”) [].

###### Corollary 4.3 (Generation of $P$ up to initial diffusion) .

Suppose $P\in\mathcal{P}_{2}$ and for some $\delta<1$,
$p_{0}=\rho_{\delta}=({\rm OU})_{0}^{\delta}(P)$
satisfies (10) for some $C_{1}$, $C_{2}$ and $L$.
With $p_{0}=\rho_{\delta}$, suppose Assumptions 1-2 hold,
then for $N$ and $C$ are as in Theorem 4.2,
the generated density $q_{0}$ of the reverse process makes
$\chi^{2}(\rho_{\delta}||q_{0})\leq C\varepsilon^{1/2}$.
Meanwhile, $\mathcal{W}_{2}(P,\rho_{\delta})\leq C_{5}\delta^{1/2}$, $C_{5}:=(M_{2}(P)+2d)^{1/2}$.

In particular, as shown in Lemma A.6, when $P$ is compactly supported (not necessarily having density), then for any $\delta>0$,
there are $C_{1}$, $C_{2}$ and $L$ such that
$p_{0}=\rho_{\delta}$ satisfies (10).
This will allow us to make $\mathcal{W}_{2}(P,\rho_{\delta})$ arbitrarily small.
Generally, the theoretical constants $C_{1}$, $C_{2}$ and $L$ may depend on $\delta$,
and consequently so does the constant $C$ in the $\chi^{2}$ bound.
At last, similarly to the comment beneath Theorem 4.2,
the $O(\varepsilon^{1/2})$-$\chi^{2}$ guarantee in Corollary 4.3 of $q_{0}\approx p_{\delta}$ implies
$O(\varepsilon^{1/2})$-KL
and $O(\varepsilon^{1/4})$-TV guarantee.

Figure: Algorithm 1 Local Flow Matching (LFM) from scratch

###### Corollary 4.3 (Generation of P 𝑃 P italic_P up to initial diffusion) .

## 5 Algorithm

### 5.1 Training LFM from scratch

The training of LFM is summarized in Algorithm 1.
We call each sub-flow FM a step or a “block”.
In each block, any FM algorithm can be adopted.
In our experiments, we consider the following choices of $I_{t}(x_{l},x_{r})$ in (4), following [, ]:
(i) Optimal Transport: $I_{t}(x_{l},x_{r})=x_{l}+t(x_{r}-x_{l})$,
(ii) Trigonometric: $I_{t}(x_{l},x_{r})=\cos(\frac{1}{2}\pi t)x_{l}+\sin(\frac{1}{2}\pi t)x_{r}$.
The selection of time stamps $\{\gamma_{n}\}_{n=1}^{N-1}$ depends on the task and is explained in Appendix B.1. In our experiments, we use $N$ up to 10.
In each sub-flow, there are no restrictions on the architecture of $\hat{v}_{n}(x,t;\theta)$. We use fully connected networks for vector data and UNets for images.
When the sub-flows are independently parametrized, we reduce the model size of each block assuming that the target flow to match is simpler than a global flow, which also reduces memory load and facilitates the innerloop training of FM.
In computing the pushforward (Line 5 in Algorithm 1) and in generation,
the numerical integration of neural ODE follows the procedure in [].

### 5.2 Distillation of LFM

Inspired by [], we propose to distill an $N$-block pre-trained LFM model into $N^{\prime}<N$ steps distilled model, where $N=N^{\prime}k$ for some integer $k$.
Each of the $N^{\prime}$ sub-models can be distilled independently if parametrized independently.
The detail is in Algorithm A.1.
If the pre-trained LFM has independently parametrized blocks, we keep the sub-model size if $k=1$ and increase the model size if the distillation combines original blocks ($k>1$).
The composition of the $N^{\prime}$ distilled sub-models generates data from noise in $N^{\prime}$ steps.

## 6 Experiments

We apply the proposed LFM to simulated and real datasets,
including tabular data (Section 6.2),
image generation (Section 6.3),
and robotic manipulation (Section 6.4).
We demonstrate the improved training efficiency and generative performance of LFM compared to (global) flow models, and we also show the advantage of LFM after distillation on image generation.

Figure: a True samples from $p$
Refer to caption: extracted/5897716/tree_true.png

### 6.1 Two-dimensional toy data

The task is to generate distributions in 2D that have no analytic form, and we consider the “tree” and “rose” examples (Figure 2).
We compare our LFM against previous FM models: FM [, Example II] and InterFlow [].
To ensure a fair comparison, we maintain an identical training scheme for all methods, including the specification of optimizers, batch size, and number of training batches.
Further details regarding the experimental setup are provided in Appendix C.
The accuracy of the trained models is evaluated using the negative log-likelihood (NLL) metric in (26).
Figure 2 shows that LFM achieves good generation of the distribution and slightly better NLL in comparison.

**Table 1: Test NLL (lower is better) on tabular data, where $d$ refers to the data dimension. The lowest NLL is in bold and the 2nd lowest NLL is underlined. All baseline values are quoted from their original publications, where “-” entries indicate this dataset was not used.**
|  | LFM | InterFlow | JKO-iFlow | AdaCat | OT-Flow | nMDMA | CPF | BNAF | FFJORD |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| POWER<br>($d=6$) | -0.67 | -0.57 | -0.40 | -0.56 | -0.30 | -1.78 | -0.52 | -0.61 | -0.46 |
| GAS<br>($d=8$) | -12.43 | -12.35 | -9.43 | -11.27 | -9.20 | -8.43 | -10.36 | -12.06 | -8.59 |
| MINIBOONE<br>($d=43$) | 9.95 | 10.42 | 10.55 | 14.14 | 10.55 | 18.60 | 10.58 | 8.95 | 10.43 |
| BSDS300<br>($d=63$) | -157.80 | -156.22 | -157.75 | - | -154.20 | - | -154.99 | -157.36 | -157.40 |

### 6.2 Tabular data generation

We apply LFM to a set of tabular datasets [], where the generation performance is quantitatively measured by test NLL.
We use the following baselines, including both discrete and continuous flows, to compare against:
InterFlow [],
JKO-iFlow [],
AdaCat [],
OT-Flow [],
nMDMA [],
CPF [], BNAF [], and FFJORD []. Additional experimental details can be found in Appendix C.
The results are shown in Table 1,
where the proposed LFM is among the two best-performing methods on all datasets.

Figure: Figure 3: Unconditional image generation by LFM on $32\times 32$ images: CIFAR10 (upper left) and Imagenet-32 (lower left); $128\times 128$ images: Flowers (upper right) and LSUN Church [] (lower right).
Refer to caption: x1.png

### 6.3 Image generation

**Table 2: FID (lower is better) comparison of LFM against InterFlow (quoted from []) under same model sizes.**
|  | FID | Batch size | # of batches |
| --- | --- | --- | --- |
| LFM | 8.45 | 200 | $\mathbf{5\times 10^{4}}$ |
| InterFlow | 10.27 | 400 | $5\times 10^{5}$ |

**Table 3: FID comparison of LFM and InterFlow before and after distillation on Flowers $128\times 128$.**
|  | Pre-distillation | Distilled @ 4 NFEs | Distilled @ 2 NFEs |
| --- | --- | --- | --- |
| LFM | 59.7 | 71.0 | 75.2 |
| InterFlow | 59.7 | 80.0 | 82.4 |

We apply LFM to unconditional image generation of $32\times 32$ and $128\times 128$ images. We compare LFM with InterFlow in terms of Frechet Inception Distance (FID) before and after distillation. To ensure a fair comparison, we maintain the same network size for both methods,
and more details of the experimental setup are provided in Appendix C.

$32\times 32$ images. We use the CIFAR-10 [] and Imagenet-32 [] datasets.
As shown in Table 2, LFM requires significantly less computation in training than InterFlow,
and it achieves lower FID values.
We present generated images by LFM in Figure 3 and additionally,
by the distilled LFM in Figure A.2 (NFE = 5) which shows an almost negligible reduction in visual quality.

$128\times 128$ images. We use the Oxford Flowers [] dataset.
We first train LFM and InterFlow to reach the same FIDs on the test set, with InterFlow requiring 1.5x more training batches to do so. Subsequently, we distill LFM using Algorithm A.1 and InterFlow according to [].
Quantitatively, Table 3 shows lower FID by LFM under 4 or 2 NFEs after distillation.
This shows that in addition to efficiency in training from scratch, LFM after distillation also achieves lower FID.
To show the qualitative results, we give generated images in Figure 3 and noise-to-image trajectories in Figure A.1. Figure A.3 further showcases high-fidelity images generated by LFM after distillation with 4 NFEs.

### 6.4 Robotic Manipulation

Figure: a Lift
Refer to caption: extracted/5897716/1st_lift.png

We consider robotic manipulation tasks from the Robomimic benchmark [], which consists of 5 tasks of controlling robot arms to perform various pick-and-place operations (Figure 4). E.g., the robot may need to pick up a square object (Figure 4) or move a soda can from one bin to another (Figure 4).
Recently, generative models that output robotic actions conditioning on the state observations (robot positions or camera image embedding) provide state-of-the-art performance on completing these tasks [].
However, it is difficult to transfer pre-trained diffusion or flow models (on natural images) to the conditional generation task here,
because the state observations are task-specific and contain nuanced details about the robots, objects, and environment, which are not present in natural images but are necessary to be understood by the model to determine the appropriate actions.
Therefore, it is often needed to train a generative model from scratch to directly learn the relation from the task-specific state observations to the actions.

**Table 4: Success rate (see (27), higher is better) of FM and LFM for robotic manipulation on Robomimic []. We evaluate both methods for 100 rollouts and report success rates at different epochs in the format as (Success rate, epochs).**
|  | Lift | Can | Square | Transport | Toolhang |
| --- | --- | --- | --- | --- | --- |
| FM | (1.00, 200) | (0.94, 200)<br>(0.98, 500) | (0.88, 200)<br>(0.94, 750) | (0.60, 200)<br>(0.81, 1500) | (0.52, 200) |
| LFM | (1.00, 200) | (0.97, 200)<br>(0.99, 500) | (0.87, 200)<br>(0.93, 750) | (0.75, 200)<br>(0.88, 1500) | (0.53, 200) |

We train a (global) FM model and the proposed LFM from scratch, where the total number of parameters is kept the same for both methods.
Additional experimental details are in Appendix C.2.
As shown in Table 4, LFM is competitive against FM in terms of the success rate,
and the convergence is faster in some cases (indicated by the higher success rate at early epochs).
The performance on the “Toolhang” task does not improve with longer training, as identified in [].

## 7 Discussion

The theoretical analysis can be extended from several angles:
First, we analyzed the $\chi^{2}$ guarantee and induced KL and TV bounds from the former. One may obtain sharper bounds by analyzing KL or TV directly and possibly under weaker assumptions.
E.g., Lemma 2.19 in [] can be used to derive a KL bound for the one-step FM
under similar assumptions as our Assumption 2(A2).
Second, we currently assume that $T_{n}^{-1}$ is exact in the reverse process (generation). The analysis can be extended to incorporate inversion error due to numerical computation in practice, possibly by following the strategy in [].
Meanwhile, the proposed methodology has the potential to be further enhanced.
In our current experiments, we learn the FM blocks independently parametrized as $\hat{v}(x,t;\theta_{n})$. If we introduce weight sharing of $\theta_{n}$ across $n$, it can impose additional time continuity of the flow model.
In addition, it would be useful to explore more distillation techniques under our framework and to extend the approach to more applications.

## Acknowledgments

The project was partially supported by National Science Foundation (NSF) DMS-2134037.
YX was also partially supported by
NSF DMS-2220495.
XC was also partially supported by
NSF DMS-2237842
and Simons Foundation (grant ID: 814643).

## Appendix A Proofs

Denote by Leb the Lebesgue measure.

###### Definition A.1 (Non-degenerate mapping) .

$T:\mathbb{R}^{d}\to\mathbb{R}^{d}$ is non-degenerate if for any set $A\subset\mathbb{R}^{d}$ s.t. $Leb(A)=0$,
then $Leb(T^{-1}(A))=0$.

###### Definition A.1 (Non-degenerate mapping) .

### A.1 Proofs in Section 4.2

###### Proof of Proposition 4.1 .

The proof is based on the descending of $\chi^{2}(p_{n}\|q)$ in each step.
Note that under (A1), by the argument following (13) below, we have $p_{0}/q\in L^{2}(q)$ i.e. $\chi^{2}(p_{0}\|q)<\infty$.

We first consider the target and learned flows on $[t_{n-1},t_{n}]$ to establish the needed descending arguments for the $n$-th step.
We shift the time interval to be $[0,T]$, $T:=\gamma>0$,
and the transport equations of $\rho$ and $\hat{\rho}$ are as in (9).
Define

$$ $G(t):=\chi^{2}(\rho_{t}\|q),\quad\hat{G}(t):=\chi^{2}(\hat{\rho}_{t}\|q).$ (12) $$

By the time endpoint values of $\rho$ and $\hat{\rho}$ as in (9), we have

$$ $G(0)=\hat{G}(0)=\chi^{2}(p_{n-1}\|q),\quad G(T)=\chi^{2}(p_{n}^{*}\|q),\quad \hat{G}(T)=\chi^{2}(p_{n}\|q).$ $$

We will show that a) $G(T)$ descends, and b) $\hat{G}(T)\approx G(T)$,
then, as a result, $\hat{G}(T)$ also descends.

We first verify the boundedness of $G(t)$ and $\hat{G}(t)$ at all time.
By definition, we have (we write integral on $\mathbb{R}^{d}$ omitting variable $x$ and $dx$ for notation brevity)

$$ $G(t)=\int\frac{(\rho_{t}-q)^{2}}{q}=\int(\frac{\rho_{t}}{q}-1)^{2}q=\int(\frac {\rho_{t}}{q})^{2}q-1,$ $$

when the involved integrals are all finite,
and similarly with $\hat{G}(t)$ and $\hat{\rho}_{t}$.
To verify integrability, observe that under (A1),

$$ $\frac{\rho_{t}^{2}}{q}(x),\,\frac{\hat{\rho}_{t}^{2}}{q}(x)\leq\frac{C_{1}^{2} e^{-\|x\|^{2}}}{c_{d}e^{-\|x\|^{2}/2}}=\frac{C_{1}^{2}}{c_{d}}e^{-\|x\|^{2}/2} ,\quad c_{d}:=(2\pi)^{-d/2}.$ (13) $$

This shows that $\rho_{t}/q,\hat{\rho}_{t}/q$ and thus $(\rho_{t}/q-1),(\hat{\rho}_{t}/q-1)$
are all in $L^{2}(q)$ for all $t$.
In particular, this applies to $\rho_{0}=\hat{\rho}_{0}=p_{0}$
by taking $n=1$ and $t=0$.
Thus, $G(t),\,\hat{G}(t)$ are always finite.
Additionally, (13) gives that

$$ $G(t),\hat{G}(t)\leq(C_{1}/c_{d})^{2}-1\leq(C_{1}/c_{d})^{2},\quad\forall t\in[ 0,T],$ (14) $$

Our descending argument is based on the following two key lemmas:

###### Lemma A.2 ( $\chi^{2}$ -contraction of OU process) .

$\chi^{2}(p_{n}^{*}\|q)\leq e^{-2\gamma}\chi^{2}(p_{n-1}\|q)$.

This implies that $G(T)\leq e^{-2T}G(0)$.

###### Lemma A.3 ( $\hat{G}(T)\approx G(T)$ ) .

$\chi^{2}(p_{n}\|q)\leq\chi^{2}(p_{n}^{*}\|q)+C_{4}\varepsilon^{1/2}$, with $C_{4}$ defined in (20).

Lemma A.2 follows standard contraction results of the diffusion process,
and Lemma A.3 utilizes the approximation of $\hat{\rho}\approx\rho$ due to the learning of the velocity field $\hat{v}\approx v$.
We postpone the proofs of the two lemmas after the proof of the proposition.

With the two lemmas in hands, we can put the $n$ steps together and prove (11).
Define

$$ $E_{n}:=\chi^{2}(p_{n}\|q),\quad\beta:=e^{-2\gamma}<1,\quad\alpha:=C_{4} \varepsilon^{1/2},$ $$

and then Lemmas A.2-A.3 give that

$$ $E_{n}\leq\beta E_{n-1}+\alpha.$ $$

By induction, one can verify that

$$ $E_{n}\leq\beta^{n}E_{0}+\frac{\alpha(1-\beta^{n})}{1-\beta}\leq\beta^{n}E_{0}+ \frac{\alpha}{1-\beta},$ $$

which proves (11) and finishes the proof of the proposition.
∎

###### Proof of Lemma A.2 .

The lemma follows the well-understood contraction of $\chi^{2}$ divergence of OU process, see e.g. [].
Specifically,
on the time interval $[0,T]$, $T=\gamma$,
the initial density $p_{n-1}$ renders $\chi^{2}(p_{n-1}\|q)<\infty$, because $p_{n-1}/q=\rho_{0}/q\in L^{2}(q)$ by the argument following (13).
For the OU process, since the equilibrium density $q\propto e^{-V}$ with $V(x)=\|x\|^{2}/2$,
we know that $q$ is strongly convex and then satisfies the Poincaré inequality (PI) with constant $C=1$.
By the argument in Eqn. (3) in [],
the PI implies the contraction claimed in the lemma.
∎

###### Proof of Lemma A.3 .

We prove the lemma by showing the closeness of $\chi^{2}(p_{n}\|q)=\hat{G}(T)$ to $\chi^{2}(p_{n}^{*}\|q)=G(T)$.
By definition (12),

$$ $G(t)=\|\frac{\rho_{t}}{q}-1\|_{L^{2}(q)}^{2},\quad\hat{G}(t)=\|\frac{\hat{\rho }_{t}}{q}-1\|_{L^{2}(q)}^{2},\quad\forall t\in[0,T].$ $$

We will use the triangle inequality $|\sqrt{G(t)}-\sqrt{\hat{G}(t)}|\leq\|\frac{\hat{\rho}_{t}}{q}-\frac{\rho_{t}}{
q}\|_{L^{2}(q)}$.
Observe that

$$ $\displaystyle\|\frac{\hat{\rho}_{t}}{q}-\frac{\rho_{t}}{q}\|_{L^{2}(q)}^{2}= \int\frac{(\hat{\rho}_{t}-\rho_{t})^{2}}{q}=\int\frac{(\hat{\rho}_{t}-\rho_{t} )^{2}}{\hat{\rho}_{t}}\frac{\hat{\rho}_{t}}{q}\leq\frac{C_{1}}{c_{d}}\int\frac {(\hat{\rho}_{t}-\rho_{t})^{2}}{\hat{\rho}_{t}},$ (15) $$

where the inequality is due to the pointwise bound
$\frac{\hat{\rho}_{t}(x)}{q(x)}\leq\frac{C_{1}}{c_{d}}$, which follows by (A1) and the expression of $q$ with $c_{d}$ as in (13).
We introduce

$$ $F(t):=\int\frac{(\hat{\rho}_{t}-\rho_{t})^{2}}{\hat{\rho}_{t}}=\int(\frac{\rho _{t}}{\hat{\rho}_{t}}-1)^{2}\hat{\rho}_{t}\geq 0,$ (16) $$

and then we have

$$ $|\sqrt{G(T)}-\sqrt{\hat{G}(T)}|\leq\|\frac{\hat{\rho}_{t}}{q}-\frac{\rho_{t}}{ q}\|_{L^{2}(q)}\leq(C_{1}/c_{d})^{1/2}\sqrt{F(T)}.$ (17) $$

Next, we will bound $F(T)$ to be $O(\varepsilon)$, and this will lead to the desired closeness of $\hat{G}(T)$ to $G(T)$.

$\bullet$ Bound of $F(T)$:

We will upper bound $F(T)$ by differentiating $F(t)$ over time. By definition (16),
$F(0)=0$, and (omitting t in $\rho_{t}$ and $\hat{\rho}_{t}$ in the equations)

$$ $\displaystyle\frac{d}{dt}F(t)$ $\displaystyle=\int(\frac{\rho}{\hat{\rho}}-1)^{2}\partial_{t}\hat{\rho}+2( \frac{\rho}{\hat{\rho}}-1)(\partial_{t}\rho-\frac{\rho}{\hat{\rho}}\partial_{t }\hat{\rho})$ $\displaystyle=\int 2\partial_{t}\rho(\frac{\rho}{\hat{\rho}}-1)-\partial_{t} \hat{\rho}((\frac{\rho}{\hat{\rho}})^{2}-1)$ $\displaystyle=\int-2\nabla\cdot(\rho v)(\frac{\rho}{\hat{\rho}}-1)+\nabla\cdot (\hat{\rho}\hat{v})((\frac{\rho}{\hat{\rho}})^{2}-1)\quad\text{(by transport equations \eqref{eq:TE-rho-hatrho})}$ $\displaystyle=\int 2(\rho v)\cdot\nabla(\frac{\rho}{\hat{\rho}})-(\hat{\rho} \hat{v})\cdot 2\frac{\rho}{\hat{\rho}}\nabla(\frac{\rho}{\hat{\rho}})$ $\displaystyle=2\int\rho(v-\hat{v})\cdot\nabla(\frac{\rho}{\hat{\rho}})$ $\displaystyle=2\int(v-\hat{v})\frac{\rho^{2}}{\hat{\rho}}\cdot(\nabla\log\rho- \nabla\log\hat{\rho}),$ $$

and the last row is by that
$\nabla(\frac{\rho}{\hat{\rho}})=\frac{\rho}{\hat{\rho}}(\nabla\log\rho-\nabla
\log\hat{\rho})$.
Thus,

$$ $\frac{1}{2}F(T)=\frac{1}{2}\int F^{\prime}(t)dt=\int_{0}^{T}\int_{\mathbb{R}^{ d}}(v-\hat{v})\cdot(\nabla\log\rho-\nabla\log\hat{\rho})\frac{\rho^{2}}{\hat{ \rho}}dxdt.$ $$

Then, by Cauchy-Schwarz, we have

$$ $\frac{1}{2}F(T)\leq\left(\int_{0}^{T}\int_{\mathbb{R}^{d}}\|v-\hat{v}\|^{2}{ \rho}dxdt\right)^{1/2}\left(\int_{0}^{T}\int_{\mathbb{R}^{d}}\|\nabla\log\rho- \nabla\log\hat{\rho})\|^{2}\frac{\rho^{3}}{\hat{\rho}^{2}}dxdt\right)^{1/2}.$ (18) $$

By the flow-matching error assumption (A0), we have the first factor in the r.h.s. of (18)
bounded by $\varepsilon$.
Meanwhile,
the technical conditions (A2)(A3) together imply that for all $t\in[0,T]$,

$$ $\int_{\mathbb{R}^{d}}\|\nabla\log\rho_{t}-\nabla\log\hat{\rho}_{t}\|^{2}({\rho _{t}^{3}}/{\hat{\rho}_{t}^{2}})dx\leq(2L)^{2}\int_{\mathbb{R}^{d}}(1+\|x\|)^{2 }({\rho_{t}^{3}}/{\hat{\rho}_{t}^{2}})(x)dx\leq(2L)^{2}C_{2},$ $$

thus the second factor in (18)
is upper bounded by $\sqrt{(2L)^{2}C_{2}T}$.
Putting together, we have

$$ $F(T)\leq 2(2L)\sqrt{C_{2}T}\varepsilon=C_{3}\varepsilon,\quad C_{3}:=4L\sqrt{C _{2}\gamma}.$ (19) $$

$\bullet$ Bound of $\hat{G}(T)-G(T)$:

With the bound (19) of $F(T)$, we are ready to go back to (17), which gives

$$ $\sqrt{\hat{G}(T)}\leq\sqrt{G(T)}+(C_{1}/c_{d})^{1/2}\sqrt{F(T)}.$ $$

Together with that $G(T)\leq(C_{1}/c_{d})^{2}$ by (14), this gives that

$$ $\hat{G}(T)\leq G(T)+2(C_{1}/c_{d})^{3/2}\sqrt{C_{3}}\varepsilon^{1/2}+(C_{1}/c _{d})C_{3}\varepsilon\leq G(T)+C_{4}\varepsilon^{1/2},$ $$

where we used that $\varepsilon<1$ by (A0) and

$$ $C_{4}:=2(C_{1}/c_{d})^{3/2}\sqrt{C_{3}}+(C_{1}/c_{d})C_{3}.$ (20) $$

The fact that $\hat{G}(T)\leq G(T)+C_{4}\varepsilon^{1/2}$ proves the lemma.
∎

###### Proof of Proposition 4.1 .

###### Lemma A.2 ( χ 2 superscript 𝜒 2 \chi^{2} italic_χ start_POSTSUPERSCRIPT 2 end_POSTSUPERSCRIPT -contraction of OU process) .

###### Lemma A.3 ( G ^ ⁢ ( T ) ≈ G ⁢ ( T ) ^ 𝐺 𝑇 𝐺 𝑇 \hat{G}(T)\approx G(T) over^ start_ARG italic_G end_ARG ( italic_T ) ≈ italic_G ( italic_T ) ) .

###### Proof of Lemma A.2 .

###### Proof of Lemma A.3 .

### A.2 Proofs in Section 4.3

###### Lemma A.4 (Bi-direction DPI) .

Let ${\rm D}_{f}$ be an $f$-divergence.
If $T:\mathbb{R}^{d}\to\mathbb{R}^{d}$ is invertible
and for two densities $p$ and $q$ on $\mathbb{R}^{d}$,
$T_{\#}p$ and $T_{\#}q$ also have densities, then

$$ $D_{f}(p||q)=D_{f}(T_{\#}p||T_{\#}q).$ $$

###### Proof of Lemma A.4 .

Let $X_{1}\sim p$, $X_{2}\sim q$, and
$Y_{1}=T(X_{1})$, $Y_{2}=T(X_{2})$.
Then $Y_{1}$ and $Y_{2}$ also have densities, $Y_{1}\sim\tilde{p}:=T_{\#}p$ and $Y_{2}\sim\tilde{q}:=T_{\#}q$.
By the classical DPI of $f$-divergence (see, e.g., the introduction of []),
we have $D_{f}(\tilde{p}||\tilde{q})\leq D_{f}(p||q)$.
In the other direction, $X_{i}=T^{-1}(Y_{i})$, $i=1,2$, then DPI also implies
$D_{f}(p||q)\leq D_{f}(\tilde{p}||\tilde{q})$.
∎

###### Proof of Theorem 4.2 .

Under the assumptions, Proposition 4.1 applies to give that

$$ $\chi^{2}(p_{N}\|q)\leq e^{-2\gamma N}\chi^{2}(p_{0}\|q)+\frac{C_{4}}{1-e^{-2 \gamma}}\varepsilon^{1/2}.$ $$

Then, whenever $e^{-2\gamma N}\chi^{2}(p_{0}\|q)\leq\varepsilon^{1/2}$ which is ensured by the $N$ in the theorem, we have
$\chi^{2}(p_{N}\|q)\leq(1+\frac{C_{4}}{1-e^{-2\gamma}})\varepsilon^{1/2}$.
Let $T_{1}^{N}:=T_{N}\circ\cdots\circ T_{1}$, which is invertible, and
$p_{N}=(T_{1}^{N})_{\#}p_{0}$, $q_{N}=(T_{1}^{N})_{\#}q_{0}$.
We will apply the bi-directional DPI Lemma A.4 to show that

$$ $\chi^{2}(p_{0}\|q_{0})=\chi^{2}((T_{1}^{N})_{\#}p_{0}\|(T_{1}^{N})_{\#}q_{0})= \chi^{2}(p_{N}\|q_{N})=\chi^{2}(p_{N}\|q),$ (21) $$

which then proves the theorem.

For Lemma A.4 to apply to show the first equality in (21),
it suffices to verify that $p_{0}$, $q_{0}$, $p_{N}$, $q_{N}$ all have densities.
$p_{0}$ has density by the theorem assumption.
From Definition A.1,
one can verify that
a transform $T$ being non-degenerate guarantees that $P$ has density $\Rightarrow T_{\#}P$ has density
(see, e.g., Lemma 3.2 of []).
Because all $T_{n}$ are non-degenerate under (A0),
from that $p_{0}$ has density we know that all $p_{n}$ has densities, including $p_{N}$.
In the reverse direction, $q_{N}=q$ which is the normal density.
By that $T_{n}^{-1}$ are all non-degenerate, we similarly have that all $q_{n}$ has densities, including $q_{0}$.
∎

###### Lemma A.5 .

For two densities $p$ and $q$ where $\chi^{2}(p\|q)<\infty$,
${\rm KL}(p\|q)\leq\chi^{2}(p\|q)$.

###### Proof.

The statement is a well-known fact and we include an elementary proof for completeness.
Because $\chi^{2}(p\|q)<\infty$, we have $(p/q-1)$ and thus $p/q\in L^{2}(q)$, i.e. $\int p^{2}/q<\infty$.
By the fact that $\log x\leq x-1$ for any $x>0$,
$\log\frac{p}{q}(x)\leq\frac{p}{q}(x)-1$, and then

$$ ${\rm KL}(p\|q)=\int p\log\frac{p}{q}\leq\int p(\frac{p}{q}-1)=\int\frac{p^{2}} {q}-1=\chi^{2}(p\|q).$ $$

∎

###### Proof of Corollary 4.3 .

Let $\rho_{t}=({\rm OU})_{0}^{t}(P)$.
Because $M_{2}(P)<\infty$, one can show that $\mathcal{W}_{2}(\rho_{t},P)^{2}=O(t)$ as $t\to 0$.
Specifically, by Lemma C.1 in [],
$\mathcal{W}_{2}(\rho_{t},P)^{2}\leq t^{2}M_{2}(P)+2td$.
Thus, for $t<1$,
$\mathcal{W}_{2}(\rho_{t},P)^{2}\leq(M_{2}(P)+2d)t$.
This proves $\mathcal{W}_{2}(P,\rho_{\delta})\leq C_{5}\delta^{1/2}$.

Meanwhile, $p_{0}=\rho_{\delta}$ satisfies the needed condition in Theorem 4.2,
the claimed bound of $\chi^{2}(\rho_{\delta}||q_{0})$ directly follows from Theorem 4.2.
∎

###### Lemma A.6 .

Suppose $P$ on $\mathbb{R}^{d}$ is compactly supported, then, $\forall t>0$, $\rho_{t}=({\rm OU})_{0}^{t}(P)$ satisfies satisfies (10) for some $C_{1}$, $C_{2}$ and $L$ (which may depend on $t$).

###### Proof of Lemma A.6 .

Suppose $P$ is supported on $B_{R}:=\{x\in\mathbb{R}^{d},\|x\|\leq R\}$ for some $R>0$.

By the property of the OU process, $\rho_{t}$ is the probability density of the random vector

$$ $Z_{t}:=e^{-t}X_{0}+\sigma_{t}Z,\quad Z\sim\mathcal{N}(0,I_{d}),\quad X_{0}\sim P ,\quad Z\perp X_{0}.$ (22) $$

where $\sigma_{t}^{2}=1-e^{-2t}$.
We will verity the first two inequalities in (10), and the third one is implied by the first one.
(The third one also has a direct proof:
$M_{2}(\rho_{\delta})=\mathbb{E}\|Z_{\delta}\|^{2}=e^{-2\delta}M_{2}(P)+\sigma_
{\delta}^{2}d<\infty$,
then the third condition in (10) holds
with $C_{2}=1+e^{-2\delta}M_{2}(P)+\sigma_{\delta}^{2}d$.)

The law of $Z_{t}$ in (22) gives that

$$ $\rho_{t}(x)=\int_{\mathbb{R}^{d}}\frac{1}{(2\pi\sigma_{t}^{2})^{d/2}}e^{-\|x-e ^{-t}y\|^{2}/(2\sigma_{t}^{2})}dP(y).$ (23) $$

This allows us to verity the first two inequalities in (10) with proper $C_{1}$ and $L$, making use of the fact that $P$ is supported on $B_{R}$.
Specifically, by definition,

$$ $\nabla\rho_{t}(x)=\int_{B_{R}}\frac{1}{(2\pi\sigma_{t}^{2})^{d/2}}\left(-\frac {x-e^{-t}y}{\sigma_{t}^{2}}\right)e^{-\|x-e^{-t}y\|^{2}/(2\sigma_{t}^{2})}dP(y),$ $$

and then, because $\|x-e^{-t}y\|\leq\|x\|+e^{-t}\|y\|\leq\|x\|+e^{-t}R$,

$$ $\|\nabla\rho_{t}(x)\|\leq\frac{\|x\|+e^{-t}R}{\sigma_{t}^{2}}\int_{B_{R}}\frac {1}{(2\pi\sigma_{t}^{2})^{d/2}}e^{-\|x-e^{-t}y\|^{2}/(2\sigma_{t}^{2})}dP(y)= \frac{\|x\|+e^{-t}R}{\sigma_{t}^{2}}\rho_{t}(x).$ $$

This means that, $\forall x\in\mathbb{R}^{d}$,

$$ $\frac{\|\nabla\rho_{t}(x)\|}{\rho_{t}(x)}\leq\frac{\|x\|+e^{-t}R}{\sigma_{t}^{ 2}},$ $$

which means that the 2nd condition in (10) holds with $L=\frac{1}{\sigma_{t}^{2}}\max\{1,e^{-t}R\}$.

To prove the first inequality, we again use the expression (23).
By that $\|x\|\leq\|x-e^{-t}y\|+\|e^{-t}y\|$, and that $\|y\|\leq R$, we have

$$ $\displaystyle\|x\|^{2}$ $\displaystyle\leq\|x-e^{-t}y\|^{2}+2\|x-e^{-t}y\|\|e^{-t}y\|+\|e^{-t}y\|^{2}$ $\displaystyle\leq\|x-e^{-t}y\|^{2}+2(\|x\|+e^{-t}R)e^{-t}R+e^{-2t}R^{2}$ $\displaystyle=\|x-e^{-t}y\|^{2}+2e^{-t}R\|x\|+3e^{-2t}R^{2},$ $$

and thus

$$ $\displaystyle e^{-\|x-e^{-t}y\|^{2}/(2\sigma_{t}^{2})}$ $\displaystyle\leq e^{-\frac{1}{2\sigma_{t}^{2}}(\|x\|^{2}-2e^{-t}R\|x\|-3e^{-2 t}R^{2})}$ $\displaystyle=e^{-\frac{1}{2\sigma_{t}^{2}}\|x\|^{2}+\frac{e^{-t}R}{\sigma_{t} ^{2}}\|x\|}e^{\frac{3e^{-2t}R^{2}}{2\sigma_{t}^{2}}}.$ $$

Inserting in (23), we have that

$$ $\rho_{t}(x)\leq\alpha_{t}e^{-\frac{1}{2\sigma_{t}^{2}}\|x\|^{2}+\frac{e^{-t}R} {\sigma_{t}^{2}}\|x\|},\quad\alpha_{t}:=\frac{e^{\frac{3e^{-2t}R^{2}}{2\sigma_ {t}^{2}}}}{(2\pi\sigma_{t}^{2})^{d/2}}.$ $$

For $\rho_{t}(x)\leq C_{1}e^{-\|x\|^{2}/2}$ to hold, it suffices to have $C_{1}$ s.t. $\forall x\in\mathbb{R}^{d}$,

$$ $\frac{C_{1}}{\alpha_{t}}\geq e^{\frac{1}{2}(1-\frac{1}{\sigma_{t}^{2}})\|x\|^{ 2}+\frac{e^{-t}R}{\sigma_{t}^{2}}\|x\|}.$ (24) $$

Because $t>0$, $1-\frac{1}{\sigma_{t}^{2}}=\frac{-e^{-2t}}{1-e^{-2t}}<0$,
the r.h.s. of (24) as a function on $\mathbb{R}^{d}$ decays faster than $e^{-c\|x\|^{2}}$ for some $c>0$ as $\|x\|\to\infty$,
and then the function is bounded on $\mathbb{R}^{d}$.
This means that there is $C_{1}>0$ to make (24) hold.
This proves that the first inequality in (10) holds for $\rho_{t}$.
∎

###### Lemma A.4 (Bi-direction DPI) .

###### Proof of Lemma A.4 .

###### Proof of Theorem 4.2 .

###### Lemma A.5 .

###### Proof.

###### Proof of Corollary 4.3 .

###### Lemma A.6 .

###### Proof of Lemma A.6 .

## Appendix B Details of LFM schedule and NLL computation

Figure: Algorithm A.1 Distillation of LFM

### B.1 Selection of LFM schedule

Given a positive integer $N\geq 1$, we specify $\{\gamma_{n}\}_{n=1}^{N}$ via the following scheme:

$$ $\gamma_{n}=\rho^{n-1}c,n=1,2,\ldots$ (25) $$

where the base time stamp $c$ and the multiplying factor $\rho$ are hyper-parameters.

### B.2 Negative log-likelihood (NLL) computation

Based on the instantaneous change-of-variable formula in neural ODE [], we know that for a trained flow model $\hat{v}(x(t),t;\theta)$
on $[0,1]$ that interpolates between $p$ and $q$, the log-likelihood of data $x\sim p$ can be expressed as

$$ $\log p(x)=\log q(x(1))+\int_{0}^{1}\nabla\cdot\hat{v}(x(s),s;\theta)ds,\quad x (0)=x,$ $$

where $\nabla\cdot\hat{v}(\cdot,s;\theta)$ is the trace of the Jacobian matrix of the network function.

Our trained LFM has $N$ sub-flows $\hat{v}_{n}(\cdot,t;\theta)$ for $n=1,\ldots,N$. The log-likelihood at test sample $x\sim p$ is then computed by

$$ $\log p(x)=\log q(x_{N})+\sum_{n=1}^{N}\int_{0}^{1}\nabla\cdot\hat{v}_{n}(x(s), s;\theta)ds,$ (26) $$

where starting from $x_{0}=x$,

$$ $x_{n}=x_{n-1}+\int_{0}^{1}\hat{v}_{n}(x(s),s;\theta)ds,\quad x(0)=x_{n-1}$ $$

Both the integration of $\hat{v}_{n}$ and $\nabla\cdot\hat{v}_{n}$ are computed using the numerical scheme of neural ODE.
We report NLL in $\log_{e}$ (known as “nats”) in all our experiments.

Figure: Figure A.1: Noise-to-image trajectories by LFM: Flowers (left) and LSUN Church (right).
Refer to caption: x2.png

## Appendix C Experimental details and additional results

To train LFM sub-flows, we use the trigonometric interpolant on 2D, tabular, and image generation experiments (Sections 6.1–6.3) and use the optimal transport interpolant on robotic manipulation (Section 6.4).

### C.1 2D, Tabular, and Image experiments

In all experiments, we use the Adam optimizer [] with the following parameters: $\beta_{1}=0.9,\beta_{2}=0.999,\epsilon=1e-8$; the learning rate is to be specified in each experiment.
Additionally, the number of training batches indicates how many batches pass through all $N$ sub-flows per Adam update.

On two-dimensional datasets and tabular datasets, we parameterize local sub-flows with fully-connected networks; the dataset detail and hyper-parameters of LFM are in Table A.1.

On image generation examples, we parameterize local sub-flows as UNets [], where the dataset details and training specifics are provided in Table A.2.

Figure: a LFM before distillation.
Refer to caption: extracted/5897716/ImageNet32_no_distill.png

Figure: a LFM distillation @ 4 NFEs. FID = 71.0.
Refer to caption: extracted/5897716/Flower_distill_ours.png

### C.2 Robotic Manipulation

In the context of generative modeling, this task of robotic manipulation can be understood as performing sequential conditional generation. Specifically, at each time step $t\geq 1$, the goal is to model the conditional distribution $A_{t}|O_{t}$, where $O_{t}\in\mathbb{R}^{O}$ denotes the states of the robots at time $t$ and $A_{t}\in\mathbb{R}^{A}$ is the action that controls the robots. During inference, the robot is controlled as we iteratively sample from $A_{t}|O_{t}$ across time steps $t$.
Past works leveraging diffusion models have reached state-of-the-art performances on this task [], where a neural network $v_{\theta}$ (e.g., CNN-based UNet []) is trained to approximate the distribution $A_{t}|O_{t}$ via DDPM []. More recently, flow-based methods have also demonstrated competitive performance with faster inference [].

We use the widely adopted success rate to examine the performance of a robot manipulator:

$$ $\text{Success rate}=\frac{\#\text{success rollouts}}{\#\text{rollouts}}\in[0,1].$ (27) $$

Specifically, starting from a given initial condition $O_{1}$ of the robot, each rollout denotes a trajectory $S=\{O_{1},A_{1},O_{2},\ldots,A_{T},O_{T}\}$ where $A_{t}|O_{t}$ is modeled by the generative. The rollout $S$ is a success if at any $t\in{1,\ldots,T}$, the robotic state $O_{t}$ meets the success criterion (e.g., successfully pick up the square as in the task “lift” in Figure 4).

**Table A.1: Hyperparameters and architecture for two-dimensional datasets and tabular datasets. The table is formatted similarly as [, Table 3].**
|  | Rose | Fractal tree | POWER | GAS | MINIBOONE | BSDS300 |
| --- | --- | --- | --- | --- | --- | --- |
| Dimension | 2 | 2 | 6 | 8 | 43 | 63 |
| # Training point | 2,000,000 | 2,000,000 | 1,615,917 | 852,174 | 29,556 | 1,000,000 |
| Batch Size | 10K | 10K | 30K | 50K | 1000 | 500 |
| Training Batches | 50K | 50K | 100K | 100K | 100K | 30K |
| Hidden layer width<br>(per sub-flow) | 256 | 256 | 256 | 362 | 362 | 512 |
| $\#$ Hidden layers | 3 | 3 | 4 | 5 | 4 | 4 |
| Activation | Softplus | Softplus | ReLU | ReLU | ReLU | ELU |
| $\#$ sub-flows $N$ | 9 | 9 | 4 | 2 | 2 | 4 |
| $(c,\rho)$ in (25) | (0.025, 1.25) | (0.025, 1.25) | (0.15, 1.3) | (0.05, 1) | (0.35, 1) | (0.25, 1) |
| Total $\#$ parameters in M<br>(all sub-flows) | 1.20 | 1.20 | 0.81 | 1.06 | 0.85 | 3.41 |
| Learning Rate (LR) | $0.0002$ | $0.0002$ | $0.005$ | $0.002$ | $0.005$ | $0.002$ |
| LR decay<br>(factor, frequency in batches) | (0.99, 1000) | (0.99, 1000) | (0.99, 1000) | (0.99, 1000) | (0.9, 4000) | (0.8, 4000) |
| Beta $\alpha,\beta$, time samples | (1.0, 1.0) | (1.0, 1.0) | (1.0, 1.0) | (1.0, 0.5) | (1.0, 1.0) | (1.0, 1.0) |

**Table A.2: Hyperparameters and architecture for image datasets. The table is formatted similarly as [, Table 4].**
|  | CIFAR-10 | Imagenet-32 | Flowers | LSUN Churches |
| --- | --- | --- | --- | --- |
| Dimension | 32$\times$32 | 32$\times$32 | 128$\times$128 | 128$\times$128 |
| # Training point | 50,000 | 1,281,167 | 8,189 | 126,227 |
| Batch Size | 200 | 256 | 40 | 40 |
| Training Batches | 5 $\times$ 10^4 | 2 $\times$ 10^5 | 4 $\times$ 10^4 | 1.2 $\times$ 10^5 |
| Hidden dim (per sub-flow) | 128 | 114 | 128 | 128 |
| $\#$ sub-flows $N$ | 4 | 5 | 4 | 4 |
| $(c,\rho)$ in (25) | (0.3, 1.1) | (0.3, 1.1) | (0.5, 1.5) | (0.4, 1.5) |
| Total $\#$ parameters in M (all sub-flows) | 160 | 120 | 464 | 464 |
| Learning Rate (LR) | $0.0001$ | $0.0001$ | $0.0002$ | $0.0002$ |
| U-Net dim mult | [1,2,2,2,2] | [1,2,2,2] | [1,1,2,3,4] | [1,1,2,3,4] |
| Beta $\alpha,\beta$, time samples | (1.0, 1.0) | (1.0, 1.0) | (1.0, 1.0) | (1.0, 1.0) |
| Learned $t$ sinusoidal embedding | Yes | Yes | Yes | Yes |
| $\#$ GPUs | 1 | 1 | 1 | 1 |

**Table A.3: Hyperparameters and architecture for robotic manipulation under state-based environment on Robomimic [].**
|  | Lift | Can | Square | Transport | Toolhang |
| --- | --- | --- | --- | --- | --- |
| Batch Size | 256 | 256 | 256 | 256 | 256 |
| Training Epochs | 200 | 500 | 750 | 1500 | 200 |
| Hidden dims<br>(per sub-flow) | [128,256,512] | [128,256,512] | [128,256,512] | [176,352,704] | [128,256,512] |
| $\#$ sub-flows $N$ | 4 | 4 | 4 | 2 | 4 |
| $(c,\rho)$ in (25) | (0.15, 1.25) | (0.2, 1.25) | (0.2, 1) | (0.5, 1) | (0.25, 1.25) |
| Total $\#$ parameters in M<br>(all sub-flows) | 66 | 66 | 66 | 67 | 67 |
| Learning Rate (LR) | $0.0001$ | $0.0001$ | $0.0001$ | $0.0001$ | $0.0001$ |
| $\#$ GPUs | 1 | 1 | 1 | 1 | 1 |

We also describe details of each of the 5 Robomimic tasks below, including dimensions of observations $O_{t}$ and actions $A_{t}$ and the success criteria. The initial condition and final successful completion were shown in Figure 4. Table A.3 contains the hyper-parameter setting in each task, where we use the same network and training procedure as in [].

Lift: The goal is for the robot arm to lift a small cube in red. Each $O_{t}$ has dimension $16\times 19$ and each $A_{t}$ has dimension $16\times 10$, representing state-action information for the next 16 time steps starting at $t$.

Can: The goal is for the robot to pick up a coke can from a large bin and place it into a smaller target bin. Each $O_{t}$ has dimension $16\times 23$ and each $A_{t}$ has dimension $16\times 10$, representing state-action information for the next 16 time steps starting at $t$.

Square: The goal is for the robot to pick up a square nut and place it onto a rod with precision. Each $O_{t}$ has dimension $16\times 23$ and each $A_{t}$ has dimension $16\times 10$, representing state-action information for the next 16 time steps starting at $t$.

Transport: The goal is for the two robot arms to transfer a hammer from a closed container on one shelf to a target bin in another shelf. Before placing the hammer, one arm has to also clear the target bin by moving away a piece of trash to the nearby receptacle. The hammer must be picked up by one arm which then hands over to the other. Each $O_{t}$ has dimension $16\times 59$ and each $A_{t}$ has dimension $16\times 20$, representing state-action information for the next 16 time steps starting at $t$.

Toolhang: The goal is for the robot to assemble a frame that includes a base piece and a hook piece by inserting the hook into the base. The robot must then hang a wrench on the hook. Each $O_{t}$ has dimension $16\times 53$ and each $A_{t}$ has dimension $16\times 10$, representing state-action information for the next 16 time steps starting at $t$.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: What does $g \bot x'_l$ mean in Eq. 6? and why do we need this constraint?
soundness: 3
strengths: The authors proposed the Local Flow Matching (LFM), which divides the Flow Matching into a sequence of local parts. The approach is intuitively correct. The authors provided theoretical analysis of the convergence of the forward process, showing that the $\chi^2$ divergence between the noise distribution and noised real data distribution decreases exponentially w.r.t. the number of sub-flows. Furthermore, the authors presented the theoretical results showing the $\chi^2$ divergence between the real distribution and the generated distribution is close w.r.t. the flow training loss $\epsilon$.
summary: This paper proposes to divide the original Flow Matching (FM) into a sequence of parts, called Local Flow Matching (LFM). The authors provide theoretical analysis of the convergence of the forward process, and the theoretical guarantee of the $\chi^2$ divergence between the real data and generated data distributions. The authors evaluated the proposed method on different tasks, tabular data generation, image generation, robotic manipulation and so on.
weaknesses: From the experiments, it turns out that the proposed LFM is not working as well as the original Flow Matching (FM) (Lipman et al., 2023). The proposed method achieved an FID of 8.45 on the CIFAR-10 dataset whereas the FM achieved a lower FID of 6.35 on the same dataset. On ImageNet 32x32 dataset, the proposed method achieved an FID of 7.0, but the FM achieved a lower FID of 5.02. The authors should compare with FM on the CIFAR-10, ImageNet 32x32, the Oxford Flowers and the LSUN Church datasets. 

The proposed LFM seems to require more training time than the original FM because for training $v_n$ we need to inference from $v_1$ to $v_{n-1}$ for each sample. 

The proposed method also need more inference time since we need to inference from $v_1$ to $v_N$.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. Can you clarify more on how the method benefits fast training based on the experimental results?
soundness: 2
strengths: The paper explores integrating flow-matching submodels into diffusion processes to enable faster and more efficient learning and inference. The approach is novel and is supported by theoretical guarantees of generation. Experimentally, the method can be applied to various tasks, including image generation, tabular data generation, and robotic manipulation.
summary: This paper presents Local Flow Matching (LFM), an enhancement of Flow Matching (FM) for faster, efficient generative modeling. LFM divides the process into sequential FM sub-models, each bridging distributions that are progressively closer from data to noise, allowing smaller models and quicker training. This stepwise approach also enables the use of distillation techniques to further accelerate data generation. Theoretically, LFM provides a generation guarantee based on $\chi^2$-divergence. Experiments show LFM’s improved training efficiency and strong performance in generating tabular data, images, and robotic manipulation policies.
weaknesses: 1. More details should be provided, such as the number of function evaluations (NFEs) and the used ODE sampler for all methods in Table 1 and Table 2, to better demonstrate LFM's effectiveness.

2. Although the tasks are diverse, the paper lacks a solid comparison with prior methods on some fundamental tasks. For instance, comparing LFM with Rectified Flow and OU diffusion on CIFAR-10 with the same amount of batches would offer a clearer understanding of LFM's training efficiency and distillation effects.

3. The advantage of LFM on fast training has not been demonstrated well except Table 2.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: 1.	In the training process, at each step, the target distribution is generated from the sample of the previous step through the Ornstein-Uhlenbeck (OU) process. What is the theoretical basis for using the OU process in this context? This essentially defines checkpoints along the probability path. Would a different process yield similar results, or is there a specific reason the OU process was chosen?

2.	In Appendix B1, how were the stepwise parameters $\rho$ and $c$ selected? Were they chosen solely based on extensive experimentation and selecting the ones that gave the best results, or was there a more principled approach to their selection?

3.	In the image generation tasks (Table 2), it would be beneficial to include comparisons with Flow Matching (FM) [1], rather than only with InterFlow, as Local Flow Matching (LFM) is based on FM. This would better illustrate the claimed improved training efficiency. Additionally, either in the main text or in the appendix, it would be helpful to provide a comparison of the total number of parameters for each method across all experiments, rather than only ensuring the same training scheme.

[1] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference on Learning Representations, 2023.
soundness: 2
strengths: 1.	The paper offers a fresh approach in the field of generative modeling, combining ideas from diffusion processes and FM. The idea of breaking down a single large flow into several smaller steps (local flows) is natural.
2.	The paper provides solid theoretical guarantees, specifically proving a generation guarantee in terms of the $\chi^2$-divergence between the generated and true data distributions. The experiments are comprehensive, covering a range of tasks from image generation to robotic manipulation, and the results demonstrate the performance of LFM.
summary: The paper introduces Local Flow Matching (LFM), a generative approach that builds upon the existing Flow Matching (FM) framework. LFM improves upon FM by dividing the global flow into multiple local flow sub-models. Each sub-model is trained to match a diffusion process between closer intermediate distributions, which reduces the model size and accelerates training.
weaknesses: 1.	FM simplifies the diffusion process into a single step, transforming a trajectory from a curve into a straight line, thus reducing training costs and improving sampling efficiency. However, this paper reverses that advantage by breaking this single step into multiple segments. The division of the trajectory into multiple steps could introduce added computational complexity, potentially negating the efficiency gains that FM originally aimed to provide.

2.	The paper presents some conceptual ambiguities regarding key terms. Continuous Normalizing Flows (CNF) and neural ODEs are two distinct models, but the paper incorrectly states that "CNF trains a neural ODE" (Line 106). In fact, FM is a method for training CNF models, not a model in itself, much like the proposed LFM method. The paper conflates these concepts in several places (e.g. Line 41, 51), which could cause confusion.

3.	The paper highlights improved training efficiency through the use of smaller local models. However, the exact speed improvements are not always well-quantified. For example, while the paper states that LFM leads to faster convergence and reduced memory usage, these claims could be supported by more detailed analysis. Providing concrete timing metrics (e.g., runtime comparisons, memory benchmarks across different hardware) would bolster the argument for LFM's practical benefits. More information can be found in Questions part.

4.	In line 127, a reference is used as the subject of the sentence but is enclosed in parentheses. Please correct this formatting issue and ensure consistency throughout the paper.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Could you provide more detailed data regarding the issues mentioned in the weaknesses, such as the specific number of parameters for each sub-model in LFM, the number of sub-models, and the parameter count of the InterFlow model (or explain why the total number of parameters across all sub-models is equal to that of InterFlow). Could you provide the amount of data required to train each sub-model?
soundness: 4
strengths: The paper provides a detailed theoretical proof that the data distribution generated by LFM is close to the true data distribution, demonstrating the exponential convergence of the forward process and the generation guarantee of the backward process. In addition, the paper conducts extensive experiments to validate the effectiveness of LFM, including generation on two-dimensional toy data, tabular data generation, image generation, and robotic manipulation. Moreover, the paper is well-written.
summary: The paper proposes a novel generative model, Local Flow Matching (LFM), which segments Flow Matching along the time dimension. The transition between two distributions within each time segment is modeled using a smaller-scale sub-model. This generative model offers the advantages of having smaller sub-models, faster convergence during training, and greater convenience for distillation.
weaknesses: Some experimental details in the paper may need further clarification. The data in Table 2 is used to demonstrate that LFM requires less computational cost, but the model sizes of both LFM and InterFlow are not provided. Additionally, it is unclear whether the training data shown in the table refers to the data used by each sub-model individually or by all sub-models collectively.

Meta-review: This paper suggests training a sequence of flow models by breaking the interval of interest to several parts and training each part using Flow Matching to match marginal probabilities defined by an OU process, in particular Variance Preserving (VP). 

It is not clear whether breaking the generation problem to a collection of smaller problems (here with VP marginal probability paths) indeed simplifies the problem and/or provides practical benefits, e.g., compared to training a flow directly using this VP marginal probability path over the full interval. The main difference is that the current approach in practice adds data couplings at each interval compared to the global coupling used for the original problem. The experimental part falls short in showing the suggested method has a strong practical benefit in terms of performance, faster training and/or sampling complexity. Furthermore, some implementation details are missing in the experimental part.

Decision: Reject