Research map / Image generation / Fast diffusion sampling
Accelerating Diffusion Transformers with Token-wise Feature Caching
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 ToCa, a training-free token-wise feature caching method for diffusion transformers that selects which tokens to recompute using self-attention influence, cross-attention entropy, cache frequency, and spatial uniformity scores, with layer-dependent cache ratios.
Problem
Feature caching for diffusion transformers treats all tokens identically, but tokens differ greatly in temporal redundancy and in how caching errors propagate through self- and cross-attention, causing large quality loss at high acceleration.
Approach
Token-level caching: at each cache step, score tokens by self-attention influence, cross-attention entropy, cache frequency, and spatial uniformity, recompute only low-cache-score-excluded high-priority tokens and reuse cached values for the rest, with cache ratios scaled by layer depth and computation type and the cache updated at every timestep.
Key results
- On OpenSora, ToCa at R=85% reaches 2.36x speedup with VBench 78.34% vs 79.13% for the original (drop 0.79), while FORA at 1.87x drops 2.22.
- On PixArt-α, ToCa (N=3, R=70%) achieves FID-30k 28.33 and CLIP 16.44 at 1.93x speedup vs. original FID 28.09, and beats FORA (N=3) on PartiPrompts CLIP (17.75 vs 17.15).
- On DiT-XL/2 (DDPM, 250 steps), ToCa (N=4, R=93%) reaches FID-50k 2.58 at 2.75x speedup, versus 2.31 for the original and 2.80 for FORA at ~2.97x.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use ToCa token-wise caching instead of FORA | PixArt-α · MS-COCO2017 | FID-30k: 29.67 -> 28.33 | Helped |
| Use ToCa token-wise caching instead of halving sampling steps | PixArt-α · MS-COCO2017 | FID-30k: 37.46 -> 28.33 | Helped |
| Use ToCa token-wise caching | OpenSora · VBench | VBench: 79.13 -> 78.34; 2.36× | Mixed |
| Use ToCa token-wise caching instead of FORA | DiT-XL/2-G · ImageNet | FID: 2.80 -> 2.58 | Helped |
| Add depth-based cache-ratio scheduling r_l (ablation) | ToCa · ImageNet | FID-5k: 9.60 -> 9.32; 0.28 FID reduction | Helped |
| Add layer-type cache-ratio scheduling r_type (ablation) | ToCa · ImageNet | FID-5k: 9.67 -> 9.32; 0.28 FID reduction | Helped |
| Add spatial uniformity score s4 (ablation) | ToCa · ImageNet | FID-5k: 9.35 -> 9.32; 0.02 FID reduction | Helped |
| Add cache frequency score s3 (ablation) | ToCa · ImageNet | FID-5k: 9.65 -> 9.32; 0.33 FID reduction | Helped |
| Use self-attention score s1 instead of random token selection (ablation) | ToCa · MS-COCO2017 | FID-30k: 28.46 -> 28.21 | Helped |
| Use cross-attention score s2 instead of random token selection (ablation) | ToCa · PartiPrompts | CLIP: 17.08 -> 17.75 | Helped |
| Use self-attention score s1 rather than cross-attention score s2 (ablation) | ToCa · PixArt-α | FID-30k: 28.33 -> 28.21 | Helped |
| Use cross-attention score s2 rather than self-attention score s1 (ablation) | ToCa · PartiPrompts | CLIP: 17.13 -> 17.75 | Helped |
Surprising findings
- Caching self-attention tokens, even partially, propagates errors so strongly that it is as harmful as caching all tokens, so ToCa caches all tokens in self-attention layers and shifts computation to MLP layers.
Limitations
- Gains are reported mainly against FORA and with hyperparameters tuned per model; sensitivity to hyperparameters is studied only on DiT.
- Acceleration on FLUX at 1.5x is modest, and the paper reports no wall-clock efficiency gains beyond these ratios.
Experimental setup
Datasets: MS-COCO 2017 PartiPrompts VBench prompts ImageNet
Benchmarks: MS-COCO 2017 FID-30k MS-COCO 2017 CLIP Score PartiPrompts CLIP Score VBench ImageNet FID-50k Image Reward (FLUX)
Metrics: FID CLIP Score VBench score sFID Precision Recall Image Reward latency FLOPs
Baselines: FORA Δ-DiT PAB T-GATE DDIM/DPM-Solver++ reduced steps
Builds on: FORA DeepCache Block Caching
Closest papers by meaning
- Learning-to-Cache: Accelerating Diffusion Transformer via Layer Caching (NeurIPS 2024 · similarity 0.92)
- HarmoniCa: Harmonizing Training and Inference for Better Feature Caching in Diffusion Transformer Acceleration (ICML 2025 · similarity 0.92)
- Accelerating Diffusion Transformer via Increment-Calibrated Caching with Channel-Aware Singular Value Decomposition (CVPR 2025 · similarity 0.89)
- SparseDiT: Token Sparsification for Efficient Diffusion Transformer (NeurIPS 2025 · similarity 0.89)
- BlockDance: Reuse Structurally Similar Spatio-Temporal Features to Accelerate Diffusion Transformers (CVPR 2025 · similarity 0.89)
- Cache Me if You Can: Accelerating Diffusion Models through Block Caching (CVPR 2024 · similarity 0.88)
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.