嵌入预测有助于图像生成

论文信息

标题: Embedding Prediction Helps Image Generation

作者: Sihan Xu, Ji Xie, Zilin Wang, et al.

发布日期: 2026-10-01

arXiv ID: 2610.02203v1

PDF 链接: 下载 PDF

3 分钟速览

• 研究问题：扩散 Transformer 的条件通常是固定的类标签或文本嵌入，在所有去噪步骤中保持不变；论文检验是否可以让条件本身从当前噪声图像中预测出来。
• 核心方法：先训练 NEPA 模型，用 Multi-Embedding Prediction 从条件和噪声图像一次性预测干净图像的全部补丁嵌入，再把这些预测嵌入作为冻结条件，在每个去噪步骤输入 DiT 生成器。
• 关键结果：在 ImageNet 256×256 类条件生成中，NEPA-DiT-XL 结合 REPA 达到 FID 1.32（SDE-250 采样），训练计算约为 REPA 的三分之一（见论文表 9 和第 4.4 节）。
• 主要局限：实验仅覆盖类条件 ImageNet 256×256；文本条件和更高分辨率仍是未来工作；推理时每步都运行额外的 NEPA 模型会增加计算量。
• 适合读者：关注扩散 Transformer、表示学习、条件机制或生成模型训练效率的研究者和工程师。

论文背景和研究动机

在主流扩散 Transformer 中，条件信号通常只被嵌入一次：类标签或文本提示经过编码后，在所有去噪步骤中被重复使用。生成器需要自己从固定的条件中推断出对当前噪声状态的含义。已有工作将预训练视觉表示引入生成流程，例如作为生成器的对齐目标、扩散潜空间或中间生成目标，但这些方法的条件本身仍然不随图像状态变化。

本文提出相反思路：条件本身也可以被预测。NEPA（Next-Embedding Predictive Autoregression）训练 Transformer 预测序列中的下一个连续嵌入，而不是重建像素或离散 token。在图像生成中，序列可以被组织为「条件 → 噪声图像 → 干净图像」，因此干净图像的嵌入就是条件和噪声图像之后的「下一个嵌入」。这个预测结果自然成为生成器可以依赖的动态条件。作者的核心问题由此变为：能否在每一步去噪时，让条件嵌入根据当前噪声输入重新计算。

核心方法和技术细节

方法分成两个阶段。第一阶段训练 NEPA 模型，第二阶段训练 DiT 生成器，生成器训练时 NEPA 模型冻结。

Multi-Embedding Prediction（MEP） 扩展了 NEPA：普通 NEPA 只预测下一个嵌入，而 MEP 让模型从同一上下文中一次性预测接下来的 $K$ 个嵌入，其损失为

$$\mathcal{L}_{\mathrm{MEP}}=\frac{1}{K}\sum_{k=1}^{K}\mathcal{D}(\mathbf{z}_{n+k},\mathbf{p}_{n+k})$$

其中 $\mathcal{D}$ 是嵌入空间中的距离。对于图像生成，论文取 $K=N$，即一次预测干净图像的全部 $N$ 个补丁嵌入。输入序列为 $[c,\langle \mathrm{boi}\rangle,\mathbf{z}_t]$，其中 $\mathbf{z}_t$ 是噪声图像的补丁嵌入，输出为干净图像嵌入 $\mathbf{z}_0$ 的预测。噪声水平 $t$ 从均匀分布 $t\sim\mathcal{U}(0,1)$ 采样，使同一模型覆盖所有去噪阶段。

注意力掩码的设计是：条件 token 和 $\langle \mathrm{boi} \rangle$ 组成因果前缀，不关注图像 token；噪声图像的补丁 tokens 关注整个前缀以及彼此，但不关注干净图像。这样每个预测都能利用当前噪声图像的全部空间上下文。由于前缀不依赖图像 token，其键值缓存可以在采样前计算一次并复用。

预测损失采用 InfoNCE：每个预测 $\mathbf{p}_i$ 与对应的干净补丁嵌入 $\mathbf{z}_{0,i}$ 作为正样本对，同一图像中的其他补丁作为负样本。实验显示，在该评测设置下，InfoNCE 得到 FID 31.36，优于负余弦相似度 37.82 和 MSE 35.23（论文表 1）。作者推测 InfoNCE 迫使模型区分同一图像中的不同补丁，因此保留了更多补丁级细节。

Embedding Conditioned Generation（ECG） 在第二阶段使用冻结的 NEPA 模型生成条件。每一步去噪时，当前噪声输入 $\mathbf{x}_t$ 先经过 NEPA 模型，得到预测嵌入 $\mathbf{p}_t$；生成器遵循 DiT 结构，但通过 SD3 的 MM-DiT 块，将预测嵌入与噪声潜在 token 分别处理并做联合注意力。生成器不再有单独的类嵌入，时间步通过 adaLN 调制输入。训练目标为加权流匹配损失：

$$\mathcal{L}_{\mathrm{ECG}}=\mathbb{E}_{\mathbf{x}_0,\boldsymbol{\epsilon},t,c}\left[w(t)\left\|\mathbf{v}_t-G_{\phi}(\mathbf{x}_t,t,\mathbf{p}_t)\right\|^2\right]$$

分类器自由引导方面，训练时以 0.1 概率将条件替换为空 token；推理时分别计算条件分支和无条件分支的速度，再按引导尺度加权。

创新点和贡献

论文提出三项主要贡献：第一，提出 Embedding Conditioned Generation，让 DiT 生成器以 NEPA 模型从条件和噪声图像中预测的干净嵌入为条件，并且在每个去噪步骤重新计算；第二，提出 Multi-Embedding Prediction，将 NEPA 的单嵌入预测扩展到一次性预测整个干净图像的嵌入；第三，在类条件 ImageNet 256×256 上系统比较了生成器的条件设计、MEP 的损失选择和模型缩放行为。

与固定条件相比，本文方法的特点是条件信号会随着噪声状态变化。与自条件化和 RIN 不同，预测信号来自一个独立训练且冻结的 NEPA 模型，而不是生成器自身的估计。论文通过控制实验显示，相同网络容量下，用 MEP 预训练并冻结的 NEPA 模型作为条件，在该实验设置下优于端到端训练或流匹配预训练的网络特征（见论文表 7）。

实验结果分析

消融实验首先考察 MEP 的损失形式和输入设计。除了 InfoNCE 优于回归损失外，对目标嵌入做归一化会让 FID 从 31.36 升至 34.76，因此论文使用原始目标；给 NEPA 模型输入时间步几乎没有影响，作者推测噪声水平已经反映在噪声补丁中（见论文第 4.1 节）。负样本库从同一图像、批次内和跨设备全量构建时，FID 分别为 31.36、32.13、31.21，论文选择同一图像内的负样本以省去跨设备通信（论文表 5）。补丁大小方面，补丁大小 4 得到 FID 31.00，优于补丁大小 2 的 31.36，且 token 数更少；补丁大小 8 则性能下降至 33.20（论文表 6）。

条件对比实验保持了生成器和训练配方不变，仅替换条件形式。类嵌入的 FID 为 36.39；类 token 加噪声补丁嵌入为 37.21；相同 NEPA 架构但端到端联合训练为 29.52；用流匹配预训练 240 epochs 后冻结为 30.86；本文 MEP 预训练后冻结为 25.04（论文表 7）。这说明条件模型的训练目标对最终生成质量有显著影响。

缩放实验显示，在该评测中，增大生成器和 NEPA 模型都会降低 FID；NEPA-XL 配 DiT-XL 在无引导、64 步 ODE 下达到 11.27，而 NEPA-B 配 DiT-XL 为 16.54（论文表 8）。最终模型结合 REPA 后，使用 96 步 ODE 和 250 步 SDE 分别达到 FID 1.57 和 1.32（论文表 9）。这些结果是在限制引导区间 $t\in[0.4,1]$、引导尺度 3.6 下得到的（论文表 10）。查询率方面，每步查询 NEPA 模型得到 FID 1.57；每 8 步查询为 1.82；只在第一步查询则升至 242.63（论文表 14）。论文报告 NEPA-DiT-XL 总参数约 1.39B，训练计算量约 $3.1\times10^{20}$ FLOPs，约为 SiT-XL/2 加 REPA 800 epochs 计算量的三分之一。

实践建议

在类似类条件图像生成的工程落地中，如果要把固定条件替换为预测嵌入，最直接的做法是复用论文的两阶段流程：先单独训练 NEPA 模型，再冻结 NEPA 模型训练生成器。这样 NEPA 模型训练一次后可以被多个生成器共享，降低整体训练开销。

推理部署时需要权衡每步查询 NEPA 模型的成本和质量。论文实验显示，每 8 步查询一次会让 FID 从 1.57 升至 1.82，而仅第一步查询会明显恶化；因此如果延迟敏感，可以尝试降低查询频率，但应在自身验证集上确认损失幅度。条件前缀的 KV 缓存已经能减少重复计算，但 NEPA 模型本身每步仍有约 92 GFLOPs 的采样开销（论文表 13）。

若复现 MEP 训练，建议在相同数据规模下优先测试 InfoNCE 损失而非回归损失，并关注补丁大小对质量和计算量的影响；本文在补丁大小 4 上取得平衡。分类器自由引导的区间和尺度配置在该实验中经过搜索得到，其他数据集或模型规模下不一定适用，需要重新调参。