摘要
https://arxiv.org/pdf/2511.13720
当今的去噪扩散模型并不以经典意义上的"去噪"方式进行工作，也就是说，它们并不直接预测干净的图像。相反，神经网络预测的是噪声或含噪量。在本文中，我们认为预测干净数据和预测含噪量在本质上是不同的。根据流形假设，自然数据应位于低维流形上，而含噪量则不然。基于这一假设，我们提倡使用直接预测干净数据的模型，这使得明显容量不足的网络仍能在非常高维的空间中有效运行。我们证明，简单的、基于大图像块的像素Transformer可以成为强大的生成模型：无需分词器、无需预训练、无需额外损失。我们的方法在概念上不过是"纯图像Transformer"，我们称之为JiT。我们在ImageNet上使用JiT配合大图像块尺寸（16和32）在256和512分辨率下报告了具有竞争力的结果，而在这些分辨率下，预测高维含噪量可能会灾难性地失败。通过让我们的网络回归到流形的基础，我们的研究回归基础，追求一种用于原始自然数据上基于Transformer的扩散的自包含范式。
1. 引言
当扩散生成模型首次被开发时[57,59,23]，其核心思想本应是去噪，即从其损坏版本中预测干净图像。然而，扩散模型演进中的两个重要里程碑最终偏离了直接预测干净图像的目标。首先，预测噪声本身（称为"ε-预测"[23]）在生成质量上产生了关键性差异，并在很大程度上普及了这些模型。后来，扩散模型通过预测流速度（"v-预测"[52]）与基于流的方法[37,38,1]联系起来，流速度是一种结合了干净数据和噪声的量。如今，实际应用中的扩散模型通常预测噪声或含噪量（例如，速度）。
大量研究[52,29,25,15]表明，预测干净图像（"x-预测"[52]）与ε-和v-预测密切相关，前提是重新制定预测损失的权重（详见第3节）。由于这种关系，人们较少关注网络应该直接预测什么，隐含地假设网络能够执行分配给它的任务。
然而，干净图像和含噪量（包括噪声本身）的角色并不等同。在机器学习中，长期以来一直假设[4,3]自然（高维）数据大致位于低维流形上[4, p.61]。根据这一流形假设，虽然干净数据可以被建模为位于低维流形上，但含噪量本质上分布在完整的高维空间中[69]（见图1）。预测干净数据与预测噪声或含噪量在本质上是不同的。
考虑一个低维流形嵌入在高维观察空间中的场景。在此高维空间中预测噪声需要高容量：网络需要保留关于噪声的所有信息。相比之下，容量有限的网络仍然可以预测干净数据，因为它只需要保留低维信息，同时过滤掉噪声。
当使用低维空间（例如，图像潜在空间[49]）时，预测噪声的难度会减轻，但同时也被隐藏而非解决。当涉及到像素或其他高维空间时，现有的扩散模型仍然难以解决维度诅咒[4]。对预训练潜在空间的严重依赖阻碍了扩散模型的自包含性。
为了追求自包含原则，人们一直强烈关注推进像素空间中的扩散建模[7,25,26,6,70]。一般来说，这些方法显式或隐式地避免了网络中的信息瓶颈，例如通过使用密集卷积或更小的图像块、增加通道数或添加长跳跃连接。我们提出，这些设计可能源于预测高维含噪量的需求。
在本文中，我们回归基本原理，让神经网络直接预测干净图像。通过这样做，我们证明了在由原始像素组成的大图像块上运行的普通Vision Transformer（ViT）[13]可以有效地用于扩散建模。我们的方法是自包含的，不依赖于任何预训练或辅助损失------没有潜在分词器[49]、没有对抗损失[16,49]、没有感知损失[77,49]（因此没有预训练分类器[56]），也没有表示对齐[74]（因此没有自监督预训练[45]）。从概念上讲，我们的模型不过是应用于扩散的"纯图像Transformer"，我们称之为JiT。
我们在ImageNet数据集[11]上以256和512分辨率进行实验，分别使用图像块大小为16和32的JiT模型。尽管这些图像块维度非常高（数百或数千），但使用x-预测的我们的模型可以轻松产生强大的结果，而ε-和v-预测则会灾难性地失败。进一步分析表明，网络宽度不必匹配或超过图像块维度；事实上，令人惊讶的是，瓶颈设计甚至可能是有益的，这呼应了经典流形学习中的观察结果。
我们的努力标志着朝着在原生数据上实现自包含的"扩散+Transformer"哲学[46]迈出了一步。超越计算机视觉，这一哲学在涉及自然数据的其他领域（例如蛋白质、分子或天气）非常理想，因为在这些领域设计分词器可能很困难。通过最小化特定领域的设计，我们希望源自计算机视觉的通用"扩散+Transformer"范式将找到更广泛的应用。
2. 相关工作
扩散模型及其预测。关于扩散模型的开创性工作[57]提出学习一个反向随机过程，其中网络预测正态分布的参数（例如，均值和标准差）。在其引入五年后，这一方法被去噪扩散概率模型（DDPM）[23]重新激发并普及：一个关键发现是将噪声作为预测目标（即，ε-预测）。
随后在[52]（最初在模型蒸馏的背景下）研究了不同预测目标之间的关系，其中还引入了v-预测的概念。他们的工作侧重于由重新参数化引入的加权效应。
同时，EDM[29]围绕去噪器函数重新制定了扩散问题，这构成了扩散模型演进中的一个重要里程碑。然而，EDM采用了预条件公式，其中网络的直接输出不是去噪图像。虽然这种公式在低维场景中是可取的，但它仍然本质上要求网络输出一个混合了数据和噪声的量（更多比较见附录）。
流匹配模型[37,38,1]可以被解释为扩散建模框架内的v-预测[52]形式。与纯噪声不同，v是数据和噪声的组合。基于流的模型与先前扩散模型之间的联系已经建立[15]。如今，扩散模型及其基于流的对应物通常在统一框架下进行研究。
去噪模型。在过去的几十年中，去噪的概念与表示学习密切相关。经典方法，以BM3D[9]和其他方法[47,14,79]为例，利用稀疏性和低维度的假设来执行图像去噪。
去噪自编码器（DAEs）[68,69]被开发为一种无监督表示学习方法，使用去噪作为其训练目标。它们利用流形假设[4]（图1）来学习近似低维数据流形的有意义表示。DAEs可以被视为一种去噪分数匹配[67]形式，而后者又与现代基于分数的扩散模型[59,60]密切相关。然而，虽然DAEs为流形学习自然地预测干净数据，但在分数匹配中，预测分数函数实际上等同于预测噪声（按比例因子缩放），即ε-预测。
流形学习。流形学习一直是一个经典领域[51,63]，专注于从观察数据中学习低维、非线性表示。一般来说，流形学习方法利用瓶颈结构[64,48,41,2]，鼓励只有有用信息通过。几项研究探索了流形学习与生成模型之间的联系[39,27,8]。潜在扩散模型（LDMs）[49]可以通过自编码器在第一阶段进行流形学习，然后在第二阶段进行扩散。
像素空间扩散 。虽然潜在扩散[49]已成为当今该领域的默认选择，但扩散模型的发展最初始于像素空间公式[59,23,44,12]。早期的像素空间扩散模型通常基于密集卷积网络，最常见的是U-Net[50]。这些模型通常使用过度完整的通道表示（例如，在第一层将H×W×3H{\times}W{\times}3H×W×3转换为H×W×128H{\times}W{\times}128H×W×128），并伴有长距离跳跃连接。虽然这些模型对ε-和v-预测效果很好，但它们的密集卷积结构通常计算成本很高。将这些卷积模型应用于高分辨率图像不会导致灾难性退化，因此这一方向的研究通常集中在噪声调度和/或加权方案[7,25,26,30]上，以进一步提高生成质量。
相比之下，将Vision Transformer（ViT）[13]直接应用于像素是一项更具挑战性的任务。标准ViT架构采用激进的图像块大小（例如，16×16像素），导致高维的令牌空间，可以与Transformer的隐藏维度相当或更大。SiD2[26]和PixelFlow[6]采用从较小图像块开始的分层设计；然而，这些模型是"FLOP-heavy"[26]，失去了标准Transformer固有的通用性和简单性。PixNerd[70]采用NeRF头[43]，整合来自Transformer输出、噪声输入和空间坐标的更多信息，并通过表示对齐[74]进一步辅助训练。
即使有这些特殊用途的设计，这些工作中的架构通常从"L"（大型）或"XL"大小开始。事实上，最近的一项工作[73]表明，对于高维度，似乎需要大的隐藏大小。
高维扩散。当使用ViT风格架构时，现代扩散模型仍然面临高维输入空间的挑战，无论是在像素还是潜在空间中。在文献[8,73,55]中，反复报告说ViT风格扩散模型在每个令牌维度增加时会迅速且灾难性地退化，无论使用像素还是潜在空间。
与我们的工作同时，一系列研究[78,34,55]诉诸于自监督预训练来解决高维扩散问题。与这些努力相比，我们证明高维扩散可以在没有任何预训练的情况下实现，仅使用Transformer。
x-预测。x-预测的公式是自然的并且不是新的；它至少可以追溯到原始DDPM[23]（参见他们的代码[24]）。然而，DDPM观察到ε-预测明显更好，后来成为首选解决方案。在后来的工作中（例如[26]），虽然分析有时更倾向于在x空间中进行，但实际预测通常在其他空间中进行，可能出于历史原因。
当涉及到由扩散解决的图像恢复应用[10,72,42]时，网络自然地预测干净数据，因为这是图像恢复的最终目标。与我们的工作同时，[18]也提倡使用x-预测，用于以先前帧为条件的生成世界模型。
我们的工作并非旨在重新发明这一基本概念；相反，我们旨在引起人们对高维数据中一个被很大程度上忽视但关键问题的关注，这些数据具有底层低维流形。
3. 扩散模型的预测输出
扩散模型可以在x、ε或v空间中公式化。空间的选择不仅决定了损失定义的位置，还决定了网络预测的内容。重要的是，损失空间和网络输出空间不必相同。这一选择会产生关键差异。
3.1. 背景：扩散和流
扩散模型可以从ODE的角度[5,60,37,58]公式化。我们从基于流的范式[37,38,1]开始我们的公式，即在v空间中，作为更简单的起点，然后讨论其他空间。
考虑数据分布x∼pdata(x)\pmb{x}\sim p_{\mathrm{data}}(\pmb{x})x∼pdata(x)和噪声分布ϵ∼pnoise(ϵ) (e.g.,ϵ∼N(0,I))\boldsymbol{\epsilon}\sim p_{\mathrm{noise}}(\boldsymbol{\epsilon})\:(e.g.,\boldsymbol{\epsilon}\sim\mathcal{N}(0,\mathbf{I}))ϵ∼pnoise(ϵ)(e.g.,ϵ∼N(0,I))。在训练期间，噪声样本ztz_{t}zt是一个插值：zt = at x + bt ϵ\pmb{z}{t}\ =\ a{t}\,\pmb{x}\:+\:b_{t}\,\boldsymbol{\epsilon}zt = atx+btϵ，其中ata_{t}at和btb_tbt是在时间t∈[0,1]t\in[0,1]t∈[0,1]预定义的噪声调度。在本文中，我们使用线性调度[37,38,1]：at=t 和 bt=1−ta_{t}=t\;\mathrm{和}\;b_{t}=1-tat=t和bt=1−t。这给出：
zt=t x+(1−t) ϵ,\boldsymbol{z}_{t}=t\,\boldsymbol{x}+(1-t)\,\boldsymbol{\epsilon},zt=tx+(1−t)ϵ,
当t=1t{=}1t=1时，这导致zt ∼ pdataz_{t}~\sim~p_{\mathrm{data}}zt ∼ pdata。我们对t使用对数正态分布[15]，即logit(t) N(μ,σ2)\mathcal{N}(\mu,\sigma^{2})N(μ,σ2)。
流速度v定义为z对时间的导数，即vt=zt′=at′x+bt′ϵ\mathbf{\boldsymbol{v}}{t}=\mathbf{\boldsymbol{z}}{t}^{\prime}=a_{t}^{\prime}\mathbf{\boldsymbol{x}}+b_{t}^{\prime}\mathbf{\boldsymbol{\epsilon}}vt=zt′=at′x+bt′ϵ。给定公式(1)，我们有：
v=x−ϵ.\boldsymbol{v}=\boldsymbol{x}-\boldsymbol{\epsilon}.v=x−ϵ.
基于流的方法[37,38,1]最小化定义为以下形式的损失函数：
L=Et,x,ϵ∥vθ(zt,t)−v∥2,\mathcal{L}=\mathbb{E}{t,\pmb{x},\pmb{\epsilon}}\|\pmb{v}{\theta}(\pmb{z}_{t},t)-\pmb{v}\|^{2},L=Et,x,ϵ∥vθ(zt,t)−v∥2,
其中vθv_θvθ是由θ参数化的函数。虽然vθv_θvθ通常是网络vθ = netθ(zt,t) [37,38,1]\pmb{v}{\theta}\,=\,\mathtt{n e t}{\theta}(\pmb{z}_{t},t)\,[37,38,1]vθ=netθ(zt,t)[37,38,1]的直接输出，但它也可以是它的转换，正如我们将详细阐述的。
给定函数vθ⋅\boldsymbol{v}{\boldsymbol{\theta}\cdot}vθ⋅，采样是通过求解z的常微分方程(ODE)[37,38,1]完成的：
dzt/dt=vθ(zt,t),d\boldsymbol{z}{t}/d t=\boldsymbol{v}{\boldsymbol{\theta}}(\boldsymbol{z}{t},t),dzt/dt=vθ(zt,t),
从z0 ∼ pnoisez_{0}\;\sim\;p_{\mathrm{noise}}z0∼pnoise开始，到t = 1t\,=\,1t=1结束。在实践中，可以使用数值求解器近似求解此ODE。默认情况下，我们使用50步Heun求解器。
1 我们在本文中的分析适用于其他调度。
3.2. 预测空间与损失空间
预测空间。网络的直接输出可以在任何空间中定义：v、x或ε。接下来，我们讨论由此产生的转换。请注意，在本文的上下文中，我们仅在当网络netθ的直接输出严格是x、ε、v时，才将其称为"x、ε、v-预测"。
给定三个未知数(x、ε、v)和一个网络输出，我们需要两个额外约束来确定所有三个未知数。这两个约束由公式(1)和(2)给出。例如，当我们让直接网络输出netθ{\mathtt{n e t}}{\theta}netθ为x时，我们求解以下方程组：
{xθ=netθzt=t xθ+(1−t) ϵθvθ=xθ−ϵθ\left\{\begin{aligned}{\mathbf{\boldsymbol{x}}{\theta}}&{{}=\mathtt{n e t}{\theta}}\\ {\mathbf{\mathit{z}}{t}}&{{}=t\,\mathbf{\boldsymbol{x}}{\theta}+(1-t)\,\mathbf{\boldsymbol{\epsilon}}{\theta}}\\ {\mathbf{\boldsymbol{v}}{\theta}}&{{}=\mathbf{\boldsymbol{x}}{\theta}-\mathbf{\boldsymbol{\epsilon}}_{\theta}}\end{aligned}\right.⎩ ⎨ ⎧xθztvθ=netθ=txθ+(1−t)ϵθ=xθ−ϵθ
这里，符号xθx_θxθ、εθε_θεθ和vθv_θvθ表明它们都是依赖于θ的预测。求解此方程组给出：KaTeX parse error: Undefined control sequence: \mathbfit at position 29: ...lon}{\theta}=(\̲m̲a̲t̲h̲b̲f̲i̲t̲{z{t}-}t\mathb...和vθ=(xθ − zt)/(1 − t)\mathbf{\boldsymbol{v}}{\theta}=(\mathbf{\boldsymbol{x}}{\theta}\!-\!\mathbf{\boldsymbol{z}}_{t})/(1\!-\!t)vθ=(xθ−zt)/(1−t)，即εθε_θεθ和vθv_θvθ都可以从ztz_tzt和网络xθx_θxθ计算得出。这些总结在表1的列(a)中。
类似地，当我们让直接网络输出net为ε或v时，我们得到其他方程组（通过替换公式(5)中的第一个方程）。转换在表1的列(b)、©中总结，分别用于ε-、v-预测。这表明当{x,ϵ,v}\{\boldsymbol{x},\boldsymbol{\epsilon},\boldsymbol{v}\}{x,ϵ,v}中的一个量被预测时，其他两个可以被推断。许多先前工作（例如[52,15]）中的推导是表1中涵盖的特殊情况。
损失空间。虽然损失通常在一个参考空间中定义（例如，公式(3)中的v-损失），但从概念上讲，可以在任何空间中定义它。已经证明[52,15]，给定从一个预测空间到另一个的重新参数化，损失实际上被重新加权。
例如，考虑表1(3)(a)中x-预测和v-损失的组合。我们有vθ=(xθ−zt)/(1−t)\mathbf{\boldsymbol{v}}{\theta}=(\mathbf{\boldsymbol{x}}{\theta}-\mathbf{\boldsymbol{z}}{t})/(1-t)vθ=(xθ−zt)/(1−t)作为预测，v = (x−zt)/(1−t)\pmb{v}\,=\,(\pmb{x}-\pmb{z{t}})/(1-t)v=(x−zt)/(1−t)作为目标。公式(3)中的v-损失变为：L=E∥vθ(zt,t)−v∥2=\mathcal{L}=\mathbb{E}\|\pmb{v}{\theta}(\pmb{z}{t},t)-\pmb{v}\|^{2}=L=E∥vθ(zt,t)−v∥2= E1(1−t)2∥xθ(zt,t) −x∥2\begin{array}{r}{\mathbb{E}\frac{1}{(1-t)^{2}}\|\pmb{x}{\theta}(\pmb{z}{t},t)\!-\pmb{x}\|^{2}}\end{array}E(1−t)21∥xθ(zt,t)−x∥2，这是x-损失的重新加权形式。可以为表1中列出的任何预测空间和任何损失空间进行这样的转换。
综合起来，考虑在{x,ϵ,v}\{\boldsymbol{x},\boldsymbol{\epsilon},\boldsymbol{v}\}{x,ϵ,v}中定义的三个未加权损失，以及三种形式的网络直接输出，有九种可能的组合（表1）。每种组合构成一个有效的公式，九种情况中没有两种在数学上等价。
生成器空间。无论使用哪种组合，在推理时执行生成时，我们总是可以将网络输出转换到v空间（表1，行(3)）并求解公式(4)中的ODE进行采样。因此，所有九种组合都是合法的生成器。
3.3. 玩具实验
根据流形假设[4]，数据x倾向于位于低维流形上（图1），而噪声ε和速度v则不在流形上。让网络直接预测干净数据x应该更易处理。我们在本节的玩具实验中验证这一假设。
我们考虑d维底层数据"埋藏"在D维空间(d<D)(d<D)(d<D)中的玩具情况。我们使用投影矩阵P∈RD×d˙P\in \mathbb{R}^{D\times\dot{d}}P∈RD×d˙合成这一场景，该矩阵是列正交的，即P⊤P = Id×dP^{\top}P\,=\,I_{d\times d}P⊤P=Id×d。这个矩阵PPP是随机创建并固定的。观察到的数据是x = Px^ ∈ RD\pmb{x}\,=\,P\hat{\pmb{x}}\,\in\,\mathbb{R}^{D}x=Px^∈RD，其中底层数据是x∈Rdx \in\mathbb{R}^{d}x∈Rd。矩阵PPP对模型是未知的，因此对模型来说，这是一个D维生成问题。
我们训练一个具有256维隐藏单元的5层ReLU MLP作为生成器，并在图2中可视化结果。我们通过使用PPP将D维生成样本投影回d维来获得这些可视化。我们研究D ∈ {2,8,16,512}D\,\in\,\{2,8,16,512\}D∈{2,8,16,512}的情况，其中d = 2d\,=\,2d=2。我们研究x、ε或v-预测，都使用v-损失，即，表1(3)(a-c)。
图2显示，只有x-预测能在D增加时产生合理结果。对于ε-/v-预测，模型在D=16D{=}16D=16时就已挣扎，当D=512D{=}512D=512时灾难性失败，此时256维MLP是不完整的。
值得注意的是，x-预测即使在模型不完整时也能很好地工作。在这里，256维MLP不可避免地会丢弃D=512−dimD{=}512{\mathrm{-dim}}D=512−dim空间中的信息。然而，由于真实数据位于低维d维空间中，x-预测仍然可以表现良好，因为理想输出在隐式上是d维的。我们在ImageNet上的真实数据案例中得出类似观察，如下所示。
4. 用于扩散的"纯图像Transformer"
基于上述分析，我们证明在像素上运行的普通Vision Transformers（ViT）[13]可以出奇地好，只需使用x-预测。
4.1. 纯图像Transformer
ViT[13]的核心思想是"在图像块上的Transformer"（ToP）。我们的架构设计遵循这一哲学。
形式上，考虑H×W×CH{\times}W{\times}CH×W×C维图像数据(C=3)(C{=}3)(C=3)。所有x、ε、v和ztz_tzt共享这一相同维度。给定图像，我们将其划分为大小为p×pp\times pp×p的非重叠图像块，产生长度为Hp×Wp{\frac{H}{p}}\times{\frac{W}{p}}pH×pW的序列。每个图像块是一个p×p×3p{\times}p{\times}3p×p×3维向量。该序列由线性嵌入投影处理，添加位置嵌入[66]，并由Transformer块堆栈[66]映射。输出层是一个线性预测器，将每个令牌投影回p×p×3p\times p\times 3p×p×3维图像块。见图3。
按照标准实践，架构以时间t和给定类别标签为条件。我们使用adaLN-Zero[46]进行条件设置，稍后将讨论其他选项。从概念上讲，此架构相当于直接应用于像素图像块的扩散Transformer（DiT）[46]。
整体架构不过是"纯图像Transformer"，我们称之为JiT。例如，我们研究JiT/16（即，图像块大小p=16p{=}16p=16[13])在256×256图像上，以及JiT/32(p=32)(p{=}32)(p=32)在512 × 512512\!\times\!512512×512图像上。这些设置分别导致每个图像块768维(16 × 16 × 3)(16\!\times\!16\!\times\!3)(16×16×3)和3072维(32×32×3)。这种高维图像块可以由x-预测处理。
4.2. 网络应预测什么？
我们已在表1中总结了损失空间和预测空间的九种可能组合。对于这些组合中的每一种，我们训练一个'基础'[13]模型（JiT-B），每个令牌具有768维隐藏大小。我们在表2(a)中研究256×256分辨率下的JiT-B/16。作为参考，我们在表2(b)中检查JiT-B/4（即，p=4p{=}4p=4)在64×64分辨率下。在这两种设置中，序列长度相同(16×16)。
我们得出以下观察：
x-预测至关重要。在表2(a)的JiT-B/16中，只有x-预测表现良好，它在所有三种损失中都有效。这里，图像块是768-d(16×16×3)，与JiT-B中768的隐藏大小巧合。虽然这可能看起来"足够"，但在实践中，模型可能需要额外容量，例如处理位置嵌入。对于ε-/v-预测，模型没有足够容量来分离和保留含噪量。这些观察与玩具案例（图2）中的观察相似。
作为比较，我们在64×64分辨率下检查JiT-B/4（表2(b)）。在这里，所有情况都表现相当好：九种组合之间的准确度差距很小，不是决定性的。每个图像块的维度是48(4×4×3)，远低于JiT-B中768的隐藏大小，这解释了为什么所有组合都能合理工作。我们注意到，许多先前的潜在扩散模型具有类似的低输入维度，因此没有暴露于我们在此讨论的问题。
损失加权不足以解决问题。我们的工作不是首次枚举相关因素的组合。在[52]中，它探索了损失加权和网络预测的组合。他们的实验是在低维CIFAR-10数据集上使用U-Net进行的。他们的观察更接近我们在ImageNet 64×64上的观察。
然而，表2(a)在ImageNet 256×256上表明，损失加权不是全部。一方面，在表2(a)中，无论损失空间如何，ε-和v-预测都会灾难性失败，这对应于不同损失空间中的不同有效加权（如所讨论的）。另一方面，x-预测在所有三种损失空间中都有效：由v-损失引起的损失加权是可取的，但不是关键的。
噪声级移位不足以解决问题。先前工作[7,25,26]表明，增加噪声级对高分辨率像素基础扩散有益。我们在表3中使用JiT-B/16对此进行检查。由于我们使用logit-正态分布[15]对t进行采样（见附录），可以通过更改此分布的参数μ来移位噪声级：直观上，将μ向负方向移位会导致更小的t，从而增加噪声级（公式(1)）。
表3显示，当模型已经表现良好时（这里，x-预测），适当高的噪声是有益的，这与先前观察[7,25,26]一致。然而，仅调整噪声级无法弥补ε-或v-预测：它们的失败本质上源于无法传播高维信息。
顺带一提，根据表3，我们在ImageNet 256×256的其他实验中设置μ=−0.8\mu=-0.8μ=−0.8。
增加隐藏单元不是必需的 。由于容量可能受到网络宽度（即隐藏单元数量）的限制，一个自然的想法是增加它。然而，当观察维度非常高时，这一补救措施既不原则化也不可行。我们证明，在x-预测的情况下，这不是必需的。
在下一节的表5和表6中，我们展示了在512分辨率下使用JiT/32和在1024分辨率下使用JiT/64的结果，使用成比例的大图像块大小p=32p{=}32p=32或p=64p{=}64p=64。这相当于每个图像块3072维（即，32×32×332{\times}32{\times}332×32×3)或12288维，远大于B、L和H模型的隐藏大小（在[13]中定义）。尽管如此，x-预测工作得很好；事实上，除了按比例缩放噪声外，它无需任何修改就能工作KaTeX parse error: Undefined control sequence: \ensuremath at position 7: (e.g.,\̲e̲n̲s̲u̲r̲e̲m̲a̲t̲h̲{\mathsf{b y}}\...和4×4\times4×在分辨率512和1024；见附录)。
这一证据表明，网络设计可以在很大程度上与观察维度解耦，就像在许多其他神经网络应用中一样。增加隐藏单元数量可能是有益的（正如深度学习中广泛观察到的），但它不是决定性的。
瓶颈可能有益。更令人惊讶的是，我们发现，相反，引入减少网络中维度的瓶颈可能是有益的。
具体来说，我们将线性图像块嵌入层转换为低秩线性层，方法是将其替换为一对瓶颈（但仍然是线性的）层。第一层将维度减少到d′d'd′，第二层将其扩展到Transformer的隐藏大小。两层都是线性的，作为低秩重新参数化。
图4绘制了FID与瓶颈维度d′d'd′的关系，使用JiT-B/16（每原始图像块768-dim）。即使将瓶颈维度减少到16-d这样小的值，也不会导致灾难性失败。事实上，跨广泛范围(32到512)的瓶颈维度可以提高质量，幅度相当大，最高可达~1.3 FID。
从表示学习的更广泛角度来看，这一观察并不完全出乎意料。通常引入瓶颈设计以鼓励学习本质上低维的表示[64,48,41,2]。
`
4.3. 我们的算法
我们的最终算法采用x-预测和v-损失，对应于表1(3)(a)。形式上，我们优化：
L=Et,x,ϵ∥vθ(zt,t)−v∥2,\mathcal{L}=\mathbb{E}{t,\boldsymbol{x},\boldsymbol{\epsilon}}\Big\|\boldsymbol{v}{\theta}(\boldsymbol{z}{t},t)-\boldsymbol{v}\Big\|^{2},L=Et,x,ϵ vθ(zt,t)−v 2,
其中:vθ(zt,t)=(netθ(zt,t)−zt)/(1−t).\mathrm{其中}:\quad\mathbf{\boldsymbol{v}}{\theta}(\boldsymbol{z}{t},t)=(\mathtt{n e t}{\theta}(\boldsymbol{z}{t},t)-\boldsymbol{z}{t})/(1-t).其中:vθ(zt,t)=(netθ(zt,t)−zt)/(1−t).
算法1显示了训练步骤的伪代码，算法2显示了采样步骤（欧拉求解器；可以扩展到Heun或其他求解器）。为简洁起见，省略了类别条件和CFG[22]，但两者遵循标准实践。为防止1/(1−t)1/(1{-}t)1/(1−t)中的零除，我们剪裁其分母（默认情况下，0.05），无论何时计算此除法。
4.4. "仅高级"Transformer
通用Transformer[66]的优势部分在于，当其设计与特定任务解耦时，它可以受益于在其他应用中开发的架构进步。这一属性支撑了使用任务无关Transformer公式化扩散的优势。
遵循[73]，我们合并流行的通用改进：SwiGLU[54]、RMSNorm[75]、RoPE[62]、qk-norm[19]，所有这些最初都是为语言模型开发的。我们还探索了上下文内类别条件：但与原始ViT[13]将一个类别令牌附加到序列不同，我们附加多个此类令牌（默认情况下，32个；见附录），遵循[35]。表4报告了这些组件的效果。
5 我们先前部分中的基线都使用SwiGLU + RMSNorm。移除它们略有退化：FID从7.48上升到7.89。
5. 比较
像素上的高分辨率生成 。在表5中，我们进一步报告了在ImageNet上256、512甚至1024分辨率下我们基础大小模型（JiT-B）的结果。我们使用与图像大小成比例的图像块大小，因此不同分辨率下的序列长度保持相同。每个图像块的维度可以高达3072或12288，没有常见模型会有足够的隐藏单元。
表5显示，我们的模型在各种分辨率下表现良好。所有模型具有相似的参数数量和计算成本，仅在输入/输出图像块嵌入上有所不同。我们的方法不受观察维度诅咒的影响。
可扩展性 。将Transformer设计与任务解耦的核心目标是利用可扩展性潜力。表6提供了在ImageNet 256和512上使用四种模型大小的结果（注意在512分辨率下，这些模型中没有一个的隐藏单元多于图像块维度）。模型大小和浮点运算次数显示在表7和8中：我们在256分辨率下的模型与其在512分辨率下的对应物具有相似成本。我们的方法受益于扩展。
有趣的是，随着模型变大，256和512分辨率之间的FID差异变小。对于JiT-G，在512分辨率下的FID甚至更低。对于ImageNet上非常大的模型，FID性能在很大程度上取决于过拟合，而在512分辨率下进行去噪是一项更具挑战性的任务，使其不易过拟合。
来自先前工作的参考结果。作为参考，我们在表7和8中与先前结果进行比较。我们标记了每种方法涉及的预训练组件。与其他基于像素的方法相比，我们的方法纯粹由普通、通用Transformer驱动。我们的模型计算友好，避免了分辨率加倍时二次扩展的昂贵计算（见表8中的浮点运算次数）。
我们的方法不使用额外损失或预训练，这可能会带来进一步收益（附录中有一个示例）。这些方向留待未来工作。
6. 讨论与结论
噪声与自然数据本质不同。多年来，扩散模型的发展主要集中在概率公式上，而较少关注所使用神经网络的能力（和限制）。然而，神经网络的能力不是无限的，它们能更好地利用其能力来建模数据而非噪声。从这些角度来看，我们关于x-预测的发现，事后看来，是一个自然结果。
我们的工作采用极简主义和自包含设计。通过减少特定领域的归纳偏差，我们希望我们的方法可以推广到难以获取分词器的其他领域。这一特性对于涉及原始高维自然数据的科学应用特别理想。我们设想，通用"扩散+Transformer"范式将成为其他领域的潜在基础。
附录
A. 实现细节
我们的实现紧密遵循DiT[46]和SiT[40]的公共代码库。我们的配置在表9中总结。我们描述细节如下。
时间分布 。遵循[15]，在训练期间，我们对t采用logit-正态分布[15]：logit(t) ∼N(μ,σ2)\sim \mathcal{N}(\mu,\sigma^{2})∼N(μ,σ2)。具体来说，我们采样s ∼ N(μ,σ2)s\,\sim\,\mathcal{N}(\mu,\sigma^{2})s∼N(μ,σ2)并让t=sigmoid(s)。超参数μ确定噪声级（见表3），默认情况下，我们在ImageNet 256（或512,1024）分辨率上设置μ=−0.8\mu=-0.8μ=−0.8，并将σ固定为0.8。
ImageNet 512×512和1024×1024 。我们在ImageNet 512 × 512512\!\times\!512512×512上采用JiT/32（即，图像块大小为32）。该模型导致256=16×16个图像块的序列，与ImageNet 256 × 256256\!\times\!256256×256上的JiT/16相同。因此，JiT/32仅在输入/输出图像块维度上与JiT/16不同，从每个图像块768-d增加到3072-d；所有其他计算和成本完全相同。
为了在512×512图像上重用ImageNet 256 × 256256\!\times\!256256×256的确切相同配方，我们将噪声ε的幅度按2×2\times2×重新缩放：即，ϵ∼N(0,22I){\boldsymbol{\epsilon}}\sim{\mathcal{N}}(0,2^{2}\mathbf{I})ϵ∼N(0,22I)。这一简单修改近似保持了256×256和512×512分辨率之间的信噪比（SNR）[25,7,26]。不需要或应用对ImageNet 256×256配置的其他更改。
对于ImageNet 1024 × 10241024\!\times\!10241024×1024，我们使用模型JiT/64并将噪声ε按4×缩放。不需要其他更改。
上下文内类别条件。标准DiT[46]通过adaLN-Zero执行类别条件。在表4中，我们进一步探索上下文内类别条件。
具体来说，遵循ViT[13]，可以在图像块序列前添加一个类别令牌。这在DiT[46]中被称为"上下文内条件"。请注意，我们与默认的adaLN-Zero条件联合使用上下文内条件，与DiT不同。此外，遵循MAR[35]，我们在序列前进一步添加多个此类令牌。这些令牌是同一类别令牌的重复实例，添加了不同的位置嵌入。我们添加32个此类令牌。此外，我们发现将它们添加到Transformer输入的后面块可能更有益（见表9中的"上下文内起始块"）。表4显示，我们实现的上下文内条件将FID提高了~1.2。
Dropout和提前停止。我们在JiT-H和G中应用dropout[61]以减轻过拟合风险。具体来说，我们将dropout应用于Transformer块的中间一半。对于具有dropout的Transformer块，我们在注意力块和MLP块中都应用它。
由于G大小模型在我们当前的dropout设置下仍然倾向于过拟合，当监控的FID开始退化时，我们应用提前停止。这对于JiT-G/16和JiT-G/32都发生在约320个周期。
EMA和CFG。我们的研究涵盖广泛的配置，包括损失和预测空间的变化、模型大小和架构组件。CFG比例[22]和EMA（指数移动平均）衰减的最佳值因情况而异，固定它们可能导致不完整或误导性观察。由于维护这些超参数配置相对便宜，我们努力采用最佳值。
具体来说，对于CFG比例ω[22]，我们通过在推理时搜索一系列候选比例来确定最佳值，这是现有工作中的常见做法。对于EMA衰减，我们在训练期间维护移动平均参数的多个副本，这带来了可忽略的计算开销。为避免内存开销，不同的EMA副本可以存储在不同设备上（例如，GPU）。因此，CFG比例和EMA衰减本质上都是推理时决策。
我们的CFG比例候选范围从1.0到4.0，步长为0.1。CFG的影响在图6中显示，适用于JiT-L/16模型。我们的EMA衰减候选为0.9996、0.9998和0.9999，使用1024的批量大小进行评估。对于每个模型（包括消融中的任何模型），我们使用8K样本搜索最佳设置，然后应用它来评估50K样本。
评估。遵循常见做法，我们针对ImageNet训练集评估FID[21]。我们使用50K生成图像评估FID，每个1000个ImageNet类别有50个样本。我们在相同的50K图像上评估Inception Score(IS)[53]。
B. 附加实验
B.1. 训练损失和去噪图像
在表2(a)中，ε-/v-预测的失败是由从有限容量网络预测高维输出的固有能力不足引起的。这一失败可以从训练损失曲线中看出。
在图7（顶部），我们在相同的v-损失下比较训练损失曲线，定义为L=E∥vθ(zt,t)−v∥2\mathcal{L}=\mathbb{E}\|\pmb{v}{\theta}(\pmb{z}{t},t)\pmb{-v}\|^{2}L=E∥vθ(zt,t)−v∥2，使用v-预测(i.e., vθ=netθ)(i.e.,\,v_{\theta}={\tt n e t}{\theta})(i.e.,vθ=netθ)与x-预测（即，vθ=\boldsymbol{v}{\boldsymbol{\theta}}=vθ= (netθ − zt)/(1 − t))(\mathtt{n e t}{\theta}\!-\!z{t})/(1\!-\!t))(netθ−zt)/(1−t))。由于损失在同一空间中计算，且只有参数化不同，比较损失值是合法的。
图7（顶部）显示，v-预测产生明显更高的损失值（约25%）比x-预测，即使v-预测似乎是v-损失的"原生"参数化。这一比较表明x-预测的任务本质上更容易，因为数据位于低维流形上。我们还观察到ε-预测（此处未显示）的损失高约3倍且不稳定，这可能解释了其在表2(a)中的失败。
图7（底部）比较了对应于两条训练曲线的去噪图像。对于x-预测，去噪图像是简单的xθ=netθ\boldsymbol{x}{\boldsymbol{\theta}}=\mathtt{n e t}{\boldsymbol{\theta}}xθ=netθ；对于v-预测，去噪图像是xθ=(1−t)netθ+zt\mathbf{\boldsymbol{x}}{\theta}{=}(1{-}t)\mathtt{n e t}{\theta}{+}\mathbf{\boldsymbol{z}}{t}xθ=(1−t)netθ+zt，其中vθ=netθ\pmb{v{\theta}}{=}\mathtt{n e t_{\theta}}vθ=netθ（见表1©(1)）。图7（顶部）中v-预测的较高损失可以通过图7（底部）中其明显的伪影反映出来。
请注意，图7（底部）中的伪影是单个去噪步骤的伪影。在生成过程中，此错误可以在多步ODE求解器中累积，导致表2(a)中的灾难性失败。
B.2. 预条件器
在EDM[29]中，额外的"预条件器"应用于包装网络的直接输出。使用我们论文的符号，预条件器公式可以写为：KaTeX parse error: Undefined control sequence: \mathbfit at position 34: ...l{x}}{\theta}(\̲m̲a̲t̲h̲b̲f̲i̲t̲{z}{t},t)=c_{\...·netξθ(zt,t){\mathbf\xi}{\theta}({\boldsymbol z}{t},t)ξθ(zt,t)，其中cskipc_{\mathrm{skip}}cskip和coutc_{\mathrm{out}}cout是预定义系数。此方程表明，除非cskip≡0c_{\mathrm{skip}}\equiv0cskip≡0，否则预条件器中的网络输出不得执行x-预测。根据流形假设，这一公式不应弥补我们考虑的问题，如下所述。
预条件器的公式 。给定预条件器的定义，我们可以将"预条件x-预测"重写为：
{xθ=cskip⋅zt+cout⋅netθϵθ=(zt−txθ)/(1−t)vθ=(xθ−zt)/(1−t)\left\{\begin{aligned}{\mathbf{\boldsymbol{x}}{\theta}}&{{}=c{\mathrm{skip}}\cdot\mathbf{\boldsymbol{z}}{t}+c{\mathrm{out}}\cdot\mathtt{n e t}{\theta}}\\ {\mathbf{\boldsymbol{\epsilon}}{\theta}}&{{}=(\mathbf{\boldsymbol{z}}{t}-t\mathbf{\boldsymbol{x}}{\theta})/(1-t)}\\ {\mathbf{\boldsymbol{v}}{\theta}}&{{}=(\mathbf{\boldsymbol{x}}{\theta}-\mathbf{\boldsymbol{z}}_{t})/(1-t)}\end{aligned}\right.⎩ ⎨ ⎧xθϵθvθ=cskip⋅zt+cout⋅netθ=(zt−txθ)/(1−t)=(xθ−zt)/(1−t)
相应地，公式(6)中的目标（v-损失）写为：
L=Et,x,ϵ∥vθ(zt,t)−v∥2,where : vθ(zt,t)=(xθ(zt,t)−zt)/(1−t)andxθ(zt,t)=cskip⋅zt+cout⋅netθ(zt,t).\begin{array}{r}{\mathcal{L}=\mathbb{E}{t,\boldsymbol{x},\boldsymbol{\epsilon}}\Big\|\boldsymbol{v}{\theta}(\boldsymbol{z}{t},t)-\boldsymbol{v}\Big\|^{2},\qquad\qquad\qquad\qquad\qquad}\\ {\mathrm{where}\colon\ \ \boldsymbol{v}{\theta}(\boldsymbol{z}{t},t)=(\boldsymbol{x}{\theta}(\boldsymbol{z}{t},t)-\boldsymbol{z}{t})/(1-t)\qquad\qquad\qquad}\\ {\mathrm{and}\quad\boldsymbol{x}{\theta}(\boldsymbol{z}{t},t)=c_{\mathrm{skip}}\cdot\boldsymbol{z}{t}+c{\mathrm{out}}\cdot\mathrm{net}{\theta}(\boldsymbol{z}{t},t).}\end{array}L=Et,x,ϵ vθ(zt,t)−v 2,where: vθ(zt,t)=(xθ(zt,t)−zt)/(1−t)andxθ(zt,t)=cskip⋅zt+cout⋅netθ(zt,t).
与公式(6)相比，唯一区别在于我们如何从网络计算xθ\pmb{x}_{\theta}xθ。
由于EDM[29]使用方差爆炸调度（即，zt=x+σtϵ)\pmb{z_{t}}{=}\pmb{x}{+}\sigma_{t}\pmb{\epsilon})zt=x+σtϵ)），而我们使用（大致）方差保持版本（即，zt=tx+(1−t)ϵ)\scriptstyle z_{t}=t x+(1-t)\epsilon)zt=tx+(1−t)ϵ)，完全等效转换是不可能的。为了在我们的情况下拥有预条件器，我们将我们的版本重写为1txt=x+1−ttϵ\textstyle{\frac{1}{t}}{\pmb x}{t}={\pmb x}+{\frac{1-t}{t}}{\pmb\epsilon}t1xt=x+t1−tϵ。因此，我们设置σt = 1−tt\begin{array}{r}{\sigma{t}\;=\;\frac{1-t}{t}}\end{array}σt=t1−t，这是添加到未缩放图像的噪声，类似于EDM的。有了这个，我们可以将EDM[29]定义的系数重写为：cskip = 1tσdata2σdata2+σt2\begin{array}{r}{c_{\mathrm{skip}}\;=\;\frac{1}{t}\frac{\sigma_{\mathrm{data}}^{2}}{\sigma_{\mathrm{data}}^{2}+\sigma_{t}^{2}}}\end{array}cskip=t1σdata2+σt2σdata2和cout = σdataσtσdata2+σt2c_{\mathrm{out}}\,=\,\frac{\sigma_{\mathrm{data}}\sigma_{t}}{\sqrt{\sigma_{\mathrm{data}}^{2}+\sigma_{t}^{2}}}cout=σdata2+σt2 σdataσt，其中σdata\sigma_{\mathrm{data}}σdata是设置为0.5[29]的数据标准差。我们选择cin = t (= 1σt+1)\begin{array}{r}{c_{\mathrm{in}}\,=\,t\;(=\,\frac{1}{\sigma_{t}+1})}\end{array}cin=t(=σt+11)，因此网络的直接输入（即，cin⋅1τztc_{\mathrm{in}}\cdot{\textstyle\frac{1}{\it\tau}}z_{t}cin⋅τ1zt)仍然是ztz_{t}zt。可以证明，仅当t→∁t\to\complementt→∁时，我们有：σt→+∞, cskip→0\sigma_{t}{\rightarrow}+\infty,\,c_{\mathrm{skip}}{\rightarrow}0σt→+∞,cskip→0，cout→1c_{\mathrm{out}}{\rightarrow}1cout→1，这接近x-预测。我们还考虑更简单的线性预条件器：cskip=t and cout=1−tc_{\mathrm{skip}}=t\;\mathrm{and}\;c_{\mathrm{out}}=1-tcskip=tandcout=1−t，这也仅在t=0t=0t=0时执行x-预测。
预条件器的结果 。表10显示，预条件版本全部灾难性失败，表明在高维空间中偏离x-预测是不理想的。有趣的是，预条件版本比表2(a)中的ε-/v-预测好得多。我们假设这是因为当t→0t{\to}0t→0时，它们更接近x-预测，这缓解了这一问题。
B.3. 探索：分类损失
我们的论文侧重于极简主义设计，我们有意不使用任何额外损失。然而，我们注意到，基于潜在的方法[49]通常依赖于使用对抗和感知损失训练的分词器，因此它们的生成过程不是纯粹由扩散驱动的。接下来，我们讨论在我们的像素基础模型上使用附加分类损失的简单扩展。
形式上，我们在特定Transformer块之后附加一个分类器头（JiT-B中为第4个，JiT-L中为第8个）。分类器由全局平均池化后跟线性层组成，并对1000类ImageNet分类任务应用交叉熵损失。此分类损失Lcls\mathcal{L}_{\mathrm{cls}}Lcls由权重λ（例如，100）缩放，并添加到基于ℓ2\ell_2ℓ2的（即，平方误差的元素和）回归损失中。为防止标签泄露，我们为分类器头之前的所有层禁用类别条件。我们注意到，此修改仍然是最小的，不依赖于任何预训练分类器，与感知损失[76]不同。
这一微小修改带来了相当大的改进，如表11所示。这一探索表明，将我们的简单方法与附加损失项结合的进一步潜力，我们将其留待未来工作。
尽管有改进，我们在本文呈现的其他实验中不使用此或任何附加损失。
B.4. 跨分辨率生成
在一个分辨率上训练的模型可以通过简单地下采样或上采样生成的图像应用到另一个分辨率。我们称之为跨分辨率生成。在我们的设置中，256分辨率下的JiT/16和512分辨率下的JiT/32具有可比的参数和计算量，使它们的跨分辨率比较有意义。结果在表12中。
将512模型生成的图像下采样到256分辨率，产生1.84的FID@256。这一结果在与256分辨率专家（FID@256为1.82）比较时仍然具有竞争力，同时保持相似的计算成本并获得在512分辨率下生成的附加能力。
另一方面，将256模型生成的图像上采样到512分辨率，与512分辨率专家的FID@512 1.78相比，FID@512明显更差，为2.45。这一退化是由上采样导致的高频细节损失引起的。
B.5. 附加指标
为完整起见，我们在表13中报告ImageNet 256×256上的精确度和召回率[32]，与常用的DiT和SiT基线以及最新的RAE进行比较：
C. 定性结果
在图8到11中，我们在ImageNet 256×256上提供了更多未精选的示例。
致谢。我们感谢Google TPU研究云(TRC)授予我们TPU访问权限，以及MIT ORCD种子基金拨款支持GPU资源。