Exploring Diffusion Transformer Designs via Grafting
Keshigeyan Chandrasegaran, Michael Poli, Daniel Fu, Dongjun Kim, Lea M. Hadzic, Manling Li, Agrim Gupta, Stefano Massaroli, Azalia Mirhoseini, Juan Carlos Niebles, Stefano Ermon, Fei-Fei Li
Model architecture design requires decisions such as selecting operators (e.g., attention, convolution) and configurations (e.g., depth, width). However, understanding quality-efficiency tradeoffs of such decisions typically requires costly pretraining, limiting architectural exploration.We present *grafting*, a simple approach for editing pretrained diffusion transformers (DiTs) to materialize new architectures by replacing expensive operators (e.g., self-attention, MLPs) with efficient alternatives. Informed by our analysis of activation behavior and attention locality, we construct a DiT-XL/2-based testbed to study grafting’s impact on model quality. Using this testbed, we develop a family of hybrid designs via grafting:replacing softmax attention with gated convolution, local, and linear attention; and MLPs with variable-width and convolutional variants. Notably, many hybrid designs achieve competitive quality (FID: 2.38–2.64 vs. 2.27 for DiT-XL/2).Next, we graft a text-to-image model (PixArt-$\Sigma$), achieving a 43% speedup with less than a 2% drop in GenEval score. Finally, we present a case study that restructures sequential computation into parallel in DiT-XL/2 via grafting, reducing model depth by 2$\times$, achieving better quality (FID: 2.84) than other models of comparable depth.Our work shows that new diffusion model designs can be explored via grafting pretrained DiTs, with edits ranging from operator replacement to restructuring computation.Pytorch code and grafted models are provided.