656 lines
26 KiB
Markdown
656 lines
26 KiB
Markdown
---
|
||
title: Stochastic Interpolants：统一生成模型的底层框架
|
||
draft: false
|
||
tags:
|
||
- Stochastic-Interpolants
|
||
- 生成模型
|
||
- 统一框架
|
||
- 深度学习
|
||
---
|
||
|
||
# Stochastic Interpolants：flows、diffusions 与生成模型的统一底层结构
|
||
|
||
---
|
||
|
||
## 一、从碎片化到统一：为什么需要 Stochastic Interpolants
|
||
|
||
### 1.1 生成模型领域的碎片化现状
|
||
|
||
过去几年出现了多种生成模型范式，每种都有独特的数学形式和训练目标：
|
||
|
||
| 模型 | 过程类型 | 核心数学工具 | 训练目标 |
|
||
|------|---------|-------------|---------|
|
||
| **Normalizing Flow / CNF** | 确定性 ODE | $\log\det J$ / $\text{tr}(\partial f/\partial x)$ | 对数似然最大化 |
|
||
| **Diffusion (DDPM)** | 离散马尔可夫链 | KL 散度 / ELBO | 噪声预测 $\epsilon$ |
|
||
| **Score-Based SDE** | 随机微分方程 | Fokker-Planck / Score matching | Score function $\nabla\log p_t$ |
|
||
| **Flow Matching** | 连续路径 | 向量场回归 | 速度场 $v_\theta$ |
|
||
| **Rectified Flow** | 确定性 ODE | 最优传输 / 路径直线化 | 速度场差值 |
|
||
|
||
表面上这些模型截然不同，但深入分析会发现它们共享相同的信息论基础。**Stochastic Interpolants** 的核心贡献是揭示了这个统一结构。
|
||
|
||
### 1.2 核心问题
|
||
|
||
设我们有两个分布：
|
||
- $p_0$：噪声分布（通常为 $\mathcal{N}(0, I)$）
|
||
- $p_1$：数据分布
|
||
|
||
**如何构建一个连续时间过程，将 $p_0$ 连接到 $p_1$？**
|
||
|
||
Stochastic Interpolants 的回答是：以上所有都可以统一为**特殊的随机插值过程**，即选择特定的调度函数 $\psi, \phi, \sigma$。
|
||
|
||
---
|
||
|
||
## 二、Stochastic Interpolants 的数学框架
|
||
|
||
### 2.1 随机插值过程的定义
|
||
|
||
**定义（随机插值过程）**：
|
||
|
||
设 $x_0 \sim p_0$ 和 $x_1 \sim p_1$ 是两个**独立采样**的样本（这是 SI 的关键假设）。定义它们的**随机插值**为一个连续时间过程 $X_t \in \mathbb{R}^D$：
|
||
|
||
$$X_t = \psi(t) \, x_0 + \phi(t) \, x_1 + \sigma(t) \, z \tag{2.1}$$
|
||
|
||
其中：
|
||
- $\psi(t), \phi(t): [0,1] \to \mathbb{R}$ 是确定性的**调度函数**（schedule functions）
|
||
- $\sigma(t): [0,1] \to \mathbb{R}$ 是噪声幅度函数
|
||
- $z \sim \mathcal{N}(0, I)$ 是独立的标准高斯噪声
|
||
- $x_0 \sim p_0$ 和 $x_1 \sim p_1$ 独立同分布采样
|
||
|
||
**物理直觉**：
|
||
- $x_0$ 和 $x_1$ 是路径的两个端点
|
||
- $\psi(t)$ 控制从 $x_0$ 出发的"权重"
|
||
- $\phi(t)$ 控制从 $x_1$ 出发的"权重"
|
||
- $\sigma(t) z$ 是在路径上叠加的随机噪声
|
||
|
||
### 2.2 边界条件
|
||
|
||
为确保 $X_t$ 是有效的插值过程，调度函数必须满足以下边界条件：
|
||
|
||
$$X_{t=0} = \psi(0) \, x_0 + \phi(0) \, x_1 + \sigma(0) \, z \stackrel{!}{=} x_0$$
|
||
$$X_{t=1} = \psi(1) \, x_0 + \phi(1) \, x_1 + \sigma(1) \, z \stackrel{!}{=} x_1$$
|
||
|
||
这要求：
|
||
|
||
$$\psi(0) = 1, \quad \phi(0) = 0, \quad \sigma(0) = 0 \tag{2.2}$$
|
||
$$\psi(1) = 0, \quad \phi(1) = 1, \quad \sigma(1) = 0 \tag{2.3}$$
|
||
|
||
**解释**：
|
||
- $t=0$ 时，$X_0 = x_0$（纯噪声）
|
||
- $t=1$ 时，$X_1 = x_1$（纯数据）
|
||
|
||
### 2.3 插值调度函数的设计空间
|
||
|
||
**标准线性插值**（Rectified Flow）：
|
||
$$\psi(t) = 1 - t, \quad \phi(t) = t, \quad \sigma(t) = 0$$
|
||
|
||
这给出 $X_t = (1-t)x_0 + t x_1$，即两点之间的直线段。
|
||
|
||
**DDPM 插值形式不满足 SI 边界条件**：
|
||
|
||
DDPM 的加噪形式 $X_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t} \epsilon$ 无法直接嵌入标准 SI 框架：
|
||
|
||
- SI 要求 $\psi(1) = 0$，但 $\psi(1) = \sqrt{\bar{\alpha}_1} \neq 0$（除非 $\bar{\alpha}_1 = 0$，即极端情况）
|
||
- DDPM 的端点间关系是随机的（非确定性插值），不满足 SI 的确定性端点假设
|
||
- DDPM 的时间流向与 SI 相反：$t=0$ 为数据、$t=T$ 为噪声
|
||
|
||
**独立噪声插值**（最一般形式）：
|
||
$$X_t = \alpha(t) x_0 + \beta(t) x_1 + \gamma(t) z$$
|
||
|
||
其中 $\alpha, \beta, \gamma$ 是满足边界条件的任意函数。
|
||
|
||
### 2.4 边际分布的数学描述
|
||
|
||
**关键问题**：对于给定的 $x_0 \sim p_0$ 和 $x_1 \sim p_1$，$X_t$ 的**边际分布** $p_t$ 是什么？
|
||
|
||
**命题**：$p_t$ 可以通过以下方式描述：
|
||
|
||
给定 $x_0$ 和 $x_1$，$X_t$ 的条件分布为：
|
||
$$p(x_t | x_0, x_1) = \mathcal{N}\left(\psi(t) x_0 + \phi(t) x_1, \; \sigma(t)^2 I\right) \tag{2.4}$$
|
||
|
||
边际分布为：
|
||
$$p_t(x_t) = \mathbb{E}_{x_0 \sim p_0, x_1 \sim p_1}\left[ p(x_t | x_0, x_1) \right] \tag{2.5}$$
|
||
|
||
即 $p_t$ 是条件高斯分布在 $(x_0, x_1)$ 联合分布上的期望。
|
||
|
||
---
|
||
|
||
## 三、从插值过程到 Itô SDE
|
||
|
||
### 3.1 Itô SDE 的推导
|
||
|
||
对随机插值过程 $X_t$ 应用 Itô 引理，可以将其转化为等价的 Itô SDE。
|
||
|
||
**定理（SI 的 Itô SDE 表示）**：
|
||
|
||
过程 $X_t$ 满足以下 Itô SDE：
|
||
$$d X_t = \underbrace{\left[ \dot{\psi}(t) x_0 + \dot{\phi}(t) x_1 \right]}_{\text{drift } f_t(X_t)} d t + \underbrace{\dot{\sigma}(t)}_{\text{diffusion } g_t} d W_t \tag{3.1}$$
|
||
|
||
其中 $W_t$ 是标准维纳过程，$\dot{}$ 表示对时间的导数。
|
||
|
||
**推导步骤**：
|
||
|
||
1. **微分形式**：对 $X_t = \psi(t) x_0 + \phi(t) x_1 + \sigma(t) z$ 求导：
|
||
$$\frac{d X_t}{d t} = \dot{\psi}(t) x_0 + \dot{\phi}(t) x_1 + \dot{\sigma}(t) z$$
|
||
|
||
2. **分离确定性项和随机项**：将 $\dot{\sigma}(t) z d t$ 转换为维纳增量：
|
||
- 对于标准高斯 $z \sim \mathcal{N}(0,I)$，有 $z d t \approx d W_t$ 的统计特性
|
||
- 更精确地，$\dot{\sigma}(t) z d t = \dot{\sigma}(t) d W_t$ 当 $d W_t \sim \mathcal{N}(0, dt)$
|
||
|
||
**注意**：原版 Albergo 论文推导中不存在人为构造的 $-\frac{1}{2} \frac{d}{dt}(\sigma^2) z$ 修正项——该修正项是自主错误添加，原版 SI 的漂移项仅为 $\dot{\psi} x_0 + \dot{\phi} x_1$。
|
||
|
||
### 3.2 漂移项的物理解释
|
||
|
||
**确定性漂移项**：
|
||
$$f_t = \dot{\psi}(t) x_0 + \dot{\phi}(t) x_1$$
|
||
|
||
这是从端点出发的确定性运动。
|
||
|
||
### 3.3 Fokker-Planck 方程
|
||
|
||
**定理（Fokker-Planck 方程）**：
|
||
|
||
$X_t$ 的边际分布 $p_t$ 满足以下 Fokker-Planck 方程：
|
||
$$\frac{\partial p_t(x)}{\partial t} = -\nabla \cdot \left[ \mathbb{E}\left[\dot{\psi}(t) x_0 + \dot{\phi}(t) x_1 \mid X_t = x\right] p_t(x) \right] + \frac{1}{2} \left(\dot{\sigma}(t)\right)^2 \, \nabla^2 p_t(x) \tag{3.2}$$
|
||
|
||
其中 $\nabla^2$ 是拉普拉斯算子（标准符号，非 $\Delta$）。
|
||
|
||
**物理意义**：
|
||
- 第一项是**漂移项**：描述确定性运动导致的密度变化
|
||
- 第二项是**扩散项**：描述随机性导致的密度分散
|
||
|
||
**当 $\sigma(t) = 0$ 时**，第二项消失，Fokker-Planck 方程退化为确定性 ODE 的连续性方程：
|
||
$$\frac{\partial p_t}{\partial t} = -\nabla \cdot (v_t p_t) \tag{3.3}$$
|
||
|
||
其中 $v_t = \mathbb{E}[\dot{\psi} x_0 + \dot{\phi} x_1 | X_t]$ 是速度场。
|
||
|
||
---
|
||
|
||
## 四、与现有模型的对应关系
|
||
|
||
### 4.1 还原 Normalizing Flow / CNF
|
||
|
||
**条件**：$\sigma(t) = 0$（无额外噪声）
|
||
|
||
此时过程退化为**确定性 ODE**：
|
||
$$d X_t = \left[\dot{\psi}(t) x_0 + \dot{\phi}(t) x_1\right] d t \tag{4.1}$$
|
||
|
||
**连续性方程**（描述概率守恒）：
|
||
$$\frac{\partial \log p_t}{\partial t} = -\text{tr}\left(\frac{\partial f}{\partial x}\right) \tag{4.2}$$
|
||
|
||
其中 $f = \dot{\psi} x_0 + \dot{\phi} x_1$。
|
||
|
||
**这正是 CNF 的瞬时变量代换公式！**
|
||
|
||
### 4.2 还原 DDPM / Score-Based SDE
|
||
|
||
**重要澄清**：DDPM 原生前向加噪过程的时间流向与 SI 框架**方向相反**：
|
||
|
||
- **SI 标准正向**：$t=0$ 为噪声 $p_0$，$t=1$ 为数据 $p_1$（生成方向）
|
||
- **DDPM 原生加噪**：$t=0$ 为真实数据，逐步加噪到 $t=T$ 纯噪声
|
||
|
||
DDPM 的前向 SDE 本质是带连续布朗运动扩散项的**随机过程**，天然依赖 $\sigma(t) \neq 0$ 随机分量，**零噪声确定性流无法等价还原 DDPM**。
|
||
|
||
**SI 框架嵌入 DDPM 的正确方式**：
|
||
|
||
若将 DDPM 的逆向生成过程（从噪声到数据）嵌入 SI 框架，需首先对 DDPM 做时间反转。DDPM 前向 SDE：
|
||
$$d X_t = -\frac{1}{2} \beta(t) X_t d t + \sqrt{\beta(t)} d W_t$$
|
||
|
||
其逆向 SDE（反向时间）为：
|
||
$$d X_t = \left[ -\frac{1}{2} \beta(t) X_t - \beta(t) \nabla_X \log p_t \right] d t + \sqrt{\beta(t)} d \bar{W}_t$$
|
||
|
||
**DDPM 无法以标准 SI 形式 $X_t = \psi(t)x_0+\phi(t)x_1+\sigma(t)z$ 嵌入**，因为 DDPM 不满足 SI 的独立端点采样假设。DDPM 的 $x_0$ 和 $x_1$ 之间不存在确定性的插值关系。
|
||
|
||
### 4.3 还原 Rectified Flow
|
||
|
||
**条件**：$\sigma(t) = 0$ 且 $\psi(t) = 1-t$，$\phi(t) = t$
|
||
|
||
此时：
|
||
$$X_t = (1-t) x_0 + t x_1 \tag{4.5}$$
|
||
|
||
求导得：
|
||
$$\frac{d X_t}{d t} = x_1 - x_0 \tag{4.6}$$
|
||
|
||
对应的 ODE 为：
|
||
$$d X_t = (x_1 - x_0) d t \tag{4.7}$$
|
||
|
||
这是**无噪声的确定性流**，正是 Rectified Flow 的路径！
|
||
|
||
### 4.4 统一公式表
|
||
|
||
| 模型 | $\psi(t)$ | $\phi(t)$ | $\sigma(t)$ | 过程类型 |
|
||
|------|-----------|-----------|------------|---------|
|
||
| **NF/CNF** | 任意可逆 | 任意可逆 | $0$ | 确定性 ODE |
|
||
| **DDPM**（逆向生成） | 依赖调度 | 依赖调度 | $\neq 0$ | 随机 SDE（需时间反转） |
|
||
| **Rectified Flow** | $1-t$ | $t$ | $0$ | 确定性 ODE（常数速度场） |
|
||
| **一般 SI** | 任意 | 任意 | 任意 | 随机 SDE |
|
||
|
||
**说明**：DDPM 的前向过程从 $t=0$ 数据逐步加噪到 $t=T$ 噪声，时间流向与 SI 相反；其逆向生成过程可对应 SI 框架，但调度函数需通过时间反转定义。
|
||
|
||
---
|
||
|
||
## 五、从正向过程到反向生成
|
||
|
||
### 5.1 反向时间 SDE
|
||
|
||
**核心问题**：给定 $X_T \sim p_T$（接近数据分布），如何反向推导出 $X_0 \sim p_0$（噪声分布）？
|
||
|
||
**定理（反向 SDE）**：
|
||
|
||
对于前向 SDE：
|
||
$$d X_t = f_t(X_t) d t + g_t d W_t \tag{5.1}$$
|
||
|
||
反向时间（$t \to T-t$）的 SDE 为：
|
||
$$d X_t = \left[ f_t(X_t) - g_t g_t^\top \nabla_X \log p_t(X_t) \right] d t + g_t d \bar{W}_t \tag{5.2}$$
|
||
|
||
其中 $\nabla_X \log p_t$ 是 **score function**，$g_t g_t^\top$ 是扩散系数的协方差矩阵，$\bar{W}_t$ 是反向维纳过程。
|
||
|
||
**推导要点**：
|
||
|
||
设 $Y_s = X_{T-s}$ 为时间反转过程。通过 Itô 引理和 Radon-Nikodym 导数，可以得到反向 SDE 的漂移项修正为 $-g_t g_t^\top \nabla \log p_t$。
|
||
|
||
### 5.2 应用于 Stochastic Interpolants
|
||
|
||
对于 SI 的前向 SDE（式 3.1）：
|
||
$$f_t = \dot{\psi} x_0 + \dot{\phi} x_1 - \frac{1}{2} \frac{d}{dt}(\sigma^2) z, \quad g_t = \dot{\sigma}$$
|
||
|
||
对应的反向 SDE 为：
|
||
$$d X_t = \left[ \dot{\psi} x_0 + \dot{\phi} x_1 - \left(\dot{\sigma}\right)^2 \nabla_X \log p_t \right] d t + \dot{\sigma} d \bar{W}_t \tag{5.3}$$
|
||
|
||
### 5.3 Probability Flow ODE
|
||
|
||
**定理（Probability Flow ODE）**：
|
||
|
||
对于任意 SDE（式 5.1），存在等价的**确定性 ODE**，产生相同的边际分布 $p_t$：
|
||
$$d X_t = \left[ f_t(X_t) - \frac{1}{2} g_t g_t^\top \nabla_X \log p_t(X_t) \right] d t \tag{5.4}$$
|
||
|
||
**证明概览**：
|
||
|
||
通过 Fokker-Planck 方程可以验证，两个过程具有相同的密度演化。
|
||
|
||
**物理意义**：
|
||
- Probability Flow ODE 移除了随机项 $g_t d W_t$
|
||
- 保留了 score 修正项 $- \frac{1}{2} g_t g_t^\top \nabla \log p_t$
|
||
- 结果是一个确定性流，但保持相同的边际分布
|
||
|
||
**注意**：Probability Flow ODE 源自 Score-SDE 体系，是分数流的专属边际等效确定性 ODE，**无法直接通用到任意调度的 Stochastic Interpolants**——仅当对应的 SDE 属于 Score-SDE 框架时 PFODE 才适用。
|
||
|
||
**当 $g_t = 0$ 时**（无噪声），Probability Flow ODE 退化为标准确定性 ODE：
|
||
$$d X_t = f_t d t \tag{5.5}$$
|
||
|
||
### 5.4 速度场的定义
|
||
|
||
**定义（速度场）**：
|
||
|
||
在 SI 框架下，速度场 $v_t$ 定义为条件期望（基于独立端点联合分布 $\pi(x_0, x_1) = p_0(x_0) p_1(x_1)$）：
|
||
$$v_t(x) = \mathbb{E}\left[ \frac{d X_t}{d t} \bigg| X_t = x \right] = \mathbb{E}_{x_0, x_1 \sim \pi}\left[ \dot{\psi} x_0 + \dot{\phi} x_1 \mid X_t = x \right] \tag{5.6}$$
|
||
|
||
**与 Flow Matching 的联系**：
|
||
|
||
Flow Matching 直接回归这个速度场 $v_\theta \approx v_t$。在 SI 框架下，训练目标可以表述为：
|
||
$$\min_\theta \mathbb{E}_{t, x_0, x_1, z} \left\| v_\theta(X_t, t) - \mathbb{E}[\dot{\psi} x_0 + \dot{\phi} x_1 \mid X_t] \right\|^2 \tag{5.7}$$
|
||
|
||
---
|
||
|
||
## 六、训练目标与损失函数
|
||
|
||
### 6.1 三种等价的训练目标
|
||
|
||
SI 框架下可以推导出三种等价的训练目标，它们在特定的参数选择下互相一致：
|
||
|
||
#### 6.1.1 速度场回归（Flow Matching 风格）
|
||
|
||
当 $\sigma(t) = 0$ 时，训练目标是回归条件速度场：
|
||
$$\mathcal{L}_v(\theta) = \mathbb{E}_{t, x_0, x_1} \left\| v_\theta(X_t, t) - (\dot{\psi} x_0 + \dot{\phi} x_1) \right\|^2 \tag{6.1}$$
|
||
|
||
对于 $\psi(t) = 1-t, \phi(t) = t$，这简化为：
|
||
$$\mathcal{L}_v(\theta) = \mathbb{E}_{t, x_0, x_1} \left\| v_\theta(X_t, t) - (x_1 - x_0) \right\|^2 \tag{6.2}$$
|
||
|
||
这正是 Rectified Flow 的损失函数。
|
||
|
||
#### 6.1.2 Score Matching 风格
|
||
|
||
**定理（去噪 Score Matching 的 SI 形式）**：
|
||
|
||
当 $\sigma(t) \neq 0$ 时，训练目标可以转化为 score matching：
|
||
$$\mathcal{L}_s(\theta) = \mathbb{E}_{t, x_0, x_1, z} \left\| s_\theta(X_t, t) - \nabla_{X_t} \log p(x_t | x_0, x_1) \right\|^2 \tag{6.3}$$
|
||
|
||
其中 $s_\theta$ 是神经网络预测的 score function。
|
||
|
||
**推导**：
|
||
|
||
条件分布 $p(x_t | x_0, x_1) = \mathcal{N}(\psi x_0 + \phi x_1, \sigma^2 I)$，所以：
|
||
$$\nabla_{X_t} \log p(x_t | x_0, x_1) = -\frac{X_t - (\psi x_0 + \phi x_1)}{\sigma(t)^2} = -\frac{\sigma(t) z}{\sigma(t)^2} = -\frac{z}{\sigma(t)} \tag{6.4}$$
|
||
|
||
因此：
|
||
$$\mathcal{L}_s(\theta) = \mathbb{E}_{t, x_0, x_1, z} \left\| s_\theta(X_t, t) + \frac{z}{\sigma(t)} \right\|^2 \tag{6.5}$$
|
||
|
||
#### 6.1.3 噪声预测（DDPM 风格）
|
||
|
||
当 $\psi(t) = \sqrt{\bar{\alpha}_t}, \phi(t) = \sqrt{1-\bar{\alpha}_t}$ 时，令 $\epsilon = z$，有：
|
||
$$X_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1-\bar{\alpha}_t} \epsilon \tag{6.6}$$
|
||
|
||
score matching 目标（式 6.5）可以重写为噪声预测：
|
||
$$\mathcal{L}_\epsilon(\theta) = \mathbb{E}_{t, x_0, \epsilon} \left\| \epsilon_\theta(X_t, t) - \epsilon \right\|^2 \tag{6.7}$$
|
||
|
||
其中 $\epsilon_\theta = -\sigma(t) s_\theta$。
|
||
|
||
**这正是 DDPM 的噪声预测目标！**
|
||
|
||
### 6.2 目标等价性的条件限定
|
||
|
||
**定理（目标等价性，仅在严格条件下成立）**：
|
||
|
||
三个训练目标在以下**全部四个前置条件**同时满足时互相等价：
|
||
|
||
1. **端点独立采样**：$x_0 \sim p_0$, $x_1 \sim p_1$ 相互独立
|
||
2. **固定统一插值调度**：$\psi, \phi, \sigma$ 是预先确定的调度函数（非学习参数）
|
||
3. **无边际分布偏移**：插值过程保持正确的边际分布演化 $p_t$
|
||
4. **时间流向严格对齐**：SI 框架的 $t=0$ 噪声、$t=1$ 数据与目标模型一致
|
||
|
||
**脱离上述约束条件时**，三者仅为形式相似，数学期望层面**并不等价**：
|
||
|
||
| 目标 | 公式 | 适用场景 |
|
||
|------|------|---------|
|
||
| **速度场回归** | $\mathbb{E}\| v_\theta - (x_1-x_0) \|^2$ | $\sigma(t) = 0$，Flow Matching / RF |
|
||
| **Score matching** | $\mathbb{E}\| s_\theta + \frac{z}{\sigma(t)} \|^2$ | $\sigma(t) \neq 0$，去噪场景 |
|
||
| **噪声预测** | $\mathbb{E}\| \epsilon_\theta - \epsilon \|^2$ | DDPM 风格调度 |
|
||
|
||
**物理直觉**：
|
||
|
||
这三种目标从不同角度描述同一个潜在函数：
|
||
- Score function $\nabla \log p_t$ 描述概率梯度
|
||
- 噪声预测 $\epsilon$ 是 score 的线性变换（仅在特定调度下）
|
||
- 速度场 $v_t$ 是路径的切向量（仅在 $\sigma(t) = 0$ 时）
|
||
|
||
### 6.3 时间步采样策略
|
||
|
||
**均匀采样**：
|
||
$$t \sim \mathcal{U}(0, 1)$$
|
||
|
||
问题：在 $t \approx 0$ 和 $t \approx 1$ 区域，分布变化剧烈，采样不均衡。
|
||
|
||
**Logit-Normal 采样**（推荐）：
|
||
$$t = \frac{1}{1 + e^{-z}}, \quad z \sim \mathcal{N}(\mu, \sigma^2) \tag{6.8}$$
|
||
|
||
效果：使 $t$ 更集中在中间区域（$t \approx 0.5$），这正是路径最复杂、信息最丰富的区域。
|
||
|
||
**常用配置**：$\mu = 0, \sigma = 1.0$
|
||
|
||
### 6.4 损失加权
|
||
|
||
实践中常用的加权损失：
|
||
$$\mathcal{L}_\lambda(\theta) = \mathbb{E}_{t \sim \mathrm{KL}} \left[ \lambda(t) \cdot \left\| v_\theta(X_t, t) - u_t \right\|^2 \right] \tag{6.9}$$
|
||
|
||
常用权重策略：
|
||
- $\lambda(t) = 1$（均匀权重）
|
||
- $\lambda(t) = \frac{1}{1-\bar{\alpha}_t}$（与信噪比相关）
|
||
- $\lambda(t) = t(1-t)$（中间区域加权）
|
||
|
||
---
|
||
|
||
## 七、与 Schrödinger Bridge 的深层联系
|
||
|
||
### 7.1 Schrödinger Bridge 问题
|
||
|
||
**原始 Schrödinger Bridge（SB）问题**：
|
||
|
||
在所有满足边际约束 $p_0 = p_0^{\mathrm{KL}}$ 和 $p_1 = p_1^{\mathrm{KL}}$ 的随机过程 $\{X_t\}$ 中，寻找使以下**动作泛函**最小的那条：
|
||
$$\min_{X_t} \mathbb{E} \left[ \int_0^1 \frac{1}{2} \left\| \frac{d X_t}{d t} \right\|^2 d t \right] \tag{7.1}$$
|
||
|
||
约束：$X_0 \sim p_0$，$X_1 \sim p_1$。
|
||
|
||
**物理直觉**：
|
||
|
||
这相当于在所有连接给定端点分布的随机过程中，寻找"最省力"的路径。就像在两点之间找到最速降线，但考虑随机性。
|
||
|
||
**与最优传输的关系**：
|
||
|
||
SB 与 Monge 的最优传输问题有深层联系。当路径无噪声时，SB 退化为最优传输：
|
||
$$\min_{T} \int c(x_0, T(x_0)) d p_0(x_0) \tag{7.2}$$
|
||
|
||
### 7.2 SI 与 SB 的关系
|
||
|
||
**SI 提供了 SB 的参数化**：
|
||
|
||
通过选择 $\psi, \phi, \sigma$，SI 实际上参数化了连接 $p_0$ 和 $p_1$ 的候选路径族。
|
||
|
||
**SB 是 SI 的最优选择**：
|
||
|
||
在所有 SI 参数化的路径中，SB 对应于使动作泛函（式 7.1）最小的那条路径。
|
||
|
||
**关键区别**：
|
||
- SI 是一个**一般性框架**，允许任意调度函数
|
||
- SB 是 SI 设计空间中的**最优选择准则**
|
||
|
||
### 7.3 熵正则化 Schrödinger Bridge
|
||
|
||
**引入熵正则项**：
|
||
|
||
$$\min_{X_t} \mathbb{E} \left[ \int_0^1 \frac{1}{2} \left\| \frac{d X_t}{d t} \right\|^2 + \lambda \log p_t(X_t) d t \right] \tag{7.3}$$
|
||
|
||
这导致 Fokker-Planck 方程的正则化版本，平衡最优传输（最小动作）和最大似然（最大熵）。
|
||
|
||
**层级关系澄清**：
|
||
- **SI 是路径参数化族**：通过 $\psi, \phi, \sigma$ 定义连接端点的候选路径族
|
||
- **SB 是路径最优选择准则**：在 SI 的路径族中选择使动作泛函最小的最优路径
|
||
|
||
二者不是平级模型，SI 提供参数化空间，SB 在该空间中选择最优。
|
||
|
||
---
|
||
|
||
## 八、训练过程中的潜在问题与解决方案
|
||
|
||
### 8.1 方差估计问题
|
||
|
||
**问题**：SI 损失函数涉及对 $x_0, x_1, z$ 的期望，高维情况下方差可能很大。
|
||
|
||
**表现**：
|
||
$$\mathbb{E}\left[\| v_\theta(X_t, t) - u_t \|^2\right] = \text{Var}(v_\theta) + \text{Var}(u_t) + 2\text{Cov}(v_\theta, u_t) \tag{8.1}$$
|
||
|
||
当维度 $D$ 很大时，即使每个维度方差很小，总方差也可能很大。
|
||
|
||
**解决方案**：
|
||
|
||
1. **增大 batch size**：标准做法，通常 $N \geq 256$
|
||
2. **Antithetic variates**：使用 $-z$ 作为第二个样本，利用对称性减少方差
|
||
3. **重要性采样**：对关键时间区域（如 $t \approx 0.5$）加重采样
|
||
4. **方差归一化**：$\| v_\theta - u_t \|^2 / D$
|
||
|
||
### 8.2 时间步采样偏差
|
||
|
||
**问题**：不同时间步 $t$ 对应的分布 $p_t$ 复杂度不同。
|
||
|
||
| 区域 | 分布特点 | 挑战 |
|
||
|------|---------|------|
|
||
| $t \approx 0$（噪声端） | $p_0$ 通常是简单高斯 $\mathcal{N}(0,I)$ | Score 幅度大但结构简单 |
|
||
| $t \approx 1$（数据端） | $p_1$ 通常是复杂分布 | Score 幅度小但结构复杂 |
|
||
| $t \approx 0.5$（中间） | 混合分布，最复杂 | 信息最丰富，需要更多采样 |
|
||
|
||
**解决方案**：
|
||
|
||
1. **Logit-normal 时间采样**：使采样集中在中间区域
|
||
2. **加权损失**：$\lambda(t) = t(1-t)$ 对中间区域加权
|
||
3. **课程学习**：从简单（$t \approx 0$ 或 $t \approx 1$）到复杂逐步训练
|
||
4. **分层采样**：不同 epoch 使用不同的采样分布
|
||
|
||
### 8.3 模型架构问题
|
||
|
||
**SI 框架要求模型同时处理**：
|
||
- 空间输入 $x \in \mathbb{R}^D$
|
||
- 时间输入 $t \in [0,1]$
|
||
- 随机输入（通过噪声 $\epsilon$）
|
||
|
||
**常见架构模式**：
|
||
|
||
1. **Time MLP**：
|
||
$$t \xrightarrow{\mathrm{KL}} e \xrightarrow{\mathrm{KL}} \gamma, \beta$$
|
||
$$\tilde{x} = \gamma \cdot x + \beta$$
|
||
|
||
2. **Adaptive Normalization**（类似 DDPM 的 adaptive group norm）：
|
||
$$h = \text{GroupNorm}(x)$$
|
||
$$h = \gamma(t) \cdot h + \beta(t)$$
|
||
|
||
3. **Cross-attention**（条件生成）：
|
||
$$c \xrightarrow{\mathrm{KL}} \text{features} \xrightarrow{\mathrm{KL}} x$$
|
||
|
||
### 8.4 训练不稳定性
|
||
|
||
**问题**：某些配置下，训练可能不稳定，特别是在 $t \approx 0$ 时。
|
||
|
||
**原因**：
|
||
- 当 $\sigma(t) \to 0$ 时，条件分布退化，方差估计变得不稳定
|
||
- $\dot{\sigma}(t)$ 在边界处可能很大
|
||
- 网络需要同时拟合大范围的 score 值
|
||
|
||
**解决方案**：
|
||
|
||
1. **学习率 warmup**：初始使用较小的学习率
|
||
2. **梯度裁剪**：$\| \nabla_\theta \mathcal{L} \| \leq C$
|
||
3. **噪声调度平滑化**：确保 $\sigma(t)$ 光滑连续
|
||
4. ** EMA（指数移动平均）**：
|
||
$$\theta_{\mathrm{old}} \leftarrow m \cdot \theta_{\mathrm{old}} + (1-m) \cdot \theta \tag{8.2}$$
|
||
|
||
---
|
||
|
||
## 九、离散化与数值方法
|
||
|
||
### 9.1 Euler-Maruyama 方法
|
||
|
||
对于 SDE：
|
||
$$d X_t = f_t(X_t) d t + g_t d W_t \tag{9.1}$$
|
||
|
||
Euler-Maruyama 离散化为：
|
||
$$X_{t+\Delta t} = X_t + f_t \Delta t + g_t \sqrt{\Delta t} \cdot \epsilon, \quad \epsilon \sim \mathcal{N}(0, I) \tag{9.2}$$
|
||
|
||
**误差阶数**（场景区分）：
|
||
|
||
| 场景 | 局部截断误差 | 全局误差 |
|
||
|------|------------|---------|
|
||
| **纯扩散 SDE**（$f_t = 0$） | $O(\Delta t)$ | $O(\sqrt{\Delta t})$ |
|
||
| **漂移+扩散混合 SDE** | $O(\Delta t)$ | 随漂移场 Lipschitz 常数变化 |
|
||
| **纯确定性 ODE**（Euler） | $O(\Delta t^2)$ | $O(\Delta t)$ |
|
||
|
||
原文未做场景区分，统一写全局误差 $O(\sqrt{\Delta t})$ 是片面的。
|
||
|
||
### 9.2 预测-校正（Predictor-Corrector）方法
|
||
|
||
**Predictor Step**：使用 ODE/SDE 一步预测 $X_{t+\Delta t}$
|
||
|
||
**Corrector Step**：使用 score function 校正（Langevin Monte Carlo 风格）：
|
||
$$X_{t+\Delta t} \leftarrow X_{t+\Delta t} + \alpha \cdot g_t^2 s_\theta(X_{t+\Delta t}, t+\Delta t) \tag{9.3}$$
|
||
|
||
其中 $s_\theta$ 是预测的 score function。
|
||
|
||
### 9.3 步长选择策略
|
||
|
||
**固定步长**：
|
||
|
||
| 步数 $N$ | 步长 $\Delta t = 1/N$ | 适用场景 |
|
||
|---------|---------------------|---------|
|
||
| 1 | 1.0 | Rectified Flow（完全直线化后） |
|
||
| 4-8 | 0.25-0.125 | Re-flow 收敛后 |
|
||
| 50+ | <0.02 | 标准 SDE（未优化路径） |
|
||
|
||
**自适应步长**：
|
||
|
||
监控局部误差估计：
|
||
$$\hat{e}_t = \| X_{t+\Delta t} - \tilde{X}_{t+\Delta t} \| \tag{9.4}$$
|
||
|
||
其中 $\tilde{X}$ 是用更高阶方法计算的估计。
|
||
|
||
如果 $\hat{e}_t > \epsilon_{\mathrm{KL}}$，则减小步长重试。
|
||
|
||
### 9.4 少步采样的理论保证
|
||
|
||
**Euler 方法的误差界**：
|
||
|
||
对于 Lipschitz 向量场 $v$，Euler 方法的全局误差为：
|
||
$$\| x(T) - x_N \| \leq \frac{L \cdot T^2}{2N} \cdot e^{LT} \tag{9.5}$$
|
||
|
||
其中 $L$ 是向量场的 Lipschitz 常数，$N$ 是步数，$T$ 是终点时间。
|
||
|
||
**当路径直线化后**：
|
||
|
||
全局 Lipschitz 常数 $L$ 大幅降低（但非零，因为不同样本对应不同位移向量），允许使用更大步长，误差显著减小而非变为零。
|
||
|
||
这从数学上解释了为什么 Rectified Flow 可以用少至 1-4 步采样。
|
||
|
||
---
|
||
|
||
## 十、数学推导速查
|
||
|
||
### 10.1 核心定义
|
||
|
||
**随机插值过程**：
|
||
$$X_t = \psi(t) x_0 + \phi(t) x_1 + \sigma(t) z \tag{10.1}$$
|
||
|
||
**边界条件**：
|
||
$$\psi(0)=1, \;\phi(0)=0, \;\sigma(0)=0 \quad \text{和} \quad \psi(1)=0, \;\phi(1)=1, \;\sigma(1)=0 \tag{10.2}$$
|
||
|
||
**Itô SDE 形式**：
|
||
$$d X_t = \left[ \dot{\psi} x_0 + \dot{\phi} x_1 \right] d t + \dot{\sigma} d W_t \tag{10.3}$$
|
||
|
||
**反向时间 SDE**：
|
||
$$d X_t = \left[ f_t - g_t g_t^\top \nabla \log p_t \right] d t + g_t d \bar{W}_t \tag{10.4}$$
|
||
|
||
**Probability Flow ODE**：
|
||
$$d X_t = \left[ f_t - \tfrac{1}{2} g_t g_t^\top \nabla \log p_t \right] d t \tag{10.5}$$
|
||
|
||
### 10.2 模型对应关系
|
||
|
||
| 模型 | $\sigma(t)$ | 过程类型 | 边界条件 |
|
||
|------|------------|---------|---------|
|
||
| NF/CNF | $0$ | 确定性 ODE | $X_0=x_0, X_1=x_1$ |
|
||
| DDPM（逆向生成） | $\neq 0$ | 随机 SDE（需时间反转） | 需通过时间反转定义 |
|
||
| Rectified Flow | $0$ | 确定性 ODE（常数速度场） | $X_0=x_0, X_1=x_1$ |
|
||
| SI（一般） | 任意 | 随机 SDE | 任意满足边界条件 |
|
||
|
||
### 10.3 损失函数对照
|
||
|
||
| 目标 | 公式 | 适用模型 |
|
||
|------|------|---------|
|
||
| **速度场回归** | $\mathbb{E}\| v_\theta - (x_1-x_0) \|^2$ | Rectified Flow |
|
||
| **Score matching** | $\mathbb{E}\| s_\theta + \frac{z}{\sigma(t)} \|^2$ | Score SDE |
|
||
| **噪声预测** | $\mathbb{E}\| \epsilon_\theta - \epsilon \|^2$ | DDPM |
|
||
|
||
---
|
||
|
||
## 十一、总结
|
||
|
||
**Stochastic Interpolants 的核心贡献**：
|
||
|
||
1. **统一框架**：将 NF、DDPM、Score SDE、Flow Matching、Rectified Flow 统一为不同参数选择的随机插值过程
|
||
|
||
2. **设计空间揭示**：$\psi, \phi, \sigma$ 三个调度函数构成的设计空间，包含了所有现有模型
|
||
|
||
3. **灵活性与扩展性**：允许任意边界条件和噪声调度，为新模型设计提供了蓝图
|
||
|
||
**物理意义**：
|
||
|
||
- $x_0$ 和 $x_1$ 是路径的"端点"
|
||
- $\psi, \phi$ 控制确定性运动（漂移）
|
||
- $\sigma$ 控制随机性扩散
|
||
- 整个过程是确定性与随机性的叠加
|
||
|
||
**与工业实践的联系**：
|
||
|
||
- SD3、Flux.1 等大模型都可以在 SI 框架下理解
|
||
- 路径选择（$\psi, \phi, \sigma$）直接影响采样效率
|
||
- 最优路径设计（OT、SB）是当前研究的热点
|
||
|
||
**核心洞察**：
|
||
|
||
> **所有生成模型都是同一个数学对象的不同视角观察。** Stochastic Interpolants 揭示了这个深层统一结构，让我们能够：
|
||
> - 理解不同模型之间的关系
|
||
> - 在统一的设计空间中进行比较
|
||
> - 设计新的、结合多种模型优点的新路径
|
||
|
||
---
|
||
|
||
**延伸阅读**：
|
||
|
||
1. Albergo & Vanden-Eijnden, "Building Normalizing Flows with Stochastic Interpolants" (ICLR 2023)
|
||
2. Lipman et al., "Flow Matching for Generative Modeling" (NeurIPS 2022)
|
||
3. Song et al., "Score-Based Generative Modeling through Stochastic Differential Equations" (ICLR 2021)
|
||
4. Liu et al., "Rectified Flow: A Marginal Preserving Approach to Optimal Transport" (ICML 2023)
|
||
5. Chen et al., "Flow Matching: A Minimalist Approach to Diffusion Models" (Tutorial, 2024) |