Video World Models · 2026-10-21

MoDDiT: Simple Token Routing in Diffusion Transformers

Jarek Liesen, Matthew T. Jackson, Benjamin Ellis, Shimon Whiteson, Jakob N. Foerster

AI-assisted summary based on the paper. Consult the linked source for methods and limitations. Benchmark numbers appear after editor verification.

TL;DR

MoDDiT modifies diffusion transformers to route token computation through both MLP and self-attention operations. The MLP routing uses dynamic, token- and timestep-dependent compute allocation, while self-attention uses fixed-capacity routing. Training and inference use fewer tokens/FLOPs while improving quality–compute tradeoffs on ImageNet, improving 1X video prediction PSNR at 12.5% capacity, and enabling sparse distillation of Wan2.1 with modest VBench score reduction.

Why it matters

The authors address that diffusion transformers process all tokens uniformly and therefore do not exploit token-level sparsity in images/videos. MoDDiT provides dynamic allocation in MLPs (varying by tokens, images, and denoising timesteps) while maintaining efficient GPU computation via fixed-capacity attention routing. The paper reports FLOPs reductions alongside quality improvements, including results for video generation and sparse distillation.

Method

Propose MoDDiT: route tokens through both self-attention and MLP blocks with separate routing for each (dynamic compute in MLPs; fixed-capacity top-K in self-attention).
Use a tiny per-token router (two-layer MLP) trained with stochastic Gumbel top-K during training; at inference track a logit threshold to match target capacity.
For MLP training, merge batch and token dimensions and select top-(B*K) tokens across each GPU to enable compute allocation depending on class and diffusion step; avoid ragged batches via fixed attention capacity.

Abstract (from arXiv)

Real-time video generation, interactive world models, and on-device image synthesis are almost a reality. Diffusion Transformers (DiT) achieve excellent sample quality, shifting the bottleneck to the throughput of the models themselves. Natural images are inherently sparse, as large portions of an image are often background, contribute little to the scene content, and are easy to denoise. DiTs fail to exploit this sparsity, as they process image tokens statically and uniformly. We propose MoDDiT, a simple architecture modification that learns to route tokens through both MLP and self-attention operations. By routing tokens through both blocks independently, MoDDiT allows for dynamic compute allocation in MLPs that varies across tokens, images, and denoising timesteps. Using MoDDiT to train image generation models from scratch significantly reduces computational costs of training and inference, processing as few as 12.5% of the tokens. By increasing sparsity and model depth, MoDDiT reduces costs and improves generation performance over unrouted baselines. The learned allocation is interpretable, revealing how computation varies with visual content and noise level. We also use MoDDiT to train video generation models, where allocation correlates with motion, and to distill a dense bidirectional video model into a sparse autoregressive student.

Publication record

Bibliographic evidence describes where and when a record appeared. It does not validate scientific claims.

First public release: 2026-10-20 · Indexed by Pub-AI: 2026-10-21 UTC

preprint