TREAD: Token Routing for Efficient Architecture-agnostic Diffusion Training

TL;DR

TREAD introduces a token routing mechanism that enhances training efficiency and generative performance of diffusion models across architectures, reducing costs and boosting quality on ImageNet-256.

Contribution

This work proposes a novel token routing method that improves diffusion model training efficiency and performance without architectural modifications or extra parameters.

Findings

14x faster convergence at 400K iterations

37x training speedup over DiT at 7M iterations

Achieves state-of-the-art FID scores on ImageNet-256

Abstract

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.