TREAD: Token Routing for Efficient Architecture-agnostic Diffusion Training

TL;DR

TREAD is a diffusion-training strategy that randomly selects a subset of tokens at an early layer, routes them around all intermediate layers, and reintroduces them at a deeper layer — for training only, with no inference-time or architectural changes. The routed tokens skip computation along the way (cheaper per-step) and the reintroduction of clean-ish noise-space information into deep layers accelerates convergence (fewer steps needed). On ImageNet-1K 256² class-conditional DiT-XL/2, TREAD achieves a 25.39× convergence speedup over the standard DiT baseline at matched FID, and works without modification on state-space backbones (RWKV) too. The recipe is one extra hyperparameter (selection rate) plus a multi-route ensemble loss — no extra components, no auxiliary networks.

Key claims

A single random route that skips ~50% of tokens from layer $L_1$ to layer $L_2$ during training matches vanilla DiT-XL/2’s FID 9.62 in roughly 41 A100 hours versus thousands for the baseline — a 25.39× wall-clock convergence speedup at the best-reported-FID anchor [§5.2, Fig. 1].
Multi-route training (two consecutive routes) nearly halves FID at 400K steps on DiT-B/2 (43.47 → 28.80 → further gains with the ensemble loss); the second route must blend its reintroduced tokens with the preceding layer’s representation via a linear combination, not replace it [§4 Eq. 12, §5.4 Table 7].
Token selection rate has a sweet spot: 0.5 outperforms both 0.25 (FID 39.30 vs 35.78) and 0.75 (72.30, worse than vanilla), indicating that aggressive token withholding (75%) degrades performance below baseline while moderate withholding (50%) gives both cost and quality wins [§5.4 Table 5].
Route location inside the network is largely arbitrary — a grid over $(L_1, L_2)$ pairs shows near-equal FID across most configurations, with the sole exception of routes ending too close to the output layer (insufficient capacity to process reintroduced tokens) [§5.4 Table 6].
The benefit generalizes off the Transformer: applying TREAD to Diffusion-RWKV (a state-space backbone) with row-selection instead of random selection improves over the Diffusion-RWKV baseline, demonstrating architecture-agnostic gains [§5.2 Table 2].
The multi-route ensemble loss (Eq. 15) — applying the masked denoising loss plus an MAE-style auxiliary loss separately per route, then summing — is necessary; ablating it on TREAD-2 collapses FID from 28.80 back to 44.54, worse than vanilla DiT-B/2 [§5.4 Table 7].
Per-iteration speed wins compose with convergence wins: TREAD reaches 7.10 it/s versus MaskDiT 6.81, vanilla DiT 4.30, and MDT 3.40 on a single A100 at batch 64, so the 5.88× iteration-count speedup becomes 9.55× in wall-clock at 400K steps and 25.39× at the DiT-XL/2 best-FID anchor [§5.2, Fig. 4].

Method

TREAD defines a route $r = (L_1 \to L_2)$ as a connection that randomly selects a fraction of tokens at layer $L_1$ and reintroduces them at layer $L_2$ , with intermediate layers operating on the un-routed subset only. Tokens carried by the route skip all computation between $L_1$ and $L_2$ , so every intermediate layer processes a smaller token set. At $L_2$ , the routed tokens are merged back with the in-flight representation. The training loss is decomposed into a standard denoising score-matching term on the visible (un-routed) tokens plus a Masked-AutoEncoder-style auxiliary term that asks the network to reconstruct the held-out tokens from the visible ones — the same loss-decomposition trick MaskDiT uses, generalized to multiple routes.

For multi-route training with routes $r_1, r_2, \dots$ , the key design choice is what to feed layer $L_2^{(r_2)}$ once $r_1$ has already deposited its tokens: the paper finds that a linear combination of the freshly-reintroduced clean-noise tokens and the existing partial-representation tokens (Eq. 12) is essential — replacing one with the other degrades performance. Each route contributes its own masked DSM + MAE loss to the final objective (Eq. 15). Inference is unchanged: at test time the model runs vanilla DiT inference with no routing, no extra parameters, no architectural special cases. The only training-time hyperparameter beyond vanilla DiT is the selection rate per route, fixed at 0.5 for the headline experiments.

Results

ImageNet-1K 256² class-conditional, DiT-XL/2: TREAD-1 reaches FID 6.58 at 256 A100 hours versus MDT 2,724 A100 hours for comparable results; TREAD-2 reaches FID 10.63 at 400K steps versus DiT-XL/2’s 19.47 at the same steps — within striking distance of DiT-XL/2’s best-reported 9.62 at 7M steps [§5.2, Table 3].
9.55× convergence speedup at the 400K-step anchor (TREAD-2 vs DiT-XL/2) and 25.39× at the DiT-XL/2 best-FID anchor; 5.88× of that comes from fewer iterations and the rest from faster per-iteration throughput [§5.2 Table 3, Fig. 1].
Diffusion-RWKV: TREAD with row-selection improves over the Diffusion-RWKV baseline at matched settings (specific FIDs in Table 2), confirming the method works on non-Transformer backbones with one small adaptation (row vs random selection) [§5.2 Table 2].
Iterations-per-second on A100 batch-64: TREAD 7.10, MaskDiT 6.81, DiT 4.30, MDT 3.40 — TREAD is the fastest per-step despite carrying the multi-route ensemble loss [§5.2 Fig. 4].
Ablations: selection rate 0.5 (FID 35.78) > 0.25 (39.30) > 0.75 (72.30, worse than vanilla); route location robust except when $L_2$ is too close to the output; multi-route ensemble loss is necessary (removing it gives FID 44.54, worse than vanilla DiT-B/2 43.47) [§5.4 Tables 5–7].

Why it’s interesting

TREAD is the January 2025 paper (now ICCV 2025) that anchors the “skip computation by routing/masking tokens during training” axis of Diffusion training efficiency. It is the closest cousin to the masking-based MaskDiT and SD-DiT lineage cited in the related work, but with the key twist that withheld tokens come back rather than being thrown away — a structural difference that lets TREAD claim performance gains on top of cost reductions, where most masking-based methods accept a small quality hit. This complements No Other Representation Component Is Needed: Diffusion Transformers Can Provide Representation Guidance by Themselves (SRA) (SRA) and Self-Flow: Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis (Self-Flow), which attack the same convergence-speed problem from the internal-self-distillation angle (align early-noisy and late-clean features inside the same network) rather than from the token-budget angle TREAD takes — and unlike those, TREAD’s costs are concentrated entirely at training time, with no inference-time teacher forward. It also overlaps with the spatial-token-reduction story in Foveated Diffusion: Efficient Spatially Adaptive Image and Video Generation (Foveated Diffusion) but flips the axis from spatial-resolution to temporal-during-training: TREAD reduces tokens per-iteration during training, foveation reduces tokens per-position at inference. The two are likely composable.