Title: Exploring Diffusion Transformer Designs via Grafting
Venue: NeurIPS 2025
Accepted: yes
Keywords: Diffusion Transformers

Model Grafting

Architectural Editing

Hybrid Models
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.43$\times$ 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 2$\times$ 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.

Exploring Diffusion Transformer Designs via Grafting

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Σ ), achieving a 1.43 × 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 2 × 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: grafting.stanford.edu.

1 Introduction

Model architecture design plays a central role in machine learning, alongside data, algorithms, compute, and benchmarks. It defines a learnable function and entails key decisions, including the choice of operators (e.g., attention, convolution) and configurations (e.g., model depth, width). Despite this, insight into architectures-what works and what doesn't-is difficult to obtain due to the prohibitive costs of training models from scratch, especially in today's foundation model era. As a result, studying new architectures remains a challenge, particularly for generative models. Much like how new software is built on existing code rather than written from scratch, can pretrained models serve as scaffolds for studying new architectures? In this work, we investigate architectural editing of pretrained models to study new architecture designs. We focus on diffusion transformers (DiTs), a class of generative transformers widely used for image and video generation [1, 2, 3].

1

(a) Model Architecture Design via Grafting

(b) Grafting Procedure

Regression

Limited data

End-to-end

Finetuning

(i) Stage 1: Activation Distillation

(ii) Stage 2: Lightweight Finetuning

Idea: Architecture editing

Limited

data

High Quality Model

(e) Depth       Width

Grafting

(c) Image Generation with Hybrid Architectures obtained via Grafting

A fantasy illustration of stoic

robot named Rob, a …

Mamba-2

FID=2.55

SWA

FID=2.62

Hyena-Y

FID=2.61

FID=2.64

r

=6

FID=2.38

r

=3

FID=2.53

A.R=164.6

FID=2.77

Full show of a Cairn Terrier

wearing a Yellow Plaid …

Ancient misty moor, forest,

overgrown, wide, highly ...

A young rabbit in adventure

clothes sets off for a new …

(d) Text-to-Image Generation with Hybrid Architectures obtained via Grafting

Train from

scratch

Grafting

S

A pretrained model implements a computational graph to perform tasks such as image or video generation. Given a new architectural idea and a pretrained model, we investigate whether the idea can be materialized by modifying its computational graph under small compute budgets. For example, one might hypothesize that a convolutional design could replace Multi-Head Attention (MHA) or Multi-Layer Perceptron (MLP) in a DiT. A simple way to materialize this idea is to replace MHA or MLP operators with a convolutional operator, while preserving model quality. This raises two key questions: (Q1) operator initialization : How to initialize a new operator before integrating it into the computational graph? (Q2) error accumulation : How to mitigate error propagation as multiple operators are integrated into the computational graph?

To address these questions, we present grafting 1 , a simple two-stage approach to architecture editing (Fig. 1). Grafting proceeds as follows: (i) activation distillation : This stage transfers the functionality of the original operator to the new one by distilling its activations using a regression objective. (ii) lightweight finetuning : This stage mitigates error propagation caused by integrating multiple new operators by finetuning using limited data. Architectural editing spans multiple strategies-adding, removing, and replacing [5, 6, 7] operators. We focus on operator replacement as the core strategy: swapping one operator for another. Other strategies can be viewed as special cases of replacement.

The space of architectural editing is vast, raising a practical question: what types of replacements should we study? We first establish a self-grafting baseline , where we replace all MHA and MLP operators in DiT-XL/2 with randomly initialized counterparts. Despite the scale of this intervention, our grafting procedure achieves near-baseline model quality using under 1% of pretraining compute. Building on this, we focus on replacing existing operators with efficient alternatives, aiming to reduce model FLOPs while preserving quality. We also explore replacements that increase model FLOPs to examine broader design choices. To study this systematically, we construct a testbed based on DiT-XL/2 and define a set of architectural edits to evaluate how different grafting schemes affect model quality. We organize our design space along four axes: (1) which operator to replace (e.g., MHA, MLP); (2) what to replace it with (e.g., convolutions); (3) how to select layers for replacement (e.g., all layers); and (4) replacement ratio (full vs. partial). We focus on replacing MHA and MLP operators, as they account for a large fraction of model FLOPs. Replacements for MHA and MLP operators are motivated by empirical findings and prior architectural designs: our locality analysis supports local operators for MHA, while for MLP, we adopt ideas from prior work [8, 9, 10].

We validate our grafting approach in increasingly challenging generative modeling setups:

Result I: Grafting yields hybrid architecture designs with good quality for class-conditional image generation (Sec. 4.2). We validate grafting using our testbed. For MHA (softmax attention), we explore several alternatives: local gated convolution (Hyena-SE, and our proposed Hyena-X/ Hyena-Y), local attention (sliding window), and linear attention (Mamba-2). For MLPs, alternatives include MLPs with variable expansion ratio (ratios=3, 6), and a convolutional variant (HyenaX). Interestingly, several interleaved hybrid architecture designs achieve FID scores between 2.38 and 2.64 (DiT-XL/2 256x256 baseline: 2.27), showing that grafting can construct good quality hybrids (Tab. 4) 2 . Grafting is simple and lightweight: each experiment completes in under 24 hours on 8×H100 GPUs, using less than 2% of pretraining compute.

Result II: We construct efficient hybrid architectures for high-resolution text-to-image (T2I) generation via grafting (Sec. 5). We validate grafting in a challenging, real-world setting: 2048 × 2048 resolution T2I generation using PixArtΣ (DiT) [11]. This setting reflects key challenges: it operates on long sequences (16,384 tokens), involves a multimodal setup with text conditioning, and lacks training data. We target MHA operators for grafting, as they account for over 62% of generation latency. Using 12k synthetic data, our grafted model achieves a 1.43× speedup with < 2% drop in GenEval score (47.78 vs. 49.75), showing that grafting scales to high-resolution, T2I generation.

Case Study: Converting model depth to width via grafting (Sec. 6). Motivated by our MLP grafting results, we rewire DiT-XL/2 by parallelizing every pair of transformer blocks, as modern GPUs favor parallel over sequential computation. This reduces model depth by 2 × (28 → 14). The grafted model achieves FID=2.77, outperforming other models of comparable depth. To our knowledge, this is the first attempt to convert sequential transformer blocks into parallel in pretrained DiTs, enabling architectures to be restructured.

2 Prerequisites

Diffusion models (DMs). DMs generate data samples by iteratively denoising random noise. This sampling process inversely mirrors the forward data corruption mechanism: z t = α t z + σ t ϵ where z = E ( x ) ∼ q ( z ) with E representing a pretrained encoder and x the data variable. The noise term ϵ follows the prior distribution N (0 , I ) . The transition kernel from time 0 to t is given by q t ( z t | z ) = N ( z t ; α t z , σ 2 t I ) . The choice of α t and σ t defines the diffusion variant, such as variancepreserving [12], or flow matching [13]. The training objective [12] is as follows:

where q ( t ) : time sampling distribution, and q ( z , c ) : joint distribution of latent z and condition c .

Diffusion transformers (DiTs). DiTs model the diffusion process by patchifying the input-noised images or latent-into a sequence of 1D tokens with positional embeddings. These tokens are processed through transformer blocks comprising self-attention, feedforward layers, residual connections, and normalization layers. DiTs also incorporate conditioning signals, such as noise timestep ( t ), class labels ( c ), or natural language prompts, enabling controllable generation [1, 14].

Datasets and evaluation metrics. For class-conditional image generation, we use ImageNet-1K [15]. We follow [1] and report Inception Score (IS), FID, sFID, Precision, and Recall using 50k generated samples (250 steps DDPM, cfg=1.5). For text-to-image generation, we report GenEval score [16].

3 Grafting Diffusion Transformers

3.1 Two-Stage Grafting Approach

Grafting aims to materialize new architectures by editing a pretrained model's computational graph. Given that we focus on replacing existing operators with alternatives, this raises two questions:

(Q1) How should a new operator be initialized before being integrated into the computational graph? Stage 1: Activation distillation. We cast initialization as a regression task. Operators in a DiT block process [ B,N,D ] inputs (batch, sequence, hidden) and output tensors of the same shape. Given a

pretrained operator f l ϕ at layer l , we learn a new operator g l θ that approximates f l ϕ [17]. Since DiT activations are continuous and smooth, this can be posed as a regression problem:

where q ( z , c ) is the joint distribution of latent representation z and condition c , q ( t ) is the time sampling distribution, and q t ( z t | z ) is the transition kernel from time 0 to t . L reg is a regression objective such as L 2 . In practice, a good initialization requires as few as 8k samples.

(Q2) How can we mitigate error propagation as multiple operators are integrated into the computational graph? Stage 2: Lightweight finetuning. As more operators are replaced, initialization errors propagate, leading to deviations from the pretrained model's behavior. We apply end-to-end finetuning with limited data to mitigate cumulative errors from stage 1. The fine-tuning objective is given in Equation 1. In practice, we find that competitive performance can be recovered using only 10% of the training data, even when replacing all MHA or MLP layers in DiT-XL/2.

3.2 Self-grafting Baseline

Prior to studying new architectural designs, we introduce self-grafting , a simple control setup where existing operators (e.g., MHA, MLP) are replaced with identical operators whose weights are randomly initialized. This preserves the computational graph's structure-operator types, receptive fields, and parameter count-while altering the computation performed. Self-grafting serves three purposes: (1) to assess the grafting procedure without architectural changes, (2) to provide a baseline for comparing replacements, and (3) to study factors affecting performance, such as data scale, regression objectives, and hyperparameters.

3.3 Activation Behavior Analysis and Self-grafting Results

We begin by analyzing the activation behavior of MHA and MLP operators across all layers in DiT-XL/2. In both cases, we observe large variance in activation values, particularly in deeper layers (Tab. 1 (i, ii)). When using regression-based distillation for Stage 1, these outliers affect optimization, particularly under the commonly used L 2 objective which penalizes all errors quadratically. This motivates a closer look at regression objectives. We study three regression objectives with different level of sensitivity to outliersL 2 , L 1 , and Huber [18]-using a self-grafting setup. We select five representative layers ( l = 1 , 8 , 17 , 27 , 28 ) for both MHA and MLP, spanning a range of activation

(i) Range of MHA and MLP Activation Values for all layers (DiT-XL/2)

Range of activation values

27 28

1

8

17

f

MLP

,

l

ϕ

x

7

14

21

28

0

5000

10000

15000

20000

27 28

1

8

17

DiT-XL/2 Layer

Activations

log(Counts)

f

MHA

,

l

ϕ

Activations

log(Counts)

DiT-XL/2 Layer

f

MLP

,

l

ϕ

f

MHA

,

l

ϕ

x

x

7

14

21

28

18.0

20.0

15.0

750

1500

2250

3000

1400

0069

-1400

2B

12700

(iii) MHA initialization

(ii) Activation Distribution for Selected Layers

(iv) MLP initialization

values. Each operator is trained with 8K ImageNet-1K [15] samples, for 200 epochs with batch size 64 and learning rate 1e -4 . We use δ = 1 . 0 for Huber objective. We then integrate the initialized operators into the pretrained DiT-XL/2 and evaluate quality without any finetuning.

High-quality initialization can be achieved by choosing operator-specific regression objectives.

As shown in Tab. 1 (iii,iv), the choice of the regression objective affects performance. For MHA, L 1 achieves the best FID (2.51), followed by Huber (2.55) and L 2 (2.58). For MLPs, L 2 performs best (2.33), while L 1 underperforms (2.83); notably, MLPs have 2× more parameters than MHA which explains its robustness to outliers [19]. This shows that high-quality initialization requires tailored, activation-aware strategies. Further, we evaluate validation loss on held-out samples. For MHA, L 1 achieves the lowest loss; for MLPs, L 2 achieves the lowest loss for all blocks (See Sec. C.3).

Full self-grafting with 10% data achieves near-baseline performance. We extend our study to replace all MHA and MLP operators in DiT-XL/2 under

the self-grafting setup and evaluate the effect of data on recovery (Tab. 2). For MHA, replacing all 28 layers without adaptation results in a noticeable performance drop, but Stage 2 (lightweight fine-tuning) is highly effective: using just 10% of the training data (128k samples), we achieve an FID of 2.53 vs. 2.27 for the baseline. Similarly, full MLP self-grafting with 10% data yields an FID of 2.54. We use batch size 256, learning rate 1e -4 , and 30k iterations. In both cases, the quality is within 0.3 FID of the baseline, showing that full self-grafting is feasible under modest data and compute budgets.

3.4 Locality Analysis of Self-attention

MHAscales quadratically with sequence length, making it a computational bottleneck. A natural idea is to replace it with local operators, such as convolution or local attention. However, this will fail if the model relies on long-range dependencies: for example, replacing all MHA operators in DiT-XL/2 with a sliding window attention degrades FID from 2.27 to 53.9. To guide grafting, we quantify MHA locality using a simple bandk metric. Given an attention matrix A ∈ R N × N and a band size of k , we define a bi-directional band indicator matrix B k ∈ R N × N as:

Then, locality within a band of size k is computed as:

We compute L k for all 28 MHA operators in DiT-XL/2 using 50-step DDIM sampling (250 ImageNet samples, sequence length 256, cfg scale 1.5), averaging across timesteps and samples. As shown in Fig. 2, MHA is largely local: on average, for k =32 , 15 out of 28 layers attend to more than 50% attention mass within the band. The first few layers ( l =1,2) display non-local attention patterns. Our analysis provides guidance for replacing MHA operators with efficient local operators.

4 Experiments I: Hybrid Architectures via Grafting

4.1 Testbed and Experiment setup

Building on our self-grafting results, we now ask: can we maintain model quality when existing operators are replaced with efficient alternatives? To investigate this, we study the grafting procedure along four design axes:

operator type to replace - MHA or MLP

replacement operator type - such as convolutions

layer selection strategy - replace operators in all layers or use heuristic-based selection

replacement ratio - full or partial

We construct a testbed to systematically evaluate how design decisions affect generative quality under grafting. We focus on efficient replacements that reduce FLOPs, but also include higher-FLOP variants to explore a broader range of architectural edits. We target MHA and MLP operators, which account for a significant portion of FLOPs in DiTs compared to other operators (e.g., normalization, activation, residuals). The rationale for replacing MHA or MLP operators is grounded in both empirical and architectural considerations: for MHA, our attention locality analysis (Fig. 2) motivates the use of local operators; for MLP, we leverage prior architecture ideas [8, 20, 21, 22, 9, 10]. Given a replacement operator, the decision to graft it to a model with L transformer layers spans a space of 2 L configurations. To make this tractable, we study two layer selection strategies: full (replace all operators) and interleaved (replace operators in a repeating pattern) strategies. The latter is inspired by striped transformer designs [23, 24, 25]. Our testbed is detailed in Tab. 3.

We introduce Hyena-X and Hyena-Y-two new efficient gated convolution operators designed as drop-in replacements for MHA. While our study includes several off-the-shelf efficient alternatives, we also contribute new operator designs motivated by our MHA locality analysis. This allows us to test novel architectural ideas via grafting, broadening our study. Both Hyena-X and Hyena-Y are local gated convolutions composed of dense, short causal depth-wise 1D convolutions. Fig. 3 (left) illustrates their structure. We also adapt Hyena-X as an MLP alternative by applying it along the channel dimension. Hyena-X and Hyena-Y scale linearly with sequence length, compared to the quadratic scaling of MHA. Operator details are provided in Sec. G. We provide FLOP calculation for both operators in Sec. H.1.

x

v

k

q

x

v

k

q

Hyena-Y

Short Explicit

Dense

Hadamard Gating

Convolution

Experiment setup. For our hybrid experiments, we mostly use the hyperparameters determined from our self-grafting studies (Sec. 3.2).

Stage 1: Operator initialization. For each new operator, we perform activation distillation using 8K ImageNet-1K samples. Each operator is trained for 200 epochs with a batch size of 64 and an initial learning rate of 1e -4 . We pre-extract and store all regression features. All operators can be initialized in parallel. Each operator's training completes in under 30 minutes on a single H100 GPU. Experiment details for stage 1 are included in Sec. C.1.

Stage 2: Lightweight finetuning. For all experiments in Table 3, we use 10% of the ImageNet-1K training data and train for 50K steps. We use a batch size of 256, linearly warming up the learning rate to 1e -4 over 1000 steps. Experiments typically complete in under 10 hours on 8 × H100 GPUs. For specific ablations on increasing data, such as those involving 20% data or 100K steps, runtimes extend up to 24 hours ( < 2% pretraining compute). We provide experiment details in Sec. C.1.

Table 3: Grafting testbed with configurations. This table defines the core design axes used in our study: operator type, efficient alternatives, layer selection strategy, and replacement ratio. For each alternative, we report configurations, including kernel size K , window size w , state size d s , expand factor E , and MLP expansion ratio r . The baseline DiT-XL/2 operator uses H =16 attention heads and MLP expansion ratio r =4 . Operator variants marked with (Ours) are proposed in this work.

4.2 Results and Insights

MHAresults. Replacing MHA operators in DiT-XL/2 via grafting yields strong quality-efficiency tradeoffs. We discuss our key insights below:

Surprising effectiveness of operators with smaller receptive fields under interleaved grafting . Our findings highlight that at 50% interleaved replacement, several alternatives-including SWA, Hyena-X/Y, and Mamba-2-consistently achieve FID scores within 0.5 of the baseline (2.27). The minimal FID drop observed especially with the SWA and Hyena variants, despite their limited receptive field ( K =4, w =4), aligns with our locality analysis (Section 3.4).

Replacement strategy: Interleaved vs. Full. Performance generally declines when increasing interleaved replacement from 50% to 75%. However, SWA remains effective at 75% interleaved replacement (FID=3.09). At 100% replacement, performance sharply degrades (all FIDs > 75). This trend aligns with our locality analysis, indicating that only a subset of layers are local and amenable to grafting.

Ablations on data scale and layer selection. We study two factors under 50% MHA replacement. (i) Increasing fine-tuning data from 10% to 20% improves FID across all variants (e.g., HyenaX: 2.74 → 2.61; SWA: 2.67 → 2.62, Mamba-

(a) Ablation 1: Data scaling (MHA)

Full self-grafting

Hyena-Y

SWA

Mamba-2

2.55

2.62

2.61

2.61

2.61

2.46

2.65

2.67

2.72

2.74

2.73

2.49

10% Data

20% Data

FID

4.00

3.18

3.02

2.74

FID

Interleaved

(50%)

Top-local

(50%)

Low-local

(50%)

Deep

(50%)

Full

(100%)

FID>140

2: 2.65 → 2.55) (Fig. 4 (a)). (ii) Under 50% replacement, we compare Hyena-X (interleaved) to three targeted heuristics: top-local (layers with highest bandk values), low-local (layers with lowest bandk values), and deep (last 14 layers). Interleaved yields the best FID (2.74), followed by top-local (3.02), low-local (3.18), and deep (4.00). These results confirm that interleaving is effective, and our bandk metric identifies layers that are more amenable to grafting (Fig. 4 (b)).

MLP results. Replacing MLP operators via grafting is effective. We discuss our key insights below:

Variable expansion ratio MLPs are effective under full replacement. MLP alternatives with expansion ratio r =3 and r =6 demonstrate good quality under all replacement ratios. Even under full (100%) replacement, both variants maintain good performance, with r =3 achieving FID=2.66. This highlights that MLP width is a robust dimension for grafting.

Convolutional alternatives. Hyena-X which combines dense and local channel mixing, performs competitively at 50% replacement (FID=2.63) but degrades at higher ratios, suggesting that such operators are only effective at moderate ratios.

Table 4: Generation quality and efficiency metrics for MHA and MLP grafting. We report quality (IS, FID, sFID, Precision, Recall) and efficiency ( ∆ FLOPs and ∆ Param) results. Baseline refers to DiT-XL/2. For each alternative, setups that maintain FID within 0.5 of the baseline and offer the largest FLOPs reduction (or smallest FLOPs increase) are highlighted . ✗ denotes setups with poor generation (FID > 50). ∆ FLOPs and ∆ Param denote the percentage change in operator FLOPs and parameters, respectively. For MHA, total cost is split into ∆ FLOPsop. (softmax attention, gating) and ∆ FLOPsft. (QKV/output projections, featurizers). We do not use this decomposition for MLP variants. Mamba-2 incurs higher ∆ FLOPsft. due to additional projections. FLOP expressions are provided in Sec. H.1. Key result: Many interleaved designs achieve good quality generation (FID within 0.5 of baseline). All experiments use 10% training data and < 1% pretraining compute.

5 Experiments II: Grafting Text-to-Image Diffusion Transformers

We apply grafting to a more challenging setting: high-resolution text-to-image generation with PixArtΣ [11]. This presents three challenges: (1) long sequences (16,384 tokens for 2048×2048 resolution), (2) a multimodal setup with text conditioning, and (3) lack of publicly available training data. These factors make PixArtΣ a representative setting for evaluating grafting under real-world constraints. PixArtΣ contains 28 transformer layers similar to DiT-XL/2.

Experiment setup. We replace MHA operators in PixArtΣ with Hyena-X via grafting, as MHA accounts for over 62% of generation latency. Hyena-X was chosen based on its good quality-efficiency tradeoff in the ImageNet-1K setup, achieving FID 2.61 with 20% data (see Fig. 4 (b)). Interleaved grafting is applied for layers 8, 10, 12, 14, 16, 18, and 20-27; empirically, we found that layers 20-27 can be replaced without significant quality drop. For grafting, we created a small, uncurated synthetic dataset of 12k image-text pairs. The text prompts for this dataset were sampled from the 30k publicly released evaluation set. Stage 1 (activation distillation): 8k uncurated synthetic image-text pairs are used to initialize Hyena-X blocks. We use the L1 regression objective, as we observe similar MHA activation behavior in PixArtΣ (Fig. E.1). Stage 2 (finetuning): We use LoRA (rank=64) [29] for finetuning. LoRA enables efficient finetuning by managing the high memory demands associated with long sequences (16,384 tokens). The full 12k synthetic dataset is used in this stage. We use 20 step DPM Solver [30] for generation. Experiment details are provided in Sec. E.2.

Results. The grafted model achieves a 1.43× speedup in wall-clock time, with a small drop in GenEval score (47.78 vs. 49.75). Attribute-specific metrics remain comparable, and qualitative samples show good alignment and quality. Some localized artifacts are observed in textured regions likely due to LoRA's adaptation capacity and low-quality synthetic data (see failure cases in Fig. E.3, E.4).

6 Case Study: Converting Model Depth to Width via Grafting

Can we rewire two sequential transformer blocks to run in parallel? Our MLP grafting results showed that MLPs are amenable to grafting, even at 100% replacement with an expansion ratio of r = 6 , demonstrating that wider computation within an operator is feasible. This success, combined with the fact that modern GPUs favor parallel over sequential computation, motivates a broader question: can we convert deeper, sequential DiT computations into wider, parallel ones via grafting while maintaining quality? To explore this, we rewire DiT-XL/2 by parallelizing every pair of sequential transformer blocks-each pair receives the same input, and their outputs are merged via a linear projection. This reduces model depth by 2 × (28 → 14) with a 6% increase in parameters.

Transformer

Layer

N

Transformer

Layer

N+1

Transformer

Layer

N

Transformer

Layer

N+1

(a) Sequential Transformer Blocks

(b) Rewiring Transformer Blocks in Parallel

via Grafting

Input

Output

Output

Input

Linear

Experiment Setup. The rewiring schematic is shown in Fig. 5. We use DiT-XL/2. Stage 1: Activation distillation. Each parallel pair was initialized via activation distillation using L1 regression. The weights for each block in the parallel pair were initialized from their corresponding pre-trained weights, rather than random initialization. Similar to our previous experiments, 8k ImageNet1K samples were used for this stage. Stage 2: Lightweight finetuning. Given the architectural restructuring, finetuning was performed using 25% of the training data. The learning rate was linearly warmed up to 1e-4 and halved at 75k and 150k iterations. Additional details can be found in Sec. D.

Results. The goal of this study is to evaluate generative quality (FID) vs. model depth. We report results in Tab. 6. To contextualize our findings, we compare against two categories: (i) DiTs trained from scratch at lower depth, and (ii) pruning methods [31, 32]. Our 14-layer grafted model achieves an FID of 2.77-surpassing DiT variants trained from scratch with similar or increased depth, including DiT-L/2 (depth 24, FID 3.73) and U-ViT-L (depth 21, FID 3.44). It also outperforms pruning baselines such as TinyDiT-D14 with masked knowledge distillation (depth 14, FID 2.86) and BK-SDM (depth 14, FID 7.43), though these baselines have fewer parameters (340M) compared to the grafted variants (712M).

Table 6: Generative quality vs. model depth. We report generative quality metrics (IS, FID, sFID, Precision, and Recall). A.R. (Aspect Ratio) is defined as model width divided by depth (e.g., 1152/14 = 82.3). For pruning and grafting setups, we report speedup with respect to DiT-XL/2 (depth=28). Off-the-shelf DiT-L/2, U-ViT-L, and DiT-B/2 scores, along with pruning baselines (BK-SDM, TinyDiT-D14, and TinyDiT-D14 w/ MKD), are sourced from [32]. MKD refers to Masked Knowledge Distillation, a recovery method used in [32]. ¶ Speedup is measured for a single forward pass on an Nvidia H100 (batch size=2). (details in Sec. D). Key result. Our grafted models achieve better generative quality at depth=14, surpassing baselines in FID, IS, Precision, and Recall.

7 Conclusion and Discussion

In this work, we introduced grafting , a simple approach to architecture editing. We constructed hybrid models by replacing self-attention and MLPs with efficient alternatives, achieving competitive quality (FID 2.38-2.64 vs. 2.27 baseline). We then applied grafting to a high-resolution text-to-image model (PixArtΣ ), yielding a 43% speedup with less than a 2% drop in GenEval score. We then used grafting to restructure DiT-XL/2, converting every pair of sequential transformer blocks into parallel, reducing model depth by half and yielding better quality (FID 2.77) among 14-layer DiTs. These results demonstrate grafting's utility in both short- and long-context settings (e.g., ImageNet-1K and PixArtΣ , respectively), and for architecture restructuring. Overall, grafting proves to be a lightweight approach for materializing diffusion transformer designs under small compute budgets.

Related work and discussion. Due to page limit, we discuss related work in Supp. A. Further, to demonstrate the generalization of grafting, we graft an LLM (Qwen3-4B [34])-a generative model (autoregressive), model architecture, and data modality distinct from diffusion-based image generation with DiTs (Supp. F). We discuss broader impact, limitations and applications in Supp. I.

Acknowledgments. We thank Liquid AI for sponsoring compute for this project. We also thank Armin W. Thomas, Garyk Brixi, Kyle Sargent, Karthik Dharmarajan, Stephen Tian, Cristobal Eyzaguirre, and Aryaman Arora for their feedback on the manuscript.

References

William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 4195-4205, 2023. 1, 3, 10

Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo, Li Jing, David Schnurr, Joe Taylor, Troy Luhman, Eric Luhman, Clarence Ng, Ricky Wang, and Aditya Ramesh. Video generation models as world simulators. 2024. 1

Agrim Gupta, Lijun Yu, Kihyuk Sohn, Xiuye Gu, Meera Hahn, Li Fei-Fei, Irfan Essa, Lu Jiang, and José Lezama. Photorealistic video generation with diffusion models. arXiv preprint arXiv:2312.06662 , 2023. 1

Eliezer Goldschmidt. Plant grafting: new mechanisms, evolutionary implications. Frontiers in Plant Science , 5:727, 12 2014. 2

Michael Zhang, Simran Arora, Rahul Chalamala, Benjamin Frederick Spector, Alan Wu, Krithik Ramesh, Aaryan Singhal, and Christopher Re. Lolcats: On low-rank linearizing of large language models. In The Thirteenth International Conference on Learning Representations , 2025. 2, 22, 29

Junxiong Wang, Daniele Paliotta, Avner May, Alexander M Rush, and Tri Dao. The mamba in the llama: Distilling and accelerating hybrid models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems , 2024. 2, 22

Aviv Bick, Kevin Li, Eric Xing, J Zico Kolter, and Albert Gu. Transformers to ssms: Distilling quadratic knowledge to subquadratic models. Advances in Neural Information Processing Systems , 37:31788-31812, 2024. 2, 22

Dan Fu, Simran Arora, Jessica Grogan, Isys Johnson, Evan Sabri Eyuboglu, Armin Thomas, Benjamin Spector, Michael Poli, Atri Rudra, and Christopher Ré. Monarch mixer: A simple sub-quadratic gemmbased architecture. Advances in Neural Information Processing Systems , 36:77546-77603, 2023. 2, 6

Aran Komatsuzaki, Joan Puigcerver, James Lee-Thorp, Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslie, Yi Tay, Mostafa Dehghani, and Neil Houlsby. Sparse upcycling: Training mixture-of-experts from dense checkpoints. In The Eleventh International Conference on Learning Representations , 2023. 2, 6

Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 , 2020. 2, 6

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, 2024. 3, 9

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems , 33:6840-6851, 2020. 3

Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747 , 2022. 3

Junsong Chen, YU Jincheng, GE Chongjian, Lewei Yao, Enze Xie, Zhongdao Wang, James Kwok, Ping Luo, Huchuan Lu, and Zhenguo Li. Pixartα : Fast training of diffusion transformer for photorealistic text-to-image synthesis. In The Twelfth International Conference on Learning Representations , 2023. 3

Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , pages 248-255, 2009. 3, 5

Dhruba Ghosh, Hannaneh Hajishirzi, and Ludwig Schmidt. Geneval: An object-focused framework for evaluating text-to-image alignment. Advances in Neural Information Processing Systems , 36:52132-52152, 2023. 3

Geoffrey Hinton, Oriol Vinyals, and Jeffrey Dean. Distilling the knowledge in a neural network. In NeurIPS Deep Learning and Representation Learning Workshop , 2015. 4

Peter J. Huber. Robust Estimation of a Location Parameter. The Annals of Mathematical Statistics , 35(1):73 - 101, 1964. 4

Peter L Bartlett, Philip M Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences , 117(48):30063-30070, 2020. 5

Sanghoon Kim, Dahyun Kim, Chanjun Park, Wonsung Lee, Wonho Song, Yunsu Kim, Hyeonwoo Kim, Yungi Kim, Hyeonju Lee, Jihoo Kim, et al. Solar 10.7 b: Scaling large language models with simple yet effective depth up-scaling. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 6: Industry Track) , pages 23-35, 2024. 6

Shikai Qiu, Andres Potapczynski, Marc Anton Finzi, Micah Goldblum, and Andrew Gordon Wilson. Compute better spent: Replacing dense layers with structured matrices. In International Conference on Machine Learning , pages 41698-41716. PMLR, 2024. 6

Jackson Petty, Sjoerd Steenkiste, Ishita Dasgupta, Fei Sha, Dan Garrette, and Tal Linzen. The impact of depth on compositional generalization in transformer language models. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers) , pages 7232-7245, 2024. 6

Michael Poli, Armin W Thomas, Eric Nguyen, Pragaash Ponnusamy, Björn Deiseroth, Kristian Kersting, Taiji Suzuki, Brian Hie, Stefano Ermon, Christopher Ré, et al. Mechanistic design and scaling of hybrid architectures. arXiv preprint arXiv:2403.17844 , 2024. 6

Jerome Ku, Eric Nguyen, David W Romero, Garyk Brixi, Brandon Yang, Anton Vorontsov, Ali Taghibakhshi, Amy X Lu, Dave P Burke, Greg Brockman, et al. Systems and algorithms for convolutional multi-hybrid language models at scale. arXiv preprint arXiv:2503.01868 , 2025. 6, 7

Garyk Brixi, Matthew G. Durrant, Jerome Ku, Michael Poli, Greg Brockman, Daniel Chang, Gabriel A. Gonzalez, Samuel H. King, David B. Li, Aditi T. Merchant, Mohsen Naghipourfar, Eric Nguyen, Chiara Ricci-Tam, David W. Romero, Gwanggyu Sun, Ali Taghibakshi, Anton Vorontsov, Brandon Yang, Myra Deng, Liv Gorton, Nam Nguyen, Nicholas K. Wang, Etowah Adams, Stephen A. Baccus, Steven Dillmann, Stefano Ermon, Daniel Guo, Rajesh Ilango, Ken Janik, Amy X. Lu, Reshma Mehta, Mohammad R.K. Mofrad, Madelena Y. Ng, Jaspreet Pannu, Christopher Ré, Jonathan C. Schmok, John St. John, Jeremy Sullivan, Kevin Zhu, Greg Zynda, Daniel Balsam, Patrick Collison, Anthony B. Costa, Tina HernandezBoussard, Eric Ho, Ming-Yu Liu, Thomas McGrath, Kimberly Powell, Dave P. Burke, Hani Goodarzi, Patrick D. Hsu, and Brian L. Hie. Genome modeling and design across all domains of life with evo 2. bioRxiv , 2025. 6

Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv:2004.05150 , 2020. 7

Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509 , 2019. 7

Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. In Forty-first International Conference on Machine Learning , 2024. 7

Edward J Hu, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. In International Conference on Learning Representations , 2022. 9

Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. Advances in Neural Information Processing Systems , 35:5775-5787, 2022. 9

Bo-Kyeong Kim, Hyoung-Kyu Song, Thibault Castells, and Shinkook Choi. Bk-sdm: A lightweight, fast, and cheap version of stable diffusion. In European Conference on Computer Vision , pages 381-399. Springer, 2024. 10

Gongfan Fang, Kunjun Li, Xinyin Ma, and Xinchao Wang. Tinyfusion: Diffusion transformers learned shallow. arXiv preprint arXiv:2412.01199 , 2024. 10

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 Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 22669-22679, 2023. 10

An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint arXiv:2505.09388 , 2025. 10, 29

Enze Xie, Junsong Chen, Junyu Chen, Han Cai, Haotian Tang, Yujun Lin, Zhekai Zhang, Muyang Li, Ligeng Zhu, Yao Lu, et al. Sana: Efficient high-resolution image synthesis with linear diffusion transformers. arXiv preprint arXiv:2410.10629 , 2024. 22

Jing Nathan Yan, Jiatao Gu, and Alexander M Rush. Diffusion models without attention. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 8239-8249, 2024. 22

Zhengcong Fei, Mingyuan Fan, Changqian Yu, and Junshi Huang. Scalable diffusion models with state space backbone. arXiv preprint arXiv:2402.05608 , 2024. 22

Vincent Tao Hu, Stefan Andreas Baumann, Ming Gui, Olga Grebenkova, Pingchuan Ma, Johannes Fischer, and Bjorn Ommer. Zigma: Zigzag mamba diffusion model. arXiv preprint arXiv:2403.13802 , 2024. 22

Yao Teng, Yue Wu, Han Shi, Xuefei Ning, Guohao Dai, Yu Wang, Zhenguo Li, and Xihui Liu. Dim: Diffusion mamba for efficient high-resolution image synthesis. arXiv preprint arXiv:2405.14224 , 2024. 22

Lianghui Zhu, Zilong Huang, Bencheng Liao, Jun Hao Liew, Hanshu Yan, Jiashi Feng, and Xinggang Wang. Dig: Scalable and efficient diffusion models with gated linear attention. arXiv preprint arXiv:2405.18428 , 2024. 22

Team Seawead, Ceyuan Yang, Zhijie Lin, Yang Zhao, Shanchuan Lin, Zhibei Ma, Haoyuan Guo, Hao Chen, Lu Qi, Sen Wang, Feng Cheng, Feilong Zuo Xuejiao Zeng, Ziyan Yang, Fangyuan Kong, Zhiwu Qing, Fei Xiao, Meng Wei, Tuyen Hoang, Siyu Zhang, Peihao Zhu, Qi Zhao, Jiangqiao Yan, Liangke Gui, Sheng Bi, Jiashi Li, Yuxi Ren, Rui Wang, Huixia Li, Xuefeng Xiao, Shu Liu, Feng Ling, Heng Zhang, Houmin Wei, Huafeng Kuang, Jerry Duncan, Junda Zhang, Junru Zheng, Li Sun, Manlin Zhang, Renfei Sun, Xiao Zhuang, Xiaojie Li, Xin Xia, Xuyan Chi, Yanghua Peng, Yuping Wang, Yuxuan Wang, Zhongkai Zhao, Zhuo Chen, Zuquan Song, Zhenheng Yang, Jiashi Feng, Jianchao Yang, and Lu Jiang. Seaweed-7b: Cost-effective training of video generation foundation model. 2025. 22

Zhen Xing, Qijun Feng, Haoran Chen, Qi Dai, Han Hu, Hang Xu, Zuxuan Wu, and Yu-Gang Jiang. A survey on video diffusion models. ACM Comput. Surv. , 57(2), November 2024. 22

Yu Gao, Jiancheng Huang, Xiaopeng Sun, Zequn Jie, Yujie Zhong, and Lin Ma. Matten: Video generation with mamba-attention. arXiv preprint arXiv:2405.03025 , 2024. 22

Hongjie Wang, Chih-Yao Ma, Yen-Cheng Liu, Ji Hou, Tao Xu, Jialiang Wang, Felix Juefei-Xu, Yaqiao Luo, Peizhao Zhang, Tingbo Hou, et al. Lingen: Towards high-resolution minute-length text-to-video generation with linear computational complexity. arXiv preprint arXiv:2412.09856 , 2024. 22

Zhengcong Fei, Mingyuan Fan, Changqian Yu, Debang Li, and Junshi Huang. Scaling diffusion transformers to 16 billion parameters. arXiv preprint arXiv:2407.11633 , 2024. 22

Armin W Thomas, Rom Parnichkun, Alexander Amini, Stefano Massaroli, and Michael Poli. Star: Synthesis of tailored architectures. In The Thirteenth International Conference on Learning Representations , 2025. 22

Song Bian, Minghao Yan, and Shivaram Venkataraman. Scaling inference-efficient language models. In Forty-second International Conference on Machine Learning . 22

Songhua Liu, Zhenxiong Tan, and Xinchao Wang. Clear: Conv-like linearization revs pre-trained diffusion transformers up. arXiv preprint arXiv:2412.16112 , 2024. 22

Songhua Liu, Weihao Yu, Zhenxiong Tan, and Xinchao Wang. Linfusion: 1 gpu, 1 minute, 16k image. arXiv preprint arXiv:2409.02097 , 2024. 22

Philipp Becker, Abhinav Mehrotra, Ruchika Chavhan, Malcolm Chadwick, Luca Morreale, Mehdi Noroozi, Alberto Gil Ramos, and Sourav Bhattacharya. Edit: Efficient diffusion transformers with linear compressed attention. arXiv preprint arXiv:2503.16726 , 2025. 22

Akhiad Bercovich, Mohammad Dabbah, Omri Puny, Ido Galil, Amnon Geifman, Yonatan Geifman, Izhak Golan, Ehud Karpas, Itay Levy, Zach Moshe, et al. Ffn fusion: Rethinking sequential computation in large language models. arXiv preprint arXiv:2503.18908 , 2025. 22

Chengchao Shen, Xinchao Wang, Youtan Yin, Jie Song, Sihui Luo, and Mingli Song. Progressive network grafting for few-shot knowledge distillation. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, pages 2541-2549, 2021. 22

Xingyi Yang, Daquan Zhou, Songhua Liu, Jingwen Ye, and Xinchao Wang. Deep model reassembly. Advances in neural information processing systems , 35:25739-25753, 2022. 22

Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence , volume 34, pages 7432-7439, 2020. 29

Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv preprint arXiv:1803.05457 , 2018. 29

Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics , pages 4791-4800, 2019. 29

Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM , 64(9):99-106, 2021. 29

Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. In International Conference on Learning Representations . 29

Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. In International Conference on Machine Learning , pages 28043-28078. PMLR, 2023. 29

Keshigeyan Chandrasegaran, Agrim Gupta, Lea M. Hadzic, Taran Kota, Jimming He, Cristobal Eyzaguirre, Zane Durante, Manling Li, Jiajun Wu, and Fei-Fei Li. Hourvideo: 1-hour video-language understanding. In Advances in Neural Information Processing Systems , volume 37, 2024. 31

Guo Chen, Zhiqi Li, Shihao Wang, Jindong Jiang, Yicheng Liu, Lidong Lu, De-An Huang, Wonmin Byeon, Matthieu Le, Tuomas Rintamaki, et al. Eagle 2.5: Boosting long-context post-training for frontier vision-language models. arXiv preprint arXiv:2504.15271 , 2025. 31

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 contributions are clearly described in the abstract and introduction.

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: We describe our limitations in Supp. I.

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: [NA]

Justification: N/A

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: All training details, hyperparameters, code, and instructions for reproducing our results are included in the Supplementary.

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: All code, and instructions for reproducing our results are included in the supplementary material including hyperparameters for each setup.

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: All training details and test details are provided in the Supplementary.

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: [Yes]

Justification: Standard deviation for selected experiments are provided in Supp. B.

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: We provide compute details in Supp J.

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: Yes, we abide by 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: We discuss broader impact in Supp. I.

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: [Yes]

Justification: We release only modified versions of existing pretrained models (e.g., DiT-XL/2) that are widely available and already publicly deployed. Our method-grafting-operates at the architectural level and does not introduce new model capabilities beyond those of the original models. To mitigate potential misuse, we follow responsible release practices: we (i) release models for research use only under appropriate licenses, (ii) provide usage guidance in our code repository, and (iii) avoid releasing models trained or finetuned on private or sensitive datasets. Since our focus is on architectural designs rather than model scaling, we assess the risk to be low and provide safeguards accordingly.

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: We use publicly released models (e.g., DiT-XL/2, PixArtΣ ) and datasets (e.g., ImageNet-1K) under their respective licenses. All assets are properly cited in the paper, and their licenses are respected.

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: [Yes]

Justification: We release new grafted model variants and accompanying code to support reproducibility. All assets are documented with model configurations, training details, and usage instructions.

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: [NA]

Justification: The paper does not involve crowdsourcing or research with human subjects.

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: This work does not involve human subjects or any procedures requiring IRB approval.

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: LLMs were not used as a component of the core methodology. Any use of LLMs was limited to minor writing support and did not influence the scientific contributions of the paper.

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.

Supplementary Material

Section A : Related Work

Section B : Standard Deviation of Experiments

Section C : Hybrid Architecture Experiments: Additional details

-Section C.1 : Experiment details and additional samples

-Section C.2 : Modulated Regression Targets for MHA

-Section C.3 : Validation Loss Curves for Self-grafting Experiments

Section D : Depth to Width Grafting Experiments: Additional details

Section E : Text-to-Image Generation Experiments: Additional details

-Section E.1 : MHA activation plots

-Section E.2 : Experiment details

-Section E.3 : Generated samples and failure cases

Section F: Grafting Autoregressive Large Language Models

Section G : Hyena-X and Hyena-Y operators: Additional details

Section H : FLOP calculation

Section I : Broader Impact, Limitations and Applications

Section J : Compute details

A Related Work

Diffusion model architectures. Recently, many architectural innovations have been proposed for diffusion models for image and video generation [35, 36, 37, 38, 39, 40, 41, 42, 43, 44]. Many recent works focus on improving the attention mechanism in diffusion models to enhance efficiency and scalability. One major direction is the use of modern linear attention variants, such as DiffuSSM [36], DiS [37], Zigma [38], DiM [39], and DIG [40]. Recently, text-to-image diffusion models such as SANA [35] have also adapted linear attention variants to support high-resolution generation. Another recent direction explores the mixture-of-experts (MoE) idea. DiT-MoE [45] introduces sparse diffusion transformers with shared expert routing and expert-level balance loss, enabling efficient scaling to 16.5B parameters while achieving competitive performance. We note that methods like STAR [46] have also successfully discovered architectures via evolutionary methods for autoregressive language modeling. Recently, Bian et al. [47] showed that pretraining wider and shallower autoregressive language models can yield efficiency gains while preserving accuracy. While effective, these approaches require training from scratch, making such studies expensive and inaccessible to practitioners. In contrast, grafting focuses on architecture editing of pretrained models to materialize new architectures under small compute budgets.

Architectural editing of pretrained generative models. Another line of work focuses on linearizing large language models by replacing softmax attention with efficient operators, such as linear attention [5, 6, 7]. Similar ideas have also been adopted for diffusion models in [48, 49, 50], though these works focus only on ultra-high-resolution settings. These prior efforts typically focus on replacing a single operator type (primarily attention) or are limited to specific application domains. Grafting presents a more general and comprehensive approach for architectural editing. It extends beyond single-operator replacement to enable modifying multiple operator types, exploring diverse architectural alternatives (e.g., both MHA and MLP replacements), and restructuring architectures (e.g., converting model depth to width). Recently, FFN Fusion [51] explored parallelizing transformer blocks in LLMs, aiming to reduce sequential computation. While our work focuses on diffusion-based generative modeling with transformers, prior work on image classification has explored model reuse and modification, such as few-shot knowledge distillation via progressive network grafting [52] and modular recombination of pretrained model components [53].

B Standard Deviation of Experiments

To compute variance associated with our reported results, we repeat two representative experiments-MHA (Hyena-Y) and MLP (width=6)-using three different random seeds ( seed = 0, 200, 300 ). We follow the exact grafting setup used in the main paper for these experiments. We report the mean and standard deviation of IS, FID, sFID, Precision and Recall in Tab. B.1. We observe that the standard deviations are within an acceptable range.

C Hybrid Architecture Experiments: Additional details

C.1 Experiment details and additional samples

We provide all hyperparameters used for the experiments in Tab. C.1. To ensure a fair comparison, we used identical hyperparameters across every hybrid experiment. We include additional qualitative samples generated using our hybrid architectures obtained via grafting in Fig. C.1.

C.2 Modulated Regression Targets for MHA

For Stage 1, we explored a modulation-aware regression variant for MHA experiments that incorporates the learned scalar ( gate_msa ) applied to the attention output. In the standard setup, we regress from input x to the raw output of the attention block y = MHA ( · ) . In the modulation-aware formulation, the target becomes y = gate_msa ⊙ MHA ( · ) . Tab. C.2 compares these two variants with L1 and L2 loss. Modulation-aware regression increases target scale, which adversely affects L2 loss performance due to its sensitivity to large values. L1 performs similarly in both settings. We adopted the standard (modulation-agnostic) formulation for all experiments for simplicity.

C.3 Validation Loss Curves for Self-grafting Experiments

To support the trends reported in the main paper (Sec.3.2, Table 2), we include validation loss curves in Fig.C.2 for five representative layers. Loss is computed using L 2 on a held-out set of 8k samples. For MHA layers (top row), L 1 consistently achieves lower validation loss in deeper layers, reflecting robustness to high activation variance. For MLP layers (bottom row), L 2 generalizes best across all layers. This contrast may be explained by parameter count: MLPs have roughly 2× more parameters than MHA layers, making them less sensitive to outliers and better suited to L 2 regression.

Hyena-Y

SWA

Mamba-2

r-3↓

Hyena-XMLP

Operator Initialization Results (Self-grafting/ MHA)

L

1

LHuber

,

δ

=1.0

L

2

0

100

200

0.01

1

2

0

100

200

1

10

20

0

100

200

1

10

40

0

100

200

10

20

100

0

100

200

5

10

100

val

.

loss

val

.

loss

val

.

loss

val

.

loss

val

.

loss

Layer 1

Layer 17

Layer 8

Layer 27

Layer 28

L

1

LHuber

,

δ

=1.0

L

2

0

100

200

0.001

1

2

0

100

200

1

10

20

0

100

200

0.1

10

20

0

100

200

0.1

10

20

0

100

200

0.1

10

100

val

.

loss

Operator Initialization Results (Self-grafting/ MLP)

val

.

loss

val

.

loss

val

.

loss

val

.

loss

Layer 1

Layer 17

Layer 8

Layer 27

Layer 28

D Depth to Width Grafting Experiments: Additional details

We provide all hyperparameters used for these experiments in Tab. D.1. We show additional samples in Fig. D.1.

Implementation details. The generative quality metrics (FID, IS, sFID, Precision, Recall) reported in Tab. 6 correspond to the exact implementation of rewiring presented in Fig. 5. For speedup measurements, a simple fused version of our rewired version was used. It is important to note that the reported speedup values are expected to decrease with large batch sizes, primarily due to the model parameter count (712M). Future work will focus on exploring hardware-aware / optimized implementations to achieve consistent speedup across a wider range of batch sizes.

E Text-to-Image Generation Experiments: Additional details

E.1 MHAactivation plots

We show activation distribution for five representative layers ( l = 15 , 17 , 19 , 21 , 23 ) in Fig. E.1.

E.2 Experiment details

We provide all hyperparameters used in our PixArtΣ grafting experiments in Tab. E.1.

(

Activation Values

DiT-XL/2 Layer

log(Counts)

x

PixArt

Σ

Activation Distribution for Selected Layers

f

MHA

l

(

x

)

Activation Values

log(Counts)

DiT-XL/2 Layer

18.0

19.0

14.2

9.5

9.0

4.8

4.5

OS.

23

-400

-

)

E.3 Generated samples and failure cases

We show additional high-resolution samples generated by the grafted PixArtΣ model in Fig. E.2, illustrating the model's ability to preserve generative quality across diverse prompts despite substantial architectural edits. Figure E.3 illustrates two types of failure modes observed in grafted PixArtΣ outputs. Each column pair shows the output of PixArtΣ (left) and the grafted model (right) for the same text prompt. In the top row, the original model generates images that are reasonably aligned with the prompts, while the grafted model fails to preserve this alignment-indicating limitations during the LoRA-based finetuning stage. In the bottom row, the synthetic supervision itself is of low quality, resulting in poor outputs from both the original and grafted models. To better understand this issue, Figure E.4 presents additional examples of low-quality synthetic data produced by PixArtΣ and used for grafting. These samples often exhibit artifacts and unrealistic physics. While synthetic data enables low-cost adaptation, these results highlight the importance of improved data curation and filtering to avoid propagating errors during the grafting process.

The image captures a scene of stark

beauty in the ...

One watercolor illustration of a

man on a white back …

The image captures a closeup view

of a vibrant red …

A very large cartoon panda wear a

baseball cap, in the style of …

In the professional oil painting

studio there is an …

Handmade vintage sock monkey

doll, attention to ...

Today marks the beginning of

autumn, summer has not …

Fresh coriander leaves plant on

white background, …

The image captures a moment of

tranquility in the deep blue …

Paint a captivating watercolor

clipart of a Lovebird, …

The image presents a delightful assortment of four donuts, each with its

own unique topping, set …

Playing classical guitar, Aesthetic flower, Full body shot, anime style,

vector …

A dog chasing its tail looks cute and funny, abstract photography,

stippling, UHD, high ...

This image captures a moment on a city street, frozen in black and white.

Dominating the left …

不

OA

F Grafting Autoregressive Large Language Models

To further validate the generality of grafting, we apply our two-stage procedure to a large language model (Qwen3-4B [34]) for autoregressive language modeling-a task, architecture, and modality distinct from diffusion-based image generation with DiTs.

Setup. We apply our two-stage grafting procedure to an autoregressive transformer, Qwen3-4B, for next-token prediction. This experiment replaces the Multi-Head Attention (MHA) operators with Sliding Window Attention (SWA) using a window size of 256. Stage 1 (activation distillation) and Stage 2 (lightweight fine-tuning) are performed identically to the diffusion experiments, with context lengths of 1024 and 8192 tokens respectively. We use 50k instruction-following examples from the Alpaca-cleaned dataset 3 following [5]. We evaluate on standard reasoning and commonsense benchmarks-PiQA [54], ARC-e [55], ARC-c [55], HellaSwag [56], Winogrande [57], and MMLU (5-shot) [58]-following the evaluation setup in [5].

Results. As shown in Table F.1, our grafted Qwen3-4B model achieves a 1 . 4 × decode throughput (8k context length) while maintaining average performance within 1% of the baseline. This result demonstrates that grafting extends beyond diffusion-based image generation, generalizing to new generative modeling task, model architecture, and data modality.

G Hyena-X and Hyena-Y operators: Additional details

Informed by our band-k analysis of MHA operators, we introduce a collection of efficient operators designed to exploit the locality in attention matrices. Given an input x ∈ R ℓ × d , a generic Hyena operator performs the following transformation:

where W,U,P,M ∈ R d × d are parametrized as dense or low-rank matrices, and T, H, K, G ∈ R ℓ × ℓ are Toeplitz matrices corresponding to convolutions with the filters h T , h H , h K , h G , respectively. In the original formulation [59], the filters h T , h H , h K are short and explicitly parametrized, whereas h G is implicitly parametrized.

We build on this formulation and propose Hyena-X and Hyena-Y , two Hyena operators designed for grafting. Hyena-X removes the implicit convolution entirely by setting G = I . In contrast, Hyena-Y introduces two changes: (i) it removes all three featurizer convolutions ( T , H , K ), and (ii) replaces

the implicit long convolution in G with a short, explicit convolution. This modified structure preserves local inductive bias while significantly reducing computational cost. An illustration is provided in main paper. These operators allows us to realize speedups across a range of inputs resolutions: both Hyena-X and Hyena-Y are faster than Mamba-2 operators on all input sequence lengths, including lower resolution regimes.

H FLOP calculation

Notations are provided in Tab. H.1.

H.1 MHA

Input projections (Q, K, V) : 6 LD 2

Softmax attention : 4 L 2 D +2 HL 2

Output projection : 2 LD 2

H.2 SWA

Input projections (Q, K, V) : 6 LD 2

Sliding window attention (Bidirectional) : 4 L (2 w +1) D +2 HL (2 w +1)

Output projection : 2 LD 2

H.3 Hyena-SE

Input projections : 6 LD

2

Featurizer : 3 LDK × 2

Inner filter convolution : LDK × 2

gates : LD × 2

Output projection : 2 LD 2

H.4 Hyena-X

Input projections : 6 LD 2

Featurizer : 3 LDK × 2

gates : LD × 2

Output projection : 2 LD

2

H.5 Hyena-Y

Input projections : 6 LD 2

Inner filter convolution : LDK × 2

gates : LD × 2

Output projection : 2 LD

2

H.6 Hyena-X (MLP)

Dense input projections : 6 LD 2 r

Featurizer : 3 LDK × 2

Gates :

LD × 2

Dense output projections : 2 LD 2 r

H.7 Mamba-2

Projections :

8 LD 2 E

Short convolution :

6 LDE

Featurization :

2 LDE (1 + 2 d state ) + 2 LDE

Associative scan :

2 LDEd state

Output layer :

2 LD 2 E

I Broader Impact, Limitations and Applications

Broader Impact. This work explores whether new architectural ideas can be materialized by editing pretrained diffusion transformers rather than training from scratch. By showing that architectural exploration is possible through grafting, our findings make model design more accessible and computeefficient, enabling broader participation in generative modeling research. Grafting facilitates practical improvements-such as efficient operator replacement and structural reorganization-offering a simple tool for exploring generative model architectures. However, as more efficient methods lower the barrier to developing powerful models, they also increase risks of misuse, including the generation of misleading or harmful content. We encourage responsible use, transparency, and content safeguards when applying these techniques.

Limitations. This work primarily focuses on architectural editing of pretrained Diffusion Transformers (DiTs), specifically targeting self-attention and MLP operators. Other architectural components, such as normalization layers and activation functions, will be explored in future work. We note that these are latent diffusion models, and grafting components in their corresponding V AEs remains an area for future study. The PixArtΣ setup used synthetic data for grafting, which may propagate artifacts and biases into the grafted models. While this work focuses on architectural editing, it remains an open question whether architectures that perform well under grafting also perform well when trained from scratch. Finally, grafting requires access to a pretrained model.

Applications and future work. Grafting holds promise for diverse applications where model customization and efficiency are important. This includes adapting models from low-resolution to high-resolution settings, extending capabilities from short-form video understanding/generation to long-form [60, 61], or improving user experience in interactive applications like image editing where even modest gains (e.g., 10% speedup) are highly valued. We hope that our testbed, insights, and results will encourage the community to actively explore new architecture designs. Code and grafted models: grafting.stanford.edu.

J Compute details

We report the total compute used to run all experiments in this work, including self-grafting baselines, hybrid architectures, ablations, PixArtΣ adaptation, and depth-to-width restructuring. All experiments were conducted using 8×H100 GPUs. Table J.1 provides a detailed breakdown. We also include time spent on feature extraction and tuning.

Grafting enables efficient architectural exploration without pretraining. All experiments combined required approximately 6,050 H100 GPU hours, which remains significantly lower than the cost of training large diffusion transformers from scratch. Each hybrid experiment for DiT-XL/2 completes within 12-24 hours on 8×H100 GPUs (both Stage 1 and Stage 2).

Table J.1: Estimated compute usage across all experiments. Total GPU hours account for training, sampling, and evaluation. Each experiment was conducted using 8×H100 unless otherwise specified.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper considers to explore diffusion transformer via grafting. Specifically, it attempts to replace expensive operators (e.g., self-attention, MLPs) with efficient alternatives and restructures sequential computation into parallel.  Grafting includes two stages: Activation distillation and Lightweight fine-tuning. The authors conducts experiments on class-conditioned DiT and text-to-image PixArt, verifying the effectiveness of the proposed approach.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\nThe exploration is interesting\n\nwell-organized and easy to follow\n\nThe results are reasonable\n\nWeaknesses\n\nSince obtaining a model with better performance is one of the most important goals in architecture design, this method lacks such a variant to illustrate the superiority of the method in this aspect. So far, it seems only be helpful in inference acceleration at a cost of performance. This method also may increase the parameters, for example in Table 6.\n\nConsidering the advantage of DiT in image and video generation,  conducting experiments only in image generation may be insufficient, it is highly recommended to employ grafting in video generation models. Similarly, can the proposed method enhance the model  performance in video generation.\n\nFor text-to-image generation, only using GenEval as the metric may not be sufficient.\n\nQUESTIONS:\n1、Why MHA and MLP prefer different regression loss?\n\n2、Actually, I may not agree with the author's opinion: 'We trade depth for width via grafting, achieving better quality than models of comparable depth.' Though the depth may be similar, e.g., DiT-B, TinyDiT-D14, and Grafting,  but the number of parameters is not. Is there a possibility that the better performance of Grafting is attributed to more model parameters?\n\n3 See Weaknesses above.\n\nLIMITATIONS:\nThe authors has discussed the limitations", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces grafting, a technique for architectural editing of pretrained diffusion transformers. it allows to replace expensive operators (such as self-attention or MLP layers) with more efficient alternatives in a two-stage process: (1) activation distillation and (2) lightweight fine-tuning. The method enables experimentation with hybrid architectures without full retraining, making it useful to explore new model designs with limited compute.\n\nSTRENGTHS AND WEAKNESSES:\n*Strengths:*\n\n1. The paper addresses a scalability challenge: architectural exploration in large generative models is expensive, which bottlenecks innovation.\n2. The authors test a variety of edit types (operator swaps, depth/width changes) on large-scale diffusion transformers, with clear benchmarks on FID, and compute costs. The depth/width example of grafting shows that parallel computation (shallower, wider models) by preserving the quality of the model is possible.\n\n*Weaknesses:*\n\n1. While finetuning is lightweight, recovery of good performance may still require significant data/compute for certain tasks or very deep editing. Whether it always outperforms a fresh, small retrain is not clear for all cases.\n2. Grafting is proposed as a method for exploring new architectures with limited compute. However, the authors have not trained their final proposed model from scratch for comparison against the grafted counterpart. As a result, it is unclear how much of the performance gain can be attributed to the distillation process itself. Additionally, unless I have misunderstood (and please feel free to clarify), it appears that the authors have not yet identified an architecture that outperforms their base model.\n\nQUESTIONS:\nPlease see weaknesses. I also have the following questions and suggestions:\n\n1. In scenarios where multiple types of operators are to be replaced, is there a recommended strategy? Should they be replaced all at once, or sequentially? Additionally, does the effectiveness of Stage 1 distillation degrade when replacing several components?\n\n2. Are the hyperparameters from the original (base) model typically retained, or is it generally necessary to re-tune them after grafting?\n\nSuggestion: For Table 4, I suggest that the authors consider adding the original FID values to plot (b), perhaps using dotted lines for visual reference. This would help contextualise the deltas and make the plot more informative.\n\nLIMITATIONS:\nPlease see weaknesses and questions.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis article discusses the author's exploration of Diffusion Transformers (DiT) structure design using the Grafting technique. Grafting is employed for structural design research and demonstrates significant efficiency advantages compared to previous retraining-based approaches. Additionally, the author uses this technique to conduct a series of interesting explorations on diffusion model structures, such as replacing softmax attention with gated convolution, linear attention, and other notably advantageous alternative modules. Although this is an inspiring piece of work, there are still many confusing aspects and areas for improvement.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**\n1. Grafting provides an efficient solution for exploring DiT structures and conducts explorations of certain value.  \n2. The article is well-written.  \n3. Experiments were conducted on DiT-XL/2 and PixArt-sigma, and the results are somewhat convincing.  \n\n**Weakness:**\n\nMajor Weakness:  \n\n1. Fine-tuning on pre-trained models leads to worse performance. The author employs various settings and uses the grafting method for fine-tuning, but almost all results in the tables show that the proposed methods underperform compared to the original pre-trained models. This diminishes the practical reference value of this work. While I understand the author's intention to explore interesting ideas, preserving generation quality is crucial for this task. The author should at least provide more analysis of the quality degradation.  \n\n2. Lack of comparison with existing work. Grafting is a technique that still requires minor fine-tuning. Since the acceleration effect is mentioned here, there are already numerous training-free acceleration techniques in the field that achieve DiT acceleration at almost no cost, such as the excellent works ToCa[1], TeaCache[2], PAB[3], and TaylorSeer[4]. When discussing acceleration, the author should compare with these methods or discuss their compatibility with such training-free approaches. \n\n\n3. Concerns about generalization. The author's experiments were only conducted on DiT-XL/2 and PixArt-sigma, which do not represent the cutting-edge level of the current field, thereby reducing the persuasiveness of the author's claims. Migrating the proposed methods to more open-source models with different architectures, such as the text-to-image model FLUX[5] and video generation models like HunyuanVideo[6] and Wan2.1[7], and discussing them alongside existing acceleration methods from other directions (e.g., those mentioned in point 2) to demonstrate the generalization of the proposed approach is strongly encouraged.   \n\n[1] [ICLR2025]Accelerating Diffusion Transformers with Token-wise Feature Caching: https://arxiv.org/abs/2410.05317  \n[2] [CVPR2025]Timestep Embedding Tells: It's Time to Cache for Video Diffusion Model: https://arxiv.org/abs/2411.19108  \n[3] [ICLR2025] Real-Time Video Generation with Pyramid Attention Broadcast: https://arxiv.org/abs/2408.12588  \n[4] [ICCV2025] From Reusing to Forecasting: Accelerating Diffusion Models with TaylorSeers: https://arxiv.org/abs/2503.06923  \n[5] FLUX: https://github.com/black-forest-labs/flux  \n[6] HunyuanVideo: A Systematic Framework For Large Video Generative Models: https://arxiv.org/abs/2412.03603  \n[7] Wan2.1: Wan: Open and Advanced Large-Scale Video Generative Models: https://arxiv.org/abs/2503.20314  \n\nMinor Weakness:  \n\nThe results in Table 6 should perhaps report the actual inference latency of the model. Additionally, the results in Table 6 are still not entirely satisfactory, and increasing the number of fine-tuning iterations might lead to more satisfactory outcomes.\n\nQUESTIONS:\nPlease see Weakness.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a method called Grafting, which is designed for editing a pretrained diffusion model to evaluate new network architectures with minimal computational cost. The approach consists of two stages: first, performing activation distillation for the new operator to be integrated; second, conducting lightweight finetuning after integration. The authors explore several architectural modifications and achieve comparable or even superior performance to the pretrained model by using this method, suggesting that new model designs can be efficiently explored by grafting pretrained model.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n- The proposed method is simple and efficient, involving only two low-cost stages: activation distillation and lightweight finetuning. The methodology is clearly explained and with code and grafted models provided, making it easy to follow.\n\n- The experiments in the paper are quite diverse and have good results. Most of the claims in the paper are supported by these experiments.\n- The method has a wide range of potential applications, which is also supported by the experiments.\n\n**Weaknesses**\n\n- The method appears to be quite general and not inherently specific to diffusion models. But the experiments are conducted solely on diffusion-based image generation using the DiT architecture. Therefore, the generalizability of the approach remains unproven.\n- One of the paper’s claims is that grafting can serve as an alternative method to exploring architectural effectiveness. However, it's unclear whether different operator might rely on their surrounding operators—since grafting requires freezing the context parameters. It also remains uncertain whether architectures that perform well under the grafting setting would yield similarly results when trained from scratch.\n\nQUESTIONS:\n- In Table 4(c), the interleaved strategy performs best, even outperforming the top-local strategy. Is there any explanation for this result? Does it suggest that the effectiveness of grafting is sensitive to the contextual operators surrounding the integrated operator?\n- The paper uses synthetic data for two-stage training and observes some localized artifacts. Would such artifacts also appear when training on real data? Are there methods to solve this problem?\n- Additional experiments on architectures beyond DiT or tasks beyond image generation would help demonstrate the broader applicability of the proposed method.\n- It would be helpful to include from-scratch experiments on small-scale architectures to verify whether the relative performance of different designs is consistent with that observed under the grafting setting. This would further support the paper’s claim.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n- This work proposes grafting, a simple approach to edit pre-trained diffusion transformers (DiTs) to materialize new architectures. They do so by replacing expensive operators like Self Attention and MLP layers with efficient alternatives like Hyena, MAMBA etc. \n- Authors perform an activation behavior and attention locality analysis on pre-trained DiT and study the impact of grafting on image quality using a DiT-XL/2 test bed.\n- Authors successfully replace the multi-head attention with gated convolution and linear attention. Similarly they replace the MLPs with convolutional variants. \n- Results show that many hybrid variants achieve competitive FID scores and on text-to-image models, grafting improves the speed of PixArt-$\\Sigma$ by 43%.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n- The activation behavior and locality analysis proposed by the authors are very interesting and can be useful diagnostic tools for understanding issues with pre-trained models. This kind of analysis can be very useful to the community. For example, from the locality analysis, one can replace the full attention in later parts of the model with local attention with little to no drop in quality. \n- The writing is crisp and the experimental section is exhaustive. The experiments clearly corroborate the claims made in the paper.\n\n**Weaknesses**\n- The result that MHA grafting above 50% results in significant drop in performance tells us that the operators identified by grafting can never fully replace MHA limiting the applicability of grafting. I wonder if training from scratch with operators from grafting, can further limit the expressiveness of the model. While I understand that requesting to train a model from scratch is not within the scope of this work but understanding how the expressiveness of the model changes with the new operators is essential to understand the effectiveness of grafting. I believe this experiment is important because pre-trained weights are not always available for most models and for this work to be universally applicable, the dynamics of grafting when pre-trained from scratch is very important.\n\nQUESTIONS:\n- Why was LoRA used for training PixArt-$\\Sigma$? What is exactly being done here? Why do you need LoRA when you are finetuning all the replaced operators in the 2nd stage?\n- What is the difference between Hyena SE and the two variants introduced in this work? Please add the differences between them in Fig. 3 Left. Why do you introduce causality in Hyena X and Y? (L205).\n- If the emphasis of this work is to make operators efficient and speed up training/inference, have the authors considered edge applications? Does grafting only improve speed of operators or does it (with minor tweaks) also reduce the memory of the full network?\n\nLIMITATIONS:\nyes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-05"}