论文状态:已完成

Back to Basics: Let Denoising Generative Models Denoise

发表:2025/11/18
原文链接PDF 下载
价格:0.100000
价格:0.100000
价格:0.100000
已有 12 人读过
本分析由 AI 生成,可能不完全准确,请以原文为准。

TL;DR 精炼摘要

本文提出回归根本的去噪生成模型方法,强调直接预测清晰图像而非噪声。作者基于流形假设,展示了一个简化的无预训练、无分词器的Transformer模型JiT,证明其在高维数据上表现优越,能有效生成高质量图像。

摘要

Today's denoising diffusion models do not "denoise" in the classical sense, i.e., they do not directly predict clean images. Rather, the neural networks predict noise or a noised quantity. In this paper, we suggest that predicting clean data and predicting noised quantities are fundamentally different. According to the manifold assumption, natural data should lie on a low-dimensional manifold, whereas noised quantities do not. With this assumption, we advocate for models that directly predict clean data, which allows apparently under-capacity networks to operate effectively in very high-dimensional spaces. We show that simple, large-patch Transformers on pixels can be strong generative models: using no tokenizer, no pre-training, and no extra loss. Our approach is conceptually nothing more than "Just image Transformers\textbf{Just image Transformers}", or JiT\textbf{JiT}, as we call it. We report competitive results using JiT with large patch sizes of 16 and 32 on ImageNet at resolutions of 256 and 512, where predicting high-dimensional noised quantities can fail catastrophically. With our networks mapping back to the basics of the manifold, our research goes back to basics and pursues a self-contained paradigm for Transformer-based diffusion on raw natural data.

思维导图

论文精读

中文精读

1. 论文基本信息

1.1. 标题

回到本源:让去噪生成模型真正去噪 (Back to Basics: Let Denoising Generative Models Denoise)

论文标题直截了当地提出了其核心论点:当前的“去噪”扩散模型实际上并没有在经典意义上进行“去噪”(即直接预测干净图像),而是预测噪声或与噪声相关的量。作者倡导回归到直接预测干净数据的本源,认为这才是更基本、更有效的方法。

1.2. 作者

  • Tianhong Li, Kaiming He: 两位作者均来自麻省理工学院 (MIT)。
  • Kaiming He (何恺明) 是计算机视觉领域享誉全球的杰出科学家,以其在深度学习和计算机视觉领域的开创性工作而闻名,包括残差网络 (ResNet)、Mask R-CNN、动量对比学习 (MoCo) 等。他的参与为这篇论文的质量和影响力提供了强有力的背书。

1.3. 发表期刊/会议

该论文以预印本 (preprint) 形式发布在 arXiv 上。arXiv 是一个开放获取的学术论文发布平台,允许研究者在同行评审前分享他们的研究成果。这篇论文尚未经过正式的同行评审,但通常是顶级会议(如 CVPR, ICCV, NeurIPS, ICLR)的投搞前奏。

1.4. 发表年份

论文在 arXiv 上发布的日期为 2025 年 11 月 17 日。这个未来的日期很可能是 arXiv 系统中的占位符或提交时出现的笔误,根据论文内容和引用,可以推断其为 2024 年末至 2025 年初的研究成果。

1.5. 摘要

摘要概括了论文的核心思想。当前主流的扩散模型并不直接预测干净图像,而是预测噪声 εε 或一个混合了噪声的量(如流速 vv)。本文认为,预测干净数据和预测含噪量是根本不同的任务。基于流形假设 (manifold assumption)——即自然数据位于一个低维流形上,而噪声则分布在整个高维空间——作者主张模型应该直接预测位于流形上的干净数据。这种方法使得看似容量不足的网络也能在极高维空间中有效工作。

作者展示了一个简单的、基于大图像块 (large patch) 的 Transformer 模型,它直接在像素上操作,无需分词器 (tokenizer)、无需预训练、也无需额外的损失函数,就能成为一个强大的生成模型。他们将这个极简方法命名为 JiT (Just image Transformers)。实验证明,JiT 在 256x256 和 512x512 分辨率的 ImageNet 数据集上取得了有竞争力的结果,而传统的噪声预测方法在同样设置下会灾难性地失败。这项研究旨在回归本源,为基于 Transformer 的原生数据扩散模型建立一个自洽的范式。

1.6. 原文链接

2. 整体概括

2.1. 研究背景与动机

2.1.1. 核心问题

当前最先进的扩散模型,如 DDPM 及其变体,虽然名为“去噪扩散模型”,但其核心神经网络的预测目标并非干净的图像本身。它们主要采用以下两种预测方式:

  1. ε-prediction (噪声预测): 网络预测添加到干净图像上的噪声 εε。这是 DDPM 取得成功的关键,并成为后续研究的主流。

  2. v-prediction (速度预测): 网络预测一个结合了干净数据和噪声的“流速” vv,这在连接扩散模型与流匹配 (Flow Matching) 的工作中被广泛使用。

    作者敏锐地指出,这种做法与一个经典的机器学习理论——流形假设 (Manifold Assumption)——存在潜在的矛盾。

2.1.2. 现有研究的挑战与空白 (Gap)

  • 流形假设的启示: 该假设认为,像自然图像这样的高维数据,其本质信息实际上分布在一个嵌入在高维像素空间中的低维流形上。例如,所有猫的图像,尽管像素维度很高(如 256x256x3),但它们在特征空间中构成的结构维度远低于像素维度。
  • 预测目标的内在差异:
    • 干净图像 xx: 遵循流形假设,位于低维流形上。
    • 噪声 εε 或速度 vv: 是高斯噪声或其与数据的线性组合,本质上是随机的、遍布整个高维空间的,即流形外 (off-manifold) 的。
  • 挑战 (Challenge): 要求一个神经网络去预测一个高维、随机、无结构的目标(如噪声 εε),本质上是让网络去拟合整个高维空间,这对网络的容量 (capacity) 提出了极高的要求。尤其当网络结构存在信息瓶颈时(例如 Vision Transformer 中,图像块被映射到维度可能更低的隐藏空间),强行预测高维噪声可能会导致信息丢失,从而性能下降。

2.1.3. 论文的切入点与创新思路

本文的切入点正是上述矛盾。作者提出一个简单而深刻的观点:为什么不让模型直接做它名字里说的事——去噪,即预测干净图像 xx

  • 创新思路: 如果模型的目标是预测流形上的干净数据 xx,那么网络只需要学习从一个流形外的噪声点映射回流形上的投影。这个任务本质上是降维结构恢复,理论上对网络容量的要求更低。网络可以专注于学习数据流形的内在结构,而过滤掉高维噪声。

    下图(原文 Figure 1)直观地展示了这个核心思想。干净图像 xx 位于低维流形上,而噪声 εε 和速度 vv 则散布在流形外的整个高维空间中。因此,训练网络去预测 xx (x-prediction) 与预测 εεvv (ε/v-prediction) 是两种性质完全不同的任务。

    Figure 1. The Manifold Assumption \[4\] hypothesizes that natural images lie on a low-dimensional manifold within the highdimensional pixel space. While a clean image \(_ { \\pmb { x } }\) can be modeled as on-manifold, the noise \(\\epsilon\) or flow velocity `_ { v }` (e.g., \({ \\pmb v } = { \\pmb x } - { \\pmb \\epsilon } )\) is inherently off-manifold. Training a neural network to predict a clean image (i.e., \(_ { \\textbf { \\em x } }\) -prediction) is fundamentally different from training it to predict noise or a noised quantity (i.e., \(\\epsilon / v\) -prediction). 该图像是示意图,展示了图像流形假设及其在噪声预测中的应用。图中表示了输入、干净图像预测(xx-pred)、噪声预测(vv-pred)及其关系。干净图像 x\pmb{x} 被视为流形上的点,而噪声 ϵ\epsilon 和流速 v\pmb{v} 则位于流形之外。此示意图强调预测干净图像与预测噪声之间的根本不同性。

2.2. 核心贡献/主要发现

  1. 理论贡献:重新审视了扩散模型的预测目标。 论文从流形假设的角度,系统性地论证了直接预测干净数据 (x-prediction) 在处理高维自然数据时,相比于预测噪声 (ε/v-prediction) 具有根本性优势。
  2. 方法贡献:提出了一个极简、自洽的生成模型 JiT (Just image Transformers)。
    • 架构简单: 采用一个标准的 Vision Transformer (ViT),直接在图像像素块上操作。
    • 完全自洽 (Self-contained): 无需任何外部组件,如 VAE 分词器、预训练的特征提取器(用于感知损失)、或对抗性判别器。它仅通过扩散目标进行端到端训练。
    • 高效: 在高分辨率图像生成任务上,JiT 在计算成本上远低于其他基于像素的卷积模型。
  3. 实践发现:验证了 x-prediction 在高维场景下的有效性。
    • 实验表明,当使用大尺寸图像块(导致每个词元 token 的维度非常高)时,x-prediction 能够稳定地生成高质量图像,而 ε/v-prediction 则会彻底失败。
    • 一个惊人的发现是,在网络中引入瓶颈 (bottleneck) 结构(即降低特征维度)反而能提升性能,这与经典的流形学习理论相符,进一步佐证了论文的核心假设。

3. 预备知识与相关工作

3.1. 基础概念

3.1.1. 去噪扩散模型 (Denoising Diffusion Models)

扩散模型是一类生成模型,其核心思想包含两个过程:

  • 前向过程 (Forward Process): 也称为扩散过程。从一张真实的干净图像 x0x_0 开始,通过一系列微小的步骤,逐步向其添加高斯噪声。经过足够多的步骤(比如 TT 步),图像最终会变成一张纯粹的、无规律的高斯噪声图。这个过程是固定的,不需要学习。
  • 反向过程 (Reverse Process): 也称为去噪过程。这是模型学习的核心。模型的目标是从一张纯噪声图开始,通过一系列学习到的去噪步骤,逐步地将噪声去除,最终还原出一张清晰的、看起来真实的图像。在每一步,神经网络都会接收当前的含噪图像和时间步 tt,并预测出用于去噪的信息。

3.1.2. 流形假设 (Manifold Assumption)

这是一个在机器学习和数据科学中非常基础且重要的假设。它指出,我们现实世界中观察到的高维数据(如图像、声音、文本),其内在结构实际上是低维的。这些数据点并非随机散布在整个高维空间,而是集中在一个或多个低维的流形 (manifold) 上。

  • 直观理解: 想象一下在三维空间中画一条曲线。这条曲线上的点虽然有三个坐标 (x, y, z),但它的内在维度只有一维,因为它只需要一个参数(如沿着曲线的距离)就能确定一个点的位置。这条曲线就是一个一维流形。同样,所有猫的图像虽然像素维度极高,但它们共享的“猫”的特征结构,使得它们在特征空间中形成一个低维流形。
  • 本文中的应用: 干净的自然图像 xx 被认为处于这样的低维流形上,而随机噪声 εε 则不具备任何结构,均匀地分布在整个高维像素空间中。

3.1.3. Vision Transformer (ViT)

ViT 是一种将 Transformer 架构成功应用于计算机视觉任务的模型。其核心思想是:

  1. 图像分块 (Image Patching): 将输入图像分割成一系列固定大小、不重叠的小方块,称为图像块 (patches)。例如,一张 256x256 的图像可以被分成 256 个 16x16 的图像块。

  2. 块嵌入 (Patch Embedding): 将每个图像块展平成一个一维向量,并通过一个线性层将其映射到一个固定维度的词元嵌入 (token embedding)

  3. 位置编码 (Positional Encoding): 由于 Transformer 本身不具备处理序列顺序的能力,需要向每个词元嵌入中添加位置信息,告诉模型每个图像块在原图中的位置。

  4. Transformer 编码器: 将带有位置信息的词元序列送入一个标准的 Transformer 编码器。编码器由多个堆叠的块组成,每个块包含一个多头自注意力 (Multi-Head Self-Attention) 层和一个前馈神经网络 (Feed-Forward Network) 层。自注意力机制允许每个图像块关注并整合来自所有其他图像块的信息,从而捕捉全局依赖关系。

    Attention 机制的核心计算公式如下,它计算查询 (Query) 与一组键值对 (Key-Value) 的关系: Attention(Q,K,V)=softmax(QKTdk)V \mathrm{Attention}(Q, K, V) = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V

  • Q, K, V 分别代表查询、键和值矩阵。
  • dkd_k 是键向量的维度,用于缩放,防止梯度过小。
  • softmax\mathrm{softmax} 函数将注意力权重归一化。

3.2. 前人工作

  • DDPM (Denoising Diffusion Probabilistic Models): 这篇开创性工作发现,让模型预测噪声 εε (ε-prediction) 比预测均值(与 x-prediction 相关)效果好得多,从而确立了 ε-prediction 在扩散模型领域的主导地位。
  • Progressive Distillation for Fast Sampling: 这项工作系统地分析了不同预测目标(xx, εε, vv)之间的关系,并引入了 v-prediction。它指出,这些预测目标可以通过数学变换相互转换,其主要区别在于对不同噪声水平的损失进行了不同的加权 (re-weighting)。但该研究主要在低维数据集上进行,未能揭示高维数据带来的根本性挑战。
  • EDM (Elucidating the Design Space of Diffusion-based Generative Models): EDM 提出了一种被称为预处理器 (pre-conditioner) 的网络设计,它将网络的直接输出与输入噪声图进行加权组合,得到最终的去噪结果。这种设计虽然强大,但网络的直接输出仍然是一个混合了数据和噪声的量,并未实现纯粹的 x-prediction
  • Latent Diffusion Models (LDM): 如 Stable Diffusion,是当前最主流的范式。它们通过一个预训练的 VAE 将图像压缩到一个低维的潜空间 (latent space) 中,然后在潜空间里进行扩散。这巧妙地规避了在像素空间处理高维数据的难题。然而,这种方法依赖于一个强大的、额外训练的 VAE (tokenizer),使其不是一个自洽 (self-contained) 的系统。
  • Pixel-space Diffusion Models: 在 LDM 出现之前和之后,都有研究在像素空间进行扩散。如 ADM 使用了计算量巨大的 U-Net 架构;SiD/SiD2 和 PixelFlow 等工作探索了更复杂的 Transformer 架构(如分层设计),但通常计算成本高昂,或需要辅助技术(如表示对齐)来稳定训练。

3.3. 技术演进

扩散模型的技术路线大致如下:

  1. 早期思想: 核心是去噪,预测干净数据 xx

  2. DDPM 时代: 发现 ε-prediction 效果更好,成为标准范式。

  3. 理论发展: v-prediction 和流匹配的出现,统一了扩散模型和基于流的方法。

  4. 效率革命: LDM 通过在低维潜空间进行扩散,大幅降低了计算成本,引爆了 AIGC 领域。

  5. 像素空间再探索: 一些工作坚持在像素空间进行探索,但往往需要复杂的架构和高昂的计算。

    本文的工作处于技术脉络的一个回归点。它没有沿着 LDM 的路线走,而是回到了像素空间,并指出早期 x-prediction 失败的原因可能不是像素空间本身的问题,而是模型架构和预测目标没有正确匹配。

3.4. 差异化分析

本文与相关工作的核心区别在于其理念的简洁性和方法的纯粹性

  • vs. LDM: LDM 依赖一个外部的、预训练的 VAE。JiT 则是一个端到端的、在原始像素数据上训练的单一模型,更加自洽和通用。
  • vs. 早期 x-prediction: 早期尝试失败了,但本文指出,结合现代的、具有合适归纳偏置的 Transformer 架构 (ViT),x-prediction 不仅可行,而且是处理高维数据的更优选择。
  • vs. 其他 Pixel-space 模型: 其他像素模型通常使用计算密集的卷积网络 (U-Net) 或复杂的分层 Transformer。JiT 使用的是一个标准、非分层的 ViT,结构更简单、更通用,且计算效率更高。最重要的是,其他模型大多仍采用 ε/v-prediction,而 JiT 的成功完全建立在 x-prediction 的基础上。

4. 方法论

4.1. 方法原理

本文方法论的基石是流形假设。其核心直觉是:

  • 一个容量有限的神经网络,在面对一个高维输入时,必然会丢失一部分信息,这构成了一个信息瓶颈 (information bottleneck)

  • 当任务是预测噪声 (ε/v-prediction) 时,输入 ztz_t 中的噪声 εε 是遍布整个高维空间的。为了准确预测 εεvv,网络必须尽可能地保留所有高维信息,任何信息的丢失都可能导致预测错误。因此,信息瓶颈是有害的。

  • 当任务是预测干净图像 (x-prediction) 时,理想的输出 xx 位于一个低维流形上。网络的目标是将高维输入 ztz_t 投影回这个低维流形。在这个过程中,高维的噪声信息是需要被丢弃的,而低维的流形结构信息是需要被保留的。因此,信息瓶颈不仅无害,甚至可能是有益的,因为它强迫网络学习只保留最重要的、与流形相关的信息。

    下图(原文 Figure 2)的玩具实验生动地展示了这一原理。实验中,一个 2 维的螺旋线数据被随机投影到一个更高维的 D 维空间中。模型(一个简单的 MLP)在 D 维空间中进行训练,但它并不知道这个随机投影矩阵。

  • 结果: 随着观测维度 D 的增加(从 2 到 512),ε/v-prediction 的效果迅速恶化,在 D=512 时完全失败。而 x-prediction 即使在 D=512(远超 MLP 隐藏层维度 256)的极端情况下,依然能很好地恢复出原始的 2 维螺旋结构。

    该图像是多幅水母的照片集,展示了不同种类和颜色的水母,背景色彩丰富,呈现了它们在水中漂浮和游动的状态。这些图像能够突出水母的优雅和神秘。 该图像是多幅水母的照片集,展示了不同种类和颜色的水母,背景色彩丰富,呈现了它们在水中漂浮和游动的状态。这些图像能够突出水母的优雅和神秘。

4.2. 核心方法详解 (逐层深入)

4.2.1. 扩散模型与流的统一表述

本文从流匹配 (Flow Matching) 的视角来构建扩散模型,因为它在数学上更简洁。

  1. 构造含噪样本: 在训练过程中,我们从数据分布 pdatap_{data} 中采样一个干净图像 x\pmb{x},从标准正态分布中采样一个噪声 ϵ\pmb{\epsilon}。然后,通过线性插值构造一个在时间 t[0,1]t \in [0, 1] 的含噪样本 ztz_tzt=tx+(1t)ϵ z_t = t \pmb{x} + (1 - t) \pmb{\epsilon}

    • tt: 时间变量。当 t=0t=0 时,z0=ϵz_0 = \pmb{\epsilon},样本是纯噪声。当 t=1t=1 时,z1=xz_1 = \pmb{x},样本是干净数据。
    • x\pmb{x}: 真实的干净图像。
    • ϵ\pmb{\epsilon}: 与 x\pmb{x} 维度相同的高斯噪声。
  2. 定义流速 (Flow Velocity): 流速 v\pmb{v} 定义为 ztz_t 对时间 tt 的导数。根据上式,我们可以得到一个非常简洁的表达式: v=dztdt=xϵ \pmb{v} = \frac{d z_t}{dt} = \pmb{x} - \pmb{\epsilon} 流速 v\pmb{v} 直观地表示了从噪声 ϵ\pmb{\epsilon} 指向干净数据 x\pmb{x} 的方向向量。

  3. 训练目标 (v-loss): 模型的任务是学习一个函数 vθ(zt,t)\pmb{v}_{\theta}(z_t, t) 来预测真实流速 v\pmb{v}。训练的目标是最小化预测值和真实值之间的均方误差: L=Et,x,ϵvθ(zt,t)v2 \mathcal{L} = \mathbb{E}_{t, \boldsymbol{x}, \boldsymbol{\epsilon}} \| \boldsymbol{v}_{\boldsymbol{\theta}}(\boldsymbol{z}_t, t) - \boldsymbol{v} \|^2 其中 E\mathbb{E} 表示对所有可能的时间 tt、数据 x\pmb{x} 和噪声 ϵ\pmb{\epsilon} 求期望。

4.2.2. 预测空间与损失空间的解耦

这是本文最核心的理论部分。作者指出,模型的直接输出(预测空间)损失函数的定义(损失空间)是两个可以解耦的概念。一共有 3 种预测空间(xx, εε, vv)和 3 种损失空间,组合起来有 9 种不同的训练范式。

下表(原文 Table 1)清晰地展示了这种关系。表格的列定义了网络的直接输出是什么(例如,(a)列 x-pred 表示 netθnet_θ 直接输出 xθ\pmb{x}_\theta),行定义了损失函数在哪一个空间计算。

(a) x-pred
xθ := netθ(zt, t)
(b) ε-pred
εθ := netθ(zt, t)
(c) v-pred
vθ := netθ(zt, t)
(1) x-loss: E‖xθ − x‖2 xθ xθ = (zt − (1−t)εθ)/t xθ = (1−t)vθ + zt
(2) ε-loss: E‖εθ − ε‖2 εθ = (zt − txθ)/(1−t) εθ εθ = zt − tvθ
(3) v-loss: E‖vθ − v‖2 vθ = (xθ − zt)/(1−t) vθ = (zt − εθ)/t vθ

深入解读此表:

  • 基础关系: 任何两个变量都可以通过 zt=tx+(1t)ϵz_t = t \pmb{x} + (1 - t) \pmb{\epsilon}v=xϵ\pmb{v} = \pmb{x} - \pmb{\epsilon} 这两个基础关系式来确定第三个。
  • x-prediction 为例((a)列):
    • 网络输出: netθnet_θ 直接输出对干净图像的预测 xθ\pmb{x}_\theta
    • 计算其他量: 我们可以用 xθ\pmb{x}_\theta 和输入 ztz_t 来推导出对应的噪声预测 ϵθ\pmb{\epsilon}_\theta 和速度预测 vθ\pmb{v}_\theta
      • zt=txθ+(1t)ϵθz_t = t \pmb{x}_\theta + (1-t) \pmb{\epsilon}_\theta 解出 ϵθ=(zttxθ)/(1t)\pmb{\epsilon}_\theta = (z_t - t\pmb{x}_\theta)/(1-t)。(对应表格 (2)(a))
      • vθ=xθϵθ\pmb{v}_\theta = \pmb{x}_\theta - \pmb{\epsilon}_\theta 代入上面解出的 ϵθ\pmb{\epsilon}_\theta,得到 vθ=(xθzt)/(1t)\pmb{v}_\theta = (\pmb{x}_\theta - z_t)/(1-t)。(对应表格 (3)(a))
  • 损失函数的加权效应: 不同的组合在数学上并不等价。例如,我们来分析表格中 (3)(a) 这一格,即 x-prediction + v-loss
    • 损失函数是 L=Evθv2\mathcal{L} = \mathbb{E} \| \pmb{v}_\theta - \pmb{v} \|^2
    • vθ=(xθzt)/(1t)\pmb{v}_\theta = (\pmb{x}_\theta - z_t)/(1-t)v=(xzt)/(1t)\pmb{v} = (\pmb{x} - z_t)/(1-t) 代入,得到: L=Exθzt1txzt1t2=E1(1t)2xθx2 \mathcal{L} = \mathbb{E} \left\| \frac{\pmb{x}_\theta - z_t}{1-t} - \frac{\pmb{x} - z_t}{1-t} \right\|^2 = \mathbb{E} \frac{1}{(1-t)^2} \| \pmb{x}_\theta - \pmb{x} \|^2 这表明,v-loss 作用于 x-prediction,等价于一个带 1(1t)2\frac{1}{(1-t)^2} 权重项的 x-loss。这个权重项会放大在 tt 接近 1(低噪声)时的误差,这与先前研究中发现的有利加权方案是一致的。但本文强调,即使加权很重要,但让网络直接输出 xx 才是解决高维问题的关键

4.2.3. JiT 架构:Just Image Transformers

作者提出的 JiT 架构非常简洁,其核心是“在像素块上应用 Transformer (Transformer on Patches)”。 下图(原文 Figure 3)展示了其结构:

该图像是一个展示多种鸟屋的插图,呈现了不同形状和颜色的鸟屋设计。这些鸟屋以创意和平面布置的方式排列,展现了丰富多彩的鸟类栖息环境。 该图像是一个展示多种鸟屋的插图,呈现了不同形状和颜色的鸟屋设计。这些鸟屋以创意和平面布置的方式排列,展现了丰富多彩的鸟类栖息环境。

  1. 输入: 一张含噪图像 ztz_t 被切分成多个不重叠的 p×pp \times p 像素块。例如,对于 256x256 的图像和 p=16p=16 的块大小,会得到 (256/16)×(256/16)=16×16=256(256/16) \times (256/16) = 16 \times 16 = 256 个块。每个块是一个 16×16×3=76816 \times 16 \times 3 = 768 维的向量。
  2. 线性嵌入: 每个 768 维的向量通过一个线性层映射到 Transformer 的隐藏维度(例如,JiT-B 模型中也是 768 维)。同时加入位置编码。
  3. Transformer 模块: 嵌入后的块序列被送入标准的 Transformer 编码器。
  4. 线性预测器: Transformer 的输出序列再通过一个线性层,将每个词元 token 映射回原始的 p×p×3p \times p \times 3 维像素块。
  5. 输出: 将所有预测的像素块拼接起来,就得到了网络对干净图像的预测 xθ=netθ(zt,t)\pmb{x}_\theta = \mathrm{net}_\theta(z_t, t)
  • 条件注入: 时间 tt 和类别标签等条件信息通过 adaLN-Zero 技术注入到 Transformer 块中,这是一种动态调整归一化层参数的方法。

4.2.4. 最终算法

综合以上分析,作者最终选择的算法是**x-prediction + v-loss** 的组合(表格 1 中的 (3)(a)),因为它在实验中表现最好。

训练步骤(原文 Algorithm 1 伪代码解析):

# net(z, t): JiT network
# x: training batch

t = sample_t()              # 随机采样一个时间步 t
e = randn_like(x)           # 生成一个与 x 同样大小的标准高斯噪声
z = t * x + (1 - t) * e     # 构造含噪样本 z_t

x_pred = net(z, t)          # 网络 net 直接预测干净图像 x_pred
v_pred = (x_pred - z) / (1 - t) # 根据 x_pred 计算出对应的流速 v_pred
v_true = x - e              # 计算真实的流速 v_true

loss = l2_loss(v_pred - v_true) # 在 v 空间计算 L2 损失
# ... (反向传播,更新网络参数) ...

采样步骤(原文 Algorithm 2 伪代码解析,以最简单的欧拉法为例): 采样过程是通过求解一个常微分方程 (ODE) 来实现的:dzt/dt=vθ(zt,t)d z_t / dt = \pmb{v}_\theta(z_t, t)

# z: current samples at t (初始时 z_0 是纯高斯噪声)
# 从 t=0 到 t=1 分成 N 步

for i from 0 to N-1:
    t = i / N
    t_next = (i + 1) / N
    
    x_pred = net(z, t)               # 网络从当前噪声图像 z 预测出干净图像 x_pred
    v_pred = (x_pred - z) / (1 - t)  # 根据 x_pred 计算出当前步的流速 v_pred
    
    # 欧拉法更新:沿着 v_pred 的方向走一小步
    z_next = z + (t_next - t) * v_pred 
    z = z_next

# 循环结束后,z 就是生成的图像

5. 实验设置

5.1. 数据集

  • ImageNet: 实验主要在 ImageNet (ILSVRC 2012) 数据集上进行。这是一个包含超过 120 万张训练图像、5 万张验证图像的大规模图像分类数据集,涵盖 1000 个类别。由于其规模和多样性,它已成为衡量生成模型性能的标准基准之一。

  • 分辨率: 实验在 256x256、512x512 和 1024x1024 三种分辨率上进行,以测试模型在高维像素空间中的性能和可扩展性。

  • 样本示例: 下图(原文 Figure 9)展示了在 ImageNet 数据集上生成的样本,类别为“豹”(leopard)、“蜜蜂”(bee)、“书柜”(bookcase)等,直观地展示了数据的多样性。

    该图像是插图,展示了不同姿态和角度的豹,包括其特征斑点和自然栖息环境。这些图像强调了豹子的多样性和优雅性。 该图像是插图,展示了不同姿态和角度的豹,包括其特征斑点和自然栖息环境。这些图像强调了豹子的多样性和优雅性。

    该图像是插图,展示了多种蜜蜂和花朵的特写图像,呈现出自然界中生物的多样性和生态关系。 该图像是插图,展示了多种蜜蜂和花朵的特写图像,呈现出自然界中生物的多样性和生态关系。

    该图像是书架的集合,展示了不同类型和样式的书架,书籍整齐摆放。每个书架的设计和排列各具特色,体现出不同的存储和展示方式。 该图像是书架的集合,展示了不同类型和样式的书架,书籍整齐摆放。每个书架的设计和排列各具特色,体现出不同的存储和展示方式。

5.2. 评估指标

5.2.1. Fréchet Inception Distance (FID)

  • 概念定义: FID 是评估生成模型性能最常用的指标之一。它通过比较生成图像集与真实图像集在 Inception-V3 网络(一个预训练的图像分类模型)的特征空间中的统计分布来衡量二者之间的相似度。FID 分数越低,表示生成图像的质量和多样性越接近真实图像。它同时考虑了保真度(图像是否清晰真实)和多样性(能否生成多种多样的图像)。
  • 数学公式: FID(x,g)=μxμg22+Tr(Σx+Σg2(ΣxΣg)1/2) \mathrm{FID}(x, g) = \|\mu_x - \mu_g\|^2_2 + \mathrm{Tr}\left(\Sigma_x + \Sigma_g - 2(\Sigma_x\Sigma_g)^{1/2}\right)
  • 符号解释:
    • xxgg 分别代表真实图像分布和生成图像分布。
    • μx\mu_xμg\mu_g 是真实图像和生成图像在 Inception 网络某一层的激活特征的均值向量。
    • Σx\Sigma_xΣg\Sigma_g 是这些特征的协方差矩阵。
    • 22\|\cdot\|^2_2 表示向量的平方欧几里得范数。
    • Tr()\mathrm{Tr}(\cdot) 表示矩阵的迹(对角线元素之和)。

5.2.2. Inception Score (IS)

  • 概念定义: IS 是另一个经典的生成模型评估指标。它主要衡量两个方面:1) 生成图像的清晰度/质量:对于一张清晰的图像,分类网络应该能以高置信度将其分到某一类(即条件概率分布 p(yx)p(y|x) 的熵低);2) 生成图像的多样性:模型应该能生成所有类别的图像,且各类别的数量均衡(即边缘概率分布 p(y) 的熵高)。IS 结合了这两点,分数越高越好。
  • 数学公式: IS(G)=exp(ExpgDKL(p(yx)p(y))) \mathrm{IS}(G) = \exp\left(\mathbb{E}_{x \sim p_g} D_{KL}(p(y|x) \| p(y))\right)
  • 符号解释:
    • xpgx \sim p_g 表示从生成器 G 生成的图像。
    • p(yx)p(y|x) 是 Inception 网络对生成图像 xx 的类别预测概率分布。
    • p(y)=p(yx)pg(x)dxp(y) = \int p(y|x)p_g(x)dx 是所有生成图像的平均类别概率分布(边缘分布)。
    • DKL()D_{KL}(\cdot \| \cdot) 是 KL 散度,用于衡量两个概率分布的差异。

5.2.3. Precision and Recall

  • 概念定义: 这组指标旨在更精细地解构生成模型的性能,将保真度 (fidelity) 和多样性 (diversity) 分开评估。
    • Precision (精确率): 衡量生成图像的真实性。高精确率意味着生成样本中有很高比例是位于真实数据流形上的。
    • Recall (召回率): 衡量生成模型覆盖真实数据分布的能力。高召回率意味着模型能够生成真实数据中绝大多数种类的样本。
  • 计算方式: 同样是在一个预训练模型的特征空间中,通过计算真实样本和生成样本特征分布之间的重叠情况来量化。

5.3. 对比基线

论文将 JiT 与两类最先进的生成模型进行了比较:

  1. 潜空间扩散模型 (Latent-space Diffusion):
    • DiT (Diffusion Transformer): 第一个将 Transformer 成功用于潜空间扩散的模型。
    • SiT, DDT, RAE: DiT 的后续改进版本,通常结合了更先进的架构或自监督预训练技术。
    • 特点: 这些模型性能强大,但都依赖于一个预训练的 VAE (tokenizer) 和其他辅助损失(如感知损失)。
  2. 像素空间扩散模型 (Pixel-space Diffusion):
    • ADM (Ablated Diffusion Model): 基于 U-Net 的强大像素空间模型,但计算成本极高。
    • SiD/SiD2SiD / SiD2: 尝试用 Transformer 在像素空间建模,但架构比 JiT 更复杂。
    • PixelFlow, PixNerd: 其他先进的像素空间模型,通常也依赖复杂的架构或外部预训练。
    • 特点: 这些是 JiT 的直接竞争者,但 JiT 在架构简单性、计算效率和自洽性上具有明显优势。

6. 实验结果与分析

6.1. 核心结果分析

6.1.1. 预测目标是关键 (Ablation Study on Prediction Target)

这是论文中最具说服力的实验。作者在 ImageNet 256x256 分辨率上,使用 JiT-B/16 模型(块大小 16x16)测试了 Table 1 中所有 9 种预测与损失组合。

关键设置: 在这个设置下,每个图像块的维度是 16×16×3=76816 \times 16 \times 3 = 768 维,这与 JiT-B 模型的隐藏维度 768 恰好相等。这是一个典型的信息瓶颈场景。

以下是原文 Table 2(a) 的结果:

x-pred ε-pred v-pred
x-loss 10.14 379.21 107.55
ε-loss 10.45 394.58 126.88
v-loss 8.62 372.38 96.53
  • 结果分析:
    • x-prediction (第一列): 无论使用哪种损失函数,模型都能稳定训练并取得非常好的结果 (FID 约 8-10)。其中 v-loss 效果最好,FID 达到 8.62。
    • ε-predictionv-prediction (第二、三列): 这两种预测方法都灾难性地失败了。ε-prediction 的 FID 高达 370-390,生成的图像是纯粹的噪声。v-prediction 稍好一些,但 FID 仍在 100 左右,远非可用。
  • 结论: 这个实验清晰地证明了,当网络存在信息瓶颈时(隐藏维度 ≈ 输入词元维度),直接预测流形外的 εεvv 会导致模型无法学习,而预测流形上的 xx 则完全可行。这强有力地支持了论文的核心假设。

6.1.2. 引入瓶颈反而有利 (Bottleneck Analysis)

为了进一步验证流形假设,作者在 x-prediction 的基础上,对 JiT-B/16 模型的输入嵌入层进行了修改,将其从一个 768 -> 768 的线性层,变成了一个 768 -> dd' -> 768 的瓶颈结构,其中 d<768d' < 768

下图(原文 Figure 4)展示了 FID 分数随瓶颈维度 dd' 变化的情况:

该图像是书架的集合,展示了不同类型和样式的书架,书籍整齐摆放。每个书架的设计和排列各具特色,体现出不同的存储和展示方式。 该图像是书架的集合,展示了不同类型和样式的书架,书籍整齐摆放。每个书架的设计和排列各具特色,体现出不同的存储和展示方式。

  • 结果分析:
    • 令人惊讶的提升: 当引入瓶颈,并将维度 dd' 从 768 降低时,FID 分数不升反降。在 dd' 约为 32 到 512 的很宽范围内,性能都优于没有瓶颈的基线模型。最佳性能出现在 d128d' \approx 128 左右。
    • 鲁棒性: 即使将维度压缩到极低的 32 甚至 16,模型性能也只是略有下降,并未崩溃。
  • 结论: 这个结果非常有力。它表明网络确实在学习一个低维表示。强制通过一个低维瓶颈,相当于一个强烈的正则化,迫使模型过滤掉无关的高维噪声信息,专注于学习数据流形的内在结构,从而提升了生成质量。这与经典流形学习和自编码器的思想不谋而合。

6.1.3. 模型可扩展性 (Scalability)

JiT 的一个巨大优势是其架构和分辨率解耦。通过调整块大小,可以在不同分辨率下使用几乎相同计算量的模型。

以下是原文 Table 6 的结果,展示了不同尺寸的模型 (B, L, H, G) 在 256 和 512 分辨率下的表现:

256×256 512×512
200-ep 600-ep 200-ep 600-ep
JiT-B/16 4.37 3.66 JiT-B/32 4.64 4.02
JiT-L/16 2.79 2.36 JiT-L/32 3.06 2.53
JiT-H/16 2.29 1.86 JiT-H/32 2.51 1.94
JiT-G/16 2.15 1.82 JiT-G/32 2.11 1.78
  • 结果分析:
    • 模型尺寸扩展: 从 B (Base) 到 G (Giant) 模型,FID 稳步下降,表明 JiT 能够很好地从增加的模型容量中受益。
    • 分辨率扩展: 对于同一尺寸的模型,从 256 分辨率 (JiT-X/16) 切换到 512 分辨率 (JiT-X/32),计算量几乎不变,而 FID 分数也保持在同一水平,甚至在 G 模型上 512 分辨率的 FID (1.78) 还优于 256 分辨率 (1.82)。这证明了该方法对高维观测空间的鲁棒性。
  • 结论: JiT 方法具有良好的可扩展性,无论是模型大小还是图像分辨率,都表现出稳定且可预测的性能提升。

6.1.4. 与最先进模型的比较

以下是原文 Table 8 中 512x512 分辨率下与 SOTA 模型的比较结果:

ImgNet 512×512 pre-training params Gflops FID↓ IS↑
token perc. self-sup.
Latent-space Diffusion
DiT-XL/2 [46] SD-VAE VGG 675+49M 525 3.04 240.8
RAE [78], DiTDH-XL/2 RAE VGG DINOv2 839+415M 642 1.13 259.6
Pixel-space Diffusion
ADM-G [12] - - 559M 1983 7.72 172.7
SiD2 [26], UViT/2 - - N/A 653 1.48 -
JiT-H/32 - - 956M 183 1.94 309.1
JiT-G/32 - - 2B 384 1.78 306.8
  • 结果分析:
    • vs. Latent-space: 虽然最强的潜空间模型(如 RAE,FID=1.13)在 FID 上仍然领先,但它们严重依赖外部预训练模型(RAE, VGG, DINOv2)。JiT 是完全自洽的,不依赖任何外部模型。
    • vs. Pixel-space: JiT 的表现非常出色。例如,JiT-G/32 (FID=1.78) 显著优于计算成本极高的 ADM-G (FID=7.72, Gflops=1983),并与复杂的 SiD2 (FID=1.48) 相当,但计算成本 (384 Gflops) 远低于 SiD2 (653 Gflops)。
  • 结论: JiT 在完全不使用预训练和额外损失的情况下,以更低的计算成本,达到了与其他先进的、更复杂的像素空间模型相当甚至更好的性能,展示了其方法的简洁性、高效性和有效性

7. 总结与思考

7.1. 结论总结

这篇论文通过回归“去噪”的本源,对扩散模型的预测目标进行了深刻的重新审视。其主要结论可以总结为:

  1. 预测目标至关重要: 基于流形假设,直接预测位于低维流形上的干净图像 (x-prediction),比预测遍布高维空间的噪声 (ε/v-prediction),是一个更适定 (well-posed)、更容易学习的任务,尤其是在网络容量受限的情况下。
  2. 简单即是强大: 一个简单的、标准的 Vision Transformer (JiT),只要采用了正确的预测目标 (x-prediction),就可以成为一个非常强大的像素空间生成模型。
  3. 自洽范式的潜力: JiT 的成功表明,我们可以在不依赖复杂外部组件(如 VAE tokenizer、预训练分类器等)的情况下,实现高质量的图像生成。这为构建更通用、更简洁的生成模型范式铺平了道路。

7.2. 局限性与未来工作

尽管论文取得了显著成功,但仍存在一些潜在的局限性和值得探索的方向:

  • 性能上限: 从实验结果看,最顶尖的、依赖大量预训练的潜空间模型(如 RAE)在 FID 指标上仍然有优势。JiT 虽然高效简洁,但在绝对性能上尚未达到 SOTA。未来的工作可以探索如何将 JiT 的简洁理念与其他技术(如本文附录中提到的额外分类损失)温和地结合,以进一步提升性能。
  • 数据类型的普适性: 论文主要在自然图像上进行了验证。x-prediction 的优势是否能推广到其他类型的高维数据(如音频、视频、科学模拟数据),还有待进一步研究。
  • 采样效率: 论文使用了标准的 ODE 求解器进行采样,通常需要几十个步骤。如何结合更快的采样方法是提升其实用性的一个重要方向。

7.3. 个人启发与批判

  • 启发:
    1. 回归第一性原理: 这篇论文是“回归第一性原理”思考的绝佳范例。当一个领域的技术变得越来越复杂时,回头审视那些最基本的假设(如“去噪模型应该预测什么?”),往往能带来颠覆性的突破。
    2. “奥卡姆剃刀”原则: JiT 的成功完美诠释了“如无必要,勿增实体”的原则。它表明,许多现有模型中的复杂组件,可能只是为了弥补一个 изначально 不那么优化的任务设定(即 ε-prediction)。
    3. 对通用人工智能的意义: JiT 这种自洽、不依赖特定领域知识(如 VAE tokenizer)的模型,更接近一个通用的学习范式。它在处理那些难以设计 tokenizer 的新领域数据时,可能具有巨大的潜力,例如在科学发现中处理原始传感器数据。
  • 批判性思考:
    • “灾难性失败”的根源: 论文将 ε/v-prediction 的失败归因于网络容量无法处理高维输出。这是一个非常合理的解释,但可能还存在其他优化上的困难。例如,εε 作为一个高频、无结构的目标,其损失曲面可能非常复杂,导致优化器难以收敛。
    • 潜空间的价值: 本文强调了像素空间的优势,但这并不意味着潜空间没有价值。一个精心设计的、能够解耦变异因素的潜空间,在可控生成、风格编辑等任务上依然具有不可替代的优势。未来的最佳模型,或许是能结合 JiT 的 x-prediction 思想和潜空间结构化优势的混合体。
    • 效率的全面考量: 论文中的 Gflops 对比的是单次前向传播的计算量。一个全面的效率比较还应考虑总的训练时间、收敛速度等因素。虽然 JiT 的 Gflops 很低,但其训练达到同样 FID 所需的总计算量与 SOTA 潜空间模型相比如何,是一个值得进一步探讨的问题。

相似论文推荐

基于向量语义检索推荐的相关论文。

暂时没有找到相似论文。