跳转至

Temporal Alignment Guidance: On-manifold Sampling in Diffusion Models

会议: NeurIPS 2026
arXiv: 2510.11057
领域: 图像生成
关键词: 扩散模型, 引导生成, 流形偏离, 时间对齐引导, 训练无关采样

一句话总结

针对扩散模型在外部引导与加速采样时易偏离数据流形导致保真度崩塌的问题,本文提出时间对齐引导(TAG),利用轻量时间预测器估计时间关联得分,在反向扩散每一步将偏离样本主动拉回对应时刻的数据流形。

研究背景与动机

扩散模型在图像、音频、分子与视频生成等多样化模态中展现出卓越的生成能力,其核心优势之一在于能够在反向去噪采样过程中灵活注入外部条件引导,如无分类器引导(CFG)以及基于 Tweedie 公式的免训练引导(Training-Free Guidance, TFG / DPS)。随着应用拓展至反问题求解、黑盒优化、个性化生成与多目标约束控制,采样流程中频繁叠加任意形式的外部能量梯度或属性预测损失。然而,这种人为调整反向漂移项的操作极易将去噪轨迹推离原始数据分布,迫使样本落入低概率密度区域。

在低密度区域中,扩散模型原本学习的得分函数 \(\nabla_{\mathbf{x}}\log p_t(\mathbf{x})\) 变得极不可靠,导致严重的得分近似误差。更致命的是,这类误差在离散时间步的反向迭代中不断复合累积,引发严重的“偏离流形”(off-manifold)现象,最终导致生成样本出现结构畸变、伪影甚至语义崩塌。尽管现有方法尝试通过输入扰动、得分范数缩放、硬性时间重分配(如 TCS)或时间扰动对比(如 TSG/SG)来缓解暴露偏差,但它们大多依赖启发式规则,或者在遭遇强对抗性外部引导时自身也会失效,无法在理论与实际中提供统一的流形回正机制。

面对采样轨迹偏离流形的核心矛盾,本文重新审视了时间步变量的角色:反向过程中时间不应仅被视为静态固定的索引,而应被看作带有后验不确定性的条件变量。核心 idea:利用轻量时间预测器构建时间关联得分(Time-Linked Score, TLS),在每个反向扩散步引入时间对齐引导(TAG),排斥非目标时刻流形并主动将样本锚定回当前时刻的真实数据流形。

方法详解

整体框架

时间对齐引导(TAG)的核心理念是在每个反向扩散时间步 \(t\),不仅依靠基础扩散模型预测去噪方向,还借助一个轻量级时间预测网络 \(\boldsymbol{\phi}\) 评估当前中间样本 \(\mathbf{x}_t\) 属于各个时间步的后验分布 \(p_\phi(t \mid \mathbf{x}_t)\)。当外部引导(如反问题测量约束或属性优化损失)将样本推入低密度非目标区域时,TAG 能够感知该偏离并通过计算对数时间后验关于样本的梯度,施加一个正交/矫正拉力将样本拉回目标时刻流形 \(\mathcal{M}_t\)。

整个流水线在每个反向采样步串行执行两阶段校正:首先计算时间关联得分并更新中间样本得到流形对齐状态 \(\tilde{\mathbf{x}}_t\),随后再基于基础扩散模型的得分估计与外部任务引导项完成向 \(t-1\) 步的反向转移。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["初始高斯噪声 x_T ~ N(0, I)"] --> B["时间关联得分估计<br/>轻量时间预测器计算对数后验梯度"]
    B --> C["时间对齐引导矫正<br/>样本拉回当前时刻流形 M_t"]
    C --> D["反向去噪步更新<br/>基础扩散得分 + 任务外部引导"]
    D -->|未到 t=0| B
    D -->|到达 t=0| E["高质量流形内生成样本 x_0"]

关键设计

1. 时间关联得分(TLS)与目标流形吸引机制 传统扩散采样将当前时间步 \(t\) 硬性输入得分网络,一旦样本因外部漂移偏离流形,模型便无法感知其在时间维度上的真实状态。为了量化样本相对目标时刻流形 \(\mathcal{M}_t\) 的偏移,本文提出了时间关联得分(Time-Linked Score, TLS),定义为时间变量关于样本的对数后验梯度: $\(\mathrm{TLS}(\mathbf{x}, t) := \nabla_{\mathbf{x}} \log p(t \mid \mathbf{x})\)$ 通过贝叶斯定理分解,该项在离散时间步设定下展现出极其优雅的物理意义: $\(\nabla_{\mathbf{x}} \log p(t_i \mid \mathbf{x}) = \sum_{k \neq i} \frac{p_{t_k}(\mathbf{x})}{p_{\mathrm{tot}}(\mathbf{x})} \left( \nabla_{\mathbf{x}} \log p_{t_i}(\mathbf{x}) - \nabla_{\mathbf{x}} \log p_{t_k}(\mathbf{x}) \right)\)$ 其中 \(p_{\mathrm{tot}}(\mathbf{x}) = \sum_j p_{t_j}(\mathbf{x})\)。该公式表明,当样本在当前目标时刻 \(t_i\) 处于极低密度区域(即严重偏离 \(t_i\) 流形时,权重 \(p_{t_k}(\mathbf{x}) / p_{\mathrm{tot}}(\mathbf{x})\) 显著增大),TLS 会自动形成两重力场:正向拉力 \(\nabla_{\mathbf{x}} \log p_{t_i}(\mathbf{x})\) 强烈吸引样本回归 \(t_i\) 对应的高密度流形,而负向斥力 \(-\nabla_{\mathbf{x}} \log p_{t_k}(\mathbf{x})\) 则有力排斥其它竞争时间步的流形,从而协助样本快速逃离错误的低密度区域。

2. 轻量级时间预测网络设计与引导更新 由于复杂连续数据流形上的真实得分与时间后验分布解析不可求,本文训练了一个辅助轻量级时间预测器 \(\boldsymbol{\phi}\)。该分类器以加噪样本 \(\mathbf{x}_t\) 为输入,输出各时间步的独热标签预测分布 \(\hat{\mathbf{p}}_\phi(\mathbf{x}_t)\),并通过简单的交叉熵损失进行训练: $\(\mathcal{L}_{\text{tp}}(\boldsymbol{\phi}) = -\mathbb{E}_{t, \mathbf{x}_0} \left[ \log \left( \hat{\mathbf{p}}_\phi(\mathbf{x}_t)_t \right) \right]\)$ 该时间预测器仅采用极简的 CNN/MLP 架构,参数量和计算开销远小于主扩散 U-Net/DiT 骨干网络。在采样反向迭代时,综合考虑扩散得分与时间对齐项,TAG 构造的修正得分为: $\(\mathrm{TAG}(\mathbf{x}, t) = \nabla_{\mathbf{x}} \log p_t(\mathbf{x}) + \omega_t \cdot \nabla_{\mathbf{x}} \log p_\phi(t \mid \mathbf{x})\)$ 其中引导强度超参数设为递减调度 \(\omega_t = \omega_0 \sqrt{1 - \bar{\alpha}_t}\)。在算法实现中,先对当前样本执行流形投影更新 \(\tilde{\mathbf{x}}_t \leftarrow \mathbf{x}_t + \omega_t \cdot \nabla \log p_\phi(t \mid \mathbf{x}_t)\),再代入标准反向扩散步,不仅避免了修改主网络权重,而且计算开销极低。

3. 多条件引导重参数化逼近 在复杂的受控生成任务中,常常需要同时施加多个异构条件(例如分子的多种物理化学性质,或人脸的性别、年龄与发色)。若直接为每种条件组合训练联合时间预测器 \(p(t \mid \mathbf{x}_t, \mathbf{c}_1, \dots, \mathbf{c}_m)\),会导致组合爆炸。本文从分层贝叶斯推断出发,证明了可以通过中间潜变量的一阶泰勒展开,将多条件时间后验近似重参数化: $\(\mathbf{x}_t' \approx \mathbf{x}_t - \eta_t^2 \nabla_{\mathbf{x}_t} \ell_1(\mathcal{A}_1(\hat{\mathbf{x}}_0), \mathbf{c}_1)\)$ 使得 \(p(t \mid \mathbf{x}_t, \mathbf{c}_1, \mathbf{c}_2) \approx p(t \mid \mathbf{x}_t', \mathbf{c}_2)\)。进一步推导表明,甚至可以直接利用无条件时间预测器,通过链式减去各条件的指导梯度构建虚拟样本 \(\mathbf{x}_t''\),从而完全复用单个单条件或无条件时间预测器,在多条件复杂场景下以零额外训练代价实现精准的流形对齐。

损失函数 / 训练策略

时间预测器 \(\boldsymbol{\phi}\) 在离散时间步集合 \(\{1, \dots, T\}\) 上独立预训练,利用正向加噪轨迹生成的成对数据 \((\mathbf{x}_t, t)\) 进行多分类监督,训练几万步即可在极短时间内收敛。在反向采样推断阶段,基础扩散模型参数 \(\boldsymbol{\theta}\) 完全冻结,无需任何微调;TAG 作为通用即插即用采样器,支持直接插入 DDPM、DDIM、DPS、TFG 以及 DAS 等各类求解器中。

实验关键数据

主实验

论文在 TFG 标准基准(涵盖图像去模糊、超分辨率、类别条件生成、音频去削波与修复、3D 分子生成等)上全面对比了 DPS 与 TFG 基线及其与 TAG 结合的效果,同时横向对比了 TCS、时间步引导(TSG)与自引导(SG)。在图像保真度(FID)、感知损失(LPIPS)、分类准确率以及音频 FAD/DTW 上均取得显著增益。

任务 / 数据集 评价指标 DPS 基线 DPS + TAG (本文) 相对改善 TFG 基线 TFG + TAG (本文) 相对改善
去模糊 (Deblur / Cat) FID ↓ / LPIPS ↓ 139.7 / 0.613 128.9 / 0.570 +7.7% / +7.0% 64.2 / 0.154 62.7 / 0.151 +2.3% / +1.9%
超分辨率 (SR / Cat) FID ↓ / LPIPS ↓ 139.0 / 0.614 128.3 / 0.572 +7.7% / +6.8% 65.5 / 0.187 64.7 / 0.175 +1.2% / +6.4%
类别引导 (CIFAR-10) FID ↓ / Acc. ↑ (%) 217.1 / 57.5 190.4 / 63.2 +12.3% / +9.9% 114.1 / 55.8 102.7 / 61.5 +10.0% / +10.2%
类别引导 (ImageNet) FID ↓ / Acc. ↑ (%) 196.9 / 24.5 192.2 / 22.9 +2.4% / -6.5% 231.0 / 14.3 219.4 / 17.8 +5.0% / +24.5%
音频去削波 (Audio) FAD ↓ / DTW ↓ 2.41 / 191 2.33 / 189 +3.3% / +1.0% 1.42 / 256 0.74 / 120 +47.9% / +53.1%
音频修复 (Audio Inpaint) FAD ↓ / DTW ↓ 2.26 / 176 2.25 / 157 +0.4% / +10.8% 0.52 / 74 0.42 / 51 +19.3% / +31.1%

在受控 3D 分子生成(QM9 属性控制)任务中,TAG 对原子稳定性(Stability, Stab.)与目标属性绝对平均误差(MAE)的改善更为悬殊,彻底解决了引导梯度破坏分子化学键合结构的顽疾:

控制属性 评估指标 DPS 基线 DPS + TAG (本文) TFG 基线 TFG + TAG (本文) TCS [Jung 2024] TSG [Sadat 2024]
极化率 \(\alpha\) MAE ↓ / Stab. ↑ (%) 13.33 / 28.4 7.96 / 96.4 8.91 / 19.2 4.46 / 43.6 11.44 / 15.3 25.07 / 70.2
偶极矩 \(\mu\) MAE ↓ / Stab. ↑ (%) 4779.92 / 34.4 1.48 / 97.2 2.41 / 26.3 1.28 / 94.3 1.60 / 6.3 N/A / N/A
热容 \(C_v\) MAE ↓ / Stab. ↑ (%) 3.47 / 36.2 3.03 / 93.0 2.65 / 96.2 2.67 / 96.7 3.17 / 19.6 4.18 / 82.9
能隙 \(\epsilon_\Delta\) MAE ↓ / Stab. ↑ (%) 1.65 / 10.6 1.29 / 93.5 1.40 / 16.1 0.78 / 82.8 1.58 / 13.9 1.39 / 48.7

消融实验

论文设计了人工噪声注入实验,系统分析不同噪声扰动等级 \(\sigma\) 与 TAG 引导强度 \(\omega\) 下的生成质量(FID / IS)及时间偏移间距(Time-Gap, TG):

噪声扰动 \(\sigma\) TAG 强度 \(\omega\) Time-Gap (TG) ↓ FID ↓ Inception Score (IS) ↑
0.1 0.0 (无 TAG) 104.1 193.6 2.37
0.1 0.5 47.9 127.7 3.65
0.1 1.0 41.8 120.9 3.69
0.1 2.0 39.0 132.6 3.33
0.1 4.0 44.4 159.8 3.00
0.2 0.0 (无 TAG) 229.6 351.4 1.50
0.2 1.0 175.5 323.7 1.61
0.2 4.0 103.4 246.3 1.70
0.3 0.0 (无 TAG) 274.0 410.1 1.28
0.3 4.0 197.8 361.2 1.31

在少步数加速采样(Few-step Generation)中,离散化误差严重偏离流形,TAG 在无额外引导下同样大幅压低 FID(CIFAR-10 1步采样从 460.0 骤降至 271.1,降低 41.1%;3步采样从 234.1 降至 160.5;ImageNet 1步采样从 430.3 降至 352.8)。在大规模文生图(Stable Diffusion v1.5 + DAS 奖励对齐)中,TAG 将美学评分从 7.948 提升至 9.087,Time-Gap 从 90.04 压降至 28.84;多目标对齐中 Time-Gap 从 20.73 减半至 9.765。

关键发现

  • Time-Gap 与生成质量高度单调负相关:定义指标 \(\text{Time-Gap} = \frac{1}{T}\sum_{t=1}^T |\arg\max \boldsymbol{\phi}(\mathbf{x}_t) - t|\) 能精准量化采样轨迹偏离流形的程度;增大 \(\omega\) 能直接收敛 Time-Gap 并同步显著压低 FID。
  • 对抗强引导漂移鲁棒性强:当 DPS 引导强度从 1.0 激增至 5.0 时,普通 DPS 彻底崩溃(CIFAR-10 FID 从 217.1 恶化至 384.8,准确率暴跌至 29.4%),而 DPS+TAG 将 FID 维持在 347.8 且准确率保持在 41.0%,分子热容控制更是完全避免了崩溃。
  • 单条件/无条件预测器即可泛化至多条件:基于理论推导的重参数化更新,无条件时间分类器在 6 属性联合分子生成中达成 85.9% 稳定性,与专门训练的多条件分类器(91.2%)相当,远超基线独立相加组合的 56.0%。

亮点与洞察

  • 将时间步从反向扩散中的“被动输入”重塑为“主动后验条件”,揭示了时间梯度具备双向排斥与吸引流形的数学本质,理论上通过修正 Langevin 势能面给出了全变差距离(TV distance)上界的收缩保证。
  • 轻量化即插即用:时间预测器仅需几层卷积或感知机,训练耗时数分钟,推断时无需改动基础扩散大模型或重训适配层,打破了过往缓解离流形依赖重型微调或昂贵轨迹搜索的范式。
  • 广泛通用性:无缝统一了反问题逆求解、无微调多条件合成、强化学习奖励测试期对齐(DAS)以及加速离散化采样等多类不同下游场景下的流形崩塌问题。

局限与展望

  • 在时间步接近完全纯噪声的大时刻(\(t \to T\))或低维数据上,由于加噪分布重叠严重,时间预测器的分类准确度有所下降,流形对齐拉力在该区间略有衰减。
  • 引导超参数 \(\omega_0\) 在不同任务间仍需简单的网格搜索调节;若强度过大(如 Table 1 中 \(\omega=4.0\) 在弱噪声下)可能会引入轻微过矫正。
  • 未来的探索方向包括将 TAG 推广到离散扩散模型(Discrete Diffusion)以及基于扩散模型的强化学习轨迹规划策略中。

相关工作与启发

  • vs TCS [Jung et al., 2024]: TCS 采用时间预测器的预测值 \(\tilde{t} = \arg\max \boldsymbol{\phi}(\mathbf{x}_t)\) 直接替换得分网络的时间输入进行硬性时间重分配;TAG 则保持原时间步不变,将时间梯度的软性后验作为外力加回样本自身,提供了更平滑、包含全时间步排斥效应的约束。
  • vs TSG / Self-Guidance [Sadat et al., 2024; Li et al., 2024b]: 扰动时间输入依赖得分模型自身局部的对称时间敏感性;TAG 借助独立时间预测器提供非对称的全局时间后验,在严重离流形区域能够提供更强烈的纠偏动力。
  • vs Classifier Guidance / CFG: CFG 主要解决条件偏好与多样性的权衡,无法感知几何流形偏移;TAG 正交于 CFG,专注于解决因任意引导项引入的得分失效与流形漂移。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 概念清晰优雅,将时间视为后验变量推导出 TLS 流形回正项并给出理论保证。
  • 实验充分度: ⭐⭐⭐⭐⭐ 跨越图像去模糊/超分/类别生成、3D分子生成、音频修复及大模型文生图测试期对齐,验证全面。
  • 写作质量: ⭐⭐⭐⭐⭐ 理论推导完备自洽,动机阐述层次分明。
  • 价值: ⭐⭐⭐⭐⭐ 对扩散模型免训练引导、逆问题求解及测试期对齐具有重要普适意义。