Exploring Diffusion Transformer Designs via Grafting

2025-06-05

AI Summary

To address the high computational cost of pretraining diffusion Transformers (DiTs), this work introduces “model grafting”—the first method enabling low-cost, editable architectural modifications to pretrained DiTs. Grounded in activation analysis and insights into attention locality, the approach supports fine-grained edits ranging from operator substitution (e.g., Softmax → gated convolutions or linear attention) to structural reorganization (e.g., serial → parallel blocks). The resulting hybrid architecture integrates gated convolutions, local/linear attention, variable-expansion-ratio MLPs, and convolutional MLPs. On DiT-XL/2, grafting consumes <2% of the original pretraining FLOPs while achieving FID scores of 2.38–2.64. Applied to PixArt-Sigma, it yields a 1.43× inference speedup with <2% degradation in GenEval. A depth-halved parallel variant achieves FID = 2.77—surpassing same-parameter baselines.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Neural Architectures and Foundation Models
• Natural Language Processing: Language Grounding & Multi-modal NLP

Application Category

• Web Mining and Content Analysis: Large pretrained models with web data
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG

Abstract

Designing model architectures requires decisions such as selecting operators (e.g., attention, convolution) and configurations (e.g., depth, width). However, evaluating the impact of these decisions on model quality requires costly pretraining, limiting architectural investigation. Inspired by how new software is built on existing code, we ask: can new architecture designs be studied using pretrained models? To this end, we present grafting, a simple approach for editing pretrained diffusion transformers (DiTs) to materialize new architectures under small compute budgets. Informed by our analysis of activation behavior and attention locality, we construct a testbed based on the DiT-XL/2 design to study the impact of grafting on model quality. Using this testbed, we develop a family of hybrid designs via grafting: replacing softmax attention with gated convolution, local attention, and linear attention, and replacing MLPs with variable expansion ratio and convolutional variants. Notably, many hybrid designs achieve good quality (FID: 2.38-2.64 vs. 2.27 for DiT-XL/2) using<2% pretraining compute. We then graft a text-to-image model (PixArt-Sigma), achieving a 1.43x speedup with less than a 2% drop in GenEval score. Finally, we present a case study that restructures DiT-XL/2 by converting every pair of sequential transformer blocks into parallel blocks via grafting. This reduces model depth by 2x and yields better quality (FID: 2.77) than other models of comparable depth. Together, we show that new diffusion model designs can be explored by grafting pretrained DiTs, with edits ranging from operator replacement to architecture restructuring. Code and grafted models: https://grafting.stanford.edu

Problem

• Evaluating architecture designs without costly pretraining
• Editing pretrained diffusion transformers for new architectures
• Improving model quality and speed via grafting techniques

Innovation

• Grafting pretrained diffusion transformers for new designs
• Replacing softmax attention with gated convolution variants
• Restructuring architecture via parallel block conversion