[Paper] 归一化轨迹模型

Source: arXiv - 2605.08078v1

概述

本文介绍了 Normalizing Trajectory Models (NTM)，这是一种在不放弃使其在研究和生产中具有吸引力的 精确似然保证 前提下，加速 diffusion‑based generative models 的新方法。通过将每个逆扩散步骤视为 conditional normalizing flow，NTM 能够在仅 四步 的情况下生成 high‑quality images，同时仍可使用 principled likelihood objective 进行训练。

关键贡献

• 条件归一化流逆向步骤：每个去噪转移都使用表达能力强的可逆块建模，保持精确的似然计算。
• 混合架构：每一步的浅层可逆层 + 一个深层并行预测器，在整个轨迹上共享信息，实现从头端到端的训练。
• 通过轨迹似然进行自蒸馏：精确的似然使得可以在模型自身的分数上训练轻量级去噪器，仅用四个步骤即可获得高保真样本。
• 实证性能：在标准的文本到图像基准上，NTM 的表现匹配或超越强大的扩散基线，同时使用的采样步骤显著更少。
• 兼容预训练的流匹配模型：NTM 可以从已有的流匹配检查点初始化，降低了采用门槛。

方法论

传统的扩散模型通过迭代地消除微小的高斯噪声增量来生成数据，这需要数百步。NTM 将每个逆向步骤重新表述为 条件归一化流：给定当前的噪声潜变量，可逆变换预测噪声更少的前一步。

该架构由两部分组成：

1. 浅层可逆块（例如耦合层），在每个时间步内局部操作，保证雅可比行列式——从而保证精确似然——能够高效计算。
2. 并行预测网络，一次前向传播处理整个轨迹，为每一步的流提供全局上下文（如文本提示）。

训练通过最大化完整逆向轨迹的精确对数似然进行，这与依赖蒸馏或对抗损失的先前少步方法形成对比。由于似然是可求解的，作者还实现了 自蒸馏：他们在完整 NTM 生成的分数上训练一个小型去噪器，得到一个仍然遵循学习分布的快速采样器。

结果与发现

• 四步采样：NTM 生成的图像可与通常需要 50–100 步的最先进扩散模型相媲美。
• 似然保持：与许多加速扩散技术不同，NTM 在整个生成路径上保留有效的概率密度，使得需要精确评分的下游任务（例如不确定性估计）成为可能。
• 基准性能：在流行的文本到图像数据集（如 MS‑COCO、LAION）上，NTM 的 FID 和 CLIP‑Score 指标与 DDIM、DPM‑Solver 以及蒸馏扩散模型等基线持平或更佳。
• 训练灵活性：从预训练的流匹配检查点初始化的模型收敛更快，且相较于从头训练可获得略高的样本质量。

实际意义

• 更快的生产推理：将采样从数百步减少到少数几步，大幅降低延迟，使高质量的扩散生成在实时应用（例如交互式设计工具、设备端图像合成）中可行。
• 精确似然度开启新用例：开发者现在可以将扩散式生成与概率推理相结合——例如基于似然度的异常检测、贝叶斯模型选择或对生成样本进行梯度优化。
• 部署简化：由于逆向步骤是可逆的，内存高效的实现（例如无检查点的反向传播）更容易实现，这对资源受限的边缘或云环境非常有价值。
• 兼容现有流水线：NTM 可以直接嵌入当前的扩散工作流，复用预训练的文本编码器、CLIP 嵌入或扩散先验，同时提供在无需重新训练大型教师模型的情况下加速的明确路径。

Limitations & Future Work

• 模型规模与速度的权衡：并行预测器增加了深度，因此整体参数量可能大于最小化的扩散基线，进而可能提升训练成本。
• 对超高分辨率的可扩展性：实验聚焦于标准基准分辨率（256–512 像素）。将 NTM 扩展到 1024 像素以上的图像可能需要更复杂的可逆块或层次化设计。
• 超出文本到图像的泛化能力：虽然论文展示了在图像生成上的强劲结果，但将 NTM 应用于其他模态（音频、视频、3‑D）仍是一个未解之题。
• 自蒸馏质量上限：轻量去噪器能够匹配四步 NTM，但仍略逊于完整模型；未来工作可以探索多阶段蒸馏或自适应步数调度。

总体而言，归一化轨迹模型在似然‑基生成建模的理论严谨性与快速高质量采样的实际需求之间架起了一座有力的桥梁——这种组合有望加速扩散技术在各类面向开发者的 AI 产品中的应用。

作者

• Jiatao Gu
• Tianrong Chen
• Ying Shen
• David Berthelot
• Shuangfei Zhai
• Josh Susskind

论文信息

• arXiv ID: 2605.08078v1
• 分类: cs.CV, cs.LG
• 发布日期: 2026年5月8日
• PDF: Download PDF