何恺明重磅新作：Just image Transformers让去噪模型回归基本功

机器之心报道

编辑：冷猫、+0

大家都知道，图像生成和去噪扩散模型是密不可分的。高质量的图像生成都通过扩散模型实现。

但有没有一种可能，「去噪扩散模型」实际上并没有做到「去噪」？

ResNet 之父，超70万引用的 AI 大神何恺明的新论文敏锐地捕捉了这一现象。

当扩散生成模型最初被提出时 ，核心思想本应是「去噪」，即从被破坏的图像中预测出干净图像。然而，在扩散模型的发展过程中出现的两个重要里程碑，却偏离了直接预测干净图像这一目标。首先，预测噪声本身（称为 「ϵ-prediction」）在生成质量上带来了关键改变，并极大推动了这些模型的普及。随后，扩散模型被连接到基于流的生成方法，通过预测流速度（「-prediction」 ），这一量结合了干净数据与噪声。

流形假设（Manifold Assumption）认为，自然图像在高维像素空间中位于一个低维流形上。干净图像 可以建模为处于流形上（on-manifold），而噪声 ϵ 或流速度 （例如 = − ϵ）则本质上处于流形之外（off-manifold）。因此，让神经网络预测干净图像（即 -prediction）在本质上不同于让其预测噪声或带噪的量（即 ϵ/-prediction）。

当今的去噪扩散模型并不按照经典意义上的「去噪」方式工作，这些神经网络预测的是噪声或带噪的量。

预测干净数据与预测带噪量在本质上是不同的。根据流形假设（manifold assumption），自然数据应当位于一个低维流形上，而带噪的量则不在其上。

基于这一假设，何恺明新论文主张使用直接预测干净数据的模型，这使得看似容量不足的网络也能够在极高维空间中有效运行。

这一研究使去噪扩散模型回归本源，探索一种在原始自然数据上构建基于 Transformer 的扩散模型的自洽范式。

• 论文标题：Back to Basics: Let Denoising Generative Models Denoise
• 论文链接：https://arxiv.org/abs/2511.13720v1
• Github 链接：https://github.com/LTH14/JiT

论文提出「Just image Transformers（JiT）」架构，仅使用简单的大 patch 像素级 Transformer 就能成为强大的生成模型，无需 tokenizer、无需预训练，也无需额外损失项。

该架构在 ImageNet 上实现了有竞争力的像素空间图像生成，在256×256分辨率下达到1.82 FID，在512×512分辨率下达到1.78 FID，且无需外部组件。

本文的工作朝着一种基于原生数据的自洽 「Diffusion + Transformer」 原则 更进一步。超越计算机视觉领域，这种原则在其他涉及自然数据的领域（如蛋白质、分子或天气）中也同样具有潜在价值，因为这些领域的 tokenizer 通常难以设计。通过尽量减少特定领域的建模设计，作者希望源自计算机视觉的通用 「Diffusion + Transformer」 范式能够在更广泛的领域中发挥作用。

用于扩散的「Just image Transformers」

Just image Transformers

ViT 的核心理念是 基于 Patch 的 Transformer（ToP），本研究的架构设计遵循这一哲学。

按照标准做法，该架构以时间 t 和给定的类别标签为条件。研究者使用adaLN-Zero进行条件化处理。从概念上讲，这种架构相当于直接应用于像素 patch 的 Diffusion Transformer (DiT)。

整体架构是 Just image Transformer（纯图像 Transformer），将其简称为JiT。例如，在 256×256 的图像上研究 JiT/16（即 patch 大小 p=16），在 512×512 的图像上研究 JiT/32（p=32）。这些设置分别导致每个 patch 的维度为 768 (16×16×3) 和 3072 (32×32×3)。如此高维的 patch 可以通过 -prediction 来处理。

网络应该预测什么？

表 1 中总结了损失空间和预测空间的九种可能组合。针对每一种组合，研究者训练了一个 Base 模型（JiT-B），其每个 token 的隐藏层大小为 768 维。

研究者在表 2(a) 中研究了 256×256 分辨率下的 JiT-B/16。作为参考，他们在表 2(b) 中检查了 64×64 辨率下的 JiT-B/4（即 p=4）。在这两种设置中，序列长度是相同的 (16×16)。

观察结果如下：

• -prediction 至关重要

在表 2(a) 的 JiT-B/16 中，只有 -prediction 表现良好，并且它在所有三种损失函数下均有效。在这里，一个 patch 是 768 维 (16×16×3)，这与 JiT-B 中 768 的隐藏层大小相吻合。虽然这看起来「刚好足够」，但在实践中，模型可能需要额外的容量，例如用于处理位置嵌入。对于 ϵ-prediction 或 -prediction，模型没有足够的容量来分离并保留噪声量。这些观察结果与 Toy 案例（图 2）中的观察结果相似。

作为对比，研究者检查了 64×64 分辨率下的 JiT-B/4（表 2(b)）。在这里，所有情况都表现得相当不错：九种组合之间的准确率差距是边缘性的，并不具有决定性。每个 patch 的维度是 48 (4×4×3)，远低于 JiT-B 中 768 的隐藏层大小，这解释了为什么所有组合都工作得相当好。研究者注意到，许多之前的潜在扩散模型具有类似的小输入维度，因此未曾暴露在此处讨论的问题中。

• 仅靠损失加权是不够的

该工作并不是第一个枚举相关因素组合的，之前的研究探索了损失加权和网络预测的组合。他们的实验是在低维的 CIFAR-10 数据集上使用 U-net 完成的。他们的观察结果与研究者在 ImageNet 64×64 上的结果更为接近。

然而，表 2(a) 在 ImageNet 256×256 上的结果表明，损失加权并非全部。一方面，无论损失空间如何，ϵ-prediction 或 -prediction 在表 2(a) 中都遭受了灾难性的失败，这对应于不同损失空间中的不同有效加权（如前所述）。另一方面，-prediction 在所有三种损失空间中都有效：虽然由 -loss 引起的损失加权更可取，但并非至关重要。

• 仅靠噪声水平偏移是不够的

先前的工作建议增加噪声水平对于高分辨率的基于像素的扩散是有用的。研究者在表 3 中使用 JiT-B/16 对此进行了检查。由于研究者使用logit-normal 分布来采样 t，可以通过改变该分布的参数 µ 来偏移噪声水平：直观地说，将 µ 向负侧偏移会导致更小的 t，从而增加噪声水平（公式 (1)）。

表 3 显示，当模型已经表现不错时（此处为 x-pred），适当的高噪声是有益的，这与先前的观察结果一致。然而，仅调整噪声水平无法补救 ϵ-prediction 或 -prediction：它们的失败本质上源于无法传播高维信息。

根据表 3，研究者在 ImageNet 256×256 的其他实验中将 µ 设置为 -0.8。

• 增加隐藏单元并非必要

既然容量可能受到网络宽度（即隐藏单元数量）的限制，一个自然的想法是增加它。然而，当观测维度非常高时，这种补救措施既不符合原则也不可行。研究者表明，在 -prediction 的情况下，这是没有必要的。

在表 5 和表 6 中，研究者展示了 JiT/32 在分辨率 512 和 JiT/64 在分辨率 1024 下的结果，使用了成比例的大 patch 尺寸 p=32 或 p=64。这相当于每个 patch 为 3072 维（即 32×32×3）或 12288 维，远大于 B、L 和 H 模型的隐藏层大小。尽管如此，-prediction 仍然运作良好；事实上，除了按比例缩放噪声（例如，在分辨率 512 和 1024 时分别缩放 2× 和 4×）之外，它无需任何修改即可工作。

这一证据表明，网络设计可以在很大程度上与观测维度解耦，就像许多其他神经网络应用中的情况一样。增加隐藏单元的数量可能是有益的（如在深度学习中广泛观察到的那样），但这并非决定性因素。

• 瓶颈结构可能有益

更令人惊讶的是，研究者发现，反过来说，在网络中引入降低维度的瓶颈结构可能是有益的。

具体来说，研究者将线性 patch 嵌入层转换为低秩线性层，方法是用一对瓶颈（但仍是线性）层替换它。第一层将维度降低到 d'，第二层将其扩展到 Transformer 的隐藏层大小。这两层都是线性的，并充当低秩重参数化

图 4 绘制了使用 JiT-B/16（原始 patch 为 768 维）时 FID 与瓶颈维度 d' 的关系。减小瓶颈维度，即使小到 16 维，也不会导致灾难性的失败。事实上，在大范围（32 到 512）内的瓶颈维度可以提高质量，FID 的提升幅度高达约 1.3。

从表征学习的更广阔视角来看，这一观察并非完全出乎意料。引入瓶颈设计通常是为了鼓励学习内在的低维表征。

算法

最终算法采用 -prediction 和 -loss，这对应于表 1(3)(a)。形式上，研究者优化：

算法 1 展示了训练步骤的伪代码，算法 2 是采样步骤的伪代码（欧拉求解器；可扩展到 Heun 或其他求解器）。为简洁起见，类别条件化和 CFG 被省略了，但两者都遵循标准做法。为了防止在计算 1/(1-t) 时出现除以零的情况，我们在计算此除法时会截断其分母（默认值为 0.05）。

更高级的 Transformer

通用 Transformer 的优势部分在于，当其设计与具体任务解耦时，它可以从其他应用中开发的架构进步中受益。这一特性通过使用任务无关的 Transformer 来构建扩散模型，巩固了其优势。

研究者整合了流行的通用改进：SwiGLU、RMSNorm、RoPE、qk-norm，所有这些最初都是为语言模型开发的。

研究者还探索了上下文内类别条件化：但不同于原始 ViT 仅向序列追加一个类别 token，他们追加了多个此类 token（默认为 32 个）。表 4 报告了这些组件的效果。

实验与对比

图 5：定性结果。使用 JiT-H/32 在 ImageNet 512×512 上生成的部分示例。

像素空间中的高分辨率生成

在表 5 中，研究者进一步汇报了基础规模模型（JiT-B）在 ImageNet 上分辨率为 512 以及甚至 1024 时的结果。研究者使用与图像尺寸成比例的 patch 大小，因此在不同分辨率下序列长度保持不变。单个 patch 的维度可以高达 3072 或 12288，而常见模型都没有足够多的隐藏单元来匹配这些维度。

表 5 显示，模型在不同分辨率下都取得了不错的性能。所有模型的参数量和计算成本相近，唯一的差别来自输入 / 输出 patch 的 embedding。本方法不会遭遇观测维度灾难。

可扩展性

将 Transformer 的结构与任务本身解耦的核心目标之一，是利用其潜在的可扩展能力。表 6 给出了四种模型规模在 ImageNet 256 与 512 分辨率下的结果（注意在 512 分辨率下，这些模型都没有隐藏单元数量超过 patch 维度）。模型规模和 FLOPs 在表 7 与 8 中列出：本文的模型在 256 分辨率下的成本与在 512 分辨率下的对应模型相似。

有趣的是，随着模型规模变大，256 与 512 分辨率之间的 FID 差距变得更小。对于 JiT-G，512 分辨率下的 FID 甚至更低。对于在 ImageNet 上的非常大的模型来说，FID 表现在很大程度上取决于过拟合，而 512 分辨率的去噪任务更具挑战性，使其不太容易过拟合。

以往工作的参考结果

作为参考，研究者们在表 7 与表 8 中与过往结果进行了比较，并标注了各方法所涉及的预训练组件。与其他像素空间方法相比，本文的方法完全依赖于朴素的、通用型的 Transformer。本文的模型计算友好，并避免了分辨率加倍时常见的二次计算成本增长（参见表 8 中的 FLOPs）。

表7：在 ImageNet 256×256 上的参考结果。评估指标为 5 万张样本的 FID 和 IS 。「pre-training」 列列出了获得这些结果所需的外部预训练模型。参数量统计包括生成器和 tokenizer 的解码器（推理阶段使用），但不包括其它预训练组件。Giga-flops 的测量基于一次前向传播（不包含 tokenizer），并且大致与训练和推理过程中每次迭代的计算成本成正比。

表 8： 在 ImageNet 512×512 上的参考结果。JiT 采用了更激进的 patch 大小，并能够在较小的计算量下取得强劲的结果。符号说明与表 7 类似。

本文的方法没有使用额外损失或预训练，这些可能会进一步带来性能提升。

总结

主要集中在概率建模的形式上，而对所使用的神经网络的能力（以及局限性）关注较少。然而，神经网络并非具有无限能力，它们更擅长利用自身容量去建模数据，而不是建模噪声。从这一视角来看，关于-prediction的发现是一个自然的结果。

论文的工作采用了一种极简且自洽的设计。通过减少特定领域的归纳偏置，研究者希望这种方法能够推广到那些难以获取 tokenizer 的其他领域。这一性质在涉及原始、高维自然数据的科学应用中尤为重要。

研究者设想，这种通用的 「Diffusion + Transformer」 范式，将有潜力成为其他领域的基础方法之一。

更多信息请参阅原论文。