NTM：打破高斯瓶颈，开启 4 步精确似然生成新时代

TL;DR

扩散模型正在统治生成 AI，但其“快不起来”的致命伤源于少步采样下的高斯假设失效。Apple 团队推出的 Normalizing Trajectory Models (NTM) 巧妙地将每一采样步建模为一个极具表达力的 Normalizing Flow。它不仅在 4 步内达成了 SOTA 级别的图像质量，更在 Normalizing Flow 领域首次实现了与顶级扩散模型抗衡的扩展性，同时保留了宝贵的精确对数似然（Exact Log-likelihood）。

痛点深挖：消失的高斯近似

在经典的扩散模型中，我们假设从噪声到图像的每一步微小逆向过程都是高斯分布。当步数高达 50-100 步时，这个假设是成立的。然而，为了追求效率，我们希望将步数压缩到 4 步甚至更少。此时，每一步跨度极大，真实的分布会变成极为复杂的多峰、长尾分布。

如果你强行用单高斯去拟合这个复杂分布，结果就是生成图像的模糊、伪影和语义缺失。前人尝试用 GAN 或是蒸馏（Distillation）来修补，但前者训练极其不稳定，后者则让模型变成了黑盒，失去了计算概率密度的能力。

核心方法论：NTM 的“预测-传输”架构

NTM 的直觉在于：如果在原始像素空间不好建模非高斯分布，那就找一个容易建模的潜在空间。

1. 架构拆解

NTM 由两个关键组件构成：

• Invertible Transporter ($f_{\mathcal{T}}$)：一个浅层的可逆变换层（通常使用 TarFlow 风格的因果自回归块）。它的任务是将复杂的非高斯分布“拉直”成更易处理的形式。
• Deep Predictor ($f_{\mathcal{P}}$)：一个深层 Transformer 架构。它在轨迹维度进行并行推理，预测潜在空间中的高斯均值和方差。

图注：NTM 概览。Transporter $f_{\mathcal{T}}$ 处理局部非高斯性，Predictor $f_{\mathcal{P}}$ 负责捕捉跨时间步的复杂关联。

2. 精确似然与训练策略

通过变量变换公式（Change-of-variables），NTM 的损失函数直接对应轨迹的精确负对数似然： ${\mathcal{L}}_{NTM}=-\log p_{\mathcal{P}}(u_s|u_t)-\log |\det J_{f_{\mathcal{T}}}|$ 这种设计确保了训练的高稳定性，并允许模型不仅能生成、还能评估生成内容的分数。

实验与战绩：少即是多

NTM 在 ImageNet 和大规模 Text-to-Image 任务上展现了惊人的效率：

• 极速采样：仅需 4 步，GenEval 得分即达 0.82，甚至在图像质量和语义对齐上反超了需要 50 步的 FLUX.1 等强基线。
• 微调能力：通过“恒等映射初始化”和“均值对齐辅助损失”，NTM 可以无缝吸收预训练 Flow Matching 模型（如 FLUX 系列）的先验知识，将其转化为高效的采样器。

图注：NTM 的 4 步生成效果。左侧为 256x256 从零训练，右侧为 512x512 微调结果，细节表现力的保留非常出色。

深度洞察：为什么 1 步还是不行？

作者在文中坦诚讨论了失败案例：当步数减少到 $T=1$ 时，效果会剧烈下滑。这是因为单步生成要求 Transporter 独立承担捕获整个数据分布复杂性的重任，而当前的浅层架构难以支撑这种复杂度。这预示着未来若要实现完美的单步采样，可能需要更深层的自适应可逆架构。

总结

NTM 是生成模型领域的一次重要回归——回归到概率分布的本质，同时拥抱了 Transformer 的扩展性。它填补了 Normalizing Flow（高学术价值但难训练） 与 Diffusion Models（强生成能力但采样慢） 之间的鸿沟。

Takeaway：对于追求工业级推理速度且不愿放弃模型可解释性、可评估性的团队来说，NTM 提供的“非高斯轨迹建模”方案是目前最值得关注的路径。

编者注：该工作由 Apple 团队完成，代码已开源在 https://github.com/apple/ml-starflow。