Diffusion models have emerged as the mainstream approach for visual generation. However, these models usually suffer from sample inefficiency and high training costs. This issue is particularly pronounced in the standard diffusion transformer architecture due to its quadratic complexity relative to input length. Recent works have addressed this by reducing the number of tokens processed in the model, often through masking. In contrast, this work aims to improve the training efficiency of the diffusion backbone by using predefined routes that store this information until it is reintroduced to deeper layers of the model, rather than discarding these tokens entirely. Further, we combine multiple routes and introduce an adapted auxiliary loss that accounts for all applied routes. Our method is not limited to the common transformer-based model - it can also be applied to state-space models. Unlike most current approaches, TREAD achieves this without architectural modifications. Finally, we show that our method reduces the computational cost and simultaneously boosts model performance on the standard benchmark ImageNet-1K 256 x 256 in class-conditional synthesis. Both of these benefits multiply to a convergence speedup of 9.55x at 400K training iterations compared to DiT and 25.39x compared to the best benchmark performance of DiT at 7M training iterations.

本研究解决了扩散模型在样本效率和高训练成本方面的不足，特别是在标准扩散变换器架构中的二次复杂性问题。通过采用预定义的路由机制有效地存储令牌信息并在更深层次重新引入，本文提出了一种无需架构修改的高效训练方法，结果在ImageNet-1K标准基准测试中显著降低了计算成本，提升了模型性能，实现了比现有最佳方法高出9.55倍的收敛加速。

TREAD：高效的架构无关扩散训练的令牌路由