Sora背后的DiT架构：Transformer如何重塑扩散模型的未来
当OpenAI发布Sora时，整个AI社区都在惊叹其生成视频的质量和连贯性。很少有人注意到，支撑这一突破的核心技术之一——DiT（Diffusion Transformer）架构，实际上早在2022年就已初露锋芒。DiT的成功并非偶然，它揭示了Transformer在生成式AI领域被长期忽视的潜力，同时也为扩散模型的发展指明了一条全新的道路。
1. 从U-Net到Transformer：扩散模型架构的范式转移
传统扩散模型几乎清一色采用U-Net作为主干网络，这一选择源于U-Net在图像处理任务中表现出的强大归纳偏置。U-Net的编码器-解码器结构天生适合捕捉图像的局部特征和层次化表示，其跳跃连接更是确保了细节信息的有效传递。然而，DiT论文的作者们提出了一个大胆的假设：这些看似不可或缺的特性，可能并非扩散模型成功的关键。
U-Net的三大局限在实验中逐渐显现：
- 局部感受野：卷积核的固定尺寸限制了长程依赖的建模能力
- 架构僵化：层级结构难以灵活调整深度和宽度
- 计算效率瓶颈：下采样-上采样过程带来不可忽视的开销
相比之下，Transformer展现出了截然不同的特质。通过自注意力机制，每个"图像块"（patch）都能直接与全局任何位置交互，彻底打破了局部感受野的限制。更关键的是，Transformer的模块化设计允许模型规模几乎无限制地扩展——只需增加层数或隐藏维度，就能线性提升模型容量。
实验数据清晰地展示了这种架构转变带来的优势。在ImageNet 256×256生成任务上，DiT-XL将FID分数从之前的最佳3.85（ADM-U）提升至2.27，同时训练效率提高了近40%。这一突破并非来自某个精巧的设计技巧，而是纯粹源于Transformer架构本身的可扩展性优势。
2. DiT的核心设计：当扩散过程遇见自注意力
DiT的架构看似简单——它本质上是一个经过适当修改的Vision Transformer（ViT），但其中的设计选择却处处体现着对扩散任务特性的深刻理解。
2.1 图像表示的转换艺术
DiT处理输入的方式独具匠心：
# 典型DiT的patch嵌入过程 def patchify(x, patch_size=4): B, C, H, W = x.shape x = x.reshape(B, C, H//patch_size, patch_size, W//patch_size, patch_size) x = x.permute(0, 2, 4, 1, 3, 5).reshape(B, -1, C*patch_size*patch_size) return x # 输出形状：(batch_size, num_patches, patch_dim)这种处理将32×32×4的潜在空间表示转换为256个16维的token序列，既保留了足够的空间信息，又适应了Transformer的序列处理范式。值得注意的是，patch大小的选择实际上成为了控制模型计算量的"旋钮"——将p从8减小到2，Gflops会激增16倍，这为后续的scaling研究提供了精确的控制维度。
2.2 条件注入的四种范式
如何将时间步和类别信息融入扩散过程，是DiT设计的核心挑战。论文系统比较了四种策略：
| 方法 | Gflops增加 | FID得分 | 特点描述 |
|---|---|---|---|
| In-Context | <1% | 12.23 | 类似ViT的cls token，最轻量级 |
| Cross-Attention | +15% | 9.12 | 类似原始Transformer设计 |
| adaLN | <1% | 7.89 | 动态调节归一化参数 |
| adaLN-Zero | <1% | 6.82 | 带零初始化的adaLN变体 |
令人惊讶的是，计算代价最低的adaLN-Zero反而取得了最佳效果。这一发现颠覆了"更多计算=更好性能"的直觉，表明恰当的条件机制比单纯增加计算量更重要。adaLN-Zero的成功秘诀在于其初始化策略——通过将残差路径初始化为恒等函数，确保了训练初期的稳定性，这与扩散模型需要精确控制噪声水平的特性完美契合。
3. Scaling Law：DiT的"秘密武器"
如果说U-Net像精心调校的跑车，那么DiT则更像可以无限扩容的货运列车。这种可扩展性不是渐进式的改进，而是质的飞跃。
3.1 计算量与性能的线性奇迹
DiT论文中最引人注目的发现莫过于Gflops与FID之间近乎完美的对数线性关系。当模型规模从DiT-S(0.3Gflops)扩展到DiT-XL(118.6Gflops)时，FID从68.4单调下降至2.27，没有任何饱和迹象。这种可预测的缩放行为在生成模型中极为罕见，它意味着：
- 资源分配变得可规划：增加10倍计算预算，就能预期确定性的质量提升
- 架构瓶颈被消除：不再受限于局部感受野或层级信息损失
- 训练效率最大化：大模型不仅最终效果更好，单位计算产生的收益也更高
这一发现直接影响了Sora的设计哲学——与其精心设计特定于视频的架构，不如构建一个足够通用的Transformer框架，然后放心地扩大规模。
3.2 数据规模与模型规模的协同效应
DiT的另一个反直觉特性是：更大的模型在相同计算预算下训练更高效。具体表现为：
- 用1/10的训练步数，DiT-XL就能超越充分训练的DiT-S
- 在相同Gflops预算下，大模型+少步数 > 小模型+多步数
- patch尺寸的影响独立于模型规模存在
这些现象暗示，Transformer在扩散模型中展现出了与LLMs类似的"预训练优势"——模型容量本身成为一种学习效率的保障，而不仅仅是最终性能的决定因素。
4. DiT对生成式AI未来的启示
DiT的成功绝非仅限于技术层面的突破，它实际上重新定义了我们对生成模型架构的认知边界。
传统认知的三大颠覆：
- 归纳偏置并非必需：U-Net的局部性假设可以被全局注意力取代
- 架构统一成为可能：同一主干网络可处理图像、视频甚至3D生成
- 规模优先于技巧：精心设计的小模型不如简单扩展的大模型
在实际应用中，DiT类架构已经展现出独特优势。例如，在需要高分辨率生成的场景中，传统U-Net会因为层级过深而面临梯度消失问题，而DiT只需增加注意力头的数量就能维持稳定的训练动态。又如在多模态生成任务中，Transformer天然的序列处理能力使其能更自然地融合文本、图像等异构输入。
不过，DiT也带来了新的挑战。其内存消耗随序列长度平方增长的特性，使得直接处理高分辨率图像仍然困难。一些新兴的解决方案如：
- 分块注意力：将图像划分为子区域分别处理
- 线性注意力：近似标准注意力但降低计算复杂度
- 层次化DiT：在不同分辨率层次应用Transformer
这些创新正在不断拓展DiT的应用边界，而Sora的出现已经证明，当DiT遇上足够大的数据和计算规模，生成式AI的能力边界将被重新定义。
