Переосмысление межслойной маршрутизации информации в диффузионных трансформерах

Аннотация

Диффузионные трансформеры (Diffusion Transformers, DiTs) стали де-факто основой современных систем генерации изображений; практически все ключевые аспекты их архитектуры — токенизация, механизм внимания, conditioning, целевые функции и латентные автоэнкодеры — были многократно пересмотрены. Однако остаточный поток (residual stream), управляющий накоплением информации между слоями, был напрямую унаследован от оригинального трансформера. В данной работе мы проводим систематический эмпирический анализ межслойного информационного потока в DiTs, рассматривая его совместно по глубине сети и временному шагу шумоподавления (denoising timestep), и выявляем три конкретных симптома традиционного остаточного сложения: монотонное увеличение амплитуды сигнала при прямом проходе, резкое затухание градиента при обратном распространении и выраженную избыточность на уровне блоков. Основываясь на этом диагнозе, мы предлагаем метод Diffusion-Adaptive Routing (DAR) — взаимозаменяемую замену остаточных связей, выполняющую обучаемую, адаптивную по временному шагу и неинкрементальную агрегацию истории выходов подслоёв. Предложенный DAR совместим со многими современными методами улучшения трансформеров, такими как REPA. На наборе данных ImageNet 256×256 DAR улучшает показатель FID для модели SiT-XL/2 на 2,11 (7,56 против 9,67) и достигает качества, сопоставимого с базовым уровнем сходимости, при сокращении количества итераций обучения в 8,75 раза. В сочетании с REPA DAR обеспечивает двукратное ускорение обучения на начальном этапе, что указывает на то, что маршрутизация межслойной информации является малоизученным аспектом проектирования диффузионных моделей, действующим ортогонально существующим целям выравнивания представлений. Помимо предварительного обучения, DAR также может быть применён на этапе дообучения крупномасштабных моделей T2I, сохраняя высокочастотные детали при дистилляции с согласованием распределений (Distribution Matching Distillation).

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.