Repenser le routage d'information inter-couches dans les Transformers de diffusion

Résumé

Les Transformers de Diffusion (DiTs) sont devenus une architecture de facto pour la génération visuelle moderne, et presque tous les axes majeurs de leur conception — tokenisation, attention, conditionnement, objectifs et autoencodeurs latents — ont été largement revisités. Le flux résiduel qui régit la manière dont l'information s'accumule à travers les couches a cependant été directement hérité du Transformer original. Dans cet article, nous présentons une analyse empirique systématique du flux d'information inter-couches dans les DiTs, à la fois en profondeur et en fonction du pas de temps de débruitage, et identifions trois symptômes concrets de l'addition résiduelle traditionnelle, à savoir l'inflation monotone de la magnitude vers l'avant, la décroissance brutale du gradient vers l'arrière et une redondance marquée au niveau des blocs. Motivés par ce diagnostic, nous proposons le Routage Adaptatif à la Diffusion (DAR), un remplacement résiduel prêt à l'emploi qui effectue une agrégation apprenable, adaptative en fonction du pas de temps et non incrémentale de l'historique des sorties des sous-couches. De plus, le DAR proposé est compatible avec de nombreuses méthodes modernes d'amélioration des Transformers, telles que REPA. Sur ImageNet 256×256, DAR améliore SiT-XL/2 de 2,11 FID (7,56 contre 9,67) et atteint la qualité convergée de la référence avec 8,75 fois moins d'itérations d'entraînement. Superposé à REPA, il offre une accélération de l'entraînement d'un facteur 2 en phase précoce, suggérant que le routage d'information inter-couches constitue un axe de conception sous-exploré dans la modélisation par diffusion, opérant orthogonalement aux objectifs d'alignement de représentation existants. Au-delà du pré-entraînement, DAR peut également être appliqué lors de la phase de fine-tuning de modèles T2I à grande échelle et préserve les détails haute fréquence durant la Distillation par Appariement de Distribution.

English

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 (DAR), a drop-in residual replacement that performs learnable, timestep-adaptive, and non-incremental aggregation over the history of sublayer outputs. Moreover, the proposed DAR is compatible with many modern Transformer enhancement methods, such as REPA. On ImageNet 256times256, DAR improves SiT-XL/2 by 2.11 FID (7.56 vs.\ 9.67) and matches the baseline's converged quality with 8.75times fewer training iterations. Stacked on top of REPA, it yields a 2times 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, DAR can also be applied during the fine-tuning stage of large-scale T2I models and preserves high-frequency details during Distribution Matching Distillation.