Exploring Diffusion Transformer Designs via Grafting
Authors: Keshigeyan Chandrasegaran
Michael Poli
Daniel Y Fu
Dongjun Kim
Lea M. Hadzic
Manling Li
Agrim Gupta
Stefano Massaroli
Azalia Mirhoseini
Juan Carlos Niebles
Stefano Ermon
Li Fei-Fei
Venue: NeurIPS 2025 oral ()
Keywords: Diffusion Transformers
Model Grafting
Architectural Editing
Hybrid Models
Primary area: deep_learning
TL;DR: We propose grafting, a simple approach to materialize new architectures by editing pretrained diffusion transformers. It enables architectural exploration under small compute budgets.

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
Keshigeyan Chandrasegaran +11.2 Michael Poli +1,2 Daniel Y. Fu 3,4 Dongjun Kim
Lea M. Hadzic Manling Li 1,5
Agrim Gupta " Stefano Massaroli 2.7
Azalia Mirhoseini Juan Carlos Niebles 1,8 Stefano Ermon It Li Fei-Fei
Stanford University
Liquid AI 3 Together Al 4 UC San Diego
5 Northwestern University
Google DeepMind 7 RIKEN Salesforce Research
grafting stanford.edu
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 testbed based on the DiT-XL/2 design to study the impact of grafting 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. 2.64 2.27 for DiT-XL/2) using- < 2% pretraining compute. We then graft text-to-image model (PixArt-E.), achieving 1 1.43x speedup with less than 2% drop in GenEval score. Finally, we present a case study that restructures DiT-X1/2 by converting every pair of sequential transformer blocks into parallel blocks via grafting. This reduces model depth by 2x and yields better quality (FID: 2.77) than other models of comparable depth. Together, we show that new diffusion model designs can be explored by grafting pretrained DiTs, with edits ranging from operator replacement architecture restructuring. Code and grafted models: grafting.stanford stanford. edu.
Introduction
Model architecture design plays central role in machine learning, alongside data, algorithms, compute, and benchmarks. It defines 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 works and what doesn't difficult obtain due to the prohibitive costs of training models from scratch, especially in today's foundation model era. As result, studying new architectures remains 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, investigate architectural editing of pretrained models to study new architecture designs. We focus on diffusion transformers (DiTs), class of generative transformers widely used for image and video generation [1, 2, 3].
Equal contribution. I Equal senior authorship. Part of this work done at Liquid Al. Correspondence (keahik, poli]@stanff addd.edu
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

(a) Model rechitecree Design via Grafting
(e) Depth - W Witt Grafting
(c) Image Generation with Hybrid A rchitectures obtained via Grafting
from scratch
I High Quality Model
Oyenz-Y
Random Init.
Graftin
SWA
r-
FIl
Pretrained Modal
High Quality Model
Mamaha 2
Elyena-X
(b) Grafting Procedure
Idea Architecture editing
Pretained Model
(d) Text-to-Image Generation with Hybrid Archituctures obtained via Grafting
(i Stage 1: Artivation Distillation
fantasy illasraton sooe nobliazned Rob,
Pull donw a Carm Temeer woating Yallow Plid...
Axciest misty mont, foont CoMMpONaL widk, Hilks
vourg is odbentarre now
Regression
Limited dela
(10 Stage 2: Lightwoight Finetuning
End to end Fineluning
I+I
Limited cota
High Quality Moče)
Figure Grafting overview. (a.b) Model architecture design via grafting. Studying new model architecture designs requires costly pretraining. Grafting materializes new architectures by editing pretrained models under small compute budgets (Sec. 3). (c) Class-conditonall image generation. Samples generated by hybrid architectures obtained via grafting (Sec. 4). (d) High-resolution text-lo- image generation, 2048 x 2048 samples generated using a grafted model (Sec. 5). (e) Depth - width case study. Samples generated using a model restructured via grafting (depth: 28 - 14) (Sec. 6).
pretrained model implements computational graph to perform tasks such as image or video generation. Given new architectural idea and 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 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 convolutional operator, while preserving model quality. This raises two key questions: (Q1) operator initialization: How to initialize a new operator before integrating into the computational graph? (Q2) error accumulation: How to mitigate error propagation multiple operators are integrated into the computational graph?
To address these questions, we present grafting',
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) lighteeight finetuning: This stage mitigates error propagation caused by integrating multiple new operators by finetuning using limited data. Architectural editing spans multiple
addiigg 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 grffftngg 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 FLOPPs to examine broader design choices. To study this systematically, we construct a testbed based on DiT-XL/2 and define 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, they account for 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].
Grafting draws inspiration from horticultural grafting, where efficient components (scions) are integrated into established systems (rootstock) to enhance functionality, such as yield and disease resistance [4].

---

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 testhed. For MHA (softmax attention), we explore several alternatives: local gated convolution (Hyena-SF, 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 (Hyena- X). 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 8xH100 GPUs, using less than 2% of pretraining compute.
Result II: We construct efficient hybrid architectures for high-resolution text-to-image (T21) gen- eration via grafting (Sec. 5). We validate grafting in challenging. real-world setting: x048x2002 resolution 121 generation using PixArt-E (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.43x 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 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 2x (28- 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: 2, = Q, + drE where z E(x) (z) with É representing pretrained encoder and x the data variable. The noise term follows the prior distribution N (0, I). The transition kernel from time 0 to & is given by qo (z. 2) = (Z, atz, I). The choice of a, and o, defines the diffusion variant, such as variance- preserving [12], or flow matching [13]. The training objective [12] is as follows:
LDM (Ø) = W(et
EZ t, 1)
(1)
where (t): time sampling distribution, and q(z, c): joint distribution of latent and condition
Diffusion transformers (DiTs). DiTs model the diffusion process by patchifying the input-noised images or iatent-into a sequence of ID tokens with positional embeddings. These tokens are pro- cessed through transformer blocks comprising self-attention, feedforward layers, residual connections, and normalization layers. DiTs also incorporate conditioning signals, such noise timestep (t), class labels (c), or natural language prompts, enabling controllable generation [1, 14].
Datasets and evaluation metrics. For class-conditions 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, report GenEval score [16].
3 Grafting Diffusson 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:
(QI) 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, DJ inputs (batch, sequence, hidden) and output tensors of the same shape. Given a
Strictly speaking, variable expansion ratio MLPs constitute a heterogeneous design rather than hybrid (i.e. they do not introduce operator class): use 'hybrid' throughout the paper for simplicity.
3

---

pretrained operator f/ at layer 1, we learn new operator g that approximates f' [17]. Since DiT activations are continuous and smooth, this can posed regression problem:
(( = (g(())((..))gg x) [eeg(g Ze, t, c), (z c))]
(2)
where q(z, c) is the joint distribution of latent representation and condition c, q(l) is the time sampling distribution, and qı(2,) (ı(zz/2) is the transition kernel from time 0 to L. Greg is a regression objective such as Lg. In practice, good initialization requires as few as 8k samples.
(Q2) How can we mitigate error propagation as multiple operators are integrated into the compu- tational 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 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- 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 baseline for comparing replacements, and (3) to study factors affecting performance, such 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 deeper layers (Tab. (i, ii)). When using regression based distillation for Stage 1. these outliers affect optimization, particularly under the commonly used L2 objective which penalizes all errors quadratically. This motivates closer look at regression objectives. We study three regression objectives with different self- self-grafting setup. We select five representative layers (1 1, 8, 17, 27, 28) for both MHA and MLP, spanning range of activation (lI) A ctivation Distribution for Selected Layers MILAJ
level of sensitivity to outliers L2, L1, and Huber
() Range of MHA and ML P. ctivation Values for all layers (DIT-XL/2)
MLPJ
17
750
3500
2150
14
2:
25
MLP
17
26
DIT Laper
DIT Layer
5000
10000
19305
70000
Honge activation
(iii) MHA initialization
(iv) MLP initialization
FID
sFID
Prec. f
Rec. 1
IS
FID 1
sFID
Prec.
Rec.
Baseline
278.20
2.27
4.60
0.83
0.57
Baseline
278.20
2.27
4.60
0.83
0.57
Random Init.
40.86
76,27
10,33
0.33
0.52
Random Init.
2.18
297.34
161.76
0.01
0.00
L2
269.31
2.58
5.75
0.82
0.58
L2
265.34
2.33
4.38
0.81
0.59
Huber (6=11))
271.30
2.55
5.44
0.82
0.57
Huber (6=1.0)
262.93
2.38
4.49
0.81
0.59
273.03
2.51
0.83
0.58
235.54
2.83
4.69
0.77
0.61
Table 1: Activation statistics and self-grafting (Stage 1) results (DiT-XL/2). (i) Activation ranges (max-min) across all 28 MHA and MLP operators, computed using 1,000 samples. Deeper layers exhibit higher variance in activation values. (ii) Activation distributions (log-scale histograms) for five selected layers (1, 8, 17, 27, 28), used in our initialization study. MLP layers show higher variance in activations than MHA, especially in deeper layers. (iii, iv) Stage results for these layers using L2, Huber, and LI regression. LI yields the best FID for MHA (2.51), while L2 performs best for MIP (2.33), which contains 2x more parameters than MHA (10.6M 5.3M). This study shows that high-quality initialization can be achieved by choosing operator-specific regression objectives.

---

values. Each operator is trained with 8K ImageNet-IIK IK [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-XI/22 and evaluate quality without any finetuning.
High-quality initialization
be achieved by choosing operator-specific regression objectives.
As shown in Tab. (ii,iv), the choice of the regression objective affects per- formance. For MHA, L1 achieves the best FID (2.51), followed by Huber (2.55) and L2 (2.58). For MLPs, L2 performs best (2.33), while L1 under-
Stage 1 Stage
1S
FID
sFID
Prec. 1
Rec.
Baseline
278.20
2.27
4.60
0.83
0.57
MHA (Full Self-grafting)
Random Init.
1.66
289.23
154,00
0.00
0.00
0.63%
117.68
16.78
13.69
0.60
0.61
performs (2.83); notably, MLPs have 2x more parameters than MHA which explains its robustness to outliers [19]. This shows that high-quality initializa- tion requires tailored, activation-aware strategies. Further, we evaluate vali- dation loss on held-out samples. For MHA, L1 achieves the lowest loss; for MLPs, L2 achieves the lowest loss for all blocks (See Sec. C.3).
0.63%
0.63%
148.56
11.26
11.10
0.66
0.60
0.63%
5.0%
270.39
2.70
5.46
0.81
0.57
0.63 %
10.0%
287.81
2.49
4.71
0.83
0.56
MLP (Full Self-
Random Init.
1.27
314.72
204.99
0.00
0.00
0.63% 10.0%
277.72
2,54
4.52
0.83
0.57
Table Full self-grafting (Stage results (DiT-XL/2).
We report results after replacing all 28 MHA and MLP operators using different amounts of training data. As we increase the training data from 0.63% (8k) to 10.0% (128k), FID improves consistently. Using only 10% of the training data, near-baseline performance is achieved: FID 2.49 for MHA and 2.54 for MLP.
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 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 gaff-gafting with 10% data yields an FID of 2.54. We batch size 256, learning rate 1e and 30k iterations. In both cases, the quality within 0.3 FID of the baseline, showing that full self self-grafting is feasible under modest data and compute budgets.
3.4 Locality Analysis of Self-attention
MHA scales quadratically with sequence length, mak- ing it a computational bottleneck. A natural idea is replace it with local operators, such convolu- tion or local attention. However, this will fail if the model relies on long-range dependencies: for exam- ple. replacing all MHA operators in DiT-XL/2 with sliding window attention degrades FID from 2.27 to 53.9. To guide grafting, we quantify MHA local- ity using simple band-k metric. Given an attention matrix A E R X and band size of k, we define a bi-directional band indicator matrix B₁ E N as:
k-8
k-16
k-32
k-64
14
0.90
0.45
0.23
DIT-XL/2 Layer
0.00
16 19
28
Figure 2: Locality of self-attention in DiT- XL/2. We plot La values for all 28 MHA operators, averaged over timesteps and sam- ples. At A=32, 15 out of 28 layers exhibit values exceeding 0.5, indicating that several MHA operators model local interactions.
(Bk)i = =00, otherwise
[1.
if
Then, locality within a band of size A is computed as:
L
1 N
Σ ABx
(3)
We compute L3 for all 28 MHA operators in DIT-XL/S 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 (1=1,2) display non-local attention patterns. Our analysis provides guidance for replacing MHA operators with efficient local operators.
5

---

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:
1. operator type to replace- MHA or MLP
2. replacement operator type such as convolutions
3. layer selection strategy replace operators in all layers or use ceuristic-based selection
4. 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-FI.OP variants to explore broader range of architectural edits. We target MHA and ML.P operators, which account for significant portion of FLOPs in DiTs compared to other operators (c.g., normalization, activation, residuals). The rationale for replacing MHA or MLP operators is grounded in both empirical and architectural considerations: for MHA, the use of local operators; for MLP, we leverage prior architecture ideas [8, 20, 21, 22, 9, 10]. Given replacement operator, the decision graft it to a model with L transformer layers spans space of 2 configurations. To make this tractable, we study two layer selection strategies: full (replace all operators) and interleaved (replace operators in repeating pattern) strategies. The lattrr is inspired by striped transformer designs [23, 24, 25]. Our testhed is detailed in Tab. 3.
attention locality analysis (Fig. 2) motivates
We introduce Hyena-X and Hyena-Y- Y-ona new efficient gated convolution operators designed drop-in replacements for MHA. While our study includes several off-the-sheds 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 ID convolutions. Fig. 3 (left) illustrates their structure. We also adapt Hyena-X an MLP alternative by applying along the channel dimen- sion. 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.
Hyena-X
Hyena-Y
Short Explicit
Convolation
Dense
Hadamard Gating
Figure Our proposed Hyena-X X and Hyena-Y efficient local gated convolution operators used as drop-in replacements for MHA.
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-IK IK samples. Each operator is trained for 200 epochs with batch size of 64 and an initial learning rate le-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 single H100 GPU. Experiment details for stage are included in Sec. C.I.
Stage 2: Lightweight finetuning. For all experiments in Table 3, we use 10% of the ImageNet- K training data and train for 50K steps. We use batch size of 256, linearly warming up the learning rate 1e-4 over 1000 steps. Experiments typically complete in under 10 hours on 8xH100 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.I.

---

Operator Type: Which operator types are we replacing?
MHA, MI.P
Efficient Alternative: Whut do we replace it with?
Convolutions: Hyena-SE [24], Hyena-X/ X/ Hyena-Y Y (Ours)
K=4, causal
MHA
Local Atlention: Sliding Window Attention (SWA) [26, 27]
4,4, bidirectional
Linear Attention: Mamba-2 [28]
d,=64, E=2
Variable expansion ratio
r=3,6
MLP
Hyena-X X (Ours)
r= K=4, causal, mix channels
Layer Selection: In which layers is the operator replaced?
Full
Replace the operator in all layers
Interleaved
Replace the operator in a repeating pattern (c.g., every 2 or 3 out of 4)
Replacement Ratio:
What percentage of operators are replaced?
50%, 75%, 100%
Table 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,, expand factor E, and MLP expansion ratio 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
(a) Ablation 1: Data scaling (MHA)
MHA results. Replacing MHA operators in DīT-XL/2 via grafting yields strong -fffity-fficciency tradeoffs. We discuss our key insights below:
Data
Full self-grafting
2.49
2.46
Hyema-SF
Surprising effectiveness of operators with smaller receptive fields under interleaved grafting. Our find- ings highlight that at 50% interleaved replacement, several ceaeativee-iicluuui SWA, Hyena-X/Y, and
Hyena-X
2.61
2.74
Hyena-Y
2.61
SWA
2.67
achieve FIDD 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).
Mamba-2
2.55
2.65
FID4
(b) Ablation 2: Layer selection heuristics (MHA Hyena-X)
FID4
4,00
Replacement strategy: Interleaved vs. Full. Per- formance generally declines when increasing inter- leaved replacement from 50% to 75%. However, SWA remains effective at 75% interleaved replace- ment (FID=3.09), At 100% replacement, perfor- mance 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.
3.18
3,02
2:744
FID-140
X
Interleaved Top-local Low-local (50%)
(50%)
(50%)
(50%)
(100%)
Figure 4: Ablation studies. (a) Data scale:
Ablations data scale and layer selection. We study two factors under 50% MHA replacement. (i) Increasing fine-tuning data from 10% 20% improves FID across all variants (c.g.. Hyena- X: 2.74 - 2.61; SWA: 2.67 + 2.62, Mamba- 2: 2.65 2.55) (Fig. 4 (a)). (ii) Under 50% replacement, we compare Hyena-X (interleaved) three targeted heuristies: top-local (layers with highest band-k values), low-local (layers with lowest band-k 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 band-k metric identifies layers that are more amenable to grafting (Fig. 4 (b)).
Increasing fine-tuning data from 10% to 20%
improves FID. (b) Layer selection strategies: Interleaved replacement outperforms other heuristics.

---

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 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.
Grafting is effective for constructing efficient hybrid architectures with good generative quality under small compute budgets. Interleaved designs are particularly effective.
Takeaway 1:
Ratio
IS
FID
sFID
Prec. 1
Rec.
AFLOPs... Sop.
AFLOPsJ.
AParam
Baseline
278.20
2.27
4.60
0.83
0.57
MHA Grafting
Random Init.
100%
1.66
289.23
154.00
0.00
0.00
Self-gratting
100%
287.81
2.49
4.71
0.83
0.56
50%
274.73
2.73
5.05
0.82
0.56
-49.52%
+0.13%
+0.22%
Hyena-SE (K==4
75%
231.15
3.62
6.04
0.81
0.54
-74.27%
+0.20%
+0.33%
100%
X
X
x
X
-99.03%
+0.26%
+0.43%
50%
273.30
2.74
5.03
0.83
0.56
-49 90%
+0.13%
+0.16%
Hyena-X (K=4)
75%
229.11
3.69
6.10
0.81
0.53
-74.85%
+0.20%
+0.24%
100%
X
X
-99.81%
+0.26%
+0.33%
273.37
2.72
5.02
0.83
0.55
-49.52%
0.00%
+0.05%
Hyena-Y
(K=4)
75%
228.99
3.66
5.95
0.81
0.53
-74.27%
0.00%
+0.08%
100%
X
X
X
X
-99.03%
0.00%
+0.1
50%
280.62
2.67
4.90
0.83
0.56
-48.24%
0.00%
0.00%
SWA
(w=4)
75%
249.99
3.09
5.54
0.82
0.55
-72.36%
0.00%
0.00%
100%
X
X
-96.48%
0.00%
0.00%
285.08
2.65
4.84
0.83
-37.59%
+77.89%
+28.02%
Mambba22 (d,=64, B=2)
75%
257.66
3.02
5.48
0.82
0.53
-56.38%
+1 +116.83%
+42.04%
100%%
X
X
x
X
-75.17%
+155.77%
+56.05%
MLP Grafting
Random Init.
100%
1.27
314.72
204.99
0.00
0.00
Self-grafting
100%
277.72
2.54
4.52
0.83
0.57
50%
272.14
2.53
4.51
0.83
0.57
-12.50%
0.00%
-12.50%
Exp. ratio +
(r=3)
75%
279.72
2.61
4.61
0.83
0.56
-18.75%
0.00%
-18.75%
100%
252.11
2.66
4.57
0.81
0.57
-25.00%
0.00%
-25.00%
50%
278.00
2.38
4.50
0.83
0,58
+25.00%
0.00%
+25.00%
Exp. ratio
(r=6)
277.94
2.37
4.48
0.82
+37.50%
0.00%
+37,50%
100%
276.86
2.42
4.50
0.82
0.58
+50.00%
0.00%
+50.00%
50%
265.60
2.64
4.66
0.83
0.56
+0.01%
0.00%
+0.02%
Hyena-X (r=2,K=4)
75%
226.13
3,26
4.79
0.81
0.55
+0.02%
0.00%
+0.03%
100%
X
X
X
+0.02%
0.00%
+0.03%
Table 4: Generation quality and efficiency metrics for MHA and MLP MLP grafting. We report quality (IS, FID, sFID, Precision, Recall) and efficiency (AFLOPs and AParam) results. Baseline refers to D/T-X1./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. X denotes setups with poor generation (FID 50). AFLOPs and Param denote the percentage change in operator FLOPs and parameters, respectively. For MHA, total cost is split into AFLOPSpp (softmax attention, gating) and AFLOPSp. (QKV/output projections, featurizers). We do not use this decomposition for MLP variants. Mamba-2 incurs higher AFLOPse. 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 < % pretraining compute.

---

Experiments II: Grafting Text-to-lmage Diffusion Transformers
We apply grafting to a more challenging setting: high-resolution text-to-image generation with PixArt-E [11]. This presents three challenges: (1) long sequences (16,384 tokens for 2048x2048 resolution), (2) multimodal setup with text conditioning, and (3) lack of publicly available training data. These factors make PixArt-E a representative setting for evaluating grafting under real-world constraints. PixArt-E contains 28 transformer layers similar to DiT-XII/2.
Experiment setup. We replace MHA operators Pix with Hyena-X via grafting, as MHA accounts for over 62% of generation latency. Hyena-X X was chosen based on its good quality-efficinnyy tradeoff in the ImageNet- IK setup, achieving FID 2.61 with 20% data (see Fig. (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 LI regression objective, activation behavior PixAnt- (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 (1 (16,384 tokens). The full 12k synthetic dataset is used in this stage. We use 20 step DPM Solver [30] for generation. Experiment details provided in Sec. E.2.
observe similar MHA
Results. The grafted model achieves a 1.43x speedup in wall-clock time, with a small drop GenEval score (47.78 vs. 49.75). Attribute-specin 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 low-quality synthetic data (see failure cases in Fig. E.3, E.4).
Takeaway 2:
We graft high-resolution ext-to--mage DiTs, constructing hybrid architectures with meaningful speedups and minimal quality drop.
Model
Ratio
Obj(1)
Obj(2)
Count
Colors
Pos
Color Attr.
Overall 1
Latency (ms)
Baseline
81.45
61.62
46.25
77.13
10.75
21.50
49.75
235.46
Hyena-X X
29%
80.31
59.34
49.69
68.62
11.50
18.75
48.04
194.95 (1.21x)
Hyena-X 50%
80.00
57.07
48.13
13
70.74
11.25
19.50
47.78
164.58 (1 (1.43x)
Table 5: GenEval results and latency for PixArt-E and the grafted variants. The 50% grafted model achieves 1.43x speedup while retaining strong text-image alignment (GenEval overall score: 47.78 vs. 49.75). Attribute-specific scores remain comparable across models. Latency is measured for single forward pass on an Nvidia H100 (batch size=2).
Case Study: Converting Model Depth to Width via Grafting
Can we rewire two sequential transformer blocks to in parallel? Our MLP grafting results showed that MLPs are amenable to grafting, even at 100% replacement with expansion ratio of r 6, demon- strating that wider computation within an operator is feasible. This success, combined with the fact that modern GPUs favor parallel over sequential compu- tation, motivates broader question: can we convert deeper, sequential DiT computations into wider, par- allel ones via grafting while maintaining quality? To explore this, we rewire DiT-XL/2 by parallelizing ev- ery 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 2x (28 + 14) with 6% increase in parameters.
(a) Sequential Transformer Blocks
Transfomer Layer N
Transformer Output
Layer N+1
(b) Rewiring Transformer Blocks in Parallel via Grafting
Transformer Layer
Input
Linerr
Output
Transformer
Layer N
Figure 5: Convert model depth -> width via grafting: (a) Two sequential transformer lay- ers. (b) Rewiring in parallel via grafting (includes skip connections).

---

Experiment Setup. The rewiring schematic is shown in Fig. 5. We use DiT-XL/2. Stage 1: Activation distillatton. Each parallel pair was initialized via activation distillation using LI 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 ImageNet- 1K 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 le-4 and halved at 75k and 10kk 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 lower depth, and (ii) pruning methods [31. 32]. Our 14-layer grafted model achieves an FID of 3777-surpssing DiT variants trained from scratch with similar or increased depth, including DiT-L/2 (depth 24, FIID 3.73) and U-VIT-L (depth 21, FID 3.44). It also outperforms pruning baselines such as TinyDiT-D144 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).
Takeaway& 3: Grafting enables architectural restrcttring at the transformer block level, allowing model depth to be traded for width.
Method
Depth
A.R
Iters
IS f
FID L
sFID
Prec. t
Recall
Speedup t Params
DIT-L/2 [1]
24
42.7
1,000K
196.26
3.73
4.62
0.82
0.54
458M
[33)
48.8
300K
221.29
3.44
6.58
0.83
0.52
287M
DiT-B/2 [1]
12
64.0
1000K
119.63
10.12
5.39
0.73
0.55
130M
BK-SDM [31]
14
82.3
100K
141.18
7.43
6.09
0.75
0.55
2x
340M
TinDDITDD14 [32]
14
82.3
500K
198.85
3,92
5.69
0.78
0.58
340M
TinDDiT-D14 w/ MKD [32]
14
82.3
500K
234.50
2.86
4.75
0.82
0.55
340M
DIT-XL/2 [1]
28
41.4
7,000K
278.20
2.27
4.60
0.83
0.57
1x
675M
Graftnn (Ours)
14
164.6
100K
231.91
3.12
4.71
0.82
0,55
2x
712M
Grafting (Ours)
14
164.6
230K
251.77
2.77
4.87
0.82
0.56
2x1
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 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 (8 Off-the-shed 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]. 9 Speedup is measured for 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-E) -), yielding 43% speedup with less than 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 DiT's These results demonstrate grafting's 's utility in both short- and long-context settings (e.g.. ImageNet- K and PixArt-E, respectively), and for architecture restructuring. Overall, grafting proves be lightweight approach for materializing diffusion transformer designs under small compute budgets.
Related work and discussion. Due to page limit, discuss related work in Supp. A. Further, to demonstrate the generalization of grafting, we graft an LLM (Qwen3-4B [34]] 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.
10

---

References
[1] William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference Computer Vision, pages 4195-4205, 2023. 1,3, 10
[2] 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.
[3] 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 Stti:2312.06662 2023. 1
[4] Eliezer Goldschmidr. Plant grafting: new mechanisms, evolutionary implications. Frontiers in Plant Science, 5:727, 12 2014. 2
[5] 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
[6] Junxiong Wang, Daniele Paliotta, Avner May, Alexander M Rush, and Tri Dao. The mamba in the Ilama: Distilling and accelerating hybrid models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. 22
[7] 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. 22
[8] Dan Fu, Simran Arora, Jessica Grogan, Isys Johnson, Evan Sabri Eyuboglu, Armin Thomas, Benjamin Spector, Michael Poli, Atri Rudra, and Christopher Ré, Monarch mixer: simple sub-quadratic gemm- based architecture. Advances in Neural Information Processing Systems, 36:77546-77603, 2023. 2,
[9] Aran Komatsuzaki, Joan Puigeerver, James Thorph Carlos Riquelme Ruiz, Basil Mustafa, Joshua Ainslic, 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
[10] 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.08360 2020. 2,
[11] Junsong Chen, Chongjian Ge, Enze Xie, Yue Wu, Lewei Yao, Xiaozhe Ren, Zhongdao Wang, Ping Luo, Huchuan Lu, and Zhenguo Li. P: Weak-to-strong training diffusion transformer for 4k text-to-image generation, 2024. 3.9
[12] Jonathan Ho, Ajay Jain, and Pieter Abbeel, Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020. 3
[13] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
[14] Junsong Chen, YU Jincheng, GE Chongjian, Lewei Yao, Enze Xie, Zhongdao Wang, James Kwok, Ping Luo, Huchuan Lu, and Zhenguo Li. Pixart-ac Fast training of diffusion transformer for photorealistic text-to-image synthesis. The Twelfth International Conference on Learning Representations, 2023. 3
[15] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEFE conference on computer vision and pattern recognition, pages 248-255, 2009. 3,5 5
[16] Dhruba Ghosh, Hannanch Hajishirzi. and Ludwig Schmidt. Geneval: An object- focused framework for evaluating text-to- text-to-image alignment. Advances in Neural Information Processing Systems, 36:52132 52152, 2023. 3
[17] Geoffrey Hinton, Oriol Vinyals, and Jeffrey Dean. Distilling the knowledge in neural network. In NeurIPS Deep Learning and Representation Learning Workshop. 2015,
[18] Peter J. Huber. Robust Estimation of Location Parameter. The Annals of Mathematical Statistics, 35(1):73 101, 1964. 4
[19] Peter L Bartlett, Philip M Long, Gábor Lugosi, and Alexander Tsigler. Benign overfitting in linear regression, Proceedings of the National Academy of Sciences, 10008)33300033000000000 2020. 5
11

---

[20] Sanghoon Kim, Dahyun Kim, Chanjun Park, Wonsung Lee. Wonho Song, Yunsu Kim, Hyeonwoo Kim. Yungi Kim, Hyeonju Lee, Jihoo Kim, 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
[21] Shikai Qiu, Andres Potapezynski, Mare 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
[22] 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 Linguisties: Human Language Technologies (Volume 1: Long Papers), pages 7232-7245, 2024.
[23] Michael Poli, Armin W Thomas, Eric Nguyen, Pragaash Ponnusamy, Bjöm 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
[24] 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 convo- lutional multi- hablid language models at scale, arXiv preprint arXiv:2503.01 2025. 6,7
[25] 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, Fric 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 Hernandez- Boussard, Eric Ho, Ming- Yu Liu, Thomas McGrath, Kimberly Powell, Dave P. Burke, Hani Goodarzi, Patrick D. Hsu, and Brian L. Hic. Genome modeling and design across all domains of life with evo 2. bioRxiv, 2025.
[26] Iz Beltagy, Matthew Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv:2004.05150, 2020. 7
[27] Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint IrXiv:1904.105550 2019. 7
[28] 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 7
[29] Edward 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
[30] 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.
[31] Bo-Kyeong Kim, Hyoung- Kyu Song, Thibault Castells, and Shinkook Choi. Bk-sdm: A lightweight, fast, and chcap version of stable diffusion. In European Conference on Computer Vision, pages 381-399. Springer, 2024. 10
[32] Gongfan Fang, Kunjun Xinyin Ma, and Xinchao Wang. Tinyfusion: Diffusion transformers learned shallow. arxiv preprint arXiv:2412.0/ 199, 2024.
[33] Fan Bao, Shen Nic, Kaiwen Xue, Yue Cao, Chongxuan Li, Hang Su, and Jun Zhu. All are worth words: vit backbone for diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 22669-2200 22679, 2023. 10
[34] 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.09:: 2025. 10, 29
[35] 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-resoluse image synthesis with linear diffusion transformers. arXiv preprint 2rX1v:2410.10629, 2024. 22
12

---

[36] 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
[37] Zhengcong Fei, Mingyuan Fan, Changqian Yu, and Junshi Huang. Scalable diffusion models with state space backbone. arXiv preprint 2rXi:2200.00608. 2024. 22
[38] Vincent Tao Hu, Stefan Andreas Baumann, Ming Gui, Olga Grebenkova, Pingchuan Ma, Johannes Fischer, and Bjom Ommer. Zigma: Zigzag mamba diffusion model. arXiv preprint arXiv:2403, 13802, 2024. 22
[39] 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. 142244 2024, 22
[40] 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
[41] 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, Jianggiao 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-7h: Cost-effective training of video generation foundation model. 2025. 22
[42] Zhen Xing, Qijun Feng, Haoran Chen, Qi Dai, Han Hu, Hang Xu, Zaxuan Wu, and Yu-Gang Jiang. survey on video diffusion models. ACM Comput. Surv. 57(2), November 2024.
[43] Yu Gao, Jiancheng Huang, Xiaopeng Sun, Zequn Jic. Yujic Zhong, and Lin Ma. Matten: Video generation with mamba-attention. arXiv preprint arXiv:2405.03000 03025, 2024, 22
[44] Hongjie Wang. Chih-Yao Ma, Yen-Cheng Liu, Ji Hou, Tao Xu, Jialiang Wang, Felix Juefei-Xu, Yaqiao Lao, Peizhao Zhang, Tingho Hou, al. Lingen: Towards high-resolution minutc-length text-to-video generation with linear computational complexity. arXiv preprint ar 241.:2241..08856, 2024. 22
[45] Zhengcong Fei, Mingyuan Fan, Changqian Yu, Debang Li, and Junshi Huang. Scaling diffusion transform- ers to billion parameters. arXiv preprint arXiv:2407. 1163,, 2024. 22
[46] Armin W Thomas, Rom Parnichkun, Alexander Amini, Stefano Massaroli, and Michael Poli. Star: Syn- thesis tailored architectures. In The Thirleenth International Conference on Learning Representations. 2025. 22
[47] Song Bian, Minghao Yan, and Shivaram Venkataraman. Scaling inference-efficient language models. In Forty-second International Conference on Machine Learning. 22
[48] Songhua Liu, Zheuxiong Tan, and Xinchao Wang. Clear: Conv-like lincarization revs pre- pre-trained diffusion transformers up. arXiv preprint 24Xiv:2412.1611 2024. 22
[49] Songhua Liu, Weihao Yu, Zhenxiong Tan, and Xinchao Wang. Linfusion: 1 gpu, minute, 16k image. arXiv preprint wwxXi::2499000999 2024. 22
[50] 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 2rXiv:2503.16726 2025. 22
[51] Akhiad Bercovich, Mohammad Dabbah, Omri Puny, Ido Galil, Amnon Geifiman, Yonatan Geifman, Izhak Golan, Ehud Karpas, Itay Levy, Zach Moshe, et al. Fin fusion: Rethinking sequential computation in large language models. arXiv preprint arXiv:2503. 18008, 2025. 22
[52] 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-2555 1-549, 2021. 22
[53] Xingyi Yang, Daquan Zhou Songhua Liu, Jingwen Ye, and Xinchao Wang. Deep model reassembly. Advancey in neural information processing systems, 35:25739-25755 2022. 22
[54] Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical common- sense in natural language. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pages 1432-7439 2020. 29
13

---

[55] 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: 05457, 2018.
[56] Rowan Zellers, Ari Holtzman, Yonatan Bisk. Ali Farhadi, and Yejin Choi. Hellaswag: Can machine really linish your sentence? Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 4791-4800, 2019. 29
[57] Keisuke Sakaguchi, Ronan Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99- 106, 2021. 29
[58] 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
[59] 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
[60] Keshigeyan Chandrasegaran, Agrim Gupta, Lea M. Hadzie, 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.
[61] Guo Chen, Zhiqi Li, Shihao Wang. Jindong Jiang, Yicheng Liu, Lidong Lu, De-An Huang, Wonmin Byeon, Matthicu Le, Tuomas Rintamaki, al. Eagle 2.5: Boosting long-context post-training for frontier laniun-angguae models, arxiv preprint 250: 15271, 2025. 31
14

---

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 perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long 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 not discussed in the paper.
The authors are encouraged to create 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 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might 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 address problems of privacy and faimess.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant 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 complete (and correct) proof?
Answer: [NA]
15

---

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 cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide 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 supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental 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)?
Question:
Answer: [Yes]
Justification: All training details, hyperparameters, code, and instructions for reproducing results are included in the Supplementary.
Guidelines:
The answer NA means that the paper does include experiments.
If the paper includes experiments, 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 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 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 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 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 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 submis- sions 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 reproduce that algorithm. (b) If the contribution is primarily
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 way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe 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., 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 instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
16

---

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 (uttps://nips.. ec/
olicy) for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ /nips... public/guides "CodeSubmisss sssissiin wwsisssssiiiyyyyy for more details.
The authors should provide instructions 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. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All training details and test details 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 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, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Standard deviation for selected experiments are provided in Supp.
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, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall 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).
17

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
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 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 figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
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 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 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the csGuidel NeurIPS Code of Ethics https: ///ner///susseee cc. htce//ublic/eee tel ines?
Answer: [Yes]
Justification: Yes, we abide by the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration 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 socictal impact 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.
18

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is 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 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
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 gettood-gaaftng-
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 code repository, and (iii) avoid releasing models trained or finctuned on private or sensitive datasets. Since our focus on architcctural 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 such risks.
Released models that have 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 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 challenging, and many papers do not require this, but encourage authors to take this into account and make 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-E) and datasets (e.g.. ImageNet- 1K) under their respective licenses. All assets properly cited in the paper, and their licenses are respected.
Guidelines:
The answer NA means that the paper does not 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 used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
19

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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 available online, the authors encouraged to reach out the asset's 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 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 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 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 the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, 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 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:
20

---

The answer NA means that the paper does not involve crowdsoureing 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 expect authors adhere 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
Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: LLMs were not used 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: / /nnuriiss.....// cc/Conf / /LLM) for what should or should be described.
21

---

Supplementary Material
Section Related Work
Section Standard Deviation of Experiments
Section C: Hybrid Architecture Experiments: Additional details
Section C.I 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: FL.OP 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 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-MoF [45] introduces sparse diffu- sion 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 el al. [47] showed that pretraining wider and shallower autoregressive lan- guage 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 materialize new architectures under small compute budgets.
Architectural editing of pretrained generative models. Another line of work focuses on lineariz- ing large language models by replacing softmax attention with efficient operators, such linear attention [5, 6, 7]. Similar ideas have also been adopted for diffusion models in [48, 49, 50], though these works focus only on ulrra-high-resolltton settings. These prior efforts typically focus on replacing a single operator type (primarily attention) or are limited to specific application domains. Grafting presents more general and comprehensive approach for architectural editing. extends beyond single-operaton 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 LLMss, aiming to reduce sequential computation. While our work focuses on diffusion-ban generative modeling with transformers, prior work 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 experi- MHa (Hyena-Y) and MLP 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.
Setup
IS
FID
sFID
Precision
Recall
MHA/ Hyena-Y
273.19 + 0.46
2.73 0.01
5.06 + 0.04
0.83 - 0.00
0.55 0.00
MLP/ higher width (r 277.91 0.95 2.41 + 0.01
4.48 + 0.02
0.82 0.00
0.58 0.00
Table B.1: Mean and standard deviation of IS, FID, sFID, Precision and Recall calculated for three runs with different random seeds ,0,200,300,
C Hybrid Architecture Experiments: Additional details
C.1 Experiment details and additional samples
We provide all hyperparameters used for the experiments in Tab. C.I. To ensure a fair comparison, used identical hyperparameters across every hybrid experiment. We include additional qualitative samples generated using our hybrid architectures obtained via grafting in Fig. C.1.
Stage 1: Activation Distillation
Initial Learning Rate
x 10
Weight Decay
0
Epochs
200
Batch Size
64
Clip Norm
10.0
Optimizer
AdamW (betas (0.9, 0.999))
Loss Function
L1 (MHA), L2 (MLP)
Stage 2: Lightweight Finetuning
Initial Learning Rate
1 x 10-4
Weight Decay
5 x 10-5
Iterations
50,000 (100 epochs)
Batch Size
256
Optimizer
AdamW (betas (0.9, 0.999))
Scheduler
Linear Warmup over 1K steps, then constant
Training Data
10% of ImageNet- 1K (128k samples)
Table C.1: Experiment details for MHA/MLP grafting experiments using DiT-XL/2 (ImageNet-1K)
C.2 Modulated Regression Targets for MHA
For Stage 1, we explored a modulation-awain -ware regression variant for MHA experiments that incor- porates the learned scalar (gate_mst (_maa)mmsaaaae applied to the attention output. In the standard setup, we regress from input x to the raw output of the attention block MHA(.). In the modulation-awan formulation, the target becomes y = gate_me msa @ (-A Tab. C.2 compares these two variants with L1 and 1.2 loss. Modulation-awanl -warre regression increases target scale, which adversely affects 1.2 loss performance due to sensitivity large values. I.1 performs similarly both settings. We adopted the standard (modulation-agnooiil formulation for all experiments for simplicity.
23

---

Modulation-aware
18
FID
sFID
Precision
Recall
Baseline
278.20
2.27
4.60
0.83
0.57
246.17
3.00
7.11
0.79
0.58
L2
X
269.31
2.58
5.75
0.82
0.58
272.86
2.51
5.29
0.82
0.57
L1
X
273.03
2.51
5.48
0.83
0.58
Table C.2: Comparison of modulation-awan and standard regression targets for Stage 1. The modulation-awe -warre setup includes the learned scalar (gate_msaa as a multiplicative factor in the regression target. L2 loss is sensitive the amplified target scale and performs worse, while L1 loss remains robust and performs similarly both cases. We adopt the standard formulation by default.
C.3 Validation Loss Curves for Self-grafting Experiments
To support the trends reported in the main paper (Sec.3.2, Table 2), include validation loss curves in Fig.C.2 C.2 for five representative layers. Loss is computed using L2 on a held-out set of 8k samples. For MHA layers (top row). L1 consistently achieves lower validation loss in deeper layers, reflecting robustness to high activation variance. For MLP layers (bottom row), L2 generalizes best across all layers. This contrast may be explained by parameter count: MLPs have roughly 2x more parameters than MHA layers, making them less sensitive to outliers and better suited to L2 regression.
Hhyna-X
Swa
Mamba-2
r
Figure C.I: Additional samples generated by grafted DiT-XL/2 variants. Each row corresponds to different hybrid. We report FID scores (lower is better, ImageNet- 1K 256x256) for each hybrid. MHA variants (top 4 rows): Hyena-X (2.61), Hyena-Y (2.61), SWA (2.62), Mamba-2 (2.55). MLP variants (bottom 3 rows): Lower width (2.53), Higher width (2.38), Hyena-X MLP (2.64). These results highlight the flexibility of grafting in constructing high-quality hybrid architectures by replacing MHA or MLP operators.

---

Operator Initialization Results (Self-grafting/ MHA)
Layer
Layer
Layer 17
Layer 27
Layer 28
10
vol. loss
Loss
20
via.
104
0.02
10
200
200
200
Huneer, 60 1.0
Operator Initialization Results (Self-grafting/ MLP)
Laycr
20
Layer 8
Layer 17
20-
Layer 27
100
Layer 28
10-
Val. loss
vol. Loxs
vol. Toss
vol.
vol.
0.1
100 200
100
200
100
200
100
200
100
200
H1.0er,==1.0
Figure C.2: Validation loss curves for MHA (top) and MLP (bottom) operator distillation showing the training dynamics for three regression objectives. As one can observe, LI shows better generalization for MHA and 1.2 shows better generalization for MI.P
D Depth to Width Grafting Experiments: Additional details
We provide all hyperparameters used for these experiments in Tab, D.I. We show additional samples in Fig. D.1.
Figure D.1: Depth-to-width grafting samples. Samples from a DiT-XL/2 model in which every pair of transformer block is converted into parallel block, effectively reducing depth by 2x (FID=2.77).
Implementation details. The generative quality metries (FID, IS, sFID, Precision, Recall) reported in Tab. 6 correspond to the exact implementation of rewiring presented in Fig. 5. For speedup measurements, simple fused version of our rewired version was used. is important to note that the reported speedup values 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 MHA activation plots
We show activation distribution for five representative layers (1 15, 17, 19, 21, 23) in Fig. E.1.
E.2
Experiment details
We provide all hyperparameters used in PixArt-\ grafting experiments in Tab. E.I.
25

---

Stage 1: Activation Distillation
Initial Learning Rate
10
Regression Objective
L1
Epochs
200
Batch Size
64
Optimizer
AdamW (betas (..9)).
Stage 2: Lightweight Finetuning
Learning Rate
x 10
Weight Decay
Iterations
230k
Batch Size
256
Optimizer
AdamW (belas = (0.9, 0.95))
Scheduler
Warmup over 1K steps, half every 75k steps
Training Data
25% of ImageNet-1K (320k)
Table D.I: Experiment details for depth-to-width grafting experiments using DiT-XL/2 (ImageNet-
1K).
Activation Distribution for Selected Layers (PixArt- E)
(H)
180
bagi Conmta)
135
142
6.
Antvatiin Valies
100
Arthaation valun
25
DIT-XL2 2 Layer
22 2 R
t
Day MAPXL Layer
4
22
275
130
Figure E.1: Similar DiT-XL/2, we observe high activation variance in PixArt-E MHA operators. We show input (left) and output (right) activation values corresponding to five representative layers (15, 17, 19, 21, 23) in PixArt-
E.3 Generated samples and failure cases
We show additional 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-E outputs. Each column pair shows the output of Pix (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- sesgmment-iddicatnng 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 avoid propagating errors during the grafting process.
samples generated by the grafted PixArt-Y model in Fig. E.2,
26

---

The saage capturer sceeo de stack beauty the...
Oze wateroalor illaraaton a man cz white back...
The image raprurts sleseap veew ot vairzt net..
vep lapge cattonn panda weurs baseball oop, the aty
in the procesicoa oú puinticg stuto taee is 32...
Handmmade vallage sock monicy dill, ahanbon
Today marks the beginning of sutuma, summer has nit...
Prest codander leares plant oz white beskyrrunn,
The image captins ruument of Iniquiliy the deop Muee
Paint captrsting wateroloc chper of Levbird.
Figure E.2: Additional 2048 2048 samples generated using our grafted PixArt-E model.
image presents delightful domuts, own unique topping, set
Playing classical guitar, Aesthetic nower, body sty vectoo
dog chasing tail looks cute funny, abstract photography, stippling, UHD, high
Tiss image captures a moment city street, frozen black white. Dominating the left
Figure E.3: Text-to-image generation failure cases. Each pair shows outputs from PixAn--E (left) and the grafted model (right) for the same prompt. In the top row, the prompt specifies four donuts with unique toppings and full-body anime-style character playing classical guitar. The grafted outputs deviate from these prompts-sie showing-showing incorrect object counts (e.g., five donuts) and degraded structure (e.g.. distorted hands), reflecting text-image misalignment and visual artifacts introduced during grafting. In the bottom row, the supervision itself is poor: prompts such as 'dog chasing its tail UHD stippling style' and 'a black-and-whin street photo' are not faithfully captured by either model. These examples highlight challenges arising both from LORA finetuning and low-quality synthetic data.
27

---

Figure E.4: Examples of low-quality samples generated by PixArt-E, used for grafting. These images contain unrealistic features, inconsistent physics, and visual artifacts. Their presence in the grafting dataset can degrade generation quality of grafted models, highlighting the importance of data curation when using synthetic data.
Stage 1: Activation Distillation
Initial Learning Rate
1 x 10-4
Weight Decay
1 10-5
Epochs
100
Clip Norm Value
0.1 (Layers 20-27), 0.01 (Other layers)
Batch Size
16
Optimizer
AdamW
Scheduler
Half Ir epochs 50
Stage 2: Lightweight Finetuning
Initial Learning Rate
1 x 10-5
Weight Decay
0
Iterations
18k
Batch Size
64 (with gradient accumulation)
Optimizer
AdamW
Scheduler
linear warmup (500 steps), then constant
LoRA rank
64
Table E.I: Experiment details for PixArt-E grafting experiments.
28

---

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 ceffusioo-basee 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 (activation distillation) and Stage 2 (lightweight fine-tuning) lengths of 1024 and 8192 tokens respectively. We use 50k instruction- following examples from the Alpaca-cleaned dataset 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) following the evaluation setup in [5].
performed identically to the diffusion experiments, with context
Results. As shown in Table F.I, our grafted Qwen3-4B model achieves a 1.4x 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.
Model
Stage
Ratio
PIQA
ARC-e
ARC-e
HellaSwag
Winogrande
MMLU
Avg.
Speedup f
Baseline
74.9
80.5
54.0
68.5
66.0
70.1
69.0
Grafting (Ours)
Random Init.
64.3
56.5
29.7
37.9
50.9
25.9
44.2
1.4x
Grafting (Ours)
Stage 1
50%
74.8
80.4
53.4
68.2
65.5
66.9
68.2
1.xx
Grafting (Ours)
Stage 2
50%
75.5
80.3
52.3
69.6
67.0
66.9
68.6
1.4x
Table F.1: Qwen3-4B [34] grafting results. We evaluate grafting Qwen3-4B by replacing Multi- Head Attention (MHA) with Sliding Window Attention (SWA, window=256). Stage denotes activation distillation; Stage 2 denotes lightweight fine-tuning. Performance is reported on standard commonsense and reasoning benchmarks. Our grafted model achieves 1.4x decode throughput at 8k context length (single Nvidia H100) with less than 1% drop in average performance relative the baseline.
G Hyena-X and Hyena-Y Y operators: Additional details
Informed by our band-k analysis of MHA operators, introduce collection of efficient operators designed to exploit the locality in attention matrices. Given an input x E Rexd, a generic Hyena operator performs the following transformation:
Σ TE
Σ
x wée
kg = Σ
H
88'
5 = Σ
Ke
pe'e
as!
v. ΣΣ (g G
where W.U,P. M E Rddd are parametrized dense or low-rank matrices, and T. H, K, G E R are Toeplitz matrices corresponding to convolutions with the filters li, by hx, hGs respectively. In the original formulation [59], the filters hr, hn, hg are short and explicitly parametrized, whereas he is implicitly parametrized.
We build on this formulation and propose Hyena-X X and Hyena-Y, two Hyena operators designed for grafting. Hyena-X X removes the implicit convolution entirely by setting G = 1. In contrast, Hyena-Y
introduces two changes: (i) it removes all three featurizer convolutions (T, H, K), and (ii) replaces
hnttps://hugginggaass http:///c...ccttttttt///eeeeeeee
cleane
29

---

the implicit long convolution in G with short, explicit convolution. This modified structure preserves local inductive bias while significantly reducing computational cost. An illustration provided in main paper. These operators allows us to realize speedups across 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.I.
H.1
MHA
Input projections (Q, V): 6LD
Softmax attention: 4IA + 2H1?
Output projection: 2LD
H.2
SWA
Input projections (Q, K, V): 6LD
Sliding window attention (Bidirectional): 4L(2w + 1)D + L + +1)
Output projection: 2LD
H.3
Hyena-SE
Input projections: 6LD
Featurizer: 3LDK
Inner filter convolution: LDK x
gates: LD x
Output projection: 2LD
H.4 Hyena-X
Input projections: 6LD
Featurizer: 3LDK x
gates: LD x
Output projection: 2LD
H.S Hyena-Y
Input projections: 6LD2
Inner filter convolution: LDK
gates: LD x
Output projection: 2LD²
H.6
Hyena-X (MLP)
Dense input projections: r
Featurizer: 3LDK
Gates: LD x
Dense output projections: 2
30

---

H.7
Mamba-2
Projections: 8LDD²
Short convolution: 6LDE
Featurization: 2LDE(1 + 2dstate)) + 2LDE
Associative scan: 2dDEdaate
Output layer: 2LDºE
Symbol Description
Sequence length
D Hidden dimension
I1
Number of attention heads
K
Kernel size for convolutions
w
Window size for sliding window attention
MLP expansion ratio
E
Expansion factor in Mamba-2
dstate
State size Mamba-2
Table H.1: Notation for FLOP calculation.
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 compute- efficient, enabling broader participation in generative modeling research. Grafting facilitates practical improvements-S55 such as efficient operator replacement and structural 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 harmful content. We encourage responsible use, transparency, and content safeguards when applying these techniques.
citiiiaattooooffrriin a
Linitations. This work primarily focuses architectural editing of pretrained Diffusion Transform- ers (DiTs), specifically targeting self-attention and MLP operators. Other architectural components, such 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 VAEs remains an area for future study. The PixArt-E 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 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/geeennii long-form [60, 61], or improving user experience in interactive applications like image editing where modest gains (e.g., 10% speedup) highly valued. We hope that testbed, insights, and results will encourage the community to actively explore new architecture designs. Code and grafted models: grafting stanford.edu.
to
31

---

Compute details
We report the total compute used to run all experiments this work, including self-grafing base- lines, hybrid architectures, ablations. PixArt-E adaptation, and depth-to-wied to-hthhh restructuring. All experiments were conducted using 8xH100 GPUs. Table J.1 provides 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-X1/2 completes within 12-24 hours on 8xH100 GPUs (both Stage and Stage 2).
Experiment
H100 hours
# Experiments
Total H100 hours
Self-grafting Experiments (Tables & 2)
12 hrs x 8
4
384
Hybrid Experiments (Table 4)
12 hrs x 8
26
2,496
Ablations
24 hrs
6
1,152
PixArt-E experiments
63 hrs x 8
2
1,008 (Table 5) 960 (Table 6) 50
Depth-to-widit restructuring
120 hrs x
1
Other (feature extraction, tuning)
50 hrs
Total
6,050
Table J.1: Estimated compute usage across all experiments. Total GPU hours account for training, sampling, and evaluation. Each experiment conducted using 8xH100 unless otherwise specified.
32


Reviews:
[{"review_id": "rCd1U7Aqgw", "reviewer": "Reviewer_qERF", "summary": "- 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%.", "questions": "- 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?", "limitations": "yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "I thank the authors for their detailed rebuttal. They have answered my questions and concerns satisfactorily. Particularly my concern about general applicability of grafting has been answered properly. After reading through other reviews and author rebuttals, I vote to retain my score.", "rebuttal": "We thank the AC and the reviewer for their valuable time and feedback. The reviewer appreciated our activation behavior and locality analysis, testbed and experiments. The reviewer raised concerns about training hybrid architectures from scratch, general questions about the experiment setup, and applications of grafting. Below, we address all comments.\n\n$~$\n> W1: Training from scratch with operators from grafting.\n\nThank you for the thoughtful question. As noted by the reviewer, pretraining is beyond the scope of this work, as we focus on post-training architectural editing of diffusion transformers. While an interesting direction, even training a 64x64 DiT takes more than 12 days on 32×A100s [A], making pretraining infeasible on an academic budget. That said, we agree this is an intellectually interesting question. To address it, we highlight independent evidence from language modeling showing that hybrid architectures mixing attention with Hyena-style operators can be trained from scratch and remain competitive.\n\n**Evidence: StripedHyena2: 7B-scale language models trained with Hyena-SE [B]:**\nStripedHyena2 [B] trains large-scale hybrid models from scratch using setups which mix attention and Hyena operators, including Hyena-SE. At the 7B scale, these models achieve lower perplexity than transformer baselines (2.83 vs. 3.09) after 400B tokens of pretraining—demonstrating that hybrid architectures remain expressive at scale.\n\n$~$\n> Q1: Why was LoRA used for training PixArt-? What is exactly being done here? Why do you need LoRA when you are finetuning all the replaced operators in the 2nd stage?\n\nWe used LoRA in our high-resolution PixArt-Σ experiments to reduce memory and compute overhead. As noted in lines 254–255, these experiments operate over long sequences (e.g., shape [B,16384,1152]), which makes standard finetuning expensive. In Stage 2 (lightweight finetuning), we apply LoRA to self-attention (q/k/v/o), cross-attention (q, kv), and MLP (fc1/fc2), and we update the 1D causal convolutions in Hyena-X directly (added to the optimizer; no LoRA, since these are 1D convs with a small parameter count). We will clarify this in the final version.\n\n$~$\n> Q2 (1) Difference between Hyena-SE, Hyena-X and Hyena-Y operators\n\nWe summarize the differences below. The key distinctions lie in the use of convolutional featurizers and the presence or absence of the inner filter. As suggested by the reviewer, we will add Hyena-SE to Fig. 3 in the final version. We have also provided the operator equations in Supp. E. \n\n|          | Conv. Featurizers | Inner Filter               | FLOP equation |\n| -------- | ----------------- | -------------------------- | ------------- |\n| Hyena-SE | ✅\t               | Short Explicit Convolution | Supp F.3      |\n| Hyena-X  | ✅\t               | ❌                    | Supp F.4      |\n| Hyena-Y  | ❌           | Short Explicit Convolution | Supp F.5      |\n\n$~$\n> Q2 (2) Question reg. causality \n\nWhile our testbed (Table 3) already includes bidirectional operators (e.g., self-grafting, sliding window attention), we introduced causal Hyena operators to study causal inductive biases. Notably, causal variants perform competitively—Hyena-Y achieves a FID of 2.61, compared to 2.62 for Sliding Window Attention (SWA). If the reviewer prefers, we can include a bidirectional version of Hyena-X or Hyena-Y in the final version. This requires only a simple change: removing left-padding in the inner convolution and instead padding the input on the right.\n\n$~$\n> Q3 (2) Memory footprint of hybrid architectures\n\nThe reviewer is correct in noting that efficient operators can reduce memory usage of the full model in addition to speed, depending on their hardware-aware implementation. We provide a reference: At a sequence length of 2048, peak memory is ~2.70 GB for a GQA Transformer++ baseline, versus ~2.55 GB for the corresponding Hyena-Y hybrid architecture (Δ −150 MB, ≈ −5.6%).\n\n$~$\n> Q3 (1) Edge applications\n\nThe reviewer is correct—while this work studies grafting as a general-purpose architectural editing approach, replacing expensive components with cheaper alternatives naturally leads to faster inference/ lower memory usage, which are particularly valuable in edge scenarios. \n\nTo further demonstrate broader applicability, we also applied grafting to large autoregressive language models. These experiments were conducted to address grafting’s generalization to other domains (different task/architectures) suggested by Reviewers r7Mh, VANM and JnND. We believe the results may be of interest here, especially in the context of real-time or edge deployment. By grafting a 36-layer **Qwen3 4 billion parameter language model with Sliding Window Attention operators (w=256, grafting ratio=50%)**, we obtain **1.4x decode throughput (at 8K context length, Nvidia H100 with KV cache) with less than 1% overall performance drop.**\n\n| Model       | Stage        | Grafting Ratio | PiQA | ARC-e | ARC-c (acc. norm) | HellaSwag (acc. norm) | Winogrande | MMLU (5-shot) | Avg. | Speedup |\n| ----------- | ------------ | -------------- | ---- | ----- | ----------------- | --------------------- | ---------- | ------------- | ---- | ------- |\n| Qwen3       | Baseline     | N/A            | 74.9 | 80.5  | 54.0              | 68.5                  | 66.0       | 70.1          | 69.0 | \\-      |\n| SWA (w=256) | Random Init. | 50%            | 64.3 | 56.5  | 29.7              | 37.9                  | 50.9       | 25.9          | 44.2 | 1.4x    |\n| SWA (w=256) | 1            | 50%            | 74.8 | 80.4  | 53.4              | 68.2                  | 65.5       | 66.9          | 68.2 | 1.4x    |\n| SWA (w=256) | 2            | 50%            | 75.5 | 80.3  | 52.3              | 69.6                  | 67.0       | 66.9          | 68.6 | 1.4x    |\n\n\n**We hope our response addresses the reviewer's questions, and we’re happy to provide further details if helpful. Thank you for the thoughtful feedback.**\n\n===\n\n[A] Karras, Tero, et al. \"Elucidating the design space of diffusion-based generative models.\" NeurIPS 2022\n\n[B] Ku, Jerome, et al. \"Systems and algorithms for convolutional multi-hybrid language models at scale.\" arXiv preprint arXiv:2503.01868 (2025)."}, {"review_id": "wFtotqXXHw", "reviewer": "Reviewer_r7Mh", "summary": "This 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.", "questions": "- 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.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "I appreciate the authors’ comprehensive rebuttal, which has satisfactorily addressed the majority of my concerns. Considering both the feedback provided by the other reviewers and the authors’ clarifications, I remain inclined to maintain my overall positive assessment of this submission.", "rebuttal": "We thank the AC and the reviewer for their constructive feedback. The reviewer highlighted the simplicity and efficiency of our approach and our experiments. The reviewer’s questions center on: (1) generalization of grafting beyond diffusion-based image generation using DiTs, (2) how operators interact with surrounding layers, (3) pretraining hybrid architectures, (4) potential artifacts from synthetic data, and (5) clarifications on experimental results. We respond in detail below:\n\n$~$\n> W1+Q3: Generalization beyond diffusion-based image generation using the DiT architecture.\n\nThank you for your feedback. To address reviewer concerns, we evaluate grafting in an **autoregressive language modeling setting—a task, architecture, and modality distinct from diffusion-based image generation with DiTs.**\n\n**Setup:**\n\n- **Task:** Diffusion-based image generation → Autoregressive language modeling (next-token prediction)\n- **Architecture & Scale:** DiT-XL/2 (0.7B) → Qwen3 (4B)\n- **Operator:** MHA → Sliding Window Attention (window=256)\n- **Dataset**: Alpaca-cleaned (50k instruction-following examples)\n- **Evaluation:** PiQA, ARC-e, ARC-c, HellaSwag, Winogrande, MMLU following [E]\n- **Experiment details:** We apply our exact two-stage grafting procedure—activation distillation followed by lightweight fine-tuning. Stage 1 uses context length 1024; Stage 2 uses 8192.\n\n**Results:**\n\nOur grafted Qwen3-4B model achieves a **1.4x decode throughput** (at 8K context length with KV cache, Nvidia H100) with **less than 1% average performance drop**, demonstrating that grafting generalizes to new tasks, architectures, and modalities. \n\n| Model       | Stage        | Grafting Ratio | PiQA | ARC-e | ARC-c (acc. norm) | HellaSwag (acc. norm) | Winogrande | MMLU (5-shot) | Avg. | Speedup |\n| ----------- | ------------ | -------------- | ---- | ----- | ----------------- | --------------------- | ---------- | ------------- | ---- | ------- |\n| Qwen3       | Baseline     | N/A            | 74.9 | 80.5  | 54.0              | 68.5                  | 66.0       | 70.1          | 69.0 | \\-      |\n| SWA (w=256) | Random Init. | 50%            | 64.3 | 56.5  | 29.7              | 37.9                  | 50.9       | 25.9          | 44.2 | 1.4x    |\n| SWA (w=256) | 1            | 50%            | 74.8 | 80.4  | 53.4              | 68.2                  | 65.5       | 66.9          | 68.2 | 1.4x    |\n| SWA (w=256) | 2            | 50%            | 75.5 | 80.3  | 52.3              | 69.6                  | 67.0       | 66.9          | 68.6 | 1.4x    |\n\n\n$~$\n> W2 (1) However, it's unclear whether different operator might rely on their surrounding operators—since grafting requires freezing the context parameters. \n\nThe reviewer is correct: operator initialization relies on the surrounding operators. We freeze surrounding parameters only in Stage 1 (activation distillation). In Stage 2 (lightweight fine-tuning) we update all parameters end-to-end (see Fig. 1, Tables 4 and 6). For PixArt-Σ, we perform LoRA finetuning for memory efficiency (long sequence length = 16,384). We apply LoRA to self-attention (q/k/v/o), cross-attention (q, kv), and MLP (fc1/fc2), and update 1D causal convolutions in Hyena-X directly.\n\n$~$\n> W2 (2) + Q4: It also remains uncertain whether architectures that perform well under the grafting setting would yield similarly results when trained from scratch.\n\nThank you for the thoughtful question. We note that pretraining is beyond the scope of this work, as we focus on **post-training architectural editing** of diffusion transformers. While an interesting direction, even training a 64x64 DiT takes more than 12 days on 32×A100s [A], making pretraining infeasible on an academic budget. That said, we agree this is an intellectually interesting question. To address it, we highlight independent evidence from language modeling showing that **hybrid architectures mixing attention with Hyena-style operators (including Hyena-SE) can be trained from scratch and remain competitive.**\n\n**Evidence: StripedHyena2: 7B-scale language models trained with Hyena-SE [B]:** StripedHyena2 [B] trains large-scale hybrid models from scratch using setups which mix attention and Hyena operators, including Hyena-SE. At the 7B scale, these models achieve **lower perplexity than transformer baselines** (2.83 vs. 3.09) after 400B tokens of pretraining—demonstrating that hybrid architectures are competitive at scale.\n\n$~$\n> Q1: 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\nThank you for the thoughtful question. We believe the strong performance of the interleaved strategy stems from its ability to balance local and global receptive fields. In contrast, stacking multiple local operators (top-local) may restrict the model’s effective receptive field. Alternating with global operators helps retain broader context. Similar interleaving strategies have recently been shown effective in the pretraining of language models [B] and biological foundation models [C].\n\n$~$\n> Q2: 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\nThank you for the question. We do not expect such artifacts to appear when using real data. In Supplementary Fig. D.4 and D.4, we show that some low-quality synthetic samples can introduce artifacts during grafting. A practical way to mitigate this is to curate or filter synthetic data before use. Recent works [D] have explored automated data curation pipelines using VLMs for improving data quality.\n\n\n**We hope our response addresses the reviewer's questions, and we’re happy to provide further details if helpful. Thank you for the thoughtful feedback.**\n\n####  \n\n####  \n\n[A] Karras, Tero, et al. \"Elucidating the design space of diffusion-based generative models.\" *NeurIPS 2022*\n\n[B] Ku, Jerome, et al. \"Systems and algorithms for convolutional multi-hybrid language models at scale.\" *arXiv preprint arXiv:2503.01868* (2025).\n\n[C] Brixi, Garyk, et al. ‘Genome Modeling and Design across All Domains of Life with Evo 2’. *bioRxiv*, Cold Spring Harbor Laboratory, 2025.\n\n[D] Cai, Shengqu, et al. \"Diffusion self-distillation for zero-shot customized image generation.\" CVPR 2025.\n\n[E] Zhang, Michael, et al. \"LoLCATs: On Low-Rank Linearizing of Large Language Models.\" ICLR 2025"}, {"review_id": "0XunkuqsAs", "reviewer": "Reviewer_VANM", "summary": "This 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.", "questions": "Please see Weakness.", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "The authors have addressed most of my concerns in the rebuttal. I believe this work deserves to be accepted.", "rebuttal": "We thank the AC and reviewer for their valuable time and feedback. The reviewer raised several concerns: (1) generation quality gap and the practical relevance of grafting, (2) lack of comparison to caching methods, (3) limited generalization beyond DiT architectures and image generation, and (4) suggestions regarding latency reporting and improved results in Table 6. Below, we address all points in detail.\n\n$~$\n> W1: Generation quality gap and practical relevance\n\nThe reviewer is correct, we agree that generation quality is essential for practical utility. However, the goal of this work is not to outperform the base model, but to investigate whether pretrained diffusion transformer architectures can be edited post-training under small compute budgets, and to characterize the resulting quality–efficiency tradeoffs.\nPrior to this work, it was unclear whether such post-training architecture editing could retain generation quality at all. Our results show that they can: across all experiments, the grafted architectures demonstrate good generation quality, despite significant architectural editing. For example:\n\n- **Table 4 & Fig. 3 (Right):** We obtain hybrid architectures that are within 0.5 FID of the DiT-XL/2 baseline, while grafting 50% of MHA layers (Hyena-SE/X/Y, SWA) and 100% of MLPs (MLP r=3).\n\n- **Table 5:** On PixArt-Σ, we obtain a 1.43× faster variant at 2048×2048 resolution with <2% GenEval score drop.\n\n- **Table 6:** We restructure DiT-XL/2 by trading depth for width, producing a 2× shallower model (2× faster with fused implementation, bs=1) that achieves better FID (2.77) and outperforms pruning baselines and shallow-from-scratch trained models at depth=14.\nWhile some quality drop is expected when replacing expensive operators with cheaper ones, our study quantifies these tradeoffs and demonstrates that high-quality generation is achievable when editing DiT architectures post-training.\n\n$~$\n> W2: Comparison to caching methods\n\nThank you for raising this point. Caching methods like ToCa, TeaCache, PAB, and TaylorSeer accelerate inference by reusing intermediate activations at runtime. These techniques alter execution but leave the model’s architecture unchanged. In contrast, grafting edits the model’s architecture itself—replacing expensive operators/restructuring the computational graph post-training. **The two approaches are orthogonal and complementary. We will include a discussion of these works in the final version. Our submission also compares grafting with SOTA pruning methods in Section 6 (Table 6).**\n\n\n$~$\n> W3: Generalization beyond image generation\n\nWe note that Generalization was a recurring theme raised by two other reviewers.\n\n- Reviewer r7Mh suggested experiments beyond “diffusion-based image generation” and “architectures beyond DiT”. \n\n- Reviewer JnND mentioned “experiments only in image generation may be insufficient” with video generation as one possible experiment.\n\n**To address these concerns collectively, we evaluate grafting in the language modeling setting—a task, architecture, and modality distinct from diffusion-based image generation with DiTs.**\n\n**Setup:**\n\n- **Task:** Diffusion-based image generation → Autoregressive language modeling (next-token prediction)\n- **Architecture & Scale:** DiT-XL/2 (0.7B) → Qwen3 (4B)\n- **Operator:** MHA → Sliding Window Attention (window=256)\n- **Dataset**: Alpaca-cleaned (50k instruction-following examples)\n- **Evaluation:** PiQA, ARC-e, ARC-c, HellaSwag, Winogrande, MMLU following [A]\n- **Experiment details:** We apply our exact two-stage grafting procedure—activation distillation followed by lightweight fine-tuning. Stage 1 uses context length 1024; Stage 2 uses 8192.\n\n**Results:**\n\nOur grafted Qwen3-4B model achieves a **1.4x decode throughput** (at 8K context length with KV cache, Nvidia H100) with **less than 1% average performance drop**, demonstrating that grafting generalizes to new tasks, architectures, and modalities. \n\n| Model       | Stage        | Grafting Ratio | PiQA | ARC-e | ARC-c (acc. norm) | HellaSwag (acc. norm) | Winogrande | MMLU (5-shot) | Avg. | Speedup |\n| ----------- | ------------ | -------------- | ---- | ----- | ----------------- | --------------------- | ---------- | ------------- | ---- | ------- |\n| Qwen3       | Baseline     | N/A            | 74.9 | 80.5  | 54.0              | 68.5                  | 66.0       | 70.1          | 69.0 | \\-      |\n| SWA (w=256) | Random Init. | 50%            | 64.3 | 56.5  | 29.7              | 37.9                  | 50.9       | 25.9          | 44.2 | 1.4x    |\n| SWA (w=256) | 1            | 50%            | 74.8 | 80.4  | 53.4              | 68.2                  | 65.5       | 66.9          | 68.2 | 1.4x    |\n| SWA (w=256) | 2            | 50%            | 75.5 | 80.3  | 52.3              | 69.6                  | 67.0       | 66.9          | 68.6 | 1.4x    |\n\n\n$~$\n> Q1: Table 6 latency and improved FID results.\n\nWe thank the reviewer for this suggestion. In the final version, we will add latency/ speedup relative to the baseline. As per reviewer’s suggestion, we have trained an improved variant that achieves a better FID (2.77) using 230k iterations. Below we include the updated table for your reference (Speedup measured using a fused implementation, bs=1).\n\n| **Method**    | **Depth** | **A.R** | **Iters** | **IS ↑** | **FID ↓** | **sFID ↓** | **Prec. ↑** | **Recall ↑** | **Speedup ↑** | **Params ↓** |\n| ----------------------------------------- | --------: | ------: | --------: | -------: | --------: | ---------: | ----------: | -----------: | ------------: | -----------: |\n| DiT-L/2 |  24 |  42.7 |    1,000K |   196.26 | 3.73 |  4.62 | 0.82 | 0.54 | — |  458M |\n| U-ViT-L |  21 |  48.8 |      300K |   221.29 |  3.44 |  6.58 | 0.83 |  0.52 |  — |  287M |\n| DiT-B/2 | 12 |   64.0 |    1,000K |   119.63 | 10.12 |  5.39 | 0.73 |0.55 |  — | 130M |\n| BK-SDM| 14 |  82.3 |      100K |   141.18 | 7.43 |  6.09 | 0.75 | 0.55 |       2× |         340M |\n| TinyDiT-D14 |  14 |    82.3 | 500K |   198.85 | 3.92 | 5.69 | 0.78 |         0.58 |       2×|         340M |\n| TinyDiT-D14 w/ MKD |        14 |    82.3 | 500K |  234.50 | 2.86 |       4.75 |        0.82 |         0.55 |       2× |         340M |\n| DiT-XL/2 |  28 | 41.4 | 7,000K |   278.20 |  2.27 | 4.60 | 0.83 |         0.57 |            1× |         675M |\n| **Grafting** |14 | 164.6 | 100K |   231.91 | 3.12 |  4.71 | 0.82 |         0.55 |       2× |         712M |\n| **Grafting** |**14** | **164.6** |  **230K** |   **251.77** |  **2.77** | **4.87** |**0.82** |  **0.56** |       **2×** |         **712M** |\n\n\n**We hope our response addresses the reviewer's questions, and we’re happy to provide further details if helpful. Thank you for the thoughtful feedback.**\n\n===\n\n[A] Zhang, Michael, et al. \"LoLCATs: On Low-Rank Linearizing of Large Language Models.\" ICLR 2025"}, {"review_id": "T5x4dH9eiO", "reviewer": "Reviewer_w2ZL", "summary": "This 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.", "questions": "Please 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.", "limitations": "Please see weaknesses and questions.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "*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.", "was_revised": false, "final_justification": "", "rebuttal": "We thank the AC and the reviewer for their helpful feedback. The reviewer appreciated our operator replacement and depth to width restructuring studies. The reviewer’s questions focus on: (1) whether lightweight fine-tuning remains modest in compute for deep edits and how this compares to small from-scratch retrains; (2) role of distillation and the feasibility of pretraining hybrids; (3) experimental details (replacement strategy, hyperparameters); and (4) presentation suggestions. Below, we address each point in detail.\n\n$~$\n> W1: 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.\n\nWe agree that very deep edits can require more updates. Empirically, our **depth→width restructuring** reached **FID = 2.77 after 230k iterations** **(new result obtained during rebuttal)**, which is still tiny compared to the **~7M iterations** used to train the base model—i.e., even for deep edits we expect compute to remain **at least an order of magnitude lower than pretraining**. On “fresh small retrain”: the reviewer is correct to raise this. We note that **pretraining is out of scope** in our work as we focus on **post-training architectural editing** to study the resulting **quality–efficiency trade-offs**. Studying the relationship to from-scratch training is left to future work due to compute demands. Even pretraining a 64×64 ImageNet-1K DiT takes >12 days on 32×A100s [A] which is not feasible in academic compute budgets. We have noted this in our limitation section (Supp Sec. I).\n\n\n$~$\n> W2 (1): 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. \n\nThe reviewer is correct—**distillation is important** for grafting, and **isolating its contribution** would require **training the hybrid architectures from scratch** for a head-to-head comparison. We note that pretraining is beyond the scope of this work, as we focus on **post-training architectural editing** of diffusion transformers. While an interesting direction, even training a 64x64 DiT takes more than 12 days on 32×A100s [A], making pretraining infeasible on an academic budget. That said, we agree this is an intellectually important question. To address it, we highlight independent evidence from language modeling showing that **hybrid architectures mixing attention with Hyena-style operators (including Hyena-SE) can be trained from scratch and remain competitive.**\n\n**Evidence: StripedHyena2: 7B-scale language models trained with Hyena-SE [B]:** StripedHyena2 [B] trains large-scale hybrid models from scratch using setups which mix attention and Hyena operators, including Hyena-SE. At the 7B scale, these models achieve **lower perplexity than transformer baselines** (2.83 vs. 3.09) after 400B tokens of pretraining—demonstrating that hybrid architectures remain expressive at scale.\n\n$~$\n> W2 (2) 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\nWe thank the reviewer for the thoughtful question. For clarity, **the goal of this work is not to outperform the base model but to investigate whether pretrained DiT architectures can be edited post-training under small compute budgets and to characterize the resulting quality–efficiency trade-offs.**\n\nPrior to this work, it was unknown whether such post-training architecture editing could preserve generation quality at all. Our results provide the first systematic study of post-training architectural edits on diffusion transformers (Tables 4, 5, and 6), mapping the quality–efficiency frontier in this setting.\n\nTo summarize, the key outcomes of our work:\n\n- **Table 4 & Fig. 3 (Right):** We obtain post-training hybrid architectures that are within 0.5 FID of the DiT-XL/2 baseline, while grafting 50% of MHA operators (Hyena-SE/X/Y, SWA) and 100% of MLP operators (MLP r=3) with cheaper alternatives (<2% pretraining compute).\n- **Table 5:** On PixArt-Sigma, we obtain a 1.43× faster hybrid variant via operator replacement, with less than 2% GenEval score drop at 2048×2048 resolution.\n- **Table 6:** We restructure DiT-XL/2 by trading depth for width, producing a **2× shallower model** (2× faster with fused implementation, bs=1) that achieves **better FID (2.77)** and outperforms pruning baselines and shallow-from-scratch trained models at depth=14. \n\nFuture work will explore pretraining and scaling of hybrid DiT architectures to improve model quality.\n\n$~$\n> Q1: 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\nReplace **at once is recommended**; we observed only **minor improvements** over sequential replacements (we tried this—**FID improvements were < 0.07 for full MHA self-grafting**). Reviewer is correct, replacing more components can propagate error because stage 1 regression loss is not zero for every operator, but **Stage-2 (lightweight fine-tuning)** handles this. For example, under **full self-grafting** of MHA and MLP we achieve **near-baseline** performance (Table 2) after stage 2.\n\n\n$~$\n> Q2: Are the hyperparameters from the original (base) model typically retained, or is it generally necessary to re-tune them after grafting?\n\nYes, base hyperparameters are reused. The only changes are (1) a simple linear warmup to the original base learning rate during Stage 2 (e.g., warming up to 1e‑4), and (2) a weight decay of 1e‑5 in stage 2. After warmup, the learning rate remains constant, following the original DiT recipe. Full hyperparameters for both Stage 1 and Stage 2 are provided in Supplement B.1.\n\n####  \n$~$\n> Suggestion: 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\nThank you for the suggestion. We will include this in the final version.\n\n\n**We hope our response addresses the reviewer's questions, and we’re happy to provide further details if helpful.**\n\n====\n\n[A] Karras, Tero, et al. \"Elucidating the design space of diffusion-based generative models.\" *NeurIPS 2022*\n\n[B] Ku, Jerome, et al. \"Systems and algorithms for convolutional multi-hybrid language models at scale.\" *arXiv preprint arXiv:2503.01868* (2025)."}, {"review_id": "8I9IG9ODFA", "reviewer": "Reviewer_JnND", "summary": "This 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.", "questions": "1、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.", "limitations": "The authors has discussed the limitations", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths\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.", "was_revised": true, "final_justification": "No", "rebuttal": "We thank the AC and the reviewer for their valuable time and feedback. In what follows, (1) we clarify the scope/goal of this work and address (2) requests for grafting experiments beyond image generation, (3) suggestions to include additional metrics beyond GenEval for T2I evaluation, and (4) clarifications on experimental details. \n\n$~$\n> W1: Since 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. \n\nWe thank the reviewer for the thoughtful question. We agree that model quality is one of the important goals in architecture design. However, the goal of this work is not to outperform the base model, but to investigate whether diffusion transformer architectures can be edited post-training under small compute budgets, and to characterize the resulting quality–efficiency tradeoffs.\nPrior to this work, it was unknown whether such post-training architecture editing could preserve generation quality at all. Our results provide the first systematic study of post-training architectural edits on diffusion transformers (Tables 4, 5, and 6), mapping the quality–efficiency frontier in this setting.\n\nWhile some quality drop is expected when replacing expensive operators with cheaper alternatives, we quantify how much can be traded off, and demonstrate that architectural editing can preserve or closely match baseline generation quality.\n\nTo summarize, the key outcomes of our work:\n\n- **Table 4 & Fig. 3 (Right):** We obtain post-training hybrid architectures that are within 0.5 FID of the DiT-XL/2 baseline, while grafting 50% of MHA operators (Hyena-SE/X/Y, SWA) and 100% of MLP operators (MLP r=3) with cheaper alternatives.\n\n- **Table 5:** On PixArt-Sigma, we obtain a 1.43× faster hybrid variant via operator replacement, with less than 2% GenEval score drop at 2048×2048 resolution.\n\n- **Table 6:** We restructure DiT-XL/2 by trading depth for width, producing a 2× shallower model (2× faster with fused implementation, bs=1) that achieves better FID (2.77) and outperforms pruning baselines and shallow-from-scratch trained models at depth=14. To our knowledge, this is the first attempt to convert sequential computation into parallel in DiTs post-training, enabling computation to be restructured.\n\n\n$~$\n> W2: Generalization beyond image generation\n\nWe note that Generalization was a recurring theme raised by two other reviewers.\n\n- Reviewer r7Mh suggested experiments beyond “diffusion-based image generation” and “architectures beyond DiT”. \n- Reviewer VANM suggested “more open-source models with different architectures” with video generation as one possible experiment.\n\n**To address these concerns collectively, we evaluate grafting in the language modeling setting—a task, architecture, and modality distinct from diffusion-based image generation with DiTs.**\n\n**Setup:**\n\n- **Task:** Diffusion-based image generation → Autoregressive language modeling (next-token prediction)\n- **Architecture & Scale:** DiT-XL/2 (0.7B) → Qwen3 (4B)\n- **Operator:** MHA → Sliding Window Attention (window=256)\n- **Dataset**: Alpaca-cleaned (50k instruction-following examples)\n- **Evaluation:** PiQA, ARC-e, ARC-c, HellaSwag, Winogrande, MMLU following [F]\n- **Experiment details:** We apply our exact two-stage grafting procedure—activation distillation followed by lightweight fine-tuning. Stage 1 uses context length 1024; Stage 2 uses 8192.\n\n**Results:**\n\nOur grafted Qwen3-4B model achieves a **1.4x decode throughput** (at 8K context length with KV cache, Nvidia H100) with **less than 1% average performance drop**, demonstrating that grafting generalizes to new tasks, architectures, and modalities. We leave video generation to future work.\n\n| Model       | Stage        | Grafting Ratio | PiQA | ARC-e | ARC-c (acc. norm) | HellaSwag (acc. norm) | Winogrande | MMLU (5-shot) | Avg. | Decode throughput |\n| ----------- | ------------ | -------------- | ---- | ----- | ----------------- | --------------------- | ---------- | ------------- | ---- | ------- |\n| Qwen3       | Baseline     | N/A            | 74.9 | 80.5  | 54.0              | 68.5                  | 66.0       | 70.1          | 69.0 | \\-      |\n| SWA (w=256) | Random Init. | 50%            | 64.3 | 56.5  | 29.7              | 37.9                  | 50.9       | 25.9          | 44.2 | 1.4x    |\n| SWA (w=256) | 1            | 50%            | 74.8 | 80.4  | 53.4              | 68.2                  | 65.5       | 66.9          | 68.2 | 1.4x    |\n| SWA (w=256) | 2            | 50%            | 75.5 | 80.3  | 52.3              | 69.6                  | 67.0       | 66.9          | 68.6 | 1.4x    |\n\n\n$~$\n> W3: For text-to-image generation, only using GenEval as the metric may not be sufficient.\n\nWe thank the reviewer for this suggestion. While the reviewer did not specify a preferred metric, we note that GenEval is a widely adopted objective metric for evaluating prompt–image alignment in text-to-image generation [A,B]. Other common metrics such as IS and FID are less appropriate for high-resolution text-to-image tasks, as they rely on low-resolution ImageNet-1K features.\nTo follow best practices, our paper already includes GenEval scores and qualitative samples (see Supplementary). To further strengthen the evaluation, we now include a small-scale human study, following recent work [C, D].\nWe randomly sample 350 text–image pairs and ask three expert annotators to rate each image along two axes:\n- Prompt–image alignment: Does the image accurately reflect the text?\n- Realism: Does the image appear visually realistic?\n\n|          | Alignment | Realism | Speedup |\n| -------- | --------- | ------- | ------- |\n| Baseline | 0.87      | 0.85    | 1x      |\n| Hyena-X  | 0.82      | 0.79    | 1.43x   |\n\nThese results show that the grafted variant (Hyena-X) maintains good quality across both axes. Together, we provide both objective (GenEval) and human (alignment/realism) evaluations of generation quality.\n\n\n\n$~$\n> Q1: Why MHA and MLP prefer different regression loss?\n\nThis difference arises from benign overfitting. MLPs have 2× more parameters than MHA (10.6M vs. 5.3M), which makes them robust to outliers [E]. This behavior is also reflected in our validation loss curves (Supplementary D.3). For MHA, L1 regression consistently yields lower validation loss than L2, while the opposite holds for MLPs. We therefore adopt L1 for MHA and L2 for MLP. We will make this clear in the final version.\n\n$~$\n> Q2: 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\nThe reviewer is correct that the depth -> width restructured model has 6% more parameters (712M vs. 675M for DiT-XL/2) due to linear projections. However, **the goal of this case study is not to reduce parameter count, but to investigate whether deep, sequential transformer blocks can be restructured post-training into shallower, wider configurations that improve inference time while preserving quality.**\n\nIn our restructuring, every two sequential DiT blocks are rewired post-training to run in parallel—reducing the model’s depth from 28 to 14 and achieving a 2× speedup (fused implementation, bs=1). Despite the slightly higher parameter count, the model achieves better quality (FID = 2.77) while benefiting from increased parallelism—demonstrating a favorable depth–width tradeoff.\nThis experiment is not focused on parameter efficiency, but rather on post-training architectural restructuring—demonstrating, to our knowledge, **the first instance of block-level rewiring of pretrained DiTs**. For completeness, we include Table 6 below with speedup and FID metrics.\n\n| **Method**    | **Depth** | **A.R** | **Iters** | **IS ↑** | **FID ↓** | **sFID ↓** | **Prec. ↑** | **Recall ↑** | **Speedup ↑** | **Params ↓** |\n| ----------------------------------------- | --------: | ------: | --------: | -------: | --------: | ---------: | ----------: | -----------: | ------------: | -----------: |\n| DiT-L/2 |  24 |  42.7 | 1,000K | 196.26 | 3.73 |  4.62 | 0.82 | 0.54 | — |  458M |\n| U-ViT-L |  21 |  48.8 | 300K | 221.29 |  3.44 |  6.58 | 0.83 |  0.52 |  — |  287M |\n| DiT-B/2 | 12 |   64.0 | 1,000K |119.63 | 10.12 |  5.39 | 0.73 |0.55 |  — | 130M |\n| BK-SDM| 14 |  82.3 | 100K |141.18 | 7.43 |  6.09 | 0.75 | 0.55 |  2× |  340M |\n| TinyDiT-D14 |  14 |    82.3 | 500K |   198.85 | 3.92 | 5.69 | 0.78 |  0.58 | 2×| 340M |\n| TinyDiT-D14 w/ MKD |        14 |    82.3 | 500K |  234.50 | 2.86 | 4.75 | 0.82 |  0.55 |  2× | 340M |\n| DiT-XL/2 |  28 | 41.4 | 7,000K |   278.20 |  2.27 | 4.60 | 0.83 | 0.57 | 1× | 675M |\n| **Grafting** |14 | 164.6 | 100K |   231.91 | 3.12 |  4.71 | 0.82 |0.55 |2× | 712M |\n| **Grafting** |**14** | **164.6** |  **230K** |   **251.77** |  **2.77** | **4.87** |**0.82** |  **0.56** |       **2×** |         **712M** |\n\n**We hope our response addresses the reviewer's questions, and we’re happy to provide further details if helpful.**\n\n====\n\n[A] Yao, Yuan, et al. \"Diffusion Transformer-to-Mamba Distillation for High-Resolution Image Generation.\" arXiv (2025) \n\n[B] Tang, Bingda, et al. \"Exploring the Deep Fusion of Large Language Models and Diffusion Transformers for Text-to-Image Synthesis.\" CVPR 2025\n\n[C] Cai, Shengqu, et al. \"Diffusion self-distillation for zero-shot customized image generation.\" CVPR. 2025.\n\n[D] Chen, Junsong, et al. \"Pixart-Sigma\" ECCV, 2024.\n\n[E] Bartlett, Peter L., et al. \"Benign overfitting in linear regression.\" PNAS 2020.\n\n[F] Zhang, Michael, et al. \"LoLCATs: On Low-Rank Linearizing of Large Language Models.\" ICLR 2025"}]

Author rebuttal:

Meta-review: The paper introduces grafting, a technique to edit trained models in terms of replacing a set of their operators with (potentially lower-weight) alternatives. Whereas the method likely has some generality to it, the paper focuses on diffusion transformers. Empirically, it shows that grafting can speed up various models with fairly low drops in performance. The strength of the paper certainly is a very innovative approach with a lot of potential for impact, especially as it is foreseeable to be applicable to other types of generation beyond what is investigated in the paper. The main weakness of the paper seems to be that the edited operators seem to come at a performance cost and don't seem to be able to truly replace the original. However, this cost seems more than acceptable and the innovativeness of the idea outshines the weakness by far. 

Above sentiment is reflected by all reviewers, who, following the rebuttal, unanimously agree to accept the paper. They generally commend the paper for its well written nature and presentation, the novelty of the idea and the potential impact of grafting as a technique. Initial concerns all seemed to revolve around requests for further empirical evaluation and concerns over the observed drop in performance. However, following rebuttal discussions, all concerns were resolved and reviewers agree that the potential of the idea outweighs any remaining concerns. A particular concern on the experimentation, voiced by reviewer VANM, revolved around inclusion of further baselines and other method comparisons. In several rounds of discussions the authors have added a significant amount of further empirical evidence, broadening the scope of their analysis and further improving an already in-depth investigation. The reviewer has respectively decided to raise their score to a full accept recommendation and the AC views this discussion as also addressing some of the concerns of the other reviewers (who are already positive anyhow).

Overall, the paper presents an interesting and novel take on a very important problem. Given the well-written nature of the paper, the careful experimentation, and the high level of expected interest of a large community the AC recommends for the paper to be accepted with an oral presentation.
Decision: Accept (oral)