Rethinking Cross-Layer Information Routing in Diffusion Transformers

Conference: NeurIPS2026
arXiv: 2605.20708
Area: Image Generation
Keywords: diffusion Transformer, cross-layer information routing, timestep conditioning, attention residuals, chunked aggregation

TL;DR

DAR replaces incremental residual addition in diffusion Transformers with timestep-aware attention over historical sublayer outputs, achieving an unguided ODE FID of 7.56 after 600K steps on ImageNet 256×256, a 2.11 improvement over the SiT baseline trained for 1.75M steps, while remaining compatible with REPA's representation-alignment loss.

Background & Motivation

Tokenization, attention, conditioning, training objectives, and latent encoders in diffusion Transformers (DiTs) have received extensive attention, yet cross-layer communication usually retains the standard Transformer design: each attention or MLP sublayer adds its output to a single residual stream. Expanding this recurrence assigns unit weights to the input embedding and every historical sublayer output. Once shallow information enters the stream, deeper layers can keep adding to it but cannot explicitly reconsider which sources should be retained or suppressed.

The paper diagnoses this behavior in SiT-XL/2 trained for 600K steps, using 4096 ImageNet samples. At timestep \(t=1.0\), block-output RMS grows from approximately 15.5 to approximately 1576; backward gradients decay sharply after the first five blocks; and token-level cosine similarity between neighboring deep blocks stays above 0.9. The authors connect these symptoms to PreNorm dilution: the residual stream grows while new branches still receive normalized inputs, making deeper updates less influential on the accumulated state. These are diagnostic associations, however, not proof that every highly similar representation constitutes useless computation.

Diffusion introduces a control dimension absent from ordinary language models: the denoising timestep. Different noise levels require different mixtures of shallow and deep information, which fixed weights cannot directly accommodate. Hand-designed U-Net skips restore shallow features but predefine layer pairings. The paper attaches measurement-only scalar gates initialized to 1 to the baseline's historical sources, keeps the forward pass unchanged, and examines loss gradients with respect to these gates, finding timestep-dependent source preferences. Core Idea: replace an equally weighted, cumulatively compressed residual stream with timestep-aware depth attention that reselects historical sources at each sublayer, using chunked source retention to control overhead.

Method

Overall Architecture

Diffusion-Adaptive Routing (DAR) preserves SiT's attention, MLP, and conditioning computations, replacing residual aggregation between these sublayers. Given the current noisy latent, timestep, and class condition, the network first obtains an input embedding. Each sublayer then aggregates available historical sources with softmax weights, executes its original transformation, and adds the new output to the sources accessible downstream. This depth attention selects among sublayer outputs; it does not replace spatial self-attention between tokens.

Three design groups determine routing: how historical sources are weighted, how the query receives timestep information, and which sources remain accessible after chunking. Ordinary DAR also uses a dedicated final aggregator before the prediction head, giving the last chunk's raw sublayer outputs direct access to the prediction. Combining DAR with REPA introduces a separate parameter-sharing adjustment for final aggregation. Inference still uses normal denoising sampling and does not require a DINOv2 teacher; that teacher supplies supervision only during REPA-enabled training.

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["Noisy latent<br/>timestep and condition"] --> B["Depth-attention aggregation"]
    B --> C["Timestep-aware query"]
    C --> D["Chunked sources and final readout"]
    D --> E["Original sublayers and prediction head<br/>velocity output"]
    E -->|Inference: repeated sampler calls| F["Generated image"]
    G["Training: velocity target<br/>optional DINOv2 teacher"] -.->|MSE; optional REPA alignment| E
    D -.->|Historical sources for later sublayers| B

The first three nodes describe design relationships inside the router, not three additional sequential networks. The query scores sources for depth attention, while chunking determines the candidates; dashed supervision applies only during training.

Key Designs

1. Depth-attention aggregation: reselect historical information instead of continuing equal-weight accumulation

Treat attention and MLP as individual sublayers, denote the input embedding by \(v_0\), and the raw output of sublayer \(i\) by \(v_i\). Standard residuals sum all preceding sources into the input of sublayer \(l\). DAR instead normalizes weights along depth and mixes the sources. RMSNorm produces each source's key; query–key dot products determine relevance, while the weighted values remain the source outputs themselves:

\[ h_l=\sum_{i=0}^{l-1}\alpha_{i\to l}(t)v_i,\qquad \alpha_{i\to l}(t)=\frac{\exp\!\left(q_l(t)^\top k_i/\sqrt{d}\right)}{\sum_{j=0}^{l-1}\exp\!\left(q_l(t)^\top k_j/\sqrt{d}\right)},\qquad k_i=\mathrm{RMSNorm}(v_i). \]

Softmax constrains weights to sum to 1, so adding more sources no longer automatically increases a sum of unit coefficients. Routing can suppress irrelevant sources and emphasize a small useful subset. Non-incremental aggregation means rebuilding the input from accessible sources rather than merely adding to the preceding accumulated state. It does not remove sequential layer execution or historical dependencies.

The modification preserves a homogeneous Transformer stack without manually pairing shallow and deep layers. DAR weights concentrate on a few sources and change with timestep, echoing the baseline's counterfactual gate-gradient diagnostic. Gate gradients and softmax weights are different quantities, however, so their values or visualization colors should not be treated as directly comparable probabilities.

2. Timestep-aware query: distinguish static query parameters from static routing weights

DAR compares three queries. Pure static uses one learnable vector per sublayer; dynamic linearly projects the most recent sublayer output; static with explicit timestep injection reuses the existing timestep embedding without adding a query-projection matrix. The variants are:

\[ q_l(t)= \begin{cases} w_l,&\text{pure static},\\ W_q^{(l)}v_{l-1},&\text{dynamic},\\ w_l+e(t),&\text{static with explicit timestep injection}. \end{cases} \]

The dynamic query inherits timestep and content information from sublayer outputs influenced by the noisy input and adaLN-Zero conditioning. Explicit injection instead adds the timestep embedding directly to the query. Zero initialization of the timestep embedder's final layer recovers the pure-static variant at the start of training. Static here describes the base query parameterization, not fixed routing: even a pure-static query scores keys and values that vary with the noisy input and network conditioning.

Table 2 shows that explicit timestep injection approaches dynamic routing and even achieves lower FID at 400K steps. The principal gain therefore cannot simply be attributed to content-dependent queries or additional parameters. A linear probe on the dynamic model also reports timestep-decoding \(R^2\) above 0.95 within the first five blocks and close to 1.0 deeper in the stack. Nevertheless, the probe paragraph refers to aggregated states \(h_l\), whereas the query equation uses the most recent raw output \(v_{l-1}\). This difference in the measured object is retained rather than assuming the two tensors are identical.

3. Chunked sources and final readout: retain current details while compressing completed chunks

Retaining every sublayer output makes the source cache grow linearly with depth. DAR divides \(L\) sublayers into chunks of size \(S\). Each completed chunk retains only its last sublayer output as its summary, \(c_n=v_{nS}\), with the input embedding represented by \(c_0\). The summary is neither a sum of chunk outputs nor a separate pooling operation over the chunk. Its dependence on earlier information comes from the routed input already consumed by that last sublayer.

Sublayer \(l\) in chunk \(n\) can access the input embedding, the last-output summaries of earlier chunks, and all preceding raw outputs within its current chunk:

\[ \mathcal{S}_l=\{c_0,c_1,\ldots,c_{n-1}\}\cup\{v_{(n-1)S+1},\ldots,v_{l-1}\},\qquad c_n=v_{nS}. \]

The number of candidate sources is at most \(S+N\), where \(N=L/S\). Accounting for the feature dimension per token, the paper reduces source-cache storage from \(O(Ld)\) to \(O((S+N)d)\); actual memory must also include batch and token dimensions. \(S=1\) makes all historical sublayers accessible. Increasing \(S\) compresses more old outputs into fewer summaries but can discard historical detail.

The final prediction head uses a slightly different source set: the input embedding, the first \(N-1\) chunk summaries, and all raw outputs of the final chunk, including \(v_L\). Compared with AttnRes reading only chunk summaries, this preserves a direct outlet for task-specific details from the latest layers. The appendix reports an approximately 2-point FID gain after 200K steps but provides no complete numerical table for this comparison. With REPA, the final aggregation does not use independent aggregator parameters; it reuses the last chunk's MLP-aggregator query and per-source RMSNorm parameters. This is an architectural implementation difference, not a new definition of the REPA loss.

A Worked Example

SiT-XL/2 has 28 Transformer blocks, each containing attention and MLP sublayers, giving \(L=56\). With \(S=4\), there are 14 chunks. Thus, c4 means four sublayers, not four Transformer blocks.

The first sublayer of the second chunk accesses the input embedding \(c_0\) and the first chunk's last output \(c_1=v_4\). At the seventh sublayer, \(v_5\) and \(v_6\) from the current chunk are also accessible. The old chunk's \(v_1\), \(v_2\), and \(v_3\) are no longer individually cached.

Before the prediction head, the source set contains \(c_0\) through \(c_{13}\) and the final chunk's \(v_{53}\) through \(v_{56}\), totaling 18 sources. Changing the timestep changes scores through explicit injection or the dynamic query, allowing reselection within the available source set. No specific routing weight is invented for this example.

Loss & Training

DAR changes cross-layer aggregation without prescribing an additional generative loss. ImageNet experiments retain SiT's velocity-prediction MSE and data processing, with global batch size 1024, learning rate \(1\times10^{-4}\), and bfloat16. The authors rerun the SiT and REPA baselines in the same experimental environment.

With REPA, DINOv2-B serves as the pretrained visual encoder, with representation-alignment coefficient 0.5 applied at the eighth layer. DAR governs how historical states are combined; REPA governs how intermediate representations are shaped by supervision. Compatibility does not imply an absence of implementation interactions: the final parameter reuse described above is specific to their combined configuration.

The implementation fuses RMSNorm, query–key dot products, softmax, and the weighted sum in Triton, recomputing some intermediates during back-propagation. This reduces kernel launches and HBM traffic but does not eliminate the memory cost of storing sources. Iteration reduction, per-step throughput, and peak memory must therefore be evaluated separately.

Key Experimental Results

Main Results

ImageNet-1K at 256×256, 50,000 generated samples, and 250 function evaluations by default; classifier-free guidance (CFG) uses scale 1.5. These XL/2 system results come from Table 1. Appendix Table 5 further specifies explicit timestep injection for static c4.

Method	Training steps	Parameters	Unguided ODE FID	Unguided SDE FID	Guided ODE FID	Guided SDE FID
SiT	1.75M	675M	9.67	8.61	2.15	2.06
SiT-Plus	1M	752M	10.85	10.02	2.36	2.34
DAR Static c4	600K	675M	7.56	6.92	2.08	2.23
DAR Dynamic c4	500K	751M	8.07	7.39	2.05	2.17

These are system comparisons with different training budgets, not architectural ablations at identical iteration counts. DAR improves both unguided ODE and SDE results. With guided SDE sampling, SiT's 2.06 is better than static DAR's 2.23 and dynamic DAR's 2.17, so superiority across every sampling configuration is not supported.

Ablation Study

The following timestep-query comparison follows Table 2 and reports FID at different training steps. It most directly separates additional parameters from timestep visibility.

Query configuration	100K	200K	400K
Static, no explicit timestep injection	22.36	15.47	11.51
Dynamic	13.95	9.29	8.10
Static, explicit timestep injection	17.39	10.12	7.97

The corresponding paragraph in Section 5.3 cites this timestep ablation as “Table 3,” but the data appear in Table 2. Table 3 contains the REPA combination below. This cross-reference inconsistency is explicitly retained rather than mixing the two tables.

REPA combination	100K FID	200K FID	300K FID
SiT + REPA	9.89	6.89	6.29
DAR + REPA	7.09	5.92	5.68

DAR+REPA's 7.09 at 100K approaches, but does not outperform, REPA's 6.89 at 200K. The paper's approximately 2× early-stage acceleration describes a similar-quality threshold, not exactly equal FID values.

Key Findings

• At 400K, explicit timestep injection reduces pure-static FID from 11.51 to 7.97, while dynamic obtains 8.10. This supports timestep awareness but does not establish dynamic as the best variant at every stage.
• At 300K without CFG, the chunk-size ablation reports FID 10.41/8.39/11.14 for \(S=1/4/8\). An intermediate chunk size performs best. Table 4's \(S=4\) value differs from the 300K static-c4 value of 8.62 in Appendix Table 5, and Table 4 does not clearly specify its query configuration; they must not be treated as the same checkpoint.
• The authors report an 8.75× iteration reduction to the baseline FID target of 9.67. With fused implementations, SiT and static DAR c4 run at 1.83/1.73 steps/s, implying an estimated 8.27× wall-clock speedup rather than 8.75×; peak memory is 54.56/69.97 GB. The exact threshold crossing is not directly tabulated, so static DAR's 200K FID of 10.12 must not be described as already reaching 9.67.
• Directly transferred AttnRes achieves its best reported FID of 8.71 at 700K; static DAR c4 achieves 7.56 at 600K, an improvement of 1.15. This supports diffusion-specific adaptation but does not isolate the contribution of each individual modification.
• Four-step Qwen-Image DMD on GenEval2 prompts reduces teacher-relative LPIPS from 0.538 to 0.512 and RAPSD high-frequency-ratio deviation from +0.340 to +0.215. The latter compares high-frequency content above frequency threshold 0.2, but the source does not provide its complete computation formula. It is not GenEval2 semantic accuracy.

Highlights & Insights

• The residual stream can be interpreted as an implicit historical retriever, not merely an optimization-stability device. Replacing fixed addition with learned selection lets deeper layers reorganize information without changing the spatial-attention backbone.
• Static parameters do not imply static behavior. Explicit timestep embeddings make a low-overhead query responsive to noise level, while keys also depend on the current input; these two aspects of routing dynamics should remain distinct.
• Chunking decides which information may be compressed, rather than simply retaining fewer tensors. Last-output summaries for old chunks and raw outputs for the final chunk jointly address cost control and the final outlet for detail.
• A transferable research direction is to inject control signals from other conditional generative models into depth queries and test complementarity with existing representation losses. This is a hypothesis motivated by the results, not a cross-task conclusion already established here.

Limitations & Future Work

• The authors identify systematic multi-billion-parameter T2I/T2V pretraining as future work. Qwen-Image post-training provides preliminary evidence only and cannot establish effectiveness for video or all preference-optimization procedures.
• The chunk-size theory assumes a surrogate cost combining routing entropy and compression distortion. Its square-root depth-scaling prediction is not an unconditional theorem about actual FID, and experiments across different depths are absent.
• Operator-level acceleration of fused kernels over naive DAR is not end-to-end acceleration over SiT. Actual training retains approximately 5.8% per-step overhead, and the reported static configuration uses more memory.
• DMD uses LoRA rank 64, student/fake-branch learning rates \(5\times10^{-6}\)/\(2\times10^{-6}\), four denoising steps, guidance 4.0, resolution 1024², and per-GPU batch size 1. Teacher-relative metrics do not replace independent evaluation of perceptual quality, semantic alignment, or user preference.
• The most informative follow-up would standardize quality thresholds and wall-clock budgets, report routing across depth and noise level, and separately ablate timestep injection, chunk summaries, final aggregation, and REPA parameter sharing.

Related Work & Insights

• vs Attention Residuals: Both use depth-wise softmax attention. DAR differs in diffusion-timestep queries, last-sublayer summaries, chunk-size choices, and final raw-output readout. Its contribution is diffusion-specific adaptation and empirical diagnosis, not the first introduction of depth attention.
• vs U-ViT / U-DiT: Long skips restore multilevel features with manual shallow–deep pairing, whereas DAR preserves a homogeneous stack through learned source selection. Different parameters, training budgets, and sampling protocols prevent attributing every cross-model difference to routing topology.
• vs REPA: REPA supervises intermediate representations with a visual teacher; DAR changes their cross-layer flow. Combined gains support complementarity, but the special final-aggregation implementation means compatibility still requires architecture-specific validation.
• vs layer-caching acceleration: Caching mainly reuses features across neighboring denoising steps or layers to reduce inference computation. DAR reconstructs depth-wise source mixtures within each network call, targeting convergence and generation quality rather than directly reducing sampling steps.

Rating

• Novelty: 4/5 — Depth attention builds on AttnRes, but timestep diagnosis and diffusion-specific adaptation provide a clear extension.
• Experimental Thoroughness: 4/5 — Query, chunking, REPA, direct AttnRes, and DMD evidence are included; large-scale pretraining and complete factor isolation are missing.
• Writing Quality: 3/5 — The mechanism is broadly clear, but table cross-references and probe-tensor descriptions are ambiguous.
• Value: 4/5 — Provides an architectural axis complementary to training objectives, with benefits requiring assessment alongside throughput, memory, and sampling conditions.