Scaling Diffusion Transformers Efficiently via μP¶

会议: NeurIPS 2025
arXiv: 2505.15270
代码: 有
领域: Image Generation
关键词: 扩散 Transformer, μP, 超参数迁移, 模型缩放, 高效训练

一句话总结¶

将 Maximal Update Parametrization (μP) 从标准 Transformer 推广到扩散 Transformer（DiT、PixArt-α、MMDiT 等），证明其超参数可从小模型稳定迁移到大模型，显著降低大规模扩散模型的调参成本。

扩散 Transformer 已成为视觉生成模型的基础架构，广泛应用于图像和视频生成。然而随着模型规模增长到数十亿参数，超参数（HP）调优变得极其昂贵，往往阻碍模型发挥全部潜力。

μP 之前已被提出用于标准 Transformer（如 LLM），能使小模型搜索到的最优超参数直接迁移到大模型，大幅降低调参成本。但扩散 Transformer 与标准 Transformer 存在根本差异：

架构差异：扩散 Transformer 包含额外组件（adaLN、cross-attention 等）来整合文本和时间步信息

训练范式差异：基于迭代去噪的生成框架，不同于自回归范式

因此现有 μP 理论能否直接应用于扩散 Transformer 是一个未解之谜。本文系统性地解决了这一问题。

本文的方法论分三步：(1) 理论证明扩散 Transformer 的 μP 形式与标准 Transformer 一致；(2) 验证 DiT-μP 的超参数迁移性；(3) 在大规模文生图任务上验证 μTransfer 的效率。

μP 的理论推广（Theorem 3.1）：利用 Tensor Programs 技术，严格证明主流扩散 Transformer（U-ViT、DiT、PixArt-α、MMDiT）的前向传播可以在 Ne⊗or⊤ Program 框架下表示。关键在于证明扩散 Transformer 特有的模块（如 adaLN 块）也能用该框架的三个算子重写。这意味着 μP 的 abc-参数化规则直接适用：
- 输入权重：\(a_W=0, b_W=0, c_W=0\)
- 隐藏权重：\(a_W=0, b_W=1/2, c_W=1\)（学习率需要 \(\eta \cdot n_{base}/n\) 缩放）
- 输出权重：\(a_W=1, b_W=0, c_W=0\)（零初始化）
宽度缩放策略：固定注意力头维度，通过增加头数来缩放宽度。理论上（Bordelon et al., 2024），增大头维度会导致多头注意力退化为单头注意力，丧失注意力模式多样性。实践中也与 LLM 主流做法一致。
μTransfer 算法：代理模型（小宽度、小 batch、短训练）搜索最优基础超参数，然后直接迁移到目标大模型。代理和目标模型使用相同的 \(n_{base}\) 进行 μP 参数化，从而共享相同的最优基础超参数。

DiT 实验：固定头维度 72，base width 288（4 头），AdamW 优化器，无学习率调度和权重衰减
PixArt-α：代理模型 0.04B（4 头），目标模型 0.61B（16 头），代理训练 5 epoch
MMDiT-18B：代理模型 0.18B（宽度 512），对 4 个基础超参数（学习率、梯度裁剪、REPA 损失权重、warm-up 步数）进行 80 次随机搜索

配置	关键发现
学习率跨宽度迁移	最优基础学习率 \(2^{-10}\) 在宽度 144-1152 间稳定迁移
学习率跨 batch 迁移	256-1024 batch size 共享最优学习率
学习率跨训练步数迁移	150K-400K steps 共享最优学习率
MMDiT 超参数搜索	学习率影响最大；最优梯度裁剪为 1（非传统的 0.1）；warm-up 影响可忽略