Research map / Architectures and efficiency / Attention and transformers
Layer- and Timestep-Adaptive Differentiable Token Compression Ratios for Efficient Diffusion Transformers
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes DiffCR, which fine-tunes pretrained DiTs (PixArt-Σ, Lazy Diffusion) with per-token routers and differentiable layer- and timestep-wise compression ratios to skip computation for unimportant tokens, improving quality-efficiency trade-offs over ToMe and AT-EDM.
Problem
Diffusion Transformers apply uniform computation across all image tokens, layers and timesteps, causing high latency and memory use that hinders deployment; existing token reduction methods rely on heuristics and uniform compression ratios.
Approach
Adapt mixture-of-depths token routing to DiTs and make the compression ratio itself learnable: each layer and timestep region holds a continuous ratio interpolated between two discrete MoD ratio bins, trained jointly with the model via a target-ratio MSE loss, so layers and timesteps learn their own compression.
Key results
- On T2I with PixArt-Σ at 20% compression, DiffCR-LT reaches FID 10.68 versus 345.91 for ToMe and 251.79 for AT-EDM, with latency and memory comparable to the baselines.
- On inpainting with Lazy Diffusion at 30% compression, DiffCR-LT reaches FID 13.42 versus 279.36 for ToMe and 203.35 for AT-EDM.
- Learned ratios show higher compression for noisier timesteps and redundant middle layers, and the average ratio converges to the target.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add MoD token routing with a uniform compression ratio (ablation) | PixArt-Σ · LAION-5B, 10K images | FID: 11.93 -> 22.78 | Hurt |
| Use layer-wise differentiable compression ratios instead of uniform-ratio MoD (ablation) | PixArt-Σ w/ MoD · LAION-5B, 10K images | FID: 22.78 -> 12.28 | Helped |
| Add timestep-wise adaptive ratios to layer-wise DiffCR (ablation) | PixArt-Σ w/ DiffCR-L · LAION-5B, 10K images | FID: 12.28 -> 10.68 | Helped |
| Add MoD token routing with a uniform compression ratio (ablation) | Lazy Diffusion · LAION-400M, 10K images, 512² mask | FID: 10.90 -> 18.34 | Hurt |
| Use layer-wise differentiable compression ratios instead of uniform-ratio MoD (ablation) | LD w/ MoD · LAION-400M, 10K images, 512² mask | FID: 18.34 -> 13.53 | Helped |
| Add timestep-wise adaptive ratios to layer-wise DiffCR (ablation) | LD w/ DiffCR-L · LAION-400M, 10K images, 512² mask | FID: 13.53 -> 13.42 | Helped |
| Use layer- and timestep-wise DiffCR instead of fine-tuned PixArt-Σ | PixArt-Σ · LAION-5B, 10K images | FID: 11.93 -> 10.68 | Helped |
| Use layer- and timestep-wise DiffCR instead of Lazy Diffusion | Lazy Diffusion · LAION-400M, 10K images, 512² mask | FID: 10.90 -> 13.42 | Hurt |
Surprising findings
- Naive token-level pruning (ToMe, AT-EDM) degrades DiT generation badly, with FIDs above 200 even at modest compression, while MoD-based routing with fine-tuning recovers quality.
Limitations
- Training uses an internal 220M-image inpainting dataset that is not released, limiting reproducibility.
- Differentiable ratio training doubles forward-pass cost during fine-tuning, and inpainting is more sensitive to pruning and needs longer fine-tuning.
- Few-step distillation is not explored and is left as an orthogonal direction.
Experimental setup
Datasets: LAION-5B LAION-400M internal inpainting dataset of 220M images
Benchmarks: LAION-5B LAION-400M
Metrics: FID CLIP score FLOPs latency memory human preference score (HPSv2)
Baselines: ToMe AT-EDM MoD RegenerateImage RegenerateCrop Lazy Diffusion SDXL Model Pruning
Builds on: mixture-of-depths Lazy Diffusion PixArt-Σ ToMe AT-EDM
Closest papers by meaning
- SparseDiT: Token Sparsification for Efficient Diffusion Transformer (NeurIPS 2025 · similarity 0.93)
- FasterDiT: Towards Faster Diffusion Transformers Training without Architecture Modification (NeurIPS 2024 · similarity 0.90)
- SD-DiT: Unleashing the Power of Self-supervised Discrimination in Diffusion Transformer (CVPR 2024 · similarity 0.90)
- Dynamic Diffusion Transformer (ICLR 2025 · similarity 0.90)
- Faster Diffusion: Rethinking the Role of the Encoder for Diffusion Model Inference (NeurIPS 2024 · similarity 0.90)
- UDT: Reconciling U-Nets and Diffusion Transformers with Data-Adaptive Token Reduction (NeurIPS 2026 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.