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 major architectural components, but their inherited residual stream remains insufficiently studied as a mechanism for cross-layer information flow. The paper analyzes this routing and introduces Diffusion-Adaptive Routing, improving SiT-XL/2 by 2.11 FID while matching baseline quality with 8.75× fewer iterations.

  • Problem

    The cross-layer information flow governed by DiTs’ inherited pre-normalized residual stream remains insufficiently investigated despite extensive study of other architectural components.

  • Method

    Diffusion-Adaptive Routing replaces standard residual addition with learnable, timestep-adaptive, non-incremental aggregation over sublayer outputs.

  • Results

    2.11 FID improvement over vanilla SiT-XL/2 accompanies matching baseline converged quality with 8.75× fewer iterations, while DAR adds 2× early-stage acceleration over REPA alone.

  • Takeaways & Limitations

    Cross-layer information routing emerges as an underexplored diffusion-modeling design axis that complements existing representation-alignment objectives.

  • Takeaways & Limitations

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

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

This work identifies cross-layer information routing as an underexplored DiT design axis, showing that default residual addition causes depth- and timestep-dependent inefficiencies. It proposes Diffusion-Adaptive Routing (DAR), a learnable, timestep-adaptive, non-incremental residual replacement that improves convergence and quality while complementing existing methods such as REPA.

  • Method: DAR elevates cross-layer information routing from an inherited convention to an explicit DiT design axis while operating along depth and preserving a homogeneous Transformer stack.The method is also compatible with modern Transformer enhancements such as REPA.
  • Diagnosis: A systematic analysis along depth and denoising timestep finds three residual-addition symptoms: monotonic hidden-state magnitude inflation, sharp backward-gradient decay, and increasing block-wise redundancy.These symptoms persist throughout training and vary systematically with noise level.
  • Method: DAR replaces running residual accumulation with learnable, timestep-adaptive, non-incremental aggregation over preceding sublayer outputs.It uses softmax attention whose query comes from the current adaLN-modulated hidden state, inheriting content and timestep dependence from DiT conditioning.
  • Results: 7.56 FID with SiT-XL/2 improves over the baseline by 2.11 FID on ImageNet 256×256.DAR also matches the baseline’s converged quality in roughly 8.75× fewer training iterations.
  • Results: 2× training acceleration in the early stage results when DAR is stacked on REPA, over REPA alone.This supports the reported complementarity between depth-wise routing and representation-alignment objectives.

2 Related Work

This section situates the paper within prior work on cross-layer information routing and Diffusion Transformer architectures. Prior routing research primarily modifies residual strength or propagation structure, while DiT research develops scalable latent-space Transformer denoisers.

  • Evolution of Cross-Layer Information Routing: Cross-layer routing began with fixed additive residual connections and later introduced gated or scaled variants such as ReZero, LayerScale, and DeepNorm.These methods improve optimization stability by adjusting residual strength without fundamentally changing the routing topology.
  • Evolution of Diffusion Transformers: U-ViT established Transformer denoisers that treat noisy image patches, timesteps, and conditions as tokens while retaining long skip connections.This work represented an early evolution from ViT-style U-Net replacements toward Transformer-based diffusion modeling.
  • Evolution of Diffusion Transformers: DiTs simplified U-ViT’s design into a pure latent-space Transformer and established clear scaling behavior.Subsequent DiT research progressed mainly through architectural improvements for multimodal conditioning and related design directions, as indicated in the passage.

3 Diagnosing Cross-Layer Information Flow in DiTs

An empirical analysis of vanilla SiT-XL/2 and static DAR diagnoses three depth-related symptoms of standard residual routing: forward magnitude inflation, backward gradient decay, and block-wise representational redundancy. These findings indicate that standard residuals propagate information but lack adaptive control over which previous representations to emphasize.

  • Experimental setup: The analysis compares vanilla SiT-XL/2 with static DAR using checkpoints after 600K iterations and diagnostics from 4096 ImageNet samples across 28 Transformer blocks.For each block, the study records forward hidden-state magnitude, backward gradient magnitude, and block-wise similarity.
  • Depth-wise symptoms: ~100× inflation: standard-residual forward hidden-state magnitude rises monotonically from ~15.5 at block 1 to ~1576 at block 28.The magnitude increase is reported as roughly 100× across depth.
  • Depth-wise symptoms: More than an order of magnitude lower: later-block backward gradients fall sharply after the first five blocks from early signals of ~5 × 10−7 and remain near zero.The pattern indicates substantially weaker optimization signals for deeper layers.
  • Depth-wise symptoms: >0.9 similarity: per-token cosine similarity between consecutive block outputs stays above 0.9 throughout the deep stack, indicating highly similar and redundant representations.The similarity pattern is associated with substantial representational redundancy under standard residual routing.
  • Diagnosis: Standard residual routing is associated with PreNorm dilution, imbalanced gradient propagation across depth, and high feature similarity and redundancy.The findings suggest that standard residuals provide cross-layer propagation but lack adaptive control over which previous representations should be emphasized.

4 Exploring Cross-Layer Interaction Spaces in DiTs

This section frames cross-layer routing as a fundamental DiT design dimension and contrasts fixed standard residuals and manually specified U-Net-like skips with the learnable, timestep-adaptive aggregation of DAR. DAR uses softmax attention over sublayer outputs, with query conditioning and chunking tailored to diffusion models.

  • Standard residual routing: Standard DiTs route information through fixed unit-coefficient residual accumulation, preventing depth- or timestep-specific retrieval and suppression of earlier representations.This routing is already a form of cross-layer information routing, but its coefficients are not learned or adaptive.
  • U-Net-like skip routing: U-Net-like skip routing enables multi-level feature fusion but relies on manually specified connections that weaken Transformer homogeneity and scalability.The routing topology pairs deep layers with selected shallow representations through a skip-fusion operation.
  • Diffusion-Adaptive Routing: DAR replaces unit-weight residual summation with softmax-weighted aggregation over the input embedding and preceding sublayer outputs.For source vi, DAR forms a key using RMSNorm(vi), computes softmax weights over the available source set, and feeds the aggregate into the next transformation.
  • Query parameterization: DAR’s dynamic query parameterization is motivated by the diffusion timestep dimension, which is absent from the LLM setting and changes how per-layer queries should be conditioned.The design considers static and dynamic query choices, with timestep information entering dynamically through the noised latent and DiT conditioning pathway or explicitly through timestep embeddings.
  • Chunked aggregation: Chunked DAR reduces source memory from O(Ld) to O((S + N)d) by combining prior chunk summaries with full intra-chunk sources.Each chunk is summarized by its last sublayer output, while the current chunk retains all preceding sources within that chunk.

5 Experiments

Experiments show that DAR improves DiT quality and convergence, with timestep-aware routing providing strong benefits and compatibility with REPA. Additional studies support its routing mechanism, identify an effective chunk size, and show preservation of high-frequency details during distillation.

  • Quality and convergence: FID 6.92 (SDE) without CFG is achieved by static DAR after 600K iterations, versus vanilla SiT-XL/2 trained for 1.75M iterations.Dynamic DAR attains the best ODE FID with CFG at 2.05.
  • Comparison with skip pathways: 0.77 FID lower than U-DiT-L is achieved by DAR static c4 under SDE with CFG, using only 0.83× as many parameters.Under ODE, DAR dynamic c4 improves over U-ViT-H/2 by 0.24 FID.
  • Timestep awareness: Timestep-aware dynamic and explicit-injection variants substantially outperform timestep-unaware routing, supporting timestep adaptation as central to DAR.Dynamic routing derives its query from the latest sublayer output, while explicit injection adds DiT’s timestep embedding without additional parameter cost.
  • Timestep awareness: R2 above 0.95 within the first five blocks and close to 1.0 throughout the deep stack shows that timestep information is linearly decodable from DAR-Dynamic router inputs.Both attention and MLP aggregator inputs outperform the raw input latents baseline of R2 ≈0.80 at every depth.
  • Chunked aggregation: S = 4 gives the best matched-compute result in the chunk-size sweep, producing a U-shaped pattern across S ∈ {1, 4, 8}.The cost decomposition combines routing precision and source compression, yielding a unique global minimizer at S⋆.
  • Distillation: High-frequency details, including sharp edges and fine textures, are preserved when DAR is applied to Qwen-Image during Distribution Matching Distillation.The result addresses detail attenuation during aggressive few-step distillation.

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

The paper investigates cross-layer information routing in diffusion transformers and identifies limitations of the inherited pre-normalized residual stream, motivating DAR as a timestep-adaptive replacement. It also situates diffusion models within their evolution from finite-step Gaussian corruption and variational learning to continuous-time SDE and probability-flow ODE formulations.

  • 6 Conclusion: The study analyzes cross-layer information routing in DiTs jointly across network depth and denoising timestep.
  • 6 Conclusion: The inherited pre-normalized residual stream exhibits forward magnitude inflation, backward gradient decay, and block-wise redundancy.
  • 6 Conclusion: DAR replaces the residual stream with learnable, timestep-adaptive, non-incremental aggregation.
  • A.1 Diffusion Models: Diffusion models were initially formulated as finite-step Markov chains that progressively corrupt data with Gaussian noise.Their reverse process was learned via a variational bound.
  • A.1 Diffusion Models: Score matching unified the diffusion view by recasting forward and reverse processes as a continuous-time SDE.
  • A.1 Diffusion Models: The continuous-time formulation has an equivalent probability-flow ODE, while later work refined noise schedules and parameterizations.

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

The proposition proves that the chunked-aggregation cost L(S) is U-shaped with a unique global minimizer S⋆. For SiT-XL/2, the predicted optimum S = 4 agrees with empirical results, motivating its use in the main experiments.

  • B Proof of Proposition and Empirical Verification: L(S) is strictly decreasing before S⋆ and strictly increasing after S⋆, establishing a unique global minimizer at S⋆.The derivative satisfies L′(S) < 0 for S < S⋆, equals zero at S = S⋆, and is positive for S > S⋆.
  • B Proof of Proposition and Empirical Verification: The proof determines the sign of L′(S) from the derivative expression after noting that S(S2 + L) > 0 for all S > 0.This sign analysis yields the decreasing-then-increasing behavior required for the proposition.
  • B Proof of Proposition and Empirical Verification: S⋆∈[3.7, 4.9] is predicted for SiT-XL/2 with depth 28, two sublayers per chunk, L = 56, and α ∈[0.4, 0.6].This identifies S = 4 as the model-predicted optimum over the realistic α range.
  • B Proof of Proposition and Empirical Verification: S = 4 improves on both S = 1 and S = 8, respectively representing no chunking and over-compression.At S = 1, the routing-entropy term dominates; at S = 8, |S| is small but per-summary distortion is large.
  • B Proof of Proposition and Empirical Verification: The empirical agreement between the predicted and observed optimum leads the main experiments to use S = 4 throughout.The passage explicitly states that the model-predicted optimum agrees quantitatively with Tab. 4.
  • B Proof of Proposition and Empirical Verification: Substantially deeper DiT backbones are conjectured to require a proportionally larger chunk size.This is presented as a further conjecture following the empirical verification for SiT-XL/2.

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

The experiments use a standardized SiT training setup and controlled optimization settings, with compute details reported for DAR. DAR’s implementation includes a dedicated final aggregator, while REPA experiments use a parameter-sharing variant that improves performance.

  • C.1 Experimental Configuration: ImageNet experiments follow the SiT training recipe unless otherwise specified, using a global batch size of 1024, learning rate 1×10−4, and bfloat16 mixed precision.Compared models use identical optimization and data-processing settings unless explicitly stated otherwise.
  • C.1 Experimental Configuration: REPA experiments adopt the original REPA configuration with DINOv2-B as the pretrained visual encoder.The passage also states that a representation-alignment setting is specified, but the supplied text ends before its value.
  • C.2 Compute Resources: Experiments run on NVIDIA H20 GPUs with 192 CPU cores and 1024 GB of system memory.These resources define the reported experimental compute environment.
  • C.2 Compute Resources: 0.59 seconds is the approximate per-step training time for the DAR-static-c4 configuration.The figure is reported for each training step.
  • C.3 Additional Implementation Details for DAR: DAR’s dedicated final aggregator combines prior chunk summaries with all raw sublayer outputs from the last chunk before prediction.Its input is represented as Sfinal containing chunk summaries and the last chunk’s raw sublayer outputs.
  • C.3 Additional Implementation Details for DAR: Unlike AttnRes, whose final layer aggregates only chunk summaries, DAR exposes recent sublayer outputs to preserve fine-grained task-specific information.This design yields about a 2-point FID gain after 200K training iterations.
  • C.3 Additional Implementation Details for DAR: REPA combinations instead reuse the last chunk’s MLP aggregator parameters and per-source RMSNorm for final aggregation, improving performance empirically.The reused parameters include the static/dynamic query and per-source RMSNorm.

D Details of Large-Scale T2I Model Post-Training

For large-scale T2I post-training, DAR is inserted into Qwen-Image’s MM-DiT backbone and trained with LoRA-based DMD. The setup uses four denoising steps, guidance scale 4.0, 1024^2 resolution, bfloat16 mixed precision, and a per-GPU batch size of 1.

  • Post-training setup: DAR is inserted into the MM-DiT backbone of Qwen-Image for large-scale T2I post-training using DMD.The model is fine-tuned with LoRA at rank 64.
  • Optimization: 5 × 10−6 is used for the student learning rate, alongside 2 × 10−6 for the fake branch.Both rates are specified for the LoRA fine-tuning setup.
  • Distillation configuration: 4 denoising steps and guidance scale 4.0 define the distillation configuration.Training uses 1024^2 resolution, bfloat16 mixed precision, and a per-GPU batch size of 1.

E Infrastructure

The DAR infrastructure replaces separate memory- and bandwidth-bound operations with fused Triton kernels that reduce data movement and integrate downstream normalization. At the SiT-XL/2 working point, this implementation substantially improves latency and activation memory while remaining numerically equivalent to the reference path.

  • Kernel fusion: A single Triton kernel fuses DAR’s forward path, replacing separate RMSNorm, query–key dot-product, softmax, and weighted-sum kernels.The baseline materializes [N, B, T, D]-shaped intermediates in HBM and reads the source tensor four times per forward pass.
  • Performance: 11.5× forward latency reduction lowers runtime from 22.5 ms to 1.96 ms for the dynamic variant at the SiT-XL/2 working point.The backward implementation also reduces HBM reads by streaming sources twice and recomputing RMSNorm intermediates on the fly.
  • Performance: 8.5× backward latency reduction lowers runtime from 115.8 ms to 13.6 ms for the dynamic variant.The fused kernel also incorporates downstream LayerNorm and adaLN modulation.
  • Memory efficiency: 78.7% forward and 74.6% backward peak activation-memory reductions are achieved for the dynamic variant, reaching 82.1% for the static variant.Memory savings increase monotonically with the number of routed sources N.
  • Compatibility: The fused kernel is numerically equivalent to the reference PyTorch path up to floating-point reordering and serves as a drop-in replacement.The implementation keeps the chunked aggregator viable as DiT backbones become deeper.

F Limitations and Future Work

Future work should test DAR across the two dominant scaling axes of modern generative Transformers: large-scale pretraining and large-scale post-training. The paper anticipates greater benefits in deeper, multi-billion-parameter models and reports preliminary evidence that adaptive routing preserves details during distillation.

  • Future Work: DAR’s most compelling next step is systematic evaluation on large-scale pretraining and post-training for modern generative Transformers.These are identified as the two scale axes dominating current generative Transformer development.
  • Large-Scale Pretraining: Several-billion-parameter T2I and T2V backbones with deeper Transformer stacks may exhibit more severe PreNorm-dilution symptoms than SiT-XL/2.The paper specifically names MM-DiT, Qwen-Image, FLUX, and HunyuanVideo as relevant targets.
  • Large-Scale Pretraining: The paper proposes multi-billion-parameter MM-DiT and video-DiT pretraining as an informative scaling study because DAR’s headroom may widen with depth.Eq. (9) predicts that the optimal chunk size grows as L.
  • Large-Scale Post-Training: Preliminary DMD experiments on Qwen-Image suggest that DAR’s better-conditioned gradient flow preserves details during otherwise brittle distillation.This experiment is described as being equipped with DAR and reported in Appendix D.
Loading 2605.20708v2…