基于Flow Matching的概率反演

# 基于流匹配的概率反演  
来源：https://arxiv.org/html/2606.31288  

###### 摘要  
我们展示了流匹配（Flow Matching）——一种源自生成式人工智能的技术——在地球物理概率反演（如地震全波形反演）中的应用。我们将生成式人工智能中已完善的流匹配数学理论适配到概率反演场景中。通过两个案例研究评估该方法：一个简单的二维速度模型用于说明方法的一般特征，以及OpenFWI数据集用于展示该方法在更复杂地震速度模型概率反演中的能力。  

## 1 引言  
令 \(\mathbf{m} \in \mathcal{M}\) 定义一个地球物理模型（例如，地下地震速度模型及特定采集几何）。令 \(\mathbf{d}_{\text{obs}}\) 为与 \(\mathbf{m}\) 关联的观测数据。正演算子 \(f: \mathcal{M} \rightarrow \mathcal{D}\) 将模型空间 \(\mathcal{M}\) 映射到数据空间 \(\mathcal{D}\)，描述底层物理系统，使得 \(f(\mathbf{m}) = \mathbf{d}_{\text{syn}}\)，其中 \(\mathbf{d}_{\text{syn}} \in \mathcal{D}\)，且 \(\mathbf{d}_{\text{obs}} = \mathbf{d}_{\text{syn}} + \varepsilon\)。因此，观测数据 \(\mathbf{d}_{\text{obs}}\) 对应合成数据 \(\mathbf{d}_{\text{syn}}\) 加上测量中少量不相关噪声 \(\varepsilon > 0\)。**正演问题** 指确定正演算子 \(f\) 的任务。**反问题** 则寻求逆算子 \(f^{-1}: \mathcal{D} \rightarrow \mathcal{M}\)，将数据空间映射到模型空间，从而从观测数据中恢复模型参数（即地震速度分布）。在确定性物理学中，正演算子有唯一解 [14 (https://arxiv.org/html/2606.31288#bib.bib14)]。然而，逆算子通常并非如此。尽管如此，大多数现代反演框架采用**确定性反演**，它实用地假设 \(f^{-1}\) 是唯一映射，因为这样计算上更容易求解。由于反问题是病态的（与适定问题相反，适定问题存在唯一解且行为随输入连续变化），确定性反演需要一系列简化才能获得可靠方法。首先，需要给定解的初始猜测作为迭代求解器的起点。求解器需要对非线性正演算子使用泰勒级数近似以计算模型更新。此外，需要显著的正则化，通常通过施加光滑性约束。最后，观测数据可能需要预处理，例如在地震数据的全波形反演中，通常对递增的频率内容进行多次反演以避免周期跳跃 [1 (https://arxiv.org/html/2606.31288#bib.bib16),15 (https://arxiv.org/html/2606.31288#bib.bib4)]。

**概率反演** 试图通过寻找概率分布 \(p(\mathbf{m}|\mathbf{d}_{\text{syn}})\) 来表达逆算子的非唯一性，从中采样 \(\mathbf{z} \sim p(\mathbf{m}|\mathbf{d}_{\text{syn}})\) 得到可能的解，满足 \(f(\mathbf{z}) = \mathbf{d}_{\text{syn}}\)。在概率反演中，无需解的初始猜测或正则化，也无需用泰勒级数近似正演算子。尽管这种方法可以说更优雅，且允许不确定性评估，但目前并不常用。原因在于，使用传统基于采样的方法寻找分布 \(p(\mathbf{m}|\mathbf{d}_{\text{syn}})\) 需要每一步评估正演算子，其计算规模随模型参数数量大致呈指数增长 [5 (https://arxiv.org/html/2606.31288#bib.bib11)]。对于实际应用中多达数百万个模型参数的情况，这计算上过于苛刻。

最近，概率反演有了显著进展，问题被重新表述为优化问题。这主要得益于深度学习研究。使用归一化流的变分推断 [18 (https://arxiv.org/html/2606.31288#bib.bib10)] 构建（离散的）归一化流序列，将已知初始分布转换为目标分布。这些流通过最大化证据下界（即最小化先验与目标分布之间的Kullback-Leibler散度）来训练。在另一个研究领域，扩散模型被用于获取目标分布 [17 (https://arxiv.org/html/2606.31288#bib.bib13)]。在扩散模型中，随机过程将已知初始分布转换为目标分布。流匹配 [8 (https://arxiv.org/html/2606.31288#bib.bib1)] 可被视为归一化流变分推断与扩散模型之间的桥梁。它也是一种深度学习方法，迭代地将已知初始分布（如 d 维正态分布）转换为无条件目标分布（如 \(p(\mathbf{m})\)）或条件目标分布（如 \(p(\mathbf{m}|\mathbf{d}_{\text{syn}})\)）。文献中，第一种情况称为**无引导流匹配**，第二种情况称为**引导流匹配**。这种概率分布的转换通过在训练过程中学习连续归一化流（CNF）来实现。在推理时，求解一个简单的常微分方程（ODE）。接下来，将借助一个简单的二维速度模型案例研究更详细地描述流匹配理论。所用的数学框架主要基于 [8 (https://arxiv.org/html/2606.31288#bib.bib1),6 (https://arxiv.org/html/2606.31288#bib.bib12),9 (https://arxiv.org/html/2606.31288#bib.bib7)]，并针对地球物理数据概率反演的应用场景进行了适配，因为流匹配通常用于生成式 AI [3 (https://arxiv.org/html/2606.31288#bib.bib17),12 (https://arxiv.org/html/2606.31288#bib.bib20),16 (https://arxiv.org/html/2606.31288#bib.bib18),2 (https://arxiv.org/html/2606.31288#bib.bib19),7 (https://arxiv.org/html/2606.31288#bib.bib8)]，尚未在地球物理学界普遍用于概率反演。

## 2 案例研究：简单二维速度模型  

作为第一步，我们在一个最小化的概率反演场景中说明流匹配方法。我们考虑一个简单的两层二维速度模型，如图 1 (https://arxiv.org/html/2606.31288#S2.F1) 所示。第一层厚度已知为 10 米，第二层为均匀半空间。因此速度模型可由两层速度描述：\(v_1 \in [300,700] \frac{\text{m}}{\text{s}}\) 和 \(v_2 \in [s \cdot v_1, 5000] \frac{\text{m}}{\text{s}}\)，其中 \(s\) 是确保折射波射线以速度 \(v_2\) 至少传播 1 米的最小因子。为了用观测数据来条件化（引导）概率反演，计算固定距离 50 米的源-接收器对对应的折射波旅行时。正演问题（测量给定速度模型的折射波旅行时）是唯一的，但反问题（给定观测到的折射波旅行时，得到可能的速度场）不唯一，因为不同的 \(v_1, v_2\) 组合可能产生相同的折射波旅行时。

> 参见标题  
> 图 1：简单二维速度模型。

### 2.1 初始分布与目标分布  

虽然任何初始分布都可使用（只要存在连续归一化流将其转换到目标分布），但实践中 d 维正态分布最受欢迎，因为它易于采样且便于构造流。因此，我们关注初始分布 \(p_{\text{init}} = \mathcal{N}(0, \mathbf{I}_d)\)。在当前场景中，地震速度模型有两个参数，故 \(d=2\)。目标分布未知，但我们假设有有限个样本来自该分布，这些样本将作为训练数据集。在我们的示例中，目标分布是所有可能地震速度模型的分布（对于无引导流匹配），或者是所有可能速度场与所有可能观测数据点的联合分布（对于引导流匹配）。我们将目标分布记为 \(p_{\text{target}}\)，其样本记为 \(\mathbf{m} \sim p_{\text{target}}\)，其中 \(\mathbf{m} \in \mathbb{R}^d\)。

### 2.2 连续归一化流  

在流匹配中，初始分布的样本沿一个随时间变化的轨迹 \(X_t: [0,1] \rightarrow \mathbb{R}^d\) 转换为目标分布的样本。对于任意时间 \(t \in [0,1]\)，\(X_t\) 返回 \(\mathbb{R}^d\) 中的一个位置。此处时间 \(t\) 代表轨迹上的一个时间步，在 \(t=0\) 时轨迹上的样本位于初始分布中，在 \(t=1\) 时位于目标分布中。该轨迹是如下 ODE 的解：
\[
\frac{\text{d}}{\text{d}t} X_t = u_t(X_t) \tag{1}
\]
其中 \(u_t(X_t)\) 是一个时间相关的向量场 \(u_t: [0,1] \times \mathbb{R}^d \rightarrow \mathbb{R}^d\)，在给定时间 \(t\) 和位置 \(\mathbf{x} \in \mathbb{R}^d\) 时指定空间中的速度。方程 1 (https://arxiv.org/html/2606.31288#S2.E1) 的初始条件为 \(X_0 = \epsilon\)，其中 \(\epsilon \sim p_{\text{init}}\) 是轨迹的起点。通常，轨迹需要依赖于初始分布中的特定样本，从而为不同样本产生不同轨迹。为此，我们使用一个流 \(\psi_t: [0,1] \times \mathbb{R}^d \rightarrow \mathbb{R}^d\)，它是一个函数，需要时间 \(t\) 和某个位置 \(\mathbf{x} \in \mathbb{R}^d\)，并返回 \(\mathbb{R}^d\) 中的一个位置（类似于之前介绍的轨迹）。由于 \(t \in [0,1]\) 是连续的，我们得到一个连续归一化流（CNF）。因此，我们需要求解 ODE：
\[
\frac{\text{d}}{\text{d}t} \psi_t(\mathbf{x}) = u_t(\psi_t(\mathbf{x})) \tag{2}
\]
初始条件为 \(\psi_0(\mathbf{x}) = \mathbf{x}\)。时间相关的向量场 \(u_t\) 将用一个神经网络（NN）\(u_t^\theta\) 来近似，其中 \(\theta\) 是网络参数。

### 2.3 高斯条件概率路径  

在流匹配中，CNF 遵循概率路径，因为它们将样本从一个概率分布搬运到另一个概率分布的样本。首先，考虑（时间相关的）高斯条件概率路径 \(p_t(\cdot|\mathbf{m})\)，其中 \(\mathbf{m} \sim p_{\text{target}}\)。它条件于目标分布的一个样本，对应于从正态分布 \(p_{\text{init}} = \mathcal{N}(0, \mathbf{I}_d)\) 到狄拉克分布 \(\delta_{\mathbf{m}}\)（从 \(\delta_{\mathbf{m}}\) 采样总是返回 \(\mathbf{m}\)）的流。因此，\(p_0(\cdot|\mathbf{m}) = p_{\text{init}}\) 且 \(p_1(\cdot|\mathbf{m}) = \delta_{\mathbf{m}}\)。这样的高斯条件概率路径可以构造为：
\[
p_t(\cdot|\mathbf{m}) = \mathcal{N}(\alpha_t \mathbf{m}, \beta_t^2 \mathbf{I}_d) \tag{3}
\]
其中 \(\alpha_t\) 和 \(\beta_t\) 是时间相关的可微单调函数（噪声调度器），满足 \(\alpha_0 = \beta_1 = 0\) 且 \(\alpha_1 = \beta_0 = 1\)。在最基本的形式中，可选取 \(\alpha_t = t\) 和 \(\beta_t = 1 - t\)，本文即采用此方案。为了从方程 3 (https://arxiv.org/html/2606.31288#S2.E3) 的分布中采样，我们可以使用重参数化技巧：
\[
\mathbf{m} \sim p_{\text{target}}, \quad \epsilon \sim p_{\text{init}} = \mathcal{N}(0, \mathbf{I}_d) \quad \Rightarrow \quad \alpha_t \mathbf{m} + \beta_t \epsilon = \mathbf{x} \sim p_t \tag{4}
\]

为了训练神经网络学习高斯条件概率路径的向量场，需要计算目标向量场 \(u_t^{\text{target}}\)。为此，构造条件流为 \(\psi_t^{\text{target}}(\mathbf{x}|\mathbf{m}) = \alpha_t \mathbf{m} + \beta_t \epsilon\)，然后求解条件向量场的 ODE：
\[
\begin{aligned}
\frac{\text{d}}{\text{d}t} \psi_t^{\text{target}}(\epsilon|\mathbf{m}) &= u_t^{\text{target}}(\psi_t^{\text{target}}(\epsilon|\mathbf{m})|\mathbf{m}) \\
\underbrace{\frac{\text{d}}{\text{d}t} \alpha_t}_{=1} \mathbf{m} + \underbrace{\frac{\text{d}}{\text{d}t} \beta_t}_{=-1} \epsilon &= u_t^{\text{target}}\big( \underbrace{\alpha_t}_{=t} \mathbf{m} + \underbrace{\beta_t}_{=1-t} \epsilon \big| \mathbf{m} \big) \\
\text{重参数化 } \mathbf{x} = \alpha_t \mathbf{m} + \beta_t \epsilon &= t \mathbf{m} + (1-t) \epsilon \\
\mathbf{m} - \epsilon &= u_t^{\text{target}}(\mathbf{x}|\mathbf{m})
\end{aligned} \tag{5}
\]

这导出了条件流匹配损失：
\[
\mathcal{L}_{\text{CFM}}(\theta) = \mathbb{E}_{t \sim \text{Unif}_{[0,1]}, \mathbf{m} \sim p_{\text{target}}, \epsilon \sim \mathcal{N}(0, \mathbf{I}_d)} \left[ \| u_t^\theta(\mathbf{x}) - (\mathbf{m} - \epsilon) \|^2 \right] \tag{6}
\]

我们现在可以实现一个极简函数，用于在条件从狄拉克分布 \(\mathbf{m} \sim \delta_{\mathbf{m}}\) 采样的情况下，训练神经网络（本例中使用简单的多层感知器）基于方程 6 (https://arxiv.org/html/2606.31288#S2.E6) 中的目标函数：

**算法 1**：条件于目标分布单一模型的条件流匹配损失伪代码。  
\(\mathbf{m} \sim \delta_{\mathbf{m}}\)，神经网络 \(u^\theta(t, \mathbf{x})\)，学习率 \(\eta\)，步数 \(n\)，索引 \(i\)。  
\(i \leftarrow 0\)  
while \(i < n\) do  
 从训练数据中采样 \(\mathbf{m} \sim \delta_{\mathbf{m}}, \epsilon \sim \mathcal{N}(0, \mathbf{I}_d), t \sim \text{Unif}_{[0,1]}\)  
 计算 \(\mathbf{x} = t \mathbf{m} + (1-t) \epsilon\)  
 计算梯度 \(\nabla_\theta \| u^\theta(t, \mathbf{x}) - (\mathbf{m} - \epsilon) \|^2\)  
 更新 \(\theta \leftarrow \theta - \eta \nabla_\theta\)  
 \(i \leftarrow i + 1\)  
end while  

经过 \(n\) 步训练后，我们得到一个可用的神经网络，可作为向量场 \(u_t^\theta\) 的近似。然后可以在采样器中使用该近似。值得注意的是，对于条件目标分布 \(p_t(\cdot|\mathbf{m})\)，向量场 \(u_t^\theta\) 仅依赖于当前样本 \(\mathbf{x}\) 和当前时间 \(t\)，而不直接依赖于条件 \(\mathbf{m}\)。然而，经过训练后，向量场隐式