人民大學&字節Seed：利用μP實現Diffusion Transformers高效擴展
本文由中國人民大學高瓴人工智能學院李崇軒團隊和字節跳動Seed團隊共同完成。第一作者鄭晨宇是中國人民大學高瓴人工智能學院二年級博士生，主要研究方向為基礎模型的優化、泛化和可擴展性理論，導師為李崇軒副教授，論文為其在字節跳動Seed實習期間完成。第二作者張新雨是字節跳動研究員，主要研究方向為視覺生成模型。李崇軒副教授為唯一通訊作者。
近年來，diffusion Transformers已經成為了現代視覺生成模型的主干網絡。隨著數據量和任務復雜度的進一步增加，diffusion Transformers的規模也在快速增長。然而在模型進一步擴大的過程中，如何調得較好的超參（如學習率）已經成為了一個巨大的問題，阻礙了大規模diffusion Transformers釋放其全部的潛能。
為此，人大高瓴李崇軒團隊和字節跳動Seed團隊的研究員引入了大語言模型訓練中的μP理論，并將其擴展到diffusion Transformers的訓練中。μP通過調整網絡不同模塊的初始化和學習率，實現不同大小diffusion Transformers共享最優的超參，使得小模型上搜到的超參可以直接遷移到最終大模型上進行訓練，從而極大地減小了超參搜索的耗費。
團隊在DiT，PixArt和MMDiT（Stable Diffusion的基座）上進行了系統的大規模實驗驗證。在MMDiT的實驗中，0.18B小模型上搜得的超參成功被用在18B大模型的訓練中，并擊敗了人工專家的手調基線。其中，小模型超參搜索的計算量（FLOPs）僅是專家手調的3%左右。
團隊已在近期開放在線論文，并開源代碼。
- 論文鏈接：https://arxiv.org/abs/2505.15270
- 代碼倉庫：https://github.com/ML-GSAI/Scaling-Diffusion-Transformers-muP
μP的背景和問題
μP全稱為最大更新參數化（Maximal Update Parametrization），是Tensor Program無窮寬網絡理論系列中的里程碑之作，相關結果已被理論證明適用于標準的Transformer架構。μP的算法實現簡潔，對于應用最為廣泛的AdamW優化器而言，μP只需要調整隱藏層權重的學習率，和輸出層權重的系數以及初始化。μP在實際中被廣泛發現能夠實現不同大小的標準Transformer共享最優的超參，使得小模型上搜到的超參可以直接遷移到大模型，極大地減小了超參搜索的耗費。由于μP帶來了穩定的超參遷移性質，它近年來已經被成功使用在大語言模型（標準Transformer）的預訓練中。
然而，diffusion Transformers和標準Transformer存在較大的差異。從架構上來看，diffusion Transformers引入了額外的模塊來處理并整合文本信息，如DiT中的adaLN block。從任務目標上來看，diffusion Transformers處理的是視覺的擴散學習任務，而標準Transformer主要處理的是語言的自回歸學習任務。這兩點差異意味著已有的μP形式及其超參遷移律在視覺diffusion Transformers中不一定成立。針對這一問題，團隊從理論和實踐上進行了系統的研究。
Diffusion Transformers的μP形式
團隊首先從理論上研究了主流diffusion Transformers的μP形式，包括DiT，U-ViT，PixArt-α和MMDiT。Tensor Program理論系列中的結果表明，如果網絡架構能夠被Tensor Program中定義的算子表示，那么現有的μP形式就能成立。基于這個理論技術，我們證明了：即使主流diffusion Transformers的結構不同于標準Transformer，它們也能夠被Tensor Program表示，因此現有的μP理論和相關實踐可以被無痛遷移到這些主流diffusion Transformers上。我們的證明技術也可以被遷移到其它的diffusion Transformers做類似的分析。
總之，diffusion Transformers的μP方法論可以由下圖總結。我們首先基于μP理論，調節不同權重的系數、初始化和學習率。然后，我們在一系列小模型上搜索得到最優的超參。最后，我們將最優的超參直接遷移到大模型的訓練。
基于μP擴展Diffusion Transformers：初探
首先，我們使用DiT網絡在ImageNet數據集上系統地驗證了：當網絡寬度，數據批量大小和訓練步數足夠大時（如寬度達到144，批量大小達到256），超參便可以較為穩定地沿著不同的網絡寬度，數據批量大小和訓練步數進行遷移。這意味著我們能在網絡寬度，數據批量大小和訓練步數都更小的代理任務上搜索超參，然后遷移到最終大網絡大數據的訓練。
然后，為了驗證μP超參遷移的有效性，我們將最優的超參（學習率2^-10）直接遷移到DiT-XL-2的訓練中，我們發現，當模型訓練到2.4M步時，FID-50K就已經超過了原論文7M步最終的FID-50K結果，DiT-XL-2-μP的收斂速度是原論文的2.9倍。這向我們展現了利用μP遷移超參做擴展的良好前景。
基于μP擴展Diffusion Transformers：大規模驗證
我們進一步在大規模的文生圖任務上驗證了μP擴展diffusion Transformers的有效性。我們首先考慮了流行的開源文生圖模型PixArt-α，我們在0.04B的代理模型上搜索學習率，并遷移到最終0.61B大小PixArt-α的訓練。其中，小模型搜索超參的計算量總和（FLOPs）僅為一次訓練的5.5%。利用搜索得到的學習率，PixArt-α-μP在訓練的過程中穩定地取得了比基線更好的效果。
最后，我們考慮了SD3的基座模型MMDiT，并將驗證的規模提高到了18B的量級。為了能夠給社區帶來更多的可信的實踐經驗，我們在 4個超參（學習率，梯度裁剪值，REPA loss的權重以及warmup的步數）上進行了多達80次的隨機搜索，總搜索計算量（FLOPs）約是人工手調的3%。在0.18B模型上的超參搜索結果表明，我們學習率，梯度裁剪值，REPA loss都對結果有影響，其中學習率的影響仍是最為關鍵的。而warmup的步數則對結果影響不大。
我們將0.18B模型上搜索的超參應用在了18B模型的訓練上，不論從訓練loss的變化還是從人工評測的結果，MMDiT-μP都穩定地超過了人工專家手調的基線，而μP的超參搜索FLOPs僅是人工手調的3%！
經過這一系列系統的實驗探索，我們證明了μP是科學擴展diffusion Transformers的有效手段，我們也相信μP會是未來基礎模型擴展的必備利器。通過本工作的大量努力，我們希望讓社區了解μP理論，擁抱μP實踐，思考理論上最優的智能擴展范式（模型大小，數據量，推理時間）。我們也相信，放眼人工智能的長遠未來，類似μP的底層理論的發展仍然是必不可少的，也必將會在未來的大規模實踐中有著不可或缺的一席之地。
