# TREAD: Token Routing for Efficient Architecture-agnostic Diffusion Training

arXiv: https://arxiv.org/abs/2501.04765 · HTML: https://mlsanity.com/paper/2501.04765/

## Key facts

- Cited by 21 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 69 of its 80 references that are in the map.
- First posted on arXiv 8 Jan 2025.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 69)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Microsoft COCO: Common Objects in Context](https://mlsanity.com/paper/1405.0312/) (2014)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer](https://mlsanity.com/paper/1701.06538/) (2017)
- [DreamBooth: Fine Tuning Text-to-Image Diffusion Models for Subject-Driven Generation](https://mlsanity.com/paper/2208.12242/) (2022)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [DreamFusion: Text-to-3D using 2D Diffusion](https://mlsanity.com/paper/2209.14988/) (2022)
- [Improved Techniques for Training GANs](https://mlsanity.com/paper/1606.03498/) (2016)
- [An Image is Worth One Word: Personalizing Text-to-Image Generation using Textual Inversion](https://mlsanity.com/paper/2208.01618/) (2022)
- [Align your Latents: High-Resolution Video Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2304.08818/) (2023)
- [Improved Techniques for Training Score-Based Generative Models](https://mlsanity.com/paper/2006.09011/) (2020)
- [Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture](https://mlsanity.com/paper/2301.08243/) (2023)
- [Building Normalizing Flows with Stochastic Interpolants](https://mlsanity.com/paper/2209.15571/) (2022)

## Built on by (most cited first, 21 of 21)

- [No Other Representation Component Is Needed: Diffusion Transformers Can Provide Representation Guidance by Themselves](https://mlsanity.com/paper/2505.02831/) (2025)
- [REPA Works Until It Doesn't: Early-Stopped, Holistic Alignment Supercharges Diffusion Training](https://mlsanity.com/paper/2505.16792/) (2025)
- [Guiding a Diffusion Transformer with the Internal Dynamics of Itself](https://mlsanity.com/paper/2512.24176/) (2025)
- [Learning on the Manifold: Unlocking Standard Diffusion Transformers with Representation Encoders](https://mlsanity.com/paper/2602.10099/) (2026)
- [How far can we go with ImageNet for Text-to-Image generation?](https://mlsanity.com/paper/2502.21318/) (2025)
- [No Alignment Needed for Generation: Learning Linearly Separable Representations in Diffusion Models](https://mlsanity.com/paper/2509.21565/) (2025)
- [UniFusion: Vision-Language Model as Unified Encoder in Image Generation](https://mlsanity.com/paper/2510.12789/) (2025)
- [FSVideo: Fast Speed Video Diffusion Model in a Highly-Compressed Latent Space](https://mlsanity.com/paper/2602.02092/) (2026)
- [Sprint: Sparse-Dense Residual Fusion for Efficient Diffusion Transformers](https://mlsanity.com/paper/2510.21986/) (2025)
- [One Model, Many Budgets: Elastic Latent Interfaces for Diffusion Transformers](https://mlsanity.com/paper/2603.12245/) (2026)
- [Speedrunning ImageNet Diffusion](https://mlsanity.com/paper/2512.12386/) (2025)
- [MaskAlign: Token-Subset Representation Alignment for Efficient Diffusion Training](https://mlsanity.com/paper/2606.08788/) (2026)
- [Guiding Token-Sparse Diffusion Models](https://mlsanity.com/paper/2601.01608/) (2026)
- [Self-transcendence: Is External Feature Guidance Indispensable for Accelerating Diffusion Transformer Training?](https://mlsanity.com/paper/2601.07773/) (2026)
- [UDT: Reconciling U-Nets and Diffusion Transformers with Data-Adaptive Token Reduction](https://mlsanity.com/paper/2608.01298/) (2026)
- [Home-made Diffusion Model from Scratch to Hatch](https://mlsanity.com/paper/2509.06068/) (2025)
- [Reg4Pru: Regularisation Through Random Token Routing for Token Pruning](https://mlsanity.com/paper/2602.02163/) (2026)
- [Motif-Video 2B: Technical Report](https://mlsanity.com/paper/2604.16503/) (2026)
- [Gradient Smoothing: Coupling Layer-wise Updates for Improved Optimization](https://mlsanity.com/paper/2606.30813/) (2026)
- [Training-Free Hidden-State Refinement for Flow-Matching Image Generators](https://mlsanity.com/paper/2608.29160/) (2026)
- [Embedding Prediction Helps Image Generation](https://mlsanity.com/paper/2610.02203/) (2026)

## Abstract (arXiv)

Diffusion models have emerged as the mainstream approach for visual generation. However, these models typically suffer from sample inefficiency and high training costs. Consequently, methods for efficient finetuning, inference and personalization were quickly adopted by the community. However, training these models in the first place remains very costly. While several recent approaches - including masking, distillation, and architectural modifications - have been proposed to improve training efficiency, each of these methods comes with a tradeoff: they achieve enhanced performance at the expense of increased computational cost or vice versa. In contrast, this work aims to improve training efficiency as well as generative performance at the same time through routes that act as a transport mechanism for randomly selected tokens from early layers to deeper layers of the model. Our method is not limited to the common transformer-based model - it can also be applied to state-space models and achieves this without architectural modifications or additional parameters. Finally, we show that TREAD reduces computational cost and simultaneously boosts model performance on the standard ImageNet-256 benchmark in class-conditional synthesis. Both of these benefits multiply to a convergence speedup of 14x at 400K training iterations compared to DiT and 37x compared to the best benchmark performance of DiT at 7M training iterations. Furthermore, we achieve a competitive FID of 2.09 in a guided and 3.93 in an unguided setting, which improves upon the DiT, without architectural changes.