Structured Residual Connectivity Matters for Diffusion Transformers
Yuhe Liu, Xinyin Ma, Gongfan Fang, Songhua Liu, Xinchao Wang
cs.CV
2026-09-27
DiT residual routing over mirrored encoder layers adds <0.1% params, cuts training up to 1.73×, and drops REPA-XL/2 FID from 5.9 to 4.34 (1.39 with CFG).
U-Net diffusion models send shallow spatial detail to the decoder through long skip connections. Diffusion Transformers (DiTs) replaced that with a uniform residual stream: each block does x{l+1} = xl + fl(xl), so the final state is an equal-weight sum of every previous layer. Early spatial cues get washed out, and the backward topology is locked into one additive chain.
U-ViT and U-DiTs bring long skips back, with fusion weights that freeze after training. Attention Residuals in LLMs score sources from the input, but they route densely over all previous layers. The question from National University of Singapore and Shanghai Jiao Tong University is what kind of cross-layer path image denoising actually needs.
A DiT is split in half along depth into an encoder and a decoder, keeping the single-scale token grid. Encoder block outputs, including the patch embedding, stay in the graph as residual sources, so decoder-side gradients can flow straight back.
Before both self-attention and the MLP, a residual routing operator RMSNorms the current tokens and the chosen sources, scores them with a learned vector, and softmax-weights them into the residual. The skip can pull earlier features per token, location, and diffusion timestep, like looking up spatial cues from a shallow memory.
When every decoder layer can see all encoder sources (Figure 4), two patterns show up that differ from language models. Early-layer features keep high weight across the whole depth. The network also spontaneously prefers mirrored layer pairs. The final design, mirror routing, gives decoder layer j only encoder layer K-j+1, plus encoder-internal routing over earlier encoder states. Dense all-to-all routing has to store O(L) sources. A static U-ViT-style skip uses about 77× more extra parameters and still cannot change fusion with the input. Mirror routing writes the pairing into the topology and keeps the weights content-adaptive.
Class-conditional ImageNet 256×256, with no classifier-free guidance (CFG) unless noted. CFG mixes in an unconditional prediction at sampling time to strengthen the condition.
| Setup | Steps | FID↓ | Baseline |
| DiT-S/2 + Ours | 400K | 62.48 | 69.81 |
| DiT-B/2 + Ours | 400K | 36.33 | 44.21 |
| DiT-XL/2 + Ours | 400K | 15.07 | 18.85 |
| DiT-XL/2 + Ours | 800K | 11.40 | baseline 13.71 @800K, 11.76 @1300K |
| REPA-XL/2 + Ours | 4M+0.28M | 4.94 | continued FT 5.87 |
| REPA-XL/2 + Ours | 4M+0.35M | 4.34 | released ckpt 5.9 |
Relative FID drops grow with scale: 10.5%, 17.8%, and 20.1% on S / B / XL. Matching a given FID takes up to 1.73× fewer steps. Extra parameters are 0.019M on S and 0.13M on XL, under 0.1%; FLOPs rise 0.3%.
REPA (Representation Alignment, which matches SiT-XL/2 hidden states to an external vision encoder) slows down late: after 6.4 FID at 1M steps, another 3M steps only reach 5.9, a 0.5 drop. Continued fine-tuning for 0.28M steps barely moves (5.87). With routing, 0.28M steps reach 4.94 and 0.35M steps reach 4.34. Under guidance interval [0, 0.7], the 0.28M model gets 1.39 FID with CFG; REPA at 4M is 1.42.
On DiT-S/2 at 400K, a U-ViT-style static mirror skip scores 67.90 and tracks the baseline the whole way. Adaptive fusion on the same pairs without encoder-side routing scores 64.83. Dense Attention Residuals score 62.87 but hold 1.51 GB of decoder source activations. Mirror routing scores 62.48 with 289 MB, about 81% less. All encoder sources: 63.48. First layer only: 63.51. Last layer only: 66.05. Mirror wins this set.
CKA (Centered Kernel Alignment, a similarity score between two layers' representations) shows vanilla DiT-XL/2 is highly homogenized across depth. Routing cuts that excess similarity, makes the anti-diagonal mirror pairing clearer, and raises gradient cosine similarity on early-late mirror pairs while middle layers stay differentiated.
This is a cheap wiring change for DiT / SiT: same single-scale backbone, no new data. It helps both from-scratch training and fine-tuning a REPA checkpoint, so the gain is in the paths and the gradient structure. The extra capacity is negligible.
For a DiT-family model that has already plateaued, adding a thin module and fine-tuning for a few hundred thousand steps can still shave more than a point of FID. Against bolting U-Net skips back on, the parameter and FLOP overhead is an order of magnitude smaller and the quality is better.
It is still an incremental architecture paper. The task is class-conditional ImageNet. Whether to follow it depends on whether you are training or fine-tuning a DiT-style backbone.
Larger text-to-image and text-to-video models are untested. Mirror routing may not stay optimal once stronger multimodal conditioning or temporal dependence shows up. That is the authors' own caveat.
All numbers are ImageNet 256×256. REPA fine-tuning drops the original projection loss; the control is continued training with that loss also off and no routing. Fine-tuning with the full REPA objective plus routing was not run. The CFG gap of 1.39 vs 1.42 is small; IS falls from 305.7 to 297.4, and Precision / Recall dip slightly, so the system-level win is mostly FID. The gradient-alignment plot uses a probe loss (output norm on four random latents at t=500). That mainly describes optimization behavior induced by the mirror links, not independent causal evidence. The U-ViT baseline is concat-plus-linear inside a single-scale DiT, with no original downsampling, which may understate hierarchical skips.