跳转至

Stochastic Interpolants via Conditional Dependent Coupling

会议: NeurIPS 2026
arXiv: 2509.23122
领域: 图像生成
关键词: 随机插值, 条件相关耦合, 多阶段流匹配, 扩散Transformer, 像素空间生成

一句话总结

本文提出基于条件相关耦合(Conditional Dependent Coupling, CDC)的多阶段流匹配框架 CDC-FM,以单一统一的 Diffusion Transformer 在像素空间实现从粗到细的高保真生成与端到端优化,显著降低传输代价与推理时间并在 ImageNet-256 上取得 1.81 FID。

研究背景与动机

高维连续图像生成长期在计算效率与保真度之间面临艰难权衡。当前主流的隐空间生成范式依赖预训练变分自编码器(VAE)将图像压缩至低维隐空间,虽然大幅削减了 Diffusion 或 Flow 模型的计算开销,但不可逆的有损压缩必然导致精细纹理丢失,且两阶段割裂的训练流程阻碍了生成主干网络与编码器的联合端到端优化。相反,直接在原始像素空间建模虽然保留了全部高频细节,但面临维度灾难,导致训练极慢且采样推理成本高昂。

为缓解这一矛盾,级联/多阶段生成方法被提出,试图将合成过程分解为从低分辨率到高分辨率的粗到细演进轨迹。然而,现有多阶段架构大多依赖离散割裂的分辨率专属子模型,不仅无法实现参数共享与端到端协同优化,更严重的是,以往方法在中间阶段往往维持高噪声目标,导致中间层实际拟合的是结构模糊、缺乏语义的近高斯分布,使条件分支间的 KL 散度极小,进而引发各阶段分布学习不准与累积误差失真。

本文的切入角度是:若能在像素级多阶段轨迹中,使每个阶段的源分布直接与目标高分辨率真实图像建立结构化依赖,同时让各阶段的目标分布严格对应各分辨率下的真实降采样图像,即可在保证准确分布拟合的同时大幅缩短概率流路径。本文的核心 idea 是:提出条件相关耦合(Conditional Dependent Coupling)策略,将多阶段随机插值分解为马尔可夫链式轨迹,利用单套 DiT 参数与多分辨率嵌入实现统一像素级端到端流匹配,在显著降低最优传输代价的同时保证各阶段准确的分布学习。

方法详解

整体框架

CDC-FM 将全分辨率图像生成分解为 \(K\) 个粗到细阶段(如 \(K=4\) 对应 \(64 \to 128 \to 256\) 渐进精化)。阶段 \(k=1\) 负责从标准高斯先验中生成基础低分辨率图像;后续阶段 \(k > 1\) 则通过条件相关耦合(CDC),将前一阶段输出经确定性上采样并注入极低阶阶梯衰减噪声作为当前阶段起点,向更高分辨率目标演化。整个生成过程参数化为单一共享权重的 Diffusion Transformer(DiT),通过分辨率嵌入与时间重缩放协同驱动。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["初始高斯噪声<br/>x_0^(1) ~ N(0, σ² I)"] --> B["阶段 1:低分辨率生成<br/>噪声到图像插值求解"]
    B --> C["阶段输出 x̂_1^(1)<br/>低分辨率基础表征 (d_1)"]
    C --> D["条件相关耦合<br/>U_2 上采样 + 衰减噪声 σ_2 ζ_2"]
    D --> E["阶段 2:中低分辨率精化<br/>图像到图像流匹配求解"]
    E --> F["条件相关耦合迭代<br/>U_k 上采样 + 递减噪声 σ_k ζ_k"]
    F --> G["阶段 K:全分辨率高频精化<br/>极小传输代价下生成细节"]
    G --> H["最终高保真图像 x̂_1^(K)"]

关键设计

1. 条件相关耦合与马尔可夫链式分解:消除独立噪声的高额传输代价

传统扩散模型与流匹配在各阶段独立抽样高斯噪声,导致从高斯先验到高维全分辨率数据间的动力学传输路径极长、传输能耗居高不下。CDC 重新设计了各阶段的联合分布耦合。在给定最终高分辨率图像 \(x^{(K)}\) 的条件下,各阶段的数据对偶满足条件独立性:\(\rho^{(k)}(x_0^{(k)}, x_1^{(k)} \mid x^{(K)}) \perp\!\!\!\perp \rho^{(k^*)}(x_0^{(k^*)}, x_1^{(k^*)} \mid x^{(K)})\)。在精化阶段 \(k > 1\),当前阶段的源分布样本 \(x_0^{(k)}\) 直接由对应目标样本构建: $\(x_0^{(k)} = m_k(x_1^{(k)}) + \sigma_k \zeta_k = U_k(D_k(x_1^{(k)})) + \sigma_k \zeta_k\)$ 其中 \(D_k\) 为邻域均值降采样,\(U_k\) 为最近邻复制上采样,\(\zeta_k \sim \mathcal{N}(0, I_{d_k})\) 为注入扰动。由此推导出的推断联合分布构成严格的马尔可夫链:\(\rho(\hat{x}_0^{(1)}, \dots, \hat{x}_1^{(K)}) = \rho^{(1)}(\hat{x}_0^{(1)}, \hat{x}_1^{(1)}) \prod_{k=2}^K \rho^{(k)}(\hat{x}_0^{(k)}, \hat{x}_1^{(k)} \mid \hat{x}_1^{(k-1)})\)。理论证明(定理 2),将图像信号能量按多分辨率正交投影分解后,传统单阶段直推模型的传输代价 \(L_A = \sum_{k=1}^K \mathbb{E}[|\tilde{\epsilon}^{(k)}|^2] + d_K \sigma^2\),而 CDC 多阶段总代价 \(L_B = \sum_{k=1}^K \mathbb{E}[|\tilde{\epsilon}^{(k)}|^2] + d_1 \sigma^2 + \sum_{k=2}^K d_k \sigma_k^2\)。高维纯噪声位移项被有效消除,从而实现了从数学上严格成立的 \(L_A > L_B\)。

2. 随机性渐消策略:按尺度解耦多样性与精化高频

在连续生成模型中,图像语义、物体轮廓与整体布局等多尺度多样性几乎全部在演化的早期低分辨率阶段决定,而后续高分辨率阶段主要任务是补充高频边缘与微观纹理,不再需要大幅度的随机探索。基于此观察,CDC 引入了噪声衰减因子(diminish factor)\(\gamma \ge 1\): $\(\sigma_k = \gamma^{-(k-1)} \sigma, \quad 2 \le k \le K\)$ 默认取 \(\gamma = 2.0\)。随着阶段 \(k\) 递增,空间维度按 \(d_k = 2^{k-1} d_1\) 翻倍增长,但注入的高斯扰动方差呈几何级数衰减,使得噪声项在各阶段累加的计算负荷维持为常数量级 \(K d_1^2 \sigma^2\),远小于单阶段模型的 \(2^{2(K-1)} d_1^2 \sigma^2\)(定理 3)。这种渐消机制不仅稳定了 ODE 轨迹并使其更为平直,还显著减少了数值求解器所需的函数评估次数(NFE)。

3. 单一 DiT 统一参数化与分辨率感知嵌入:促进跨尺度知识共享

不同于传统级联扩散模型为每个分辨率分别训练一套独立 U-Net 或 DiT,CDC-FM 将全部 \(K\) 个阶段的流匹配向量场完全统一到同一个 Diffusion Transformer \(b^\theta\) 中。为使主干网络能够敏锐区分当前处理的具体特征尺度,模型提取 Patch 嵌入后的绝对特征图分辨率 \(r_k\),通过正弦位置编码映射为分辨率嵌入 \(e_k = E(r_k)\),并在输入层通过交叉注意力机制与时间步嵌入完成深度融合。同时,全局连续时间 \(t \in [0, 1]\) 按照阶段切分区间 \([t_0^k, t_1^k]\),各阶段局部时间被统一重缩放为 \(\tau = \frac{t - t_0^k}{t_1^k - t_0^k} \in [0, 1]\),从而保证共享网络在各尺度下的动力学尺度一致性与参数梯度联合优化。

4. 全阶段统一无分类器引导:避免繁琐的阶段式引导调度

在以往多阶段模型中,由于中间阶段目标分布模糊、条件信息表达不足,不得不针对每个分辨率人工精心调整专属的 CFG 引导强度。而在 CDC 框架下,每个阶段 \(k\) 的演化终点都被严格锚定为真实降采样图像分布 \(\rho_1^{(k)}(x_1^{(k)} \mid x^{(K)})\),因而模型能够精准拟合条件分布与无条件分布的差异。因此,推断时无需任何阶段性衰减或增强调度,可以直接使用统一的全局常数引导强度 \(S_{\text{cfg}}\)(如 3.0): $\(b_{\text{cfg}}^\theta(z, \tau, e_k; \mathbf{c}, S_{\text{cfg}}) = (1 - S_{\text{cfg}}) b^\theta(z, \tau, e_k, \varnothing) + S_{\text{cfg}} b^\theta(z, \tau, e_k, \mathbf{c})\)$ 在保持高保真度的同时大幅简化了调参复杂度。

一个完整示例

以 ImageNet 类别条件生成 \(256 \times 256\) 图像(\(K=4\) 阶段,\(d_k \in \{64^2, 96^2, 128^2, 256^2\}\) 或级联倍频)为例: 1. 阶段 1(\(k=1, 64 \times 64\)):采样 \(x_0^{(1)} \sim \mathcal{N}(0, \sigma^2 I)\),在局部时间 \(\tau \in [0, 1]\) 内以步长 \(N_1=6\) 调用 DiT 计算引导速度场 \(v_{\text{cfg}}\) 并通过前向欧拉法推进,得到粗粒度语义轮廓 \(\hat{x}_1^{(1)}\)。 2. 阶段 2(\(k=2, 128 \times 128\)):将 \(\hat{x}_1^{(1)}\) 经双线性/最近邻复制上采样并叠加方差为 \(\sigma_2 = 0.5 \sigma\) 的噪声,作为 \(\hat{x}_0^{(2)}\);以步长 \(N_2=6\) 求解 ODE,恢复中频物体部件结构 \(\hat{x}_1^{(2)}\)。 3. 阶段 3(\(k=3, 256 \times 256\)):再次上采样并注入微弱噪声 \(\sigma_3 = 0.25 \sigma\),在此平直轨迹上以 \(N_3=6\) 步精确合成微观纹理与高频边缘,最终输出 \(256 \times 256\) 全保真图像 \(\hat{x}_1^{(3)}\)。整个流程仅需 \(3 \times 6 \times 2 = 36\) NFE,远低于单阶段模型 250 至 500 次采样的开销。

损失函数 / 训练策略

在训练阶段,每个批次样本随机抽取全局时间步 \(t \sim \mathcal{U}[0, 1]\),定位其所在阶段 \(k = \operatorname*{arg\,k}\{t \in [t_0^k, t_1^k]\}\) 并计算重缩放局部时间 \(\tau \in [0, 1]\)。真实图像 \(x^{(K)}\) 经降采样得到目标 \(x_1^{(k)} = D_{K \to k}(x^{(K)})\)。当 \(k=1\) 时独立采样高斯噪声 \(x_0^{(1)}\);当 \(k > 1\) 时通过 CDC 构造源分布样本 \(x_0^{(k)} = U_k(D_k(x_1^{(k)})) + \sigma_k \zeta_k\)。构建线性插值轨迹及其目标速度: $\(I_\tau^{(k)} = (1 - \tau) x_0^{(k)} + \tau x_1^{(k)}, \quad \dot{I}_\tau^{(k)} = x_1^{(k)} - x_0^{(k)}\)$ 联合条件丢弃机制(概率 \(p_\varnothing\) 替换条件为 \(\varnothing\))最小化均方流匹配误差: $\(\mathcal{L}_{\text{cfg}}(\theta) = \mathbb{E}\left[ \| b^\theta(I_\tau^{(k)}, \tau, e_k, \bar{\mathbf{c}}) - \dot{I}_\tau^{(k)} \|_2^2 \right]\)$ 网络采用 28 层 DiT 架构(隐藏维度 1152,16 注意力头,参数量 677M),Patch 尺寸固定为 4,使用 AdamW 优化器(学习率 \(1 \times 10^{-4}\),\(\beta_1=0.9, \beta_2=0.95\))和 bfloat16 混合精度进行端到端优化。

实验关键数据

主实验

在 ImageNet-1k \(256 \times 256\) 类别条件生成基准上,CDC-FM 与代表性 GAN、扩散/流匹配、自回归模型及多阶段模型的对比如下:

模型 / 方法 参数量 FID ↓ IS ↑ Precision ↑ Recall ↑ NFE 单图耗时 (s)
BigGAN (GAN) 112M 6.95 224.5 0.89 0.38 1 1.46
StyleGAN (GAN) 166M 2.30 265.1 0.78 0.53 1 0.96
ADM (Diffusion) 554M 10.94 101.0 0.69 0.63 250×2 9.46
LDM-4-G (Latent Diff) 400M 3.60 247.7 - - 250×2 4.18
DiT-XL/2 (Latent Diff) 675M 2.27 278.2 0.83 0.57 250×2 3.66
SiT-XL/2 (Interpolant) 675M 2.06 270.3 0.82 0.59 250×2 3.53
VDM++ (Diffusion) 2.0B 2.12 267.7 - - - -
SiD2 (Pixel Diff) 2.0B 1.72 - - - - -
MaskGIT (Masked AR) 227M 6.18 182.1 0.80 0.51 8 0.97
MAR-H (Masked AR) 943M 1.55 303.7 0.81 0.62 256×2 1.23
VAR-d30 (Scale AR) 2.0B 1.92 323.1 0.82 0.59 10×2 1.12
PixelFlow (Multi-stage Flow) 677M 1.98 282.1 0.81 0.60 - -
PaGoDA (Multi-stage Diff) 0.9B 1.56 259.6 - 0.59 - -
CDC-FM (Ours, 16步) 677M 1.97 298.2 0.80 0.58 2×4×2 1.31
CDC-FM (Ours, 32步) 677M 1.87 301.8 0.82 0.60 4×4×2 2.66
CDC-FM (Ours, 48步) 677M 1.81 304.3 0.83 0.61 6×4×2 5.18

在高分辨率基准 ImageNet-1k \(512 \times 512\) 下,CDC-FM 依然展现出强大竞争力:

模型 / 方法 参数量 FID ↓ IS ↑
BigGAN - 8.43 177.9
ADM 554M 7.72 172.7
DiT-XL/2 675M 3.04 240.8
VDM++ 2.0B 2.65 278.1
U-ViT 501M 4.05 263.8
SiD2 2.0B 2.19 -
MaskGIT 227M 7.32 156.0
MAGVIT-v2 307M 3.07 324.3
MAR-L 481M 2.74 205.2
VAR-d36-s 2.0B 2.63 303.2
PaGoDA 0.9B 1.80 251.3
CDC-FM (Ours) 684M 2.65 307.4

消融实验与分析

针对核心超参数与设计选择的系统分析:

评估维度 参数变动范围 最优配置 指标表现 关键规律说明
衰减因子 \(\gamma\) \(\gamma \in [1.0, 5.0]\)(步长 0.5) \(\gamma = 2.0\) 最优 FID 达到波谷 随 \(\gamma\) 增大,高分辨率阶段轨迹变直,推断时间持续单调下降并收敛;但 \(\gamma > 2.5\) 后扰动过小导致多样性受限,FID 反弹。
CFG 强度 \(S_{\text{cfg}}\) \(S_{\text{cfg}} \in [1.0, 6.0]\) \(S_{\text{cfg}} = 3.0\) FID 最优,IS 保持高位 \(S_{\text{cfg}}\) 从 1.0 升至 3.0 时,样本质量与精确率明显提升;超过 4.0 后出现过度饱和与模式坍缩。
阶段步数分配 \(N_k \in \{2, 4, 6\}\) \(N_k = 6\) FID 1.81, IS 304.3 即使在极度轻量的 2 步/阶段(总计 16 NFE)下,CDC-FM 仍取得 1.97 FID,证明子轨迹极其平直。

关键发现

  • 传输平直性使少步生成具备实用性:在仅需 16 次函数评估(NFE=16,耗时 1.31 秒)的极限设置下,CDC-FM 在像素空间达到了 1.97 FID,超越了使用 500 NFE 的 DiT-XL/2(2.27 FID)与 PixelFlow(1.98 FID)。
  • 统一模型打破了多阶段容量壁垒:通过分辨率位置编码共享 677M DiT 参数,各阶段能够无缝复用浅层感知与高层语义特征,完全消除了以往级联模型中重复部署多个 U-Net 所带来的显存浪费。
  • 推断耗时与分辨率呈严格线性关系:实测在 \(64 \to 128 \to 256 \to 512\) 分辨率下的耗时曲线完美拟合了定理 3 所预测的理论线性复杂度,破解了全像素模型以往在超高分辨率下二次方膨胀的计算陷阱。

亮点与洞察

  • 多阶段高斯方差衰减律:将多尺度视觉生成的结构规律与物理扩散扰动显式绑定,利用 \(\sigma_k = \gamma^{-(k-1)} \sigma\) 在数学上大幅压缩了高维空间内的无效传输距离。
  • 阶段性真实锚定破解多阶段分布退化:通过确保每个阶段都以实际降采样的真值图像分布作为目标,一举解决了传统多阶段模型在低分辨率阶段目标分布模糊、无法学习有效类别区分度的致命缺陷。
  • 统一 DiT 无缝支持全阶段 CFG:无需任何启发式的分段引导衰减超参,以单一全局标量 \(S_{\text{cfg}}\) 实现全局高保真度控制,极大地提高了工程落地的健壮性。

局限与展望

  • 作者承认的局限:目前骨干网络主要在标准 ImageNet 类别条件生成上展开验证,尚未在海量开放文本描述条件(Text-to-Image)数据集及兆级超大参数模型上完成极致规模扩展。
  • 自行挖掘的局限:降采样算子目前采用固定的简单邻域均值池化,针对复杂拓扑结构(如细密毛发、复杂几何网格)可能产生局部走样;若能在早期阶段引入自适应滤波核,有望进一步改善高频重建一致性。
  • 未来改进思路:可将条件相关耦合与潜一致性模型(Consistency Flow)进一步结合,探索在全像素空间下实现 1~2 步超实时视频生成与高分辨率 3D 纹理合成。

相关工作与启发

  • vs PixelFlow:PixelFlow 同样探索像素级级联流匹配,但在中间阶段保留了过多的扩散噪声目标,导致语义分布区分度不足;CDC-FM 通过真值降采样锚定与 CDC 耦合,在相同 677M 参数下将 FID 从 1.98 推进至 1.81。
  • vs DiT / SiT:DiT/SiT 依赖 VAE 隐空间进行单阶段漫长传输;CDC-FM 直接在像素空间操作,彻底摆脱了 VAE 解码模糊的瓶颈,同时凭借粗到细轨迹分解获得了比隐空间扩散模型更快的采样速度。
  • vs VAR / MAR:自回归模型通过离散或连续多尺度词元级预测生成,推理受限于自回归长历史依赖;CDC-FM 采用基于 ODE 的马尔可夫链式连续动力学传输,并行度更高且生成轨迹更加平滑可控。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 巧妙利用条件相关耦合建立严格马尔可夫链式多阶段流匹配,理论与形式十分优美。
  • 实验充分度: ⭐⭐⭐⭐⭐ 在 ImageNet 256 与 512 上均给出了极详实的主指标、计算耗时及参数消融。
  • 写作质量: ⭐⭐⭐⭐⭐ 数学证明严谨,从能量正交分解到传输代价不等式推导一气呵成。
  • 价值: ⭐⭐⭐⭐⭐ 为摆脱 VAE 依赖的纯像素高分辨率端到端生成树立了兼具理论深度与实用效率的全新范式。