Scaling Diffusion Transformers Efficiently via $μ$P

TL;DR

通过$μ$P高效扩展扩散Transformer,PixArt-$α$在0.61B参数下超越基线。

cs.LG 🔴 高级 2025-05-21 38 次浏览
Chenyu Zheng Xinyu Zhang Rongzhen Wang Wei Huang Zhi Tian Weilin Huang Jun Zhu Chongxuan Li
扩散Transformer 参数化 超参数转移 生成模型 图像生成

核心发现

方法论

本文将标准的最大更新参数化($μ$P)推广到扩散Transformer中,验证其在U-ViT、DiT、PixArt-$α$和MMDiT等主流模型中的有效性。使用Tensor Programs技术证明这些模型的$μ$P与传统Transformer一致,从而可以直接应用现有的$μ$P方法。

关键结果

  • DiT-XL-2-$μ$P通过学习率转移实现了2.9倍的收敛速度提升,相比原始DiT-XL-2。
  • PixArt-$α$在从0.04B扩展到0.61B时,$μ$P模型在保持低调优成本的情况下超越基线,仅需5.5%的训练运行成本。
  • MMDiT从0.18B扩展到18B,$μ$P模型在仅3%的人类专家消耗下表现优于基线。

研究意义

研究表明,$μ$P可以作为扩展扩散Transformer的高效框架,显著降低大规模模型的超参数调优成本。这为视觉生成模型的进一步扩展提供了理论和实践支持,解决了大规模模型调优成本高的问题。

技术贡献

技术贡献在于将$μ$P理论从传统Transformer扩展到扩散Transformer,提供了理论证明和实践验证,展示了$μ$P在不同模型宽度、批量大小和训练步数下的超参数转移能力。

新颖性

这是首次将$μ$P理论应用于扩散Transformer,证明其在不同架构和生成目标下的适用性,填补了现有研究的空白。

局限性

  • 尽管$μ$P在扩散Transformer中表现出色,但其在其他类型的生成模型中的适用性尚未验证。
  • 当前研究主要集中在视觉生成任务,其他任务的效果有待进一步研究。

未来方向

未来工作可以探索$μ$P在其他生成模型中的应用,以及在不同任务和数据集上的表现。此外,可以研究如何进一步优化$μ$P的参数化策略以提高模型性能。

AI 总览摘要

扩散Transformer因其在视觉生成模型中的广泛应用而备受关注,但其扩展性受到大规模超参数调优成本的限制。本文提出将最大更新参数化($μ$P)应用于扩散Transformer,验证了其在U-ViT、DiT、PixArt-$α$和MMDiT等模型中的有效性。

通过大规模实验,研究表明,$μ$P可以显著提高模型的超参数转移能力,降低调优成本。例如,DiT-XL-2-$μ$P在学习率转移后实现了2.9倍的收敛速度提升。PixArt-$α$在从0.04B扩展到0.61B时,$μ$P模型在保持低调优成本的情况下超越基线。

这些结果表明,$μ$P是扩展扩散Transformer的一个高效框架,为未来的研究提供了新的方向。然而,$μ$P在其他类型生成模型中的适用性尚需进一步验证。未来工作可以探索$μ$P在不同任务和数据集上的表现,以及如何进一步优化其参数化策略。

深度分析

研究背景

扩散Transformer已成为现代视觉生成模型的基础,应用于图像和视频生成等任务。然而,随着数据集的增长和任务复杂性的增加,扩散Transformer的进一步扩展变得不可避免。大规模模型的超参数调优成本高昂,成为限制其扩展的主要瓶颈。

核心问题

扩散Transformer在扩展过程中面临的核心问题是超参数调优成本高昂,尤其是在模型参数达到数十亿时。这限制了模型的性能发挥,迫切需要一种有效的方法来识别最佳超参数。

核心创新

本文的核心创新在于将最大更新参数化($μ$P)理论从传统Transformer扩展到扩散Transformer。通过理论证明和实践验证,展示了$μ$P在不同模型宽度、批量大小和训练步数下的超参数转移能力。

方法详解

  • �� 使用Tensor Programs技术证明扩散Transformer的$μ$P与传统Transformer一致。
  • �� 在ImageNet数据集上进行DiT-μP的系统研究,验证其超参数转移能力。
  • �� 将$μ$P应用于PixArt-$α$和MMDiT的文本到图像生成任务,验证其在大规模模型上的有效性。

实验设计

实验设计包括在ImageNet数据集上验证DiT-μP的超参数转移能力,以及在PixArt-$α$和MMDiT的文本到图像生成任务中验证$μ$P的有效性。使用的基线包括原始的DiT、PixArt-$α$和MMDiT模型。

结果分析

实验结果显示,DiT-XL-2-$μ$P通过学习率转移实现了2.9倍的收敛速度提升。PixArt-$α$在从0.04B扩展到0.61B时,$μ$P模型在保持低调优成本的情况下超越基线。

应用场景

$μ$P可用于大规模视觉生成模型的扩展,特别是在需要低调优成本的场景中。其在文本到图像生成任务中的成功应用表明其在多模态生成任务中的潜力。

局限与展望

尽管$μ$P在扩散Transformer中表现出色,但其在其他类型的生成模型中的适用性尚未验证。此外,当前研究主要集中在视觉生成任务,其他任务的效果有待进一步研究。

通俗解读 非专业人士也能看懂

想象你在厨房里做饭。传统的做法是每次做菜都要调整调料的量,这就像传统模型需要大量的超参数调优。而$μ$P就像一个万能调料包,只需在小锅里试一下,然后直接用在大锅里。这样,你就不用每次都重新调整调料了,省时省力。

简单解释 像给14岁少年讲一样

嘿,小伙伴们!想象你在玩一个游戏,每次升级都要重新设置角色的装备和技能,这是不是很麻烦?$μ$P就像一个超级道具包,你只需在低级别试用一下,然后就可以直接用在高级别,省去了很多麻烦!是不是很酷?

术语表

最大更新参数化 ($μ$P)

一种参数化方法,允许从小模型到大模型的稳定超参数转移,显著降低调优成本。

在扩散Transformer中应用以提高扩展效率。

扩散Transformer

一种用于视觉生成的模型,结合了Transformer架构和扩散过程。

作为现代视觉生成模型的基础。

超参数转移

将小模型中搜索到的最佳超参数直接应用于大模型的过程。

通过$μ$P实现,降低大规模模型的调优成本。

PixArt-$α$

一种用于文本到图像生成的扩散Transformer模型。

在实验中用于验证$μ$P的有效性。

MMDiT

一种多模态扩散Transformer模型,结合了图像和文本模态。

在实验中用于验证$μ$P的有效性。

开放问题 这项研究留下的未解疑问

  • 1 如何将$μ$P应用于其他类型的生成模型?现有研究主要集中在视觉生成任务,其他任务的效果尚不明确。
  • 2 在不同任务和数据集上,$μ$P的参数化策略是否需要调整以提高模型性能?

应用场景

近期应用

视觉生成模型扩展

通过$μ$P降低大规模视觉生成模型的超参数调优成本,适用于需要快速扩展的场景。

远期愿景

多模态生成任务

在多模态生成任务中应用$μ$P,提高模型的扩展性和性能,推动生成模型的进一步发展。

原文摘要

Diffusion Transformers have emerged as the foundation for vision generative models, but their scalability is limited by the high cost of hyperparameter (HP) tuning at large scales. Recently, Maximal Update Parametrization ($μ$P) was proposed for vanilla Transformers, which enables stable HP transfer from small to large language models, and dramatically reduces tuning costs. However, it remains unclear whether $μ$P of vanilla Transformers extends to diffusion Transformers, which differ architecturally and objectively. In this work, we generalize standard $μ$P to diffusion Transformers and validate its effectiveness through large-scale experiments. First, we rigorously prove that $μ$P of mainstream diffusion Transformers, including U-ViT, DiT, PixArt-$α$, and MMDiT, aligns with that of the vanilla Transformer, enabling the direct application of existing $μ$P methodologies. Leveraging this result, we systematically demonstrate that DiT-$μ$P enjoys robust HP transferability. Notably, DiT-XL-2-$μ$P with transferred learning rate achieves 2.9 times faster convergence than the original DiT-XL-2. Finally, we validate the effectiveness of $μ$P on text-to-image generation by scaling PixArt-$α$ from 0.04B to 0.61B and MMDiT from 0.18B to 18B. In both cases, models under $μ$P outperform their respective baselines while requiring small tuning cost, only 5.5% of one training run for PixArt-$α$ and 3% of consumption by human experts for MMDiT-18B. These results establish $μ$P as a principled and efficient framework for scaling diffusion Transformers.

cs.LG cs.AI cs.CV