Scaling Diffusion Transformers Efficiently via $\mu$P
Authors: Chenyu Zheng
Xinyu Zhang
Rongzhen Wang
Wei Huang
Zhi Tian
Weilin Huang
Jun Zhu
Chongxuan Li
Venue: NeurIPS 2025 poster ()
Keywords: Diffusion Transformer
muP
Tensor Program
image generation
AIGC
Scaling
Primary area: deep_learning
TL;DR: Scaling Diffusion Transformers up to 18B Efficiently via $\mu$P

Abstract: 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 ($\mu$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 $\mu$P of vanilla Transformers extends to diffusion Transformers, which differ architecturally and objectively. In this work, we generalize $\mu$P to diffusion Transformers and validate its effectiveness through large-scale experiments. First, we rigorously prove that $\mu$P of mainstream diffusion Transformers, including DiT, U-ViT, PixArt-$\alpha$, and MMDiT, aligns with that of the vanilla Transformer, enabling the direct application of existing $\mu$P methodologies. Leveraging this result, we systematically demonstrate that DiT-$\mu$P enjoys robust HP transferability. Notably, DiT-XL-2-$\mu$P with transferred learning rate achieves 2.9$\times$ faster convergence than the original DiT-XL-2. Finally, we validate the effectiveness of $\mu$P on text-to-image generation by scaling PixArt-$\alpha$ from 0.04B to 0.61B and MMDiT from 0.18B to 18B. In both cases, models under $\mu$P outperform their respective baselines while requiring small tuning cost—only 5.5% of one training run for PixArt-$\alpha$ and 3% of consumption by human experts for MMDiT-18B. \textit{These results establish $\mu$P as a principled and efficient framework for scaling diffusion Transformers}.

Scaling Diffusion Transformers Efficiently via pP
Chenyu Zheng' 1,2,3 Xinyu Zhang', Rongzhen Wang 1,2.3 Wei Huang', Zhi Tian', Weilin Huang', Jun Zhu", Chongxuan Li3.2.3
Gaoling School of Artificial Intelligence, Renmin University of China
2 Beijing Key Laboratory of Research Large Models and Intelligent Governance
3 Engineering Research Center of Next-Generation Intelligent Search and Recommendation, MOE ByteDance Seed RIKEN AIP Dept. of Comp. Sci. & Tech., Tsinghua University
Abstract
Diffusion Transformers have emerged as the foundation for vision generative mod- els, 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 lan- guage models, and dramatically reduces tuning costs. However, it remains unclear whether μuP 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 jPP of mainstream diffusion Transformers, including DiT, U-ViT, PixArt-a: and MMDIT, aligns with that of the vanilla Transformer, enabling the direct application of existing JP methodologies. Leveraging this result, we systematically demonstrate that DiT-pPP enjoys robust HP transferability. Notably, DiT-XL-2-pPP with transferred learning rate achieves 2.9x faster convergence than the original DiT-XL-2. Finally, validate the effectiveness of μP text-to-image generation by scaling PixArt-a from 0.04B to 0.61B and MMDiT from 0. 08B to 18B. In both cases, models under JPP outperform their respective baselines while requiring small tuning cost-only 5.5% of one training run for PixArt-a and 3% of consumption by human experts for MMDiT-I8D These results establish μP as a principled and efficient framework for scaling diffusion Transformers.
Introduction
Owing its generality and scalability, diffusion Transformers 52 have become the backbone of modern vision generation models, with widespread applications in various tasks such as image
generation (55) 3 16 35, 20 and video generation [6] 8T 62 777 As datasets grow and task complexity increases, further scaling of diffusion Transformers has become inevitable and 75). However, as model sizes reach billions of parameters, hyperparameter (HP) tuning becomes prohibitively expensive, often hindering the model from achieving its full potential. This underscores the urgent need for principled approach to efficiently identify the optimal HPs when scaling diffusion Transformers.
attracting increasing attention [41 40
was recently proposed as promising solution to the HP selection problem for large-scale vanilla Transformer [61]. It stabilizes optimal HPs across different model widths, enabling direct transfer HPss searched from small models to large models (a.k.a,, μ I/Trnsffr) and significantly reducing tuning costs at scale. Due to its strong transferability.
Maximal Update Parametrization (nP) 73 70
μP has been applied the pretraining large language models (LLMs) 73
28
Work done during an internship at ByteDance Seed. 'Corespondence to Chongxuan Li <chongruanli wrrueeuuaanniuu.e..oooooo edu. cn>
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

(i) PixArt-a
100%
Relative FLOPs
80%
5.5%
20%
colorful donut fixed on a white clöthesline, with Bright sky and fluffy white clouds in the background
furry fittle black cat looked up and saw beautiful butterfit landing on its nose
Blue sky, calm a fuuge blue whale leaps from the sea, and the yacht below small
0%
1x Train
μP Search
(i) MMDIT-18B
100%
Relative FLOPs
80%
3%
0%
Cat Picasso's Painting Style
teddy bear holding a guitar stands in scene with star lightss and grass
Flowers grow inside the spacesuit
Manual Tuning
μPP Search
grass
(a) Samples produced by the MMDIT-pP-18BB
(h) Efficiency of jP search.
Figure 1: Visualization results and efficiency of HP search under JaP (a) Samples generated by the MMDiT- wP--1BB model exhibit strong fidelity and precision in aligning with the provided textual descriptions. (b) HP scarch for large diffusion Transformers efficient under J.P requiring only 5.5% FLOPs of single training for PixArt-a and just 3% FLOPs the human experts for MMDiT-188
differ fundamentally from vanilla Transform- ers [61]. First, their architectures incorporate additional components to integrate text information and diffusion timesteps for final vision generation. Second, they operate under distinct generative framework based on iterative denoising, in contrast to the autoregressive paradigms typically used in vanilla Transformers (e.g., LLMs) 22 80 As result, existing μP theory and its associated HP transfer properties may not directly apply to diffusion Transformers. This paper systematically investigates this issue, as detailed below.
Unfortunately, diffusion Transformers 52
First, we extend the standard juP theory from vanilla Transformers to diffusion Transformers in
Section 3.1 Using the Tensor Programs technique [72 70
67
rigorously prove that the juP
formulation of mainstream diffusion Transformers, including DiT (52) U-ViT [T] PixArt-a [8] and
MMDiT from Stable Diffusion 3 (16), matches that of vanilla Transformers (Theorem 3.1). This compatibility enables us to directly apply existing practical methodologies developed for standard μP
diffusion Transformers, as described in Section 3.2
Based on the rigorous yP result of diffusion Transformers, we then conduct systematic study of DiT-pPP the image generation task using the ImageNet dataset [33, presented in Section 4 We first verify the stable HP transferability of DiT-uP across widths, batch sizes, and training steps. We then j'Transfer the optimal learning rate searched from small models DiT-XL-22ppP. Notably, DiT- XL--2-uP trained with the transferred learning rate achieves 2.9x faster convergence than the original DiT-XL-2
[52]
suggesting that μP offers an efficient principle for scaling diffusion Transformers.
Finally, we further validate the efficiency of μP on large-scale text-to-image generation tasks. In Section 5.1 we apply the j'Transfer algorithm to PixArt-a [8], scaling the model from 0.04B to 0.61B. In Section 5.2 we apply to MMDiT [6], scaling from 08 to 18B. In both cases, diffusion Transformers under /PP consistently outperform their respective baselines with small HP tuning cost on proxy tasks. For PixArt-a. tuning consumes only 5.5% of the FLOPs required for a full pretraining run, while for MMDiT, tuning uses just 3% of the FLOPs typically consumed by human experts. These real-world experiments further confirm the scalability and reliability of jP.
2 Preliminaries
We begin by establishing the necessary background for diffusion Transformers and μP. Detailed discussion of additional related work is placed in Appendix

---

Table JaP for (difflusion) Transformers with differences between μP and standard parameterization (SP) in practice (e.g.. Kaiming initialization 25 and gray text to indicate the SP settings. Formal definitions of the weight type are provided in Appendix
optimizer. We purple to highlight the
Input weights
Hidden weights
Output weights
aw
0
(0)
bw
0
1/2
(1/2)
CW
0
(0)
0
2.1 Diffusion Transformers
Due to its superior scalability and compatibility, Transformers 152 1:16 30 have replaced CNNs
(e.g., U-Net (54] 53 63) as the backbone for advanced diffusion models [55]
3 16. 35
20).. U-
VIT I firstly introduces Transformers with long skip connections between shallow and deep layers for various image generation tasks and achieves remarkable performance. After that, DiT [22 proposes Transformers with adaptive layer normalization (adaLN) blocks for class-to-image generation and demonstrates strong scalability with respect to network complexity. PixArt-a [8] extends DiT for text-conditional image generation by incorporating cross-altention for text features and efficient shared adal .N-single block. MMDiT [16] further extends DiT by introducing two separate parameter sets for image and text modalities, along with joint attention mechanism to facilitate multimodal interaction. As diffusion Transformers continue to scale, increasing attention is being paid to principled approaches for scaling [4] 40 29 75
2.2 Maximal Update Parametrization
In this section, we provide practical overview of μP, with focus the widely used AdamW optimizer [47). The comprehensive review the theoretical foundations of μP is in Appendix
μP identifies a unified parameterization that applies to common architectures expressive in NE Program 72 (e.g.. vanilla Transformer [6]), offering strong guidance for practical HP transfer across model widths, batch sizes, and training steps. Under JPP, HPs can be tuned on small proxy task (e.g., 0.04B parameters and 6B tokens in (73) and directly transferred a large-scale pretraining task (e.g.. 6.7B parameters and 300B tokens in [73 ). significantly reducing tuning costs at scale. As
a result, μP has been widely adopted in the pretraining of L.LMS [73 14
82
Concretely, μP is implemented by analytically adjusting HPs with model width, typically involving the weight multiplier, initialization variance, and learning rate (a.k.a., abc- bellee-rramttrrzationnn Formally, let n denote the network width, we set each weight as W = =wwn aa W, where the trainable component W is initialized as Waj ~ (0, (0, (bbw) ), and its learning rate is 7w -w Henceforth, we call the width- hntth-independent parts (øw, ow. 7/) base HPs. As summarized Table μP identifies values of aw, bw: and ew that enable models of different widths to share the (approximately same optimal base HPs swy, ow, and l/y- jP adjustion for other HPss is placed Appendix B.4
The most related work to us is AuraFlow v0.1 [T), which applied μP empirically to an MMDiT-style model for learning rate transfer. Compared AuraFlow v0.1, our work makes additional non-trivial contributions. First, we provide a rigorous theoretical proof for mainstream diffusion Transformers, thereby justifying the validity of jPP for this family of models. Second, systematically validate the HP transferability of diffusion Transformers under jPP across multiple widths, batch sizes, and training steps. Finally, for MMDiT in particular, conduct a more extensive search over multiple HPs and scale the model up to 18B parameters, providing detailed intermediate results and comparisons.
3 Scaling Diffusion Transformers by jPP
In this section, we extend the principles of μP to scale diffusion Transformers. In Section 3.1 we prove that despite fundamental differences from vanilla Transformer, mainstream diffusion Transformers-includee standard μP formulation summarized in Table T. Then, Section 3.2 we introduce the practical methodology for applying μP to diffusion Transformers, illustrated Figure 2
U-ViT []. DIT [52] PixArt-a [8) and MMDiT [16] add-adhere to the
3

---

Input Woights
Hidden Weights
μP search base HPss small models
Output Weights
New in DFT
Noise
Cov E
Scale
HPS-3
Pointwise
Linear and Reshape
Feoelforraadd
Layer Norm
Scale, Shift
Layer Norm
uTransfer
DFT Block
Scaling to large models with the optimal base HPs
Multi Ilead Self-Attemtion
Patchify
A-Emih
Sctie Shift
Layer Norm
AdaLN
HPs-2
Noised
t-Enib
y y-Emb!
Latent
Timestep ?
Tabely
Tokena
Cancitio
(a) Implementation of DiT-jdP.
(b) p'Transfer.
Figure 2:
A overview of applying ,P to diffusion Transformers. (a) We illustrate the implementation of pP for DiT as example. The abe-pan pare-aramtterizaionn of each weight is adjusted based on its type and visualized using different colors. Modules that differ from the vanilla Transformer also highlighted. (b) We μ TTrrnsfer the optimal base HPs searched from multiple trials on small models pretrain the target large models.
3.1 /cP of Diffusion Transformers
As mentioned in Section 2 the existing uP results 72 in Table apply only architectures expressible a NE@OR program. Therefore, it is crucial to determine whether the diffusion Transformers can be represented this framework. The following theorem establishes that several are compatible with the existing pP results.
prominent diffusion Transformers [52 128. [6]
Theorem 3.1 (μP of diffusion Transformers, proof in Appendix C.. The forward passes of mainstream diffussion Transformers (U-VTT (,, DIT (52), Pixart-a 18), and MMDIT (I6]) can be represented within the NEMOR T Program. Therefore, their uP matches the standard μP presented in Table
The proof of Theorem 3.1 relies on rewriting the forward pass diffusion Transformers using the three operators defined in the NETORT Program [72]. Given that the representability of standard Transformer [61] within the NESORT Program has been established [67]. we focus on demon- strating that the novel modules specific diffusion Transformers can also be expressed in this formalism. These modules primarily serve to integrate text and diffusion timestep information for vision generation; examples include the adaLN blocks of DiT [22 in Figure 2a
Theorem 3.I ensures that existing practical methodologies developed for standard yaP apply directly to DIT, V-VIT, PixArt-a, and MMDiT. Moreover, analysis technique naturally extends to other variants of diffusion Transformers 49 19 59 48) To the best our knowledge, mainstream
diffusion Transformer variants in use can be expressed within the NE3OR Program.
3.2 Practical Methodology of μP in Diffusion Transformers
Given the rigorous jP result for diffusion Transformers established in Theorem 3.1 this section introduces how to apply μP to diffusion Transformers in practice. addition, our code is available at http::/htulllooolll.////////iiit
3.2.1 Implementation of pP
The width of a multi-head attention layer is determined by the product of the head dimension and the number of attention heads. Thus, there are two degrees of freedom when scaling the width of diffusion Transformers. Theoretically, recent theoretical advances [5] reveal an important difference: when the head dimension tends to infinity, multi-head self-attention can collapse to single-head self-attention, losing the diversity of attention patterns. In contrast, scaling the number of heads avoids this degeneracy and preserves the expressivity of the multi-head mechanism. Empirically, well-known models in real-world large-scale practice [73 22 43 28 favor increasing the number of

---

heads rather than the head dimension. Given these theoretical and empirical insights, we fix the head dimension and scale the number of heads in this work.
Given scaling the number of heads, we implement the jPP of diffusion Transformers following the standard procedure in 73 Specifically, we replace the vanilla width in the abe-parameterization (see Section 2 with the width ratio nt/Tibae when standard parameterization (SP) and μP differ in Table ! where Thass is fixed base width. This implementation remains consistent with Table I since Thase is a constant. For example, for hidden weight matrix W with width n and base HPs
Øw, ew,oly, (maybe searched from proxy model), we set W éw W with initial W ~ (00, oty /n), as in SP. In contrast, its learning rate is myaase n rather than 17/n, n, where jP differs from SP. Finally, we also follow the suggestion from [73 to initialize the output weights by zero (i.e., o 0).
3.2.2 Base Hyperparameter Transferability of jaP
Prior work has shown that μP provides strong guidance for base HP transferability in image classifi- cation and language modeling tasks 773 28 64
24: 60; 14]. Building on the rigorous μP form
for diffusion Transformers in Theorem 3.1 we further validate its HP transferability in the context of visual generation. We summarize the methodology for verifying base HP transferability across widths, batch sizes, and training steps in Algorithm 4 and 3 in Appendix [D] respectively. We describe how verify HP transfer across widths below; those for batch size and step similar.
For simplicity, we describe how to evaluate the transferability of the base learning rate v across widths; similar procedures apply to other HPs, such as the multiplier ø. We define ) and {n, 1 as the sets of widths and base learning rates used in the evaluation. Given a fixed batch size B and training iterations of T, we train PR diffusion Transformers, each parameterized by pP with base width Basae, true width 72. and base learning rate 77. Finally, given evaluation metrics, if models at different widths (approximately) share the same optimal base learning rate, we conclude that learning rate transferability holds for diffusion Transformers. In Section we verify that diffusion Transformers under the jP exhibit robust transferability of base HPs.
3.2.3
p Transfer from Proxy Task to Target Task
Once base HP transferability is validated for diffusion Transformers, we can directly apply the j'Transfer algorithm [7] (see Algorithm 4 in Appendix D to zero-shot transfer base HPs from a proxy task to target task. Specifically, both the proxy model with width Rproxy and the target model with width Rtarget are parameterized by μP using the same base width 7hasee We first search for an optimal combination of base HPs using various proxy models trained with a small batch size Broxy and a limited training steps proxy.. These optimal base HPs are then used to train the large target model with larger batch size Btarget and longer training iteration T Trarget- Experimental results in Section 5 show the superior performance of the TTrranffrr in worldwoldd vision generation.
4 Systematic Investigation for DiT-pP on ImageNet
In this section, we first empirically verify the base HP transferability of DiT (52) under pP. We then ji Transfer the optimal learning rate to train DiT-XL-2-pPP which achieves 2.9x faster convergence.
4.1 Basic Experimental Settings
To ensure fair comparison between DiT-uP and the original DiT [52] we adopt the default configurations from 52 and describe the basic setup in detail below.
Dataset. We train DiT and DiT-pP on the ImageNet training set [13, which contains 1.28M images across 1,000 classes. All images are resized to a resolution of 256 256 during training, following
standard practice in generative modeling benchmarks 52 49
Architecture of DiT The architecture of DiT- models is identical to that of DiT-XL--2, except for the width. We fix the attention head dimension at 72 (as in DiT-XL-2) and vary the number of heads. The base width Thase in the pP setup corresponds 288, which uses 4 attention heads.
To scaling head dimension, Yang et al. [73 additionally change the calculation of attention logit from k/ / to q k/d, where query g and key k have dimension d. We do not use this modification [5].
5

---

100
Width 144
90
128
100-
iteracion 1enazoon-100
288
Batch 256
FID-50K
Width-576
80
Batch size-512
FID 10-50K
90
ateration 2orrnion--000 -200
60
70
40
60
-13
-12
11
10
-9
-13
11
10
-13
-12
-11
-10
loge
log (base
logo (base
(a) HP transfer across widths.
(b) HP transfer across batch sizes. (c) HP transfer across iterations.
Figure 3: DiT-uPP enjoys base HP transferability. Unless otherwise specified, we use a model width 288, a batch size of 256, and training iteration of 200K. Missing data points indicate training instability, where the loss explodes. Under μP, the base learning rate can be transferred across model widths, batch sizes, and steps.
Training. We train DiT and DiT using the AdamW W (47. Following the original DiT setup [2] we do not apply any learning rate schedule or weight decay, and constant learning rates are used in all experiments. The original D is trained with a learning rate 10-4 and a batch size of 256.
report FID (26). sFID (50), Inception Score [56], precision, and recall (34] on 50K generated samples without classifier-free guidance (cfg), in Table 4 of (52). In the main text, we present the FID results, while the remaining metrics are provided in Appendix E.2
Evaluation metrics. To comprehensively evaluate generation performance,
4.2 Base Hyperparameters Transferability of DiT-MPP
In this section, we evaluate the HP transferability of DiT-uP across different widths, batch sizes, and training steps. We focus primarily on the base learning rate, as has the most significant
impact performance 773 28
43
Similar results for other HPP (weight multiplier) are
provided in Appendix E.2 We sweep the base learning rate over the set (2-13, 2 2-12, 1-122-111 2-1,111111 2- 10 10.2-9).. 2-
12
across various widths, batch sizes, and training steps. FID-50K results
shown in Figure and comprehensive results for other metrics are presented in Tables 8 | and 10 in Appendix E.2
As presented in Figure the optimal base learning rate 2- generally transfers across scaling dimensions when some minimum width (e-g. 144), batch size (e.g., 256), and training steps (e-g.. 1 150K) are satisfied, which verifies the base HP transferability of DīT-uP.. Interestingly, we observe that neural networks with μPP tend to favor a large learning rate close to the maximum stable value (e.g., see Figure Ba). This aligns with empirical findings and theoretical insights reported for standard neural networks trained for multiple epochs (65) rates introduce beneficial gradient noise help guide optimization towards flatter minima that generalize better. Our empirical results suggest that the optimization landscape of jPP shares certain similarities with that of SP, offering a direction for future theoretical investigation.
I. which suggest that larger learning
4.3 Scaling Performance of DiT-
We μTransfer the optimal base learning rate of 2- to train D-T-XL-2-p1P with a width of 1152. The batch size is set to 256, following [52]. We evaluate DiT-uPP every 400K steps without using cfg. Training continues until DiT surpasses the best performance reported by the original DiT at final 7M steps [52]. To enable a detailed comparison throughout the training, we also reproduce the original DiT training using its official codebase. Complete evaluation results throughout training are provided in Table 14 in Appendix
20.0
DFFX1-2
12--X--2-PP
17.5
FID 50K
15.0
12.5
2.9x faster
E.2
10.0
As shown in Figure 4 DIT-XL-2--P with the transferred base learning rate performs effectively, achieving 2.9x faster con- vergence (2.4M steps) compared to the original DiT (7M steps). These results suggest that μP offers a simple and promising approach to improve the pretraining of large-scale diffusion Transformers. In the following sections, we further validate this claim through text-to-image generation tasks.
IeeraiionM
Figure 4: μP accerlates the train- ing of diffusion Transformers. Con- sidering FID-50K, DiT-XL-2-ppP with transferred learning rate achieves 2.9x faster convergence than the original DIT-XL-2 and slightly better result.

---

5
Large-Scale Text-to-lmage Generation
In this section, we verify the efficiency of pTransfer algorithm on real-world text-to-image generation tasks. Diffusion Transformers under μP outperform the baselines while requiring small tuning cost.
5.1 Scaling PixArt-a-jaP on SA-1B
In this section, we perform yTransfer experiments the PixArt-a [ scaling from 0.04B to 0.61B parameters. Using the same pretraining setup, PixArt-a-jaP with the transferred learning rate outperforms the original PixArt-a, while incurring only 5.5% FL.OPs of one full pretraining run.
5.1.1 Experimental Settings
To ensure the fairest possible comparison between PixArt-a-uPP and the original PixArt-a [8] wc mainly adopt the original setup [8 and summarize the key components below.
Dataset. We use the SAM/SA-1B dataset 333, which contains 11M high-quality images curated for segmentation tasks with diverse object-rich scenes. For text captions, we use the SAM-LLaVA annotations released in g All images are resized to a resolution of 256 256 during training.
Architecture of PixArt-a-juP models. The target PixArt-a-pP model adopts the same architecture as PixArt-a (0.61B parameters). The proxy model also follows the same architecture, differing only in width. To construct the proxy Pix model, we fix the attention head dimension at 72 (as in PixArt-a) and reduce the number of heads from 16 4 (0.04B parameters). In the pP framework, the base width Thbse (see Section 3.2.1) is set to the proxy width of 288.
Training, PixArt-a is implemented using the official codebase and original configuration. We train the original PixArt-a and the target PixArt-a-jiP model for 30 epochs with a batch size of 176 x 32 (approximately 59K steps) | The small proxy PixArt-a--aP models trained for 5 epochs with batch size of 176 x 8 (approximately 39K steps). Notably, the model width, batch size, and training steps are all smaller than those used in the largel pretraining setling.
Hyperparameter search. this section, we focus solely on the base learning rate. We search over the candidate set {2 13 2-12 2 2 10 2 "} as in Section 4 resulting in five proxy training trials.
Ratio of tuning cost to pretraining cost. We consider the FLOPs as the metric for the computational cost, then the ratio of tuning cost to pretraining cost can be estimated as
R.Prraxy Exroxy Starget BrangetBrargee TtargetTiagge
R.prraxy Brraxyy T proxy
ratio-
Starget
where R is the number of HP search trials, S is the number of parameters, is the training epochs, B is the batch size and T is the training iteration. The ratio 5.5% here, detailed in Appendix E.3
Evaluation metrics. We evaluate text-to-image generation performance following standard prac- tice 8& [66 (66 [], including FID, CLIP Score, and GenEval [21]. Both FID and CLIP Score are computed on the aesthetic MJHQ-30K 339 and real MS-COCO-30K (45) datasets. MJHQ-30K contains 30K images generated by Midjourney, while MS-COCO-30K is randomly sampled subset of the MS-COCO [45] dataset. GenEval evaluates text-image alignment using 533 test prompts. Following the official implementation [8]. we use cfg of 4.5 to generate samples.
5.1.2 Experimental Results of PixArt-a-juP
We begin by conducting a base learning rate search using the PixArt-a-juP proxy models. The evaluation results for different base learning rates summarized in Table Details of GenEval results can be found in Table 15 jn Appendix E.3 Since overfitting is not observed in this setting, we include training loss as an additional evaluation metric. Overall, the base learning rate of 2 yields the best performance. Interestingly, this optimal learning rate matches that of DiT- (see Figure Ba). We hypothesize that this consistency arises from the architectural similarity between the two models and the fact that both the ImageNet and SAM datasets consist of real-world images. This observation suggests that optimal base HPs may exhibit some degree of transferability across different datasets and architectures.
We confirmed with the authors that this setup is reasonable; longer training may result in overfitting.
7

---

Table 2: Results of base learning rate search PixArt-a-pP proxy tasks. 0.04B proxy models with different base learning rates are trained for epochs on SAM dataset. Overall, base learning rate optimal.
log2(1 (Ir)
Training loss
GenEval f
FID-30K (MS-COCO)
FID-30K (MJHQ)
-9
NaN
NaN
NaN
NaN
-10
0.1493
0.083
47.46
47.71
-11
0.1494
0.078
49.24
46.31
-12
0,1504
0,030
66.77
63,36
-13 0.1536
0.051
60.28
60.93
Table 3: Comprehensive comparison between PixArt-o-pPP and PixArt-a. Both models trained the SAM dataset epochs. PixArt-a-juP (0.61B), using base learning rate transferred from optimal 0.04B proxy model, consistently outperforms the original baseline throughout the training process.
MJHQ
MS-COCO
Epoch
Method
GenEval
FID-30K &
CLIP Score f FID-30K
CLIP Score ?
PixArt-a
8
0.19
38.36
25.78
34.58
28.12
10
PixArt-a-pP (Ours)
0.20
33.35
26.25
29.68
28.87
PixArt-a:
8
0.20
35.68
26.54
30.13
28.81
20
PixArt-a-uP (Ours)
0.23
33.42
26.83
29.05
29.53
PixArt-a
g
0.15
42.71
26.25
37.61
28.91
30
PixArt-a-pPP (Ours)
0.26
29.96
27.13
25.84
29.58
We then apply j/Transfer by transferring the searched optimal base learning rate of 2- 10 to train the target PixArt-a. A comparison between PixArt-a and PixArt-a-pP throughout training is provided Table (with the complete results shown in Table [6 Appendix[E.3 The results demonstrate that PixArt-a-jaP consistently outperforms PixArt-a across all evaluation metrics during training, supporting jP as an efficient and robust approach for scaling diffusion Transformers. Furthermore, observe that the benchmark performance of PixArt-a degrades after 20 epochs, primarily due to overfitting. In contrast, PixArt-a-puP continues to improve, suggesting that μP enhances the model's generalization ability, which offers interesting direction for future theoretical investigation.
Specifically, we emphasize that our current experimental results validate that jP works well without cfg (DIT) and with efg (PixArt-a). Because efg is an important factor during inference and affects the optimal hyperparameters, we strongly recommend practitioners align the evaluation proxy models and that of the target model. In the following section, we further extend the current method to large-scale applications.
5.2 Scaling MMDIT to 18B
In this section, we validate the efficiency of μP in the large-scale setup. We scale up the MMDIT [6 architecture from 0.18B 18B. Under the same pretraining setup. MMDDT-pP-18B with the trans- ferred base HPs outperforms the MMDiT-18B tuned by human algorithmic experts.
5.2.1 Experimental Settings
Dataset. We train models internally constructed dataset comprising 820M high-quality image- text pairs. All images are resized to a resolution of 256 x 256 during training.
Baseline MMDiT-18B. The width and depth of MMDiT- 18B are 5,120 and 16, respectively. The training objective combines flow matching loss 46 and a representation alignment (REPA) loss (76). The model is optimized by Adam [47 with batch size of 4,096 and 200K training iterations. The learning rate schedule constant with warm-up duration. The HPs were tuned by algorithmic experts, requiring roughly 5 times the cost of full pretraining, detailed in Appendix E.4

---

Trainiug loss
0.224
Training loss
0.24
0.224
Training less
0.224
0.223
0222
Trainiag
0.222
0.222
0.220
0.220
0.220
0.220
0.218
0.218
0.218
0.218
10:4
10a
10
10
109
10 10
10-1
100
10"
10:
10th
10th 10
Base learning
Base gradient clip
REPA weightt
warmup iteration
(a) Base learning rate.
(b) Base gradient clip.
(c) Base REPA loss weight. (d) Base Warm-up steps.
Figure 5: Results of base HP search on proxy MMDiT-uP tasks. We train O. 0.18B MMDiT-pP proxy models with 80 different base HPs settings. The optimal base HPs are transferred to the training of 18B target model.
Architecture of MMDiT--pP models. The target MMDDTT-pP-III model shares the same architec- ture as SMMDiT-18B. The proxy model also follows this architecture. differing only in width by reducing the number of attention heads. It contains 0.18B parameters (1% of the target model) with a width of 512. In the jP setup, the base width Tlbae is set to 1,920 (see Section 3.2. for definition of (.base).
Training of MMDIT-4p models. The training procedure for the target LIMDiT-pP-18D is identical that of the baseline MMDiT-18B models, except for the selected HPs. The proxy models are trained for 30K steps with a batch size of 4,096, We conduct 80 searches over four base HPs. Concretely, we uniformly sample base learning rate from 2.5 10 to 2.5 10-3 gradient clipping from 0.01 to 100, weight of REPA loss from 0.1 to 1, and warm-up up iteration from to 20K. To further verify that 30K iterations are enough for the proxy task, we conducted five proxy training runs (100K steps each) using the searched optimal HPss and different learning rates, as detailed in Appendix E.4
As derived in Appendix E.4 based Equation (T), the total tuning FLOPSs under μP amounts to 1 14.5% of one full pretraining cost, and thus only 3% of the human-tuned cost. We think the batch size and iterations used during the base HP search could be further reduced to lower the tuning cost. However, due to limited resources, we are unable to explore additional setups in this work.
Evaluation metries. We use training loss to select base HPs proxy tasks, as the dataset is passed through only once, and overfitting does not occur. We follow the standard practice in the uP literature [23] [14], where the base hyperparameter is typically selected by identifying the value that aligns with the lowest envelope of the training loss plot. To assess the final text-to- sott-to-imaee generation, in addition to the standard GenEval benchmark (21), also created internal benchmark comprising 150 prompts to comprehensively evaluate text-image alignment. Each prompt includes an average of five binary alignment checks (yes no), covering wide range of concepts such as nouns, verbs, adjectives (e.g., size, color, style), and relational terms (e.g.. spatial relationships, size comparisons). Ten human experts conducted the evaluation, with each prompt generating three images, resulting in a total of 22,500 alignment tests. The final score is computed as the average correctness across all test points. The details can be found in Appendix
5.2.2 Analyzing the Results of the Random Search
The visualization of the results of the base HP search is shown in Figure First, the base learning rate has the most significant impact on training loss. In our case, we observed that the envelope near 2.5 x 10- was stable and close to optimal. so we chose 2.5 x 10-4 as the optimal value. Interestingly. unlike the DiT and PixArt setups, the optimal learning rate here not close to the maximum stable value. This highlights a key difference between multi-epoch training in traditional deep learning and the single-epoch training in large model pretraining [65]. Intuitively, since the gradient signal for any individual sample is not revisited, the training must be more conservative to maintain stability. Second, the optimal gradient clipping value is 1, which deviates from the common practice of using a small value (e.g., 0.1) to stabilize pretraining. Intuitively, jaP favors a larger clipping value, as aggressive clipping can undermine the maximal update property central to jP. Third, the optimal weight for the REPA loss is determined to be around 0.5, consistent with the experience from existing works (66.. Finally, the warm-up iteration has a negligible impact on the training loss, so we adopt the default value of IK in the pretraining of the target model.
5.2.3 Transfer Results of MMDiT-p.P
The comparisons of training losses, GenEval results, and human evaluations are shown in Figure 6 Table and Ta respectively. In addition, the complete comparison of training losses and detailed

---

Table 4:
GenEval results of pretrained MMDIT-I8B and MMDIT-,uP-1SB models. MMDiT-ppP-8B
achieves better benchmark results with only 3% of manual tuning cost.
Method
Overall 1
Single
Two
Counting
Colors
Position Color Attribution
MMDiT-18B 0.8154
99.38
93.69
81.88
88.03
57.5
68.75
MMDDTT-pP-18B
0.8218
99.38
94.44
79.69
88.83
62.25
68.5
0.175
Table Results human evaluation for text- image alignment. The alignment accuracy (acc.) is computed as the average over 22,500 human alignment tests. MMDiT-pP-18BB achieves supe- rior results with only 3% of the manual tuning cost.
MMDiT 18B
Troining loss
0.174
MMDiT 16B
0.173
0.172
0.171
Method
Alignment acc. 1
50
100
150
IteatiionKK
MMDiT-18B
0.703
Figure MMDIT-paP-188 achieves consistently lower training loss than baseline after 15K steps.
MMDiT-pP-18B (Ours)
0.715
visualizations are provided in Figure 8 and Figure 10 in Appendix E.4
As result, MMDDT-ppP-180
outperforms the baseline MMDIT- all cases, achieving this with only 3% FLOPs of the standard manual tuning cost. These results demonstrate that μP is a reliable principle for scaling diffusion Transformers. As models grow larger and standard HP tuiing becomes prohibitively expensive, jP offers scientifically grounded frumework to unlock the full potential of large models.
6
Discussion
There are several promising research directions building on this work. First, the principles of pP could be extended to diffusion models with more advanced and efficient architectures, such as linear transforners [66] and cexture-oofexpert models [7]. Second, μP can applied to more sophisticated optimization algorithms, including the Muon optimizer [1] and the warmup-stan saapp-ttableeeeeeeee soppiiseee
learning
schedule 28). Third, while our experiments suggest that proxy model width of 256-512 and proxy dataset size of 1/10-1/6 of the full pretraining data are sufficient for stable HP transfer in diffusion Transformers, determining the optimal proxy task size that balances tuning cost and target model performance remains an important avenue for future work. Finally, developing learning- theoretic framework to explain the optimization dynamics [58 5I] generalization behavior [9], and downstream performance (79; 78 of diffusion Transformers under μP would be both meaningful and impactful. In summary, this provides the vision community with an initial principled approach for scaling diffusion Transformers efficiently.
Broader Impacts and Limitations. Our work has the potential accelerate progress in generative modeling applications using diffusion Transformers, including text-to-image and video generation. scaling diffusion Transformers could also facilitate the creation of deepfakes for disinformation. Regarding the limitations of this work, although we demonstrate the efficiency of μuP in large-scale applications, we do not identify the optimal proxy task size that balances HP tuning cost and target model performance, due to limited computational resources.
However, improvements
7 Conclusion
In this paper, we extend μP from standard Transformers to diffusion Transformers. By proving that mainstream diffusion Transformers share the same jPP form as vanilla Transformers, we enable direct application of existing μP practice and verify the reliable base HP transfer from small to large diffusion Transformers. This leads to practical performance gains on DiT-XL-2, Pit-rt- and MMDiT-188 while requiring small fraction of the usual tuning effort (e-g., 3% for MMDiT-18B), Our results establish μP principled and efficient scaling strategy for diffusion Transformers.
10

---

Acknowledgments
This work was supported by the National Natural Science Foundation of China (No. 924701 18);
the Beijing Natural Science Foundation (No. L247030); the Beijing Major Science and Technology Project under Contract No. 1221000004425000; the Beijing Nova Program (No. 20230484416); the ByteDance Seed Research Fund; the Public Computing Cloud of Remmin University of China; and the fund for building world-class universities (disciplines) of Remmin University of China. Chenyu Zheng was also supported by the Outstanding Innovative Talents Cultivation Funded Programs 2024 of Renmin University of China. Finally, the authors thank Enze Xie for his helpful discussion on the experimental setup of PixArt-ar
References
[1] Fan Bao, Shen Nie, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu, All are worth words: A vit backbone for diffusion models. In CVPR, pages 22669-22679, 2023.
[2] Fan Bao, Chendong Xiang, Gang Yue, Guande He, Hongzhou Zhu, Kaiwen Zheng, Min Zhao, Shilong Liu, Yaole Wang, and Jun Zhu. Vidu: a highly consistent, dynamic and skilled text-to-video generator with diffusion models. arXiv preprint arXiv:2405.0422 2024.
[3] James Betker, Gabriel Goh, Li Jing, Tim Brooks, Jianfeng Wang, Linjie Li, Long Ouyang. Juntang Zhuang, Joyce Lee, Yufei Guo, et al. Improving image generation with better captions. Computer Science. https://edn. openai. com/papers/tall-e-33 pdf. 2(3):8, 2023.
[4] Charlie Blake, Constantin Eichenberg, Josef Dean, Lukas Balles, Luke Prince, Björn Deis- eroth, Andrés Felipe Cruz-Salinas, Carlo Luschi, Samuel Weinbach, and Douglas Orr. u-/pp: The unit-scaled maximal update parametrization. CoRR, 01/22477.17455 2024.
[5] Blake Bordelon, Hamza Tahir Chaudhry, and Cengiz. Pehlevan. Infinite limits of multi-head transformer dynamics. In NewrIPS, 2024.
[6] Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo, Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, et al. Video generation models as world simulators. OpenAl Blog, 1:8, 2024.
[7] Junsong Chen, Chongjian Ge, Enze Xie, Yue Wu, Lewei Yao, Xiaozhe Ren, Zhongdao Wang, Ping Luo, Huchuan Lu, and Zhenguo Li. Pixart-E: Weak-to-strong training of diffusion transformer for 4k text-to-image generation. In ECCV, volume 15090, pages 74-91. 2024.
[8] Junsong Chen, Jincheng Yu, Chongjian Ge, Lewei Yao, Enze Xie, Zhongdao Wang, James T. Kwok, Ping Luo, Huchuan Lu, and Zhenguo Li. Pixart-o: Fast training of diffusion transformer for photorealistic text-to-image synthesis. In ICLR, 2024.
[9] Zixiang Chen, Yuan Cao, Quanquan Gu, and Tong Zhang. A generalized neural tangent kemel analysis for two-layer neural networks. Advances in Neural Information Processing Systems, 33:13363-13373, 2020.
[10] cloneofsimo and Team fal. Introducing auraflow v0.1: An open exploration of large rectified 2024.
flow models. httpp::/bbog. fal ai/
[11] Jeremy Cohen, Simran Kaur, Yuanzhi Li, J. Zico Kolter, and Ameet Talwalkar. Gradient descent on neural networks typically occurs at the edge of stability. In ICLR, 2021.
[12] Katherine Crowson, Stefan Andreas Baumann, Alex Birch, Tanishq Mathew Abraham, Daniel Kaplan, and Enrico Shippole. Scalable rish-resslution pixel-space image synthesis with hourglass diffusion transformers. ICML, 2024.
[13] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In CVPR, pages 248-255 2009.
[14] Nolan Dey, Gurpreet Gosal, Zhiming Chen, Hemant Khachane, William Marshall, Ribhu Pathria, Marvin Tom, and Joel Hestness. Cerebras-gpt: Open compute-optimal language models trained on the cerebras scaer cluster. CoRR, abs/2304.03200 2023.

---

[15] Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat gans on image synthesis. In NeurIPS. pages 8780-8794, 2021.
[16] Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, Dustin Podell, Tim Dockhor, Zion English, and Robin Rombach. Scaling rectified flow transformers for high-resolution image synthesis. In ICMI., 2024.
[17] Zhengcong Fei, Mingyuan Fan, Changqian Yu, Debang Li, and Junshi Huang. Scaling diffusion transformers to 16 billion parameters. CoRR, abs/2407. 1163,, 2024.
[18] Peng Gao, Le Zhuo, Dongyang Liu, Ruoyi Du, Xu Luo, Longtian Qiu, Yuhang Zhang, Chen Lin, Rongjie Huang, Shijie Geng, et al. Lumina-12x: Transforming text into any modality, resolution, and duration via flow-based large diffusion transformers. arXiv preprint arXiv:2405.05945, 2024.
[19] Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Masked diffusion transformer a strong image synthesizer. In ICCV, pages 23107-231 16, 2023.
[20] Yu Gao. Lixue Gong, Qiushan Guo, Xiaoxia Hou, Zhichao Lai, Fanshi Li, Liang Li, Xi- aochen Lian, Chao Liao, Liyang Liu, al. Seedream 3.0 technical report. arXiv preprint ar5iv:2504.1 11346, 2025.
[21] Dhruba Ghosh, Hannanch Hajishirzi, and Ludwig Schmidt. Geneval: An object-focused framework for evaluating text-to-image aligmment. In Neur!PS, 2023.
[22] Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The Ilama
3 herd of models. arXiv preprint 24XX::2447.217833 2024.
[23] Ishaan Guirajani and Tatsunori B. Hashimoto. Likelihood iikelihood-based diffusion language models. In NeurlPS, 2023.
[24] Moritz Haas, Jin Xu, Volkan Cevher, and Leena Chennuru Vankadara. up² Effective sharpness aware minimization requires layerwise perturbation scaling, In NeurIPS, 2024.
[25] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In ICCV, pages 1026-1034, 2015.
[26] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. In Advances in Neural Information Processing Systems, pages 6626-6637, 2017.
[27] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Comput., 9(8): 1735-1780, 1997.
[28] Shengding Hu, Yuge Tu, Xu Han, Chaoqun Hc, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, Xinrong Zhang. Zhen Leng Thai, Kai Zhang, Chongyi Wang, Yuan Yao, Chenyang Zhao, Jie Zhou, Jie Cai, Zhongwu Zhai, Ning Ding, Chao Jia, Guoyang Zeng, Dahai Li, Zhiyuan Liu, and Maosong Sun. Minicpm: Unveiling the potential of small language models with scalable training strategies. CoRR, abs/2404.06395, 2024.
[29] Zhongzhan Huang, Pan Zhou, Shuicheng Yan, and Liang Lin. Scalelong: Towards more stable training of diffusion model via scaling network long skip connection. In NewrIPS, 2023.
[30] Jiarui Jiang, Wei Huang, Miao Zhang, Taiji Suzuki, and Liqiang Nie. Unveil benign overfitting for transformer in vision: Training dynamics, convergence, and generalization. Advances in Neural Information Processing Systems, 37:1 35444-13666222 2024.
[31] Keller Jordan, Yuchen Jin, Vlado Boza, You Jiacheng, Franz Cecista, Laker Newhouse, and Jeremy Bernstein. Muon: An optimizer for hidden layers in neural networks, 2024. URL. hutps://kellerjordaan. github. io/posts/muon, 6, 2024.
12

---

[32] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors, ICLR, 2015.
[33] Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloé Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C. Berg. Wan-Yen Yen Lo, Piotr Dollár, and Ross B. Girshick. Segment anything. In ICCV, pages 3992-4003 2023.
[34] Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. In Advances in Neural Information Processing Systems, pages 3929-3938, 2019.
wwtp:/accooooaaoooee [35] Black Forest Labs. Flux. http:////ggggio...oo///aaa--oooa forestroesstte labs maassfflu ffux 2024.
[36] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proc. IEFE, 86(11):2278 2324, 1998.
[37] Jaehoon Lee, Samuel S. Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: empirical study. In NeurIPS, 2020.
[38] Aitor Lewkowycz, Yasaman Bahri, Ethan Dyer, Jascha Sohl-Dickstein, and Guy Gur-Ari. The large learning rate phase of deep learning: the catapult mechanism. CoRR, abs/2003.02216 2020.
[39] Daiqing Li, Aleks Kamko, Ehsan Akhgari, Ali Sabet, Linmiao Xu, and Suhail Doshi. Playground v2.5: Three insights towards enhancing aesthetic quality rext-to-image generation. CoRR, abs/2402. 17245, 2024.
[40] Hao Li, Shamit Lal, Zhiheng Li, Yusheng Xie, Ying Wang, Yang Zou, Orchid Majumder, R. Manmatha, Zhuowen Tu, Stefano Ermon, Stefano Soatto, and Ashwin Swaminathan. Efficient scaling of diffusion transformers for text-to-image generation. CoRR, abs/2412.12391,, 2024.
[41] Hao Li, Yang Zou, Ying Wang, Orchid Majumder, Yusheng Xie, R. Manmatha, Ashwin Swaminathan, Zhuowen Tu, Stefano Ermon, and Stefano Soatto. On the scalability of diffusion- based texl-to-image generation. CVPR, pages 9400-9409, 2024.
[42] Xiang Li, Yiqun Yao, Xin Jiang, Xuezhi Fang, Xuying Meng, Siqi Fan, Peng Han, Jing Li, Li Du, Bowen Qin, Zheng Zhang, Aixin Sun, and Yequan Wang. FLMM-001B B: an open LLM and how to train it with $100k budget. CoRR, 0as/2309.03852. 2023.
[43] Xiang Li, Yiqun Yao, Xin Jiang, Xuezhi Fang, Chao Wang, Xinzhang Liu, Zihan Wang. Yu Zhao, Xin Wang, Yuyao Huang, Shuangyong Song, Yongxiang Li. Zheng Zhang, Bo Zhao, Aixin Sun, Yequan Wang, Zhongjiang He, Zhongyuan Wang. Xuelong Li, and Tiejun Huang. Tele-flm technical report. CoRR, abs/2404. 166454 2024.
[44] Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. In NeurlPS, pages 11669-11680, 2019.
[45] Tsung- Yi Lin, Michael Maire, Serge J. Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and Lawrence Zitnick. Microsoft COCO: common objects in context. In ECCV, volume 8693, pages 740-755, 2014.
[46] Yaron Lipman, Ricky T. Q Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In ICLR, 2023.
[47] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. ICLR. OpenRe- view.net, 2019.
[48] Zeyu Lu, Zidong Wang. Di Huang, Chengyue Wu, Xihui Liu, Wanli Ouyang, and Lei Bai. Fit: Flexible vision transformer for diffusion model. In ICML, 2024.
[49] Nanye Ma, Mark Goldstein, Michael S. Albergo, Nicholas M. Boffi, Eric Handen-Fijnded and Saining Xie. Sit: Exploring flow and diffusion- generative models with scalable interpolant transformers. In ECCV, volume 15135, pages 23-40, 2024.
13

---

[50] Charlie Nash, Jacob Menick, Sander Dieleman, and Peter W. Battaglia. Generating images with sparse representations. In ICML, volume 139, pages 7958-7968. PMLR, 2021.
[51] Atsushi Nitanda, Denny Wu, and Taiji Suzuki. Convex analysis of the mean field langevin dynamics. In AISTATS, volume 151, pages 9741-9757 2022.
[52] William Peebles and Saining Xie. Scalable diffusion models with transformers. In ICCV, pages 4172-4182. IEEE, 2023.
[53] Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhor, Jonas Müller, Joe Penna, and Robin Rombach. Sdxl: Improving latent diffusion models for high-resolution image synthesis. arXiv preprint arXiv:2307.01952, 2023.
[54] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In MICCAI, volume 9351, pages 234 244-241, 2015.
[55] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang. Emily L. Denton, Scyed Kamyar Seyed Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, Jonathan Ho, David J. Fleet, and Mohammad Norouzi. Photorealistic text-to-image diffusion models with deep language understanding. In NeurlPS, 2022.
[56] Tim Salimans, Ian J. Goodfellow, Wojeiech Zaremba, Vicki Cheung. Alec Radford, and Xi Chen. Improved techniques for training gans. In Advances in Neural Information Processing Systems, pages 1226-2234 2016.
[57] Kyle Sargent, Kyle Hsu, Justin Johnson, Li Fei-Fei, and Jiajun Wu. Flow to the mode: Mode- seeking diffusion autoencoders for state-of-the-arr image tokenization. CoRR, abs/2503. 11056, 2025.
[58] Mei Song, Andrea Montanari, and P Nguyen. A mean field view of the landscape of two-layers neural networks. Proceedings of the National Academy of Sciences, 115633):77665-7777,, 2018.
[59] Yuchuan Tian, Zhijun Tu, Hanting Chen, Jie Hu, Chao Xu, and Yunhe Wang. U-dits: Down- sample tokens in u-shaped diffusion transformers. In NeurlPS, 2024.
[60] Leena Chennuru Vankadara, Jin Xu, Moritz Haas, and Volkan Cevher. On feature learning in structured state space models. In NeurlPS, 2024.
[61] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, In Advances in Neural Informa- tion Processing Systems, pages 5998-6008, 2017.
[62] Ang Wang, Baole Ai, Bin Wen, Chaojie Mao, Chen- Wei Xie, Di Chen, Feiwu Yu, Haiming Zhao, Jianxiao Yang, Jianyuan Zeng, et al. Wan: Open and advanced large-scale video generative models. arXiv preprint arXiv:2503.2031 20314, 2025.
[63] Rongzhen Wang, Yan Zhang, Chenyu Zheng, Chongxuan Li, and Guoqiang Wu, A theory for conditional generative modeling on multiple data sources. arXiv preprint 1rX1v:2502.14583 2025.
[64] Mitchell Wortsman, Peter Liu, Lechao Xiao, Katie E. Everett, Alexander A. Alemi, Ben Adlam, John D. Co-Reyes, Izzeddin Gur, Abhishek Kumar, Roman Novak, Jeffrey Pennington, Jascha Sohl-Dickstein, Kelvin Xu, Jaehoon Lee, Justin Gilmer, and Simon Korbblith. Small- scale proxies for large-scale transformer training instabilities. In ICLR, 2024.
[65] Lechao Xiao. Rethinking conventional wisdom in machine learning: From generalization to scaling. CoRR, abs/2409. 2024.
[66] Enze Xic, Junsong Chen, Junyu Chen, Han Cai, Haotian Tang, Yujun Lin, Zhekai Zhang, Muyang Li, Ligeng Zhu, Yao Lu, and Song Han. SANA: efficient high-resolution image synthesis with linear diffusion transformers. CoRR, abs/2410. 10629, 2024.
(67] Greg Yang. Wide feedforward or recurrent neural networks of any architecture are gaussian processes. In NewrlPS, pages 9947-9960. 2019.
14

---

[68] Greg Yang. Tensor programs II: neural tangent kernel for any architecture. abs/2006. 14548, 2020.
CoRR,
[69] Greg Yang. Tensor programs III: neural matrix laws. CoRR, abs/2009. 2020.
[70] Greg Yang and Edward J. Hu. Tensor programs IV: feature learning in infinite-wine neural networks. In Marina Meila and Tong Zhang, editors, ICML, volume 139, pages 11727-11737. PMLR, 2021.
[71] Greg Yang and Etai Littwin. Tensor programs iib: Architectural universality of neural tangent kernel training dynamics. ICML, volume 139, pages 11762-11772, 2021.
[72] Greg Yang and Etai Littwin. Tensor programs ivb: Adaptive optimization in the anfinitewwidth limit. CoRR, abs/2308.018 001144 2023.
[73] Greg Yang, Edward J. Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs V: tuning large neural networks via zero-shot hyperparameter transfer. CoRR, abs/2203.03466 2022.
[74] Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Tensor programs VI: feature learning in infinite depth neural networks. In ICLR, 2024.
[75] Yuanyang Yin, Yaqi Zhao, Mingwu Zheng, Ke Lin, Jiarong Ou, Rui Chen, Victor Shea-Jay Huang, Jiahao Wang, Xin Tao, Pengfei Wan, Di Zhang, Baoqun Yin, Wentao Zhang, and Kun Gai. Towards precise scaling laws for video diffusion transformers. CoRR, abs/2411.17470, 2024.
[76] Sihyun Yu, Sangkyung Kwak, Huiwon Jang, Jongheon Jeong, Jonathan Huang, Jinwoo Shin, and Saining Xie. Representation alignment for generation: Training diffusion transformers is easier than you think. CoRR, abs/2410.06940, 2024.
[77] Min Zhao, Guande He, Yixiao Chen, Hongzhou Zhu, Chongxuan Li, and Jun Zhu, Riflex: A free lunch for length extrapolation in video diffusion transformers. arXiv preprint arXiv:2502.15894 2025.
[78] Chenyu Zheng, Guoqiang Wu, Fan Bao, Yue Cao, Chongxuan Li, and Jun Zhu, Revisiting discriminative generative classifiers: Theory and implications. In International conference on machine learning, pages 42420-42477. PMLR, 2023.
[79] Chenyu Zheng, Guoqiang Wu, and Chongxuan Li. Toward understanding generative data augmentation. Advances neural information processing systems, 36:54046-5400 2023.
[80] Chenyu Zheng, Wei Huang, Rongzhen Wang, Guoqiang Wu, Jun Zhu, and Chongxuan Li. On mesa-optimization in autoregressively trained transformers: Emergence and capability. Advances Neural Information Processing Systems, 37:49081-49123 2024.
[81] Zangwei Zheng, Xiangyu Peng. Tianji Yang, Chenhui Shen, Shenggui Li, Hongxin Liu, Yukun Zhou, Tianyi Li, and Yang You. Open-sora: Democratizing efficient video production for all. arXiv preprint ar arXiv:24/2.20400
2024.
[82] Yutao Zhu, Kun Zhou, Kelong Mao, Wentong Chen, Yiding Sun, Zhipeng Chen, Qian Cao, Yihan Wu, Yushuo Chen, Feng Wang, et al. Yulan: An open-source large language model. arXiv preprint arXiv:2406. 19853, 2024.
15

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our abstract and introduction accurately reflect the theoretical contribution in Section 3.1 and experimental contribution in Section and
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
See details in Section 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't It acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
16

---

Justification: Assumption used to simplify the proof can be found in Assumption C.1 and the completed proof can be found in Appendix C
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We introduce the necessary configurations in Section 4 Section and Ap- pendix E We also open-source our code at httts:///piiub..cmmm/Ml.//A///// Diffussion
Transformers-muP
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
17

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We also open-source our code https//gituu..cm/MLGGSAAAAA
Diffusion- Diffusion-Transfornae Transformers--muu for DiT and PixArt-a experiments. However, the data for the MMDiT experiments is not available now because it is an internal dataset. We will try to release it in the future.
Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurlPS code and data submission guidelines (https://nips.co
ooteessouubiissioooooooor for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to to reproduce the results. See the NeurIPS code and data submission guidelines (https leaiic/gguieee gui /ooeee//ooddebmiiiccc hwtrbbbmssioooooooioiooo for more details.
/
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: YYs]
Justification: We introduce the necessary configurations in Section 4 Section and Ap- pendix E We also open-source code at http:://ggtab..commMLLSAAAAAA Transformers for DiT and PixArt-a experiments.
Diffusion
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Error bars reported because it would be too computationally expensive. expensive
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
18

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for caleulating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yss]
Justification:
See details in Appendix E
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: http://neeippppe http://rriisscccpllllleeeeeeeeee cccppublic/tthicccuuiieeell.eee.
Answer: [Yes]
Justification:
We confirm the research conducted in the paper satisfies the NeurIPS Code of
Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
See details in Section
Guidelines:
The answer NA means that there is no societal impact of the work performed.
19

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
See details in Table 7 Appendix
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
20

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
batasets.papersstt
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. If this information is available online, the authors encouraged to reach out the asset's s creators.
13. New assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [No]
Justification: The internal dataset and 18B models are not available to the public now. We will try to release them in the future.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [Yes]
Justification:
See details human evaluation in Appendix E.4.5
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Our human evaluation process is safe, and the generated images are legal.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
21

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
We only use the LLMs to polish our writing.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: //tppr////eurpp..... / for what should or should be described.
22

---

Contents of Appendix
Additional Related Work
24
A.1 Scaling Diffusion Transformers
24
A.2 Applications of /uP in AIGC
24
Theoretical Background of uP
24
B. abe-Parameterization Setup
24
B.2 NE8OR Program
24
B.3 pP Any Representable Archtecture
25
B.4 Extensions to Other HPs
26
Proof of Theorrm
26
Proof of DIT.
26
Proof of PixAn--
29
C,3
Proof of U-VIT
30
Proof of MMDiT
30
D Additional Details for Section
31
Additional Experimental Details and Results
32
Ássets and Licenses)
32
Additional Details of DIT Experiments
33
E.2.1
Evaluation Implementation
................................
33
B.2.2
Additional Results of Base transferability
................................
33
E.2.3
Additional Results of DiT XL-2 Pretraining
38
F.2.4
Computational Cost
38
Additional Details of PixArt-a Experiments
39
E.3.1
FLOPs ratio of HP Search on Proxy Models
39
E.3.2
Additional Results of Base Learning Rate Search
................................
39
E.3.3
Additional Results of PixArt-a Pretraining
39
Additional Details MMDIT Experiments
40
E.4.1
HPs Tuned by Human Experts)
40
E.4.2
Additional Details of Base HP Search
40
E.4.3
FLOPs ratio of HP Search on Proxy Models
40
E.4.4
Additional Results of Training Loss Comparison
40
E.4.5
Additional Details of Human Evaluation
40
E.4.6
Additional Results of Visualization Comparison
23

---

Appendix A
Additional Related Work
A.1 Scaling Diffusion Transformers
HourglassDiT 12 introduces a hierarchical architecture with downsampling and upsampling in DIT, reducing computational complexity for high-resolution image generation. SD3 [I6] scales MMDiT to 8B parameters by increasing the backbone depth. Large-DiT [18] incorporates LLaMA's text embeddings and scales the DiT backbone, demonstrating that scaling to 7B parameters improves convergence speed. DiT-MoE [17] further scales diffusion Transformers to 16.5B parameters using a mixture-of-experts architecture. Li et al. [0] empirically investigate the scaling behavior of various diffusion Transformers, finding that U-ViT [1 scales more efficiently than cross-attention based DiT variants. Yin al. (75) systematically analyze scaling laws for video diffusion Transformers, enabling the prediction of optimal hyperparameters (HPs) for any model size and compute budget. We note that all these works adopt standard parameterization (SP), which prevents the transfer of optimal HPs across different model sizes, and thus suffer from heavy tuning costs at scale.
A.2 Applications of jP in AIGC
Recently, yP has been successfully applied to the pretraining of large language models (LLMs) 73
[4 [2] 43 23 42 82 reducing HP tuning costs and stabilizing training. For diffusion Transformers,
in addition to AuraFlow v0.1 [T], some works have also employed μ.P. For example, Gulrajani and Hashimoto
2 searches for the base learning rate under μP and transfers it to a 1B 1B-parameter diffusion language model, while Sargent al. (57 uses pP to tune the HPs of Transformer-basee diffusion autoencoder, achieving rates. However, these studies differ fundamentally from ours. First, while they assume the HP transferability property of original μP (even though it may not hold in their setups). our work provides rigorous theoretical guarantees for the yP of diffusion Transformers and systematccally verifies its HP transferability. Second, we validate the scalability of μP at up to 18B parameters, significantly larger than the models used in prior work.
art image tokenization across multiple compression
Appendix B Theoretical Background of μP
In this section, we review the theoretical foundations of μP. For simplicity, we mainly focus on the Adam/AdamW optimizer widely used in practice. Besides, the in the optimizer is assumed to be zero, which is very small in practice (c.g.. 10 10-5).
B.1 abc- Paramettritiiiiinnnn Setup
For completeness, we restate the ade-parameterizztion from Section 2 which underpins typical μP theory. The praatttrtettinntionnnn specifies the multiplier, initial variance, and learning rate for each weight 770 73.. Specifically, let n denote the network width, we parameterize each weight as
W = own W, where the trainable component W is initialized as Wij ~ (O on The learning rate for W is parameterized n The width-independent quantities ow, ow. and n/ are referred to as base HPs and can be transferred across different widths.
B.2 NETOR Program
which is a theoretical series mainly developed to express and analyze the infinite-width limits of neural networks both at initialization and during training. The latest and most general version among them the NEgOR T Program (72). which is formed by inductively generating scalars and vectors, starting from initial matriies//vectorsssellass (typically the initial weights of the neural network) and applying defined operations, including vector averaging (Avg), matrix multiplication (MatMul), and nonlinear outer product transformations (OuterNonlin). We introduce these operators, where n the model width tending to be infinite.
μP is originally proposed by Tensor Programs 67
72
24

---

Table jaP for (difflusion) Transformers with differences between JP and standard parameterization (SP) in practice (e.g.. Kaiming initialization 25). and gray text to indicate the SP settings. We do not contain the scalar weights in the main text because they do not exist in DIT.
optimizer. We purple to highlight the
Input weights
Hidden weights
Output weights
Scalar weights
aw
0
0
1 (0)
0
bw
0
1/2 (0)
0 (1/2)
0
CW 0
0
Avg We can choose a existing vector a E Rn and append to the program a scalar
Za E I.
MatMul We can choose a matrix W E R and vector a E IR" existing in the program, and append to the program vector
WoER"
W R".
OuterNonlin
For any integer k. I O, we can aggregate it existing vectors and I existing scalars to X E Raak and E R respectively. With an integer 2 0 and a pseudo-Lipschitz function vb p)(r++)++ + R (e.g., SiLU, GeLU), we append to the program vector
y R"
ya
Σ 91.09
+ X Xiss... xB,s)c),
n
where the r + is called the order of vb.
While the above operators do not directly instruct to transform scalars into scalars, previous work 772 has proved that this can be done by combining the above operators. This can be formed as follows. Lemma B.1 Secalars-to-scalle 2.6.2 in 72 If TRl + R is a ggeudo-Lipschine function and e are the collected scalars NESOR T program, then (b(c) E R can be introduced as a new scalar in the program.
transformation is representable by the NESORT Program, Lemma
B.3 μPP Any Representable Archtecture
The seminal work of Yang and Littwin 72] systematically studies the dynamics of common neural architectures trained with adaptive optimizers (e-g.. Adam [32]. AdamW 47)), where the forward pass of neural network can be expressed as NEgOR T program (e.g.. MLPs, CNNs (36). RNNs 27, standard Transformers (6]). In this setting, Yang and Littwin [72 prove the existence of unique, optimal scaling Maee-Malaaal Update Parametrization evolve independently of width and at maximal strength. Otherwise, feature evolution in some layers will either explode vanish in the infinite-width limit, rendering large-scale pretraining ineffective.
which features in each layer
Concretely, μP is implemented by analytically adjusting HPs with model width, typically involving the weight multiplier, initialization variance, and learning rate considered in aib-paarameterization The jP formulation of any architecture whose forward pass can be expressed as NE@OR program follows the same scaling rules presented in Table o (equivalent to Table [ in the main text). Specifically, adjusts the ade-parameteriiztton of each weight according to its type, where the weights are categorized into four types: input weights, hidden weights, output weights, and scalar weights. The definitions are as follows: input weights satisfy dn = O(1) and d O(n); hidden weights satisfy d (n) and dout (nn); output weights satisfy dn O(n) and dout O(1), ); and scalar weights satisfy dn = ((1) and dout = (11) where n, din dout denotes the model width, fan-in dimension and fan-out dimension, respectively. The μP result is formally summarized in the following lemma.
Lemma B.2 (nP of representable architecture, Definition 2.9.12 in 72). If the forward pass of an architecture can be represented by NE@ORT program, its μP formulation follows the scaling rules in Table ( equivalent Table 7 in the main text).
25

---

B.4 Extensions to Other HPs
The μP principle can be extended other HPs. such as common optimization-relant HPs (e.g.. warm-up iteration). We refer the reader to [73] 72 for comprehensive discussion. We mainly focus on the HPss that occur in the main text.
Gradient clip: The clip value should be held constant with respect to width.
Warm-up iteration: The warm-up iteration should be held constant with respect to width.
Weights of different losses (c.g.. REPA loss, noise schedule): The weights of different losses should be held constant with respect to width.
Weight decay: For coupled weight decay (e.g., Adam 332 the weight decay value can not transferred. For decoupled weight decay (e.g., AdamW 47), the weight decay value should be held constant with respect to width.
Appendix C Proof of Theorem 3.1
Elementary notations. We lightface (e-g., a, A), lowercase boldface (e.g., a), and uppercase boldface letters (e.g., A) to denote scalars, vectors, and matrices, respectively. For vector a, we denote its i-th element as a, For a matrix A, we use Ag:. A. and Aij to denote its k-th row, k-th column and (i,j) j)-th element, respectively. We define [n] = nu.
For the reader's convenience, restate Theorem
follows.
Theorem C.1. The forward passes of mainstream diffusion Transformers (U-VIT M1 DiT (52), Pixart-a (8), and MMDiT (I6)) can be represented within the NEλOR T Program. Therefore, their μPP matches the standard μP presented Table (equivalent to Table 7 in the main text).
Following the standard practice in the literature 70 the following assumption in all of our proofs.
72
24 to simplify the derivation, we adopt 74,
Assumption C.1. We assume that constant dimensions (fxed during scaling the width of neural networks) are which includes data dimension, label dimension, patch size, frequency embedding size of time embedding. dimension of attention head, and so on. We also assume the width of different layers is the same n and do not consider bias parameters here.
In principle, the proof technique naturally extends to any finite fixed dimension (> 1) and to hidden layers with different widths (non-uniform infinite-width setting) 72. 24]
We refer the readers to Section 2.9 in Yang and Littwin (72), which provides a detailed discussion of both these generalizations and shows how the tensor program framework naturally accommodates them.
C.1 Proof of DiT
Proof. By Lemma B.2 we can prove the theorem by proving the forward pass of DiT can be represented by the NE@OR Program.
The DiT architecture includes embedder for the input latent R. an embedder for the diffusion timestep t e IR, an embedder for the label y E R. Transformer blocks with adaLN, and a final layer with adaLN. In the following, use the NESORT Program represent the forward computation of these modules in sequence.
Embedder for the input latent The embedder for the x is one-layer CNN, which is known be representable by the NESORT T Program (e.g., see Program 7 in [67] for general CNNs with pooling). We write the program for clarity here. In our case, the one-layer CNN's parameters can be denoted by w
CNN E R and the operation can implemented by
embed := (w
CNN. CNN. ;x), (OuterNonlin)
where (w CNN, x) = wo CNN
x.
Positional embedding. Given a number of patches (1 in our simplified case), the positional embed- ding is initialized with the sine-cosine version and fixed during the training. We use the Pos to denote the positional embedding, and adding it to the x embed can be written as
embed
:= (/x clbeed cibeed Pos Ja:), (OuterNonlin)
26

---

where wi( ti([x x
zembed + 2h Here, we reuse the notation of ambed of the final embedding for the imput latent a without confusion.
a:)
Zambed
Embedder for the diffusion timestep f. The embedder for the diffusion timestep is a frequency embedding (one dimension here) followed two-layer MLP. First, for the frequency embedding, we can represent it as
freq
((). (Lemma B.1
Second, the MLP is known to be representable by the NE@ORT Program (e.g., Program 67 We write it here for completeness. We use w(1) E R" and W R=x" to denote the weights in the first layer and the second layer of the MLP, respectively. We can derive
e 4h (w ((ee),
NonterNonlin
( = 2z(t t tembed W
(OuterNonlin)
(MatMul)
where 1 (wa))
(1), ffre)) wa! ffre and (h2 S)
Embedder for the label Denoting the weights of embedding as wmbed E R", then its computa- tion can be written directly as
ya embed
= (w y), (OuterNonlin)
where (wmabeed :y) embeddy y.
Merging the embeddings of and y. In DiT, the embeddings of and y as the condition for the following calculation, which can be written as
== ([[tebeed yembedJa:)) (OuterNonlin)
where ([tombed cmmbed] ]a:) gembed yo cmmeed
Transtormer block with adaLN. First, the adaLN block maps the condition c to the gate, shift, and scale values (0attn Buttn yattu Omip, 3ulpp, for attention layer and MLP layer in the Transformer block. We denote the parameters of the adaLN block by (WD-atin W 8-aastn W W W WT-mlp)) each of them is in Rax=, Then
the forward pass of adaLN block can be written as
C (ca),
(OuterNonlin)
Øattn W 10--tta
c,
(MatMul)
gattn W atttn
C
(MatMul)
7 attn = W -
(MatMul)
mmlp 9 9-mipg č
(MatMul)
gmlp W B WA-mipg C
(MatMul)
mlp W W
(MatMul)
where db is the SiLU activation.
Now, we denote the input feature of the Transformer block as x E R". A layer norm without learnable parameters first normalizes it, which is also known to be representable by the NESORT Program (e.g., Program 9 in 677). We rewrite it in our simplified case.
E E|x]
1
Σ
Fo.
(Avg)
n
2.2 a
= wi (
E(x)):
(OuterNonlin)
V/x] 1
(Avg)
n
norm = (za [B(z], V[øl)),
(OuterNonlin)
27

---

where 1 (#a; B[z)) (s. E[x:)?, 12 [[a,], V[=l]) (z, - H|cl\///[
After that, the normalized feature interacts with the condition information via the output adaLN (Battn yaten, which can be represented follows.
Tnorm ((x (orm Buttn yatt" Ja:), (OuterNonlin)
where ([a Batta yattu lo:) Ta atta
za + gattn.
Now, the processed feature znorm is put to the multi-heate attention layer, which is known to be representable by the NEλORT Program (e.g., see Program 10 in [6]]) We write it for completeness here. We denote W Wº, W [R"X" the attention layer's s key, query, and value matrix, each row representing head with dimension 1. The forward pass be written follows.
k = W k x morm
(MatMul)
9 Wº Q muorm
(MatMul)
V WV muorm
(MatMul)
aa attu == (ik, q, v)o:),
(OuterNonlin)
where ((ke q. v).:)) = la
Before discussing the MLP layer, still have residual connection and interaction from the gate value gattn which can be represented by
ha == (/x gattn x 2)).))
(OuterNonlin)
where ([x gattu x]l:) Za gattm.ratan
The remaining part is very similar the above, which includes the layer norm, interaction with the outputs of adaLN, the MLP layer, and the residual connection. We write them in sequence for clarity.
For the layer norm, write
Ehh] == 1
har
(Avg)
n
622 a
= en (hai E(h)),
(OuterNonlin)
V(h) n
(Avg)
hnorm := 1/2 ( [E(h], V(hl)),
(OuterNonlin)
where uh (ha; B(hh) (ha- E[h))), 422 (has [E[h), V(hl)) = (ha - (((t)////VIII +c. For the interaction with scale and shift value from adaLN, we have
Nnorm := < norm ,3mln, , m))))):),
(OuterNonlin)
where ((h Gulp mlp == 7a hilpp/noo 3mlp,
For the two-layer MIP in the Transformer block, we use ( ( e Raxn to denote the weights in its first layer and the second layer, respectively. We have
h0 (1]
(MatMul)
T(()) = (11))),
(OuterNonlin)
h mlp == (1)
(MatMul)
where the db is the GeLU activation function.
For the final residual connection, it can be represented by
ax
Thlock
(1h mlp mllp h)).)),
(OuterNonlin)
28

---

where ( @mip, hlaa)) = ha + 0mlphplp
Final layer with adaLN. The final layer includes the adaL.N, layer norm, and linear projection. We discuss them in sequence.
The adaLN block has parameters (W---fnal final W W and receive the condition c. Its forward pass can be written as follows.
E ((a)): g W W-B-iinal
(OuterNonlin)
C;
(MatMul)
Sinal W
C,
(MatMul)
where &a is the SiLU activation.
The layer norm layer normalizes the output of the transformer block Tblock, which the same the above layer norm. We omit the derivation here for simplicity. We denote the output of the layer norm by zuor R and parameter of the linear projection by final then the remains calculation can be derived as
&a
znorm = 4/1 ([z norm Gffnal 7 la:),
finaly
(OuterNonlin)
zfinal = V2(w final gmorm la:),
(OuterNonlin)
final
2a
ffnal
(Avg)
where 1 ([æ G
7 final] la:) == %a 2a
final ,orm + fffinal and (22(w @uorm] Ja:) = wa unninalgooom
Therefore, forward pass from the inputs (x, y,, 1) to output ,inaal can be represented by the NE@ORT Program, which means that The pP of DiT is the same as the standard μP in Table 6 The proof completed.
C.2 Proof of PixArt-a
Proof. Since most parts of the PixArt-a the same those in DIT, only discuss the different modules from DiT for simplicity. These modules are mainly induced from the condition process.
Embedder for the diffusion timestep I. The embedder for the diffusion timestep / not only out- puts gembed but also gate, shift, and scale values based on gembed (named adal.N-single in (8), which are shared by the following Transformer block. We denote the additional parameters by (Wtttattn 8 tattattn 0-tlp, S W2-tmlp), each of them is in pR Then
its forward pass can be written as
To == ((aabbed)
(OuterNonlin)
glatta W d-ttin
(MatMul)
cattn 8-B-atca f
(MatMul)
tattn W Wtttattn t
(MatMul)
ttmpp W tmipp
(MatMul)
gaupp
W 8 -
(MatMul)
tmlp = W mmpp
(MatMul)
where vb is the SiLU activation.
Transformer block with gate-shift-scale table. in DiT block with the gate-shift-scale table, @xattn 3 grastn xattt @xmlp gemip , xmmpp E IR". The gate-shift-scale table interacts with the shared embeddings of in the following way.
PixArt-a replaces the adaLN block which parameters are denoted by
gatta := v0 @lattn) )
(OuterNonlin) (OuterNonlin)
gattn (l gtattm]

---

aam = (hyxatin y la:);
(OuterNonlin)
gmlp eb (1oxmlp, @umlp) la:);
(OuterNonlin)
gmip ([Bxmlp, Buulp] la:).
(OuterNonlin)
mip (byxml,, و y
(OuterNonlin)
where v' attention layer and MLP layer, just like DIT.
The above outputted gate, shift, and scale values are used in the following
One more difference between the Transformer block of PixArt-a and DiT is the existence of the additional cross- cross-attention layer, which integrates the o of the self-attention layer and the text embedding y ebeed We denote its parameters by W W WeV W E Ruxn It can be written as follows.
k. Wok y ombedd
(MatMul)
q 8o h.
(MatMul)
V = W embbdd
(MatMul)
ha v (1k, q. v| u)a:),
(OuterNonlin)
h croes W h"
(MatMul)
where (ik, q, vl:)) 2a The following pre-LN MIP with gate-shift-scale table is a trivial extension of the pre-LN MLP layer in the DiT block, like what we did above, so we omit them here.
Final layer with shift-scale table. The only difference is that we replace the adaLN block with shift-scale table, just like what did for the attention layer, so we omit the derivation here.
C.3 Proof of U-VIT
Proof. We only discuss the modules that do not occur in DiT and PixArt-a. We find that the only new operator in U-VIT is the long skip connection, which concatenates the current main branch h" th I and the long skip branch h² E R" and then performs linear projection to reduce the dimension. We denote the parameters of the linear layer by W [W"m Wt, where Wm W E Rax. then the long skip connection can be represented by
mm ""h"
(MatMul)
h W"h"
(MatMul)
h, = vb (h h. m.:)
(OuterNonlin)
where "hh" h2 (.:) = Tm + Teg Because other operators in U-ViT (CNN embeddings/decoder, layer norm, attention layer, MI.P layer, residual connection) have been studied in DiT, we finish the proof.
C.4 Proof of MMDiT
Proof. For simplicity, we only discuss the modules that do not occur in the above models. The new module that emerges in the MMDiT is the joint attention with the QK-Norm for the latent feature x and text condition c. For the latent feature, we denote the key/query/valuee parameters by k Thr WVz Buxn the parameters of QK-norm by y Kx 2 E RR.. and the parameters of linear projection by w E Raxn We can define the similar parameters for the text condition c WAF W We Ko g0 Woe The forward pass of the joint attention layer be written as follows.
ka Whz
(MatMul)
q² = W E,
(MatMul)
v" = WV= x,
(MatMul)
he = CC
(MatMul)
30

---

q = c
(MatMul)
2 = WVcc,
(MatMul)
ka A.morm vi(( . ka] la:),
(OuterNonlin)
anorm + (1q .7 Q la:);
(OuterNonlin)
kenorm = 41 (1ke, y Ja:):
(OuterNonlin)
Lo cuorm = vi(l"" 909 la:)
(OuterNonlin)
t = (>2 ( onorm k = v² k cuorm v (a))
(OuterNonlin)
O = 2/2 (1a cuorm h amoor v², k cnorm v0' Ja:))
(OuterNonlin)
v rout Wo0y
(MatMul)
cout W
(MatMul)
where (1, b) = ab/ al and (2 b, ed abc + ade. The proof is completed.
Appendix D
Additional Details for Section
In this section, we summarize the methodology verifying base HP transferability across widths, batch sizes, and training steps in Algorithms 2 and respectively, We also present the HTransfer algorithm [73 in Algorithm for the completeness.
Algorithm Validate the base HP Transferability of μP across widths
Input: of widths {n,)g=1r set of base HP validate (Ay) other fixed base HPs 0, batch
size B, training iteration T, base width Tbaas
2: Output: whether model with μP enjoy base HP transferability across widths
3: for = 1 to P do
4:
for j = to do
5:
actual HP for n;-width model actual HP n.-width model My f train n. widtt model with b.s. B, iter. T, HPs Azg, O,
Ài < μx(()j, Thbse, 71:)
6:
0, <- μP(0, rbaaee n1))
7:
8:
Sij 4 evaluate score for Mij
9:
end for
10: X 4- arg maxx, (1
Obtain the optimal base HP for widwhd model
11:
end for
12: return whether X;s are approximately same
Algorithm 2 Validate the base HP Transferability of μP across batch sizes
1: Input: of batch sizes {B:)e set of base HP validate (Ai) other fixed base HPs O,
width n, training iteration T', base width 7abase
2: Output: whether model with jP enjoy base HP transferability across batch sizes
3: for = 1 to P do
4: for j to do
5:
Ai < μιP(}j. Rbbse, 71)
actual HP for n--iddth model
6:
0, < pP(0, Mbase. n1)
actual HP for n-width model
7:
M,j + train n-width model with b.s. B,, iter. T, HPs Xij. 0
8:
8ij f- evaluate score for Mij
9:
end for
X; +- arg maxx, ((
Obtain the optimal base HP for n-width model
end for
12:
return whether X/s are approximately same
31

---

Algorithm 3 Validate the base HP Transferability of μP across training steps
1: Input: of training steps {Ti) set of base HP validate (Ay) other fixed base HPs 0,
batch size B, width n. base width Rhase
2: Output: whether model with μP enjoy base HP transferability across training steps
3: for = 1 to P' do
4:
for j to R do
5:
X; < μdP(}j. Rbbse. 11)
actual HP for n-widthh model
6:
actual HP for n-width model Mij < train n-width model with b.s. B, iter. Tú HPs Xij. 0,
0, < pP(0, Mbaaee n)
7:
8:
8ij f- evaluate score for Mij
9:
end for
X f- arg max ((
Obtain the optimal base HP for n-width model
end for return whether X are approximately same
12:
Algorithm 4 p/Transfer (Algorithm in (3)
Input: base/proxyttarget width Rbass//ppros///earggr proxy/target batch size Bproxy/Baarge proxy/target training steps prooey//T set of base HPs to search {0:
2: Output: Trained target model Mtarget
3: for j
do
4: θ proxy 6 jaP(O), Rbaas "proxy)
> actual HPs of proxy model
5: M2 train Aprrxxy-Wiid width model with batch size B Bproxy, training steps proxy.. HPs 0 proxy.
6:
8 t evaluate score for Mi
7: end for
8: 0" +- arg maxe ({5))
9: O target + pP NBbass: Rtanget)
actual HPs of target model
10: Mtarget f- train Ptarget- -Width model with batch size Brarget, training steps Tiagett HPs 0 Orarget
11: return Mtarget
Appendix E Additional Experimental Details and Results
In this section, present the additional experimental details and results which neglected by the main text.
E.1 Assets and Licenses
All used assets (datasets and codes) and their licenses are listed in Table
Table 7: The used assets and licenses.
URL
Citation
License
http::/wwwwimagg-ntt.org
13
non-commercial
33
Apache-2.0
http://wngg..oooodnnsscoo//////n
Commons Attribution 4.0
http://ghhhbbboomm//gghhh/////gg
2
MIT
http://gghhbbbbomm//cccoooeeeeen
52
Link
15
MIT
Apache-2.0
http://gghhubbcommmmirrrooo//ppp
73
MIT
32

---

E.2
Additional Details of DiT Experiments
E.2.1 Evaluation Implementation
Following the Table in original DiT paper [2 we use the codebase of ADM 15 to obtain the results of FID, sFID, IS, precision, and recall without cfg.
E.2.2
Additional Results of Base HP transferability
Base learning rate. Comprehensive results of base learning rate transferability across widths, batch sizes, and training steps are presented in Tables 8 and respectively.
Table 8: DuP enjoys base learning rate transferability across widths. We use hatch size of 256 and 200K training iterations. NaN data points indicate training instability, where the loss explodes. Under pPP the base learning rate can transferred across model widths.
Width
(Ir)
FID-50K 4
SFID-SOK /
Inception Score t
Precision 1
Recall
144
-9
NaN
NaN
NaN
NaN
NaN
144
-10
88.89
16.19
14.24
0.30
0.437
144
-11
91.88
16.79
13.73
0.29
0.431
144
-12
93.61
17.94
13.41
0.28
0.406
144
-13
102.99
21.26
11.93
0.25
0.366
288
-9
NaN
NaN
NaN
NaN
NaN
288
-10
61.65
9.06
20.60
0.41
0.563
288
-11
63.85
10.32
20.59
0.40
0.561
-12
65.99
10.79
19.73
0.38
0.544
288
-13
79.17
13.04
15.92
0.33
0.503
576
-9
NaN
NaN
NaN
NaN
NaN
576
-10
43.73
6.65
28.82
0.51
0.611
576
-11
45,00
7.22
28.61
0.50
0.606
576
-12
50.30
7.87
25.67
0.47
0.602
576
-13
66.20
9.67
18.91
0.39
0.562
Base output multiplier. Likewise that for the base learning rate, we also sweep the base multiplier of the output weight over the set (2-6,2-4, 2-1,2-22 2-2 and training steps. FID-50K results are shown in Figure 7 and comprehensive results for other presented Figure 7 the optimal base output multiplier is approximately transferable across different widths, batch sizes, and training iterations.
20 22 2 26] across various widths, batch sizes
metrics are presented in Tables [1] 12] and [3] respectively.
Width 144
Datch size-120
200
-teration= -100K
140
150
Width=288
Batch size=256
Aeration= =15K
FID-50K 50K
FID 50K
120
sire=512
FID sok
150
100
100
100
lags out multiphter)
logs(base out multipHter)
logs out multiplier)
(a) HP transfer across widths.
(b) HP transfer across batch sizes. (c) HP transfer across iterations.
Figure 7: DiT-uP enjoys base HP transferability. Unless otherwise specified,
model width of 288,
batch size of 256, and training iteration 200K. Missing data points indicate training instability, Under j.P. the hase output multiplier can be transferred across model widths, batch sizes, steps.
33

---

Table 9: DiT-pPP enjoys base learning rate transferability across batch sizes. We use width of 288 and a training iteration 200K. NaN data points indicate training instability, where the loss explodes. We bold the best result and underline the second-best result. Under JP, the base learning rate can be (approximately) transferred across batch sizes.
Batch size
(Ir)(Ir)
FID-50K 4
sFID-50K
Inception Score
Precision
Recall
128
-9
NaN
NaN
NaN
NaN
NaN
128
-10
74.54
10.32
16.70
0.358
0.513
128
-11
73.47
10.33
16.83
0.366
0.523
128
-12
76.98
12.12
16.50
0.343
0.498
128
-13
88.79
14.63
14.10
0.293
0.453
256
NaN
NaN
NaN
NaN
NaN
256
-10
61.65
9.06
20.60
0.407
0.563
256
-11
63.85
10.32
20.59
0.398
0,561
256
-12
65.99
10.79
19.73
0.383
0.544
256
-13
79.17
13.04
15.92
0.327
0.503
512
-9
53.90
8.52
24.18
0.452
0.585
512
-10
53.72
9.27
24.66
0.448
0.592
512
-11
54.36
9.33
24.96
0.442
0.587
512
-12
57.60
10.13
23.37
0.420
0.577
512
-13
72.09
12.15
17.81
0.356
0,537
Table 10: DiT-uP enjoys enjoys base learning rate transferability across training steps. We width of 288 and batch size of 256. NaN data points indicate training instability, where the loss explodes. Under j.P the base learning rate can transferred across training steps.
Iteration
loga(1n (Ir)
FID-50K
sFID-50K J
Inception Score
Precision
Recall 1
100K
-9
74.40
11:37
16.75
0.363
0.504
100K
-10
75.35
10.67
16.61
0.363
0.514
100K
-11
78.57
11.69
16.17
0.338
0.495
100K
-12
84.74
12.97
14.78
0.315
0.474
100K
-13
102.79
18.37
11.74
0,243
0.363
150K
-9
70.44
10.07
18.03
0.396
0.534
150K
-10
67.48
9.45
18.83
0.387
0.553
150K
-11
69.61
10.51
18.45
0.380
0.533
150K
-12
72.72
11.33
17.52
0.357
0.515
150K
-13
88.05
14.07
14.34
0.296
0,455
200K
-9
NaN
NaN
NaN
NaN
NaN
200K
-10
61.65
9.06
20.60
0.407
0.563
200K
-11
63.85
10.32
20.59
0.398
0.561
200K
-12
65,99
10.79
19.73
0.383
0,544
200K
-13
79.17
13.04
15.92
0.327
0.503
34

---

Table 11: DiT-pPP enjoys base output multiplier transfcrability across widths. use batch size of 256 and a training iteration of 200K. NaN data points indicate training instability, where the loss explodes. Under pP, the base output multiplier can be (approximately) transferred across model widths.
Width
(oot)
FID-50K
sFID-50K +
Inception Score
Precision
Recall
144
168.81
62.70
6.30
0.157
0,061
144
119.38
27.05
10.77
0.185
0.300
144
104.26
15.65
12.17
0.262
0.419
144
0
88.89
16.19
14.24
0.297
0.437
144
89.01
16.83
13.97
0.294
0.433
144
88.68
15.71
14.19
0.303
0,441
144
100.03
19.33
12.48
0.262
0.386
288
-6
134.82
41.52
9.44
0.184
0.186
288
-4
63.93
8.98
20.31
0.395
0.552
288
63.65
9.15
20.23
0.401
0.563
288
61.65
9.06
20.60
0.407
0.563
288
64.64
9,44
19.30
0.398
0.555
288
59.85
9.39
20.95
0.415
0.566
288
NaN
NaN
NaN
NaN
NaN
576
-6
167.14
69.32
5.04
0.130
0.051
576
44.82
6,66
29.47
0,510
0.592
576
-2
42.75
6.68
30.73
0.526
0.600
576
43.73
6.65
28.82
0.509
0.611
576
41.93
6.68
29.45
0.537
0.600
576
46.87
6.74
26.78
0.501
0.594
576
52.93
7.30
24.07
0.455
0.584
35

---

Table 12: DiT-pPP enjoys base output multiplier transfcrability across batch sizes. use width of 288 and a training iteration of 200K. NaN data points indicate training instability, where the loss explodes. Under pP, the base output multiplier can be (approximately) transferred across batch sizes.
Batch Size logz(Pout)
FID-50K 4
sFID-50K +
Inception Score
Precision t
Recall
128
-6
148.06
51.95
7.76
0.170
0.126
128
-4
74.74
10.21
17.23
0.357
0.513
128
-2
77.62
11.55
16.20
0.346
0.511
128
0
74.54
10.32
16.70
0.358
0.513
128
2
73.55
10.24
16.68
0.365
0.514
128
4
73.19
9.93
16.93
0.375
0.514
128
6
83.48
11.79
14.92
0.323
0.464
256
-6
134.82
41.52
9.44
0.184
0.186
256
-4
63.93
8.98
20.31
0.395
0.552
256
-2
63.65
9.15
20.23
0.401
0.563
256
0
61.65
9.06
20.60
0.407
0.563
256
2
64.64
9,44
19.30
0.398
0.555
256
4
59.85
9.39
20.95
0.415
0.566
256
6
NaN
NaN
NaN
NaN
NaN
512
-6
82.00
14.83
16.13
0.325
0.468
512
-4
53.99
8.83
24.49
0.447
0.585
512
-2
53.85
9.11
24.88
0.443
0.584
512
0
53.72
9.27
24.66
0.448
0.592
512
2
54.62
9.14
23.97
0.449
0.593
512
4
55.68
8.49
23.29
0.441
0.597
512
6
NaN
NaN
NaN
NaN
NaN
36

---

Table 13: DiT-pP enjoys base output multiplier transferability across training steps. use width 288 and a batch size of 256. NaN data points indicate training instability, where the loss explodes. Under jP, the base output multiplier can be (approximately) transferred across training steps.
Iteration loge(óout) (Dout)
FID-50K
sFID-50K
Inception Score
Precision
Recall 1
100K
194.59
87.88
4.38
0.084
0.016
100K
75.87
10.28
16.64
0.350
0.508
100K
-2
72.09
9.43
17.15
0.352
0.523
100K
0
75.35
10.67
16.61
0.363
0.514
100K
75.43
9.87
16.57
0.361
0.529
100K
72.26
9.47
17.08
0.378
0.514
100K
NaN
NaN
NaN
NaN
NaN
150K
-6
150.86
48.06
8.17
0.163
0.143
150KK
-4
68.94
9.26
18.57
0.377
0.541
150K
-2
68.24
9.77
18.78
0.386
0.535
150K
67.48
9.45
18.83
0.387
0.553
150K
68.67
9.29
17.96
0.383
0.545
150K
67.32
9.08
18.17
0.397
0.539
150K
NaN
NaN
NaN
NaN
NaN
200K
-6
134.82
41.52
9.44
0.184
0.186
200K
63.93
8.98
20.31
0.395
0,552
200K
-2
63.65
9.15
20.23
0.401
0.563
200K
0
61.65
9.06
20.60
0.407
0.563
200K
64.64
9.44
19.30
0.398
0.555
200K
59.85
9.39
20.95
0.415
0.566
200K
NaN
NaN
NaN
NaN
NaN
37

---

E.2.3 Additional Results of DiT-XL-2 Pretraining
Complete benchmark results of DiT-XL-2 and DiT-XL-2-juP throughout the training are provided in Table[[4 DiT-XL-2-1uP, using base learning rate transferred from small proxy models, consistently outperforms the original DiT-XL-2 throughout the training process.
Table 14: Benchmark results of D and DIT-XI.-2-p.1 without classifier-free guidance. Both models trained the ImageNet x 256 dataset. Results with reported in the original paper (52), and others are reproduced by us. DiT-XI-22-pP, using a hase learning rate transferred from small proxy models, consistently outperforms the original baseline throughout the training process.
Iteration
Method
FID-50K +
SFID-5OK | Inception Score
Precision
Recall
DIT-XL-2 [52]
48.17
7.33
26.61
0.49
0.595
0.1M
DiT-XL-2-1(P
44.15
6.72
28.54
0.54
0.575
DiT-XL-2 [52]
20.38
6.37
65.02
0.63
0.635
0.4M
DiT-XI--2-1.P
18.63
5,49
67.74
0.66
0.606
DiT-XL-2 [52)
14.73
6.37
85.03
0.66
0.638
0.8M
D-T-XL-2-11PP
12,91
5.48
89.46
0.68
0.634
DIT-XL-22 [52]
12.78
6.41
95.13
0.66
0.644
1.2M
DiT-XL-2-1uP
11.25
5.56
99.89
0.68
0.638
DIT-XL-2 [52]
11.79
6.49
100.88
0.67
0.651
1.6M
D
10.16
5.59
107.58
0.69
0.649
DIT-XL-2 12
10.97
6.45
105.86
0.67
0.657
2M
2-1-XL-2-11P
9.72
5.62
111.23
0.69
0.651
2.4M
D-T-XI--2
52
10.75
6.59
108.23
0.67
0.665
7M
DiT-XL-2
[52]
9.62*
2.4M
DiT-XL-2-1PP
9.44
5.66
112.98
0.68
0.653
E.2.4
Computational Cost
takes 104 (13x8) A100-80GB hours to train the DiT-uP with a width of 288, batch size of 256, and a train iteration of 200K steps. The computational cost of other DiT-pPP proxy models can be inferred based on this situation. It takes around 224 (32x7) A100-800000 80-8 days to reproduce the pretraining of DiT-XL-2-pP.
38

---

E.3 Additional Details of PixArt-a Experiments
E.3.1 FLOPs ratio of HP Search on Proxy Models
Based on Equation in the main text, we can derive that
ratio R.Srraxy Eproxy
x 0.04B x 5
5.5%.
Starget Brarget
0.61 30
Therefore, the HP tuning requires only 5.5% FLOPs of single training run for PixArt-a.
E.3.2 Additional Results of Base Learning Rate Search
We present the detailed GenEval results 21 of different base learning rates in Table 15 base learning rate 2-10 achieves the best GenEval performance.
Overall, the
Table 15: GenEval results of base learning rate search PixArt-a-pPP proxy models. 0.04B models with different learning rates are trained for epochs on the SAM dataset. Overall, the base learning rate 2 achieves the best GenEval performance.
logg(Ir) (Ir)
Overall f Single
Two
Counting
Colors
Position
Color Attribution
-9
-10
0.083
30.63
4.04
2.5
12.23
0.5
0
-11
0.078
29.38
5.05
1.25
9.57
1,5
0.25
-12
0.030
12.81
0.31
4.79
0.25
0
-13
0.051
20.62
3.54
0.31
5.59
0.25
0
E.3.3
Additional Results of PixArt-a Pretraining
We present the complete comparison between PixArt-a and PixArt-a-pP throughout training in Table 16 The results demonstrate that PixArt-a-juP consistently outperforms PixArt-a across all evaluation metries during training.
Table 16: Benchmark comparison between PixArt-a-aP and PixArt-a. Both models are trained on the SAM dataset for 30 epochs. PixArt-a-puP with transferred base learning rate from the optimal 0.04B proxy model consistently outperforms the original baseline throughout the training process.
MJHQ
MS-COCO
Epoch
Method
GenEval
FID-30K +
CLIP Score + FID-30K 4. CLIP Score
PixArt-a 8
0.14
43.19
25.16
42.25
27.02
6
PixArt-a-uP (Ours)
0.17
34.24
25.99
32.62
28.16
PixArt-a
g
0.19
38.36
25.78
34.58
28.12
10
PixArt-a-pP (Ours)
0.20
33.35
26.25
29.68
28.87
PixArt-a
8
0.22
36.19
26.21
33.71
28.35
16
PixArt-a-juP (Ours)
0.23
32.28
26.67
28.03
29.45
PixArt-a
[
0.20
35.68
26.54
30.13
28.81
20
PixArt-a-puP
0.23
33.42
26.83
29.05
29.53
PixArt-a
8
0.18
38.39
26.44
34.98
29.17
26
PixArt-a-paP (Ours)
0.28
32.34
27.17
27.59
29.83
PixArt-a
8
0.15
42.71
26.25
37.61
28.91
30
PixArt-a-iPP (Ours)
0.26
29.96
27.13
25.84
29.58
39

---

E.4
Additional Details of MMDIT Experiments
E.4.1 HPs Tuned by Human Experts
Algorithmic experts take roughly times the cost of full pretraining to tune HPs based on their experience. The best HPs are a learning rate of 2E-4, a gradient clip of 0.1, REPA loss weight of 0.5 and warm-up iteration of 1K.
E.4.2 Additional Details of Base HP Search
Based on the optimal base HPs searched from 80 trials trained for 30K steps, we further evaluate the HP transferability across iterations. We conducted an additional proxy training runs (100K steps each) using the searched optimal HPPs and five different base learning rates, as detailed in Table 17
Table 17:
Iteration of 30K is enough for proxy task. The base learning rate of 2.5E-4 is optimal when proxy models are trained for 100K steps, consistent with the results observed at 30K steps.
Base learning rate
1E-4
2F-4
2.5E-4
4E-4
8F-4
Training loss
0.185028
0. 0.184505
0.184423
0.184436
0.185562
E.4.3 FLOPs ratio of HP Search on Proxy Models
Based on Equation I in the main text, we can derive the ratio of tuning cost to single pretraining cost
ratio RSproxy Eproxy
80 x 0.18B x 30K + 5 x 0.18B x 100K 14.5%.
Starget Etarget
18B 200K
Besides, as algorithmic experts take roughly times the cost of full pretraining, the ratio of μP tuning cost to standard human tuning cost is approximately 3%.
E.4.4 Additional Results of Training Loss Comparison
We present the complete comparison between the training loss of MMDiT- and MMDiT-pPP188 in Figure 8 MMDiT-pP-188 18B achieves consistently lower training loss than the baseline after 15K steps, and the advantage gradually increasing.
Training loss gap/10-4
baseline better
pP better
100
Iteration/K
Figure 8: Comparision between training loss. We present the training loss of MMDIT- -PP-ITT--P-1BBBB minus that of MMDiT-18B. The training loss gap less than 0 means that MMDIT-pP- 18B is better. MMDiT-uP 18B achieves consistently lower training loss than the baseline after 15K steps, and the advantage is gradually expanding.
E,4.5 Additional Details of Human Evaluation
In this section, we provide the details about the user study in Section 5.2 To evaluate the text- exxt-imaee alignment, we created an internal benchmark comprising 150 prompts, where each prompt includes an average of five binary alignment checks (yes or no), covering wide range of concepts such as nouns, verbs, adjectives (e-g., size, color, style), and relational terms (e.g., spatial relationships, size comparisons). Ten human experts conducted the evaluation, with each prompt generating three
40

---

images, resulting in a total of 22,500 alignment tests. The final score is computed as the average correctness across all test points. We present two examples in Figure
Prompt
Image
Alignment Evaluation
1. Wine glass:
One (wine glass):
One wine glass, two bottles of wine, three cans of beer
Bottles of wine:
4.
Two (bottles of wine):
Cans of beer: v
Three (cans beer):
Sydney Opera House:
The Sydney Opera House with the Eiffel tower sitting the right, and Mount Everest rising above
Eiffel tower:
3. Right:
Mount Everest:
Rising above:
Figure 9: Examples of text-image alignment test evaluated by human. Ten human experts conducted the evaluation, with each prompt containing an average of five test points and generating three images, resulting in 22,500 alignment test points. The final score is computed as the average correctness across all test points.
E.4.6 Additional Results of Visualization Comparison
In this section, we provide visual comparisons between -M-DTTpuP-18B and MMDiT-I8B baseline in Figure 10
41

---

Baseline
μP
Baseline
μP
IL92
ST AA START
The 'SARTT
white rabbit in blue jogging clothes, turtle wearing red tank top
A fugge whale soars through the clouds, with colorful unicorn standing on its back, and fittle girl riding the unicorn
Horses pulling carriage on the funar surface. The background is the
Statue Liberty and the Great Pyramid. The Tarth and two spaceships can be seen the sky
Seen from above the sea surface, a girl in swimsuit is diving underwater with group of whales cirefing around her, in long shot, cinematic, realistic style
A tall terracotta warrior wielding a spear is confronting a Western knight in metallic armor, both holding their weapons. In between them, modern person stands with arms outstretched, trying to mediate.
Figure 10: Visual comparison between MMDDT-;pP-18B and MMDIT-18B baseline. The configurations for generating images (c.g., random seed, cfg value) of two models are the same. MMDiT- MMDDTT-pP1188 1PB shows better text-image alignment than the baseline model.
42


Reviews:
[{"review_id": "E8MlZjZg5w", "reviewer": "Reviewer_AKwD", "summary": "The paper extends $\\mu P$ theory to the mainstream diffusion transformer family. The authors rigorously prove the scaling compatibility with vanilla transformers, and derives an implementation recipe for robust HP transfer. Extensive experiments on class-conditional model DiT and text-conditional models on PixArt-$\\alpha$ and MMDiT show faster convergence, lower tuning cost and better FID/GenEval than carefully tuned baselines.", "questions": "See \"weakness\" section.\nI am willing to change my score if my concerns are addressed.", "limitations": "Yes", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. The paper shows solid theoretical proof of diffusion transformers fitting the $NE \\otimes OR^{\\top}$ program, and distinguishes it from prior uses of $\\mu P$ in diffusion transformers without a theoretical guarantee. \n2. This work also spans various image generation models and evaluates the performance with standard metrics plus human expert checks. \n3. The writing is clear and easy to follow.\n\nWeaknesses:\n1. The paper fails to cite and discuss the AuraFlow v0.1 blog (cloneofsimo, Team Fal, July 2024). They have applied $\\mu P$ empirically to an MMDiT-style model for zero-shot LR transfer.[1]\n2. The warmup-straedy-decay learning rate schedule is a widely used schedule for the training of diffusion transformers. It is unclear whether $\\mu P$ interacts benignly with schedules like these schedules. It will be beneficial if the authors could at least provide intuitive thoughts to this.\n\n[1] Introducing AuraFlow v0.1, an Open Exploration of Large Rectified Flow Models. cloneofsimo, Team fal. https://blog.fal.ai/auraflow/", "was_revised": true, "final_justification": "I appreciate the authors' rebuttal and have examined the feedback from the other reviewers together with the newly provided experiments. I am inclined to recommend acceptance for this paper: 1. Although the paper is not the first academic attempt to apply $\\mu P$ to diffusion transformers operating on image data, the breadth and depth of the empirical investigation offer comprehensive insights that will benefit the broader community. 2. Although the theoretical framework represents an extension of prior works on vanilla transformer, the paper establishes a reasonable degree of theoretical rigor that supports the methodological contributions.\n\nTaking these points together, my final score falls between “borderline accept” and “accept,” leaning toward “accept,” provided that the authors incorporate the promised clarifications (including the discussion of AuraFlow and the supplementary experiments provided in response to reviewers sRBX and hm23).", "rebuttal": "We thank Reviewer AKwD for the positive score and valuable comments, which can definitely improve the quality of this paper.\n\n## Weakness 1: missing related blog\n\nThank you very much for pointing this out. We will certainly add a citation and discussion of the AuraFlow v0.1 blog in the final version. We acknowledge that while we carefully surveyed peer-reviewed publications that use µP for diffusion models [22,52], we unfortunately overlooked the insightful AuraFlow v0.1 blog.\n\nCompared to AuraFlow v0.1, our work makes additional non-trivial contributions:\n1. We provide a rigorous theoretical proof for mainstream diffusion Transformers, thereby justifying the validity of µP for this family of models.\n2. We systematically validate the hyperparameter (HP) transferability of diffusion Transformers under µP, across multiple widths, batch sizes, and training steps.\n3. For MMDiT in particular, we conduct a more extensive search over multiple HPs and scale the model up to 18B parameters, providing detailed intermediate results and comparisons with human-tuned baselines. We believe these results offer principled and reproducible insights for future large-scale scaling of diffusion Transformers.\n\nWe will explicitly include a discussion of AuraFlow v0.1 in the related works section of the final version. We thank the reviewer again for highlighting this relevant prior effort!\n\n## Weakness 2: other learning schedules\n\nWe thank the reviewer for raising this insightful question. \n\nYes. According to our theoretical result (Theorem 3.1), which rigorously identifies the µP form of diffusion Transformers, and based on existing µP theory (please see Table 2 and Figure 4 in TP-V [66]), optimizer-related schedules and HPs, including the warmup steps, schedules (e.g., constant, cosine, and warmup-steady-decay), and maximum learning rate, can in principle be treated as HPs that can be searched in the proxy model and transferred to the large model. We will try to present a new corollary for different learning schedules in the final version.\n\nDue to limited computational resources, we have not yet conducted an extensive search over schedules in this work, but we see this as a promising and scientifically meaningful direction. Our study provides a starting point for future works to investigate the optimal schedules for large-scale diffusion Transformers.\n\nWe will clarify this point and add a discussion on this possible extension in the discussion section in the final version. Thank you again for pointing it out!"}, {"review_id": "nJScgC5l3O", "reviewer": "Reviewer_hm23", "summary": "This paper addresses a critical bottleneck in developing large-scale generative models: the prohibitive cost of hyperparameter (HP) tuning. The authors propose applying Maximal Update Parametrization (µP), a technique successful in Large Language Models (LLMs), to Diffusion Transformers.\nThe work makes three core contributions. First, it provides a rigorous theoretical proof that mainstream Diffusion Transformer architectures (including DiT, U-ViT, PixArt-α, and MMDIT) share the same µP formulation as standard vanilla Transformers. This key finding allows the direct application of existing µP methodologies. Second, the paper systematically validates that under µP, optimal hyperparameters (like the learning rate) are stable and can be reliably transferred from small \"proxy\" models to large-scale models.\nFinally, the authors demonstrate significant practical benefits. A DiT-XL-2 model trained with a transferred learning rate achieves 2.9x faster convergence. Furthermore, when scaling up text-to-image models like PixArt-α and MMDiT (up to 18B parameters), the µP approach consistently outperforms baselines while requiring only a fraction of the typical tuning cost—as little as 3% of the effort of manual expert tuning.", "questions": "- Your theoretical proof confirms that Diffusion Transformers share the same µP formulation as vanilla Transformers. However, the denoising objective in diffusion models is fundamentally different from the autoregressive, next-token prediction objective in LLMs where µP was popularized. Did you observe any instances where this difference in training objective affected HP transfer stability, particularly when transferring across a large number of training steps?\n\n- The paper scales model width by increasing the number of attention heads while keeping the head dimension fixed. Does the µP framework offer any theoretical guidance on this choice versus the alternative of scaling the head dimension itself? Have you empirically explored if one scaling strategy leads to more stable HP transfer than the other?\n\n- You note an interesting difference in optimal learning rates: for DiT/PixArt-α (trained for multiple epochs), the optimal LR was near the stability limit, while for the 18B MMDIT (trained for a single pass), it was not. Could you elaborate on why this occurs?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- The paper is built on a rigorous proof (Theorem 3.1) that formally extends µP theory to cover several state-of-the-art Diffusion Transformer architectures. This provides a principled basis for the experimental work, moving beyond simple heuristics.\n\n- The method drastically reduces the cost of finding optimal hyperparameters for billion-parameter models. By transferring HPs from small proxy models, the tuning cost for the 18B MMDiT model was reduced to just 3% of the cost of traditional manual tuning by experts, making large-scale model development more accessible and efficient. The experiments show clear, substantial improvements. Achieving 2.9x faster convergence for DiT-XL-2  and outperforming baselines in large-scale text-to-image tasks  are powerful validations of the method's effectiveness.\n\n- The approach is validated across multiple modern architectures (DiT, PixArt-α, MMDiT), different scaling dimensions (model width, batch size, number of training steps), and various generative tasks, demonstrating its robustness and wide applicability.\n\nWeaknesses: \n- The formal proof in the appendix relies on an assumption (Assumption C.1) that simplifies model architectures, for instance by treating certain fixed dimensions as 1 and assuming all layers share the same width. While the paper notes this is standard for such proofs, it is a deviation from the true complexity of the models being analyzed.\n- The paper successfully uses small \"proxy\" models for tuning, but as acknowledged in the conclusion, it does not investigate the optimal design of these proxy tasks. There is no analysis of the trade-off between the size (and cost) of the proxy model and the quality of the transferred HPs, leaving this as an open question for practitioners.\n- The experiments primarily focus on transferring the learning rate. While the learning rate is often the most critical HP, the paper does not extensively explore the transferability of other important parameters like weight decay, dropout, or optimizer-specific settings (e.g., Adam's betas), which can also significantly impact performance.", "was_revised": true, "final_justification": "While I understand that the theoretical contribution could be more substantial, I tend to accept this paper since i) it is rigorous and consistent from theoretical groundings to empirical experiments; ii) It offers practical value by helping to reduce the costly trial-and-error process involved in training large-scale diffusion models; iii) the experiments are comprehensive to support the idea and also serve as a user guide for practitioners.", "rebuttal": "We thank Reviewer hm23 for the positive support and valuable comments, which can definitely improve the quality of this paper.\n\n## Weakness 1: assumption in proof\n\nThank you for raising this thoughtful point. **In principle, the proof technique naturally extends to any finite fixed dimension (>1) and to hidden layers with different widths (non-uniform infinite-width setting)**. \n\nTo show that our theory can be extended to any finite fixed dimensions, we take the ouput projection layer $\\boldsymbol{y} = \\boldsymbol{W} \\boldsymbol{x}$ as an example (corresponding to line 717-718 at Appendix C.1), where the output dimension is any finite $d$, $\\boldsymbol{W} \\in \\mathbb{R}^{d\\times n}$, and $\\boldsymbol{x} \\in \\mathbb{R}^{n}$. To represent the operation in the NE⊗OR⊤ Program, we denote the $i$-th row of $\\boldsymbol{W}$ by $\\boldsymbol{w}^{(i)} \\in \\mathbb{R}^n$, then the output projection layer can be written as\n\\begin{align*}\n \\widetilde{c}^{(i)} _\\alpha = \\psi _1([\\boldsymbol{w}^{(i)}, \\boldsymbol{x}] _{\\alpha:}) = n {w}^{(i)} _\\alpha x _\\alpha, (\\text{OuterNonlin})  \\quad \\quad\n {y} _{i} = \\frac{1}{n} \\sum _{\\alpha=1}^n \\widetilde{c}^{(i)} _\\alpha.  (\\text{Avg})\n\\end{align*}\n\nFor the non-uniform infinite-width setting, the Avg, MatMul, and OuterNonlin operators in Appendix B.2 simply need to be extended to accept input tensors with any $\\Theta(n)$ dimension rather than a single shared width $n$.\n\nFurthermore, **please see Section 2.9 in TP-4b [65], which provides a more detailed discussion of both these generalizations** and shows how the tensor program framework naturally accommodates them.\n\nWe will clarify this point in the appendix to make explicit how these assumptions can be relaxed to better reflect the practical model architecture. We appreciate the reviewer for pointing this out and helping us improve the presentation.\n\n## Weakness 2: optimal design of proxy tasks\n\nWe thank the reviewer for this nice suggestion. We note that determining the optimal design of proxy tasks is a rarely explored question in the µP literature [66, 12, 27, 40, 22, 39, 71]. In this paper, **our empirical results on three different diffusion Transformers and datasets provide a practical starting point for the vision community**:\n1. A proxy model width of approximately 256–512 works reliably for stable hyperparameter transfer.\n2. The data scale for proxy tuning typically needs to reach about 1/10 to 1/6 of the final pretraining dataset to yield robust hyperparameter transfer.\n\nIn addition, to guide future work for determining the optimal design of proxy tasks, we **suggest a concrete experiment design**:\n1. Select a target large-scale diffusion Transformer (e.g., MMDiT-18B with a width of 5120).\n2. Construct multiple proxy tasks varying in width from 64 up to 5120, and data from 10% up to 100% of the target dataset.\n3. For each proxy task, tune the base HPs and evaluate their transferability to the target task.\n4. Analyze the relationship between the proxy task scale and final performance to identify the optimal proxy scale.\n\nGiven our current computational budget, we could not conduct this sweep, but we strongly encourage future work to investigate this direction. We will add this discussion to the final version to better highlight this practical research frontier. Thank you for encouraging us to clarify this.\n\n## Weakness 3: transferability of other HP\n\nThank you for this thoughtful comment. We agree that investigating the transferability of other hyperparameters is important. **In Figure 7 at Appendix E.2.2, we also empirically verified the transferability of the output multiplier**, which is an important hyperparameter in µP theory and practice [66, 40, 5]. We will make this more visible in the main text to avoid confusion.\n\nIn addition, although we did not run an explicit transferability study for each individual base hyperparameter in MMDiT, it is worth noting that the hyperparameters found on the proxy model transferred successfully to the target 100× larger model, outperforming hand-tuned settings. **This strongly suggests that other base hyperparameters are reasonably stable under µP as well**.\n\nWe are open to the reviewer’s further feedback. Thank you again for raising this practical point and helping us clarify our contribution!\n\n## Question 1: influence of denoising objective\n\nThank you for raising this thoughtful question. Testing whether the different training objective and data domain affect the hyperparameter transferability of µP is also one of the key motivations behind our work. In fact, across all three of our experimental settings, **we did not observe any instability in base HP transfer caused by the denoising objective**, even when transferring across a large number of training steps (from 200k steps to 2400k steps in DiT, and from 30k steps to 200k steps in MMDiT).\n\nWe will clarify this empirical observation more explicitly in the final version. We appreciate the reviewer for highlighting this point.\n\n## Question 2: why scaling number of attention heads\n\nThank you for raising this insightful question. **Our choice to scale model width by increasing the number of attention heads while fixing the head dimension follows both theoretical and practical guidance**.\n\n1. Theoretically, recent theoretical advance [5] shows that both scaling strategies are consistent with the transferability condition required by µP. However, they also reveal an important difference: **when the head dimension tends to infinity, multi-head self-attention can collapse to single-head self-attention, losing the diversity of attention patterns**. In contrast, scaling the number of heads avoids this degeneracy and preserves the expressivity of the multi-head mechanism.\n2. Empirically, the TP-V paper [66] (please see Figure 4 and Figure 13) shows that both scaling strategies yield stable base HP transfer under µP. However, as we mentioned in line 120, **well-known models in real-world large-scale practice like Llama3 [21] and MiniCPM [27] favor increasing the number of heads** rather than the head dimension.\n\nGiven these theoretical and empirical insights, we chose this widely supported scaling strategy in our experiments. A systematic empirical comparison of the two strategies under µP for diffusion Transformers is an interesting direction for future work, and we will clarify this point in the discussion section in the final version.\n\n## Question 3: optimal learning rate and stability edge\n\nThank you for highlighting this intriguing point. This phenomenon has also been noted in recent studies [58] and remains an interesting open question in large-scale training.\n1. For multiple-epoch training, it is well established that the optimal learning rate often lies near the stability limit [a; b; c]. Larger learning rates is thought to introduce beneficial gradient noise, which can help guide optimization towards flatter minima that generalize better. Several theoretical perspectives attempt to explain this, including the flat minima hypothesis [d] and the edge of stability framework [f].\n2. For single-pass training, recent work [58] in the context of LLM pretraining has also pointed out that the optimal learning rate tends to be away from the stability limit. This suggests a fundamental difference in optimization dynamics between single-pass and multi-epoch regimes, although the precise reason remains an open question. Intuitively, since the gradient signal for any individual sample is not revisited, the training must be more conservative to maintain stability.\n\nOverall, this is a meaningful difference that highlights the complex interaction between optimal learning rate, model scale, and training strategy. We agree it is a valuable topic for future theoretical investigation, and we will add a brief discussion of this point in the final version.\n\n[a] Towards explaining the regularization effect of initial large learning rate in training neural networks, NeurIPS, 2019\n\n[b] The large learning rate phase of deep learning: the catapult mechanism, 2021\n\n[c] Finite versus infinite neural networks: an empirical study, NeurIPS, 2020\n\n[d] On the Relation Between the Sharpest Directions of DNN Loss and the SGD Step Length, 2018\n\n[f] Gradient Descent on Neural Networks Typically Occurs at the Edge of Stability, ICLR, 2021"}, {"review_id": "MaTnbYw0RH", "reviewer": "Reviewer_sRBx", "summary": "This paper generalizes the standard Maximal Update Parametrization, short for μP, to diffusion\nTransformers and validate their effectiveness through large-scale experiments. \nThey rigorously prove that μP of mainstream diffusion Transformers, including DiT, U-ViT, PixArt-α, and MMDiT, aligns with that of the vanilla Transformer, enabling the direct application of existing μP methodologies. The article systematically demonstrates that DiT-μP enjoys robust HP transferability. Furthermore, they 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.", "questions": "1. This investigation of the Maximal Update Parametrization method for the diffusion transformer is valuable. \nHowever, I believe the author should focus on distinguishing between vanilla Transformers and diffusion Transformers and develop specific characteristics for the Maximal Update Parametrization method that address this distinction. Otherwise, the author could consider providing a unified theoretical formulation for Maximal Update Parametrization that encompasses both vanilla Transformers and diffusion Transformers. \n\n2. I think you can verify the convergence speed for Maximal Update Parametrization method on PixArt-α and MMDiT models.\n\n3. I think you can scale the model from a small number of parameters to a large number of parameters in Dit models and test its effectiveness.\n\n4. If the experiment scales the PixArt-α-μP starting from 0.02B or 0.08B parameters, how would the results change? If it scales the MMDiT-μP starting from 0.06B or 0.5B parameters, what differences might occur?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "# Strengths\n\nThe article presents systematic experiments on the effectiveness of the Maximal Update Parametrization method for diffusion transformers. Including DiT-XL-2-μP with transferred learning rate achieves 2.9× faster convergence than the original DiT-XL-2. And only 5.5% of one training run for PixArt-α, 3% of consumption by human experts for MMDiT-18B outperformed their respective baselines. Moreover, the article is well written and organized.\n\n\n# Weaknesses\n\nThe original contributions of the article are not sufficient. The article reads like a technical report rather than a research paper. The author just verified the Maximal Update Parametrization method on the diffusion Transformer models. Although they propose a new theorem to ascertain the effectiveness of Maximal Update Parametrization on the diffusion transformer, it only demonstrates that the diffusion transformer can be represented within the NE⊗OR⊤ Program. The new theorem does not offer any new theoretical insight. Furthermore, the author stated that the DIT includes some modules different from the standard transformer, such as Scale, Shift, T-emb, and y-emb, as highlighted in Figure 2. However, these modules appear to have no impact on the proof of DIT within the NE⊗OR⊤ Program, which is shown in the appendix.\nBased on the above content, the article lacks enough contribution and novelty.", "was_revised": true, "final_justification": "The author's response to the reviewers has addressed many of my concerns. I believe this work can help the research community explore the generative model more efficiently. However, the article still faces some unresolved issues. I will list them above.\n\n# solved issues (These responses have considered the rebuttals with reviewers hm23 and AKwD)\n1.  The method drastically reduces the cost of finding optimal hyperparameters for billion-parameter models. \n\n2.  By transferring HPs from small proxy models, the tuning cost for the 18B MMDiT model was reduced to just 3% of the cost of traditional manual tuning by experts. Meanwhile, the method achieved 2.9x faster convergence for DiT-XL-2 and outperformed baselines in large-scale text-to-image tasks. These results make large-scale model development more accessible and efficient.\n\n3. The approach is validated across multiple modern architectures (DiT, PixArt-α, MMDiT), different scaling dimensions (model width, batch size, number of training steps), and various generative tasks, demonstrating its robustness and wide applicability.\n\nThe experiments show clear, substantial improvements.\n\n# unsolved issues (These responses have considered the rebuttals with reviewers sRBx and mn36)\n1. The original contribution of the work as a research article could be more substantial. The article proves the µP form in the new architecture without more theoretical insight. This seems like a straightforward extension rather than a novel contribution.\n\nAlthough a solid theoretical proof is essential, I believe that beneficial empirical results are more helpful for the community.\n\n Based on the above comments, I am leaning towards increasing my score to 4.", "rebuttal": "We thank Reviewer sRBx for the feedback and valuable comments.\n\n## Weakness: novelty\n\nThank you for raising these concerns about the contribution and novelty of our work. We respectfully offer the following clarifications.\n### Significance and contribution\n\nFirst, we would like to emphasize that **verifying a theory or method does not imply insufficient contribution**, especially when the setting is practically important and the conclusions can impact a broad audience. For example, ViT [a] successfully changed how the community builds vision models by systematically verifying that the Transformer architecture could be transferred to vision models. In this way, **introducing and verifying µP in the large-scale diffusion Transformer is also significant and non-trivial**.\n\n**Significance**. Training large-scale diffusion Transformers requires enormous compute resources (e.g., ~100M dollars for Sora), while still lacking a principled way to choose hyperparameters at scale. This bottleneck limits the development of large vision foundation models. **Our work gives a preliminary solution to this significant problem** by showing µP as a principled method to obtain good hyperparameters with small tuning cost (e.g., 3% for MMDiT-18B). \n\n**Non-trivial**. Prior to our work, µP was rarely applied in large-scale diffusion model pretraining. This was mainly because there was **no rigorous µP theory** for diffusion Transformers, **nor clear empirical evidence** that µP would ensure stable and effective hyperparameter transfer for these models. **We take a first step to address this gap, both theoretically and empirically**. First, we formally proved the µP form for diffusion Transformers. **As acknowledged by Reviewers hm23, \"this provides a principled basis for the experimental work, moving beyond simple heuristics\"**. We then conducted large-scale experiments, which demonstrate that µP indeed enables reliable and effective hyperparameter transfer for diffusion Transformers, **providing the vision community a principled approach for scaling diffusion Transformers efficiently**.\n\n### Proof of DiT\nOn the role of additional modules: we do not fully agree that modules like Scale, Shift, T-emb, and y-emb “have no impact” on the proof. **As detailed in Appendix C (e.g., lines 664–721)**, we provide **explicit derivations** showing how these components are mapped and handled within the NE⊗OR⊤ Program. These **module-specific derivations** are essential to ensure that the entire DiT architecture adheres to the same µP consistency.\n\nIn summary, we believe our paper contributes both practically and theoretically: it rigorously extends µP theory to an important new domain, and provides the community with a validated, scalable methodology that is easy to adopt. We sincerely thank the reviewer for pushing us to better clarify these contributions and will emphasize them more clearly in the final version.\n\n[a] An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale, ICLR, 2021\n\n## Question 1: unified theoretical formulation\n\nThank you very much for this constructive suggestion. In fact, our theoretical proof is built precisely on the latter idea: we **unify the µP of diffusion Transformers and the classic µP (e.g., vanilla Transformers)** by leveraging the NE⊗OR⊤ Program framework [65] to represent the forward pass of diffusion Transformers.\n\nWe will emphasize this point more clearly in the final version to highlight that our theorem indeed offers the requested unification. Thank you for pointing this out and helping us clarify this aspect of our contribution.\n\n## Question 2: convergence speed on PixArt-α and MMDiT\n\nThank you for this suggestion. In fact, **we have already included empirical evidence for the faster convergence speed of µP on PixArt-α and MMDiT in the current version**. Table 3 and Figure 6 show the performance trajectories during training, clearly indicating that µP consistently achieves better performance than the baselines throughout the training process. This directly demonstrates that µP enables faster convergence for these models as well.\n\nWe would also like to emphasize that our main goal is not merely to maximize convergence speed, **but to provide a principled, low-cost approach to find hyperparameters** that match or even surpass human-tuned baselines—**especially in large-scale training settings** (e.g., only 3% FLOPs for MMDiT-18B).\n\nWe appreciate the reviewer’s point that this could be made clearer. We will improve the visualization and discussion in the final version to more explicitly highlight the convergence speed improvements enabled by µP.\n\n## Question 3: scaling from small DiT model to large DiT model\n\nThank you for this suggestion. As an initial exploration of µP for Diffusion Transformers, we followed the classical µP theory [63, 66] by scaling width while keeping depth fixed. **We did not scale from a small DiT variant (e.g., DiT-S) to a large DiT variant (e.g., DiT-XL) because their depths differ significantly**, which does not satisfy the assumptions required for directly verifying our theoretical guarantee.\n\nWe note that the recent work [b] proposes ideas to enable hyperparameter transfer across depth in the context of LLMs, which is an interesting direction for future work on scaling the depth of diffusion Transformers. \n\nWe thank the reviewer for highlighting this point and will mention this potential extension more clearly in the final version.\n\n[b] Don't be lazy: CompleteP enables compute-efficient deep transformers, 05/2025\n\n## Question 4: scaling PixArt-α/MMDiT from different scale\n\nThank you for this practical suggestion. We run additional experiments for PixArt-α-µP here. **As a result, we found that the base learning rate searched on the 0.04B proxy model can be transferred to a larger model (e.g., 0.087B)**.  Concretely, we keep the same setup as in Table 2 in the submission but vary the number of attention heads to get different parameter scales. The results are:\n\n| head| width | Params |$\\log_2$(lr) | Training loss $\\downarrow$ | Geneval $\\uparrow$ |\n|-| -|-|-------- | -------- | -------- |\n|2   | 144 |   0.01B  |  -9  | NaN     |  NaN    |\n|2   | 144 |   0.01B  |  -10  |   0.1539   |  0.03   |\n|**2**   | **144** |   **0.01B**  | **-11**  |   **0.1537**   |  **0.046**    |\n|2   | 144 |   0.01B  |  -12  |   0.155   |  0.038    |\n|4   | 288 |   0.04B  |  -9  | NaN     |  NaN    |\n|**4**   | **288** |   **0.04B**  |  **-10**  |   **0.1493**   | **0.083**    |\n|4   | 288 |   0.04B  |  -11  |   0.1494   |  0.078    |\n|4   | 288 |   0.04B  |  -11  |   0.1494   |  0.078    |\n|6   | 432 |   0.087B  |  -9  | NaN     |  NaN    |\n|**6** | **432**   |   **0.087B**  |  **-10**  |   **0.1475**   |  **0.134**    |\n|6     | 432  |   0.087B  |  -11  |   0.1476   |  0.132    |\n|6     | 432  |   0.087B  |  -12  |   0.1487   |  0.062    |\n\nAs shown above, the optimal base learning rate for 0.087B matches the 0.04B proxy ($2^{-10}$), **confirming that our base learning rate in the main paper transfer well to wider models**. For the 0.01B proxy, the best learning rate shifts **slightly** ($2^{-11}$), which is expected because µP theory predicts that base hyperparameter transfer is stable when the width is sufficiently large (our experiments in the main paper suggest ~256–512 width works well).\n\nThis new test supports our claim that µP enables robust hyperparameter transfer. Due to computational limits, we could not repeat this for MMDiT within the rebuttal phase, but we expect similar trends to hold.\n\nWe appreciate this suggestion and will include a clearer discussion of the practical lower-bound proxy scale in the final version."}, {"review_id": "LUN7IPEvXp", "reviewer": "Reviewer_mn36", "summary": "This paper investigates the applicability of Maximal Update Parametrization (µP) to diffusion Transformers, a class of models foundational to vision generative tasks but constrained by costly hyperparameter tuning at scale. \n\nWhile µP has proven effective for vanilla Transformers in language models, its effectiveness for diffusion Transformers—architecturally and functionally distinct—remained uncertain. The authors extend µP theory to diffusion Transformers such as DiT, U-ViT, PixArt-α, and MMDiT, showing its theoretical alignment with vanilla Transformers. \n\nEmpirical results demonstrate that models like DiT-XL-2-µP achieve significantly faster convergence and efficient hyperparameter transfer. Moreover, large-scale experiments with PixArt-α and MMDiT reveal that µP enables superior performance with minimal tuning cost, establishing µP as a robust and efficient framework for scaling diffusion Transformers.", "questions": "Q1: Can the authors better justify the novelty of their work, given that the µP formulation in diffusion Transformers (e.g., DiT) appears structurally similar to vanilla Transformers with only minor conditioning changes (e.g., class + timestep via AdaLN)? What specific technical challenges or non-trivial insights distinguish this extension?\n\nQ2: Are the performance gains still significant when classifier-free guidance (CFG) is applied during generation? Please provide comparisons with and without CFG, especially for results in Figure 4 and Table 3.\n\nQ3: In Figure 5, what is the rationale for including MMDiT in the REPA comparison, and how is it related? Can the authors clarify the meaning and role of \"base warm-up interactions\" and ensure that all such terms are well-defined in the main text?\n\nQ4: The \"base learning rate\" subplot in Figure 5 appears noisy and lacks a clear justification for selecting 2.5e-4 as optimal. Can the authors provide more systematic tuning experiments or ablation studies to support this choice?\n\nQ5: While Figure 6 shows lower training loss, it is unclear how this correlates with generation quality. Can the authors provide additional evaluation on downstream generation quality (e.g., FID, CLIP score, human studies) to confirm whether improved training loss leads to better outputs during inference?", "limitations": "Yes", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n\n+ The paper explores the adaptation of Maximal Update Parametrization (µP) to diffusion Transformers, which could be beneficial for scaling vision generative models with lower hyperparameter tuning costs.\n\n+ Some experimental results, particularly on DiT, validate the applicability of µP in the vision domain, indicating a promising direction for efficient large-scale training.\n\n**Weaknesses**\n\n+ The novelty of the work is questionable. While the authors initially suggest that µP’s hyperparameter transferability may not hold for diffusion Transformers, they ultimately prove that the µP formulation aligns with that of vanilla Transformers. This result is not particularly surprising, given the structural similarities—DiT mainly comprises standard Transformer blocks with adaptive layer normalization conditioned on class and timestep embeddings. As such, applying µP here seems like a straightforward extension rather than a novel contribution.\n\n+ The reported performance gains appear marginal. For instance, in Figure 4, the improvements seem modest and are likely measured without classifier-free guidance (CFG). With CFG applied, the performance gap might diminish further. Similar trends are observed in Table 3 across several evaluation metrics.\n\n+ In Figure 5, the inclusion of MMDiT in relation to the REPA method seems unclear and insufficiently justified. Additionally, terms like “base warm-up interactions” are introduced without adequate explanation, making it difficult to understand their relevance. The “base learning rate” subplot also appears noisy and lacks clear justification for the chosen optimal value (e.g., 2.5e-4).\n\n+ Figure 6 shows lower training loss, but this does not necessarily translate to better generative performance during inference. More compelling qualitative or quantitative evidence is needed to support the claim of improved generation.", "was_revised": true, "final_justification": "After reviewing it along with the other reviews, I find that the authors have addressed some of my concerns, particularly regarding experimental clarity and their justification for µP in diffusion Transformers. \n\nHowever, I still believe the novelty remains limited, as also pointed out by Reviewer sRBx. While the theoretical extension and large-scale empirical study are appreciated, the overall contribution feels incremental. \n\nI adjusted my score from 2 to 3 to reflect the partial resolution of concerns, but I remain cautious about the paper’s overall significance.", "rebuttal": "We thank Reviewer mn36 for the feedback and valuable comments.\n\n## Weakness 1 & Question 1: novelty\n\nThank you for the question.  We respectfully offer the following clarifications.\n\nFrom a theoretical perspective, **it is not obvious that the µP form remains valid when the network architecture is modified**. In fact, recent evidence shows that architectural changes can break µP’s conditions (e.g., the Mamba architecture [54]). Therefore, a new rigorous derivation for diffusion Transformers is necessary to ensure that additional modules do not invalidate the theoretical guarantees. **As acknowledged by Reviewers hm23, \"this provides a principled basis for the experimental work, moving beyond simple heuristics\"**.\n\nBeyond the theoretical contribution, **our large-scale empirical study is also significant and non-trivial**:\n\n**Significance**. Training large-scale diffusion Transformers requires enormous compute resources (e.g., ~100M dollars for Sora), while still lacking a principled way to choose hyperparameters at scale. This bottleneck limits the development of large vision foundation models. **Our work gives a preliminary solution to this significant problem** by showing µP as a principled method to obtain good hyperparameters with small tuning cost (e.g., 3% for MMDiT-18B). \n\n**Non-trivial**. Prior to our work, µP was rarely applied in large-scale diffusion model pretraining. This was mainly because there was **no rigorous µP theory** for diffusion Transformers, **nor clear empirical evidence** that µP would ensure stable and effective hyperparameter transfer for these models. Based on our rigorous µP theory, we conducted large-scale experiments, which demonstrate that µP indeed enables reliable and effective hyperparameter transfer for diffusion Transformers, **providing the vision community a principled approach for scaling diffusion Transformers efficiently**.\n\nWe believe this bridge from theoretical consistency to robust practical evidence distinguishes our work from prior µP studies and provides meaningful value to the vision and generative modeling communities.\n\n\n## Weakness 2 & Question 2: comparisons with and without CFG\n\nThank you for the question. We respectfully offer the following clarifications.\n\n### Performance gains appear marginal\n\nWe would like to emphasize that our main goal is not to improve performance significantly, but to **provide a principled, low-cost approach to find hyperparameters** that match or even surpass human-tuned baselines—**especially in large-scale training settings** (e.g., only 3% FLOPs for MMDiT-18B).\n\n### Results with and without CFG\n\nTo clarify, **for each experiment, we followed the exact setup of the corresponding baseline paper to ensure fair comparison**. Specifically, DiT does not use CFG during training evaluations (see Table 4 in [48]), so our DiT hyperparameter search in Figure 3 also does not include CFG. In contrast, PixArt-α uses CFG during evaluation (see Table 2 in [8]), so our hyperparameter search for PixArt-α (Table 2) includes CFG as well.\n\nBecause CFG is an important factor during inference and **affects the optimal hyperparameters**, reproducing results in Figure 4 with CFG for DiT would principly require re-running the full proxy search and retraining the DiT-XL-2 under the new searched hyperparameters, which would cost roughly **300 A100 GPU days** (as detailed in Appendix E.2.4). Unfortunately, this is not feasible within the rebuttal period, but we are willing to include these additional runs in the final version.\n\nNevertheless, **our current results already cover both settings: µP works well without CFG (DiT) and with CFG (PixArt-α, MMDiT)**. In both cases, we observe stable hyperparameter transferability and performance improvements over the baseline.\n\n## Weakness 3 & Question 3: unclear REPA loss and warm-up\n\nThank you for raising this important point. Though we mention these hyperparameters in Section 5.2.1 (lines 260–264), we agree they could be explained in more detail. Below, we provide a clearer clarification:\n\n1. The REPA loss [69] is an auxiliary regularization term widely used in diffusion Transformer to accelerate training, which explicitly aligns the diffusion model representation with powerful pretrained visual representation (e.g., DINOv2). Therefore, we include REPA loss and aim to **provide practitioners a suitable REPA weight for large-scale diffusion Transformers training**.\n2. In fact, **there is no term “base warm-up interactions” in our paper**. The term used is “base warm-up iteration”, which refers to the warm-up steps used in the learning rate schedule. This is standard practice for large models, where the learning rate starts small and gradually grows up to its target value to improve stability at the start of training.\n\nThank you for highlighting this. We will revise the text to define REPA weight and warm-up iteration more explicitly in the main paper so that their meaning and role are unambiguous.\n\n## Weakness 3 & Question 4: selection of learning rate\n\nThank you for raising this practical question.\n\n1. **We follow the established practice in the µP literature** [12, 27, 66], where the base hyperparameter is typically selected by identifying the value that aligns with the lowest envelope of the training loss plot. In our case, **we observed that the envelope near 2.5e-4 was stable and close to optimal**, so we chose 2.5e-4 as a representative value within this range.\n2. As noted in lines 274–275 and detailed in Appendix Table 15, **we did perform a more systematic check**: we ran additional training with five candidate learning rates around 2.5e-4, each for 100K steps (the original sweep in Figure 5 covered 30K steps). The results confirmed that **learning rates near 2.5e-4 (2e-4~4e-4) all produced robust and strong performance**, with 2.5e-4 yielding the best overall results.\n\nWe will clarify the envelope selection logic and highlight the extended ablation in the final version.\n\n## Weakness 4 & Question 5:  evaluation on downstream generation quality\n\nThank you for raising this valid point about generation quality.\n\nTo address this, **we have already included human evaluation results in the submission**. As described in **Section 5.2.1 (lines 280–287) and Table 4**, we conducted a comprehensive human study where 10 human experts compared outputs from µP and baseline models. The µP-trained MMDiT model outperformed the baseline in human preference scores, confirming that the observed training loss reduction translates to better sample quality. **The visual comparison is also presented in Figure 10 in Appendix E.4.6**.\n\nIn addition, **we now provide the Geneval score, which is a standard quantitative metric used for text-to-image models** (e.g., Stable Diffusion 3 [15]). The **overall** Geneval score for the µP-trained MMDiT model is higher than the baseline, further supporting the claim that lower training loss yields better generative performance at inference.\n\n\n| Method       | Overall $\\uparrow$ | Single | Two | Counting | Colors | Position | Color Attribution |\n| ------------ | -------- | -------- | -|-|-|-|-|\n| MMDiT-18B     | 0.8154     | 99.38    | 93.69 | 81.88 | 88.03 | 57.50 | 68.75|\n| MMDiT-μP-18B     | **0.8218**     | 99.38    | 94.44 | 79.69 | 88.83 | 62.25| 68.50|\n\nThank you for suggesting this important clarification. We agree that these links can be presented more clearly and will highlight them alongside Figure 6 in the final version."}]

Author rebuttal:

Meta-review: This paper extends Maximal Update Parametrization (µP) to diffusion Transformers with both theoretical grounding and large-scale empirical validation. It demonstrates robust hyperparameter transfer across architectures (DiT, PixArt-α, MMDiT), reducing tuning cost by over 90% with competitive performance. While some reviewers noted limited novelty, others praised its rigor, broad applicability, and practical impact. The rebuttal addressed concerns thoroughly with new experiments, clearer discussions, and stronger justification. Overall, the paper offers a grounded approach to scaling diffusion models efficiently. I recommend acceptance.
Decision: Accept (poster)