从噪声到艺术:深入浅出Diffusion Transformer

发布于:2025-06-28 ⋅ 阅读:(19) ⋅ 点赞:(0)

1. 引言:AI绘画背后的"魔法师"

当你第一次看到Sora生成的逼真视频,或者被DALL-E创造的奇幻图像震撼时,你是否好奇过这些AI是如何"凭空"创造出如此精美的视觉内容的?在这些令人惊叹的AI应用背后,有一位默默无闻但功不可没的"魔法师"——Diffusion Transformer(DiT)。

想象一下,如果有人告诉你,计算机可以像艺术家一样从一片混沌的噪声中"雕刻"出美丽的图像,你会相信吗?这听起来像是科幻小说中的情节,但这正是Diffusion Transformer正在做的事情。它不是简单地拼接已有的图像片段,而是真正理解了图像的本质,能够从随机噪声开始,一步步"去噪",最终创造出前所未见的视觉作品。

这项技术的重要性远不止于创造美丽的图像。它代表了人工智能在理解和生成视觉内容方面的一个重大突破。从传统的基于模板的图像处理,到今天能够理解语义、遵循复杂指令的智能生成系统,DiT标志着我们正在进入一个全新的AI时代。在这个时代里,机器不再是简单的工具,而是具备了某种"创造力"的伙伴。

本文将带你从最基础的概念开始,逐步深入DiT的技术内核。我们会用最通俗易懂的语言,配合生动的比喻和实际的例子,让你不仅知其然,更知其所以然。准备好了吗?让我们一起揭开Diffusion Transformer的神秘面纱。

2. 基础概念:什么是扩散模型?

2.1 生活中的扩散现象

要理解扩散模型,我们先从一个每个人都能理解的现象开始:想象你在一杯清水中滴入一滴墨水。你会看到什么?墨水不会保持原来的形状,而是慢慢地向四周扩散,最终整杯水都变成了淡淡的灰色。这个过程是不可逆的——你无法简单地把扩散开的墨水重新聚集成原来的一滴。

这个简单的物理现象蕴含着深刻的数学原理。在物理学中,这被称为扩散过程,它遵循热力学第二定律:系统总是趋向于从有序状态向无序状态发展。有序的墨滴变成了无序的灰水,信息似乎"丢失"了。

但是,如果我们能够精确记录扩散过程中的每一个细节,理论上是可以逆转这个过程的。这就是扩散模型的核心思想:如果我们能够学会如何"逆转"扩散过程,就能从看似随机的噪声中恢复出有意义的信息。

在图像生成的语境下,我们可以把一张图片想象成那滴墨水,把随机噪声想象成扩散后的灰水。扩散模型要做的,就是学会如何从"灰水"中恢复出原来的"墨滴"——也就是从噪声中生成清晰的图像。

2.2 扩散模型的核心思想

扩散模型的工作原理可以分为两个相互对应的过程:前向过程和反向过程。

前向过程:从图像到噪声

前向过程就像我们刚才描述的墨水扩散。给定一张原始图像,我们逐步向其中添加高斯噪声。这个过程不是一步完成的,而是分成很多小步骤。每一步都只添加一点点噪声,就像墨水扩散的每一个瞬间。

想象你有一张清晰的照片,第一步,你给它加上一层非常轻微的"雪花"噪声,照片变得稍微模糊了一点,但仍然能够清楚地看出内容。第二步,再加一层噪声,照片变得更模糊了。如此反复,经过几百步之后,原本清晰的照片就变成了一片看似随机的噪声,就像完全扩散后的墨水一样。

这个过程的数学表达虽然复杂,但直观理解很简单:我们在每一步都知道添加了多少噪声,以及添加噪声的方式。这些信息为后续的"逆转"过程提供了重要的指导。

反向过程:从噪声到图像

反向过程是扩散模型的核心,也是最困难的部分。它要做的是学会如何从噪声中一步步"去噪",最终恢复出清晰的图像。这就像是一个超级复杂的"修复破损照片"任务。

但这里有一个关键的洞察:我们不需要恢复出原来的那张特定图片,而是要生成一张新的、但同样真实的图片。就像艺术家不是在复制照片,而是在创作新的作品。

神经网络在这里扮演了"艺术家"的角色。它需要学会在每一个去噪步骤中,预测应该去除多少噪声,以及去除哪些噪声。这个学习过程需要大量的训练数据和计算资源,但一旦训练完成,网络就获得了从噪声"雕刻"出图像的能力。

用一个更形象的比喻:如果说前向过程是把一座精美的雕塑逐渐埋在沙子里,那么反向过程就是学会如何从沙堆中重新挖出雕塑。不同的是,我们挖出的不是原来的那座雕塑,而是一座全新的、但同样精美的艺术品。

2.3 DDPM:扩散模型的奠基之作

2020年,来自UC Berkeley的研究者Jonathan Ho、Ajay Jain和Pieter Abbeel发表了一篇名为《Denoising Diffusion Probabilistic Models》的论文[1],这篇论文奠定了现代扩散模型的理论基础。DDPM(去噪扩散概率模型)不仅提供了扩散过程的严格数学描述,更重要的是证明了这种方法在图像生成任务上的有效性。

DDPM的核心贡献在于将扩散过程形式化为一个马尔可夫链。简单来说,就是将连续的扩散过程离散化为一系列有限的步骤,每一步的状态只依赖于前一步的状态。这种形式化使得我们可以用计算机来精确模拟和控制整个过程。

在数学层面,DDPM定义了两个关键的概率分布:一个描述如何从图像逐步添加噪声(前向过程),另一个描述如何从噪声逐步恢复图像(反向过程)。虽然前向过程是确定的,但反向过程需要通过神经网络来学习。

DDPM的另一个重要贡献是提出了一种有效的训练目标。传统的生成模型往往需要复杂的对抗训练或者变分推断,而DDPM将训练简化为一个相对简单的去噪任务:给定一个加了噪声的图像,网络需要预测原始图像是什么样的。这种训练方式不仅更稳定,而且更容易理解和实现。

为什么扩散模型如此有效?原因有几个方面。首先,逐步去噪的过程允许模型在每一步只需要做出小的修正,这比一次性生成完整图像要容易得多。其次,这种方法天然地避免了模式崩塌问题——即生成模型只能生成有限几种样本的问题。最后,扩散过程的随机性确保了生成结果的多样性,每次从不同的噪声开始,都能得到不同的结果。

DDPM的成功不仅在于其理论的优雅,更在于其实际效果的显著。在多个图像生成基准测试中,DDPM都取得了当时最好的结果,特别是在生成图像的质量和多样性方面。这为后续的研究奠定了坚实的基础,也为DiT的出现铺平了道路。

3. Transformer的崛起:从语言到视觉

3.1 Transformer简介

如果说扩散模型解决了"如何生成"的问题,那么Transformer则解决了"如何理解"的问题。要理解DiT的革命性意义,我们必须先了解Transformer这个深度学习领域的"游戏规则改变者"。

Transformer最初是为自然语言处理而设计的。想象一下,当你阅读一篇文章时,你的大脑是如何工作的?你不会逐字逐句地线性处理,而是会同时关注多个词语之间的关系,理解上下文的含义。比如在句子"银行的利率上升了"中,"银行"这个词的含义需要结合"利率"来理解,而不是孤立地看待。

这正是Transformer的核心思想——注意力机制(Attention Mechanism)。如果把注意力机制比作一盏聚光灯,那么在处理每个词语时,这盏聚光灯会照亮所有相关的词语,让模型能够"看到"它们之间的关系。更神奇的是,这盏聚光灯不是固定的,而是可以学习的——模型会自动学会在什么时候关注什么内容。

传统的循环神经网络(RNN)就像是一个只能从左到右阅读的人,必须按顺序处理每个词语。而Transformer则像是一个可以同时看到整个句子的人,能够并行地理解所有词语之间的关系。这种并行性不仅大大提高了处理效率,更重要的是让模型能够捕捉到长距离的依赖关系。

Transformer的另一个关键创新是多头注意力(Multi-Head Attention)。这就像是给模型配备了多盏不同的聚光灯,每盏灯都专注于不同类型的关系。有的灯专门关注语法关系,有的专门关注语义关系,有的专门关注情感色彩。通过组合这些不同的"视角",模型能够形成对文本更全面、更深入的理解。

3.2 Vision Transformer (ViT)

Transformer在自然语言处理领域的巨大成功自然引发了一个问题:这种强大的架构能否应用到计算机视觉领域?2020年,Google的研究团队给出了肯定的答案,他们提出了Vision Transformer(ViT)[2],将Transformer成功地应用到了图像分类任务上。

但是,图像和文本有一个根本的不同:文本天然地由离散的词语组成,而图像是连续的像素矩阵。如何将图像"翻译"成Transformer能够理解的"语言"呢?

ViT的解决方案既简单又巧妙:将图像切分成固定大小的小块(patches),然后将每个小块当作一个"词语"来处理。想象你有一幅拼图,你把它拆分成许多小块,然后将每个小块展平成一个向量。这样,一张图像就变成了一个"句子",每个patch就是这个句子中的一个"词语"。

这种方法的美妙之处在于,它保留了图像的空间信息,同时又让Transformer能够处理。每个patch不仅包含了局部的视觉信息,patch之间的注意力机制还能捕捉到全局的空间关系。就像阅读文章时需要理解词语之间的关系一样,ViT通过注意力机制理解不同图像区域之间的关系。

为了让模型知道每个patch在原图像中的位置,ViT还引入了位置编码(Positional Encoding)。这就像给每个拼图块标上坐标,确保模型在重新组合时不会搞错位置。

ViT的成功证明了一个重要观点:Transformer的强大不仅仅局限于语言,它是一种通用的序列建模架构。无论是文本中的词语序列,还是图像中的patch序列,Transformer都能有效地建模它们之间的复杂关系。

3.3 从U-Net到Transformer

在ViT证明了Transformer在视觉任务上的有效性之后,一个自然的问题出现了:既然Transformer在图像理解方面如此出色,为什么不将它应用到图像生成任务上呢?

传统的扩散模型,包括DDPM,都使用U-Net作为去噪网络的骨干架构。U-Net是一种专门为图像分割任务设计的卷积神经网络,它的特点是具有编码器-解码器结构,并且在对应层之间有跳跃连接。这种设计使得U-Net能够很好地保留图像的细节信息,因此在图像生成任务上表现出色。

但是,U-Net也有其局限性。首先,卷积操作本质上是局部的,虽然通过堆叠多层可以扩大感受野,但在捕捉长距离依赖关系方面仍然有限。其次,U-Net的架构相对固定,难以像Transformer那样通过简单地增加层数或宽度来扩展模型规模。

相比之下,Transformer具有几个显著的优势。首先是可扩展性:Transformer的架构非常规整,可以通过增加层数、增加注意力头数、或者增加隐藏维度来轻松扩展模型。这种扩展性在大模型时代显得尤为重要。其次是全局建模能力:注意力机制天然地能够捕捉任意两个位置之间的关系,这对于理解图像的全局结构非常有帮助。

更重要的是,Transformer在其他领域(特别是自然语言处理)已经展现出了惊人的扩展规律:模型越大,效果越好。这种规律被称为"scaling law",它暗示着只要有足够的计算资源和数据,我们就能通过简单地增大模型来获得更好的性能。

这些观察为DiT的诞生提供了强有力的动机:如果我们能够将Transformer成功地应用到扩散模型中,不仅能够获得更好的生成质量,还能享受到Transformer优秀的扩展性。这正是DiT要解决的核心问题。

技术演进往往有其内在的逻辑。从RNN到Transformer,从CNN到ViT,再从U-Net到DiT,每一次转变都不是偶然的,而是技术发展的必然结果。当一种新的架构在某个领域证明了其优越性,将其推广到其他领域就成为了自然而然的选择。DiT正是这种技术演进逻辑的产物,它代表了扩散模型发展的下一个阶段。

4. Diffusion Transformer:两大技术的完美融合

4.1 DiT的诞生背景

2022年12月,来自UC Berkeley的William Peebles和来自New York University的Saining Xie发表了一篇题为《Scalable Diffusion Models with Transformers》的论文[3],正式提出了Diffusion Transformer(DiT)。这篇论文不仅仅是一个技术改进,更是一个范式转换的标志。

DiT的提出并非偶然。在2022年,AI领域正经历着一场"大模型革命"。GPT-3的成功让人们看到了大规模Transformer模型的巨大潜力,而ViT在计算机视觉领域的突破则证明了Transformer的通用性。在这样的背景下,将Transformer应用到扩散模型中成为了一个自然而然的想法。

但是,想法容易,实现困难。扩散模型有其特殊性:它需要处理条件输入(如时间步和类别标签),需要在每个去噪步骤中做出精确的预测,还需要保持生成结果的高质量。如何将Transformer优雅地融入到这个框架中,是DiT需要解决的核心挑战。

Peebles和Xie的贡献不仅在于提出了一个可行的解决方案,更在于他们系统地研究了DiT的扩展性质。他们发现,DiT完美地继承了Transformer的扩展规律:更大的模型、更多的计算量,带来更好的生成质量。这一发现为后续的研究指明了方向,也为Sora等大规模应用奠定了理论基础。

4.2 DiT的核心架构

DiT的整体架构可以看作是ViT在扩散模型上的巧妙适配。让我们从输入到输出,逐步解析DiT是如何工作的。

输入处理:从图像到Token

DiT的第一步是将输入图像转换为Transformer能够处理的token序列。这个过程通过一个叫做"patchify"的操作来完成。想象你有一张256×256像素的图像,DiT会将其切分成许多16×16像素的小块(patch)。每个patch被展平成一个768维的向量,这样一张图像就变成了256个token的序列。

但这里有一个重要的细节:DiT实际上不是直接处理原始图像,而是处理图像的潜在表示(latent representation)。这是通过一个预训练的变分自编码器(VAE)来实现的。VAE将原始图像压缩到一个低维的潜在空间中,这不仅减少了计算量,还提高了训练的稳定性。

Transformer Blocks:信息处理的核心

DiT的主体由多个Transformer block组成,每个block包含多头自注意力层和前馈网络层。这些组件与标准的Transformer基本相同,但DiT在如何处理条件输入方面做了重要的创新。

在扩散模型中,网络需要知道当前处于扩散过程的哪一步(时间步t),以及要生成什么类型的图像(类别标签c)。这些条件信息对于正确的去噪至关重要。DiT通过一种叫做"自适应层归一化"(Adaptive Layer Normalization, adaLN)的机制来注入这些条件。

条件输入的巧妙处理:adaLN的设计

adaLN是DiT最重要的创新之一。传统的层归一化(Layer Normalization)是这样工作的:对于输入 x x x,计算其均值 μ μ μ和方差 σ 2 σ² σ2,然后进行标准化:

y = ( x − μ ) / σ y = (x - μ) / σ y=(xμ)/σ

而adaLN则在此基础上引入了条件信息。它不是使用固定的缩放和偏移参数,而是根据条件输入(时间步和类别)动态地生成这些参数:

y = γ ( t , c ) ∗ ( x − μ ) / σ + β ( t , c ) y = γ(t,c) * (x - μ) / σ + β(t,c) y=γ(t,c)(xμ)/σ+β(t,c)

这里, γ ( t , c ) γ(t,c) γ(t,c) β ( t , c ) β(t,c) β(t,c)是根据时间步 t t t和类别 c c c计算出的缩放和偏移参数。这种设计的巧妙之处在于,它让每个Transformer block都能够"感知"到当前的条件,从而做出相应的调整。

更重要的是,DiT采用了一种特殊的初始化策略:所有的adaLN层都被初始化为恒等函数。这意味着在训练开始时,每个Transformer block都不会改变输入,整个网络从一个"什么都不做"的状态开始学习。这种初始化方式大大提高了训练的稳定性。

4.3 关键技术创新

DiT在技术上有几个重要的创新点,每一个都对最终的性能有着关键的影响。

三种条件输入方式的探索

在设计DiT时,研究者们探索了三种不同的条件输入方式:

  1. In-context conditioning:将条件信息作为额外的token添加到输入序列中。这就像在句子中添加一个特殊的词语来表示上下文。

  2. Cross-attention:使用交叉注意力机制,让图像token能够"关注"到条件信息。这类似于机器翻译中源语言和目标语言之间的注意力。

  3. Adaptive Layer Normalization (adaLN):如前所述,通过动态调整归一化参数来注入条件信息。

通过大量的实验,研究者们发现adaLN是最有效的方法。它不仅性能最好,而且实现简单,计算开销小。这个发现对后续的研究产生了重要影响。

残差连接的特殊处理

在标准的Transformer中,残差连接是这样的:

output = x + Attention(LayerNorm(x))

而在DiT中,adaLN不仅作用于注意力层的输入,还作用于残差连接之前的激活:

output = x + adaLN(Attention(adaLN(x)))

这种设计确保了条件信息能够影响到网络的每一个部分,而不仅仅是某些特定的层。

可扩展的架构设计

DiT采用了与ViT相同的扩展策略,定义了四种不同规模的模型:

  • DiT-S (Small):12层,384维隐藏层,6个注意力头
  • DiT-B (Base):12层,768维隐藏层,12个注意力头
  • DiT-L (Large):24层,1024维隐藏层,16个注意力头
  • DiT-XL (Extra Large):28层,1152维隐藏层,16个注意力头

这种规整的扩展方式使得研究者们能够系统地研究模型规模对性能的影响。

4.4 扩展性分析

DiT最重要的发现之一是它完美地继承了Transformer的扩展规律。研究者们从两个维度来研究扩展性:模型大小和输入token数量。

模型大小的扩展

通过比较不同规模的DiT模型,研究者们发现了一个清晰的规律:更大的模型总是带来更好的性能。从DiT-S的33M参数到DiT-XL的675M参数,FID分数(越低越好)持续下降,生成质量持续提升。

更有趣的是,这种改进不仅仅是参数数量的增加,更重要的是计算量(以Gflops衡量)的增加。DiT-XL/2模型虽然参数量比DiT-XL/8略少,但由于使用了更小的patch size(更多的token),其计算量更大,性能也更好。

Token数量的扩展

DiT的另一个扩展维度是输入token的数量,这通过调整patch size来实现。更小的patch size意味着更多的token,也意味着更大的计算量。

研究者们测试了三种patch size:8×8、4×4和2×2。结果显示,更小的patch size(更多的token)总是带来更好的性能。DiT-XL/2(patch size为2×2)在所有配置中表现最佳。

计算效率的优势

虽然DiT-XL/2需要更多的计算资源,但它在计算效率方面相比传统方法有显著优势。例如,在256×256分辨率下,LDM-4模型需要103 Gflops,ADM-U需要742 Gflops,而DiT-XL/2只需要119 Gflops就能达到更好的效果。

这种效率优势在更高分辨率下更加明显。在512×512分辨率下,ADM-U需要2813 Gflops,而DiT-XL/2只需要525 Gflops。这意味着DiT不仅效果更好,而且更加实用。

扩展规律的深层含义

DiT的扩展规律揭示了一个重要的原理:在生成模型中,计算量(而不仅仅是参数量)是决定性能的关键因素。这个发现对整个领域都有重要意义,它告诉我们应该如何设计和优化生成模型。

更重要的是,DiT的扩展规律为未来的发展指明了方向。既然更大的模型、更多的计算量能够带来更好的效果,那么随着计算资源的不断增长,我们可以期待DiT在未来会有更加惊人的表现。这正是Sora等大规模应用能够实现的理论基础。

5. DiT的实际表现:数据说话

5.1 性能指标解读

在机器学习领域,特别是生成模型的评估中,我们需要客观的指标来衡量模型的性能。对于图像生成任务,最重要的指标之一是FID(Fréchet Inception Distance)分数。

FID分数衡量的是生成图像与真实图像在特征空间中的距离。想象一下,我们有一个"图像质量检测器"(实际上是一个预训练的Inception网络),它能够提取图像的高级特征。FID计算的是生成图像特征分布与真实图像特征分布之间的距离。FID分数越低,说明生成的图像越接近真实图像的质量和多样性。

DiT-XL/2在ImageNet 256×256数据集上取得了2.27的FID分数,这是一个突破性的成绩。要知道,在DiT之前,最好的结果是LDM达到的3.60。从3.60到2.27,这不仅仅是数字上的改进,而是质的飞跃。

在512×512分辨率上,DiT-XL/2同样表现出色,将FID从之前最好的3.85(ADM-U)降低到3.04。这些数字背后代表的是更清晰、更真实、更多样化的生成图像。

5.2 扩展规律的发现

DiT最重要的发现之一是确立了扩散模型中的扩展规律。通过系统性的实验,研究者们发现了几个重要的规律:

计算量决定性能

DiT的实验清楚地表明,模型的性能主要由计算量(Gflops)决定,而不仅仅是参数数量。这个发现颠覆了很多人的直觉。例如,DiT-XL/8虽然参数数量与DiT-XL/2相近,但由于计算量较小,性能明显不如后者。

这个规律的实际意义是巨大的。它告诉我们,在设计生成模型时,应该优先考虑如何增加有效的计算量,而不是盲目地增加参数数量。这为模型设计提供了明确的指导原则。

大模型的效率优势

另一个令人惊讶的发现是,大模型在训练效率上反而更有优势。DiT-XL/2虽然需要更多的计算资源,但它达到相同性能所需的训练步数更少。这意味着,如果我们有足够的计算资源,训练大模型实际上是更经济的选择。

这个现象在其他领域(如自然语言处理)也有类似的观察,但在生成模型中得到确认还是第一次。它暗示着,随着计算资源的不断增长,我们应该期待看到更大、更强的生成模型。

5.3 实际应用效果

数字是抽象的,但实际的生成效果是直观的。DiT-XL/2生成的图像在多个方面都展现出了显著的改进:

图像质量的提升

DiT生成的图像在细节丰富度、色彩准确性、纹理真实性等方面都有明显提升。特别是在处理复杂场景时,DiT能够更好地保持全局一致性,避免出现不协调的元素。

多样性的保证

生成模型的一个常见问题是模式崩塌,即模型只能生成有限几种样本。DiT通过其强大的建模能力和扩散过程的随机性,很好地避免了这个问题。即使是相同的类别标签,DiT也能生成风格迥异的图像。

潜在空间的连续性

DiT学到的潜在空间具有很好的连续性。这意味着我们可以在潜在空间中进行插值,生成平滑的过渡效果。这种特性对于图像编辑和创意应用非常有价值。

6. 技术细节深入:代码与实现

6.1 官方实现解析

Facebook Research在GitHub上开源了DiT的完整实现[5],这为我们深入理解DiT的技术细节提供了宝贵的资源。让我们从代码的角度来看看DiT是如何实现的。

核心模型结构

DiT的核心模型定义在models.py文件中。整个模型可以分为几个关键组件:

class DiT(nn.Module):
    def __init__(self, input_size=32, patch_size=2, in_channels=4, 
                 hidden_size=1152, depth=28, num_heads=16, 
                 mlp_ratio=4.0, class_dropout_prob=0.1, num_classes=1000):
        super().__init__()
        self.num_heads = num_heads
        
        # Patchify层:将输入转换为token
        self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size)
        
        # 位置编码
        self.pos_embed = nn.Parameter(torch.zeros(1, self.x_embedder.num_patches, hidden_size))
        
        # 时间步嵌入
        self.t_embedder = TimestepEmbedder(hidden_size)
        
        # 类别嵌入
        self.y_embedder = LabelEmbedder(num_classes, hidden_size, class_dropout_prob)
        
        # Transformer blocks
        self.blocks = nn.ModuleList([
            DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio) for _ in range(depth)
        ])
        
        # 最终的归一化和输出层
        self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels)

这个结构清晰地展现了DiT的设计思路:首先将输入转换为token,然后通过多个Transformer block进行处理,最后输出预测结果。

adaLN的具体实现

adaLN是DiT最重要的创新,让我们看看它是如何实现的:

def modulate(x, shift, scale):
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)

class DiTBlock(nn.Module):
    def __init__(self, hidden_size, num_heads, mlp_ratio=4.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True)
        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        mlp_hidden_dim = int(hidden_size * mlp_ratio)
        self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim)
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(),
            nn.Linear(hidden_size, 6 * hidden_size, bias=True)
        )

    def forward(self, x, c):
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
        x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa))
        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
        return x

这个实现展示了adaLN的精妙之处:条件信息c(包含时间步和类别信息)被转换为6个参数,分别控制注意力层和MLP层的归一化过程。

6.2 关键算法步骤

训练过程

DiT的训练过程遵循标准的扩散模型训练流程:

  1. 前向扩散:给定一个真实图像x₀,随机选择一个时间步t,按照预定义的噪声调度添加噪声,得到xₜ。

  2. 条件编码:将时间步t和类别标签y编码为条件向量c。

  3. 噪声预测:使用DiT网络预测添加的噪声ε。

  4. 损失计算:计算预测噪声与真实噪声之间的均方误差。

训练的目标函数可以简化为:
L = E [ ∣ ∣ ε − ε θ ( x t , t , y ) ∣ ∣ 2 ] L = E[||ε - ε_θ(x_t, t, y)||²] L=E[∣∣εεθ(xt,t,y)2]

其中 ε θ ε_θ εθ是DiT网络的输出。

推理过程

推理时,DiT使用DDPM或DDIM采样器从纯噪声开始逐步去噪:

  1. 初始化:从标准高斯分布中采样初始噪声 x T x_T xT

  2. 逐步去噪:对于每个时间步 t = T , T − 1 , . . . , 1 t = T, T-1, ..., 1 t=T,T1,...,1

    • 使用DiT预测噪声 ε θ ( x t , t , y ) ε_θ(x_t, t, y) εθ(xt,t,y)
    • 根据DDPM/DDIM公式计算 x t − 1 x_{t-1} xt1
  3. 最终输出:经过 T T T步去噪后得到 x 0 x_0 x0,即生成的图像。

7. 总结

回顾DiT的发展历程,我们可以看到它是多个技术趋势汇聚的结果。扩散模型提供了强大的生成框架,Transformer提供了优秀的建模能力,而DiT则巧妙地将两者结合,创造出了一个既强大又可扩展的生成模型。

DiT的核心贡献可以总结为几个方面:

首先是架构创新。DiT证明了Transformer可以成功地应用到扩散模型中,并且效果优于传统的U-Net架构。特别是adaLN机制的设计,为条件信息的注入提供了一个优雅的解决方案。

其次是扩展规律的发现。DiT系统地研究了模型规模与性能的关系,发现了计算量是决定性能的关键因素。这一发现为后续的研究提供了明确的指导方向。

最后是实际效果的突破。DiT在多个基准测试中都取得了当时最好的结果,证明了这种方法的有效性。更重要的是,DiT为Sora等大规模应用奠定了技术基础。


参考文献

[1] Ho, J., Jain, A., & Abbeel, P. (2020). Denoising diffusion probabilistic models. arXiv preprint arXiv:2006.11239. https://arxiv.org/abs/2006.11239

[2] Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., … & Houlsby, N. (2020). An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929. https://arxiv.org/abs/2010.11929

[3] Peebles, W., & Xie, S. (2022). Scalable diffusion models with transformers. arXiv preprint arXiv:2212.09748. https://arxiv.org/abs/2212.09748

[4] OpenAI. (2024). Video generation models as world simulators. https://openai.com/index/video-generation-models-as-world-simulators/

[5] Facebook Research. (2022). DiT: Official PyTorch Implementation of “Scalable Diffusion Models with Transformers”. https://github.com/facebookresearch/DiT


网站公告

今日签到

点亮在社区的每一天
去签到