UniDiffuser: One Transformer Fits All Distributions in Multi-Modal Diffusion at Scale

论文链接 | GitHub

作者：Fan Bao, Shen Nie, Kaiwen Xue, Chongxuan Li, Shi Pu, Yaole Wang, Gang Yue, Yue Cao, Hang Su, Jun Zhu
发表于：ICML 2023

研究背景

扩散模型（Diffusion Models）在图像生成领域取得了巨大成功，但现有方法主要基于 U-Net 架构。随着 Transformer 在各领域的成功应用，如何将 Transformer 有效地应用于多模态扩散模型成为一个重要研究方向。现有的多模态生成方法通常需要为不同任务（文生图、图生文、联合生成等）设计不同的模型架构。

研究目标

• 多模态联合建模：用一个统一的框架同时处理图像、文本等多种模态的生成任务
• 任务统一：用单一模型支持文生图、图生文、联合生成、无条件生成等多种任务
• 架构创新：设计适合多模态扩散的 Transformer 架构（U-ViT）

核心概念

扩散模型基础

扩散模型通过前向过程逐步向数据添加噪声，再通过反向过程学习去噪，从而实现生成。

多模态独立加噪

UniDiffuser 的核心思想是对不同模态独立添加噪声，使用不同的时间步 $t$ 和 $s$，通过控制时间步实现任务切换。

U-ViT 架构

将 U-Net 的长跳跃连接（Long Skip Connection）引入 Vision Transformer，在浅层和深层之间建立连接。

研究方法

1. 单模态扩散基础公式

前向过程（加噪）

$$q(x_t | x_0) = \mathcal{N}(x_t; \alpha_t x_0, \sigma_t^2 I)$$
参数说明：

• $x_0$：原始干净数据
• $x_t$：时间步 $t$ 的加噪数据
• $\alpha_t, \sigma_t$：噪声调度参数，满足 $\alpha_t^2 + \sigma_t^2 = 1$（VP-SDE 设定）
• 随着 $t$ 增大，$\alpha_t \to 0$，$\sigma_t \to 1$，数据逐渐变成纯噪声

等价的重参数化表示：
$$x_t = \alpha_t x_0 + \sigma_t \epsilon, \quad \epsilon \sim \mathcal{N}(0, I)$$
这个形式便于采样和训练，$\epsilon$ 是标准高斯噪声。

反向过程（去噪）

扩散模型学习反向过程 $p_\theta(x_{t-1}|x_t)$，通过神经网络预测噪声 $\epsilon_\theta(x_t, t)$，然后恢复 $x_0$：
$$\hat{x}_0 = \frac{x_t - \sigma_t \epsilon_\theta(x_t, t)}{\alpha_t}$$
解释：由 $x_t = \alpha_t x_0 + \sigma_t \epsilon$ 反解得到。预测出噪声后，即可估计原始数据。

2. 多模态联合扩散框架

联合前向过程

对于图文对 $(x_0, y_0)$，独立地对每个模态加噪：
$$q(x_t, y_s | x_0, y_0) = q(x_t | x_0) \cdot q(y_s | y_0)$$
展开形式：
$$x_t = \alpha_t x_0 + \sigma_t \epsilon_x, \quad y_s = \alpha_s y_0 + \sigma_s \epsilon_y$$
关键点：两个模态使用独立的时间步 $t$ 和 $s$，这是实现多任务统一的核心设计。

联合分布的边缘化

数据的联合分布 $q(x_0, y_0)$ 加噪后变为：
$$q(x_t, y_s) = \int q(x_t, y_s | x_0, y_0) q(x_0, y_0) , dx_0 dy_0$$
解释：这是对所有可能的原始数据对进行积分，得到加噪后数据的边缘分布。

3. 统一多任务的核心机制

核心洞察

通过控制 $t$ 和 $s$ 的取值，可以从联合分布中恢复各种边缘分布和条件分布：

条件生成的数学原理

当 $s=0$ 时，$y_s = y_0$（文本无噪声），此时：
$$q(x_t, y_0) = q(x_t | y_0) \cdot q(y_0)$$
解释：对 $y$ 不加噪声（$s=0$）在数学上等价于以 $y$ 为条件进行生成。这是一个优雅的设计——不需要修改模型架构，只需控制时间步即可切换任务。

4. 训练目标函数

噪声预测目标

模型 $\epsilon_\theta$ 同时预测两个模态的噪声：
$$[\hat{\epsilon}_x, \hat{\epsilon}y] = \epsilon_\theta(x_t, y_s, t, s)$$
完整训练损失：
$$\mathcal{L}(\theta) = \mathbb{E}{t, s, (x_0, y_0), \epsilon_x, \epsilon_y} \left[ \lambda_t |\epsilon_x - \hat{\epsilon}_x|^2 + \lambda_s |\epsilon_y - \hat{\epsilon}_y|^2 \right]$$
各项说明：

• $t, s \sim \mathcal{U}[0, T]$：从均匀分布采样时间步
• $(x_0, y_0) \sim q(x_0, y_0)$：从数据集采样图文对
• $\epsilon_x, \epsilon_y \sim \mathcal{N}(0, I)$：独立采样两个高斯噪声
• $\lambda_t, \lambda_s$：损失权重（通常设为 1）
• $|\cdot|^2$：均方误差损失

简化形式（实际训练中常用）：
$$\mathcal{L} = \mathbb{E}\left[|\epsilon_x - \hat{\epsilon}_x|^2 + |\epsilon_y - \hat{\epsilon}_y|^2\right]$$

5. 采样过程

联合采样（DDPM 形式）

从 $t=T$（纯噪声）开始，逐步去噪到 $t=0$：
$$x_{t-1} = \frac{1}{\sqrt{\alpha_{t|t-1}}} \left( x_t - \frac{1-\alpha_{t|t-1}}{\sigma_t} \hat{\epsilon}x \right) + \tilde{\sigma}t z$$
$$y{s-1} = \frac{1}{\sqrt{\alpha{s|s-1}}} \left( y_s - \frac{1-\alpha_{s|s-1}}{\sigma_s} \hat{\epsilon}_y \right) + \tilde{\sigma}_s z’$$
参数说明：

• $z, z’ \sim \mathcal{N}(0, I)$：采样的随机噪声（引入随机性）
• $\tilde{\sigma}_t$：后验方差，控制采样的随机程度
• $\alpha_{t|t-1} = \alpha_t / \alpha_{t-1}$：相邻时间步的比值

条件采样（文生图）

固定 $s=0$（即 $y_s = y_0$ 为输入文本），只对图像进行去噪：
$$x_{t-1} = \frac{1}{\sqrt{\alpha_{t|t-1}}} \left( x_t - \frac{1-\alpha_{t|t-1}}{\sigma_t} \hat{\epsilon}_x(x_t, y_0, t, 0) \right) + \tilde{\sigma}_t z$$
解释：文本时间步固定为 0，模型接收干净文本作为条件，只更新图像。

6. U-ViT 架构中的跳跃连接

U-ViT 在第 $l$ 层和第 $(L-l)$ 层之间添加跳跃连接：
$$h^{(L-l)} = \text{Block}^{(L-l)}\left( \text{Concat}(h^{(L-l-1)}, h^{(l)}) \right)$$
参数说明：

• $h^{(l)}$：第 $l$ 层的隐藏状态
• $L$：Transformer 总层数
• Concat：沿特征维度拼接
• 拼接后通过线性层降维回原始维度

设计动机：借鉴 U-Net 的成功经验，跳跃连接帮助保留低层的细节信息，有助于生成高质量图像。

7. 分类器自由引导（Classifier-Free Guidance, CFG）

为增强条件生成效果，使用 CFG 技术：
$$\tilde{\epsilon}_x = \epsilon_\theta(x_t, \varnothing, t, 0) + w \cdot \left( \epsilon_\theta(x_t, y_0, t, 0) - \epsilon_\theta(x_t, \varnothing, t, 0) \right)$$
参数说明：

• $\varnothing$：空文本条件（null condition）
• $w$：引导强度（guidance scale），通常 $w > 1$
• 第一项：无条件预测
• 括号内：条件预测与无条件预测的差值（条件信号）

等价形式：
$$\tilde{\epsilon}_x = (1-w) \cdot \epsilon_\theta(x_t, \varnothing, t, 0) + w \cdot \epsilon_\theta(x_t, y_0, t, 0)$$

训练技巧：训练时以一定概率（如 10%）将文本随机替换为空文本，使模型同时学习条件和无条件生成。

模型输入总结

主要发现

实验结果

MS-COCO 256×256 文生图（零样本）：

ImageNet 256×256 类条件生成：

讨论

优势

• 统一框架：单一模型支持多种生成任务，无需为每个任务单独训练
• 优雅设计：通过时间步控制实现任务切换，不需要修改架构
• 强大性能：在多个基准上达到 SOTA
• 可扩展性：在 10 亿参数规模上验证有效

局限性

• 需要大规模图文对数据进行训练
• 文本生成质量依赖于 CLIP 编码器的表示能力
• 推理速度受限于扩散模型的迭代采样

相关工作

• DALL-E：基于 VQ-VAE 的文生图模型
• Stable Diffusion：潜空间扩散模型
• DiT：Diffusion Transformer 架构[[Diffusion-Transformers-DiT]]
• GLIDE：文本引导的扩散模型

公式速查表

参考文献

• Bao, F., et al. (2023). One Transformer Fits All Distributions in Multi-Modal Diffusion at Scale. ICML 2023.
• Ho, J., et al. (2020). Denoising Diffusion Probabilistic Models. NeurIPS 2020.
• Peebles, W., & Xie, S. (2023). Scalable Diffusion Models with Transformers. ICCV 2023.

Scalable Diffusion Models with Transformers | ICCV 2023

研究背景

扩散模型(Diffusion Models)在图像生成领域取得了显著成功，但其架构设计一直沿用卷积U-Net作为主干网络。与此同时，Transformer架构已经在自然语言处理、视觉识别等多个领域取得了统治地位，并展现出优秀的可扩展性。本文探索将Transformer架构引入扩散模型，研究其在图像生成任务中的可扩展性和性能表现。

研究目标

• 突破架构局限：探索用Transformer替代传统U-Net作为扩散模型主干的可行性
• 验证可扩展性：研究Transformer在扩散模型中的可扩展性规律
• 建立性能基准：在ImageNet等基准数据集上达到SOTA性能
• 揭示计算-质量关系：分析模型计算量(Gflops)与生成质量之间的关系

核心概念

Latent Diffusion Models (LDMs)

在潜在空间而非像素空间训练扩散模型，提高计算效率：

1. VAE编码器：将图像压缩到潜在空间 $z = E(x)$
2. 扩散模型：在潜在空间 $z$ 中训练
3. VAE解码器：将生成的潜在表示解码为图像 $x = D(z)$

对于256×256×3的图像，VAE将其压缩为32×32×4的潜在表示（下采样因子为8）。

注意：这里使用的是标准VAE，输出是连续的潜在表示，而非VQ-VAE的离散codebook索引。

Patchify机制

将潜在表示分解为patch序列：

• 输入：32×32×4的潜在表示
• Patch大小：$p \times p$（$p \in {2, 4, 8}$）
• 输出序列长度：$T = (I/p)^2$

例如，$p=2$ 时，序列长度 $T = (32/2)^2 = 256$。

Classifier-Free Guidance

条件生成的采样技巧，提高生成质量：

$$
\hat{\epsilon}_\theta(x_t, c) = \epsilon_\theta(x_t, \emptyset) + s \cdot (\epsilon_\theta(x_t, c) - \epsilon_\theta(x_t, \emptyset))
$$

其中：

• $c$：条件信息（如类别标签）
• $\emptyset$：空条件（训练时随机dropout）
• $s$：guidance scale（$s > 1$增强条件控制）

研究方法

DiT架构设计

DiT基于Vision Transformer (ViT)架构，包含以下组件：

1. Patchify层：将潜在表示转换为token序列
2. 位置编码：使用正弦-余弦位置编码
3. DiT Blocks：N个Transformer block
4. 解码器：将token序列解码为噪声预测和协方差预测

条件注入机制

探索了四种将时间步 $t$ 和类别标签 $c$ 注入Transformer的方式：

1. In-Context Conditioning

将 $t$ 和 $c$ 的embedding作为额外token添加到序列中：

1

tokens = [t_embed, c_embed, patch_1, patch_2, ..., patch_T]

• 优点：无需修改标准Transformer block
• 缺点：性能较差（FID ~80）

2. Cross-Attention

通过交叉注意力机制注入条件：

1
2
3

# 标准self-attention后添加cross-attention
x = x + SelfAttention(x)
x = x + CrossAttention(x, [t_embed, c_embed])

• 优点：灵活的条件控制
• 缺点：增加15%计算量，性能中等（FID ~60）

3. Adaptive Layer Norm (adaLN)

通过自适应归一化层注入条件：

1
2

γ, β = MLP(t_embed + c_embed)
output = γ * normalize(x) + β

• 优点：计算高效，性能较好（FID ~45）
• 缺点：所有token共享���同的调制参数

4. adaLN-Zero（最优方案）

在adaLN基础上增加门控参数并零初始化：

1
2
3
4
5
6
7
8

γ, β, α = MLP(t_embed + c_embed)

# Transformer Block
x = x + α₁ * Attention(γ₁ * normalize(x) + β₁)
x = x + α₂ * FFN(γ₂ * normalize(x) + β₂)

# 初始化：MLP输出零向量 → α=0, γ=0, β=0
# 因此初始时：x = x + 0 = x（恒等函数）

为什么需要两组参数？

Transformer block有两个子层（Attention + FFN），每个子层需要独立的条件控制：

• 第一组 $(γ_1, β_1, α_1)$：用于Attention子层
• 第二组 $(γ_2, β_2, α_2)$：用于FFN子层

总共6个参数：shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp

零初始化的优势：

• 训练稳定性：初始时网络是恒等映射，梯度流动顺畅
• 更好的性能：FID显著优于其他方法（FID ~23）
• 渐进式学习：从恒等函数开始，逐步学习有用的变换

模型配置

设计四种规模的模型配置：

Point-wise FFN

Transformer中的标准组件，对序列中每个位置独立应用前馈网络：

1
2
3
4
5
6
7
8
9
10
11
12

class PointwiseFeedForward(nn.Module):
    def forward(self, x):
        # x: [batch, seq_len, d_model]
        x = Linear1(x)      # d_model → d_ff (通常4倍扩展)
        x = GELU(x)
        x = Linear2(x)      # d_ff → d_model
        return x

# 关键特点：
# 1. 每个token独立处理（point-wise）
# 2. 所有位置共享相同的权重
# 3. 可以完全并行计算

与Self-Attention的分工：

• Self-Attention：全局信息聚合（不同token间交互）
• Point-wise FFN：局部特征变换（每个token独立处理）

这是所有Transformer变体（LLM、ViT、DiT）的统一设计。

主要发现

1. Gflops与生成质量强相关

模型前向传播的计算量(Gflops)与FID呈强负相关（相关系数-0.93）：

• 增加模型深度/宽度 → 提升Gflops → 降低FID
• 减小patch大小 → 增加token数量 → 提升Gflops → 降低FID

关键洞察：参数量不是���一决定因素，计算量才是提升性能的关键。

2. adaLN-Zero显著优于其他条件注入方式

在400K训练步数时的FID对比：

3. 优秀的可扩展性

DiT展现出与ViT类似的可扩展性：

• 增加模型规模持续提升性能
• 训练高度稳定，无需学习率预热或特殊正则化
• 未观察到常见的loss spike现象

4. 计算效率优势

DiT-XL/2 (118.6 Gflops) 比传统方法更高效：

• 像素空间U-Net (ADM)：1120 Gflops（~10倍）
• 潜在空间U-Net (LDM-4)：103.6 Gflops（相近但性能更优）

5. 采样计算无法弥补模型计算不足

增加采样步数（增加测试时计算量）无法弥补模型规模不足：

• DiT-L/2 使用1000步采样：80.7 Tflops，FID=25.9
• DiT-XL/2 使用128步采样：15.2 Tflops，FID=23.7

结论：模型计算量比采样计算量更重要。

实验设计

数据集

• ImageNet：256×256和512×512分辨率
• 任务：类条件图像生成

训练配置

• 优化器：AdamW
• 学习率：$1 \times 10^{-4}$（常数，无warmup）
• 批大小：256
• 数据增强：仅水平翻转
• EMA：decay=0.9999
• 硬件：TPU v3-256 pod

评估指标

• 主要指标：FID-50K（使用250步DDPM采样）
• 次要指标：Inception Score、sFID、Precision/Recall

结果分析

ImageNet 256×256基准测试：

ImageNet 512×512基准测试：

DiT-XL/2在两个分辨率上都达到了SOTA性能。

讨论

优势

• 架构统一性：证明Transformer可以成功替代U-Net，推动生成模型架构统一化
• 优秀的可扩展性：计算量与性能呈强相关，为大规模模型发展指明方向
• 训练稳定性：无需特殊技巧即可稳定训练
• 计算效率：在相近或更少的计算量下达到更好的性能
• 更高的Recall：相比LDM，DiT在所有guidance scale下都有更高的recall值

局限性

• 依赖预训练VAE：使用Stable Diffusion的VAE，是混合架构而非纯Transformer
• 仅探索类条件生成：未涉及文生图等更复杂的条件生成任务
• 计算资源需求：大规模模型训练需要TPU集群
• patch大小的权衡：更小的patch提升性能但增加计算量

与传统U-Net的对比

相关工作

Transformer在生成模型中的应用

• 自回归模型：GPT系列、ImageGPT、DALL·E
• 掩码生成模型：MaskGIT、MAGE
• DALL·E 2：使用Transformer生成CLIP embedding

扩散模型架构

• DDPM (Ho et al., 2020)：首次引入U-Net作为扩散模型主干
• ADM (Dhariwal & Nichol, 2021)：改进U-Net设计，达到SOTA
• LDM (Rombach et al., 2022)：潜在空间扩散模型
• Concurrent work：U-ViT探索了类似的Transformer架构

架构复杂度分析

• ViT (Dosovitskiy et al., 2021)：证明Transformer在视觉任务中的可扩展性
• Scaling Laws：语言模型中计算量与性能的幂律关系

未来方向

1. 扩展到文生图任务：将DiT应用于DALL·E 2、Stable Diffusion等文生图模型
2. ���大规模的模型：继续扩展模型规模，探索scaling law
3. 纯Transformer架构：在像素空间训练DiT，摆脱VAE依赖
4. 多模态条件生成：探索更复杂的条件注入机制
5. 高效采样方法：结合DiT开发更快的采样算法
6. 架构搜索：自动化探索最优的DiT配置

技术细节补充

VAE vs VQ-VAE

Transformer Block完整结构

1
2
3
4
5
6
7
8
9
10
11
12
13
14

def dit_block(x, c):
    # 生成6个调制参数
    γ₁, β₁, α₁, γ₂, β₂, α₂ = adaLN_modulation(c)

    # 子层1: Multi-Head Self-Attention
    x = x + α₁ * Attention(adaLN(x, γ₁, β₁))

    # 子层2: Point-wise Feed-Forward
    x = x + α₂ * FFN(adaLN(x, γ₂, β₂))

    return x

def adaLN(x, γ, β):
    return γ * normalize(x) + β

计算复杂度分析

对于DiT-XL/2（$N=256$ tokens，$d=1152$）：

• Self-Attention：$O(N^2 \cdot d) \approx 75M$ ops
• Point-wise FFN：$O(N \cdot d^2 \cdot 2) \approx 680M$ ops

FFN占据了大部分计算量！

参考文献

• Peebles, W., & Xie, S. (2023). Scalable Diffusion Models with Transformers. ICCV 2023.
• Ho, J., et al. (2020). Denoising Diffusion Probabilistic Models. NeurIPS 2020.
• Rombach, R., et al. (2022). High-Resolution Image Synthesis with Latent Diffusion Models. CVPR 2022.
• Dosovitskiy, A., et al. (2021). An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale. ICLR 2021.
• Dhariwal, P., & Nichol, A. (2021). Diffusion Models Beat GANs on Image Synthesis. NeurIPS 2021.

关键要点总结

1. 架构创新：首次系统性地将纯Transformer应用于扩散模型
2. adaLN-Zero：零初始化的自适应归一化是性能关键
3. Gflops定律：计算量与生成质量强相关（-0.93）
4. SOTA性能：ImageNet 256×256达到FID 2.27
5. 统一架构：推动生成模型向Transformer统一的趋势