PaperlayerView exact revision on arXiv

Source-linked AI summary

Rethinking Cross-Layer Information Routing in Diffusion Transformers

Chao Xu, Maohua Li, Qirui Li, Yixuan Xu, Yanke Zhou, Yunhe Li, Cuifeng Shen, Hanlin Tang, Kan Liu, Tao Lan, Lin Qu, Shao-Qun Zhang

arXiv:2605.20708v2cs.CVcs.AI

TL;DR

DiTs have extensively revisited many design choices while retaining the Transformer’s residual stream. This paper studies cross-layer information flow and proposes Diffusion-Adaptive Routing (DAR), which improves SiT-XL/2 by 2.11 FID and matches baseline quality with 8.75× fewer training iterations.

  • Problem

    DiT design has extensively revisited major components, while its pre-normalized residual stream remains inherited from the original Transformer.

    Source: §1 · p. 2

  • Method

    DAR replaces standard residual accumulation with learnable, timestep-adaptive, non-incremental aggregation over sublayer outputs.

    Source: §1 · p. 3, §6 · p. 10

  • Results

    2.11 FID improvement over vanilla SiT-XL/2 (7.56 vs. 9.67), with baseline converged quality matched using 8.75× fewer training iterations.

    Source: §1 · p. 3

  • Takeaways & Limitations

    Cross-layer information routing is an underexplored diffusion-model design axis that complements existing representation-alignment objectives.

    Source: §1 · p. 3, §6 · p. 10

    The dynamic variant assumes that the most recent sublayer output retains sufficient timestep information for its routing query.

    Source: §5.3 · p. 8

Abstract

from arXiv · show

Diffusion Transformers (DiTs) have become a de facto backbone of modern visual generation, and nearly every major axis of their design -- tokenization, attention, conditioning, objectives, and latent autoencoders -- has been extensively revisited. The residual stream that governs how information accumulates across layers, however, has been directly inherited from the original Transformer. In this paper, we present a systematic empirical analysis of cross-layer information flow in DiTs, jointly along depth and denoising timestep, and identify three concrete symptoms of traditional residual addition, namely monotonic forward magnitude inflation, sharp backward gradient decay, and pronounced block-wise redundancy. Motivated by this diagnosis, we propose Diffusion-Adaptive Routing (\textsc{DAR}), a drop-in residual replacement that performs \emph{learnable, timestep-adaptive, and non-incremental} aggregation over the history of sublayer outputs. Moreover, the proposed \textsc{DAR} is compatible with many modern Transformer enhancement methods, such as REPA. On ImageNet $256\times256$, \textsc{DAR} improves SiT-XL/2 by $2.11$ FID ($7.56$ vs.\ $9.67$) and matches the baseline's converged quality with $8.75\times$ fewer training iterations. Stacked on top of REPA, it yields a $2\times$ training acceleration in the early stage, suggesting cross-layer information routing as an underexplored design axis in diffusion modeling, one that operates orthogonally to existing representation-alignment objectives. Beyond pretraining, \textsc{DAR} can also be applied during the fine-tuning stage of large-scale T2I models and preserves high-frequency details during Distribution Matching Distillation.

1 Introduction

The introduction argues that DiT residual routing is an overlooked design axis: conventional addition shows depth- and timestep-related problems, motivating a learnable, timestep-adaptive replacement. It presents DAR as a compatible method with reported gains in quality and training speed.

1 Introduction

  • Because relevant representations shift from coarse structure at high noise to fine detail at low noise, DiT aggregation should be learnable, timestep-adaptive, and non-incremental rather than fixed and uniform.Prior U-Net-style long skips [2] [4] [54] restore fixed hierarchical bias rather than providing dynamic timestep-aware aggregation.Source: §1 · p. 2
  • The paper frames cross-layer information routing as an explicit DiT design axis and reports the first systematic study of its flow jointly across depth and denoising timestep.The analysis suggests residual-stream effects extend beyond stabilizing deep training and reveal noise-dependent structure in PreNorm dilution.Source: §1 · p. 2
  • The study identifies monotonic hidden-state magnitude inflation, sharp backward-gradient decay, and increasing block redundancy, persisting through training and varying systematically with noise level.The authors describe this as a joint analysis across depth and denoising timestep.Source: §1 · pp. 2–3
  • DAR replaces each running residual with softmax attention over preceding sublayer outputs, using a query from the current adaLN-modulated hidden state to enable content- and timestep-dependent routing.It preserves the isotropic, homogeneous Transformer stack and is compatible with enhancement methods such as REPA.Source: §1 · p. 3
  • On ImageNet 256×256, DAR achieves 7.56 FID with SiT-XL/2, 2.11 lower than baseline, matches converged quality in 8.75× fewer iterations, and accelerates early REPA training 2×.The reported gains are described as complementary to representation-alignment objectives.Source: §1 · p. 3

2 Related Work

Related work traces cross-layer routing from fixed residual addition toward learned, selective aggregation, while DiT research has focused on backbone scaling, conditioning, and training objectives, leaving residual pathways largely unchanged.

Evolution of Cross-Layer Information Routing

  • Routing research progresses from fixed residual recursion [16] [51], through gated or scaled methods [1] [55] and learned multi-stream mixing [70] [60], to direct depth aggregation [20] [40] [53].Hyper-Connections [70] adds learned multi-stream mixing, refined by mHC [60] with doubly stochastic constraints; DenseNet [20], DenseFormer [40], and Attention Residuals [53] provide direct access to earlier representations or depth-wise aggregation.Source: §2 · p. 3

Evolution of Diffusion Transformers

  • DiT work has emphasized backbone design and scaling, multimodal conditioning, and generative objectives, while leaving the residual pathway largely unchanged.U-ViT [2] treats noisy patches, timesteps, and conditions as tokens; DiT [41] establishes latent-Transformer scaling. Later work advances conditioning and fusion, including PixArt [6] [8], MM-DiT [13], and LLM-based encoders, and objectives such as SiT [35], Stable Diffusion 3 [13], and REPA.Source: §2 · p. 4

3 Diagnosing Cross-Layer Information Flow in DiTs

Diagnostics of DiTs across depth and denoising timestep reveal that standard residual routing exhibits magnitude inflation, gradient decay, and representational redundancy. Source importance varies with timestep, while DAR learns timestep-shifting weights over historical representations.

3 Diagnosing Cross-Layer Information Flow in DiTs

  • The diagnostics compare vanilla SiT-XL/2 with static DAR using chunk size S = 4, checkpointed at 600K iterations and evaluated on 4096 ImageNet samples.They measure hidden-state RMS, gradient magnitude, and cosine similarity between consecutive block outputs across 28 blocks.Source: §3 · p. 4
  • ∼15.5-to-∼1576 forward hidden-state RMS across blocks 1–28 (~100×), gradients dropping by over an order of magnitude, and cosine similarity above 0.9 indicate inflation, decay, and redundancy.The baseline’s early-block gradient magnitude is ∼5 x 10^-7; later-block gradients remain close to zero.Source: §3 · pp. 4–5
  • Baseline counterfactual source importance varies with timestep, while DAR’s softmax selects a small, timestep-shifting subset of historical sources at shallow and deep blocks.The baseline’s forward pass remains numerically identical when scalar gates are added to measure counterfactual importance.Source: §3 · p. 5
  • Standard residual routing is rigid: magnitude growth is associated with PreNorm dilution [30] [38] [53], imbalanced gradient propagation [53] [60] [70], and feature redundancy [5] [22] [36] [49].The findings suggest standard residuals propagate information across layers but lack adaptive control over which previous representations to emphasize or suppress.Source: §3 · p. 5

4 Exploring Cross-Layer Interaction Spaces in DiTs

DiTs already route information across layers, but standard residuals use fixed unit coefficients and U-Net-like skips impose manual connections. DAR replaces fixed accumulation with softmax-weighted aggregation over earlier sublayer outputs, with query and memory designs tailored to diffusion.

4.1 Cross-layer Routing in DiTs

  • Standard DiTs route earlier representations with fixed unit coefficients, so they cannot explicitly retrieve or suppress sources by depth or denoising stage.Source: §4.1 · p. 5
  • U-Net-like skip routing adds multi-level feature fusion [2] [4] [54], but its manually specified topology weakens Transformer homogeneity.Source: §4.1 · p. 6

4.2 Diffusion-Adaptive Routing

  • DAR replaces unit-weight residual accumulation with softmax-weighted aggregation over prior sublayer outputs, drawing on AttnRes [53].Each source is represented by its sublayer output, and the aggregated representation enters the next sublayer transformation.Source: §4.2 · p. 6
  • DAR offers pure static, explicit timestep-injection, and dynamic query variants, reflecting the timestep dimension distinctive to DiTs.Explicit injection adds the existing timestep embedding at no additional parameter cost; zero initialization of the embedder’s final layer makes the model start as the static variant.Source: §4.2 · p. 6
  • Chunked DAR summarizes completed chunks and retains full sources within the current chunk, reducing source memory from O(Ld) to O((S + N)d).The source set combines prior chunk summaries with current intra-chunk outputs.Source: §4.2 · pp. 6–7

5 Experiments

Experiments report that DAR improves SiT-XL/2 quality and convergence, benefits from timestep-aware routing, and compounds with REPA. Additional analyses examine chunk size and use DAR in large-scale text-to-image distillation.

5.1 Setup

  • Experiments train on ImageNet-1K at 256 × 256, retrain SiT-XL/2 under the same recipe for fair comparison, and use ODE results by default.Evaluation uses 50,000 samples and 250 function evaluations by default, with both ODE and SDE samplers.Source: §5.1 · p. 7

5.2 Better Quality and Faster Convergence

  • DAR’s static variant achieves FID 6.92 (SDE) without CFG after 600K iterations, while dynamic DAR attains ODE FID 2.05 with CFG, outperforming the SiT-XL/2 baseline trained for 1.75M iterations.A widened SiT-Plus matched to dynamic c4’s 752M parameters still trails DAR despite using 2× the training budget.Source: §5.2 · p. 8
  • Under SDE with CFG, DAR static c4 beats U-DiT-L by 0.77 FID using 0.83× its parameters, while under ODE dynamic c4 improves over U-ViT-H/2 by 0.24 FID.Unlike fixed skip topologies, DAR learns timestep-adaptive aggregation while preserving the homogeneous Transformer stack.Source: §5.2 · p. 8

5.3 Timestep Awareness Is Crucial for Routing in DiTs

  • Both implicit and explicit timestep-aware routing substantially outperform timestep-blind routing at matched compute, and the timestep is linearly decodable from router inputs across all 28 blocks.Probe R2 exceeds 0.95 within the first five blocks and remains close to 1.0 deeper in the stack, versus approximately 0.80 for raw input latents.Source: §5.3 · pp. 8–9, Figure 5 · p. 8

5.4 DAR Is Orthogonal to REPA

  • At 100K iterations, DAR+REPA improves FID from 9.89 to 7.09, and at 200K from 6.89 to 5.92, with 100K DAR+REPA matching 200K REPA alone.DAR+REPA at 200K also surpasses REPA alone at 300K, indicating compounding routing- and representation-level gains.Source: Table 3 · p. 8, §5.4 · p. 9

5.5 Chunk Size Choices

  • On SiT-XL/2 ImageNet 256×256 without classifier-free guidance, the chunk-size sweep at 300K iterations shows a U-shaped pattern with S = 4 at the minimum.The proposed cost decomposition attributes the pattern to competing routing-precision and source-compression costs, yielding a unique global minimizer.Source: §5.5 · p. 9, Table 4 · p. 9, B Proof of Proposition and Empirical… · p. 15

5.6 Large-Scale T2I Model Post-Training

  • With Distribution Matching Distillation on Qwen-Image, DAR helps preserve high-frequency details, including sharp edges and fine textures, during aggressive few-step distillation.The paper defers full setup and samples to Appendix D.Source: §5.6 · p. 9

6 Conclusion · A More Discussion on Related Work · A.1 Diffusion Models

The paper identifies residual-stream problems in DiTs and proposes DAR, a timestep-adaptive, non-incremental replacement that improves training results. The related-work discussion situates diffusion models among noise-reversal, latent-space, and flow-based approaches.

6 Conclusion

  • On ImageNet 256×256, DAR attains a best FID of 6.92 on SiT-XL/2, matches the baseline with 8.75× fewer iterations, and delivers a 2× early-stage speedup over REPA.The conclusion presents cross-layer routing as an underexplored design axis that complements advances in diffusion modeling.Source: §6 · p. 10
  • The inherited pre-normalized residual stream exhibits forward magnitude inflation, backward gradient decay, and block-wise redundancy; DAR replaces it with learnable, timestep-adaptive, non-incremental aggregation.The paper’s investigation analyzes cross-layer information routing jointly across depth and denoising timestep.Source: §6 · p. 10

A.1 Diffusion Models

  • Diffusion models were formulated as finite-step Gaussian-noise Markov chains reversed via a variational bound [19] [47], then connected to score matching through continuous-time SDEs and probability-flow ODEs [50].Later work refined noise schedules and parameterizations [23] [24] [39] [48].Source: §A.1 · p. 15
  • Latent diffusion reduces pixel-space cost by denoising in pretrained-autoencoder latent space [42], while flow-based methods learn deterministic velocity fields for transporting noise to data.Flow Matching regresses conditional velocity along a prescribed path [32], Rectified Flow favors straighter transport for few-step inference [34], and SiT unifies diffusion- and flow-based objectives under one interpolant framework [35].Source: §A.1 · p. 15
  • Stable Diffusion 3 scales rectified-flow training to large text-to-image models [13].This follows the broader shift toward deterministic transport formulations for generation.Source: §A.1 · p. 15
  • Diffusion models have become a dominant generative paradigm for image generation [58] [59], image editing [11] [28] [67], video generation [56] [62], and related visual synthesis.Source: §A.1 · p. 15

B Proof of Proposition and Empirical Verification · C Additional Implementation Details

The appendix proves that chunked-aggregation cost has a unique U-shaped optimum and reports that the predicted chunk size agrees with the empirical choice for SiT-XL/2.

B Proof of Proposition and Empirical Verification

  • L(S) is strictly decreasing below S⋆ and strictly increasing above it, so S⋆ is its unique global minimizer.The derivative is negative for S < S⋆, zero at S = S⋆, and positive for S > S⋆.Source: B Proof of Proposition and Empirical… · p. 15
  • S⋆∈[3.7, 4.9] predicts S = 4 for SiT-XL/2, matching Tab. 4; S = 4 improves on both S = 1 and S = 8 and is used throughout the main experiments.For depth 28 with two sublayers per chunk, L = 56; the prediction covers α ∈[0.4, 0.6].Source: B Proof of Proposition and Empirical… · p. 15
  • Because Eq. (9) predicts S⋆ to scale with L, the authors conjecture that substantially deeper DiTs will require proportionally larger chunk sizes.This is presented as a conjecture about scaling to deeper backbones.Source: B Proof of Proposition and Empirical… · p. 15

C.1 Experimental Configuration · C.2 Compute Resources · C.3 Additional Implementation Details for DAR

The appendix specifies the shared ImageNet and REPA training setups, compute environment, and DAR’s final-aggregation implementations, including a reported FID gain from exposing raw last-chunk outputs.

C.1 Experimental Configuration

  • ImageNet experiments follow SiT [35] with global batch size 1024, learning rate 1×10−4, and bfloat16 mixed precision.Compared models use the same optimization and data-processing settings unless stated otherwise.Source: §C.1 · p. 16
  • REPA experiments use DINOv2-B, alignment coefficient 0.5, and an alignment loss on the eighth-layer hidden representation, without changes beyond replacing residual routing with DAR.The setup follows the original REPA configuration.Source: §C.1 · p. 16
  • SiT and REPA baselines are rerun under the same experimental setting rather than compared using numbers copied from prior work.This supports fair comparisons under the stated optimization and data-processing settings.Source: §C.1 · p. 16

C.2 Compute Resources

  • Experiments use NVIDIA H20 GPUs with 192 CPU cores and 1024 GB of system memory; DAR-static-c4 takes approximately 0.59 seconds per training step.Source: §C.2 · p. 16

C.3 Additional Implementation Details for DAR

  • The dedicated final aggregator receives prior chunk summaries and all raw sublayer outputs from the last chunk as input to the prediction layer.This differs from AttnRes, whose final layer aggregates strictly over the chunk summaries.Source: §C.3 · p. 16
  • About a 2-point FID gain after 200K training iterations comes from exposing raw last-chunk outputs rather than only chunk summaries, unlike AttnRes.The passage attributes the gain to preserving fine-grained, task-specific information that chunk-level summaries could discard.Source: §C.3 · p. 16
  • For DAR with REPA, the last chunk reuses its own MLP aggregator parameters and per-source RMSNorm for final aggregation, an empirically motivated choice that improves performance.The reused parameters include the static or dynamic query.Source: §C.3 · p. 16

D Details of Large-Scale T2I Model Post-Training

This section specifies the DMD post-training setup used to apply DAR to Qwen-Image [58].

D Details of Large-Scale T2I Model Post-Training

  • DMD post-training applies DAR to Qwen-Image [58] in its MM-DiT backbone using rank-64 LoRA fine-tuning.The student and fake branches use learning rates of 5 × 10−6 and 2 × 10−6, respectively; distillation uses 4 denoising steps and guidance scale 4.0, at 1024^2 resolution with bfloat16 mixed precision and per-GPU batch size 1.Source: D Details of Large-Scale T2I Model… · p. 16

E Infrastructure

The infrastructure section introduces a fused Triton implementation of DAR that reduces latency and activation memory by streaming routed sources and keeping intermediates in registers. The fused kernel is numerically equivalent to the PyTorch reference up to floating-point reordering and is used throughout the paper.

E Infrastructure

  • A single Triton kernel uses an online-softmax recurrence to stream each routed source once, keeping per-source normalization, keys, dot products, and exponentials in registers instead of materializing intermediates in HBM.The baseline launches separate kernels and reads each source four times per forward pass, making it memory- and bandwidth-bound as source count grows with depth.Source: E Infrastructure · p. 16
  • At N=57, the fused kernel reduced forward latency from 22.5 ms to 1.96 ms (11.5×), backward latency from 115.8 ms to 13.6 ms (8.5×), and peak activation memory by 78.7% forward and 74.6% backward, with up to 82.1% savings for the static variant.The savings grow monotonically with N, keeping the chunked aggregator viable for deeper DiTs.Source: Figure 6 · p. 17, E Infrastructure · p. 17
  • The backward kernel streams sources twice to recover softmax statistics and recompute RMSNorm intermediates, emitting gradients in two HBM reads instead of four to five while fusing downstream LayerNorm and adaLN modulation.The implementation is numerically equivalent to the reference PyTorch path up to floating-point reordering.Source: E Infrastructure · p. 17

F Limitations and Future Work

The authors identify large-scale pretraining and post-training as the most compelling directions for extending DAR. They propose studying whether its benefits grow with model depth and generalize across diffusion post-training objectives and backbones.

Pretraining

  • DAR’s most natural next step is a systematic scaling study on multi-billion-parameter MM-DiT and video-DiT pretraining, where its headroom may widen with depth.The authors expect PreNorm-dilution to be more severe in deeper, larger backbones and Eq. (9) predicts that optimal chunk size grows with L.Source: F Limitations and Future Work · p. 17

Post-training

  • A preliminary Qwen-Image DMD experiment suggests DAR’s better-conditioned gradient flow preserves detail in brittle distillation where the vanilla counterpart diverges.The authors plan to test this observation across supervised fine-tuning, RL-style preference optimization, and few-step distillation on multiple large-scale T2I and T2V backbones.Source: F Limitations and Future Work · p. 17

Related papers

Loading 2605.20708v2…