Notes/机器学习/深度学习模型/04-扩散模型/7-统一理论/2-Schrödinger-Bridge.md

title, draft, tags

title	draft	tags
Diffusion Schrödinger Bridge / Simulation-Free Schrödinger Bridges	false	Schrödinger-Bridge 最优传输 生成模型 扩散模型 Score-Matching Flow-Matching

Schrödinger Bridge：最优传输、熵正则与生成模型的统一数学框架

---

一、从最优传输到 Schrödinger Bridge

1.1 最优传输问题回顾

Monge 最优传输问题（1781）：

给定两个概率分布 p_0 和 $p_1$，寻找一个确定性映射 $T: \mathbb{R}^D \to \mathbb{R}^D$，使得：

\min_{T} \int_{\mathbb{R}^D} c(x, T(x)) \, d p_0(x) \quad \text{s.t.} \quad T_\# p_0 = p_1 \tag{1.1}

其中：

• c(x, y) 是代价函数（cost function）
• T_\# p_0 = p_1 表示 T 将 p_0 映射到 $p_1$（push-forward）
• 典型选择：c(x, y) = \frac{1}{2} \| x - y \|^2

物理直觉：

这相当于将一堆沙土（$p_0$）搬运成指定形状（$p_1$），寻找最"省力"的方案。这里的"省力"由代价函数衡量。

Brenier 定理（当 c(x, y) = \frac{1}{2} \| x - y \|^2 时）：

如果 p_0 和 p_1 都是绝对连续的，则最优传输映射 T^* 存在且唯一，且可以表示为：

T^*(x) = \nabla \phi(x) \tag{1.2}

其中 \phi 是一个凸函数。

1.2 Schrödinger 问题的引入

需区分三类不同的 Schrödinger Bridge：

1. 经典统计物理 SB（1931 原始）：平衡态假设，固定稳态分布，路径熵最小化
2. 动力学 SB：以路径运动动能 \int \|\dot{X}_t\|^2 dt 为优化目标
3. 生成建模非平衡 SB：当前 AI 主流，用于时序/插值分布演化，本文聚焦此版本

原始 Schrödinger 问题（1931）：

设 W_t 是维纳过程。寻找概率分布族 p_t 使得在固定两端边际分布 $p_0$、p_1 下，最小化路径相对熵：

\min_{p_t} D_{\mathrm{KL}}\left( p_{0:1} \| \mathbb{P}_W \right) \tag{1.3}

其中 p_{0:1} 是路径空间上的分布，\mathbb{P}_W 是维纳过程的路径测度。

动力学版本 SB（动作泛函形式）：

\min_{X_t} \mathbb{E}_{p_t} \left[ \int_0^1 \frac{1}{2} \left\| \frac{d X_t}{d t} \right\|^2 d t \right] \tag{1.4}

重要澄清：式(1.3)与式(1.4)仅在单位扩散系数、平衡稳态假设下弱等价，并非无条件等价的原始定义。式(1.4)是动力学版本的 SB，而非 1931 年薛定谔原始定义。

物理直觉：

在所有连接给定端点分布的随机过程中，寻找"最像布朗运动"的路径。代价函数 \int \|\dot{X}_t\|^2 d t 惩罚路径的"不平滑性"。

1.3 Schrödinger Bridge 与最优传输的联系

定理（Schrödinger Bridge 的 OT 等价性）：

当路径无噪声（即 $\sigma(t) = 0$）时，Schrödinger Bridge 退化为 Monge 最优传输问题。

证明思路：

无噪声时，路径是确定性的：$X_t = \gamma(t)$，其中 \gamma 是连接 x_0 和 x_1 的曲线。

动作泛函：

\int_0^1 \frac{1}{2} \|\dot{\gamma}(t)\|^2 d t

由积分型柯西-施瓦茨不等式：

\left\| \int_0^1 \dot{\gamma}(t) d t \right\|^2 \leq \int_0^1 1^2 d t \cdot \int_0^1 \|\dot{\gamma}(t)\|^2 d t = \int_0^1 \|\dot{\gamma}(t)\|^2 d t

因此：

\int_0^1 \frac{1}{2} \|\dot{\gamma}(t)\|^2 d t \geq \frac{1}{2} \left\| \int_0^1 \dot{\gamma}(t) d t \right\|^2 = \frac{1}{2} \| x_1 - x_0 \|^2 \tag{1.5}

等号成立当且仅当 $\dot{\gamma}(t) = \text{const}$（即匀速直线运动）。

因此，无噪声 SB 的最优路径是直线，这正是 Rectified Flow 的设计。

---

二、Schrödinger Bridge 的数学理论

2.1 路径空间与 Radon-Nikodym 导数

路径空间定义：

设 \Omega = C([0,1]; \mathbb{R}^D) 是 [0,1] 到 \mathbb{R}^D 的连续函数空间（路径空间）。

对于维纳过程 $W_t$，定义路径测度 $\mathbb{P}_W$（维纳 measure）。

** Schrödinger Bridge 是 \mathbb{P}_W 的条件版本**：

\mathbb{P}^* = \arg\min_{\mathbb{P}: \pi_0^* \mathbb{P} = p_0, \; \pi_1^* \mathbb{P} = p_1} D_{\mathrm{KL}}(\mathbb{P} \| \mathbb{P}_W) \tag{2.1}

其中 \pi_t^* 是到时刻 t 的边际投影。

2.2 Fokker-Planck 方程约束

定理（SB 的必要性条件）：

设 \mathbb{P}^* 是 Schrödinger Bridge，则 \mathbb{P}^* 在路径空间上的分布密度可以分解为：

\frac{d \mathbb{P}^*}{d \mathbb{P}_W} = \frac{\phi_0(X_0)}{\mathbb{E}[\phi_0(W_0)]} \cdot \frac{\phi_1(X_1)}{\mathbb{E}[\phi_1(W_1)]} \tag{2.2}

其中 \phi_0, \phi_1 是某个（非负）函数。

对应的 Fokker-Planck 方程（仅适用于单位扩散系数 D=1 的标准布朗运动）：

边际分布 p_t 满足：

\frac{\partial p_t}{\partial t} = -\nabla \cdot (v_t p_t) + \frac{1}{2} \nabla^2 p_t \tag{2.3}

其中 v_t 是待确定的速度场，\nabla^2 是拉普拉斯算子。通用形式应保留扩散系数 $D$：$\frac{1}{2} D \nabla^2 p_t$。

2.3 互熵最小化视角

定义（路径 KL 散度）：

设 p_{0:1} 是路径空间的联合分布密度，定义路径 KL 散度（而非互熵）：

D_{\mathrm{KL}}(p_{0:1} \| \mathbb{P}_W) = \mathbb{E}_{p_{0:1}} \left[ \log \frac{d p_{0:1}}{d \mathbb{P}_W} \right] \tag{2.4}

SB 的变分形式（最小化路径 KL 而非互熵）：

\min_{p_{0:1}} D_{\mathrm{KL}}(p_{0:1} \| \mathbb{P}_W) \quad \text{s.t.} \quad \pi_0^* p_{0:1} = p_0, \; \pi_1^* p_{0:1} = p_1 \tag{2.5}

通过拉格朗日乘子法求解：

引入边际约束的拉格朗日乘子 $\lambda_0(x_0), \lambda_1(x_1)$：

\mathcal{L} = D_{\mathrm{KL}}(p_{0:1} \| \mathbb{P}_W) + \int \lambda_0(x_0) (p_{0:1} - p_0) d x_0 + \int \lambda_1(x_1) (p_{0:1} - p_1) d x_1 \tag{2.6}

通过变分法求解，得到最优路径的分解形式（式 2.2）。

---

三、Simulation-Free Schrödinger Bridges（SFSB）

3.1 核心问题：为什么需要 Simulation-Free？

原始 SB 的计算瓶颈：

直接求解 Schrödinger Bridge 需要：

1. 在路径空间上优化分布
2. 通过 Fokker-Planck 方程反向计算 score
3. 这需要大量的 SDE 模拟（simulation），算力成本极高

"Simulation-Free" 的真正核心内涵：

SFSB 舍弃全局路径遍历，仅利用两端静态分布和时序插值完成训练，彻底摆脱传统 SB 依赖 SDE 轨迹模拟的算力瓶颈。不通过显式模拟 SDE 路径来训练，而是通过端点驱动的条件分布进行训练。

3.2 SFSB 的数学框架

关键观察：

在独立采样的假设下（x_0 \sim p_0, x_1 \sim p_1 独立），Schrödinger Bridge 可以通过端点的条件分布来描述。

定理（SFSB 的损失函数）：

设 $p(x_t | x_0, x_1) = \mathcal{N}(\psi(t) x_0 + \phi(t) x_1, \sigma(t)^2 I)$，则 SFSB 的训练目标是：

\mathcal{L}_{\mathrm{KL}}(\theta) = \mathbb{E}_{t, x_0, x_1, z} \left[ \lambda(t) \left\| s_\theta(X_t, t) - \nabla_{X_t} \log p(x_t | x_0, x_1) \right\|^2 \right] \tag{3.1}

其中：

X_t = \psi(t) x_0 + \phi(t) x_1 + \sigma(t) z \tag{3.2}

3.3 与 Score Matching 的等价性

SFSB 等价去噪分数匹配需满足强约束：

当以下全部三个前置条件同时满足时，SFSB 损失（式 3.1）等价于条件去噪 score matching：

1. 端点独立采样：x_0 \sim p_0, x_1 \sim p_1 相互独立
2. 高斯插值核：p(x_t | x_0, x_1) = \mathcal{N}(\psi x_0 + \phi x_1, \sigma^2 I)
3. 固定时序调度：\psi, \phi, \sigma 预先确定（非学习参数）

满足约束时的推导：

条件分布 $p(x_t | x_0, x_1) = \mathcal{N}(\mu, \sigma^2 I)$，其中 $\mu = \psi x_0 + \phi x_1$。

Score 函数：

\nabla_{X_t} \log p(x_t | x_0, x_1) = -\frac{X_t - \mu}{\sigma^2} = -\frac{\sigma z}{\sigma^2} = -\frac{z}{\sigma} \tag{3.4}

因此训练目标变为预测 $-z/\sigma$，即标准的去噪 score matching。

脱离约束条件时，SFSB 与 Score Matching 仅为形式相似，数学期望层面并不等价。

3.4 与 Flow Matching 的联系

当 \sigma(t) = 0 时（无噪声插值）：

SFSB 退化为条件速度场回归：

\mathcal{L}_{\mathrm{KL}}(\theta) = \mathbb{E}_{t, x_0, x_1} \left[ \lambda(t) \left\| v_\theta(X_t, t) - (\dot{\psi} x_0 + \dot{\phi} x_1) \right\|^2 \right] \tag{3.5}

这正是 Flow Matching 的条件流匹配损失！

3.5 统一视角下的模型谱系

Schrödinger Bridge（统一框架）
├── 无噪声路径（σ=0）
│   └── Flow Matching / Rectified Flow
└── 有噪声路径（σ≠0）
    ├── Score Matching（条件版本）
    └── DDPM（特殊噪声调度）

---

四、Entropy-Regularized Schrödinger Bridge

4.1 熵正则化的引入

标准 SB 的问题：

当 p_0 和 p_1 都是离散分布时，最优路径可能过于"曲折"——因为它可以集中于少数"高效"路径，导致解的稀疏性。

熵正则化 SB：

\min_{p_t} \mathbb{E}_{p_t} \left[ \int_0^1 \frac{1}{2} \left\| \frac{d X_t}{d t} \right\|^2 d t \right] + \lambda \mathbb{E}_{p_t} \left[ \log p_t(X_t) \right] \tag{4.1}

第二项是路径熵的正则化，鼓励路径更加"均匀"分布。

4.2 Sinkhorn 算法的联系

通过路径空间分解：

设 u(x_0), v(x_1) 是 (x_0, x_1) 上的势函数，则熵正则化 OT 的最优联合分布为：

\pi^*(x_0, x_1) = \pi_0(x_0) \pi_1(x_1) \exp(-c(x_0, x_1)/\lambda) \cdot \frac{v(x_0)}{u(x_1)} \tag{4.2}

其中 u, v 是通过 Sinkhorn 迭代交替更新的势函数（标准形式为 u 在分子、v 在分母，与原文相反）。

与 Schrödinger Bridge 的联系：

熵正则化 OT 的最优传输映射等价于某个熵正则化 SB 的终点映射。

4.3 训练中的应用

Practical Entropic SB 训练：

1. 初始化：通过 Sinkhorn 计算初始配对 \pi^*
2. 路径采样：沿配对 \pi^* 采样 (x_0, x_1) 对
3. 路径优化：在采样路径上训练 score/velocity 模型
4. 迭代改善：重复上述步骤直到收敛

---

五、Schrödinger Bridge 与 Diffusions 的深层联系

5.1 互为对偶的关系

定理（SB-Diffusion 对偶性）：

设 p_t 是某个扩散过程的边际分布，s_t = \nabla \log p_t 是对应的 score function。则 p_t 是以下 Schrödinger Bridge 问题的解：

\min_{q_t} D_{\mathrm{KL}}\left( q_{0:1} \| p_{0:1}^{\mathrm{ref}} \right) \quad \text{s.t.} \quad q_0 = p_0, \; q_1 = p_1 \tag{5.1}

其中 p_{0:1}^{\mathrm{ref}} 是参考过程（通常为维纳过程）的路径测度。此对偶关系需对齐正向/逆向时间流向：DDPM 前向加噪对应正向 SB，反向去噪生成对应逆向 SB。

直观理解：

• Diffusion 模型：从数据分布出发，通过加噪过程趋近于噪声分布
• Schrödinger Bridge：从噪声分布出发，通过"最像布朗运动"的路径趋近于数据分布

两者是同一个数学对象的不同方向描述。

5.2 路径测度的表示

克喇莫尔表示（仅适用于高斯端点分布）：

设 \mathbb{P}^* 是 Schrödinger Bridge，当端点分布 p_0, p_1 为高斯时，最优路径可以写成：

X_t^* = \mathbb{E}[W_t \mid X_0 \sim p_0, \; X_1 \sim p_1] \tag{5.2}

即在给定端点分布的条件下，布朗运动的条件期望。

注意：此显式表示仅严格成立于纯高斯分布场景，对于复杂多峰数据无通用显式表达式，需通过数值方法近似。

5.3 与 Score SDE 的等价性

定理（SB = Score SDE）：

对于任意的 Schrödinger Bridge $\mathbb{P}^*$，存在一个对应的 Score SDE，其解的边际分布就是 \mathbb{P}^* 的边际分布。

构造：

给定 \mathbb{P}^* 的 Fokker-Planck 方程：

\frac{\partial p_t}{\partial t} = -\nabla \cdot (v_t^{\mathrm{SB}} p_t) \tag{5.3}

其中 v_t^{\mathrm{SB}} 是 SB 对应的最优速度场。

定义 Score SDE：

d X_t = \left[ v_t^{\mathrm{SB}} + \nabla \log p_t \right] d t + d W_t \tag{5.4}

可以验证，这个 SDE 的边际分布与 \mathbb{P}^* 一致。

---

六、训练过程与算法

6.1 SFSB 的训练算法

Algorithm: Simulation-Free Schrödinger Bridge Training

输入：

• 数据分布 p_{\mathrm{data}}
• 噪声分布 p_0 = \mathcal{N}(0, I)
• 调度函数 \psi, \phi, \sigma
• 网络 $s_\theta$（score network）或 $v_\theta$（velocity network）

for iteration = 1 to N do:

1. 采样批次 ${(x_0^i, x_1^i)}$，其中 $x_0^i \sim p_0$，$x_1^i \sim p_{\mathrm{data}}$（独立）
2. 采样时间步 t \sim \text{LogitNormal}(\mu, \sigma^2)
3. 采样噪声 z^i \sim \mathcal{N}(0, I)
4. 计算插值点：X_t^i = \psi(t) x_0^i + \phi(t) x_1^i + \sigma(t) z^i
5. 计算目标 score：s^* = -z^i / \sigma(t)
6. 计算损失：\mathcal{L} = \frac{1}{B} \sum_i \lambda(t) \| s_\theta(X_t^i, t) - s^* \|^2
7. 反向传播并更新 \theta

输出：训练好的网络 s_\theta 或 v_\theta

6.2 带 OT 配对的改进版本

Algorithm: OT-Regularized SFSB

初始化：

1. 通过 Sinkhorn 算法计算 p_0 和 p_1 之间的最优配对 \pi^*
2. 令 $x_1^i = T(x_0^i)$，其中 T 是 \pi^* 对应的映射

主循环（与上面相同，但配对固定）

收敛性：

当配对是最优传输映射时，直线插值路径就是 SB 的最优路径。

6.3 多步 Rectification（Re-flow 的 SB 解释）

Re-flow 是 SB 最优路径的工程数值近似解法：

Re-flow（重边缘化整流匹配）并非等价于 SB，而是 SB 最优路径的工程数值近似实现。二者关系：

• SB：理论最优路径定义（路径熵最小化）
• Re-flow：通过迭代数值近似求解 SB 的工程落地方法

具体而言：

1. 初始化：假设独立配对 x_1 = T^{(0)}(x_0)
2. 迭代：通过当前配对训练 $v_\theta$，生成新轨迹，用轨迹终点更新配对 T^{(k+1)}
3. 收敛：当配对趋近于 SB 最优传输映射时，路径接近 SB 最优

---

七、潜在问题与解决方案

7.1 路径爆炸问题

问题：当 OT 配对不是双射时，路径可能交叉或"爆炸"。

数学表现：

• 学习到的映射 T: x_0 \to x_1 不是光滑的双射
• 某些区域密度过高或过低
• 轨迹曲率在某些点变得非常大

解决方案：

1. Unbalanced OT：允许边际分布不完全匹配

\min_{\pi} \int c(x_0, x_1) d\pi + \lambda_0 D_{\mathrm{KL}}(\pi_0 \| p_0) + \lambda_1 D_{\mathrm{KL}}(\pi_1 \| p_1) \tag{7.1}

2. 正则化：在损失中添加曲率惩罚

\mathcal{L} + \nu \mathbb{E}[\kappa(t)^2] \tag{7.2}

7.2 训练不稳定性

问题：当 \sigma(t) 在边界附近变化剧烈时，训练可能不稳定。

原因分析：

• \sigma(0) = 0 时，条件分布退化
• \dot{\sigma}(t) 在边界处可能奇异
• score 的幅度 \| s \| \sim 1/\sigma(t) 在边界处爆炸

解决方案：

1. 平滑调度：使用 \sigma(t) = \epsilon + t(1-t) 代替纯 t(1-t)
2. 截断梯度：对边界区域的梯度进行裁剪
3. 双时间步策略：分别训练 t \approx 0 和 t \approx 1 的模型

7.3 维度灾难

问题：当维度 D 很大时，OT 求解和路径采样都变得困难。

表现：

• Sinkhorn 迭代收敛慢：O(N^2 D)
• 轨迹方差随 D 线性增长
• 模型需要拟合 D 维 score 向量

解决方案：

1. 子空间建模：在低维子空间中进行 OT 和插值
2. 分块处理：将高维数据分块独立处理
3. 对比学习：使用对比损失辅助学习

---

八、与现有模型的关系总结

8.1 模型对应关系图

生成模型（概率传输）
│
├── 最优传输（OT）
│   ├── Monge 问题
│   └── Brenier 定理（凸势函数）
│
└── Schrödinger Bridge
    ├── 标准 SB（路径熵最小化）
    ├── 熵正则化 SB（+OT 正则）
    └── Simulation-Free SB
        ├── 当 σ=0：Flow Matching / Rectified Flow
        ├── 当 σ≠0：Score Matching / DDPM
        └── 中间情况：条件插值

8.2 核心公式对照

概念	数学表达	物理意义
Monge OT	\min_T \int c(x, Tx) dp_0	确定性映射，最小化传输代价
SB	\min_{p_t} D_{\mathrm{KL}}(p_{0:1}\| \mathbb{P}_W)	找到最像布朗运动的路径
熵正则化 SB	\min (D_{\mathrm{KL}} + \lambda \cdot \text{OT})	OT 与最大熵的平衡
SFSB	\min_\theta \mathbb{E}\| s_\theta - \nabla\log p\|	端点驱动的 score 训练

8.3 采样效率对比

方法	采样步数	理论依据
DDPM	50-1000	多步去噪链
Score SDE	10-1000	SDE 求解
Flow Matching	1-10	路径直线化
SB（精确）	1（理论上）	最优传输映射
SFSB	1-10	近似 SB

---

九、数学推导速查

9.1 Schrödinger Bridge 核心方程

路径 KL 散度（而非互熵）：

D_{\mathrm{KL}}(p_{0:1} \| \mathbb{P}_W) = \mathbb{E}_{p_{0:1}} \left[ \log \frac{d p_{0:1}}{d \mathbb{P}_W} \right] \tag{9.1}

SB 变分形式：

\min_{p_{0:1}} D_{\mathrm{KL}}(p_{0:1} \| \mathbb{P}_W) \quad \text{s.t.} \quad \pi_0^* p_{0:1} = p_0, \; \pi_1^* p_{0:1} = p_1 \tag{9.2}

Fokker-Planck 约束（仅适用于单位扩散系数 $D=1$）：

\frac{\partial p_t}{\partial t} = -\nabla \cdot (v_t p_t) + \frac{1}{2} \nabla^2 p_t \tag{9.3}

9.2 SFSB 损失函数

条件 score matching：

\mathcal{L}_{\mathrm{KL}} = \mathbb{E}_{t, x_0, x_1, z} \left[ \lambda(t) \left\| s_\theta(X_t, t) + \frac{z}{\sigma(t)} \right\|^2 \right] \tag{9.4}

条件 velocity matching（当 $\sigma=0$）：

\mathcal{L}_{\mathrm{KL}} = \mathbb{E}_{t, x_0, x_1} \left[ \lambda(t) \left\| v_\theta(X_t, t) - (\dot{\psi}x_0 + \dot{\phi}x_1) \right\|^2 \right] \tag{9.5}

9.3 OT 配对公式

Sinkhorn 势函数更新：

u^{(k+1)} = \frac{p_0}{K v^{(k)}}, \quad v^{(k+1)} = \frac{p_1}{K^T u^{(k)}} \tag{9.6}

最优配对：

\pi^*(x_0, x_1) = \text{diag}(u) K \text{diag}(v) \tag{9.7}

---

十、总结

Schrödinger Bridge 的核心价值：

1. 统一理论框架：将最优传输、扩散模型、score matching、F low matching 统一在一个变分框架下

2. 最优路径解释：SB 提供了"什么是最好的从噪声到数据的路径"的理论答案

3. Simulation-Free 训练：SFSB 使得在不模拟 SDE 路径的情况下训练成为可能

与生成模型实践的联系：

• SD3、Flux 等大模型都可以在 SB 框架下理解
• Re-flow 是 SB 的迭代近似求解
• 最优传输配对是 SB 与工业实践的关键桥梁

核心洞察：

Schrödinger Bridge 回答了"最优生成路径应该是什么样的"这个问题。 它告诉我们，在所有可能的随机过程中，最优的那条是"最像布朗运动"且满足端点约束的路径。这个结论将最优传输的确定性与扩散模型的随机性统一了起来。

---

延伸阅读：

1. Schrödinger, "Über die Umkehrung der Naturgesetze" (1931)
2. Föllmer, "Snell Envelope and Schrödinger Bridge" (1988)
3. Léonard, "From Schrödinger to McCann's Optimal Transport" (2014)
4. Wang et al., "Schrödinger Bridge for Generative Modeling" (2024)
5. De Bortoli et al., "Simulation-Free Schrödinger Bridges via Score and Flow Matching" (2024)