Reflected Schrödinger Bridge Matching¶

会议: NeurIPS 2026
arXiv: 2607.03626
OpenReview: https://openreview.net/forum?id=OO3amJIfRg
代码: https://github.com/viktor765/rsbm
领域: 优化/理论
关键词: 薛定谔桥, 反射扩散过程, 桥匹配, 有界域生成, 熵正则最优输运

一句话总结¶

本文提出反射薛定谔桥匹配（RSBM），将部分无模拟（partially simulation-free）的迭代马尔可夫拟合推广至单位超立方体上的反射布朗运动参考过程，通过解析截断高斯和与两步反射桥中点采样，在几乎零额外耗时下严格消除了高维受限域图像配对生成的越界伪影。

研究背景与动机¶

高维生成建模近期通过扩散模型与流匹配技术取得了显著进展，且随着整流流（Rectified Flows）等算法的引入，生成轨迹被显著拉直并趋向于最优输运。然而，确定性流匹配在有限步下缺乏严格的最优输运理论保证，而薛定谔桥（Schrödinger Bridge, SB）作为熵正则化最优输运（Entropic Optimal Transport）的动态连续推广，能够在保持边缘分布约束的同时给出严格的概率流动力学解释。随着扩散薛定谔桥匹配（DSBM）与迭代马尔可夫拟合（IMF / $\alpha$-IMF）的出现，研究者们借助 Doob $h$-变换与类似流匹配的回归损失函数，摆脱了传统前向-后向 SDE 复杂的全轨迹采样，在无约束欧氏空间中实现了高维薛定谔桥的高效求解。

然而，现实中许多关键生成任务的数据天然严格分布在有界紧致域中（例如灰度或 RGB 图像的像素值被严格限制在 $[0,1]^d$ 内）。现有的无约束薛定谔桥与扩散模型在推理时极易采样到域外点（Out-of-Bounds），必须依赖人工启发式的截断裁剪（Clipping），这不仅破坏了测度变换在边界处的理论一致性，还会引起色彩过度饱和或结构扭曲。此前针对受限域薛定谔桥的尝试依赖正倒向随机微分方程（FBSDE）理论，训练时必须在每一步模拟完整动力学轨迹，且损失函数包含高成本的散度项（Divergence），需要二次反向传播，在大规模高维图像配对任务中极其低效且难以扩展。

受限域生成建模的核心矛盾在于：如何在维持物理边界硬约束（反射动力学）与最优输运严格解的同时，保留流匹配级别的高效部分无模拟训练范式。本文的切入角度是直接将参考过程限定在单位超立方体上的反射布朗运动（Reflected Brownian Motion, RBM），并证明其在马尔可夫投影与倒易投影下的理论自洽性。核心 idea：将迭代马尔可夫拟合拓展至反射 SDE 空间，利用超立方体反射算子的独立可分性推导出高斯和形式的解析漂移回归目标，并设计基于反射逆像权重的两步精确桥中点采样算法，实现计算开销与无约束流匹配相当的高维反射薛定谔桥生成。

方法详解¶

整体框架¶

反射薛定谔桥匹配（Reflected Schrödinger Bridge Matching, RSBM）旨在两个支撑在紧致域 $\mathcal{D} = [0,1]^d$ 上的经验分布 $\pi_0$ 与 $\pi_1$ 之间，求解以反射布朗运动 $\mathbb{Q}^r$ 为参考测度的最小 KL 散度路径测度 $\mathbb{P}^*$。算法整体遵循迭代马尔可夫拟合（IMF 及在线连续更新版本 $\alpha$-RSBM）的双重投影交替循环：在倒易投影步骤中，通过从当前端点耦合采样反射桥生成配对轨迹样本；在马尔可夫投影步骤中，通过神经网络回归带漂移修正的反射 SDE 漂移项，并在前向与反向动力学上协同学习。

整个流程由端点配对生成、反射桥中点插值采样、单步神经网络正反向漂移拟合，以及欧拉-丸山（Euler-Maruyama）反射数值积分四个核心阶段构成，数据流紧凑闭环：

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入分布 π₀ 与 π₁<br/>支撑严格约束于 [0, 1]ᵈ"] --> B["反射参考过程与理论投影保证<br/>定义 RBM 参考测度与马尔可夫投影极小化"]
    B --> C["两步反射桥中点采样<br/>逆像权重采样 Y₁ + 高斯桥插值 + 反射映射 R"]
    C --> D["维度解耦的解析 Score 漂移回归<br/>截断高斯和独立分解，计算修正项 v*(t, x)"]
    D --> E["双向在线训练与反射数值模拟<br/>α-RSBM 统一参数化网络 + 反射 Euler-Maruyama 积分"]
    E --> F["输出：边界严格无越界的双向生成样本"]

关键设计¶

1. 反射参考过程与理论投影保证：将薛定谔桥扩展至具有法向反射的边界紧致域

标准流匹配与扩散模型假定参考过程为全空间自由布朗运动，当遇到 $[0,1]^d$ 等硬边界时会出现测度泄漏。RSBM 将参考过程 $\mathbb{Q}^r$ 显式设定为带局部时（Local Time）补偿的反射随机微分方程： $$d X_t = \mu(t, X_t) dt + \sigma(t) d W_t + n(X_t) d L_t$$ 其中 $n(x)$ 为边界处的内向单位法向量，标量过程 $L_t$ 为边界局部时，仅当 $X_t \in \partial \mathcal{D}$ 时递增，以法向冲量阻止粒子逃逸。作者在理论上给出了关键推广（Proposition 4.1）：设 $\mathbb{M}^* = \operatorname{proj}_{\mathcal{M}^r}(\Pi)$ 为倒易类分布 $\Pi$ 在反射马尔可夫测度集 $\mathcal{M}^r$ 上的投影，论文利用 Girsanov 定理绕过了反射边界上无穷小生成元（Infinitesimal Generator）定义域难以刻画的数学障碍，证明了 $\mathbb{M}^*$ 不仅在所有时刻精确保留边缘分布 $\mathbb{M}^*_t = \Pi_t$，并且唯一极小化路径空间 KL 散度： $$\operatorname{D_{KL}}(\Pi \mid\mid \mathbb{M}^*) = \frac{1}{2} \int_0^1 \frac{1}{\sigma^2(t)} \mathbb{E}_{\Pi_{0,t}} \left[ \left\| \sigma^2(t) \mathbb{E}_{\Pi_{1 \mid 0, t}} [\nabla \log \mathbb{Q}^r_{1 \mid t}(X_1 \mid X_t) \mid X_0, X_t] - v^*(t, X_t) \right\|^2 \right] dt$$ 该结论奠定了反射设置下采用最小二乘回归目标训练漂移修正项的理论正当性。

2. 两步反射桥中点采样：利用逆像离散权重与高斯桥实现无偏高效采样

在倒易投影步骤中，模型必须从已知端点 $(x_0, x_1)$ 采样中间时刻 $t$ 的反射布朗桥 $X_t \sim \mathbb{Q}^r_{t \mid 0, 1}(\cdot \mid X_0=x_0, X_1=x_1)$。若直接对带边界反射的 SDE 进行时间离散数值积分，将产生巨大的计算开销并引入时间步累积离散误差。论文发现反射算子 $R$ 满足轨道等价性：无约束布朗运动 $W_t$ 经折叠映射 $R(W_t)$ 后的边缘与联合测度与反射布朗运动 $X_t$ 完全相同。基于此，算法引入解析的两步采样机制（Algorithm 2）：首先，利用贝叶斯法则，在终点 $x_1$ 对应的多维逆像集合 $\mathcal{I}(x_1) = \{y \in \mathbb{R}^d \mid R(y) = x_1\}$ 上按离散概率分布采样未折叠虚拟终点 $Y_1 \sim \sum_{y_1 \in \mathcal{I}(x_1)} \mathbb{Q}_{1 \mid 0}(y_1 \mid x_0) \delta_{y_1}$；接着，利用标准自由布朗桥在 $x_0$ 与 $Y_1$ 之间解析采样过渡点 $Y_t \sim \mathcal{N}(t Y_1 + (1-t) x_0, t(1-t)\sigma^2 I_d)$；最后直接应用反射算子输出 $X_t = R(Y_t)$。该采样步骤完全是闭式计算，无需任何中间动力学步进，不仅从理论上保证了样本严格服从反射桥分布，而且执行速度与无约束高斯桥采样一致。

3. 维度解耦的解析 Score 漂移回归：基于截断高斯和的低成本目标计算

神经网络在训练阶段拟合的目标漂移项包含反射转移密度的 Score 函数 $\nabla \log \mathbb{Q}^r_{1 \mid t}(x_1 \mid x_t)$。在无约束欧氏空间中该项退化为平凡的线性漂移 $(x_1 - x_t) / (1-t)$，但在反射边界处，转移密度由沿各个反射平面的镜像高斯核叠加而成。若在整个高维空间直接展开高斯和，计算复杂度将随维度 $d$ 指数爆炸（即包含 $K^d$ 项）。论文指出，由于单位立方体各坐标轴之间相互正交且参考布朗运动具有独立增量，反射核在坐标维度上天然具有可乘分解性： $$\mathbb{Q}^r_{t \mid s}(x_t \mid x_s) = \prod_{i=1}^d \mathbb{Q}^{r, i}_{t \mid s}(x_t^i \mid x_s^i) = \prod_{i=1}^d \sum_{y_t^i \in \mathcal{I}(x_t^i)} \mathbb{Q}^i_{t \mid s}(y_t^i \mid x_s^i)$$ 因此，各维度的对数梯度可完全独立进行一维标量计算： $$(\nabla_{x_s} \log \mathbb{Q}^r_{t \mid s}(x_t \mid x_s))_i = \frac{\sum_{y_t^i \in \mathcal{I}(x_t^i)} \frac{d}{d x_s^i} \mathbb{Q}^i_{t \mid s}(y_t^i \mid x_s^i)}{\sum_{y_t^i \in \mathcal{I}(x_t^i)} \mathbb{Q}^i_{t \mid s}(y_t^i \mid x_s^i)}$$ 在实际工程实现中，每个维度仅需在 $x_t^i$ 邻域内截断保留 11 个一维反射镜像点即可达到极高的数值精度，从而将高维 Score 函数计算开销直接降至与维度成线性的 $\mathcal{O}(K \cdot d)$，避免了以往方法中复杂的网格离散化或显式散度计算。

4. 双向在线训练与反射数值模拟：单网络参数化对称流与极速边界推理

为避免标准 IMF 算法中交替重训前向网络与反向网络所带来的误差累积与高昂计算延迟，本文采用 $\alpha$-RSBM 在线连续迭代协议。系统使用单一 U-Net 骨干网络 $v_\theta(t, X_t; \operatorname{dir})$，通过在时间正弦位置编码后拼接方向嵌入（Directional Embedding）并通过 FiLM 调制网络各残差块中的 GroupNorm 层，实现前向流（$t=0 \to 1$）与后向流（$t=1 \to 0$）的联合参数化。在生成采样阶段，算法使用带边界处理的 Euler-Maruyama 离散积分：每走一步数值更新后，对处于内部和穿越边界的坐标值直接执行模运算反射折叠 $R(y)$；仅在最终时间步 $t=1$ 对极微量可能的数值截断误差执行硬投影。这保证了推理生成的整条几何轨迹在物理上严格闭锁在单位超立方体内。

损失函数 / 训练策略¶

算法综合前向漂移损失 $L_\to$ 与反向漂移损失 $L_\leftarrow$ 进行端到端优化，损失函数定义为对称联合均方误差： $$\mathcal{L}(\theta; \Pi) = \frac{1}{2} \left( \mathcal{L}_\to(\theta; \Pi) + \mathcal{L}_\leftarrow(\theta; \Pi) \right)$$ 其中前向损失为： $$\mathcal{L}_\to(\theta; \Pi) = \int_0^1 \frac{1}{\sigma^2(t)} \mathbb{E}_{\Pi_{t, 1}} \left[ \left\| \sigma^2(t) \nabla \log \mathbb{Q}^r_{1 \mid t}(X_1 \mid X_t) - v_\theta(t, X_t; \to) \right\|^2 \right] dt$$ 反向损失 $\mathcal{L}_\leftarrow$ 则通过关于时间反转的对称形式定义。训练采用两阶段在线机制： 1. 预训练阶段（Pretrain）：令端点直接采样自边缘分布的独立乘积 $(X_0, X_1) \sim \pi_0 \otimes \pi_1$，通过拟合直连反射桥提供优质的初始向量场（MNIST $\to$ EMNIST 训练 50,000 步；AFHQ 训练 75,000 步）。 2. 微调阶段（Finetune）：在每个优化步中，模型利用当前权重的指数移动平均（EMA，衰减率 0.999）数值求解正反向 SDE 获得近似端点 $\hat{X}_1$ 与 $\hat{X}_0$，重新构筑近似倒易分布并执行单步梯度更新（MNIST $\leftrightarrow$ EMNIST 进行 100,000 步；AFHQ 进行 20,000 步）。优化器采用 Adam，学习率分别为 $1 \times 10^{-4}$ 与 $2 \times 10^{-4}$，梯度裁剪范数阈值设为 1.0。

实验关键数据¶

主实验¶

实验在两个无配对图像到图像（Unpaired Image-to-Image）高维转换基准上展开：MNIST ↔ EMNIST（$28 \times 28$ 灰度手写数字与字母转换）以及 AFHQv2 cat ↔ wild（$64 \times 64$ 彩色猫与野生动物头像转换）。基线对比标准无约束的 $\alpha$-DSBM（为了计算有效评估指标，$\alpha$-DSBM 的生成样本经过显式截断处理）。评价指标包括衡量分布拟合保真度的 FID、衡量像素级起点终点对应保真度的平均平方距离（Mean Square Distance, MSD），以及在 AlexNet 特征空间度量感知一致性的 LPIPS。

数据集迁移任务	方法模型	MSD ↓	LPIPS ↓	FID ↓	训练总耗时 (GPU)
MNIST → EMNIST	$\alpha$-DSBM (无约束截断)	0.416	—	6.41	31.5 h (1× A100)
	$\alpha$-RSBM (本文)	0.356	—	5.22	32.0 h (1× A100)
EMNIST → MNIST	$\alpha$-DSBM (无约束截断)	0.415	—	8.58	31.5 h (1× A100)
	$\alpha$-RSBM (本文)	0.362	—	8.23	32.0 h (1× A100)
AFHQ: cat → wild	$\alpha$-DSBM (无约束截断)	0.131	0.239	27.6	108.0 h (4× A100)
	$\alpha$-RSBM (本文)	0.084	0.248	26.1	109.0 h (4× A100)
AFHQ: wild → cat	$\alpha$-DSBM (无约束截断)	0.200	0.302	31.5	108.0 h (4× A100)
	$\alpha$-RSBM (本文)	0.082	0.244	33.8	109.0 h (4× A100)

消融实验与越界伪影分析¶

由于 $\alpha$-RSBM 的设计初衷是彻底消除越界问题，论文专门对无约束 $\alpha$-DSBM 生成过程中的像素越界（Out-of-Bounds）严重程度进行了全面统计量化，论证了反射约束的不可或缺性：

数据集	样本越界比例	平均越界像素占比	越界像素平均偏差量	单图最大越界偏差量
MNIST / EMNIST	100.0%	34% – 46%	0.002 – 0.003	0.008 – 0.013
AFHQ (64×64)	~30.0% (越界>1%)	~2.0% (全部通道)	0.020	0.030 (图像均值)

关键发现¶

结构保真度显著提升：在全部迁移任务中，$\alpha$-RSBM 的均方距离（MSD）均显著优于 $\alpha$-DSBM（例如 AFHQ wild $\to$ cat 任务上 MSD 从 0.200 骤降至 0.082，降幅达 59%）。这表明在严格反射边界约束下，薛定谔桥保留了更加紧致的几何传输路径，生成的图像在背景细节和整体色调上与源图像对应得更加精确。
训练开销几乎零增长：相较于此前 Deng et al. (2024) 提出的基于 FBSDE 的反射薛定谔桥（必须在训练中运行高成本全轨迹采样并计算散度项），本文的 $\alpha$-RSBM 依靠一维高斯和解耦与解析反射桥采样，在 MNIST 上的实际壁钟时间仅从 31.5 小时增至 32.0 小时，在 AFHQ 上仅从 108 小时增至 109 小时，计算开销增幅不足 1.5%。
彻底根治裁剪失真：无约束模型在手写体任务中有高达 100% 的生成图片包含非法越界像素，在自然图像中有 15% 的样本超过 5% 的像素通道溢出边界；$\alpha$-RSBM 无需任何人后处理截断即可百分之百保证样本在 $[0, 1]^d$ 域内，避免了边缘过饱和伪影。

亮点与洞察¶

将反射边界与整流桥流优雅缝合：巧妙利用超立方体反射算子与无约束高斯过程的逆像测度映射，将原本极难处理的有界域动力学转化为对离散镜像点的期望计算，在保持流匹配式简易训练的同时严谨封闭了测度泄漏。
维度解耦破除了高维维数灾难：以往受限扩散模型在处理流形与紧致域时常陷入高维积分爆炸，本文指出在正交超立方体下转移密度具有坐标独立可乘性，使单步 Score 计算退化为一维标量前向，极大释放了其在实际图像生成中的工程可用性。
在理论严密性与工程吞吐之间取得极致平衡：既不同于只顾速度却牺牲边界约束的常规 Flow Matching，也不同于理论完备但计算开销高昂不可行的 FBSDE 方案，成功将反射薛定谔桥拉入工业级可承受算力范围。

局限与展望¶

参考过程无法带漂移项：目前理论与解析 Score 计算严格依赖零漂移的反射布朗运动（RBM），无法直接套用扩散模型中常用的方差保持（Variance-Preserving, VP）SDE，这限制了参考动力学的多样性设计。
几何域局限于正交超立方体：算法高度依赖立方体在各正交坐标轴上的可分独立性。若将其扩展至单纯形（Simplex）、球面或一般黎曼流形，坐标解耦性质将失效，必须依赖昂贵的神经网络桥采样器或 PDE 数值求解器。
未来方向：探索通过可微同胚（Diffeomorphism）将复杂凸几何体拉普拉斯变换至超立方体空间的可行性，以及融合自适应时间步长采样进一步加速生成推理。

评分¶

新颖性: ⭐⭐⭐⭐ [首次将无模拟桥匹配框架成功推广至反射布朗运动，闭式反射桥采样与 Score 解耦设计精巧]
实验充分度: ⭐⭐⭐⭐ [涵盖手写体与真实高维自然图像双向迁移，对越界指标、像素距离及壁钟开销分析透彻]
写作质量: ⭐⭐⭐⭐⭐ [数学推导严密规范，从测度论、随机微积分到算法落地的推演脉络清晰]
价值: ⭐⭐⭐⭐ [为受限紧致域高维概率生成与最优输运提供了高效率、即插即用的工业级算法范式]

数据集迁移任务	方法模型	MSD ↓	LPIPS ↓	FID ↓	训练总耗时 (GPU)
MNIST → EMNIST	\(\alpha\)-DSBM (无约束截断)	0.416	—	6.41	31.5 h (1× A100)
	\(\alpha\)-RSBM (本文)	0.356	—	5.22	32.0 h (1× A100)
EMNIST → MNIST	\(\alpha\)-DSBM (无约束截断)	0.415	—	8.58	31.5 h (1× A100)
	\(\alpha\)-RSBM (本文)	0.362	—	8.23	32.0 h (1× A100)
AFHQ: cat → wild	\(\alpha\)-DSBM (无约束截断)	0.131	0.239	27.6	108.0 h (4× A100)
	\(\alpha\)-RSBM (本文)	0.084	0.248	26.1	109.0 h (4× A100)
AFHQ: wild → cat	\(\alpha\)-DSBM (无约束截断)	0.200	0.302	31.5	108.0 h (4× A100)
	\(\alpha\)-RSBM (本文)	0.082	0.244	33.8	109.0 h (4× A100)