Exploring Diffusion Transformer Designs via Grafting

TL;DR

Grafting is a two-stage procedure for editing the computational graph of a pretrained DiT to materialize new architectures without training from scratch. Stage 1 (activation distillation) regresses the new operator’s outputs onto the old operator’s activations on ~8k samples; Stage 2 (lightweight finetuning) absorbs cumulative error from multiple replacements on ~10% of pretraining data. Applied to DiT-XL/2 with swaps of MHA → gated convolution / SWA / Mamba-2 and MLP → variable-expansion or conv variants, hybrid architectures reach FID 2.38–2.64 (baseline 2.27) using <2% of pretraining compute. The same recipe gives a 1.43× speedup on PixArt-Σ 2048² T2I (<2% GenEval drop) and supports a more aggressive case study that converts every pair of sequential DiT blocks into parallel blocks, halving depth at FID 2.77.

Key claims

Activation distillation initializes a replacement operator by regressing its output onto the pretrained operator’s output on ~8k samples; the choice of regression objective is operator-specific — L1 best for MHA, L2 best for MLP — driven by activation-outlier statistics [§3.3, Tab. 1].
Lightweight finetuning on 10% of ImageNet-1K training data is sufficient to absorb the cumulative error from replacing all 28 MHA or all 28 MLP operators in DiT-XL/2, recovering FID within ~0.3 of the 2.27 baseline [§3.3, Tab. 2].
DiT-XL/2 attention is largely local: under a band-α locality metric, 15 of 28 layers attend ≥50% of mass within a band of α=0.25, with non-locality concentrated in the first two layers [§3.4, Fig. 2].
Naïvely replacing all MHA with sliding-window attention degrades FID from 2.27 to 53.9; interleaved 50% replacement with SWA, Hyena-X/Y, or Mamba-2 stays within 0.5 FID of baseline [§3.4, §4.2].
At 50% interleaved replacement, the band-α locality metric is a useful layer-selection heuristic: interleaved gives FID 2.74 (Hyena-X), top-local 3.02, low-local 3.18, and “last 14 layers” only 4.00 [§4.2, Fig. 4(b)].
Variable-expansion-ratio MLPs (ratio 3, 6) are the most robust replacement: even at 100% full replacement, ratio-3 MLPs achieve FID 2.66 [§4.2].
Grafting scales to high-resolution T2I: replacing MHA in PixArt-Σ with Hyena-X via interleaved grafting (layers 8/10/12/14/16/18/20–27) on 12k synthetic image-text pairs yields 1.43× wall-clock speedup at 2048×2048 with GenEval 47.78 vs. 49.75 [§5].
Sequential transformer blocks can be converted to parallel blocks via grafting: pairing every two DiT-XL/2 blocks and merging their outputs with a linear projection reduces depth 28→14 (with 6% more parameters) and reaches FID 2.77, beating DiT-L/2 trained from scratch (FID 3.73, depth 24) and TinyDiT-D14 (FID 2.86) [§6, Tab. 6].
The full grafting procedure for a hybrid model fits in <24 h on 8×H100, totaling <2% of the DiT-XL/2 pretraining compute [§4.1].

Method

Grafting frames architecture exploration as editing the computational graph of a pretrained model. Given a pretrained DiT and a candidate operator replacement (e.g., swap MHA for Hyena-X in half the layers), the procedure has two stages.

Stage 1 — activation distillation. For each layer where an operator is being replaced, the new operator $g_θ$ is trained to match the pretrained operator $f_l$ ‘s output on the same input: $\min_θ \mathbb{E}[\|f_l(x_t, c) - g_θ(x_t, c)\|]$ under L1, L2, or Huber regression. Inputs are 8k ImageNet samples passed through the original DiT to extract per-layer activations; each replacement operator trains in <30 min on one H100 in parallel across layers. The regression objective is operator-specific: L1 is better for MHA because attention activations have large outliers in deep layers, while L2 is better for MLPs because their 2×-larger parameter count damps outlier sensitivity.

Stage 2 — lightweight finetuning. Replacements are inserted into the computational graph and the full model is finetuned end-to-end on the original diffusion loss using ~10% of pretraining data and ~10% of pretraining iterations, at lr 1e-4 with linear warmup. This absorbs the cumulative error from inserting many replacements.

The paper introduces two new operators in the process: Hyena-X and Hyena-Y, local gated convolutions composed of short causal depth-wise 1D convolutions designed as drop-in MHA replacements (linear in sequence length). Hyena-X is also adapted as an MLP alternative by applying it along the channel dimension.

A separate diagnostic — the band-α locality metric — gives a per-layer scalar measuring how concentrated attention mass is near the diagonal, used to decide which layers are safe to replace with local operators. Only layers with high band-α are good candidates for SWA/Hyena replacement.

Results

On class-conditional ImageNet-256 with DiT-XL/2 (baseline FID 2.27):

Self-grafting (replace operators with random-init identical operators, recover via grafting): full MHA replacement FID 2.53, full MLP replacement FID 2.54 on 10% data [Tab. 2].
50% interleaved MHA replacement: SWA FID 2.67, Hyena-X 2.74, Mamba-2 2.65 (10% data); SWA improves to 2.62, Hyena-X to 2.61, Mamba-2 to 2.55 with 20% data [§4.2, Fig. 4(a)].
50% interleaved MLP replacement: variable-expansion-ratio FID 2.38–2.49, Hyena-X (channel) FID 2.63 [Tab. 4].
100% MLP replacement with expansion ratio 3: FID 2.66 [Tab. 4].

On 2048×2048 PixArt-Σ T2I, grafting MHA → Hyena-X via interleaved selection of 13/28 layers with LoRA rank-64 finetuning on 12k synthetic pairs: 1.43× wall-clock speedup, GenEval 47.78 vs. 49.75 (~2 points absolute drop) [§5].

Depth-to-width restructuring case study on DiT-XL/2: pair every two sequential blocks into one parallel pair (28→14 layers, +6% params), 25% data, 250k iterations. Result: FID 2.77, beating from-scratch DiT-L/2 (FID 3.73, depth 24) and U-ViT-L (FID 3.44, depth 21), and beating TinyDiT-D14 (FID 2.86) and BK-SDM (FID 7.43) at matched depth [Tab. 6].

Each hybrid experiment completes in <24h on 8×H100 (<2% of DiT-XL/2 pretraining compute) [§4.1].

Why it’s interesting

This is the first filed paper on the wiki that treats architecture exploration itself as a finetuning-budget problem rather than a from-scratch problem — distinct from but adjacent to Diffusion training efficiency‘s catalog of recipes for reaching target FID in fewer GPU-days. The wiki currently holds two threads where grafting is directly relevant: (a) the case for internal operator-level efficiency in DiTs (sliding-window attention, Mamba-2, Hyena variants) that Foveated Diffusion: Efficient Spatially Adaptive Image and Video Generation complements at the spatial-token level rather than the operator level; and (b) the parameter-efficient-finetuning thread where Text-to-LoRA: Instant Transformer Adaption and Doc-to-LoRA: Learning to Instantly Internalize Contexts use LoRA to graft new behavior in, while this paper uses lightweight finetuning to graft new architecture in. The depth→width case study is the cleanest result for Luma’s purposes: if a 2× depth halving at FID 2.77 holds at >1B-scale T2V, the inference-latency implications are large.