论文题目:Scaling Rectified Flow Transformers for High-Resolution Image Synthesis(用于高分辨率图像合成的缩放整流变压器)

会议:ICML2024

摘要:扩散模型通过将数据向噪声方向的前向路径反转来从噪声中生成数据,并且已经成为高维、感知数据(如图像和视频)的强大生成建模技术。整流是最近出现的一种生成模型,它将数据和噪声在一条直线上联系起来。尽管它具有更好的理论性质和概念上的简单性,但它尚未被决定性地确立为标准实践。在这项工作中,我们改进了现有的噪声采样技术,通过将它们偏向于感知相关的尺度来训练整流模型。通过大规模的研究,我们证明了这种方法与已建立的高分辨率文本到图像合成的扩散公式相比具有优越的性能。此外,我们提出了一种新的基于转换器的文本到图像生成体系结构,该体系结构为两种模式使用单独的权重,并支持图像和文本标记之间的双向信息流,从而改进文本理解、排版和人类偏好评级。我们证明了这种架构遵循可预测的缩放趋势,并通过各种指标和人工评估将更低的验证损失与改进的文本到图像合成联系起来。我们最大的型号比最先进的型号性能更好。Stability AI正在考虑公开实验数据、代码和模型权重。


论文概述

今天为大家介绍一篇来自Stability AI团队发表在ICML 2024上的重要论文:《Scaling Rectified Flow Transformers for High-Resolution Image Synthesis》。这篇论文不仅在理论上推进了生成模型的发展,更在实践中展示了令人瞩目的成果——其8B参数模型在多项基准测试中超越了DALL-E 3等业界标杆。

一、背景:当前面临的挑战

1.1 Rectified Flow的潜力未被充分挖掘

Rectified Flow(修正流)是一种新型生成模型,它通过直线路径连接数据和噪声分布。相比传统扩散模型的弯曲路径,直线路径理论上可以:

  • 用单步进行模拟
  • 减少误差累积
  • 提高采样效率

然而,尽管有这些优势,Rectified Flow尚未在实践中成为标准方法。之前的研究主要局限于小规模的类条件模型,缺乏大规模文本到图像任务的验证。

1.2 现有架构的局限性

当前主流的文本到图像模型(如Stable Diffusion)通常采用:

  • 固定的文本表示通过交叉注意力输入模型
  • 文本和图像信息流是单向
  • 这种设计在文本理解、排版生成等方面存在瓶颈

1.3 扩展性研究的空白

与大语言模型(LLM)领域的scaling laws研究相比,文本到图像扩散模型的扩展规律尚未得到系统探索。

二、核心创新

2.1 改进的Rectified Flow训练方法

论文的第一个重大创新是重新设计了训练时的时间步采样策略

传统Rectified Flow的问题

经典Rectified Flow使用均匀分布采样时间步t ∈ [0,1]。但作者观察到,对于中间时刻(t ≈ 0.5),预测目标 ε - x₀ 更加困难,因为:

  • t=0时,最优预测接近p₁的均值
  • t=1时,最优预测接近p₀的均值
  • 中间时刻的不确定性最大
创新的采样策略

论文提出了三种新的时间步采样分布:

1. Logit-Normal采样

使用logit-normal分布,其密度函数为:

  • 参数m控制偏向数据端(负值)或噪声端(正值)
  • 参数s控制分布宽度
  • 最佳配置:m=0.00, s=1.00

2. Mode采样

定义映射函数:

可以控制是偏向中点(正s)还是端点(负s)。

3. CosMap

匹配cosine schedule的log-SNR:

实验验证在61种不同的训练配置对比中(包括ImageNet和CC12M两个数据集),rf/lognorm(0.00, 1.00)一致取得了优秀的排名,超越了使用均匀时间步采样的rectified flow,证实了中间时间步更为重要的假设。

特别值得注意的是,在少步采样(5步)时,Rectified Flow表现更优;而在25步及以上时,只有rf/lognorm(0.00, 1.00)能与传统的eps/linear保持竞争力。

2.2 MM-DiT:革命性的多模态架构

论文提出的MM-DiT(Multimodal Diffusion Transformer)是对DiT架构的重大改进:

核心设计理念

分离的权重参数

  • 文本和图像使用两套独立的transformer权重
  • 每个模态在自己的表示空间中工作

双向信息流

  • 在注意力操作中,文本和图像序列连接在一起
  • 允许文本→图像和图像→文本的相互影响
  • Q、K、V的计算涉及两个模态的信息
架构优势

与现有架构的对比实验显示:

架构 说明 性能
DiT 单套权重,序列连接 基线
CrossDiT 单套权重,交叉注意力 优于DiT
UViT U-Net与Transformer混合 初期快但后期落后
MM-DiT (2套权重) 分离权重,双向流 显著最优
MM-DiT (3套权重) CLIP和T5分开处理 略微提升,代价大

实验表明,MM-DiT在验证损失、CLIP分数和FID指标上都显著优于所有对比架构。

2.3 其他关键改进

改进的自编码器

使用16通道的潜在表示(相比传统的4或8通道):

  • 重建质量大幅提升
  • FID: 1.06(16通道)vs 2.41(4通道)
  • PSNR: 28.62 vs 25.12
合成标题混合

使用CogVLM生成合成标题,与原始标题50/50混合

  • GenEval整体得分从43.27提升到49.78
  • 颜色归因:11.75% → 24.75%
  • 计数能力:33.44% → 41.56%
  • 双对象:41.41% → 52.53%
高分辨率训练技术

QK归一化

  • 在Q和K计算后添加RMSNorm
  • 防止注意力logit增长不稳定
  • 支持bf16混合精度训练,性能提升2×

分辨率相关的时间步调整: 基于理论推导,对于分辨率m和n,实验确定α = √(m/n) = 3.0为最佳调整参数。

三、实验结果与发现

3.1 扩展性验证

论文进行了迄今为止最大规模的文本到图像模型扩展研究:

模型规模
  • 深度:15, 18, 21, 24, 30, 38层
  • 最大参数量8B参数(depth=38)
  • 训练规模:5×10²² FLOPs
  • 分辨率:256² → 512² → 1024²
关键发现

1. 平滑的扩展曲线

验证损失随着模型大小和训练步数平滑下降,未出现饱和:

  • 图像模型:500k步训练
  • 视频模型:140k步训练
  • 皮尔逊相关系数达到0.92以上

2. 验证损失的强预测性

验证损失与多个评估指标呈现强相关

  • GenEval:r = -0.920
  • T2I-CompBench:r = -0.898
  • 人类偏好(ELO):r = -0.982
  • 人类偏好(视频):r = -0.908

这意味着可以通过验证损失来预测模型的实际性能

3. 采样效率提升

更大的模型表现出更好的采样效率:

模型深度 5步相对下降 10步相对下降 20步相对下降 路径长度
depth=15 4.30% 0.86% 0.21% 191.13
depth=30 3.59% 0.70% 0.24% 187.96
depth=38 2.71% 0.14% 0.08% 185.96

depth=38模型的轨迹更接近直线,因此需要更少的采样步骤。

3.2 性能基准对比

GenEval基准

在最具挑战性的GenEval基准上:

模型 整体 单对象 双对象 计数 颜色 位置 颜色归因
DALL-E 3 0.67 0.96 0.83 0.87 0.47 0.43 0.45
本文(depth=38, 512²) 0.68 0.98 0.74 0.84 0.66 0.40 0.43
本文+DPO, 1024² 0.74 0.99 0.89 0.94 0.72 0.33 0.60

关键亮点

  • 基础模型在整体得分上首次超越DALL-E 3
  • 加入DPO后,整体得分提升到0.74
  • 在高分辨率(1024²)下,各项指标进一步提升
人类偏好评估

在Parti-prompts基准上进行的人类偏好评估:

vs. DALL-E 3

  • 视觉美学:约55%胜率
  • 提示遵循:约50%胜率
  • 排版生成:约65%胜率

vs. 其他SOTA模型(SDXL, Midjourney, Ideogram等):

  • 在大多数对比中获得60-80%的胜率
  • 排版生成能力尤为突出

3.3 额外发现

灵活的文本编码器

模型可以在推理时选择性使用文本编码器:

  • 移除T5:美学质量无影响(50%),提示遵循轻微下降(46%)
  • 但排版能力显著依赖T5(38%)
  • 节省4.7B参数的VRAM
零样本图像编辑

2B模型经过编辑任务微调后,展示了零样本文本操纵能力:

  • 可以在图像中编辑或插入文字
  • 即使训练数据中没有文本操纵任务
  • SDXL基线模型无法复现类似结果
数据去重效果

使用SSCD (threshold=0.5)进行去重:

  • 记忆样本数量减少5倍
  • 从8000个潜在记忆样本降至约95个(clique size=10)
  • 有效降低了版权风险

四、技术细节与洞察

4.1 训练效率优化

预计算嵌入

  • 预先计算VAE、CLIP、T5的输出
  • VAE编码时间:每样本2.45ms → 0ms
  • 训练每步节省约98%的文本编码时间
  • 代价:存储空间增加(每样本约630.7kB)

数据预处理

  • 使用NSFW检测过滤不当内容
  • 美学评分过滤低质量图像
  • SSCD聚类去重(4个阈值:0.50, 0.65, 0.75, 0.85)
  • 最终去除25-48%的训练数据

4.2 理论贡献

修正流的噪声预测目标

论文推导出修正流可以写成噪声预测形式:

这为统一理解不同扩散formulations提供了框架。

log-SNR视角

通过定义signal-to-noise ratio λ_t = log(a²_t/b²_t),可以统一分析:

  • 不同的λ_t调度对应不同的扩散模型
  • Rectified Flow: λ_t = log((1-t)²/t²)
  • Cosine: λ_t对应余弦调度
  • EDM: λ_t ~ N(-2P_m, (2P_s)²)

4.3 实现细节

优化器配置

  • AdamW优化器
  • 学习率:10⁻⁴
  • β₁=0.9, β₂=0.999
  • ε=10⁻¹⁵(用于bf16稳定性)
  • 1000步线性warmup

批次大小

  • 预训练256²:4096
  • 微调512²/1024²:根据模型大小调整
  • 使用梯度累积处理大模型

EMA

  • 衰减率:0.99
  • 每100个batch更新一次
  • 同时保存训练权重和EMA权重

五、局限性与未来方向

当前局限

  1. 计算资源需求:8B模型训练需要巨大的计算资源
  2. 长文本处理:受限于T5的77 token长度
  3. 视频生成:视频模型的扩展研究还处于初步阶段
  4. 一致性:同一概念的多图像一致性仍有提升空间

潜在改进方向

  1. 继续扩展:趋势未饱和,更大模型可能带来更多提升
  2. reflow操作:通过iterative reflow进一步直化路径
  3. 蒸馏:将大模型蒸馏到更小、更快的版本
  4. 架构创新:探索更高效的注意力机制
  5. 多模态扩展:整合音频、3D等其他模态

六、意义与影响

学术贡献

  1. 首次大规模验证了Rectified Flow在文本到图像任务上的优越性
  2. 提出了统一的框架来理解和比较不同的扩散formulations
  3. 建立了清晰的scaling laws,为未来研究提供指导
  4. 开创性的多模态架构为处理多种数据类型提供了新思路

工业影响

  1. 性能突破:开源模型首次在综合能力上超越DALL-E 3
  2. 效率提升:更少的采样步骤意味着更快的推理
  3. 可扩展性:清晰的扩展趋势降低了大模型开发的不确定性
  4. 开源承诺:Stability AI考虑公开代码、数据和权重

对AI生成内容的影响

  1. 质量飞跃:更好的提示理解和视觉质量
  2. 排版能力:显著改善的文字生成能力
  3. 创作灵活性:支持各种风格和构图
  4. 民主化:强大的开源模型降低了使用门槛

七、总结

这篇论文代表了文本到图像生成领域的一个重要里程碑。通过系统的研究和创新,作者们不仅证明了Rectified Flow的实用价值,还提出了革命性的MM-DiT架构,并通过大规模实验验证了模型的可扩展性。

核心价值

  1. 理论与实践结合:将理论优势转化为实际性能提升
  2. 系统化研究:61种配置的全面对比
  3. 可复现性:详细的技术细节和开源承诺
  4. 前瞻性:清晰的扩展路径为未来发展指明方向

对于研究者,这篇论文提供了丰富的insights和可复现的方法;对于开发者,它展示了构建SOTA文本到图像模型的完整pipeline;对于整个AI社区,它推动了开源生成模型向前迈进了一大步。

随着Stability AI承诺开源更多资源,我们有理由期待这项工作将激发更多创新,推动文本到图像生成技术进入新的发展阶段。


本文基于ICML 2024论文《Scaling Rectified Flow Transformers for High-Resolution Image Synthesis》撰写,旨在为中文读者提供详细的技术解析。

Logo

火山引擎开发者社区是火山引擎打造的AI技术生态平台,聚焦Agent与大模型开发,提供豆包系列模型(图像/视频/视觉)、智能分析与会话工具,并配套评测集、动手实验室及行业案例库。社区通过技术沙龙、挑战赛等活动促进开发者成长,新用户可领50万Tokens权益,助力构建智能应用。

更多推荐