跳转至

BridgeTwist: Twisting Schrödinger Bridges for Training-Free Conditional Sampling

会议: NeurIPS 2026
arXiv: 2610.18658
代码: https://github.com/Mengyu8042/BridgeTwist
领域: 图像生成
关键词: 薛定谔桥, 免训练条件采样, 序贯蒙特卡洛, 扭曲采样, 扩散调度

一句话总结

针对预训练无条件薛定谔桥生成模型缺乏免训练后验采样机制的空白,提出结合速度场与得分场双通道终点预测器的退火扭曲 SMC 采样框架 BridgeTwist,实现渐进精确且兼顾多样性与保真度的即插即用条件生成。

研究背景与动机

生成模型服务于下游多样化的条件生成任务(如类别引导、文本引导及图像逆问题)是现代生成学习的核心需求。在免训练(training-free)范式下,研究者将无条件预训练模型作为先验,并结合任务特定的似然函数直接推断后验分布,避免了针对每种新观测重新训练特定网络的巨大开销。目前在基于得分的扩散模型(如 DPS、TDS、FPS)和流匹配模型(如 FMPS、TFTF)中,免训练条件采样已建立了梯度引导与序贯蒙特卡洛(SMC)两大技术体系;然而,薛定谔桥(Schrödinger Bridge, SB)作为能够在任意源分布与目标分布间构建熵正则化最优随机传输的新兴生成范式,其免训练条件采样方法仍处于空白状态,现有 SB 条件方法几乎全部依赖针对特定似然任务训练专用网络。

在探索免训练 SB 采样的过程中,核心难点在于前向路径中对终点状态的准确预测。扩散模型依赖基于得分场的 Tweedie 去噪公式预测纯净终点,流匹配模型依靠速度场做确定性欧拉外推;但对于布朗运动参考下的薛定谔桥,中间状态既包含源分布与目标分布间的插值,又叠合了布朗桥的随机高斯扰动。若仅采用得分场做 Tweedie 式去噪,推断出的仅仅是中间插值状态而非目标终点;若仅沿速度场做欧拉外推,则完全丢失了修正中间扰动的得分梯度。两种单一通道在薛定谔桥设定下均发生本质失真,无法为序贯采样提供有效的超前预测(lookahead)。

为了克服这一两难困境,本文发现概率流速度场与边际得分场在布朗桥框架下天然互补,两者的线性加权组合能够精确解析还原终点的条件期望。核心 idea:利用预训练速度场与得分场构建闭式终点预测器作为扭曲代理,引入时变退火指数抑制早期不确定性,并配合与桥状方差匹配的推断扩散调度,构建了渐进精确的薛定谔桥免训练条件 SMC 采样器 BridgeTwist。

方法详解

整体框架

BridgeTwist 输入为任意连续可微的观测似然函数 \(L_y(x) = p(y \mid x)\) 及预训练无条件薛定谔桥的速度场 \(v_\theta\) 与得分场 \(s_\theta\)。在时间离散网格上,算法通过粒子群模拟带引导项的正向 SDE;在每一步中,首先通过速度-得分互补公式预测终点并评估似然,作为当前步的扭曲函数(twist);随后通过时变退火权重削弱前序不准的引导,利用 \(h\)-变换构造引导提议核,最后借助可伸缩的递增重要性权重抵消所有中间扭曲因子,以渐进精确的方式收敛到真实条件后验。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["源分布粒子初始化<br/>X_0 ~ q_0, 权重归一化"] --> B["终点预测与扭曲构建<br/>速度场 + 得分场闭式外推 μ_1|t"]
    B --> C["时变退火与方差匹配<br/>退火指数 β(t) 动态温和放缩"]
    C --> D["引导提议核推进<br/>h-变换漂移校正 + 桥状扩散注入"]
    D --> E["权重更新与伸缩消去<br/>重要性权重伸缩抵消中间扭曲"]
    E -->|ESS 低于阈值 ρK| F["自适应分层重采样<br/>复制高权粒子并重置等权"]
    E -->|ESS 充裕| G["下一时间步迭代<br/>t_n -> t_{n+1}"]
    F --> G
    G -->|达到终点 t=1| H["终点加权后验粒子<br/>高保真条件生成样本"]

关键设计

1. 终点预测与扭曲构建:双通道互补解析预测终点条件期望

针对单一通道无法准确预测薛定谔桥终点的根本缺陷,本文从布朗参考桥的连续条件扩散性质出发进行数学推导。在给定端点 \((X_0, X_1)\) 条件下,中间时刻状态服从布朗桥 \(X_t \mid (X_0, X_1) \sim \mathcal{N}(M_t, \sigma^2 t(1-t) I)\),其中 \(M_t = (1-t)X_0 + tX_1\) 为线性插值。利用边际得分对卷积边缘密度的 Tweedie 展开,中间插值均值为 \(\mathbb{E}[M_t \mid X_t=x] = x + \sigma^2 t(1-t) s^\star_t(x)\),而端点位移 \(D = X_1 - X_0\) 的条件均值为 \(\mathbb{E}[D \mid X_t=x] = v^\star_t(x) + \frac{1-2t}{2}\sigma^2 s^\star_t(x)\)。通过几何关系 \(X_1 = M_t + (1-t)D\),二者相加精准消去了非对称项,得到解析的终点条件期望:

\[\mu^\star_{1 \mid t}(x) := \mathbb{E}[X_1 \mid X_t=x] = x + (1-t)v^\star_t(x) + \frac{\sigma^2}{2}(1-t)s^\star_t(x)\]

这一公式表明,真正的终点预测器必须由速度通道(承载沿桥的确定性传输外推)与得分通道(纠正当前时刻的高斯噪声涨落)协同提供。在算法实现中,利用预训练神经网络预测场将其代入观测似然,构建即插即用的闭式扭曲代理 \(\hat{h}_t(x) = L_y(\hat{\mu}_{1 \mid t}(x))\),避免了昂贵的蒙特卡洛多步前向 rollout 模拟。

2. 时变退火与方差匹配:抑制早期预测过信度并保持终点渐进无偏

尽管终点一阶矩得以解析估计,但直接以单点代入似然 \(\hat{h}_t(x)\) 会完全忽略终点分布在早期的方差扩散,导致在靠近源分布 \(t \to 0\) 区域产生严重的虚假过信度(over-confidence)。为消除这一失真,BridgeTwist 借鉴退火重要性采样思路,引入时变退火指数 \(\beta(t) = t^a\)(\(a \ge 1\),默认 \(a=1\)),满足 \(\beta(0)=0\) 且 \(\beta(1)=1\)。在各中间时间步定义退火中间目标路径测度 \(\gamma_n^\beta(\mathrm{d}x_{t_{0:n}}) = P_{\theta, g}^{(n)}(\mathrm{d}x_{t_{0:n}}) \hat{h}_n(x_{t_n})^{\beta_n}\)。

理论分析证明,在线性高斯观测模型下,精确匹配预测方差的最优退火系数严格等于 \(\beta^\star(t) = \sigma_y^2 / (\sigma_y^2 + \sigma_{1 \mid t}^2)\),该系数随时间自 \(t=0\) 单调递增至 \(t=1\) 处的 \(1\)。由于最终目标在 \(n=N\) 时刻天然满足 \(\beta_N=1\) 且 \(\hat{h}_N(x) = L_y(x)\),中间退火因子在连续权重的累乘中由于伸缩结构(telescoping cancellation)被完全对消,因此改变退火策略仅平滑调整了粒子轨迹的方差与引导强度,绝不引入最终稳态后验的渐近偏差。

3. 引导提议核推进:桥状扩散调度与 Doob h-变换联合漂移

为引导粒子朝高似然区域移动,提议分布 \(q_n\) 基于离散化连续时间 Doob \(h\)-变换,将带有退火强度的扭曲似然梯度注入转移漂移中:

\[m_{q, n}(x_{t_{n-1}}) = x_{t_{n-1}} + \left[ b_{\theta, g}(t_{n-1}, x_{t_{n-1}}) + w_g g(t_{n-1})^2 \beta_{n-1} \nabla_x \log \hat{h}_{n-1}(x_{t_{n-1}}) \right] \Delta_n\]

在此基础上,推断阶段扩散系数 \(g(t)\) 具备保边缘不变性的自由度。传统常数扩散 \(g(t) \equiv \sigma\) 在边界附近注入了过量噪声,极易扰乱已与观测对齐的粒子。BridgeTwist 提出与布朗桥标准差同构的桥状扩散调度 \(g(t) = c_g \sigma \sqrt{t(1-t) + \delta_g}\)。该调度将提议噪声集中在路径中段,而在起点(先验已提供充分多样性)与终点(避免破坏细粒度细节)压低噪声。二者在提议漂移中形成的有效引导系数 \(g(t)^2 \beta(t) \propto t^{a+1}(1-t)\) 在默认 \(a=1\) 时于 \(t = 2/3\) 处达到峰值,实现了前期稳定探索、中期强力引导、后期平稳收敛的动力学特性。

4. 权重更新与伸缩消去:重要性权重全链消解与分层重采样保障收敛性

通过将前向离散先验核 \(M_n\) 与提议核 \(q_n\) 的概率密度比,乘上相继两个时间步的退火扭曲函数比值,构造粒子增量权重:

\[\omega_n(x_{t_{n-1}}, x_{t_n}) = \frac{M_n(x_{t_{n-1}}, x_{t_n}) \hat{h}_n(x_{t_n})^{\beta_n}}{q_n(x_{t_n} \mid x_{t_{n-1}}, y) \hat{h}_{n-1}(x_{t_{n-1}})^{\beta_{n-1}}}\]

累乘展开时,分子分母的中间扭曲项 \(\hat{h}_n^{\beta_n}\) 形成完美的可伸缩消除链(telescoping identity),最终权重的未归一化形式恰好对应真实后验路径测度。当有效粒子数 \(\mathrm{ESS}_{n-1} = 1 / \sum_k (W_{n-1}^{(k)})^2\) 低于设定阈值 \(\rho K\) 时,触发分层重采样复制优势粒子,并重置等权。严格的理论分析证明,在 \(K\) 个粒子下,对于任意有界测试函数 \(\varphi\),估计误差在 \(L_2\) 范数下以 \(\mathcal{O}(K^{-1/2})\) 速率严格收敛至离散化先验导出的真后验期望。

实验关键数据

主实验

在类别条件生成(MNIST 与 CIFAR-10,使用独立外部判别器 ACCE 评估条件遵从性,KID 评估生成真实度)以及图像逆问题(CelebA-HQ \(4\times\) 超分辨率、高斯去模糊,CIFAR-10 图像补全)上进行了全面评估。

数据集 / 任务 评估指标 本文 (BridgeTwist) 领先基线 (TFTF / PnP-Flow / DPS-SB) 相对增益 / 优势
CIFAR-10 (类别条件, \(K=64\)) ACCE ↑ / KID ↓ 0.979 / 0.0403 0.934 (TFTF) / 0.0572 (TFTF) 外部准确率 +4.5%, KID 显著降低 29.5%
MNIST (类别条件, \(K=64\)) ACCE ↑ / KID ↓ 0.952 / 0.0003 0.942 (TDS) / 0.0010 (DPS-SB) 外部准确率最高,保真度指标达到最优
CelebA-HQ (高斯去模糊) PSNR ↑ / SSIM ↑ 31.13 dB / 0.847 30.91 dB (Flower) / 0.839 (Flower) PSNR 领先 0.22 dB,结构相似度最优
CelebA-HQ (\(4\times\) 超分辨率) PSNR ↑ / SSIM ↑ 28.33 dB / 0.836 28.47 dB / 0.843 (PnP-Flow) 仅次于专用逆问题解算器,差距微弱
CIFAR-10 (Half-mask 补全) SSIM ↑ 0.655 0.623 (PnP-Flow) 结构恢复质量领先 0.032
CIFAR-10 (Random 50% 补全) SSIM ↑ 0.943 0.940 (PnP-Flow) 随机遮挡下保真度位居第一
CelebA-HQ (文本提示生成, \(K=25\)) ACCE ↑ / LPIPS 多样性 ↑ 0.827 / 0.613 0.803 / 0.501 (TFTF) 属性命中率提升 2.4%,样本多样性大幅提高 22.4%

消融实验

在相同无条件薛定谔桥先验下(CIFAR-10, \(K=16\) 及 \(K=64\)),系统检验了终点预测器通道构成、算法核心组件、以及扩散/退火超参配置。

配置 ACCL ↑ ACCE ↑ KID ↓ 说明
仅速度通道 (Velocity-only Euler) 0.526 0.485 0.2805 缺少得分纠偏,条件一致性与保真度全面崩溃
仅得分通道 (Score-only Tweedie) 0.711 0.514 0.2339 预测退化为中间插值,终点指引严重失真
BridgeTwist 完整预测器 (公式7) 0.948 0.914 0.0873 速度与得分协同,准确率与生成质量达到最优

针对算法核心机制的控制消融(CIFAR-10, \(K=64\)):

变体配置 \(h\)-变换引导 退火加权 权重递推修正 ACCE ↑ KID ↓
终端重要性采样 (IS) \(\times\) \(\times\) 仅终点加权 0.865 0.0744
无权重修正引导 (DPS-SB) \(\checkmark\) \(\times\) \(\times\) 0.585 0.1051
无 \(h\)-变换漂移 (w/o \(h\)-transform) \(\times\) \(\checkmark\) \(\checkmark\) 0.925 0.0437
无时变退火 (w/o tempering, \(a=0\)) \(\checkmark\) \(\times\) \(\checkmark\) 0.938 0.0813
完整 BridgeTwist (\(a=1\), 桥状调度) \(\checkmark\) \(\checkmark\) \(\checkmark\) 0.979 0.0403

关键发现

  • 终点预测通道的协同性至关重要:无论是仅保留速度欧拉外推还是仅保留得分 Tweedie 去噪,ACCE 均跌破 52%,证实了布朗桥中间态的非对称性必须由两通道联合补偿。
  • 引导与重要性校正缺一不可:纯引导无权重修正的 DPS-SB 由于累积不可消除的近似偏差,增加粒子数无法收敛;而纯终端重加权 (IS) 在复杂条件下由于权重退化严重,采样效率极低。
  • 桥状扩散调度展现出压倒性优势:在任意退火指数 \(a \in \{0, 1, 2, 3\}\) 下,桥状调度均显著优于传统常数调度与余弦调度;将线性退火(\(a=1\))与桥状调度组合时,有效引导在 \(t=2/3\) 达到峰值,取得了最佳的平衡。

亮点与洞察

  • 薛定谔桥终点解析解的严谨揭示:利用连续布朗桥扩散与边际密度卷积微分性质,首次在封闭形式下推导出 \(\mathbb{E}[X_1 \mid X_t]\) 必须由速度场与半方差加权得分场共同组成的数学恒等式,厘清了生成扩散与确定性流在随机最优传输中的本质交汇。
  • 中间扭曲的无偏消去机制:巧妙利用序贯蒙特卡洛中间比率的伸缩结构(telescoping cancellation),使得无论中间采用了何种退火曲线和经验近似扭曲,中间项均在全链累乘中严格相消,在理论上完整保留了渐进无偏性。
  • 两级噪声协同调度范式:将推断阶段的布朗噪声方差 \(g(t)\) 塑形为内凹桥状,与外凸退火因子 \(\beta(t)\) 发生乘积共振,自然构成了前段保护多样性、中段强力纠正轨迹、后段锁定细节的免调参动力学。

局限与展望

  • 参考测度的形式限制:当前理论推导与闭式终点预测器严格限定于布朗运动参考过程(Brownian reference),对于 Ornstein-Uhlenbeck (OU) 过程或带有流形几何先验的一般参考测度,尚需推导特定的非各向同性终点展开式。
  • 高阶曲率的未建模近似:扭曲代理本质上采用了一阶条件期望代入似然 \(\hat{h}_t(x) = L_y(\hat{\mu}_{1 \mid t}(x))\),对于高度非线性、多峰或几何狭长曲折的极端似然,忽略二阶协方差仍可能造成局部的粒子失配。
  • 预训练双头网络的开销:先验模型必须同时具备准确的速度场和得分场,相比于传统仅有单一速度或得分网络的基础模型,需要额外的训练通道(如在预训练阶段增加轻量得分头)。

相关工作与启发

  • vs DPS (Diffusion Posterior Sampling) / FMPS (Flow Matching Posterior Sampling):DPS 与 FMPS 分别针对扩散与流模型引入单步启发式引导,但均缺乏序贯粒子权重的动态校正,存在无法消除的固有近似偏差(asymptotic bias)。BridgeTwist 基于扭曲 SMC 框架,误差随粒子数按 \(\mathcal{O}(K^{-1/2})\) 严格衰减。
  • vs TDS (Twisted Diffusion Sampler) & TFTF (Training-Free Targeted Flow):TDS 仅使用 Tweedie 得分做扩散后验超前预测,TFTF 仅使用速度场欧拉步做流后验预测。BridgeTwist 证明了二者分别是薛定谔桥退化到单侧边界时的特例,确立了统一的双通道终点估计理论。
  • vs I²SB 等任务特定桥模型:I²SB 需要针对特定退化(如超分辨率倍率、固定模糊核)进行端到端的条件对训练,泛化代价高昂。BridgeTwist 无需任何下游微调,即插即用支持任意可微判别器与逆问题算子。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 首次打通薛定谔桥免训练条件采样路径,揭示速度与得分双通道终点预测恒等式。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖低维解析混合高斯验证、MNIST/CIFAR 类别生成、CelebA 文本生成及三大图像逆问题与多项细粒度消融。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论框架清晰,数学推导严密自洽,动力学直觉与实验设计环环相扣。
  • 价值: ⭐⭐⭐⭐⭐ 为最优传输与随机桥式生成模型在各种下游复杂受限控制任务中的落地提供了关键采样基础设施。