回归基础:让去噪生成模型去噪
摘要
如今的去噪扩散模型并不以经典意义上的"去噪"方式工作,即它们不直接预测干净图像 。相反,神经网络预测的是噪声或含噪声的量 。在本文中,我们认为预测干净数据与预测含噪声的量存在根本区别 。根据流形假设,自然数据应位于低维流形上,而含噪声的量则不然。基于此假设,我们提倡模型应直接预测干净数据,这使得看似容量不足的网络能够在非常高维的空间中有效运作。我们展示了在像素上使用简单的大块Transformer可以成为强大的生成模型:无需分词器、无需预训练、也无需额外损失。我们的方法在概念上无非就是 "仅仅是图像Transformer" ,我们称之为 JiT。我们报告了在ImageNet数据集上,分辨率为256和512时,使用大块大小(16和32)的JiT所取得的具有竞争力的结果,而在这些情况下,预测高维含噪声的量可能会灾难性地失败。通过让我们的网络回归到流形的基础,我们的研究也回归基础,并追求一种基于Transformer的、处理原始自然数据的、自包含的扩散范式。
流形假设 是机器学习和统计学中的一个核心概念。它认为:虽然现实世界的数据(如图像、音频)通常以高维形式呈现(例如一张图片的像素数可能达成千上万维),但这些数据实际上分布在一个维数低得多的"流形"(Manifold)上,嵌入在高维空间中。
简单理解:想象一张二维的纸(低维流形),它被卷曲、折叠后放在了一个三维的房间(高维空间)里。虽然纸上的点在三维空间中有坐标,但它们本质上只需要两个坐标(长和宽)就能描述。
自然数据的固有维度 远低于其观察维度。
1 引言
当扩散生成模型最初被提出时 [57, 59, 23],其核心思想应该是去噪,即从其损坏版本中预测干净图像 。然而,扩散模型发展过程中的两个重要里程碑却偏离了直接预测干净图像的目标 。首先,预测噪声本身(称为"ϵ\epsilonϵ-预测")[23] 在生成质量上起到了关键性作用,并极大地推广了这些模型。后来,扩散模型通过预测流速度("v\boldsymbol{v}v-预测" [52])与基于流的方法 [37, 38, 1] 联系起来,该量结合了干净数据和噪声。如今,实践中的扩散模型通常预测噪声或含噪声的量(例如,速度)。
基于流的方法 (Flow/ODE) 核心思想:
把生成过程看作是粒子在空间中的流动。假设有一个向量场(就像水流),数据点沿着这个流从噪声分布漂流到数据分布。 数学工具: 常微分方程
(ODE) 核心任务: 只要模型能预测出这个流场的"速度" v,我们就可以用积分器(Solver)算出粒子的轨迹,从而生成图像。
广泛的研究 [52, 29, 25, 15] 表明,预测干净图像("x\boldsymbol{x}x-预测" [52])与 ϵ\boldsymbol{\epsilon}ϵ- 和 v\boldsymbol{v}v-预测密切相关,前提是相应地重新表述预测损失的加权(详见第3节)。由于这种关系,网络应该直接预测什么这一问题得到的关注较少,隐含地假设网络有能力执行分配的任务。
然而,干净图像和含噪声的量(包括噪声本身)的作用并不相同。在机器学习中,长期以来有一个假设 [4, 3],即"(高维)数据(大致)位于一个低维流形上"([4, p.6])。在这个流形假设下,干净数据可以被建模为位于低维流形上,而含噪声的量本质上分布在整个高维空间中 [69](见图1)。预测干净数据与预测噪声或含噪声的量存在根本区别 。
简单来说,我们不要去预测噪声是什么样子了,也不要预测流速度了,直接预测原图是最简单的。

图 1:流形假设 [4] 假设自然图像位于高维像素空间内的一个低维流形上。虽然干净图像 x\boldsymbol{x}x可以被建模为在流形上,但噪声 ϵ\boldsymbol{\epsilon}ϵ 或流速度v\boldsymbol{v}v(例如,v=x−ϵ\boldsymbol{v}=\boldsymbol{x}-\boldsymbol{\epsilon}v=x−ϵ)本质上是脱离流形的。训练神经网络预测干净图像(即 x\boldsymbol{x}x-预测)与训练它预测噪声或含噪声的量(即ϵ\boldsymbol{\epsilon}ϵ/v\boldsymbol{v}v-预测)是根本不同的。
预测干净图像(x)就是把数据"投影"回有序的低维流形上(简单任务);而预测速度(v)或噪声(ϵ)则是在无序的高维空间里寻找一个混乱的目标(困难任务)。
考虑一个场景,其中低维流形嵌入在高维观测空间中。在这个高维空间中预测噪声需要 高容量 :网络需要保留关于噪声的所有信息。相比之下,一个 容量有限 的网络仍然可以预测干净数据,因为它只需要保留低维信息同时过滤掉噪声。
当使用低维空间(例如,图像潜空间 [49])时,预测噪声的困难得到缓解,但与此同时,这个问题被 隐藏 而非解决。当涉及到像素或其他高维空间时,现有的扩散模型仍然难以应对维度灾难[4]。对预训练潜空间的严重依赖阻碍了扩散模型成为自包含的。
为了追求自包含的原则,学术界一直强烈关注在像素空间中推进扩散建模 [7, 25, 26, 6, 70]。通常,这些方法显式或隐式地避免网络中的信息瓶颈,例如通过使用密集卷积或更小的块、增加通道数、或添加长跳跃连接。我们认为这些设计可能源于预测高维含噪声的量的需求。
在本文中,我们回归第一性原理,让神经网络 直接 预测干净图像。通过这样做,我们展示了在由原始像素组成的大图像块上操作的 朴素 视觉Transformer (ViT) [13] 可以有效地用于扩散建模。我们的方法是 自包含的 ,不依赖于任何预训练或辅助损失------无潜空间分词器 [49]、无对抗损失 [16, 49]、无感知损失 [77, 49](因此无预训练分类器 [56])、无表示对齐 [74](因此无自监督预训练 [45])。在概念上,我们的模型无非就是应用于扩散的"仅仅是图像Transformer",或我们称之为 JiT。
我们在ImageNet数据集 [11] 上以256和512的分辨率进行实验,分别使用块大小为16和32的JiT模型。尽管这些块非常高维(数百或数千维),我们使用 x\bm{x}x-预测的模型可以轻松产生强大的结果,而 ϵ\bm{\epsilon}ϵ- 和 v\bm{v}v-预测则会灾难性地失败。进一步的分析表明,网络的宽度没有必要匹配或超过块的维度;事实上,并且令人惊讶的是,_瓶颈_设计甚至可能有益,这与经典流形学习的观察相呼应。
我们的努力标志着在原生数据上迈向自包含的"扩散 + Transformer"理念 [46] 的一步。除了计算机视觉,这种理念在其他涉及自然数据的领域(例如,蛋白质、分子或天气)也非常可取,在这些领域中分词器可能难以设计。通过最小化领域特定的设计,我们希望源自计算机视觉的通用"扩散 + Transformer"范式能够找到更广泛的适用性。
2 相关工作
2.1 扩散模型及其预测
关于扩散模型的开创性工作 [57] 提出学习一个反向随机过程,其中网络预测正态分布的参数(例如,均值和标准差)。在其引入五年后,该方法通过去噪扩散概率模型 (DDPM ) [23] 得以复兴和普及:一个 关键 发现是将噪声作为预测目标 (即 ϵ\bm{\epsilon}ϵ-预测)。
随后在 [52] 中研究了不同预测目标之间的关系(最初是在模型蒸馏的背景下),其中也引入了 v\bm{v}v-预测的概念。他们的工作侧重于重新参数化引入的加权效应。
同时,EDM [29] 围绕 去噪器函数 重新表述了扩散问题,这构成了扩散模型发展的一个主要里程碑。然而,EDM 采用了 预条件 公式,其中网络的直接输出_不是_去噪后的图像。虽然这种公式在低维场景中更可取,但它本质上仍然要求网络输出一个混合了数据和噪声的量(更多比较见附录)。
流匹配模型 [37, 38, 1] 可以在扩散建模框架内被解释为一种 v\bm{v}v-预测 [52]。与纯噪声不同,v\bm{v}v 是数据和噪声的组合。基于流的模型与先前扩散模型之间的联系已被建立 [15]。如今,扩散模型及其基于流的对应模型通常在统一框架下进行研究。
2.3 去噪模型。
在过去的几十年里,去噪的概念与表示学习密切相关。经典方法,以 BM3D [9] 和其他方法 [47, 14, 79] 为例,利用稀疏性和低维度的假设来执行图像去噪。
去噪自编码器 (DAE) [68, 69] 被开发为一种无监督表示学习方法,使用去噪作为其训练目标。它们利用流形假设 [4](图1)来学习逼近低维数据流形的有意义的表示。DAE 可以被视为去噪分数匹配 [67] 的一种形式,而后者又与基于分数的现代扩散模型 [59, 60] 密切相关。然而,虽然 DAE 很自然地预测干净数据以进行流形学习,但在分数匹配中,预测分数函数实际上等同于预测噪声(相差一个缩放因子),即 ϵ\bm{\epsilon}ϵ-预测。
2.4 流形学习
流形学习 一直是一个经典领域 [51, 63],专注于从观测数据中学习低维、非线性的表示 。通常,流形学习方法利用 瓶颈 结构 [64, 48, 41, 2] 来鼓励只有有用的信息通过。一些研究探索了流形学习与生成模型之间的联系 [39, 27, 8]。潜扩散模型 (LDM) [49] 可以被视为在第一阶段通过自编码器进行流形学习,然后在第二阶段进行扩散。
2.5 像素空间扩散
虽然潜扩散 [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),并伴有长程跳跃连接。虽然这些模型对于 ϵ\epsilonϵ- 和 vvv-预测效果很好,但它们的密集卷积结构通常计算成本高昂。将这些卷积模型应用于高分辨率图像不会导致灾难性的性能下降,因此,这个方向的研究通常侧重于噪声调度和/或加权方案 [7, 25, 26, 30] 以进一步提高生成质量。
相比之下,将视觉 Transformer (ViT) [13] 直接应用于像素提出了更具挑战性的任务。标准 ViT 架构采用激进的块大小(例如,16×16 像素),导致高维的标记空间,其维度可能与 Transformer 的隐藏维度相当或更大。SiD2 [26] 和 PixelFlow [6] 采用了从较小块开始的分层设计;然而,这些模型是"FLOP-heavy" [26] 的,并且失去了标准 Transformer 固有的通用性和简单性。PixNerd [70] 采用了一个 NeRF 头 [43],该头集成了来自 Transformer 输出、噪声输入和空间坐标的信息,并且训练进一步通过表示对齐 [74] 得到辅助。
即使有这些特殊用途的设计,这些工作中的架构通常也从"L"(大)或"XL"尺寸开始。事实上,一项最新工作 [73] 表明,大的隐藏尺寸似乎对于高维性是必要的。
2.6 高维扩散
当使用 ViT 风格的架构时,现代扩散模型仍然受到高维输入空间的挑战,无论是在像素还是潜空间中。在文献 [8, 73, 55] 中,反复 报告了当每个标记的维度增加时,ViT 风格的扩散模型会迅速且 灾难性地 性能下降,无论使用的是像素还是潜空间。
高维扩散的现状:如果你想用 ViT 处理高维数据,现有的解决方案要么是做复杂的架构,要么是做自监督预训练 (Self-supervised Pre-training)(先让模型看很多图学会特征,再来做扩散)。
与我们的工作同时,一系列研究 [78, 34, 55] 诉诸于 自监督预训练 来解决高维扩散问题。与这些努力相比,我们表明高维扩散是 无需 任何预训练,并且 仅 使用 Transformer 就可以实现的。
2.7 xxx- 预测
xxx-预测的公式是自然的且并非新事物 ;它至少可以追溯到原始的 DDPM [23](参见他们的代码 [24])。然而,DDPM 观察到 ϵ\epsilonϵ- 预测明显更好,这后来成为了首选方案。在后期的工作中(例如,[26]),尽管分析有时优选在 xxx- 空间中进行,但实际的预测通常在其他空间中进行,可能是出于传统原因。
当涉及到扩散所解决的图像恢复应用 [10, 72, 42] 时,网络预测干净数据是很自然的 ,因为这是图像恢复的最终目标。与我们的工作同时,[18] 也提倡使用 xxx- 预测,用于以先前帧为条件的生成式世界模型。
我们的工作并非旨在重新发明这个基本概念;相反,我们旨在提请注意一个在具有底层低维流形的高维数据背景下被 largely 忽视但至关重要的问题。
3 论扩散模型的预测输出
扩散模型可以在 xxx、ϵ\epsilonϵ 或 vvv 空间中制定。空间的选择不仅决定了损失定义的位置,也决定了网络预测的内容。重要的是,损失空间和网络输出空间不必相同。这个选择可能产生关键性差异。
3.1 背景:扩散与流
扩散模型可以从 ODE 的角度来表述 [5, 60, 37, 58]。我们从一个更简单的起点开始我们的公式,即基于流的范式 [37, 38, 1],也就是在 **vvv-**空间中,然后讨论其他空间。
考虑一个数据分布 x∼pdata(x)\bm{x}\sim p_{\text{data}}(\bm{x})x∼pdata(x) 和一个噪声分布 ϵ∼pnoise(ϵ)\bm{\epsilon}\sim p_{\text{noise}}(\bm{\epsilon})ϵ∼pnoise(ϵ)(例如,ϵ∼N(0,I)\bm{\epsilon}\sim\mathcal{N}(\mathbf{0},\mathbf{I})ϵ∼N(0,I))。在训练期间,一个噪声样本 zt\bm{z}{t}zt 是一个插值:zt=at x + bt ϵ\bm{z}{t}=\bm{a}{t}\,\bm{x}\,+\,b{t}\,\bm{\epsilon}zt=atx+btϵ,其中 at\bm{a}{t}at 和 btb{t}bt 是在时间 t∈[0,1]t\in [0,1]t∈[0,1] 的预定义噪声调度。在本文中,我们使用线性调度 [37, 38, 1]¹:at=t\bm{a}{t}=tat=t 和 bt=1−tb{t}=1-tbt=1−t。这给出:
zt=t x+(1−t) ϵ,\bm{z}_{t}=t\,\bm{x}+(1-t)\,\bm{\epsilon},zt=tx+(1−t)ϵ,
(1)
当 ttt=1 时,这导致 zt∼pdata\bm{z}{t}\sim p{\text{data}}zt∼pdata。我们对 ttt 使用 logit-正态分布 [15],即 logit(t)∼N(μ,σ2)\text{logit}(t)\sim\mathcal{N}(\mu,\sigma^{2})logit(t)∼N(μ,σ2)。
流速度 v\bm{v}v 定义为 z\bm{z}z 的时间导数,即 vt=zt′=at′x+bt′ϵ\bm{v}{t}=\bm{z}^{\prime}{t}=\bm{a}^{\prime}{t}\bm{x}+b^{\prime}{t}\bm{\epsilon}vt=zt′=at′x+bt′ϵ。给定方程 (1),我们有:
v=x−ϵ.\bm{v}=\bm{x}-\bm{\epsilon}.v=x−ϵ.
(2)
基于流的方法 [37, 38, 1] 最小化定义为以下形式的损失函数:
L=Et,x,ϵ∥vθ(zt,t)−v∥2,\mathcal{L}=\mathbb{E}{t,\bm{x},\bm{\epsilon}}\|\bm{v}{\theta}(\bm{z}_{t},t)-\bm{v}\|^{2},L=Et,x,ϵ∥vθ(zt,t)−v∥2,
(3)
其中 vθ\bm{v}{\theta}vθ 是由 θ\thetaθ 参数化的函数。虽然 vθ\bm{v}{\theta}vθ 通常是网络 vθ = netθ(zt,t)\bm{v}{\theta}\,=\,\mathrm{net}{\theta}(\bm{z}_{t},t)vθ=netθ(zt,t) 的_直接_输出 [37, 38, 1],但它也可以是它的一个变换,正如我们将要阐述的。
给定函数 vθ\bm{v}_{\theta}vθ,通过求解 z\bm{z}z 的常微分方程 (ODE) 来完成采样 [37, 38, 1]:
dzt/dt=vθ(zt,t),d\bm{z}{t}/dt=\bm{v}{\theta}(\bm{z}_{t},t),dzt/dt=vθ(zt,t),
(4)
从 z0∼pnoise\bm{z}{0}\sim p{\text{noise}}z0∼pnoise 开始,到 t=1t=1t=1 结束。在实践中,这个 ODE 可以使用数值求解器近似求解。默认情况下,我们使用一个 50 步的 Heun 求解器。
3.2 预测空间与损失空间
预测空间。 网络的直接输出可以在任何空间中定义:v,x\bm{v},\bm{x}v,x,或 ϵ\bm{\epsilon}ϵ。接下来,我们讨论由此产生的变换。注意,在本文的上下文中,我们将其称为"x\bm{x}x、ϵ\bm{\epsilon}ϵ、v\bm{v}v**-预测** ",仅当_网络 netθ\operatorname{net}{\theta}netθ 的_直接_输出严格是 x\bm{x}x、ϵ\bm{\epsilon}ϵ、v\bm{v}v 时。
给定三个未知数(x\bm{x}x、ϵ\bm{\epsilon}ϵ、v\bm{v}v)和一个网络输出,我们需要两个额外的约束来确定所有三个未知数。这两个约束由方程 (1) 和 (2) 给出。例如 ,当我们让直接网络输出 netθ\operatorname{net}_{\theta}netθ 为 x\bm{x}x 时,我们求解以下方程组:
{xθ=netθzt=t xθ+(1−t) ϵθvθ=xθ−ϵθ\begin{cases}\bm{x}{\theta}=\operatorname{net}{\theta}\\ \bm{z}{t}=t\,\bm{x}{\theta}+(1-t)\,\bm{\epsilon}{\theta}\\ \bm{v}{\theta}=\bm{x}{\theta}-\bm{\epsilon}{\theta}\end{cases}⎩ ⎨ ⎧xθ=netθzt=txθ+(1−t)ϵθvθ=xθ−ϵθ
(5)
这里,符号 xθ\bm{x}{\theta}xθ、ϵθ\bm{\epsilon}{\theta}ϵθ 和 vθ\bm{v}{\theta}vθ 表明它们都是依赖于 θ\thetaθ 的预测。求解这个方程组得到:ϵθ=(zt−txθ)/(1−t)\bm{\epsilon}{\theta}=(\bm{z}{t}-t\bm{x}{\theta})/(1-t)ϵθ=(zt−txθ)/(1−t) 和 vθ=(xθ−zt)/(1−t)\bm{v}{\theta}=(\bm{x}{\theta}-\bm{z}{t})/(1-t)vθ=(xθ−zt)/(1−t),也就是说,ϵθ\bm{\epsilon}{\theta}ϵθ 和 vθ\bm{v}{\theta}vθ 都可以从 zt\bm{z}{t}zt 和网络 xθ\bm{x}_{\theta}xθ 计算出来。这些总结在表 1 的列 (a) 中。
类似地,当我们让直接网络输出 netθ\operatorname{net}_{\theta}netθ 为 ϵ\bm{\epsilon}ϵ 或 v\bm{v}v 时,我们得到其他方程组(通过替换方程 (5) 中的第一个方程)。这些变换总结在表 1 的 (b), © 列中,分别对应 ϵ\bm{\epsilon}ϵ-、v\bm{v}v-预测。这表明当 {x,ϵ,v}\{\bm{x},\bm{\epsilon},\bm{v}\}{x,ϵ,v} 中的一个量被预测时,其他两个可以被推断出来。许多先前工作中的推导(例如,[52, 15])是表 1 涵盖的特殊情况。
损失空间。 虽然损失通常在一个参考空间中定义(例如,方程 (3) 中的 v\bm{v}v-损失),但概念上,可以在任何空间中定义它。已经表明 [52, 15],在给定从一个预测空间到另一个预测空间的重新参数化的情况下,损失实际上被重新加权了。
例如 ,考虑表 1**(3)(a)** 中 x\bm{x}x**-预测** 和 v\bm{v}v**-损失** 的组合。我们有 vθ=(xθ−zt)/(1−t)\bm{v}{\theta}=(\bm{x}{\theta}-\bm{z}{t})/(1-t)vθ=(xθ−zt)/(1−t) 作为预测,而 v=(x−zt)/(1−t)\bm{v}=(\bm{x}-\bm{z}{t})/(1-t)v=(x−zt)/(1−t) 作为目标。方程 (3) 中的 v\bm{v}v-损失变为:L=E∥vθ(zt,t)−v∥2=E1(1−t)2∥xθ(zt,t)−x∥2\mathcal{L}=\mathbb{E}\|\bm{v}{\theta}(\bm{z}{t},t)-\bm{v}\|^{2}=\mathbb{E}\frac{1}{(1-t)^{2}}\|\bm{x}{\theta}(\bm{z}{t},t)-\bm{x}\|^{2}L=E∥vθ(zt,t)−v∥2=E(1−t)21∥xθ(zt,t)−x∥2,这是 x\bm{x}x-损失的一个_重新加权_形式。像这样的变换可以对表 1 中列出的任何预测空间和任何损失空间进行。
总结起来,考虑在 {x,ϵ,v}\{\bm{x},\bm{\epsilon},\bm{v}\}{x,ϵ,v} 中定义的三个_未加权_损失,以及网络直接输出的三种形式,总共有_九种可能的组合_(表 1)。每种组合都构成一个有效的公式,并且这九种情况中_没有_两种在数学上是等价的。
生成器空间。 无论使用哪种组合,为了在推理时执行生成,我们总是可以将网络输出转换到 v\bm{v}v-空间(表 1,第 (3) 行)并求解方程 (4) 中的 ODE 进行采样。因此,所有九种组合都是合法的生成器。
表 1:在 x\bm{x}x、v\bm{v}v 或 ϵ\bm{\epsilon}ϵ 空间中定义损失和网络预测的所有可能组合。 网络的直接输出用颜色高亮显示。对于网络输出空间与损失空间不同的任何非对角线条目,都会对网络输出应用变换。
3.3 玩具实验
根据流形假设 [4],数据 x\bm{x}x 倾向于位于低维流形上(图 1),而噪声 ϵ\bm{\epsilon}ϵ 和速度 v\bm{v}v 是脱离流形的。让网络直接预测干净数据 x\bm{x}x 应该更容易处理。我们在本节的玩具实验中验证这个假设。
我们考虑 ddd 维底层数据"埋藏"在观测到的 DDD 维空间(d<Dd<Dd<D)中的玩具情况。我们使用一个列正交的投影矩阵 P∈RD×dP\in\mathbb{R}^{D\times d}P∈RD×d 来合成这个场景,即 P⊤P=Id×dP^{\top}P=I_{d\times d}P⊤P=Id×d。这个矩阵 PPP 是随机创建并固定的。观测数据是 x=Px^∈RD\bm{x}=P\hat{\bm{x}}\in\mathbb{R}^{D}x=Px^∈RD,其中底层数据是 x^∈Rd\hat{\bm{x}}\in\mathbb{R}^{d}x^∈Rd。矩阵 PPP 对模型是未知的,因此,对于模型来说,这是一个 DDD 维的生成问题。
图 2:玩具实验 :ddd 维(d=2d=2d=2)底层数据通过一个固定的、随机的、列正交的投影矩阵"埋藏"在 DDD 维空间中。在 DDD 维空间中,我们训练一个简单的生成模型(5层 ReLU MLP,256维隐藏单元)。投影矩阵对模型是未知的,我们仅用它来可视化输出。在这个玩具实验中,随着观测维度 DDD 的增加,只有 x\bm{x}x-预测 能产生合理的结果。
我们训练一个具有 256 维隐藏单元的 5 层 ReLU MLP 作为生成器,并在图 2 中可视化结果。我们通过使用 PPP 将 DDD 维生成的样本投影回 ddd 维来获得这些可视化结果。我们研究了对于 d=2d=2d=2,D∈{2,8,16,512}D\in \{2,8,16,512\}D∈{2,8,16,512} 的情况。我们研究 x\bm{x}x、ϵ\bm{\epsilon}ϵ 或 v\bm{v}v-预测,全部使用 v\bm{v}v-损失,即表 1 (3)(a-c)。
图 2 显示,只有 x\bm{x}x*-预测在 DDD 增加时能产生合理的结果* 。对于 ϵ\bm{\epsilon}ϵ-/v\bm{v}v-预测,模型在 DDD=16 时就很困难,在 DDD=512 时则灾难性地失败,此时 256 维的 MLP 是不完备的。
值得注意的是,x\bm{x}x*-预测即使在模型是不完备的情况下也能很好地工作* 。这里,256 维的 MLP 不可避免地会丢弃 DDD=512 维空间中的信息。然而,由于真实数据位于低维的 ddd 维空间中,x\bm{x}x-预测仍然可以表现良好,因为理想的输出隐含地是 ddd 维的。我们在 ImageNet 上的真实数据案例中也得出了类似的观察结果,如下所示。
4 用于扩散的"仅仅是图像Transformer"
受上述分析的驱动,我们展示了在像素上操作的_朴素_视觉 Transformer (ViT) [13] 可以出人意料地工作得很好,只需使用 x\bm{x}x-预测。
4.1 仅仅是图像Transformer
ViT [13] 的核心思想是"在块上的Transformer (ToP)"²。我们的架构设计遵循这一理念。
形式上,考虑 H×W×CH\times W\times CH×W×C 维的图像数据(CCC=3)。所有 x\bm{x}x、ϵ\bm{\epsilon}ϵ、v\bm{v}v 和 zt\bm{z}_{t}zt 共享相同的维度。给定一张图像,我们将其划分为大小为 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。
作为标准实践,该架构以时间 ttt 和给定的类别标签为条件。我们使用 adaLN-Zero [46] 进行条件化,稍后将讨论其他选项。概念上,这种架构相当于将扩散 Transformer (DiT) [46] 直接应用于像素块。
整个架构无非就是"仅仅是图像Transformer",我们称之为 JiT 。例如,我们在 256×256256\times 256256×256 图像上研究 JiT/16 (即块大小 ppp=16,[13]),在 512×512512\times 512512×512 图像上研究 JiT/32 (ppp=32)。这些设置分别导致每个块的维度为 768(16×16×316\times 16\times 316×16×3)和 3072(32×32×332\times 32\times 332×32×3)。如此高维的块可以通过 x\bm{x}x-预测来处理。
4.2 网络应该预测什么?
我们在表 1 中总结了损失空间和预测空间的所有九种可能组合。对于每种组合,我们训练一个"Base" [13] 模型 (JiT-B),该模型每个标记的隐藏大小为 768 维。我们在表 2(a) 中研究了在 256×256 分辨率下的 JiT-B/16。作为参考,我们在表 2(b) 中检查了在 64×64 分辨率下的 JiT-B/4(即 ppp=4)。在这两种设置中,序列长度是相同的(16×16)。
我们得出以下观察结果:
x\bm{x}x-预测至关重要。 在表 2(a) 中使用 JiT-B/16 的情况下,只有 x\bm{x}x*-预测表现良好* ,并且它在所有三种损失下都有效。这里,一个块是 768 维(16×16×3),这与 JiT-B 的隐藏大小 768 一致。虽然这看起来"差不多足够",但在实践中,模型可能需要额外的容量,例如 ,来处理位置嵌入。对于 ϵ\bm{\epsilon}ϵ-/v\bm{v}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) 中,ϵ\epsilonϵ- 和 υ\upsilonυ-预测都灾难性地失败,无论损失空间如何,这对应于不同损失空间中不同的有效加权(如前所述)。另一方面,x\bm{x}x-预测在_所有_三种损失空间中都有效:由 υ\upsilonυ-损失引起的损失加权更可取,但并非关键⁴。
图 3:"仅仅是图像Transformer" (JiT) 架构 :简单地是一个在像素块上的朴素 ViT [13],用于 x\bm{x}x-预测。
噪声水平偏移并不足够。 先前的工作 [7, 25, 26] 表明,增加噪声水平对于基于像素的高分辨率扩散是有用的。我们在表 3 中使用 JiT-B/16 对此进行了检验。由于我们使用 logit-正态分布 [15] 来采样 ttt(见附录),噪声水平可以通过改变该分布的参数 μ\muμ 来偏移:直观上,将 μ\muμ 向负方向偏移会导致更小的 ttt,从而增加噪声水平(方程 (1))。
表 3 显示,当模型已经表现良好时(此处为 x\bm{x}x-预测 ),适当高的噪声是有益的,这与先前的观察结果 [7, 25, 26] 一致。然而,_仅仅_调整噪声水平_无法_补救 ϵ\epsilonϵ- 或 υ\upsilonυ-预测:它们的失败源于无法传播高维信息的固有缺陷。
另外,根据表 3,我们在 ImageNet 256×256 的其他实验中设置 μ=−0.8\mu=-\text{0.8}μ=−0.8。
增加隐藏单元并非必要。 由于容量可能受网络宽度(即隐藏单元数量)的限制,一个自然的想法是增加它。然而,当观测维度非常高时,这种补救措施既没有原则性也不可行。我们表明在 x\bm{x}x-预测的情况下,这不是必要的。
在下一节的表 5 和表 6 中,我们展示了在 512 分辨率下的 JIT/32 和在 1024 分辨率下的 JIT/64 的结果,使用了_按比例放大_的块大小 ppp=32 或 ppp=64。这相当于每个块 3072 维(即 32×32×3)或 12288 维,远大于 B、L 和 H 模型(在 [13] 中定义)的隐藏大小。尽管如此,x\bm{x}x-预测效果很好;事实上,它_无需_任何修改就能工作,除了按比例缩放噪声(例如,在分辨率 512 和 1024 时分别缩放 2 倍和 4 倍;见附录)。
这一证据表明,网络设计可以在很大程度上与观测维度_解耦_,就像在许多其他神经网络应用中一样。增加隐藏单元的数量可能是有益的(正如在深度学习中广泛观察到的),但这并非决定性因素。
瓶颈可能有益。 更令人惊讶的是,我们发现,相反地,引入一个_降低_网络维度的瓶颈可能是有益的。
具体来说,我们将线性块嵌入层转变为_低秩_线性层,通过将其替换为一对瓶颈(但仍然是线性)层来实现。第一层将维度减少到 d′d^{\prime}d′,第二层将其扩展到 Transformer 的隐藏大小。两层都是线性的,并作为低秩重新参数化。
图 4 绘制了 FID 与瓶颈维度 d′d^{\prime}d′ 的关系,使用 JiT-B/16(每个原始块 768 维)。减小瓶颈维度,即使小到 16 维,也不会导致灾难性的失败。事实上,在一个宽范围(32 到 512)内的瓶颈维度可以提高质量,提升幅度高达约 1.3 FID。
从表示学习的更广泛视角来看,这一观察并非完全出乎意料。瓶颈设计通常被引入以鼓励学习固有的低维表示 [64, 48, 41, 2]。
表 3:噪声水平偏移 (JiT-B/16,ImageNet 256×256,FID-50K)。我们通过调整 logit-正态 ttt 采样器 [15] 中的 μ\muμ 来偏移噪声水平。适当的噪声水平是有用的,但不足以解决 ϵ\epsilonϵ-/υ\upsilonυ-预测中的灾难性失败。设置(与表 2 相同):200 轮 epoch,使用 CFG。
表 2:损失空间和网络空间所有组合的结果 (见表 1),在 ImageNet 上通过 FID-50K 评估:(a) JiT-B/16 在 256 分辨率,每个块 768 维;(b) JiT-B/4 在 64 分辨率,每个块 48 维。我们用红色标注灾难性失败,用绿色标注合理结果。设置:200 轮 epoch,使用 CFG [22]。
图 4:瓶颈线性嵌入。 结果是针对 ImageNet 256×256 上的 JiT-B/16 。一个原始块是 768 维(16×16×3),并通过两个连续的线性层嵌入,中间有一个瓶颈维度 d′d^{\prime}d′(d′<768d^{\prime}<768d′<768)。这里,瓶颈嵌入通常是有益的 ,并且我们的 x\bm{x}x-预测模型即使在使用小至 32 或 16 的激进瓶颈时也能正常工作。设置(与表 3 相同):200 轮 epoch,使用 CFG。
4.3 我们的算法
我们的最终算法采用 x\bm{x}x-预测和 v\bm{v}v-损失,这对应于表 1(3)(a)。形式上,我们优化:

其中: vθ(zt,t)=(netθ(zt,t)−zt)/(1−t).\text{其中: }\ \ \bm{v}{\theta}(\bm{z}{t},t) =(\operatorname{net}{\theta}(\bm{z}{t},t)-\bm{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 报告了这些组件的影响。
表 4:具有_通用_设计的"仅仅是高级"Transformer。 所有都是用于 ImageNet 256×256 的 JiT/16,带有瓶颈块嵌入(128 维,图 4),通过 FID-50K 评估。设置:200 轮 epoch,使用 CFG(以及括号中的 CFG 区间 [33])。
5 比较
像素上的高分辨率生成。 在表 5 中,我们进一步报告了我们的_基础_尺寸模型 (JiT-B) 在 ImageNet 上分辨率为 512 甚至 1024 的结果。我们使用_与图像尺寸成比例_的块大小,因此不同分辨率下的序列长度保持不变。每个块的维度可以高达 3072 或 12288,并且_没有_一个常见模型具有足够的隐藏单元。
表 5 显示我们的模型在不同分辨率下表现良好。所有模型具有相似的参数数量和计算成本,仅在输入/输出块嵌入上有所不同。我们的方法不受观测维度灾难的影响。
可扩展性。 解耦 Transformer 设计与任务的一个核心目标是利用可扩展性的潜力。表 6 提供了在 ImageNet 256 和 512 分辨率下四种模型尺寸的结果(注意在 512 分辨率下,这些模型中没有一个的隐藏单元数超过块的维度)。模型尺寸和 flops 显示在表 7 和 8 中:我们的模型在 256 分辨率下的成本与其在 512 分辨率下的对应模型相似。我们的方法受益于缩放。
有趣的是,256 和 512 分辨率之间的 FID 差异随着模型增大而变小。对于 JiT-G,512 分辨率的 FID 甚至更低。对于 ImageNet 上的非常大的模型,FID 性能在很大程度上取决于过拟合,而在 512 分辨率下去噪提出了一个更具挑战性的任务,使其不易过拟合。
表 5:使用 JiT-B/64 的 ImageNet 1024×1024。 所有条目具有大致相同的参数数量和计算量。设置:如果此处未指定,则与表 4 相同(全部使用 CFG 区间)。
表 6:在 ImageNet 256×256 和 512×512 上的可扩展性,通过 FID-50K 评估。所有模型具有_相同_的 16×16 序列长度,因此 512 分辨率的模型与其 256 分辨率的对应模型具有几乎相同的计算量。设置:与表 5 相同。
表 7 和 8:模型尺寸和计算成本。 参数数量(百万)和 Gflops(用于单次前向传播,不包括分词器)。
来自先前工作的参考结果。 作为参考,我们在表 7 和 8 中与先前的结果进行了比较。我们标记了每种方法所涉及的_预训练_组件。与其他基于像素的方法相比,我们的方法纯粹由_朴素的_、通用的 Transformer 驱动。我们的模型是_计算友好的_,并且避免了分辨率加倍时费用的_二次_缩放(参见表 8 中的 flops)。
我们的方法不使用额外的损失或预训练,这可能会带来进一步的收益(附录中有一个例子)。这些方向留待未来工作。
6 讨论与结论
噪声本质上不同于自然数据。 多年来,扩散模型的发展主要集中于概率公式,而对所使用的神经网络的能力(和局限性)关注较少。然而,神经网络的能力并非无限,它们可以更好地利用其容量来建模数据而不是噪声。在这些视角下,我们关于 x\bm{x}x-预测的发现,事后看来,是一个自然的结果。
我们的工作采用了一种极简主义和自包含的设计。通过减少领域特定的归纳偏差,我们希望我们的方法能够推广到其他难以获得分词器的领域。这一特性对于涉及原始、高维自然数据的科学应用尤其可取。我们设想通用的"扩散 + Transformer"范式将成为其他领域的潜在基础。
表 7:ImageNet 256×256 上的参考结果。 评估了 50K 样本的 FID [21] 和 IS [53]。"预训练"列列出了获得结果所需的外部模型(注意感知损失 [77] 使用了预训练的 VGG 分类器 [56])。参数包括生成器和分词器解码器(在推理时使用),但排除其他预训练组件。Giga-flops 是针对单次前向传播测量的(不计算分词器),并且大致与训练和推理期间每次迭代的计算成本成比例(对于多尺度方法 [6],我们测量最精细的级别)。
表 8:ImageNet 512×512 上的参考结果。 评估了 50K 样本的 FID [21] 和 IS [53]。细节同表 7。
图 5:定性结果。 使用 JiT-H/32 在 ImageNet 512×512 上选取的示例。更多未经筛选的结果见附录。
附录 A 实现细节
我们的实现紧密遵循 DiT [46] 和 SiT [40] 的公开代码库。我们的配置总结在表 9 中。我们描述细节如下。
时间分布。 遵循 [15],在训练期间,我们对 ttt 采用 logit-正态分布 [15]:logit(t)∼N(μ,σ2)\text{logit}(t)\sim\mathcal{N}(\mu,\sigma^{2})logit(t)∼N(μ,σ2)。具体来说,我们采样 s∼N(μ,σ2)s\sim\mathcal{N}(\mu,\sigma^{2})s∼N(μ,σ2) 并令 ttt=sigmoid(sss)。超参数 μ\muμ 决定了噪声水平(见表 3),默认情况下,在 ImageNet 256(或 512、1024)分辨率上我们设置 μ=−0.8\mu=-0.8μ=−0.8,并固定 σ\sigmaσ 为 0.8。
ImageNet 512×512 和 1024×1024。 我们在 ImageNet 512×512 上采用 JiT/32(即块大小为 32)。该模型产生一个 256 = 16×16 块的序列,与 ImageNet 256×256 上的 JiT/16 相同。因此,JiT/32 仅在输入/输出块维度上与 JiT/16 不同,从每个块 768 维增加到 3072 维;所有其他计算和成本完全相同。
为了重用 ImageNet 256×256 的完全相同的配方,对于 512×512 图像,我们将噪声 ϵ\epsilonϵ 的幅度重新缩放 2 倍:即 ϵ∼N(0,22I)\epsilon\sim\mathcal{N}(0,2^{2}I)ϵ∼N(0,22I)。这个简单的修改近似保持了 256×256 和 512×512 分辨率之间的信噪比 (SNR) [25, 7, 26]。不需要或应用对 ImageNet 256×256 配置的其他更改。
对于 ImageNet 1024×1024,我们使用模型 JiT/64 并将噪声 ϵ\epsilonϵ 缩放 4 倍。不需要其他更改。
上下文中的类别条件化。 标准的 DiT [46] 通过 adaLN-Zero 执行类别条件化。在表 4 中,我们进一步探索了上下文中的类别条件化。
具体来说,遵循 ViT [13],可以将一个类别标记附加到块序列的前面。这在 DiT [46] 中被称为"上下文中的条件化"。注意,我们与默认的 adaLN-Zero 条件化联合使用上下文中的条件化,这与 DiT 不同。此外,遵循 MAR [35],我们进一步将多个这样的标记附加到序列前面。这些标记是相同类别标记的重复实例,添加了不同的位置嵌入。我们附加 32 个这样的标记。而且,我们发现将这些标记附加到 Transformer 的后面块(见表 9 中的"上下文起始块")可能更有益,而不是附加到 Transformer 的输入。表 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 轮 epoch。
EMA 和 CFG。 我们的研究涵盖了广泛的配置,包括损失和预测空间、模型尺寸和架构组件的变体。CFG 尺度 [22] 和 EMA(指数移动平均)衰减的最佳值因情况而异,固定它们可能导致不完整或误导性的观察结果。由于维护这些超参数配置相对成本较低,我们力求采用最优值。
具体来说,对于 CFG 尺度 ω\omegaω [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]。
表 9:实验配置。
图 6:CFG 的影响,无和有 CFG 区间(对于 ImageNet 256×256 上的 JiT-L/16,600 轮 epoch)。
附录 B 额外实验
B.1 训练损失和去噪图像
在表 2(a) 中,ϵ\bm{\epsilon}ϵ-/v\bm{v}v-预测的失败是由有限容量的网络无法预测高维输出所固有的无能引起的。这种失败可以从训练损失曲线中看出。
在图 7(顶部),我们比较了在相同的 v\bm{v}v-损失下的训练损失曲线,定义为 L=E∥vθ(zt,t)−v∥2\mathcal{L}=\mathbb{E}\|\bm{v}{\theta}(\bm{z}{t},t)-\bm{v}\|^{2}L=E∥vθ(zt,t)−v∥2,使用 v\bm{v}v-预测(即 vθ=netθ\bm{v}{\theta}=\mathrm{net}{\theta}vθ=netθ)与 x\bm{x}x-预测(即 vθ=(netθ−zt)/(1−t)\bm{v}{\theta}=(\mathrm{net}{\theta}-\bm{z}_{t})/(1-t)vθ=(netθ−zt)/(1−t))。由于损失是在同一空间中计算的,只有参数化不同,比较损失值是合理的。
图 7(顶部)显示,v\bm{v}v-预测产生的损失值显著高于 x\bm{x}x-预测(约 25%),即使 v\bm{v}v-预测似乎是 v\bm{v}v-损失的"原生"参数化。这一比较表明,x\bm{x}x-预测的任务本质上_更容易_,因为数据位于低维流形上。我们还观察到 ϵ\bm{\epsilon}ϵ-预测(此处未显示)的损失高出约 3×3\times3× 并且不稳定,这可能解释了其在表 2(a) 中的失败。
图 7(底部)比较了对应于两条训练曲线的_去噪_图像。对于 x\bm{x}x-预测,去噪图像就是 xθ=netθ\bm{x}{\theta}=\mathrm{net}{\theta}xθ=netθ;对于 v\bm{v}v-预测,去噪图像是 xθ\bm{x}{\theta}xθ=(1−t)netθ+zt(1-t)\mathrm{net}{\theta}+\bm{z}{t}(1−t)netθ+zt,其中 vθ\bm{v}{\theta}vθ=netθ\mathrm{net}_{\theta}netθ(见表 1©(1))。图 7(顶部)中 v\bm{v}v-预测的_更高_损失可以反映在图 7(底部)中其明显的_伪影_上。
注意,图 7(底部)中的伪影是_单次_去噪步骤的伪影。在生成过程中,这种误差会在多步 ODE 求解器中累积,从而导致表 2(a) 中的灾难性失败。
B.2 预条件器
在 EDM [29] 中,应用了一个额外的"预条件器"来包装网络的直接输出。使用我们论文的符号,预条件器公式可以写成:xθ(zt,t)=cskip⋅zt+cout⋅netθ(zt,t)\bm{x}{\theta}(\bm{z}{t},t)=c_{\textrm{skip}}\cdot\bm{z}{t}+c{\textrm{out}}\cdot\mathrm{net}{\theta}(\bm{z}{t},t)xθ(zt,t)=cskip⋅zt+cout⋅netθ(zt,t),其中 cskipc_{\textrm{skip}}cskip 和 coutc_{\textrm{out}}cout 是预定义的系数。这个方程表明,除非 cskip≡0c_{\textrm{skip}}\equiv 0cskip≡0,否则预条件器中的网络输出必须_不_执行 x\bm{x}x-预测。并且根据流形假设,这个公式不应该补救我们考虑的问题,正如我们接下来检验的那样。
预条件器的公式。 给定预条件器的定义,我们可以将"预条件化的 x\bm{x}x-预测"重写为:
{xθ=cskip⋅zt+cout⋅netθϵθ=(zt−txθ)/(1−t)vθ=(xθ−zt)/(1−t)\begin{cases}\bm{x}{\theta}=c{\textrm{skip}}\cdot\bm{z}{t}+c{\textrm{out}}\cdot\mathrm{net}{\theta}\\ \bm{\epsilon}{\theta}=(\bm{z}{t}-t\bm{x}{\theta})/(1-t)\\ \bm{v}{\theta}=(\bm{x}{\theta}-\bm{z}_{t})/(1-t)\end{cases}⎩ ⎨ ⎧xθ=cskip⋅zt+cout⋅netθϵθ=(zt−txθ)/(1−t)vθ=(xθ−zt)/(1−t)
(7)
相应地,方程 (6)(v\bm{v}v-损失)中的目标写为:

与方程 (6) 相比,唯一的区别在于我们如何从网络计算 xθ\bm{x}_{\theta}xθ。
由于 EDM [29] 使用方差爆炸调度(即 zt\bm{z}{t}zt=x\bm{x}x+σtϵ\sigma{t}\bm{\epsilon}σtϵ),而我们使用(大致)方差保持版本(即 zt\bm{z}{t}zt=txt\bm{x}tx+(1−t)ϵ(1-t)\bm{\epsilon}(1−t)ϵ),完全等效的转换是不可能的。为了在我们的情况下有一个预条件器,我们将我们的版本重写为 1tzt=x+1−ttϵ\frac{1}{t}\bm{z}{t}=\bm{x}+\frac{1-t}{t}\bm{\epsilon}t1zt=x+t1−tϵ。因此,我们设 σt=1−tt\sigma_{t}=\frac{1-t}{t}σt=t1−t,这是添加到未缩放图像上的噪声,类似于 EDM 的。有了这个,我们可以将 EDM [29] 定义的系数重写为:cskip=1tσdata2σdata2+σt2c_{\textrm{skip}}=\frac{1}{t}\frac{\sigma_{\textrm{data}}^{2}}{\sigma_{\textrm{data}}^{2}+\sigma_{t}^{2}}cskip=t1σdata2+σt2σdata2 和 cout=σtσdata2+σt2c_{\textrm{out}}=\frac{\sigma_{t}}{\sqrt{\sigma_{\textrm{data}}^{2}+\sigma_{t}^{2}}}cout=σdata2+σt2 σt,其中 σdata\sigma_{\textrm{data}}σdata 是数据标准差(我们设 σdata=0.5\sigma_{\textrm{data}}=0.5σdata=0.5,遵循 EDM)。这个微小的修改带来了不错的改进,如表 11 所示。这一探索表明,将我们的简单方法与额外的损失项相结合有进一步的潜力,我们将其留待未来工作。
尽管有改进,我们在本文中提出的其他实验中 没有 使用这个或任何额外的损失。
图 7:(顶部):x\bm{x}x- 和 v\bm{v}v-预测的训练损失,使用相同的 v\bm{v}v-损失空间(表 2(a),第三行)。 我们绘制了每个像素每个通道的平均损失。(底部):来自 x\bm{x}x- 和 v\bm{v}v-预测的去噪图像 ,其中 v\bm{v}v-预测的去噪输出根据表 1©(1) 可视化。来自 v\bm{v}v-预测的去噪图像有明显的伪影,这反映了更高的损失。
B.3 跨分辨率生成
在一个分辨率上训练的模型可以通过简单地对生成的图像进行下采样或上采样来应用于另一个分辨率。我们称之为_跨分辨率_生成。在我们的设置中,256 分辨率的 JiT/16 和 512 分辨率的 JiT/32 具有可比较的参数和计算量,使得它们的跨分辨率比较有意义。结果在表 12 中。
将 512 模型生成的图像_下采样_到 256 分辨率产生了不错的 FID@256 为 1.84。这个结果与 256 分辨率的专家模型(FID@256 为 1.82)相比仍然具有竞争力,同时保持了相似的计算成本并获得了额外生成 512 分辨率图像的能力。
另一方面,将 256 模型生成的图像_上采样_到 512 分辨率导致明显更差的 FID@512 为 2.45,与 512 分辨率的专家模型的 FID@512 为 1.78 相比。这种性能下降是由于上采样导致的高频细节丢失造成的。
B.4 额外指标
为完整起见,我们在表 13 中报告了 ImageNet 256×256 上的精确度和召回率 [32],与常用的 DiT 和 SiT 基线以及最新的 RAE 进行了比较:
附录 C 定性结果
在图 8 到 11 中,我们提供了 ImageNet 256×256 上额外的_未经筛选_的示例。
表 10:与预条件器的比较(FID-50K,ImageNet 256,JiT-B/16)。设置与表 2 (a) 相同。
表 11:探索:额外的分类损失。 我们在任何其他实验中_不_使用此损失。设置:ImageNet 256×256,200 轮 epoch,使用 CFG 区间。
表 12:跨分辨率生成 (灰色标注),使用在 256 分辨率训练的 JiT-G/16 和在 512 分辨率训练的 JiT-G/32,然后进行上采样(↑\uparrow↑)或下采样(↓\downarrow↓)。所有条目具有相似的参数和 flops(见表 7 和 8)。
表 13:ImageNet 256×256 上的精确度和召回率。
图 8:在 ImageNet 256×256 上使用 JiT-G/16 以指定类别为条件的_未经筛选_样本。与使用更高 CFG 可视化的常见做法不同,这里我们使用达到报告 FID 1.82 的 CFG 值(2.2)显示图像。
图 9:在 ImageNet 256×256 上使用 JiT-G/16 以指定类别为条件的_未经筛选_样本。与使用更高 CFG 可视化的常见做法不同,这里我们使用达到报告 FID 1.82 的 CFG 值(2.2)显示图像。
图 10:在 ImageNet 256×256 上使用 JiT-G/16 以指定类别为条件的_未经筛选_样本。与使用更高 CFG 可视化的常见做法不同,这里我们使用达到报告 FID 1.82 的 CFG 值(2.2)显示图像。
图 11:在 ImageNet 256×256 上使用 JiT-G/16 以指定类别为条件的_未经筛选_样本。与使用更高 CFG 可视化的常见做法不同,这里我们使用达到报告 FID 1.82 的 CFG 值(2.2)显示图像。