归一化轨迹模型

论文信息

标题: Normalizing Trajectory Models

作者: Jiatao Gu, Tianrong Chen, Ying Shen, et al.

发布日期: 2026-05-08

PDF 链接: 下载 PDF

论文背景与研究动机

生成模型在过去几年经历了革命性变化，扩散模型（Diffusion Models）凭借高质量和稳定的训练成为图像、音频等多模态生成的主流范式。扩散模型的核心思想是将复杂的数据分布转化为简单分布（如高斯噪声），再学习逆过程逐步去噪：采样时从纯噪声出发，经过数百甚至上千步微小的去噪操作，最终还原出清晰样本。每个去噪步通常被建模为高斯分布，这种“小步高斯去噪”的分解方式赋予扩散模型精确的似然估计能力，同时带来高质量的生成效果。

然而，小步长和高斯假设在需要快速生成的实际应用中成为严重瓶颈。如果将整个生成过程强行压缩到极少的粗步长（例如 4 步或 8 步），单步需要覆盖的分布变化将变得剧烈而复杂，再假设为简单高斯分布显然不再合理。为解决少步生成问题，学术界已发展出多种路线：**蒸馏（Distillation）**将教师模型的多步轨迹压缩到学生模型的少数步；**一致性训练（Consistency Training）**强制模型在轨迹上相邻点输出一致；**对抗目标（Adversarial Objectives）**使用判别器提升单步生成质量。但这些方法普遍舍弃了精确似然框架——蒸馏依赖教师模型近似，一致性训练使用替代损失，对抗训练更无从谈论似然。没有似然就无法进行基于概率的异常检测、模型评估和无偏采样，也丧失了生成轨迹的可解释性和理论保障。

论文《Normalizing Trajectory Models》正是在这种背景下，提出了保持精确似然的少步生成新范式。作者观察到，扩散模型中生成轨迹（trajectory）本身是一个从前向逐步加噪到逆向逐步去噪的连续过程，如果为每一步逆过程赋予强大的表达能力，就能在极少步数内精确建模复杂的数据变换。为此，他们引入归一化流（Normalizing Flow）来参数化每一个逆向步骤，形成所谓归一化轨迹模型（NTM）。由于归一化流天然提供精确对数似然，NTM 在任意步数下都能对整个生成轨迹进行精密似然训练，同时实现高质量少步采样。

核心方法和技术细节

NTM 的设计哲学可提炼为一句话：将扩散轨迹视为一系列可逆变换，并用归一化流为每一步赋予表达力。

传统扩散逆向过程通常为：

$p_\theta(x_{t-1}|x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \sigma_t^2 I)$

这个高斯假设在步数极少时严重受限。NTM 则将其替换为一个条件归一化流：

$x_{t-1} = g_\theta(z_t; x_t, t), \quad z_t \sim p_Z$

其中 $g_\theta$ 是相对于 $z_t$ 可逆的变换， $p_Z$ 一般为标准高斯。给定当前状态 $x_t$ 和时间 $t$ ，归一化流将基础噪声映射为下一步状态 $x_{t-1}$ 。因为归一化流是可逆的且雅可比行列式可计算，我们能精确写出单步条件似然：

$\log p_\theta(x_{t-1}|x_t) = \log p_Z(z_t) - \log\left|\det\frac{\partial g_\theta}{\partial z_t}\right|$

沿着整条轨迹 $x_T \to x_{T-1} \to \cdots \to x_0$ ，联合似然分解为各步条件似然之和，再加上前向噪声过程的已知分布，模型可通过最大似然直接端到端训练。这种设计使 NTM 成为第一个在极低步数下仍然保持精确轨迹似然的扩散类生成模型。

为实现高效且表现力强的 $g_\theta$ ，论文提出了精心设计的架构。每个逆向步骤内部由多个**浅层可逆块（shallow invertible blocks）**堆叠而成，这些块通常使用仿射耦合层（affine coupling layers）或可逆 $1\times1$ 卷积等典型归一化流构件，保证可逆性和雅可比行列式的易计算性。浅层设计避免了单步计算开销过大，因为实际生成只需 4 步左右，总推理成本依然可控。

但仅仅在每个局部步骤放置可逆块并不足够——模型还需要把握跨步骤的全局一致性。为此，NTM 引入深度并行预测器（deep parallel predictor），它横跨整个轨迹，为每个时间步骤的归一化流提供条件信息。具体来说，预测器接收当前噪声状态 $x_t$ 和时间 $t$ ，输出归一化流所需的网络参数（如耦合层中神经网络的权重和偏置）。这个预测器本身是一个深度网络，但其计算在所有时间步上共享，并可通过并行机制高效运行。这种“浅局部可逆块 + 深全局预测器”的组合，既保证了每个逆向步强大的分布建模能力，又让模型能从整个轨迹视角学习一致的去噪策略。

训练方面，NTM 支持从零开始训练（train from scratch），也可从预训练的流匹配模型（flow-matching models）初始化。流匹配是近年来流行的一种训练连续归一化流的框架，其训练目标与扩散模型类似，但侧重于学习向量场。NTM 通过离散化流匹配中的连续时间轨迹，可将预训练权重直接映射到对应离散步骤的可逆块中，大幅加速收敛并提升最终性能。这使得 NTM 能够继承大规模预训练模型的先验知识，再通过精确似然训练在少步设置下进一步微调。

NTM 的另一个关键创举是自蒸馏（self-distillation）。由于 NTM 能计算轨迹上任意点的精确得分（score），即 $\nabla_{x_t}\log p(x_t)$ ，论文在训练好的 NTM 模型基础上，训练一个轻量级去噪器去拟合这些得分。这个去噪器结构简单，生成时只需执行少数几步确定性或随机过程，却因教师信号来自 NT M 自身的精确得分，能产生极高质量的样本。实验表明，使用 4 步自蒸馏采样，NTM 的表现可以媲美甚至超越许多强基线模型。自蒸馏完全保留在精确似然框架内，与通常丢弃似然的对抗蒸馏有本质区别。

创新点和贡献

NTM 的首个子贡献在于概念突破：它抛弃了“逆向步必为高斯”的教条，将每一步赋予归一化流的强表达能力，从而将扩散模型框架与可逆生成模型的优点深度融合。这从根本上解决了少步生成下高斯假设失效的问题。

第二，精确似然的保留是论文最鲜明的旗帜。在主流少步方案纷纷牺牲似然时，NTM 反其道而行之，用归一化流天然的对数似然计算能力，保证了整个生成轨迹上每一步的概率都可以严格评估。这不仅使模型训练有明确的最大似然目标，还开启了基于似然的下游应用（如离群检测、概率校准、可控生成中的重要性采样）。

第三，架构设计的巧思——“浅可逆块 + 深并行预测器”的异构组合，实现了表达力、计算效率和轨迹一致性三者的平衡。深预测器让局部模块可以“看到”全局，但又不必在每个时间步重复繁重计算；浅可逆块则控制单步成本，使 4 步推理的实际用时完全可以接受。

第四，自蒸馏机制将模型自身的得分作为监督信号去训练小型快速采样器，相当于在不引入额外偏差的前提下，再一次压缩了采样成本。这种“自我提升”式的训练流程非常优雅，且源于 NT M 内在的似然能力。

第五，与预训练流匹配模型的兼容使得 NTM 可充分利用大规模预训练成果，在文本到图像等高难度任务上快速达到最佳性能。这一特性大大降低了研究和工程上的尝试成本。

实验结果分析

论文在文本到图像生成基准上进行了系统验证，主要与强扩散类基线（如 Stable Diffusion 及其少步变体）、一致性模型、对抗蒸馏方法等对比。评估指标包括 FID（Fréchet Inception Distance）、CLIP Score 等常见图像生成质量度量。

核心实验结果显示，仅用 4 个采样步数的 NTM，在 FID 指标上与传统需要 50～100 步的强基线持平或更优，同时在 CLIP Score 上也保持高水平，表明语义匹配度未因步数减少而损失。与同样针对少步设计的蒸馏方法相比，NTM 由于保留了精确似然训练，不仅在生成质量上有竞争力，在概率度量上还具有独特优势。当启用了自蒸馏后，4 步采样的轻量去噪器甚至能够超越 NTM 自身的直接采样效果，体现出自我提炼的有效性。

值得一提的是，NTM 在训练过程中并未像对抗训练那样遭受模式坍塌风险，也不像一致性模型那样需要精细调参。其训练稳定性和收敛速度得益于最大似然目标的良好几何性质。消融实验表明，深度并行预测器和浅可逆块的设计对最终性能至关重要，缺少任一部分都会导致生成质量明显下降。

尽管实验主体集中在文本到图像，这些结果具有很强的代表性：文本条件生成是高度多模态、高维输出的复杂任务，NTM 在此设定下证明了自己。可以预见，该框架在无条件图像生成、图像到图像转换、音频生成等其他领域同样具有广阔前景。

实践应用建议与未来发展方向

对于从事生成模型研发与应用的实践者，NTM 提供了一条既快又准且可解释的新路径。以下建议可供参考：

资源敏感场景的生成：在移动设备或实时交互系统中，推理步数是硬约束。采用 NTM 4 步采样可直接替换原有扩散模型，保持质量同时大幅降低延迟。由于 4 步过程仍可计算似然，可以对生成结果进行置信度评估，对高风险应用（如医学图像生成、自动驾驶仿真）尤为重要。
预训练模型迁移：如果已有流匹配预训练模型（例如使用条件流匹配训练的文生图模型），可以利用 NT M 的初始化策略快速获得少步生成能力，无需重复昂贵的从零训练。这降低了尝试成本，适合团队快速验证概念。
自蒸馏为其提速：在实际部署时，可先用完整 NTM 计算高质量得分，再蒸馏出一个极简网络。蒸馏后的模型完全脱离归一化流结构，可进一步进行算子融合和量化，实现极致速度。
概率驱动的应用：NTM 的精确似然可用于错例检测、领域外样本识别。生成模型在遇到异常输入时，其似然值会显著偏离正常范围，这种特性是传统蒸馏或一致性模型无法提供的，可为企业级可靠系统增加一层安全网。

未来研究方向同样令人兴奋。一是架构轻量化：当前浅可逆块的设计仍有压缩空间，可探索使用更高效的可逆变换（如神经样条流、残差流）来进一步降低单步开销。二是扩展到离散数据：文本、分子图等离散数据上的少步生成同样渴望精确似然，需要设计对应的离散归一化流轨迹模型。三是多步生成的理论深化：NTM 提供了一个用于研究少步生成过程信息损失的完备概率框架，有望指导最佳步数与计算预算的权衡。四是与其他似然自由方法的结合：虽然 NTM 强调精确似然，但也可在蒸馏阶段引入感知损失或对抗损失做微小补充，达到感知质量与概率保真度的更好平衡。

总结与展望

Normalizing Trajectory Models 通过将归一化流嵌入扩散轨迹，优雅地解决了少步生成中高斯假设失效和似然丢失的双重难题。其“浅可逆块+深并行预测器”的架构在表达力与效率间取得精妙平衡，配合自蒸馏机制，可在仅 4 步推理下交出媲美数十步扩散模型的成绩，且全流程保留正则化似然。这一工作不仅是技术上的精巧拼接，更代表一种理念上的回归：生成模型应当既能快速采样，又不丢失对数据概率结构的深刻理解。

展望未来，NTM 很可能引领一股“概率保真的少步生成”浪潮。随着可逆神经网络和流模型的持续进化，NTM 的每一步都会变得更强、更快，促使生成式 AI 从“仅关注样本美观”迈向“美观与概率严谨兼顾”的新阶段。届时，我们或许不再需要在质量和速度、似然和感知之间反复权衡——NTM 已证明，鱼与熊掌可以兼得。