Lagrangian Flow Matching: 基于最小作用原理的规范路径设计框架

# 拉格朗日流匹配：基于最小作用原理的规范路径设计框架
来源：https://arxiv.org/html/2605.15419
杜书凯¹ 张俊哲² 李一鸣¹
¹ 雪城大学数学系 ² 雪城大学电子工程与计算机科学系
{sdu113,jzhan403,yli658}@syr\.edu

###### 摘要

流匹配通过回归目标速度来训练神经速度场，该目标速度与连接简单初始分布到数据分布的预设概率路径相关联。核心设计选择在于路径本身。现有构造（包括纠正流和基于最优传输的路径）在耦合端点之间沿直线传输样本，因此仅覆盖狭窄的动力学类别。我们观察到，这对应于经典力学中*最小作用原理*的最简单情形，其中动能拉格朗日量 $\mathcal{L}=\tfrac{1}{2}\|\dot{x}\|^2$ 产生自由粒子直线轨迹。基于此观察，我们提出*拉格朗日流匹配*，这是一个基于物理的框架，其中概率路径和速度场通过最小化一般拉格朗日量的作用量确定，需满足连续性方程和预设端点。我们证明该动态问题等价于一个静态最优传输（OT）表述，从而产生一族无模拟训练目标，其中基于OT的流匹配作为动能特例被恢复，而三角方差保持扩散路径作为谐振子特例被恢复。更一般的拉格朗日量可衍生出新的概率路径和速度场，数值实验表明它们在保持与现有条件流匹配模型竞争力的同时，诱导了学习动力学中的有意义变化。

## 1 引言

生成建模的一种自然方法是从连续时间视角出发：不是直接建模数据分布，而是构造一个随时间变化的概率路径 $(p_t)_{t\in[0,1]}$，将简单初始分布 $p_{\mathrm{init}}$ 连续传输到目标分布 $p_{\mathrm{data}}$。基于此视角的两类重要生成模型是扩散模型（Sohl-Dickstein 等人，2015 (https://arxiv.org/html/2605.15419#bib.bib5)；Ho 等人，2020 (https://arxiv.org/html/2605.15419#bib.bib4)；Song 等人，2021b (https://arxiv.org/html/2605.15419#bib.bib6)）和基于流的模型（Chen 等人，2018 (https://arxiv.org/html/2605.15419#bib.bib18)；Grathwohl 等人，2019 (https://arxiv.org/html/2605.15419#bib.bib19)；Lipman 等人，2023 (https://arxiv.org/html/2605.15419#bib.bib8)；Liu 等人，2023 (https://arxiv.org/html/2605.15419#bib.bib11)；Tong 等人，2024 (https://arxiv.org/html/2605.15419#bib.bib13)），它们分别通过随机微分方程（SDE）和常微分方程（ODE）实现传输。我们关注基于流的模型。早期有影响的一类是连续归一化流（CNFs）（Chen 等人，2018 (https://arxiv.org/html/2605.15419#bib.bib18)；Grathwohl 等人，2019 (https://arxiv.org/html/2605.15419#bib.bib19)），它们具有足够的表达能力来表示广泛的连续概率路径（Song 等人，2021a (https://arxiv.org/html/2605.15419#bib.bib7)），但训练计算成本高昂：基于似然的训练需要重复的前向和反向ODE求解。*流匹配*提供了一种高效替代方案，通过直接回归神经速度场 $v_\theta(x,t)$ 到与*预设概率路径*相关联的目标速度场 $v_t(x)$（Lipman 等人，2023 (https://arxiv.org/html/2605.15419#bib.bib8)；Liu 等人，2023 (https://arxiv.org/html/2605.15419#bib.bib11)；Tong 等人，2024 (https://arxiv.org/html/2605.15419#bib.bib13)）。通过避免训练中的ODE求解，这种无模拟目标使基于流的模型能够扩展到高维数据和大规模应用。

然而，流匹配模型的性能不仅取决于回归目标本身，还取决于预设概率路径 $(p_t)$ 及其关联的速度场 $v_t$。现有方法从相对狭窄的设计空间中选择这些对象：条件高斯路径（Lipman 等人，2023 (https://arxiv.org/html/2605.15419#bib.bib8)；Tong 等人，2024 (https://arxiv.org/html/2605.15419#bib.bib13)）、耦合端点之间的仿射插值（Liu 等人，2023 (https://arxiv.org/html/2605.15419#bib.bib11)）以及基于最优传输的直线插值（Tong 等人，2024 (https://arxiv.org/html/2605.15419#bib.bib13)）。这引发了一个自然问题：流匹配中使用的概率路径和速度场能否从更一般、更规范的设计空间中选择？

| 动能拉格朗日量 $\mathcal{L}=\tfrac{1}{2}\|\dot{x}\|^2$ | 一般拉格朗日量 $\mathcal{L}(x,\dot{x},t)=K(\dot{x})-V(x)$ |
| --- | --- |
| 确定性轨迹 | 概率路径 |
| 自由粒子轨迹 | 见说明图 |
| $\gamma(t)=(1-t)x_0+tx_1$ | 最小作用轨迹 |
| $\partial\mathcal{L}/\partial x=\frac{d}{dt}\partial\mathcal{L}/\partial\dot{x}$ | |
| 基于OT的流匹配 | 见说明图 |
| 拉格朗日流匹配 | 见说明图 |
| 静态-动态等价（定理 3.1） | |

图1：拉格朗日流匹配通过最小作用原理统一了轨迹选择（上行）和概率路径构造（下行），并由静态-动态等价（定理 ̃2.1 (https://arxiv.org/html/2605.15419#S2.Thmtheorem1)）连接。动能拉格朗日量（左列）产生直线轨迹并恢复基于OT的流匹配；一般拉格朗日量（右列）产生弯曲轨迹并诱导新的概率路径。

我们通过最小作用原理（经典力学中的基础变分原理）的视角来回答这个问题（Arnold 等人，1989 (https://arxiv.org/html/2605.15419#bib.bib40)）。给定拉格朗日量 $\mathcal{L}(x,\dot{x},t)$，该原理在固定端点之间选择最小化相应作用量的轨迹。拉格朗日量的选择完全决定了动力学：动能情形产生直线自由粒子轨迹（图 ̃1 (https://arxiv.org/html/2605.15419#S1.F1)，左上），而一般拉格朗日量产生由欧拉-拉格朗日方程支配的弯曲路径（图 ̃1 (https://arxiv.org/html/2605.15419#S1.F1)，右上）。同样的原理自然适用于流匹配。基于仿射和基于OT路径的现有构造在耦合端点之间沿直线轨迹传输样本（图 ̃1 (https://arxiv.org/html/2605.15419#S1.F1)，左下），对应于动能情形。用更一般的拉格朗日量替换动能拉格朗日量，可在匹配端点之间产生弯曲的最小作用轨迹，进而诱导新的概率路径（图 ̃1 (https://arxiv.org/html/2605.15419#S1.F1)，右下）。这种从轨迹到概率路径的转移通过最优传输中的经典静态-动态等价得到严格化（Benamou and Brenier, 2000 (https://arxiv.org/html/2605.15419#bib.bib20)；Bernard and Buffoni, 2007 (https://arxiv.org/html/2605.15419#bib.bib22)；Villani, 2009 (https://arxiv.org/html/2605.15419#bib.bib1)）。近期几项工作也探索了用于学习传输动力学的变分或最小作用公式（Neklyudov 等人，2023 (https://arxiv.org/html/2605.15419#bib.bib15), 2024 (https://arxiv.org/html/2605.15419#bib.bib14)；Pooladian 等人，2024 (https://arxiv.org/html/2605.15419#bib.bib16)；Koshizuka and Sato, 2022 (https://arxiv.org/html/2605.15419#bib.bib27)）。然而，相对较少的工作旨在直接将最小作用动力学表述为流匹配问题。我们旨在通过将拉格朗日视角引入流匹配来填补这一空白：拉格朗日量定义端点代价、选择耦合、并构造显式最小作用路径。这些路径随后用于定义标准流匹配回归目标。在此意义上，我们的方法旨在将最小作用诱导的概率路径的灵活性与流匹配中无模拟回归的简洁高效相结合。

基于这一视角，我们提出*拉格朗日流匹配*，一个框架，其中概率路径和速度场通过最小化服从连续性方程和端点约束 $p_0=p_{\mathrm{init}}$、$p_1=p_{\mathrm{data}}$ 的拉格朗日作用量确定。我们的贡献有三点。首先，我们提出了一个基于物理的框架，通过最小作用原理选择流匹配概率路径和速度场。其次，我们推导了一族无模拟训练目标，并展示了若干现有构造（包括基于OT的和条件流匹配模型）作为特例被恢复。第三，我们从更一般的拉格朗日量推导出新的流匹配模型，特别是一个由频率 $\omega\in(0,\pi)$ 索引的*谐振族*，它连续变形最小作用路径，并在其两个端点处分别恢复基于OT的流匹配和三角方差保持调度。我们在合成二维数据、单细胞轨迹插值和CIFAR-10图像生成上验证了该框架，并表明不同拉格朗日量在学习动力学中诱导有意义的变化，同时与现有条件流匹配模型保持竞争力。

### 1.1 预备知识与相关工作

通常，$p_{\mathrm{init}}$ 和 $p_{\mathrm{data}}$ 是 $\mathbb{R}^d$ 上的概率密度。生成建模的目标是学习一个将 $p_{\mathrm{init}}$ 推送到 $p_{\mathrm{data}}$ 的变换，给定 $p_{\mathrm{data}}$ 的有限样本。该设定包括标准情形（$p_{\mathrm{init}}$ 易于采样，例如标准高斯分布）和更一般的情形（仅通过样本访问 $p_{\mathrm{init}}$ 和 $p_{\mathrm{data}}$）（Albergo and Vanden-Eijnden, 2023 (https://arxiv.org/html/2605.15419#bib.bib9)；Liu 等人，2023 (https://arxiv.org/html/2605.15419#bib.bib11)；Tong 等人，2024 (https://arxiv.org/html/2605.15419#bib.bib13)）。我们关注通过概率路径 $(p_t)_{t\in[0,1]}$（满足 $p_0=p_{\mathrm{init}}$、$p_1=p_{\mathrm{data}}$）实现此变换的连续时间方法。

**流匹配** 训练参数化速度场 $v_\theta(x,t)$，通过回归到目标场 $v_t(x)$：
$$
\min_\theta\; \mathbb{E}_{t\sim\mathcal{U}[0,1],\,x\sim p_t}\bigl\|v_\theta(x,t)-v_t(x)\bigr\|^2,
\tag{1}
$$
其中 $(p_t)_{t\in[0,1]}$ 满足端点条件 $p_0=p_{\mathrm{init}}$、$p_1=p_{\mathrm{data}}$，且 $v_t$ 通过连续性方程 $\partial_t p_t+\nabla\cdot(p_t v_t)=0$ 生成该路径。无穷多对 $(p_t,v_t)$ 满足这些条件，模型的行为取决于选择。

**条件流匹配** 通过将 $p_t$ 表示为关于条件变量 $z\sim q$ 的混合 $p_t(x)=\mathbb{E}_{z\sim q}[p_t(x\mid z)]$ 来解决此不可解性（Lipman 等人，2023 (https://arxiv.org/html/2605.15419#bib.bib8)；Tong 等人，2024 (https://arxiv.org/html/2605.15419#bib.bib13)）。若 $v_t(x\mid z)$ 生成 $p_t(\cdot\mid z)$，则边际速度是后验平均值 $v_t(x)=\mathbb{E}[v_t(x\mid z)\mid x_t=x]$。边际目标 (1 (https://arxiv.org/html/2605.15419#S1.E1)) 等于（相差一个加性常数）：
$$
\min_\theta\; \mathbb{E}_{t\sim\mathcal{U}[0,1],\,z\sim q,\,x\sim p_t(\cdot\mid z)}\bigl\|v_\theta(x,t)-v_t(x\mid z)\bigr\|^2,
\tag{2}
$$
当条件路径和速度具有闭式表达式时，该目标是可解的。一个代表性示例是 Lipman 等人 (2023 (https://arxiv.org/html/2605.15419#bib.bib8)) 的仿射构造：取 $z=y\sim p_{\mathrm{data}}$、$p_{\mathrm{init}}=\mathcal{N}(0,I)$，条件轨迹为 $\psi_t(x\mid y)=(1-(1-\varepsilon)t)x+ty$。我们将在 $\omega\to 0$ 极限下以条件拉格朗日流匹配恢复此构造（示例 ̃2.4 (https://arxiv.org/html/2605.15419#S2.Thmexample4)）。

**动态最优传输**。此处发展的最小作用视角建立在 Benamou and Brenier (2000 (https://arxiv.org/html/2605.15419#bib.bib20)) 的动态最优传输公式之上。$p_0$ 和 $p_1$ 之间的平方 $2$-Wasserstein 距离具有表示：
$$
W_2^2(p_0,p_1)=\inf_{(p_t,v_t)}\int_0^1\!\int_{\mathbb{R}^d}\|v_t(x)\|^2\,p_t(x)\,dx\,dt,
\tag{3}
$$
其中下确界取遍满足连续性方程和端点约束 $p_{t=0}=p_0$、$p_{t=1}=p_1$ 的对 $(p_t,v_t)$。这是下文发展的拉格朗日框架的动能特例，也是基于OT的流匹配的基本观点（Tong 等人，2024 (https://arxiv.org/html/2605.15419#bib.bib13)）。我们请读者参考附录 ̃A (https://arxiv.org/html/2605.15419#A1) 以获取更详细的相关工作讨论。

## 2 拉格朗日流匹配

本节介绍我们的主要框架。我们首先回顾拉格朗日力学中陈述最小作用原理所需的要素，然后制定*拉格朗日最优传输*问题并建立将其转化为可解训练目标的静态-动态等价（第 ̃2.1 (https://arxiv.org/html/2605.15419#S2.SS1) 节）。使框架可扩展至大规模的条件公式将在第 ̃3 (https://arxiv.org/html/2605.15419#S3) 节中发展。

| 见说明图 |
| --- |
| 图2：谐波最小作用轨迹 (4 (https://arxiv.org/html/2605.15419#S2.E4)) 的位置 $\gamma^\omega(t)$ 和速度 $\dot{\gamma}^\omega(t)$，$\omega\in(0,\pi)$，端点为 $x_0=0$、$x_1=1$。 |

拉格朗日力学通过作用量泛函提供了选择动力学的变分原理，并长期作为描述物理系统的统一语言（Arnold 等人，1989 (https://arxiv.org/html/2605.15419#bib.bib40)）。*拉格朗日量*是位置 $x\in\mathbb{R}^d$、速度 $v\in\mathbb{R}^d$ 和时间 $t\in[0,1]$ 的函数 $\mathcal{L}(x,v,t)$，编码系统的动力学。给定光滑曲线 $\gamma:[0,1]\to\mathbb{R}^d$，其*作用量*为：
$$
S[\gamma] = \int_0^1 \mathcal{L}\bigl(\gamma(t),\dot{\gamma}(t),t\bigr)\,dt.
$$
*最小作用原理*在两个固定端点之间选择