Rethinking Cross-Layer Information Routing in Diffusion Transformers

arXiv:2605.20708 · cs.CV, cs.AI · Submitted 2026-05-20 · Read on arXiv

Listen

Radio episode about this paper

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom, and with me are Jane, Lu, senior AI researcher at Tsinghua, Meng, lead engineer at a mysterious AI startup and Lalam, the in-house Large Language Model.

Jane: Today's paper: "Rethinking Cross-Layer Information Routing in Diffusion Transformers".

Tom: Diffusion Transformers (DiTs) are central to modern visual generation, yet their fundamental residual stream inherited from standard Transformers remains largely unchanged, leading to underexplored issues in cross-layer information flow.

Jane: First, who's behind it and why it matters.

Title and authors: Tom: Alright team, let's start with what this paper is actually called—"Rethinking Cross-Layer Information Routing in Diffusion Transformers"—and who wrote it. It’s clear from the title that they are focusing on how information gets routed between different layers within these Transformer models.

Jane: And the authors are a solid group, with names like Chao Xu, Maohua Li, and others listed—it shows this is a collaborative effort coming from some major institutions.

Lu: The implication right there is that they aren't just tweaking one small part; they are systematically analyzing the whole flow structure of these models.

Meng: So, what does "rethinking" actually mean in this context? Are they suggesting a complete overhaul of the residual connection, or just a refined way to use it?

Lalam: It suggests that the traditional way we handle information flow is too rigid, and they are proposing a new approach to make that process more flexible.

The paper's summary: Tom: Moving into the summary of "Rethinking Cross-Layer Information Routing in Diffusion Transformers," the main point is that they systematically analyzed two models—a standard SiT-XL/two baseline and a static variant of DAR—and identified three specific problems with standard residual addition.

Jane: Those problems are quite technical: they found "monotonic forward magnitude inflation," which means the information signal just keeps growing without limit as it moves deeper into the layers, plus "sharp backward gradient decay" and "pronounced block-wise redundancy."

Lu: Those symptoms are really telling; the inflation suggests a lack of proper normalization, while the redundancy points to unnecessary repetition in how layers are interacting.

Meng: From an engineering standpoint, monotonic inflation is worrying because it usually means instability or difficulty in training deeper parts of the network effectively.

Lalam: It sounds like they’ve pinpointed specific failure modes that plague standard residual routing, which gives us a clear target for improvement when we look at model architectures.

The paper's improvements: Tom: Now, the core contribution is their proposed solution: Diffusion-Adaptive Routing, or DAR. They describe this as a drop-in replacement that uses learnable, timestep-adaptive aggregation over the history of sublayer outputs instead of just a simple sum.

Jane: The mechanism involves using a softmax attention where the query comes from the AdaLN-modulated hidden state to decide which past representations to emphasize or suppress based on the current noise level.

Lu: What I find really clever is how they make this routing mechanism inherit information about both content and timestep dependence directly from the existing conditioning pathway of the DiT architecture, which keeps it compatible with things like REPA sixty-six.

Meng: That compatibility is key for us; if we can drop in a solution that works orthogonally to other alignment techniques, that simplifies our development pipeline immensely.

Lalam: I think the fact that this routing is timestep-adaptive means the model can dynamically adjust its focus depending on whether it’s denoising noise or trying to capture fine details.

Conclusion: Tom: So, to wrap up, this paper suggests that cross-layer information routing is a really underexplored area in diffusion modeling because standard methods suffer from those inflation and redundancy issues they identified.

Jane: The main takeaway is that Diffusion-Adaptive Routing addresses these rigidity problems by introducing adaptive control over which previous representations are emphasized or suppressed based on the noise level.

Lu: The results show that DAR actually helps both final quality, achieving a FID of seven point five six with SiT-XL/two and it speeds up training significantly, matching the baseline's quality in about eight point seven five times fewer iterations.

Meng: That eight point seven five times fewer iterations is significant for the computational cost of training these large models; that kind of acceleration makes scaling much more feasible.

Lalam: For our culture here, I see this as a way to ensure that the AI systems we build are not just fast, but also dynamically intelligent about what information they choose to focus on during complex tasks like image synthesis.

Tom: That’s a fantastic summary of where they landed with DAR; it seems like a solid architectural adjustment for these visual generation models.

Jane: Indeed, the ability to match baseline quality while reducing training time is a very practical demonstration of the paper's impact on efficiency.

Nanjing University · Alibaba Group

cs.CV, cs.AI

Submitted: 2026-05-20

Updated: 2026-09-30

Code: https://github.com/black-forest-labs/flux

Importance score: 86/100

The gist: Diffusion Transformers (DiTs) are central to modern visual generation, yet their fundamental residual stream inherited from standard Transformers remains largely unchanged, leading to underexplored

Key concepts

Monotonic Forward Magnitude Inflation
This occurs when the hidden state magnitude grows steadily across deeper layers, similar to issues in language models. It suggests that standard residual connections allow too much information to accumulate without proper control, leading to unstable scaling and potentially poor optimization signals for deeper parts of the network.
Diffusion-Adaptive Routing (DAR)
DAR replaces standard addition with a weighted aggregation mechanism. It uses a softmax function based on a query parameter that adapts to the current denoising timestep. This allows the model to selectively emphasize or suppress specific historical layer outputs, tailoring information flow dynamically.
Chunked Aggregation
To handle large models efficiently, DAR uses chunked aggregation. Instead of routing every single previous sublayer output, it groups them into chunks. This reduces computational cost while still performing the adaptive routing necessary for controlling information flow across the network depth.

Terminology

Summary

Diffusion Transformers (DiTs) are central to modern visual generation, yet their fundamental residual stream inherited from standard Transformers remains largely unchanged, leading to underexplored issues in cross-layer information flow. This paper systematically investigates this flow along both depth and denoising timestep, identifying three concrete symptoms of traditional residual addition: monotonic forward magnitude inflation, sharp backward gradient decay, and pronounced block-wise redundancy. Motivated by these findings, the authors propose Diffusion-Adaptive Routing (DAR), a drop-in residual replacement that performs learnable, timestep-adaptive, and non-incremental aggregation over the history of sublayer outputs. This architectural change is shown to improve convergence speed and final quality in diffusion models by operating orthogonally to existing representation-alignment objectives.

Diagnostic Symptoms of Standard Residual Routing

The authors conducted the first systematic study of cross-layer information flow in DiTs, decomposed jointly by depth and denoising timestep. They analyzed two models: a vanilla SiT-XL/2 baseline and a static variant of DAR. The diagnostic analysis revealed that these symptoms persist throughout training and vary systematically with the noise level. Specifically, for the standard residual baseline, they observed:

  1. forward hidden-state magnitude grows monotonically from ∼ 15.5 at block 1 to ∼ 1576 at block 28, corresponding to roughly 100× inflation, which echoes the PreNorm dilution phenomenon seen in LLMs.

  2. backward gradient magnitude drops sharply after the first five blocks, suggesting limited control over gradient flow and weaker optimization signals for deeper layers.

  3. adjacent transformer blocks become increasingly redundant, as indicated by a high block-wise similarity, suggesting substantial representational redundancy under standard residual routing.

Diffusion-Adaptive Routing (DAR) Mechanism

DAR is proposed as a drop-in residual replacement that performs learnable, timestep-adaptive, and non-incremental aggregation over the history of sublayer outputs. It replaces the unweighted sum in the standard residual stream with a softmax-weighted aggregation:

hl = Xl−1i=0 αi→l(t) vi with αi→l(t) = exp(ql(t)⊤ki / √d))

where:

  1. The query parameterization, denoted as Query parameterization, has three variants: pure static, explicit timestep injection (using the existing time-embedder), and a dynamic variant that derives its query from the most recent sublayer output, allowing it to inherit timestep information implicitly.

  2. The design preserves the isotropic and homogeneous Transformer stack by operating purely along the depth dimension, avoiding manually specified layer pairing or U-Net-like skip routing topologies.

Empirical Performance and Benefits

The empirical results demonstrate that DAR significantly improves both final quality and convergence speed on ImageNet 256×256. Key findings include:

  1. Quality Improvement: On SiT, the static variant of DAR achieved a substantially better FID of 6.92 (SDE) without CFG while training for only 600K iterations, and the dynamic variant attained the best ODE FID with CFG (2.05).

  2. Convergence Speed: The static DAR variant matches the baseline's converged quality with 8.75× fewer training iterations.

  3. Orthogonality to Other Objectives: Combining DAR with REPA yields a 2× training acceleration in the early stage over REPA alone, demonstrating that routing-level and representation-level accelerations compound rather than offset each other.

Design Choices for Scalability

To manage computational cost as DiTs scale deeper, the paper introduces chunked aggregation. This method partitions sublayers into chunks of size S, where each chunk is summarized by a single representation:

The aggregated hl then enters the sublayer transformation following vl = fl(hl;t).

  1. The authors identify a U-shaped pattern in the cost function for chunked aggregation, showing that the optimal chunk size is found at S ≈ 4, which minimizes a decomposed cost term related to routing entropy and rate-distortion compression.

  2. They detail an optimized infrastructure using Triton kernels that fuse the normalization constant with the weighted accumulator in one streaming loop over N sources, achieving significant latency speedup and activation-memory savings as the number of routed sources (N) increases.

Conclusion

The paper concludes that cross-layer information routing is a promising and underexplored design axis for diffusion modeling. DAR successfully addresses the inherent rigidity of standard residual routing by introducing adaptive control over which previous representations should be emphasized or suppressed, proving that timestep-adaptive cross-layer routing is a "latent need of the DiT residual pathway.

Improvements for AI systems

As a fastidious and diligent researcher, I have analyzed this paper, Rethinking Cross-Layer Information Routing in Diffusion Transformers, and identified several concrete architectural and training improvements that can be implemented to enhance Diffusion Transformer (DiT) systems.

Here are the specific improvements and the resulting capabilities:

)

  1. Improve cross-layer information flow by replacing standard residual addition with a learnable, timestep-adaptive, non-incremental aggregation mechanism called Diffusion-Adaptive Routing (DAR).

  2. Implement DAR as a drop-in residual replacement that performs softmax attention over preceding sublayer outputs, where the query is computed from the current AdaLN-modulated hidden state.

  3. Enable timestep awareness in routing by using a dynamic query parameterization (or explicit injection) that allows the router to selectively emphasize or suppress specific historical representations based on the current noise level.

  4. Introduce a chunked aggregation strategy for DAR, which partitions sublayers into chunks and aggregates sources via softmax over chunk summaries (instead of all individual sources), optimizing memory and computational costs while maintaining performance.

  5. Integrate DAR synergistically with existing representation-alignment objectives (like REPA) to achieve compounded training acceleration, suggesting that architectural routing improvements are orthogonal to representation alignment methods.

)

The improved AI system can achieve the following specific capabilities:

  1. Enhanced Fidelity in Image/Video Generation: By mitigating the forward magnitude inflation and block-wise redundancy symptoms inherent in standard DiTs, the system will generate images and videos with significantly higher fidelity, specifically preserving high-frequency details like sharp edges and fine textures, even during aggressive few-step distillation or low-iteration training.

  2. Faster Convergence: The system can achieve convergence with significantly fewer training iterations (e.g., 8.75x faster than the baseline on ImageNet 256x256), drastically reducing the computational cost and time required to train large diffusion models.

  3. Superior Feature Selection Based on Noise Level: Because DAR is timestep-adaptive, the system will dynamically select which prior representations are most relevant for denoising at any given noise level (e.g., focusing on coarse structures at high noise and fine details at low noise), leading to more accurate and temporally coherent generation across the entire denoising trajectory.

  4. Scalable Performance in Large Models: The chunked aggregation strategy allows DAR to scale effectively to substantially deeper backbones (multi-billion parameter models) by managing the quadratic growth of source memory, ensuring that performance gains are maintained even as model depth increases, avoiding the saturation points seen in naive residual structures.

  5. Optimized Post-Training Performance: When used in post-training distillation (e.g., on Qwen-Image), DAR allows the model to better preserve fine details during the process, leading to a more robust and high-quality distilled model that outperforms vanilla counterparts under distillation schemes where they typically diverge.

Sources

Related papers