Title: Scaling Diffusion Transformers Efficiently via $\mu$P
Venue: NeurIPS 2025
Accepted: yes
Keywords: Diffusion Transformer

muP

Tensor Program

image generation

AIGC

Scaling
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 µ 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 ( µ P) was proposed for vanilla Transformers, which enables stable HP transfer from small to large language models, and dramatically reduces tuning costs. However, it remains unclear whether µ P of vanilla Transformers extends to diffusion Transformers, which differ architecturally and objectively. In this work, we generalize standard µ P to diffusion Transformers and validate its effectiveness through large-scale experiments. First, we rigorously prove that µ P of mainstream diffusion Transformers, including DiT, U-ViT, PixArtα , and MMDiT, aligns with that of the vanilla Transformer, enabling the direct application of existing µ P methodologies. Leveraging this result, we systematically demonstrate that DiTµ P enjoys robust HP transferability. Notably, DiT-XL-2µ P with transferred learning rate achieves 2.9 × faster convergence than the original DiT-XL-2. Finally, we validate the effectiveness of µ P on text-to-image generation by scaling PixArtα from 0.04B to 0.61B and MMDiT from 0.18B to 18B. In both cases, models under µ P outperform their respective baselines while requiring small tuning cost-only 5 . 5% of one training run for PixArtα and 3% of consumption by human experts for MMDiT-18B. These results establish µ P as a principled and efficient framework for scaling diffusion Transformers .

1 Introduction

Owing to its generality and scalability, diffusion Transformers [52; 1] 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; 81; 62; 2; 77]. As datasets grow and task complexity increases, further scaling of diffusion Transformers has become inevitable and is attracting increasing attention [41; 40; 29; 17; 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 a principled approach to efficiently identify the optimal HPs when scaling diffusion Transformers.

Maximal Update Parametrization ( µ P) [73; 70; 72] was recently proposed as a promising solution to the HP selection problem for large-scale vanilla Transformer [61]. It stabilizes optimal HPs across different model widths, enabling direct transfer of HPs searched from small models to large models (a.k.a., µ Transfer) and significantly reducing tuning costs at scale. Due to its strong transferability, µ P has been applied to the pretraining of large language models (LLMs) [73; 14; 28; 43; 42; 82].

1

×

Train

𝝁

P Search

100%

80%

⋮

20%

0%

5.5%

( i ) PixArt-

𝜶

( ii ) MMDiT-18B

Manual

Tuning

𝝁

P Search

3%

100%

80%

⋮

20%

0%

Relative FLOPs

Relative FLOPs

(b) Efficiency of

µ

P search.

Acolorfuldonutfixed

Afurrylittleblack

Bluesky,calmsea,

onawhiteclothesline,catCookedupandsaw

a huge blue whale

withabrightskyand

hfingmfinpg

leaps from the sea,

fluffy white cloudsin

Canding on its nose

andtheyacht

thebackground

belowissmall

?

G

Acuteteddy bear

Painting Style

holding a guitar

inside the spacesuit

standsinascenewith

star lights and grass

Unfortunately, diffusion Transformers [52; 1; 16; 8] differ fundamentally from vanilla Transformers [61]. First, their architectures incorporate additional components to integrate text information and diffusion timesteps for final vision generation. Second, they operate under a distinct generative framework based on iterative denoising, in contrast to the autoregressive paradigms typically used in vanilla Transformers (e.g., LLMs) [22; 80]. As a 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.

First, we extend the standard µ P theory from vanilla Transformers to diffusion Transformers in Section 3.1. Using the Tensor Programs technique [72; 70; 67], we rigorously prove that the µ P formulation of mainstream diffusion Transformers, including DiT [52], U-ViT [1], PixArtα [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 to diffusion Transformers, as described in Section 3.2.

Based on the rigorous µ P result of diffusion Transformers, we then conduct a systematic study of DiTµ P on the image generation task using the ImageNet dataset [13], presented in Section 4. We first verify the stable HP transferability of DiTµ P across widths, batch sizes, and training steps. We then µ Transfer the optimal learning rate searched from small models to DiT-XL-2µ P. Notably, DiTXL-2µ P trained with the transferred learning rate achieves 2.9 × 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 µ Transfer algorithm to PixArtα [8], scaling the model from 0.04B to 0.61B. In Section 5.2, we apply it to MMDiT [16], scaling from 0.18B to 18B. In both cases, diffusion Transformers under µ P consistently outperform their respective baselines with small HP tuning cost on proxy tasks. For PixArtα , 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 µ P.

2 Preliminaries

We begin by establishing the necessary background for diffusion Transformers and µ P. Detailed discussion of additional related work is placed in Appendix A.

2.1 Diffusion Transformers

Due to its superior scalability and compatibility, Transformers [52; 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]. UViT [1] firstly introduces Transformers with long skip connections between shallow and deep layers for various image generation tasks and achieves remarkable performance. After that, DiT [52] proposes Transformers with adaptive layer normalization (adaLN) blocks for class-to-image generation and demonstrates strong scalability with respect to network complexity. PixArtα [8] extends DiT for text-conditional image generation by incorporating cross-attention for text features and an efficient shared adaLN-single block. MMDiT [16] further extends DiT by introducing two separate parameter sets for image and text modalities, along with a joint attention mechanism to facilitate multimodal interaction. As diffusion Transformers continue to scale, increasing attention is being paid to principled approaches for scaling [41; 40; 29; 75].

2.2 Maximal Update Parametrization

In this section, we provide a practical overview of µ P, with a focus on the widely used AdamW optimizer [47]. The comprehensive review of the theoretical foundations of µ P is in Appendix B.

µ P identifies a unified parameterization that applies to common architectures expressive in NE ⊗ OR glyph[latticetop] Program [72] (e.g., vanilla Transformer [61]), offering strong guidance for practical HP transfer across model widths, batch sizes, and training steps. Under µ P, HPs can be tuned on a small proxy task (e.g., 0.04B parameters and 6B tokens in [73]) and directly transferred to 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 LLMs [73; 14; 28; 43; 23; 42; 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 -parameterization). Formally, let n denote the network width, we set each weight as W = φ W n -a W ˜ W , where the trainable component ˜ W is initialized as ˜ W ij ∼ N (0 , σ 2 W n -2 b W ) , and its learning rate is η W n -c W . Henceforth, we call the width-independent parts ( φ W , σ W , η W ) as base HPs . As summarized in Table 1, µ P identifies values of a W , b W , and c W that enable models of different widths to share the (approximately) same optimal base HPs φ ∗ W , σ ∗ W , and η ∗ W . µ P adjustion for other HPs is placed in Appendix B.4.

The most related work to us is AuraFlow v0.1 [10], which applied µ P empirically to an MMDiT-style model for learning rate transfer. Compared to 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 µ P for this family of models. Second, we systematically validate the HPtransferability of diffusion Transformers under µ P, across multiple widths, batch sizes, and training steps. Finally, 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.

3 Scaling Diffusion Transformers by µ P

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-including U-ViT [1], DiT [52], PixArtα [8], and MMDiT [16]-adhere to the standard µ P formulation summarized in Table 1. Then, in Section 3.2, we introduce the practical methodology for applying µ P to diffusion Transformers, as illustrated in Figure 2.

+

Noised

Latent

Patchify

Timestep

𝑡

𝑡

-Emb 1

DiT Block

Layer Norm

Noise

Σ

Label

𝑦

𝑡

-Emb 2

𝑦-

Emb 1

N

×

Layer Norm

Scale, Shift

Multi-Head

Self-Attention

Scale

Layer Norm

Scale, Shift

Pointwise

Feedforward

Scale

Condition

adaLN

+

+

Input Weights

Hidden Weights

Output Weights

New in DiT

(a) Implementation of DiT-

µ

P.

𝝁

P search base HPs on

small models

HPs-1

HPs-2

HPs-3

HPs-2

Scaling to large models

with the optimal base HPs

𝝁

Transfer

3.1 µ P of Diffusion Transformers

As mentioned in Section 2, the existing µ P results [72; 73] in Table 1 apply only to architectures expressible as a NE ⊗ OR glyph[latticetop] program. Therefore, it is crucial to determine whether the diffusion Transformers can be represented in this framework. The following theorem establishes that several prominent diffusion Transformers [52; 1; 8; 16] are compatible with the existing µ P results.

Theorem 3.1 ( µ P of diffusion Transformers, proof in Appendix C) . The forward passes of mainstream diffusion Transformers (U-ViT [1], DiT [52], Pixartα [8], and MMDiT [16]) can be represented within the NE ⊗ OR glyph[latticetop] Program. Therefore, their µ P matches the standard µ P presented in Table 1.

The proof of Theorem 3.1 relies on rewriting the forward pass of diffusion Transformers using the three operators defined in the NE ⊗ OR glyph[latticetop] Program [72]. Given that the representability of standard Transformer [61] within the NE ⊗ OR glyph[latticetop] Program has been established [67], we focus on demonstrating that the novel modules specific to 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 [52] in Figure 2a.

Theorem 3.1 ensures that existing practical methodologies developed for standard µ P apply directly to DiT, U-ViT, PixArtα , and MMDiT. Moreover, our analysis technique naturally extends to other variants of diffusion Transformers [49; 7; 19; 59; 48]. To the best of our knowledge, mainstream diffusion Transformer variants in use can be expressed within the NE ⊗ OR glyph[latticetop] Program.

3.2 Practical Methodology of µ P in Diffusion Transformers

Given the rigorous µ P result for diffusion Transformers established in Theorem 3.1, this section introduces how to apply µ P to diffusion Transformers in practice. In addition, our code is available at https://github.com/ML-GSAI/Scaling-Diffusion-Transformers-muP .

3.2.1 Implementation of µ P

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 µ P of diffusion Transformers following the standard procedure in [73]. 3 Specifically, we replace the vanilla width n in the abc -parameterization (see Section 2) with the width ratio n/n base when standard parameterization (SP) and µ P differ in Table 1, where n base is a fixed base width. This implementation remains consistent with Table 1 since n base is a constant. For example, for a hidden weight matrix W with width n and base HPs φ W , σ 2 W , η (maybe searched from proxy model), we set W = φ W ˜ W with initial ˜ W ∼ N (0 , σ 2 W /n ) , as in SP. In contrast, its learning rate is ηn base /n rather than η/n , where µ P differs from SP. Finally, we also follow the suggestion from [73] to initialize the output weights by zero (i.e., σ 2 out = 0 ).

3.2.2 Base Hyperparameter Transferability of µ P

Prior work has shown that µ P provides strong guidance for base HP transferability in image classification and language modeling tasks [73; 28; 64; 4; 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 1, 2, and 3 in Appendix D, respectively. We describe how to verify HP transfer across widths below; those for batch size and step are similar.

For simplicity, we describe how to evaluate the transferability of the base learning rate η across widths; similar procedures apply to other HPs, such as the multiplier φ . We define { n i } P i =1 and { η j } R j =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 µ P with base width n base , true width n i , and base learning rate η j . 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 4, we verify that diffusion Transformers under the µ P exhibit robust transferability of base HPs.

3.2.3 µ Transfer from Proxy Task to Target Task

Once base HP transferability is validated for diffusion Transformers, we can directly apply the µ Transfer algorithm [73] (see Algorithm 4 in Appendix D) to zero-shot transfer base HPs from a proxy task to a target task. Specifically, both the proxy model with width n proxy and the target model with width n target are parameterized by µ P using the same base width n base . We first search for an optimal combination of base HPs using various proxy models trained with a small batch size B proxy and a limited training steps T proxy . These optimal base HPs are then used to train the large target model with a larger batch size B target and longer training iteration T target . Experimental results in Section 5 show the superior performance of the µ Transfer in real-world vision generation.

4 Systematic Investigation for DiTµ P on ImageNet

In this section, we first empirically verify the base HP transferability of DiT [52] under µ P. We then µ Transfer the optimal learning rate to train DiT-XL-2µ P, which achieves 2.9 × faster convergence.

4.1 Basic Experimental Settings

To ensure a fair comparison between DiTµ P 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µ P 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; 1; 49].

Architecture of DiTµ P. The architecture of DiTµ P 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 n base in the µ P setup corresponds to 288, which uses 4 attention heads.

90

100

Width=144

100

Iteration=100K

Width=288

Batchsize=256

Iteration=150K

80

-50K

Width=576

-50K

Batchsize=512

FID-50K

90

Iteration=200K

80

70

80

60

60

70

40

60

-13

12

-11

-10

-9

-13

-12

11

-10

-9

13

12

11

-10

6-

log2(base Ir)

log2(base lr)

log2(base lr)

(a) HP transfer across widths.

(b) HP transfer across batch sizes. (c) HP transfer across iterations.

Training. We train DiT and DiTµ P using the AdamW [47]. Following the original DiT setup [52], we do not apply any learning rate schedule or weight decay, and constant learning rates are used in all experiments. The original DiT-XL-2 is trained with a learning rate 10 -4 and a batch size of 256.

Evaluation metrics. To comprehensively evaluate generation performance, we report FID [26], sFID [50], Inception Score [56], precision, and recall [34] on 50K generated samples without classifier-free guidance (cfg), as in Table 4 of [52]. In the main text, we present the FID results, while the remaining metrics are provided in Appendix E.2.

4.2 Base Hyperparameters Transferability of DiTµ P

In this section, we evaluate the HP transferability of DiTµ P across different widths, batch sizes, and training steps. We focus primarily on the base learning rate, as it has the most significant impact on performance [73; 28; 14; 43]. Similar results for other HP (weight multiplier) are provided in Appendix E.2 . Wesweep the base learning rate over the set { 2 -13 , 2 -12 , 2 -11 , 2 -10 , 2 -9 } across various widths, batch sizes, and training steps. FID-50K results are shown in Figure 3, and comprehensive results for other metrics are presented in Tables 8, 9, and 10 in Appendix E.2.

As presented in Figure 3, the optimal base learning rate 2 -10 generally transfers across scaling dimensions when some minimum width (e.g. 144), batch size (e.g., 256), and training steps (e.g., 150K) are satisfied, which verifies the base HP transferability of DiTµ P. Interestingly, we observe that neural networks with µ P tend to favor a large learning rate close to the maximum stable value (e.g., see Figure 3a). This aligns with empirical findings and theoretical insights reported for standard neural networks trained for multiple epochs [65; 44; 37; 38; 11], which suggest that larger learning rates introduce beneficial gradient noise to help guide optimization towards flatter minima that generalize better. Our empirical results suggest that the optimization landscape of µ P shares certain similarities with that of SP, offering a direction for future theoretical investigation.

4.3 Scaling Performance of DiTµ P

We µ Transfer the optimal base learning rate of 2 -10 to train DiT-XL-2µ P with a width of 1152. The batch size is set to 256, following [52]. We evaluate DiTµ P every 400K steps without using cfg. Training continues until DiTµ P 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 E.2.

As shown in Figure 4, DiT-XL-2µ P with the transferred base learning rate performs effectively, achieving 2.9 × faster convergence (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.

20.0

DiT-XL-2

DiT-XL-2-μP

17.5

-50K

15.0

12.5

2.9xfaster

10.0

1

3

5

7

Iteration/M

5 Large-Scale Text-to-Image Generation

In this section, we verify the efficiency of µ Transfer 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α -µ P on SA-1B

In this section, we perform µ Transfer experiments on the PixArtα [8], scaling from 0.04B to 0.61B parameters. Using the same pretraining setup, PixArtα -µ P with the transferred learning rate outperforms the original PixArtα , while incurring only 5 . 5% FLOPs of one full pretraining run.

5.1.1 Experimental Settings

To ensure the fairest possible comparison between PixArtα -µ P and the original PixArtα [8], we mainly adopt the original setup [8] and summarize the key components below.

Dataset. We use the SAM/SA-1B dataset [33], 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 [8]. All images are resized to a resolution of 256 × 256 during training.

Architecture of PixArtα -µ P models. The target PixArtα -µ P model adopts the same architecture as PixArtα (0.61B parameters). The proxy model also follows the same architecture, differing only in width. To construct the proxy PixArtα -µ P model, we fix the attention head dimension at 72 (as in PixArtα ) and reduce the number of heads from 16 to 4 (0.04B parameters). In the µ P framework, the base width n base (see Section 3.2.1) is set to the proxy width of 288.

Training. PixArtα is implemented using the official codebase and original configuration. We train the original PixArtα and the target PixArtα -µ P model for 30 epochs with a batch size of 176 × 32 (approximately 59K steps). 4 The small proxy PixArtα -µ P models are trained for 5 epochs with a batch size of 176 × 8 (approximately 39K steps). Notably, the model width, batch size, and training steps are all smaller than those used in the target pretraining setting.

Hyperparameter search. In this section, we focus solely on the base learning rate. We search over the candidate set { 2 -13 , 2 -12 , 2 -11 , 2 -10 , 2 -9 } 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

where R is the number of HP search trials, S is the number of parameters, E is the training epochs, B is the batch size and T is the training iteration. The ratio ≈ 5 . 5% here, as detailed in Appendix E.3.

Evaluation metrics. We evaluate text-to-image generation performance following standard practice [8; 16; 66; 1], including FID, CLIP Score, and GenEval [21]. Both FID and CLIP Score are computed on the aesthetic MJHQ-30K [39] and real MS-COCO-30K [45] datasets. MJHQ-30K contains 30K images generated by Midjourney, while MS-COCO-30K is a 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 a cfg of 4.5 to generate samples.

5.1.2 Experimental Results of PixArtα -µ P

We begin by conducting a base learning rate search using the PixArtα -µ P proxy models. The evaluation results for different base learning rates are summarized in Table 2. Details of GenEval results can be found in Table 15 in 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 -10 yields the best performance. Interestingly, this optimal learning rate matches that of DiTµ P (see Figure 3a). 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 then apply µ Transfer by transferring the searched optimal base learning rate of 2 -10 to train the target PixArtα . A comparison between PixArtα and PixArtα -µ P throughout training is provided in Table 3 (with the complete results shown in Table 16 in Appendix E.3). The results demonstrate that PixArtα -µ P consistently outperforms PixArtα across all evaluation metrics during training, supporting µ P as an efficient and robust approach for scaling diffusion Transformers. Furthermore, we observe that the benchmark performance of PixArtα degrades after 20 epochs, primarily due to overfitting. In contrast, PixArtα -µ P continues to improve, suggesting that µ P enhances the model's generalization ability, which offers an interesting direction for future theoretical investigation.

Specifically, we emphasize that our current experimental results validate that µ P works well without cfg (DiT) and with cfg (PixArtα ). Because cfg is an important factor during inference and affects the optimal hyperparameters, we strongly recommend practitioners align the evaluation of 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 [16] architecture from 0.18B to 18B. Under the same pretraining setup, MMDiTµ P-18B with the transferred base HPs outperforms the MMDiT-18B tuned by human algorithmic experts.

5.2.1 Experimental Settings

Dataset. We train models on an internally constructed dataset comprising 820M high-quality imagetext pairs. All images are resized to a resolution of 256 × 256 during training.

Baseline MMDiT-18B. The width and depth of MMDiT-18B are 5,120 and 16, respectively. The training objective combines a flow matching loss [46; 49] and a representation alignment (REPA) loss [76]. The model is optimized by AdamW [47], with a batch size of 4,096 and 200K training iterations. The learning rate schedule is constant with a warm-up duration. The HPs were tuned by algorithmic experts, requiring roughly 5 times the cost of full pretraining, as detailed in Appendix E.4.

(a) Base learning rate.

(b) Base gradient clip.

(c) Base REPA loss weight.

S

0.224

0.222

0.222

0.224

0.224

0.222

S

S

0.222

0.224

g

6

0.220

0.220

0.220

0.220

C

0.218

0.218

0.218

0.218

10-4

10-3

10-2

10-1

100

101

102

10-1

100

100

101

102

103

104

Baselearningrate

Base gradient clip

BaseREPAweight

Basewarmupiteration

Architecture of MMDiTµ P models. The target MMDiTµ P-18B model shares the same architecture as MMDiT-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 µ P setup, the base width n base is set to 1,920 (see Section 3.2.1 for definition of n base ).

Training of MMDiTµ P models. The training procedure for the target MMDiTµ P-18B is identical to 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 -5 to 2 . 5 × 10 -3 , gradient clipping from 0.01 to 100, weight of REPA loss from 0.1 to 1, and warm-up iteration from 1 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 HPs and different learning rates, as detailed in Appendix E.4.

As derived in Appendix E.4 based on Equation (1), the total tuning FLOPs under µ P amounts to 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 metrics. We use training loss to select base HPs on proxy tasks, as the dataset is passed through only once, and overfitting does not occur. We follow the standard practice in the µP literature [73; 28; 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-image generation, in addition to the standard GenEval benchmark [21], we also created an internal benchmark comprising 150 prompts to comprehensively evaluate text-image alignment. Each prompt includes an average of five binary alignment checks (yes or no), covering a 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 E.4.

5.2.2 Analyzing the Results of the Random Search

The visualization of the results of the base HP search is shown in Figure 5. First, the base learning rate has the most significant impact on training loss. In our case, we observed that the envelope near 2 . 5 × 10 -4 was stable and close to optimal, so we chose 2 . 5 × 10 -4 as the optimal value. Interestingly, unlike the DiT and PixArt setups, the optimal learning rate here is 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, µ P favors a larger clipping value, as aggressive clipping can undermine the maximal update property central to µ P. Third, the optimal weight for the REPA loss is determined to be around 0.5, consistent with the experience from existing works [76]. Finally, the warm-up iteration has a negligible impact on the training loss, so we adopt the default value of 1K in the pretraining of the target model.

5.2.3 µ Transfer Results of MMDiTµ P

The comparisons of training losses, GenEval results, and human evaluations are shown in Figure 6, Table 4, and Table 5, respectively. In addition, the complete comparison of training losses and detailed

visualizations are provided in Figure 8 and Figure 10 in Appendix E.4. As a result, MMDiTµ P-18B outperforms the baseline MMDiT-18B in 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 tuning becomes prohibitively expensive, µ P offers a scientifically grounded framework to unlock the full potential of large models .

6 Discussion

There are several promising research directions building on this work. First, the principles of µ P could be extended to diffusion models with more advanced and efficient architectures, such as linear transformers [66] and mixture-of-expert models [17]. Second, µ P can be applied to more sophisticated optimization algorithms, including the Muon optimizer [31] and the warmup-stable-decay learning schedule [28]. Third, while our experiments suggest that a proxy model width of 256-512 and a 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 a learningtheoretic framework to explain the optimization dynamics [58; 51], 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 to accelerate progress in generative modeling applications using diffusion Transformers, including text-to-image and video generation. However, improvements in scaling diffusion Transformers could also facilitate the creation of deepfakes for disinformation. Regarding the limitations of this work, although we demonstrate the efficiency of µ P 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.

7 Conclusion

In this paper, we extend µ P from standard Transformers to diffusion Transformers. By proving that mainstream diffusion Transformers share the same µ P 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, PixArtα , and MMDiT-18B, while requiring a small fraction of the usual tuning effort (e.g., 3% for MMDiT-18B). Our results establish µ P as a principled and efficient scaling strategy for diffusion Transformers.

Acknowledgments

This work was supported by the National Natural Science Foundation of China (No. 92470118); the Beijing Natural Science Foundation (No. L247030); the Beijing Major Science and Technology Project under Contract No. Z251100008425002; the Beijing Nova Program (No. 20230484416); the ByteDance Seed Research Fund; the Public Computing Cloud of Renmin University of China; and the fund for building world-class universities (disciplines) of Renmin 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α .

References

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.

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.04233 , 2024.

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://cdn. openai. com/papers/dall-e-3. pdf , 2(3):8, 2023.

Charlie Blake, Constantin Eichenberg, Josef Dean, Lukas Balles, Luke Y. Prince, Björn Deiseroth, Andrés Felipe Cruz-Salinas, Carlo Luschi, Samuel Weinbach, and Douglas Orr. uµ p: The unit-scaled maximal update parametrization. CoRR , abs/2407.17465, 2024.

Blake Bordelon, Hamza Tahir Chaudhry, and Cengiz Pehlevan. Infinite limits of multi-head transformer dynamics. In NeurIPS , 2024.

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. OpenAI Blog , 1:8, 2024.

Junsong Chen, Chongjian Ge, Enze Xie, Yue Wu, Lewei Yao, Xiaozhe Ren, Zhongdao Wang, Ping Luo, Huchuan Lu, and Zhenguo Li. PixartΣ : Weak-to-strong training of diffusion transformer for 4k text-to-image generation. In ECCV , volume 15090, pages 74-91, 2024.

Junsong Chen, Jincheng Yu, Chongjian Ge, Lewei Yao, Enze Xie, Zhongdao Wang, James T. Kwok, Ping Luo, Huchuan Lu, and Zhenguo Li. Pixartα : Fast training of diffusion transformer for photorealistic text-to-image synthesis. In ICLR , 2024.

Zixiang Chen, Yuan Cao, Quanquan Gu, and Tong Zhang. A generalized neural tangent kernel analysis for two-layer neural networks. Advances in Neural Information Processing Systems , 33:13363-13373, 2020.

cloneofsimo and Team fal. Introducing auraflow v0.1: An open exploration of large rectified flow models. https://blog.fal.ai/auraflow/ , 2024.

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.

Katherine Crowson, Stefan Andreas Baumann, Alex Birch, Tanishq Mathew Abraham, Daniel Z. Kaplan, and Enrico Shippole. Scalable high-resolution pixel-space image synthesis with hourglass diffusion transformers. In ICML , 2024.

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.

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 wafer-scale cluster. CoRR , abs/2304.03208, 2023.

Prafulla Dhariwal and Alexander Quinn Nichol. Diffusion models beat gans on image synthesis. In NeurIPS , pages 8780-8794, 2021.

Patrick Esser, Sumith Kulal, Andreas Blattmann, Rahim Entezari, Jonas Müller, Harry Saini, Yam Levi, Dominik Lorenz, Axel Sauer, Frederic Boesel, Dustin Podell, Tim Dockhorn, Zion English, and Robin Rombach. Scaling rectified flow transformers for high-resolution image synthesis. In ICML , 2024.

Zhengcong Fei, Mingyuan Fan, Changqian Yu, Debang Li, and Junshi Huang. Scaling diffusion transformers to 16 billion parameters. CoRR , abs/2407.11633, 2024.

Peng Gao, Le Zhuo, Dongyang Liu, Ruoyi Du, Xu Luo, Longtian Qiu, Yuhang Zhang, Chen Lin, Rongjie Huang, Shijie Geng, et al. Lumina-t2x: Transforming text into any modality, resolution, and duration via flow-based large diffusion transformers. arXiv preprint arXiv:2405.05945 , 2024.

Shanghua Gao, Pan Zhou, Ming-Ming Cheng, and Shuicheng Yan. Masked diffusion transformer is a strong image synthesizer. In ICCV , pages 23107-23116, 2023.

Yu Gao, Lixue Gong, Qiushan Guo, Xiaoxia Hou, Zhichao Lai, Fanshi Li, Liang Li, Xiaochen Lian, Chao Liao, Liyang Liu, et al. Seedream 3.0 technical report. arXiv preprint arXiv:2504.11346 , 2025.

Dhruba Ghosh, Hannaneh Hajishirzi, and Ludwig Schmidt. Geneval: An object-focused framework for evaluating text-to-image alignment. In NeurIPS , 2023.

Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783 , 2024.

Ishaan Gulrajani and Tatsunori B. Hashimoto. Likelihood-based diffusion language models. In NeurIPS , 2023.

Moritz Haas, Jin Xu, Volkan Cevher, and Leena Chennuru Vankadara. µ p 2 : Effective sharpness aware minimization requires layerwise perturbation scaling. In NeurIPS , 2024.

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.

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.

Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Comput. , 9(8): 1735-1780, 1997.

Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, 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.

Zhongzhan Huang, Pan Zhou, Shuicheng Yan, and Liang Lin. Scalelong: Towards more stable training of diffusion model via scaling network long skip connection. In NeurIPS , 2023.

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:135464-135625, 2024.

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 https://kellerjordan. github. io/posts/muon , 6, 2024.

Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun, editors, ICLR , 2015.

Alexander Kirillov, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloé Rolland, Laura Gustafson, Tete Xiao, Spencer Whitehead, Alexander C. Berg, Wan-Yen Lo, Piotr Dollár, and Ross B. Girshick. Segment anything. In ICCV , pages 3992-4003, 2023.

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.

Black Forest Labs. Flux. https://github.com/black-forest-labs/flux , 2024.

Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proc. IEEE , 86(11):2278-2324, 1998.

Jaehoon Lee, Samuel S. Schoenholz, Jeffrey Pennington, Ben Adlam, Lechao Xiao, Roman Novak, and Jascha Sohl-Dickstein. Finite versus infinite neural networks: an empirical study. In NeurIPS , 2020.

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.02218, 2020.

Daiqing Li, Aleks Kamko, Ehsan Akhgari, Ali Sabet, Linmiao Xu, and Suhail Doshi. Playground v2.5: Three insights towards enhancing aesthetic quality in text-to-image generation. CoRR , abs/2402.17245, 2024.

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.

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 diffusionbased text-to-image generation. In CVPR , pages 9400-9409, 2024.

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. FLM-101B: an open LLM and how to train it with $100k budget. CoRR , abs/2309.03852, 2023.

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.16645, 2024.

Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. In NeurIPS , pages 11669-11680, 2019.

Tsung-Yi Lin, Michael Maire, Serge J. Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C. Lawrence Zitnick. Microsoft COCO: common objects in context. In ECCV , volume 8693, pages 740-755, 2014.

Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In ICLR , 2023.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In ICLR . OpenReview.net, 2019.

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.

Nanye Ma, Mark Goldstein, Michael S. Albergo, Nicholas M. Boffi, Eric Vanden-Eijnden, and Saining Xie. Sit: Exploring flow and diffusion-based generative models with scalable interpolant transformers. In ECCV , volume 15135, pages 23-40, 2024.

Charlie Nash, Jacob Menick, Sander Dieleman, and Peter W. Battaglia. Generating images with sparse representations. In ICML , volume 139, pages 7958-7968. PMLR, 2021.

Atsushi Nitanda, Denny Wu, and Taiji Suzuki. Convex analysis of the mean field langevin dynamics. In AISTATS , volume 151, pages 9741-9757, 2022.

William Peebles and Saining Xie. Scalable diffusion models with transformers. In ICCV , pages 4172-4182. IEEE, 2023.

Dustin Podell, Zion English, Kyle Lacey, Andreas Blattmann, Tim Dockhorn, Jonas Müller, Joe Penna, and Robin Rombach. Sdxl: Improving latent diffusion models for high-resolution image synthesis. arXiv preprint arXiv:2307.01952 , 2023.

Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In MICCAI , volume 9351, pages 234-241, 2015.

Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily L. Denton, Seyed 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 NeurIPS , 2022.

Tim Salimans, Ian J. Goodfellow, Wojciech Zaremba, Vicki Cheung, Alec Radford, and Xi Chen. Improved techniques for training gans. In Advances in Neural Information Processing Systems , pages 2226-2234, 2016.

Kyle Sargent, Kyle Hsu, Justin Johnson, Li Fei-Fei, and Jiajun Wu. Flow to the mode: Modeseeking diffusion autoencoders for state-of-the-art image tokenization. CoRR , abs/2503.11056, 2025.

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 , 115(33):E7665-E7671, 2018.

Yuchuan Tian, Zhijun Tu, Hanting Chen, Jie Hu, Chao Xu, and Yunhe Wang. U-dits: Downsample tokens in u-shaped diffusion transformers. In NeurIPS , 2024.

Leena Chennuru Vankadara, Jin Xu, Moritz Haas, and Volkan Cevher. On feature learning in structured state space models. In NeurIPS , 2024.

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 Information Processing Systems , pages 5998-6008, 2017.

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.20314 , 2025.

Rongzhen Wang, Yan Zhang, Chenyu Zheng, Chongxuan Li, and Guoqiang Wu. A theory for conditional generative modeling on multiple data sources. arXiv preprint arXiv:2502.14583 , 2025.

Mitchell Wortsman, Peter J. 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 Kornblith. Smallscale proxies for large-scale transformer training instabilities. In ICLR , 2024.

Lechao Xiao. Rethinking conventional wisdom in machine learning: From generalization to scaling. CoRR , abs/2409.15156, 2024.

Enze Xie, 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.

Greg Yang. Wide feedforward or recurrent neural networks of any architecture are gaussian processes. In NeurIPS , pages 9947-9960, 2019.

Greg Yang. Tensor programs II: neural tangent kernel for any architecture. CoRR , abs/2006.14548, 2020.

Greg Yang. Tensor programs III: neural matrix laws. CoRR , abs/2009.10685, 2020.

Greg Yang and Edward J. Hu. Tensor programs IV: feature learning in infinite-width neural networks. In Marina Meila and Tong Zhang, editors, ICML , volume 139, pages 11727-11737. PMLR, 2021.

Greg Yang and Etai Littwin. Tensor programs iib: Architectural universality of neural tangent kernel training dynamics. In ICML , volume 139, pages 11762-11772, 2021.

Greg Yang and Etai Littwin. Tensor programs ivb: Adaptive optimization in the infinite-width limit. CoRR , abs/2308.01814, 2023.

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.

Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Tensor programs VI: feature learning in infinite depth neural networks. In ICLR , 2024.

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.

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.

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.

Chenyu Zheng, Guoqiang Wu, Fan Bao, Yue Cao, Chongxuan Li, and Jun Zhu. Revisiting discriminative vs. generative classifiers: Theory and implications. In International conference on machine learning , pages 42420-42477. PMLR, 2023.

Chenyu Zheng, Guoqiang Wu, and Chongxuan Li. Toward understanding generative data augmentation. Advances in neural information processing systems , 36:54046-54060, 2023.

Chenyu Zheng, Wei Huang, Rongzhen Wang, Guoqiang Wu, Jun Zhu, and Chongxuan Li. On mesa-optimization in autoregressively trained transformers: Emergence and capability. Advances in Neural Information Processing Systems , 37:49081-49129, 2024.

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 arXiv:2412.20404 , 2024.

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.

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 4 and 5.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in 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 or 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.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of 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 a separate "Limitations" section in 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 on 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 a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory assumptions and proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to 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 experimental results of the paper to the extent that it 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 5 and Appendix E. We also open-source our code at https://github.com/ML-GSAI/Scaling-DiffusionTransformers-muP .

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to 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 is a specific model and empirical evaluation, it 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 a 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 all submissions 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 to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a 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 case 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 to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We also open-source our code at https://github.com/ML-GSAI/ScalingDiffusion-Transformers-muP for DiT and PixArtα 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 NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an 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 run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on 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. If only a 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 in supplemental material (appended to 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, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We introduce the necessary configurations in Section 4, Section 5 and Appendix E. We also open-source our code at https://github.com/ML-GSAI/Scaling-DiffusionTransformers-muP for DiT and PixArtα experiments.

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

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [No]

Justification: Error bars are not 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 error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

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 calculating 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 or 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 a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments compute resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

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 as 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 it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

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 not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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 6.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

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 a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 a 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

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 a 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 a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: See details in Table 7 in Appendix E.

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 or dataset.

The authors should state which version of the asset is used and, if possible, include a 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.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

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 not available online, the authors are encouraged to reach out to the asset's creators.

13. New assets

Question: Are new 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 as 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. Crowdsourcing 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 as details about compensation (if any)?

Answer: [Yes]

Justification: See details in 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 in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or 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 not involve crowdsourcing nor research with human subjects.

Depending on 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.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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

Question: 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, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.

Answer: [NA]

Justification: We only use the LLMs to polish our writing.

Guidelines:

The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

Contents of Appendix

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 [16] 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. [40] empirically investigate the scaling behavior of various diffusion Transformers, finding that U-ViT [1] scales more efficiently than cross-attention-based DiT variants. Yin et 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 µ P in AIGC

Recently, µ P has been successfully applied to the pretraining of large language models (LLMs) [73; 14; 28; 43; 23; 42; 82], reducing HP tuning costs and stabilizing training. For diffusion Transformers, in addition to AuraFlow v0.1 [10], some works have also employed µ P. For example, Gulrajani and Hashimoto [23] searches for the base learning rate under µ P and transfers it to a 1B-parameter diffusion language model, while Sargent et al. [57] uses µ P to tune the HPs of a Transformer-based diffusion autoencoder, achieving state-of-the-art image tokenization across multiple compression 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 µ P of diffusion Transformers and systematically 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.

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 glyph[epsilon1] in the optimizer is assumed to be zero, which is very small in practice (e.g., 10 -8 ).

B.1 abc -Parameterization Setup

For completeness, we restate the abc -parameterization from Section 2, which underpins typical µ P theory. The abc -parameterization specifies the multiplier, initial variance, and learning rate for each weight [70; 73]. Specifically, let n denote the network width, we parameterize each weight as W = φ W n -a W ˜ W , where the trainable component ˜ W is initialized as ˜ W ij ∼ N (0 , σ 2 W n -2 b W ) . The learning rate for ˜ W is parameterized as η W n -c W . The width-independent quantities φ W , σ W , and η W are referred to as base HPs and can be transferred across different widths.

B.2 NE ⊗ OR glyph[latticetop] Program

µ P is originally proposed by Tensor Programs [67; 68; 71; 69; 70; 72; 74], 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 is the NE ⊗ OR glyph[latticetop] Program [72], which is formed by inductively generating scalars and vectors, starting from initial matrices/vectors/scalars (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 is the model width tending to be infinite.

Avg We can choose a existing vector x ∈ R n and append to the program a scalar

MatMul We can choose a matrix W ∈ R n × n and vector x ∈ R n existing in the program, and append to the program a vector

OuterNonlin For any integer k, l ≥ 0 , we can aggregate k existing vectors and l existing scalars to X ∈ R n × k and c ∈ R l , respectively. With an integer r ≥ 0 and a pseudo-Lipschitz function ψ : R k ( r +1)+ l → R (e.g., SiLU, GeLU), we append to the program a vector

where the r +1 is called the order of ψ .

While the above operators do not directly instruct to transform scalars into scalars, previous work [72] has proved that this can be done by combining the above operators. This can be formed as follows. Lemma B.1 (Scalars-to-scalars transformation is representable by the NE ⊗ OR glyph[latticetop] Program, Lemma 2.6.2 in [72]) . If ψ : R l → R is a pseudo-Lipschitz function and c are the collected scalars in a NE ⊗ OR glyph[latticetop] program, then ψ ( c ) ∈ R can be introduced as a new scalar in the program.

B.3 µ P of 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 a NE ⊗ OR glyph[latticetop] program (e.g., MLPs, CNNs [36], RNNs [27], standard Transformers [61]). In this setting, Yang and Littwin [72] prove the existence of a unique, optimal scaling rule-Maximal Update Parametrization ( µ P)-under which features in each layer evolve independently of width and at maximal strength. Otherwise, feature evolution in some layers will either explode or vanish in the infinite-width limit, rendering large-scale pretraining ineffective.

Concretely, µ P is implemented by analytically adjusting HPs with model width, typically involving the weight multiplier, initialization variance, and learning rate considered in abc -parameterization. The µ P formulation of any architecture whose forward pass can be expressed as a NE ⊗ OR glyph[latticetop] program follows the same scaling rules presented in Table 6 (equivalent to Table 1 in the main text). Specifically, it adjusts the abc -parameterization 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 d in = Θ(1) and d out = Θ( n ) ; hidden weights satisfy d in = Θ( n ) and d out = Θ( n ) ; output weights satisfy d in = Θ( n ) and d out = Θ(1) ; and scalar weights satisfy d in = Θ(1) and d out = Θ(1) , where n, d in , d out 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 ( µ P of representable architecture, Definition 2.9.12 in [72]) . If the forward pass of an architecture can be represented by a NE ⊗ OR glyph[latticetop] program, its µ P formulation follows the scaling rules in Table 6 (equivalent to Table 1 in the main text).

B.4 Extensions to Other HPs

The µ P principle can be extended to other HPs, such as common optimization-related HPs (e.g., warm-up iteration). We refer the reader to [73; 72] for a comprehensive discussion. We mainly focus on the HPs 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 (e.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 [32]), the weight decay value can not be 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 use 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 a vector a , we denote its i -th element as a i . For a matrix A , we use A k : , A : k and A ij to denote its k -th row, k -th column and ( i, j ) -th element, respectively. We define [ n ] = { 1 , 2 , . . . , n } .

For the reader's convenience, we restate Theorem 3.1 as follows.

Theorem C.1. The forward passes of mainstream diffusion Transformers (U-ViT [1], DiT [52], Pixartα [8], and MMDiT [16]) can be represented within the NE ⊗ OR glyph[latticetop] Program. Therefore, their µ P matches the standard µ P presented in Table 6 (equivalent to Table 1 in the main text).

Following the standard practice in the literature [70; 72; 74; 24] to simplify the derivation, we adopt the following assumption in all of our proofs.

Assumption C.1. We assume that constant dimensions (fixed during scaling the width of neural networks) are 1, 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) [67; 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 glyph[latticetop] Program.

The DiT architecture includes an embedder for the input latent x ∈ R , an embedder for the diffusion timestep t ∈ R , an embedder for the label y ∈ R , Transformer blocks with adaLN, and a final layer with adaLN. In the following, we use the NE ⊗ OR glyph[latticetop] Program to represent the forward computation of these modules in sequence.

Embedder for the input latent x . The embedder for the x is a one-layer CNN, which is known to be representable by the NE ⊗ OR glyph[latticetop] 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 ∈ R n , and the operation can be implemented by

Positional embedding. Given a number of patches (1 in our simplified case), the positional embedding is initialized with the sine-cosine version and fixed during the training. We use the x pos to denote the positional embedding, and adding it to the x embed can be written as

where ψ ([ x embed , x pos ] α : ) = x embed α + x pos α . Here, we reuse the notation of x embed of the final embedding for the input latent x without confusion.

Embedder for the diffusion timestep t . The embedder for the diffusion timestep t is a frequency embedding (one dimension here) followed by a two-layer MLP. First, for the frequency embedding, we can represent it as

Second, the MLP is known to be representable by the NE ⊗ OR glyph[latticetop] Program (e.g., Program 1 in [67]). We write it here for completeness. We use w (1) ∈ R n and W (2) ∈ R n × n to denote the weights in the first layer and the second layer of the MLP, respectively. We can derive

where ψ 1 ( w (1) α ; t freq ) = w (1) α t freq and ψ 2 ( t (1) α ) = SiLU( t (1) α ) .

Embedder for the label y . Denoting the weights of embedding as w embed ∈ R n , then its computation can be written directly as

where ψ ( w embed α ; y ) = w embed α y .

Merging the embeddings of t and y . In DiT, we sum the embeddings of t and y as the condition c for the following calculation, which can be written as where ψ ([ t embed , y embed ] : ) = t embed + y embed

Transformer block with adaLN. First, the adaLN block maps the condition c to the gate, shift, and scale values ( θ attn , β attn , γ attn , θ mlp , β mlp , γ mlp ) for attention layer and MLP layer in the Transformer block. We denote the parameters of the adaLN block by ( W θ -attn , W β -attn , W γ -attn , W θ -mlp , W β -mlp , W γ -mlp ) , each of them is in R n × n . Then the forward pass of adaLN block can be written as

where ψ is the SiLU activation.

Now, we denote the input feature of the Transformer block as x ∈ R n . A layer norm without learnable parameters first normalizes it, which is also known to be representable by the NE ⊗ OR glyph[latticetop] Program (e.g., Program 9 in [67]). We rewrite it in our simplified case.

After that, the normalized feature interacts with the condition information via the output of adaLN ( β attn , γ attn ), which can be represented as follows.

Now, the processed feature ˜ x norm is put to the multi-head attention layer, which is known to be representable by the NE ⊗ OR glyph[latticetop] Program (e.g., see Program 10 in [67]). We write it for completeness here. We denote W K , W Q , W V ∈ R n × n the attention layer's key, query, and value matrix, each row representing a head with dimension 1. The forward pass can be written as follows.

where ψ ([ k , q , v ] α : ) = k α q α v α .

Before discussing the MLP layer, we still have a residual connection and interaction from the gate value θ attn , which can be represented by

The remaining part is very similar to 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, we can write

For the interaction with scale and shift value from adaLN, we have

For the two-layer MLP in the Transformer block, we use W (1) , W (2) ∈ R n × n to denote the weights in its first layer and the second layer, respectively. We have

where the ψ is the GeLU activation function.

For the final residual connection, it can be represented by

Final layer with adaLN. The final layer includes the adaLN, layer norm, and a linear projection. We discuss them in sequence.

The adaLN block has parameters ( W β -final , W γ -final ) and receive the condition c . Its forward pass can be written as follows.

where ψ is the SiLU activation.

The layer norm layer normalizes the output of the transformer block h Tblock , which is the same as the above layer norm. We omit the derivation here for simplicity. We denote the output of the layer norm by z norm ∈ R n and the parameter of the linear projection by w final , then the remains calculation can be derived as

where ψ 1 ([ z norm , β final , γ final ] α : ) := γ final α z norm α + β final α and ψ 2 ([ w final , ˜ z norm ] α : ) = w final α ˜ z norm . Therefore, the forward pass from the inputs ( x, y, t ) to output z final can be represented by the NE ⊗ OR glyph[latticetop] Program, which means that The µ P of DiT is the same as the standard µ P in Table 6. The proof is completed.

C.2 Proof of PixArtα

Proof. Since most parts of the PixArtα are the same as those in DiT, we only discuss the different modules from DiT for simplicity. These modules are mainly induced from the condition process.

Embedder for the diffusion timestep t . The embedder for the diffusion timestep t not only outputs t embed but also gate, shift, and scale values based on t embed (named adaLN-single in [8]), which are shared by the following Transformer block. We denote the additional parameters by ( W θ -tattn , W β -tattn , W γ -tattn , W θ -tmlp , W β -tmlp , W γ -tmlp ) , each of them is in R n × n . Then its forward pass can be written as

where ψ is the SiLU activation.

Transformer block with gate-shift-scale table. PixArtα replaces the adaLN block in DiT block with the gate-shift-scale table, which parameters are denoted by θ xattn , β xattn , γ xattn , θ xmlp , β xmlp , γ xmlp ∈ R n . The gate-shift-scale table interacts with the shared embeddings of t in the following way.

where ψ ( a, b ) = a + b . The above outputted gate, shift, and scale values are used in the following attention layer and MLP layer, just like DiT.

One more difference between the Transformer block of PixArtα and DiT is the existence of the additional cross-attention layer, which integrates the output of the self-attention layer h and the text embedding y embed . We denote its parameters by W cK , W cQ , W cV , W proj ∈ R n × n It can be written as follows.

where ψ ([ k , q , v ] α : ) = k α q α v α . The following pre-LN MLP 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 a shift-scale table, just like what we 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α . We find that the only new operator in U-ViT is the long skip connection, which concatenates the current main branch h m ∈ R n and the long skip branch h s ∈ R n , and then performs a linear projection to reduce the dimension. We denote the parameters of the linear layer by W = [ W m , W s ] , where W m , W s ∈ R n × n , then the long skip connection can be represented by

where ψ ([ ˜ h m , ˜ h s ] α : ) = ˜ h m α + ˜ h s α . Because other operators in U-ViT (CNN embeddings/decoder, layer norm, attention layer, MLP 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/value parameters by W Kx , W Qx , W V x ∈ R n × n , the parameters of QK-norm by γ Kx , γ Qx ∈ R n , and the parameters of linear projection by W Ox ∈ R n × n . We can define the similar parameters for the text condition c as W Kc , W Qc , W V c , γ Kc , γ Qc , W Oc . The forward pass of the joint attention layer can be written as follows.

where ψ 1 ( a, b ) = ab/ | a | and ψ 2 ( a, b, c, d, e ) = abc + ade . The proof is completed.

Appendix D Additional Details for Section 3

In this section, we summarize the methodology for verifying base HP transferability across widths, batch sizes, and training steps in Algorithms 1, 2, and 3, respectively. We also present the µ Transfer algorithm [73] in Algorithm 4 for the completeness.

Algorithm 1 Validate the base HP Transferability of µ P across widths

Algorithm 2 Validate the base HP Transferability of µ P across batch sizes

,

Algorithm 3 Validate the base HP Transferability of µ P across training steps

1: Input: set of training steps { T i } P i =1 , set of base HP to validate { λ j } R j =1 , other fixed base HPs θ , batch size B , width n , base width n base

2: Output: whether model with µ P enjoy base HP transferability across training steps

3: for i = 1 to P do

4: for j = 1 to R do

) glyph[triangleright] actual HP for n -width model

glyph[triangleright]

5: λ ij ← µ P ( λ j , n base , n

6: θ i ← µ P ( θ , n base , n )

actual HP for

7: M ij ← train n -width model with b.s. B , iter. T i , HPs λ ij , θ i

8: s ij ← evaluate score for M ij

9: end for

10: ) glyph[triangleright] Obtain the optimal base HP for n -width model

λ ∗ i ← arg max λ j ( { s ij } R j =1

11: end for

12: return whether λ ∗ i s are approximately same

Algorithm 4 µ Transfer (Algorithm 1 in [73])

1: Input: base/proxy/target width n base / n proxy / n target , proxy/target batch size B proxy / B target , proxy/target training steps T proxy / T target , set of base HPs to search { θ i } R

j =1 2: Output: Trained target model M target 3: for j = 1 to R do 4: θ proxy ,j ← µ P ( θ j , n base , n proxy ) glyph[triangleright] actual HPs of proxy model 5: M j ← train n proxy -width model with batch size B proxy , training steps T proxy , HPs θ proxy ,j 6: s j ← evaluate score for M j 7: end for 8: θ ∗ ← arg max θ j ( { s j } R j =1 ) 9: θ target ← µ P ( θ ∗ , n base , n target ) glyph[triangleright] actual HPs of target model 10: M target ← train n target -width model with batch size B target , training steps T target , HPs θ target 11: return M target

Appendix E Additional Experimental Details and Results

In this section, we present the additional experimental details and results which are neglected by the main text.

E.1 Assets and Licenses

All used assets (datasets and codes) and their licenses are listed in Table 7.

n

-width model

E.2 Additional Details of DiT Experiments

E.2.1 Evaluation Implementation

Following the Table 4 in original DiT paper [52], 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, 9, and 10, respectively.

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 -2 , 2 0 , 2 2 , 2 4 , 2 6 } across various widths, batch sizes and training steps. FID-50K results are shown in Figure 7, and comprehensive results for other metrics are presented in Tables 11, 12, and 13, respectively. As presented in Figure 7, the optimal base output multiplier is approximately transferable across different widths, batch sizes, and training iterations.

Width=144

140

Batchsize=128

200

Iteration=100K

150

Width=288

Batchsize=256

Iteration=150K

FID-50K

Width=576

FID-50K

120

Batch size=512

FID-50K

150

Iteration=200K

100

100

80

100

50

60

50

-6

-4

-2

0

2

4

6

-6

-4

-2

0

2

4

6

-6

-4

-2

0

2

4

6

log2(base out multiplier)

log2(base out multiplier)

log2(base out multiplier)

(a) HP transfer across widths.

(b) HP transfer across batch sizes.

(c) HP transfer across iterations.

E.2.3 Additional Results of DiT-XL-2 Pretraining

Complete benchmark results of DiT-XL-2 and DiT-XL-2µ P throughout the training are provided in Table 14. DiT-XL-2µ P, using a base learning rate transferred from small proxy models, consistently outperforms the original DiT-XL-2 throughout the training process.

E.2.4 Computational Cost

It takes 104 (13 × 8) A100-80GB hours to train the DiTµ P with a width of 288, a batch size of 256, and a train iteration of 200K steps. The computational cost of other DiTµ P proxy models can be inferred based on this situation. It takes around 224 (32 × 7) A100-80GB days to reproduce the pretraining of DiT-XL-2µ P.

E.3 Additional Details of PixArtα Experiments

E.3.1 FLOPs ratio of HP Search on Proxy Models

Based on Equation (1) in the main text, we can derive that

Therefore, the HP tuning requires only 5 . 5% FLOPs of a single training run for PixArtα .

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. Overall, the base learning rate 2 -10 achieves the best GenEval performance.

E.3.3 Additional Results of PixArtα Pretraining

We present the complete comparison between PixArtα and PixArtα -µ P throughout training in Table 16. The results demonstrate that PixArtα -µ P consistently outperforms PixArtα across all evaluation metrics during training.

E.4 Additional Details of MMDiT Experiments

E.4.1 HPs Tuned by Human Experts

Algorithmic experts take roughly 5 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, a REPA loss weight of 0.5 and a 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 5 proxy training runs (100K steps each) using the searched optimal HPs and five different base learning rates, as detailed in Table 17.

E.4.3 FLOPs ratio of HP Search on Proxy Models

Based on Equation (1) in the main text, we can derive the ratio of tuning cost to single pretraining cost as

Besides, as algorithmic experts take roughly 5 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-18B and MMDiTµ P-18B in Figure 8. MMDiTµ P-18B achieves consistently lower training loss than the baseline after 15K steps, and the advantage is gradually increasing.

gap/10-4

2

1

baselinebetter

0

loss

μPbetter

Training l

-1

-2

-3

0

50

100

150

Iteration/K

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-image alignment, we created an internal benchmark comprising 150 prompts, where each prompt includes an average of five binary alignment checks (yes or no), covering a 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. We present two examples in Figure 9.

Alignment Evaluation

Image

Prompt

1.

Wine glass:

✅

2.

One (wine glass):

✅

3.

Bottles of wine:

✅

4.

Two (bottles of wine):

✅

5.

Cans of beer:

✅

6.

Three (cans of beer):

❌

One

wine

glass,

two

bottles

of

wine, three cans of beer

1.

Sydney Opera House:

✅

2.

Eiffel tower:

✅

3.

Right:

✅

4.

Mount Everest:

✅

5.

Rising above:

✅

The Sydney Opera House with the

Eiffel

tower

sitting

on the

right,

and Mount Everest rising above

E.4.6 Additional Results of Visualization Comparison

In this section, we provide visual comparisons between MMDiTµ P-18B and MMDiT-18B baseline in Figure 10.

Baseline

𝜇 P

STAATSTART

Baseline

𝜇 P

12926

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n**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.\n\nQUESTIONS:\nQ1: 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 5, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-20"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n# 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

comments: 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.

STRENGTHS AND WEAKNESSES:
Strengths:
1. 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. 
2. This work also spans various image generation models and evaluates the performance with standard metrics plus human expert checks. 
3. The writing is clear and easy to follow.

Weaknesses:
1. 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]
2. 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.

[1] Introducing AuraFlow v0.1, an Open Exploration of Large Rectified Flow Models. cloneofsimo, Team fal. https://blog.fal.ai/auraflow/

QUESTIONS:
See "weakness" section.
I am willing to change my score if my concerns are addressed.

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 4
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-02