这篇论文的关键价值在于,它跳出了传统扩散模型“使用 U-Net 结合交叉注意力进行单向文本引导”以及“沿着弯曲轨迹去噪”的旧范式。作者通过引入具有双向信息交互的双流多模态 Transformer 架构(MM-DiT),并结合改进的 Rectified Flow(修正流)直线轨迹与重加权采样策略,论文报告其缓解了模型对复杂文本遵循困难(尤其是图像内的文字拼写)以及高分辨率生成推理成本过高的痛点。
论文的核心贡献
- 提出针对 Rectified Flow 的新型噪声采样策略 (Tailored SNR Samplers): 论文采用 Logit-Normal 等分布在训练时对时间步进行采样,从而为中间的去噪时间步赋予更大的权重。这解决了标准 Rectified Flow 均匀采样导致难度最高的中间生成阶段学习不足的问题,在保持直线轨迹优势(少步推理)的同时,论文报告其采样质量提升,并在作者实验设置下优于传统扩散模型目标函数。
- 设计双流多模态扩散 Transformer 架构 (MM-DiT): 作者放弃了在 U-Net 中用固定文本特征进行交叉注意力的做法,而是为图像和文本标记(Tokens)分配了两组独立的网络权重,并在同一个自注意力块中进行拼接计算。这种双向信息流动让图像和文本在同一个特征空间中深度混合,论文报告其增强了模型对复杂 Prompt 的理解能力和图像内排版文字生成能力。
- 验证可预测的模型扩展规律 (Predictable Scaling Trends): 论文系统地将模型参数规模扩展至 8B,并展示验证损失(Validation Loss)的平滑下降与各项生成质量指标(如自动评测指标及人类偏好评分)呈现出较强相关性。这为大规模生成模型的训练提供了明确的扩展证据,表明该架构在作者实验范围内仍有继续扩展的潜力。
复杂 Pipeline 深度解析

图中元素对照解读
- 图中左侧 (a) 顶部粉色 Caption 框及向下三条箭头: 对应论文中的文本提示输入;它在 pipeline 中的作用是同时送入多种预训练文本编码器,为后续图像生成提供语义条件。
- 图中左上绿色 CLIP-G/14、CLIP-L/14、T5 XXL 三个横向框: 对应论文中的三路冻结文本编码器;它在 pipeline 中的作用是分别提取全局/局部文本语义,并把 caption 转成后续 Transformer 可处理的条件特征。
- 图中左中 “77 + 77 tokens” 和 “4096 channel” 旁的堆叠矩形: 对应论文中的文本 token 序列和 T5 高维通道特征;它在 pipeline 中的作用是保留细粒度词序列信息,而不只依赖单个 pooled 文本向量。
- 图中左下粉色 Timestep、绿色 Sinusoidal Encoding 与上下两个 MLP: 对应论文中的时间步嵌入和条件调制信号生成;它在 pipeline 中的作用是告诉网络当前处于 rectified flow 轨迹的哪个位置,并生成后续层归一化/残差分支使用的调制参数。
- 图中右侧 (a) 粉色 Noised Latent 以及 Patching、Linear、Positional Embedding、加号节点: 对应论文中的加噪图像潜变量 patch 化和位置注入步骤;它在 pipeline 中的作用是把二维潜变量转成带空间位置信息的图像 token 流。
- 图中 (a) 中央浅紫色大框内的 MM-DiT-Block 1、MM-DiT-Block 2、...、MM-DiT-Block d: 对应论文中重复堆叠的多模态扩散 Transformer 主干;它在 pipeline 中的作用是反复更新文本流 c 和图像流 x,使图像 token 在去噪过程中持续吸收文本条件。
- 图中 (a) 底部的 Modulation、Linear、Unpatching、粉色 Output: 对应论文中的最终图像流读出头;它在 pipeline 中的作用是将经过 d 个 MM-DiT block 的图像 token 重新映射并还原为潜空间输出,用于预测 rectified flow 的生成方向。
- 图中右半 (b) 左右两条竖向分支、粉色 c/x 节点和绿色 Layernorm/Mod 框: 对应论文中“文本流 c”和“图像流 x”使用独立权重处理的设计;它在 pipeline 中的作用是让两种模态在归一化、调制、线性投影和 MLP 中保持各自参数,同时在注意力处交互。
- *图中 (b) 中央绿色 Q K V Attention 大框、圆圈 ⊙ 与星号 :* 对应论文中的联合注意力、拼接和逐元素乘法操作;它在 pipeline 中的作用是把文本与图像的 Q/K/V 信息送入同一个注意力计算,使图像 token 可直接参考文本 token,caption 标注也说明 ⊙ 表示 concatenation、 表示 element-wise multiplication。
1. 多模态特征预编码 (Pre-encoding)
- 输入: 文本描述 (Caption) 和加噪的图像潜变量 (Noised Latent)。
- 操作: 文本通过冻结的预训练语言模型(CLIP-G/14, CLIP-L/14, T5 XXL)提取特征,获取池化后的全局向量以及细粒度的序列 Token。图像潜变量被分割成多个 Patch,并展平为序列。
- 输出: 文本序列向量、池化文本向量以及图像 Patch 序列向量。
- 作用: 利用强大的专家模型获取富含语义的文本表示和高效的视觉压缩表示,作为后续 Transformer 处理的底层原料。
2. 模态独立对齐与位置注入 (Modality-Specific Embedding)
- 输入: 文本序列、图像序列、当前时间步 (Timestep)。
- 操作: 时间步通过正弦编码和 MLP 提取特征。图像序列加入位置编码 (Positional Embedding)。随后,文本和图像序列分别通过不同的线性层映射到相同的隐藏维度。
- 输出: 维度对齐的文本特征流 ($c$) 和图像特征流 ($x$)。
- 作用: 确保两种截然不同的模态被投射到一个公共特征空间内,同时注入时间和空间位置信息,让模型知道当前处于去噪的哪个阶段以及每个 Patch 在原图中的位置关系。
3. 双向信息交互 (MM-DiT Block 计算)
- 输入: 对齐后的文本流和图像流,以及由时间步和池化文本组合生成的调制信号。
- 操作:
- 调制 (Modulation): 文本流和图像流分别经历层归一化 (LayerNorm),并根据调制信号进行独立的缩放和偏移。
- 特征投影: 两股流分别通过各自独立的线性层生成 Queries (Q)、Keys (K) 和 Values (V)。
- 联合注意力计算: 将文本和图像的 Q 和 K 拼接起来,计算一个巨大的联合注意力矩阵,然后分别乘回各自的 V。
- 独立前馈 (MLP): 文本和图像流分离后,分别经过各自独立的 MLP 层进行特征变换。
- 输出: 更新后的文本特征流和图像特征流。
- 作用: 这是 Pipeline 最核心的机制。通过在同一个注意力矩阵中联合计算,图像 Token 可以直接“看到”文本 Token,文本 Token 也可以“感知”图像的生成状态,实现了深度的双向通信(Bidirectional flow);而独立计算的 QKV 映射和 MLP 又保证了两种模态各自特定的内在逻辑不会被生硬地破坏。
4. 预测输出 (Output Generation)
- 输入: 经过 $d$ 个 MM-DiT Block 处理后的最终图像特征流。
- 操作: 经历最后一次调制后,通过线性层降维,并执行 Unpatching 操作还原回潜变量的 2D 空间结构。
- 输出: 预测的速度场 (Velocity) 或噪声。
- 作用: 将 Transformer 的一维序列输出解码为图像潜空间的预测目标,用于在推理时指导常微分方程 (ODE) 求解器逐步去噪生成最终图像。
核心 Pipeline 对应论文中的 Figure 2. Our model architecture,主要包含 (a) Overview of all components 和 (b) One MM-DiT block。我们可以按照图中的数据流顺序将其拆解为以下几个关键步骤:
关键术语解释
- Rectified Flow (修正流):
- 功能与设计动机: 是一种将纯噪声和真实数据分布连接在一条直线上的生成模型数学框架。在这篇论文中,它替代了传统扩散模型中极其弯曲的去噪路径。设计动机是:直线路径的常微分方程 (ODE) 在推理求解时更为平缓,所需的积分步数更少、误差累积更低。它让模型能够以更少的采样步数快速生成极高质量的图像。
- Logit-Normal Sampling (Logit-正态采样):
- 功能与设计动机: 是训练 Rectified Flow 时分配给不同时间步的概率密度函数。设计动机是:在直线去噪路径的中间点(即噪声和图像混杂一半的状态),网络预测正确方向的难度最大。传统的均匀采样忽略了这种难度差异,而 Logit-Normal 采样人为地在训练时提高了中间困难步骤的采样频率。这迫使模型将更多算力倾注于攻克去噪瓶颈,论文报告这提升了最终的生成表现。
- MM-DiT (Multimodal Diffusion Transformer):
- 功能与设计动机: 论文提出的多模态扩散网络架构,用 Transformer 主干替代传统 U-Net 式扩散骨架。设计动机是:旧的交叉注意力机制通常将文本作为额外的控制条件“单向”注入图像,复杂的指令关系和文字细节容易丢失。MM-DiT 保留了双流独立的权重体系,但在 Self-Attention 环节打通模态交互,使得文本与图像能够动态、双向沟通;论文报告这有助于提升复杂 prompt 遵循和图像内文字生成。
核心结构总结
Figure 2 是整篇论文的核心架构图,因为它展示了文本到图像生成主干如何从文本编码、图像 patch 化、时间步调制进入重复堆叠的 MM-DiT 架构。这套基于 Transformer 的结构为论文后续讨论的 Scaling Trends 提供了模型基础;同时,MM-DiT 提供的高容量、双向多模态交互能力用于拟合经过采样策略训练的 Rectified Flow 直线轨迹。因此,这幅图不仅是一张数据流向说明,也概括了论文从“单向文本注入”转向“双流多模态交互”的主要架构设计。