归一化轨迹模型：用归一化流建模扩散反向条件

用归一化流精确建模扩散反向条件，实现4步生成且保留似然

前置知识

扩散模型与流匹配

扩散模型通过前向加噪 $q(x_t|x_0)=\mathcal{N}((1-t)x_0,t^2I)$ 把数据逐步变噪声，反向用单高斯 $p_\theta(x_s|x_t)$ 复原。流匹配学习速度场 $v_\theta(x_t,t)$，ODE 积分得样本。少步时反向的高斯假设失效。

本文所有问题都源自反向条件的高斯假设在少步生成时失效，不理解扩散/FM 框架就读不懂 NTM 的动机。

归一化流与 TarFlow

归一化流把数据 $x$ 通过可逆映射映到潜变量 $z\sim N(0,I)$，由变量替换公式给出精确对数似然。TarFlow 用因果 Transformer 实现自回归仿射耦合，使流模型能扩展到高分辨率图像，是 STARFlow/NTM 的核心组件。

NTM 的 transporter 直接复用 TarFlow 块，是其能精确计算对数似然的核心组件。

高斯近似瓶颈

理论上反向条件 $p(x_s|x_t)=\int p(x_s|x_t,x_0)p(x_0|x_t)dx_0$ 是关于 $x_0$ 的高斯混合，因为先验 $p(x_0|x_t)$ 对自然图像很复杂。少步生成时混合分量的峰更多、拖尾更重，单高斯无法覆盖。

这是 NTM 要解决的根本瓶颈——把单高斯替换成可表达条件分布但仍具精确似然的归一化流。

轨迹分数去噪

$\nabla_x\log p(x)$ 是去噪方向。STARFlow 用 NF 似然梯度逐样本去噪。NTM 把整条轨迹视作联合随机变量，用马尔可夫协方差 $S$ 一次反传修正所有时间步，$[S]_{ij}=\min(t_i,t_j)^2(1-\max)/(1-\min)$。

这是 NTM 相对 STARFlow 的关键扩展，也是它能把训练好的模型蒸馏成 4 步采样器的核心机制。

研究动机

当前主流扩散/流匹配模型在反向过程中默认每步条件 $p(x_s|x_t)$ 是单高斯，这一假设只有在步数足够多时（通常 20–50 步甚至更多）才近似成立。把步数压缩到 4 步时，单步反向需要跨越很大噪声区间，真实反向条件变成多模态、重尾的高斯混合，单高斯拟合误差急剧增大，导致图像模糊、结构错乱。已有少步方法如一致性模型、渐进式蒸馏、DMD、DDGAN 等都另辟蹊径：蒸馏和一致性模型舍弃似然框架，DDGAN 用对抗训练替代高斯但牺牲密度估计且训练不稳定。因此业界缺乏「少步 + 精确似然 + 可表达反向条件」三全其美的方法。

本文的目标是本文提出归一化轨迹模型（NTM），目标是把每个反向条件 $p(x_s|x_t)$ 用条件归一化流精确建模，在端到端负对数似然 $-\log p(x_s|x_t)$ 训练下，让 4 步反向就能产出高质量样本，同时保留整条生成轨迹的精确对数似然与可解析的 score 信号。模型既能完全从零训练，又能从任意预训练流匹配 checkpoint 微调；进一步把反向迭代过程中的 trajectory score denoise 蒸馏成一个轻量 denoiser，让推理时只需 predictor+denoiser 两次前向就能 4 步生成。文生图与类条件 ImageNet 实验覆盖 256×256 / 512×512 两个尺度和自监督与微调两条路线，目标都是用最少的反向步数逼近甚至超越 SOTA 扩散模型的合成质量，同时不让出似然框架。

与已有工作不同的是，与自监督表征学习的 encoder–predictor 架构（如 BYOL、I-JEPA）不同，本文让 encoder 变成严格可逆、等维的 transporter，从而把表征学习目标无缝转写成归一化流的负对数似然；同时引入轨迹级协方差加权的 score 去噪并蒸馏成一个单独的 denoiser，把反向迭代和 backprop 去噪都收敛到 4 个前向 pass。这是「用归一化流显式建模反向条件」的首次可扩展尝试。

核心方法

直觉上，NTM 想要找到一种「隐空间」，使得在那个空间里少步反向变简单。结构上它由两个组件构成：(1) 一个共享的可逆 transporter $f_T$，把任意 $x_t$ 和 $x_s$ 映射到同一个 $u$-空间潜变量；(2) 一个深度 predictor $f_P$，在 $u$-空间内从 $u_t$ 和随机潜变量 $z\sim\mathcal{N}(0,I)$ 预测 $\hat u_s$。因为 $f_T$ 可逆，$x_s\leftrightarrow u_s\leftrightarrow z$ 整体构成一个条件归一化流，端到端最小化 NLL；多步组合可同时对所有时间步计算联合轨迹似然。可从零训练（用随机前向轨迹 + T 个不同长度的多轨迹训练），也可以从任意预训练流匹配 checkpoint 微调（transporter 设成恒等、predictor 用后验均值初始化，并加均值对齐辅助损失）。

核心创新是把 BYOL/I-JEPA 风格的 predictor–encoder 框架严格「归一化流化」：$f_T$ 必须是同维可逆映射而非压缩式编码器，这样变化量替换公式给出每一步的精确似然；$f_P$ 退化成简单的仿射高斯，$f_T$ 负责把混合高斯「抹平」成单高斯。第二个本质区别是轨迹分数去噪：NTM 把整条 $(\hat x_{t_0},...,\hat x_{t_T})$ 当作一个被前向过程噪声污染的马尔可夫序列，其协方差矩阵 $S$（元素 $[S]_{ij}=\min(t_i,t_j)^2(1-\max(t_i,t_j))/(1-\min(t_i,t_j))$）让一次 backprop 同时修正所有时间步，相关性强过逐样本去噪。第三个创新是均值对齐辅助损失 $\mathcal{L}_{\rm aux}=\|\mu_P-\mu_{FM}\|^2$，把微调中的 predictor 锚定到预训练速度场，避免灾难性漂移。

方法步骤详情

训练：(1) 按 $\alpha=\frac{1-t}{1-s}$、$\sigma=\sqrt{t^2-\alpha^2 s^2}$ 在 $t_0<\cdots<t_T$ 上采样轨迹。(2) 共享 transporter $f_T$（2 块 TarFlow、4 层）映 $x_t,x_s$ 到 $u_t,u_s$。(3) predictor $f_P$（24 层 Transformer）输出 $\mu_P,\sigma_P$，采 $\hat u_s=\mu_P+\sigma_P z$。(4) 端到端损失 $\mathcal{L}_{NTM}=\sum\tfrac12\|z_k\|^2+\sum\log\sigma_P+\sum\log\sigma_T$，每 batch 独立采 $T\in\{4,8,16\}$。微调时 $f_T$ 设恒等、$f_P$ 用后验均值 $\mu_{post}$ 初始化，并加辅助损失 $\|\mu_P-\mu_{FM}\|^2$ 锚定。推理可选完整 NF+trajectory denoise 或训练轻量 $g_\phi$ 做一次前向去噪。

技术新颖性

新颖性体现在三处：(i) 把 NLL 训练目标无缝嵌入 predictor–encoder 自监督范式，使可逆表征学习与精确似然生成首次合流；(ii) 首次利用马尔可夫前向协方差 $S$ 在反向轨迹上做 joint score denoising，比 STARFlow 的逐样本去噪更高效，并能蒸馏进单步前向 denoiser；(iii) 均值对齐辅助损失让从大模型 FM checkpoint 微调成为可能，避免 NLL 单独优化时的灾难性遗忘，从而让 4B 参数 FLUX.2-klein backbone 在 4 步下达到 0.76 GenEval/83.38 DPG。

**Figure 3**: NTM overview. Shared transporter f_T maps x_t, x_s to representations u_t, u_s with a tractable Jacobian. The predictor f_P takes u_t and latent z to produce û_s. D measures the distance between prediction and target at distribution level.

**Figure 4**: Finetuning diagram. L_aux aligns μ_P with frozen μ_FM; L_NTM trains the full model.

**Figure 5**: Denoiser training via trajectory score denoising. Frozen NTM computes NLL and its gradient refines every position producing x_den targets; a denoiser g_φ predicts x^den_t0 directly from u_t0.

**Figure 8**: Comparison between TARFlows, Normalizing Trajectory Models, and Diffusion Models. Subscripts denote spatial patch indices, superscripts denote trajectory timesteps, with x_t1, u_t1 pure noise and x_t0, u_t0 clean data.

实验结果

从零训练在 256×256、约 7000 万图文对（含 CC12M）、64 块 H100、batch 1024 下，4 步 NTM 在 GenEval 0.82、DPG-Bench 79.64，相对需 256 步 AR 的 STARFlow（0.56）大幅提升，逼近需多步的扩散强基线。类条件 ImageNet 上 16 步 NTM 取得 FID 2.80、4 步 3.83，与 STARFlow(FAE) 2.67/256 步可比，证明精确 NLL 即可得强少步模型。微调 FLUX.2-klein 4B 在 512×512，4 步 NTM 达 GenEval 0.76、DPG 83.38，与 FLUX.1-dev（0.66/83.84）相当但步数从 50 降到 4。Table 2 显示完整 NF+traj denoise 0.20 img/s、predictor+learned denoiser 1.88 img/s，约 9× 加速同时 LPIPS 仅 0.121。消融（图 6、7）证 T=4 最佳、无辅助损失早发散、T=1 严重退化（图 9）。

**Table 1**: T2I Evaluation. GenEval overall score and DPG-Bench percentage.

**Table 2**: Score denoising vs. learned denoiser (finetuned setting).

**Figure 1**: Text-to-image generation with NTM in 4 denoising steps. Samples from models trained from scratch at 256×256 and from finetuning pretrained flow-matching checkpoints at 512×512.

**Figure 6**: Ablation: multi-trajectory training. Comparison of NTM evaluated with T=4, T=8, T=16 and baseline FLUX (50 steps).

**Figure 7**: Ablations. (a) Finetuning directly with NF loss diverges. (b) Adding mean-alignment loss stabilizes training. (c) Comparison of denoising approaches.

查看结构化数据

任务	指标	本文	基线	提升
文生图（从零训练 256×256）	GenEval 整体分	0.82（4 步）	STARFlow 0.56（256 AR 步）	在 1/64 步数下显著超越 SOTA 归一化流文生图模型
文生图（从零训练 256×256）	DPG-Bench	79.64%（4 步）	FLUX.1-dev 83.84%（50 步）	以少 12× 步数达到接近水平
文生图（FLUX.2-klein 微调 512×512）	GenEval	0.76（4 步）	FLUX.1-dev 0.66（50 步）	在 1/12 步数下超越 FLUX.1-dev 0.10 分
文生图（FLUX.2-klein 微调 512×512）	DPG-Bench	83.38%（4 步）	FLUX.1-dev 83.84%（50 步）	在 1/12 步数下基本追平
类条件 ImageNet 256×256（从零）	FID	2.80（16 步）/ 3.83（4 步）	STARFlow(FAE) 2.67（256 AR 步）	在 1/64 步数下质量相当（4 步 3.83 略逊）
微调模型推理加速	img/s	1.88 img/s（predictor+learned denoiser）	0.20 img/s（完整 NF + trajectory denoise）	约 9× 加速，LPIPS 仅 0.121 损失少量逼真度

局限与改进

作者明确指出：(1) 单步 T=1 时 NTM 严重退化（图 9），原因不在训练而在容量——浅 transporter（2 块 × 4 层）需把所有非高斯结构压缩到一层，像素级细节难以承载；(2) 512×512 微调模型在位置绑定、属性绑定子任务上仍有挑战，作者归因为微调时长不足或预训练 backbone 选择问题；(3) 轨迹 score denoise 虽效果最好但仍需一次 backprop，相对 learned denoiser 的纯前向在延迟敏感场景下不占优；(4) 完全依赖精确 NLL 训练，没有用对抗或感知损失，对极致细节和文本渲染或不如大扩散模型。观察：训练资源 64 块 H100、batch 1024、内部分布式数据 70M 对，多数研究者难以复现；4 步配置下 DPG-Bench 与 FLUX.1-dev 仍有 ~4.5 个百分点差距，提示短步与大模型的扩展曲线有待进一步研究。

独立分析的弱点

(1) 单步 T=1 性能崩塌——可加深的 transporter 配合 hierarchical 自适应深度或把 transporter 嵌入到 predictor 每层，而非固定 2 块浅层；可探索 per-token 或 per-region 的自适应深度 transporter。(2) 微调模型在位置/属性绑定子任务上落后基线，可尝试更长微调 + 引入分布级对抗/感知损失补足 NLL 在高频细节上的盲点。(3) 轨迹 score denoise 仍需 backprop，可改成无梯度近似（如 Neumann 级数截断 / Fisher 评分近似）或把 backprop 完全折进一个大 denoiser。(4) DPG-Bench 上从零 4 步落后 FLUX.1-dev 约 4 个百分点，可考虑与判别式 reward（如 HPSv2、PickScore）结合做 RL 后训练。(5) 隐空间维度受 TarFlow 块限制 32 维（FAE latents），更高分辨率/更高维潜空间对 transporter 可逆性的算力压力未充分讨论。

未来方向

作者方向：(a) 引入对抗 / 感知 / CLIP reward 等分布级后训练目标提升少步质量；(b) 扩展到 1024+ 分辨率、文生视频、个性化等场景；(c) 探索自适应深度 transporter 推动 1–2 步精确似然生成。延伸方向：(d) 用 NTM 显式似然做似然引导编辑、inpainting、reward guidance，无需额外训练 reward model；(e) 把 learned denoiser 与 LCM、NFM 等蒸馏框架混合，进一步把 4 步压到 1–2 步；(f) NTM 在跨模态（如文生 3D、文生音频）上的扩展潜力；(g) 探索 trajectory score denoise 在 RL 偏好对齐下的可微替代。

复现评估

代码：作者在文末公开 https://github.com/apple/ml-starflow，与 STARFlow 同源，方便扩展。数据：使用内部约 7000 万图文对（含 CC12M），未完全公开；外部研究者只能用 CC12M 类似语料加 LAION 等；类条件 ImageNet 256×256 是公开可复现的。算力：从零训练 64 块 H100、batch 1024、bfloat16 + FSDP，单次完整训练需数十万美元，普通实验室难以负担；微调 4B 参数同样需数十块 H100。超参（$\lambda=2.5$、CFG 10% dropout、默认 T=4、T∈{4,8,16} 多轨迹训练）大部分在正文和附录给出。难点：TarFlow 块的可逆性需实现因果掩码 Transformer；denoiser 后训练需稳定 NLL 通过反向 predictor 链路反传。

归一化轨迹模型：用归一化流建模扩散反向条件 Normalizing Trajectory Models