Exploring Diffusion Transformer Designs via Grafting

• URL: http://arxiv.org/abs/2506.05340v2
• Date: Fri, 06 Jun 2025 17:59:47 GMT
• Title: Exploring Diffusion Transformer Designs via Grafting
• Authors: Keshigeyan Chandrasegaran, Michael Poli, Daniel Y. Fu, Dongjun Kim, Lea M. Hadzic, Manling Li, Agrim Gupta, Stefano Massaroli, Azalia Mirhoseini, Juan Carlos Niebles, Stefano Ermon, Li Fei-Fei,
• Abstract summary: We present grafting, a simple approach for editing pretrained diffusion transformers (DiTs) to materialize new architectures under small compute budgets.<n>We show that new diffusion model designs can be explored by grafting pretrained DiTs, with edits ranging from operator replacement to architecture restructuring.
• Score: 82.91123758506876
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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

Related papers

• Growing Transformers: Modular Composition and Layer-wise Expansion on a Frozen Substrate [0.0]
  This paper explores an alternative, constructive approach to model development, built upon the foundation of non-trainable, deterministic input embeddings.<n>We show that specialist models trained on disparate datasets can be merged into a single, more capable Mixture-of-Experts model.<n>We introduce a layer-wise constructive training methodology, where a deep Transformer is "grown" by progressively stacking and training one layer at a time.
  arXiv  Detail & Related papers  (2025-07-08T20:01:15Z)
• ArchComplete: Autoregressive 3D Architectural Design Generation with Hierarchical Diffusion-Based Upsampling [0.0]
  ArchComplete is a two-stage voxel-based 3D generative pipeline consisting of a vector-quantised model.<n>Key to our pipeline is (i) learning a contextually rich codebook of local patch embeddings, optimised alongside a 2.5D perceptual loss.<n> ArchComplete autoregressively generates models at the resolution of $643$ and progressively refines them up to $5123$, with voxel sizes as small as $ approx 9textcm$.
  arXiv  Detail & Related papers  (2024-12-23T20:13:27Z)
• STAR: Synthesis of Tailored Architectures [61.080157488857516]
  We propose a new approach for the synthesis of tailored architectures (STAR)<n>Our approach combines a novel search space based on the theory of linear input-varying systems, supporting a hierarchical numerical encoding into architecture genomes. STAR genomes are automatically refined and recombined with gradient-free, evolutionary algorithms to optimize for multiple model quality and efficiency metrics.<n>Using STAR, we optimize large populations of new architectures, leveraging diverse computational units and interconnection patterns, improving over highly-optimized Transformers and striped hybrid models on the frontier of quality, parameter size, and inference cache for autoregressive language modeling.
  arXiv  Detail & Related papers  (2024-11-26T18:42:42Z)
• Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models [92.36510016591782]
  We present a method that is able to distill a pretrained Transformer architecture into alternative architectures such as state space models (SSMs)<n>Our method, called MOHAWK, is able to distill a Mamba-2 variant based on the Phi-1.5 architecture using only 3B tokens and a hybrid version (Hybrid Phi-Mamba) using 5B tokens.<n>Despite using less than 1% of the training data typically used to train models from scratch, Phi-Mamba boasts substantially stronger performance compared to all past open-source non-Transformer models.
  arXiv  Detail & Related papers  (2024-08-19T17:48:11Z)
• Cross-Architecture Transfer Learning for Linear-Cost Inference Transformers [1.1499643186017316]
  We propose Cross-Architecture Transfer Learning (XATL) to improve efficiency of Transformer Language Models. Methodabbr significantly reduces the training time up to 2.5x times and converges to a better minimum with up to 2.6% stronger model on the LM benchmarks within the same compute budget.
  arXiv  Detail & Related papers  (2024-04-03T12:27:36Z)
• MatFormer: Nested Transformer for Elastic Inference [91.45687988953435]
  MatFormer is a novel Transformer architecture designed to provide elastic inference across diverse deployment constraints.<n>MatFormer achieves this by incorporating a nested Feed Forward Network (FFN) block structure within a standard Transformer model.<n>We show that a 850M decoder-only MatFormer language model (MatLM) allows us to extract multiple smaller models spanning from 582M to 850M parameters.
  arXiv  Detail & Related papers  (2023-10-11T17:57:14Z)
• The Languini Kitchen: Enabling Language Modelling Research at Different Scales of Compute [66.84421705029624]
  We introduce an experimental protocol that enables model comparisons based on equivalent compute, measured in accelerator hours. We pre-process an existing large, diverse, and high-quality dataset of books that surpasses existing academic benchmarks in quality, diversity, and document length. This work also provides two baseline models: a feed-forward model derived from the GPT-2 architecture and a recurrent model in the form of a novel LSTM with ten-fold throughput.
  arXiv  Detail & Related papers  (2023-09-20T10:31:17Z)
• Attention over pre-trained Sentence Embeddings for Long Document Classification [4.38566347001872]
  transformers are often limited to short sequences due to their quadratic attention complexity on the number of tokens. We suggest to take advantage of pre-trained sentence transformers to start from semantically meaningful embeddings of the individual sentences. We report the results obtained by this simple architecture on three standard document classification datasets.
  arXiv  Detail & Related papers  (2023-07-18T09:06:35Z)
• Full Stack Optimization of Transformer Inference: a Survey [58.55475772110702]
  Transformer models achieve superior accuracy across a wide range of applications. The amount of compute and bandwidth required for inference of recent Transformer models is growing at a significant rate. There has been an increased focus on making Transformer models more efficient.
  arXiv  Detail & Related papers  (2023-02-27T18:18:13Z)
• Learning Augmentation Distributions using Transformed Risk Minimization [47.236227685707526]
  We propose a new emphTransformed Risk Minimization (TRM) framework as an extension of classical risk minimization. As a key application, we focus on learning augmentations to improve classification performance with a given class of predictors.
  arXiv  Detail & Related papers  (2021-11-16T02:07:20Z)
• Predicting Attention Sparsity in Transformers [0.9786690381850356]
  We propose Sparsefinder, a model trained to identify the sparsity pattern of entmax attention before computing it. Our work provides a new angle to study model efficiency by doing extensive analysis of the tradeoff between the sparsity and recall of the predicted attention graph.
  arXiv  Detail & Related papers  (2021-09-24T20:51:21Z)
• Generating Diverse Structure for Image Inpainting With Hierarchical VQ-VAE [74.29384873537587]
  We propose a two-stage model for diverse inpainting, where the first stage generates multiple coarse results each of which has a different structure, and the second stage refines each coarse result separately by augmenting texture. Experimental results on CelebA-HQ, Places2, and ImageNet datasets show that our method not only enhances the diversity of the inpainting solutions but also improves the visual quality of the generated multiple images.
  arXiv  Detail & Related papers  (2021-03-18T05:10:49Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.