Rethinking Cross-Layer Information Routing in Diffusion Transformers

Abstract: 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.

• The paper introduces DAR, a learnable, timestep-adaptive routing mechanism that replaces standard residual addition in Diffusion Transformers.
• It empirically identifies degradation issues such as forward magnitude inflation, backward gradient decay, and block-wise redundancy in the residual stream.
• Experimental results show DAR improves ImageNet FID scores and accelerates convergence, with further gains when integrated with methods like REPA.

Rethinking Cross-Layer Information Routing in Diffusion Transformers

Overview and Motivation

The paper "Rethinking Cross-Layer Information Routing in Diffusion Transformers" (2605.20708) interrogates a longstanding architectural assumption in Diffusion Transformers (DiTs): the practice of standard residual addition for cross-layer information propagation. The authors note that, while most DiT design axes (tokenization, attention, conditioning, objectives, autoencoding) have seen extensive study and innovation, the residual stream inherited from NLP Transformers has remained essentially unaltered. Through rigorous empirical diagnosis, three degradation phenomena are identified: monotonic inflation of the forward hidden-state magnitude, sharp decay in backward gradients, and pronounced block-wise redundancy. These issues, reminiscent of PreNorm dilution in LLMs, point to the need for a fundamentally different approach to cross-layer information aggregation in the time-varying context of visual diffusion modeling.

Empirical Diagnosis of Standard Residuals

A comprehensive study decomposed the cross-layer information flow across both depth and denoising timestep in DiTs, particularly SiT-XL/2. The analysis reveals:

• Forward Magnitude Inflation: The RMS of block outputs increases by approximately two orders of magnitude across depth, resulting in the need for deeper layers to emit increasingly large outputs to retain influence on the normalized residual stream.
• Backward Gradient Decay: Gradients with respect to deep block outputs decay sharply, leaving late blocks with highly attenuated optimization signals.
• Block-wise Redundancy: High cosine similarity is maintained between adjacent block outputs throughout the stack, indicating substantial feature redundancy and representational stalling.

Furthermore, when probed along the denoising timestep axis, the routing importance of historical sublayer outputs varies systematically, suggesting that timestep-aware, adaptive aggregation is a latent requirement rather than an externally imposed bias.

Diffusion-Adaptive Routing (DAR): Architectural Innovation

Grounded in these findings, the authors propose Diffusion-Adaptive Routing (DAR)—a learnable, timestep-adaptive, and non-incremental residual mechanism. DAR replaces conventional additive residuals with a softmax-attention-based aggregation over preceding sublayer outputs:

$$h_l = \sum_{i=0}^{l-1} \alpha_i^{(l)}(t) v_i \quad \text{with} \quad \alpha_i^{(l)}(t) = \frac{\exp(q_l(t)^\top k_i / \sqrt{d})}{\sum_j \exp(q_l(t)^\top k_j / \sqrt{d})}$$

• Query Parameterization: DAR supports static, explicit timestep-injected, and dynamic query variants, with dynamic queries leveraging content and timestep information from the current hidden state.
• Chunked Aggregation: To mitigate memory scaling, DAR supports chunked aggregation, partitioning the sublayers into chunks and summarizing their outputs, with empirical and theoretical justification (optimal chunk size $S = 4$ for SiT-XL/2).
• Isotropy and Homogeneity: The DAR mechanism is inherently compatible with Transformer enhancement methods (e.g., REPA), preserves the homogeneous stack property, and avoids manually specified skip connections.

Experimental Results

DAR demonstrates both enhanced sample quality and dramatic training acceleration:

• FID Improvement: On ImageNet 256×256, DAR yields a best FID of 6.92 (SDE; static variant), improving SiT-XL/2 baseline by 2.11 points.
• Convergence Speedup: DAR matches baseline output quality with 8.75x fewer training iterations.
• Compatibility with REPA: Stacking DAR on REPA delivers up to 2x further early-stage training acceleration, suggesting synergistic rather than additive effects.
• Chunking Analysis: Optimal chunk size for aggregation is analytically and empirically identified; larger chunk sizes may benefit deeper Transformer stacks.

Ablations underscore the criticality of timestep awareness for routing efficacy—both explicit and dynamic query parameterizations outperform the static variant.

Theoretical and Practical Implications

The findings demonstrate that the default residual stream inherited from NLP Transformers is suboptimal for DiTs due to the explicit time-varying requirements of the denoising process. DAR's learnable, timestep-adaptive aggregation directly addresses forward magnitude inflation, backward gradient decay, and block-wise redundancy, while preserving architectural scalability and compatibility with prevailing representation-alignment objectives.

Practically, DAR proves effective not only during pretraining but also in large-scale T2I post-training, particularly with distribution matching distillation, where high-frequency details are preserved. The detailed infrastructure benchmarks reveal that DAR can be operationalized efficiently, with substantial latency and memory savings using a fused kernel implementation.

Future Directions

Extending DAR to multi-billion-parameter DiTs and video-generation backbones is a compelling direction, with theoretical predictions suggesting increasing chunk sizes may be optimal as depth scales. Furthermore, the improved gradient flow via adaptive routing may generalize to various post-training objectives, including RL-based preference optimization and supervised fine-tuning, establishing DAR as a foundational technique for future generative modeling pipelines.

Conclusion

"Rethinking Cross-Layer Information Routing in Diffusion Transformers" delivers a methodical assessment and remedy for a previously unexamined architectural bottleneck in DiTs. DAR offers a principled, empirically validated alternative to standard residual aggregation, yielding measurable improvements in training efficiency and final sample quality, and opening a new orthogonal axis for model improvement in diffusion-based generative modeling.