(论文速读)Scaling Rectified Flow Transformers:Rectified Flow + MM-DiT 的高分辨率文生图路线

发布时间:2026/9/4 17:02:37
(论文速读)Scaling Rectified Flow Transformers:Rectified Flow + MM-DiT 的高分辨率文生图路线 论文题目Scaling Rectified Flow Transformers for High-Resolution Image Synthesis中文翻译用于高分辨率图像合成的可扩展 Rectified Flow Transformer会议ICML 2024摘要扩散模型通过逆转数据向噪声演化的前向路径从噪声中生成数据并已经成为处理图像、视频等高维感知数据的重要生成建模方法。Rectified Flow 是一种较新的生成模型形式它通过直线连接数据与噪声。尽管其理论性质更好、概念也更加简单但尚未被明确确立为标准实践。本文通过让训练时的噪声采样更加偏向感知上重要的噪声尺度改进了现有 Rectified Flow 的噪声采样方法。通过大规模研究作者证明这种方法在高分辨率文生图任务上优于已有扩散模型形式。此外论文提出了一种新的 Transformer 文生图架构文本和图像两种模态分别使用不同的参数同时允许两者之间进行双向信息流动从而改善文本理解、文字生成以及人类偏好评分。实验进一步表明该架构具有可预测的 Scaling 趋势并且更低的验证损失与更好的文生图指标和人类评价高度相关。最大的模型超过了当时的多种先进模型。一、研究背景这篇论文真正想解决的不只是“把模型做大”这篇论文主要想回答三个问题。第一个问题是扩散模型从 Data 走向 Noise 时到底应该选择一条什么样的路径传统扩散模型通过逐渐加入噪声构造 Forward Process生成时再反过来一步一步求解。但不同 Forward Path 会直接影响最终的反向生成过程。如果路径弯曲数值求解时就需要更多积分步而且误差会不断累积如果数据到噪声之间的路径足够直那么理论上就可以用更少的 Sampling Steps 完成生成。Rectified Flow 的思路恰好非常直接让 Data 和 Noise 之间尽量沿直线连接。第二个问题是文生图明明包含 Text 和 Image 两种完全不同的模态为什么还要让它们共享几乎相同的 Transformer 参数传统文生图模型常见的做法是用 Cross-Attention把固定的文本表示作为条件注入图像网络。但作者认为这种单向“Text → Image”的条件方式没有充分利用文本和图像之间的交互。于是论文提出 MM-DiT让 Text Tokens 和 Image Tokens 分别拥有自己的参数空间同时在 Attention 中真正进行双向信息交换。第三个问题是 Scaling当 Rectified Flow、网络结构和训练数据都设计好以后继续扩大 Transformer性能是否还能稳定增长论文最终把模型扩展到了8B 参数、约 5 × 10²² Training FLOPs并发现 Validation Loss 随模型规模和训练计算量平滑下降而且这种下降能够预测 GenEval、T2I-CompBench 和 Human Preference 的改善。所以如果把整篇论文压缩成一条主线其实就是更合理的生成路径 → 更合理的时间采样 → 更合理的多模态 Transformer → 更大的模型 → 更强的文生图能力。二、核心方法一Rectified Flow 为什么要把 Data 和 Noise 连成一条直线2.1 先理解 Flow 到底在学什么论文将生成过程看成一个 ODE其中 () 是神经网络预测的 Velocity Field。简单来说模型每看到当前位置 ()就需要告诉我们“接下来应该朝哪个方向走”如果把这个速度场一路积分就可以从噪声分布走回数据分布。传统 Flow 的麻烦是如果直接通过 ODE Solver 来训练整个速度场计算量很高。因此 Flow Matching 的核心思想是人为构造一条从 Data 到 Noise 的路径然后直接监督网络学习这条路径上的速度。论文统一写成其中 () 是真实数据() 是 Gaussian Noise(t) 控制当前处于数据和噪声之间的什么位置。2.2 Rectified Flow最简单的路径就是直线Rectified Flow 选择了一个非常简单的 Forward Process当 (t0) 时()也就是真实图片当 (t1) 时()变成纯噪声。于是整个过程就是真实图片 () → 线性插值 → 噪声 ()而因为这是一条直线它对应的理想速度也非常简单本质上就是从 () 指向 () 的方向。这也是 Rectified Flow 最吸引人的地方。假设传统 Diffusion 的 Sampling Path 是一条弯弯曲曲的道路那么 Sampling 时必须频繁修正方向Rectified Flow 则希望把道路“拉直”。道路越直用少量 Euler Steps 就越容易逼近真实轨迹。所以这里真正优化的并不是单纯的网络速度而是让模型所学习的生成动力学本身更加容易被数值积分。这也是为什么论文后面的实验会不断关注 Few-Step Sampling。三、核心方法二Uniform Sampling 并不合理作者为什么改成 Logit-Normal仅仅使用 Rectified Flow 还不够。原始 RF 通常从 () 中均匀采样训练时间也就是说靠近 Data 的位置、中间位置和靠近 Noise 的位置被训练到的概率一样。作者认为这并不合理。3.1 最难学的其实是中间状态对于 Rectified Flow模型的 Velocity Prediction Target 与 () 有关。靠近两端时输入更接近纯数据或者纯噪声相对容易判断真正困难的是中间区域——此时 Data 和 Noise 高度混合模型更难判断应该如何运动。所以问题就变成既然中间时间步更难、更重要为什么训练时还要把所有 t 一视同仁作者的思路很直接多采中间少采两端。3.2 Logit-Normal Sampling论文重点研究了一种 Logit-Normal 分布。可以先不管它复杂的概率密度公式只需要记住它的作用Uniformt0 ───────────────────────── t1每个位置差不多一样重要。Logit-Normalt0 ────────████████──────── t1更多训练样本集中在中间区域。它还有两个参数 (m) 和 (s)(m) 控制概率分布往 Data 还是 Noise 一侧移动(s) 控制分布有多宽。论文 Figure 8 对 Logit-Normal 和另一种 Mode Sampling 的分布形状进行了可视化。论文 Figure 8Mode Sampling 与 Logit-Normal Sampling 的时间步分布3.3 这个看起来很小的改动实际上非常关键作者一口气训练了61 种不同的 Diffusion / Rectified Flow 配置并在 ImageNet、CC12M、不同 Sampling Steps、不同 Guidance 和 EMA 设置下进行比较。Table 1 中普通 RF 的平均 Rank 是 5.67而rf/lognorm(0.00, 1.00) 的 Rank 达到 1.54为所有配置中最佳。更重要的是它在 5 Steps 和 50 Steps 两种区域都保持很好的排名不是只对某一种 Sampling Setting 特化。Table 2 进一步说明一个很有意思的问题有些配置可以拿到某一个 FID 或 CLIP 的最好结果却很容易在另外的数据集或指标上变差相比之下rf/lognorm(0.00, 1.00) 并非所有单项都第一但跨数据集、跨指标表现非常稳定。论文 Figure 3不同生成形式在 Sampling Steps 减少时的 FID 变化Figure 3 更能体现 RF 的价值当 Sampling Steps 不断减少时Rectified Flow 的性能下降通常更慢而在步数增加到 25 以上以后经过 Logit-Normal 改进的 RF 依然可以与强力的 eps/linear 配置竞争。所以这一部分最值得记住的不是“Logit-Normal 数学形式有多复杂”而是一个很实用的训练思想训练预算不应该平均花在所有噪声等级上而应该更多分配给模型真正难学、真正影响生成质量的区域。四、核心方法三MM-DiT为什么 Text 和 Image 要“分开算、一起注意”如果说 Rectified Flow 改的是“生成路径”那么 MM-DiT 改的就是“负责走这条路径的神经网络”。4.1 传统 DiT 的问题在哪里DiT 把图像 Latent 切成 Patch变成 Token再用 Transformer 进行处理。但 Text-to-Image 不只有图像。这里实际上存在两种性质差异很大的 TokenImage Token描述空间结构、纹理、颜色和局部视觉信息Text Token描述语言语义、对象关系、数量、位置以及文字内容。如果它们完全共享一套 Transformer 参数相当于要求同一个 Feature Space 同时处理两种统计性质完全不同的数据。作者认为并不理想。4.2 MM-DiT 的核心参数分开Attention 合起来论文 Figure 2MM-DiT 整体结构和单个 Block图像先经过 Autoencoder 得到 Latent然后以 2 × 2 Patch 切分成 Image Tokens文本则由预训练、冻结的 Text Encoders 编码。论文使用 CLIP-L、OpenCLIP bigG 和 T5 相关表示。真正关键发生在 MM-DiT Block 内。Text Stream 有自己的一套 LayerNorm、Linear、MLP 等参数Image Stream 也有另外一套参数。因此可以粗略理解成Text → Text Transformer ┐├→ Joint Attention → 分回各自 StreamImage → Image Transformer ┘论文的原话是这相当于存在两个独立 Transformer但在 Attention 时把两个模态连接起来。这样一来Text 可以在适合语言的表示空间里处理Image 可以在适合视觉的表示空间里处理Attention 又允许两边互相“看见”对方。这与传统 Cross-Attention 一个很明显的区别是信息流方向。传统 Cross-Attention 更接近Text → Image而 MM-DiT 强调Text ↔ Image于是 Text Representation 不再只是一个静态 Condition而是在网络中随着 Image Representation 一起被不断更新和交互。4.3 实验真的需要这么做吗作者比较了 Vanilla DiT、CrossDiT、UViT 和 MM-DiT。Figure 4 中Vanilla DiT 最弱CrossDiT 明显改善UViT 早期收敛较快而最终MM-DiT 在 Validation Loss、CLIP 和 FID 上表现最好。论文还尝试把 CLIP Token 与 T5 Token 再拆成第三套参数但收益很小、参数和 VRAM 却继续增加因此后续采用两套主要 Modality Parameters。论文 Figure 4DiT、CrossDiT、UViT 与 MM-DiT 的训练曲线这组实验说明 MM-DiT 的提升不只是“把 Transformer 做大”而是文本和图像应该保留各自的建模空间但需要在 Attention 层真正建立双向交流。五、实验与工程优化真正把模型 Scale 到 8B还需要解决哪些问题这篇论文一个很明显的特点是它不是只提出 Rectified Flow MM-DiT然后直接报告最终 SOTA。作者实际上把 Autoencoder、Caption、模型架构、High-Resolution Training、Preference Alignment 和 Sampling Efficiency 都做了一轮系统实验。5.1 AutoencoderLatent Channel 从 4 增加到 16Latent Diffusion 最终生成质量有一个很现实的上限如果 Autoencoder 自己就无法准确重建图片那么后面的生成模型再强也无法弥补被 VAE 丢掉的信息。Table 3 比较了 4、8、16 个 Latent ChannelsLatent ChannelsFID ↓Perceptual Similarity ↓SSIM ↑PSNR ↑42.410.850.7525.1281.560.680.7926.40161.060.450.8628.62增加到 16 Channels 后Autoencoder Reconstruction 明显提升。但这里有代价Latent Dimension 越高后面的生成任务本身也越难。Figure 7 显示小模型使用 16 Channels 时不一定占优势但随着模型 Capacity 增大两者差距不断缩小到 depth 22 左右8-channel 和 16-channel 的差距已经接近消失。由于作者最终目标就是大规模 Scaling因此选择 16-channel Autoencoder。这其实体现了一个典型 Scaling 思想更丰富的 Representation 本身可能更难学但当模型容量足够以后更高的信息上限最终会转化成更高的生成质量。5.2 Caption人类写的 Caption 不一定是最好的训练标签很多真实数据集中的 Caption 只描述主体比如“一个人在街上”却可能完全忽略背景、对象之间的位置、数量、画面中文字等细节。作者因此使用 CogVLM 生成 Synthetic Caption再与 Original Caption 按50% / 50%混合训练。Table 4 的 GenEval Overall Score 从43.27 提升到 49.78其中 Color Attribution 从 11.75 提升到 24.75Position 从 6.50 提升到 18.00Two Objects 从 41.41 提升到 52.53。并非所有项目都提高例如 Colors 和 Single Object 略有下降所以更准确的结论是更详细的 Synthetic Caption 显著增强了复杂组合关系和属性理解能力。5.3 ScalingValidation Loss 真的可以预测最终画质吗论文 Figure 5模型规模、Training FLOPs 与 Validation Loss / 下游指标的 Scaling 曲线作者在 256 × 256 上用 Batch Size 4096 训练不同尺寸 MM-DiT并将模型扩展到 8B 参数。随着模型深度增加以及 Training FLOPs 增加Validation Loss 持续、平滑下降。更重要的是Validation Loss 并不只是一个“训练看起来更好了”的指标。它与以下指标都高度相关GenEvalT2I-CompBenchHuman Preference视频模型中的 Human Preference。而且作者没有观察到明显的 Saturation。这意味着在这套方法和计算范围内Model Size ↑ / Compute ↑ → Validation Loss ↓ → 文生图综合能力 ↑因此Validation Loss 可以成为一个相对便宜的 Scaling Proxy而不用每训练一点就进行昂贵的人类评测。5.4 最终结果GenEval Overall 0.74Table 5 给出了比较完整的 GenEval 结果。其中 DALL-E 3 的 Overall Score 为 0.67depth 38 的 512² 模型为 0.68经过 DPO 后达到 0.71最后使用 1024² DPO 的最大模型达到0.74。同时也必须注意边界它并不是每个子指标都绝对第一。例如 DALL-E 3 在部分单项上仍然很有竞争力。因此论文更准确的结论应该是最大的模型在 GenEval 综合分数上超过了当时比较的开放模型和 DALL-E 3并且在多数复杂 Prompt Understanding 项目中表现很强。论文 Figure 6 又通过 PartiPrompts Human Preference 比较 Visual Aesthetics、Prompt Following 和 Typography。作者报告 d 38 的 8B 模型在这些人类评价上对当时多种开放及闭源模型具有竞争优势。论文 Figure 68B 模型与其他文生图模型的人类偏好比较Figure 1 的高分辨率案例也能非常直观地看到论文重点强调的能力Typography、复杂 Prompt Following、Spatial Reasoning 和细节生成。5.5 一个很有意思的结果模型越大反而需要越少 Sampling StepsTable 6 是这篇论文里一个非常值得注意的结果。以 50 Steps 的 CLIP Score 为参考depth 15 使用 5 Steps 时下降 4.30%depth 30 下降 3.59%depth 38 只下降2.71%。如果使用 10 Stepsdepth 38 相对 50 Steps 只下降0.14%。与此同时估计出的 Sampling Path Length 也从 depth 15 的 191.13 降低到 depth 38 的 185.96。因此作者认为更大的模型不仅 Prediction 更准同时能够更好地拟合 Rectified Flow 所希望的 Straight Path。也就是说模型变大不仅提高单步预测质量还可能让整个生成轨迹变得更容易数值积分。这就形成一个比较有意思的结果Model Size ↑ → Quality ↑同时 Sampling Efficiency 也 ↑。5.6 高分辨率训练还需要两个关键工程技巧首先是QK-Normalization。作者发现模型从 256² Fine-tune 到高分辨率以后Mixed Precision Training 容易失稳。Figure 9 显示如果不处理Transformer 后几层的 Attention Logit 会不断变大Attention Entropy 最终塌缩并出现 NaN。对 Q 和 K 使用 RMSNorm 后这种增长被明显抑制从而可以继续使用高效的 bf16 Mixed Precision Training。论文 Figure 9QK-Normalization 对 Attention Logit 和 Entropy 的影响其次是Resolution-Dependent Timestep Shifting。作者观察到分辨率越高一张图片包含的信息越多因此同样的噪声强度未必能够像低分辨率那样充分破坏原图信号。所以从 256² 扩展到 1024² 时应该把 Sampling / Training Schedule 往更高噪声区域移动。Figure 10 的 Human Preference 显示 Shift 大于 1.5 后明显更好最终 1024 × 1024 模型采用α 3.0。论文 Figure 10High-Resolution Timestep Shifting5.7 DPO生成能力之外还要做人类偏好对齐高分辨率训练完成后作者进一步对 2B 和 8B 模型进行 DPO Preference Finetuning。具体没有重新训练全部参数而是在 Linear Layers 上加入 Rank 128 的 LoRA2B 和 8B 模型分别 Fine-tune 4k 和 2k Iterations。Figure 11 的样例显示 DPO 后图像通常更加美观文字拼写也更准确Figure 12 的人类评价则明确偏好 DPO 模型无论是 Prompt Following 还是 General Quality 都得到改善。论文 Figure 11、Figure 12Base Model 与 DPO Model 的样例和人类评价因此最终系统实际上不是单纯Rectified Flow MM-DiT而更完整地说是Rectified Flow Logit-Normal Sampling 16-channel Autoencoder Better Caption MM-DiT Scaling High-Resolution Finetuning DPO这也是为什么只看 MM-DiT 结构图还不能完全解释最终模型的效果。5.8 Appendix 里还有几个很值得看的结果Figure 13 研究了 T5 的作用。由于 T5-XXL 本身有约 4.7B 参数Inference Memory 成本很高。实验发现对普通 Prompt 去掉 T5 的影响并不大相对于完整模型No-T5 在 Aesthetic Quality 上仍有约 50% Win RatePrompt Adherence 为 46%但 Typography 只有 38%说明T5 对长 Prompt、复杂场景和文字生成明显更重要。Figure 14 则把 Model Size 和 Training Steps 直接画成样例矩阵从 50k → 200k → 350k → 500k同时从 depth 15 → 30 → 38复杂 Prompt 的对象关系和细节逐渐改善是 Scaling Effect 最直观的定性证据。Figure 15 还展示了基于 2B 模型 Fine-tune 的 Image Editing Model。比较有意思的是它没有专门在 Text Manipulation 数据上训练却出现了 Zero-Shot Text Replacement / Insertion 能力。作者还表示同样的数据配置下没有在 SDXL-Based Editing Model 上复现类似结果。最后Figure 1617 讨论 Data Deduplication。论文利用 SSCD 去除近重复训练图片并专门通过 Memorization Detection 验证去重效果。在 Clique Size 10 的判断标准下SSCD 0.5 的去重训练集使潜在 Memorized Examples减少约 5 倍。Table 7 则解释了为什么大规模训练时预先计算 VAE、CLIP 和 T5 Embedding 很重要。例如 T5 单次 Forward 开销很高如果每一个 Training Step 都重新运行冻结的 Text Encoder会浪费大量计算提前 Precompute 可以节约显存和后续训练计算但代价是 Storage 增加而且每个 Epoch 无法再随机做 Image Augmentation。六、总结与思考这篇论文真正值得记住的是什么如果只把这篇论文理解成“8B Transformer 文生图模型”其实会错过它最有价值的部分。它真正展示的是一套比较完整的Scaling Recipe。第一作者先从Generative Formulation入手。与其接受传统 Diffusion 的弯曲轨迹不如使用 Rectified Flow让 Data 和 Noise 之间尽可能沿 Straight Path 运动从而天然获得 Few-Step Sampling 的优势。第二作者发现原始 Rectified Flow 的Uniform Timestep Sampling 并不最优。中间噪声区域更难学习因此通过 Logit-Normal Sampling 把更多训练预算放到这些位置最终在 61 种不同训练配置中得到最稳定的综合表现。第三作者没有简单沿用 U-Net 或普通 DiT而是针对文生图本身的Multimodal Nature设计 MM-DiTText 和 Image 分开使用自己的 Transformer 参数但在 Attention 内双向交互。可以把它压缩成一句很好记的话参数分开Attention 合起来。第四Scaling 并不是单纯“多堆几层 Transformer”。16-channel Autoencoder 决定更高的视觉信息上限Synthetic Caption 提高复杂 Prompt SupervisionQK-Norm 保证高分辨率 Mixed Precision 稳定Timestep Shifting 适配 Resolution ChangeDPO 最后再进行 Human Preference Alignment。最终这些模块才共同支撑了 8B 模型的效果。论文的结论是模型规模与训练计算量增加后 Validation Loss 持续下降而且这种下降能够很好地预测 GenEval、T2I-CompBench 和人类偏好在本文所覆盖的规模范围内Scaling Curve 仍没有明显饱和。如果把这篇论文最后压缩成一句话我会这样理解作者不是单独发明了一个“更大的文生图 Transformer”而是从生成路径、时间采样、多模态网络、Latent 表示、训练数据到高分辨率 Scaling把整条文生图训练链重新设计了一遍MM-DiT 是网络核心Rectified Flow 是生成核心而 Scaling 实验则证明这套组合在模型继续做大时依然有效。从研究思路上看最值得借鉴的可能也正是这一点真正有效的 Scaling往往不是简单地“参数量 ↑”而是先保证 Representation、Objective、Architecture、Data 和 Optimization 都能够跟得上模型规模。