PaperlayerView exact revision on arXiv

Source-linked AI summary

Accelerating Diffusion Transformers with Token-wise Feature Caching

Chang Zou, Xuyang Liu, Ting Liu, Siteng Huang, Linfeng Zhang

arXiv:2410.05317v4cs.LGcs.AIcs.CV

TL;DR

Diffusion transformers are computationally costly, and tokens differ in their sensitivity to feature caching. ToCa adaptively selects tokens and varies caching ratios across layer types and depths; experiments report improved acceleration over prior caching methods in image and video generation, including 2.36× acceleration on OpenSora.

  • Problem

    Diffusion transformers have high computational costs, and tokens differ in their sensitivity to feature caching, making token selection a relevant gap for acceleration.

    Source: §1 · pp. 1–2

  • Method

    ToCa uses four scores to select tokens for caching and supports different caching ratios across layers of different types and depths.

    Source: §1 · p. 2

  • Results

    ToCa improves acceleration over previous caching methods in image and video generation; OpenSora achieves 2.36× acceleration without training.

    Source: §1 · p. 2, §5 · p. 10

  • Takeaways & Limitations

    Token-wise differences provide an optimization axis for improving caching performance in diffusion transformers across image and video generation.

    Source: §5 · p. 10

Abstract

from arXiv · show

Diffusion transformers have shown significant effectiveness in both image and video synthesis at the expense of huge computation costs. To address this problem, feature caching methods have been introduced to accelerate diffusion transformers by caching the features in previous timesteps and reusing them in the following timesteps. However, previous caching methods ignore that different tokens exhibit different sensitivities to feature caching, and feature caching on some tokens may lead to 10$\times$ more destruction to the overall generation quality compared with other tokens. In this paper, we introduce token-wise feature caching, allowing us to adaptively select the most suitable tokens for caching, and further enable us to apply different caching ratios to neural layers in different types and depths. Extensive experiments on PixArt-$\alpha$, OpenSora, and DiT demonstrate our effectiveness in both image and video generation with no requirements for training. For instance, 2.36$\times$ and 1.93$\times$ acceleration are achieved on OpenSora and PixArt-$\alpha$ with almost no drop in generation quality.

1 INTRODUCTION

Diffusion transformers are costly to run, while caching faces token-level differences in temporal redundancy and error propagation. ToCa selects tokens and layer-specific caching ratios to accelerate generation while maintaining quality.

1 INTRODUCTION

  • Diffusion transformers' high computational costs slow inference and hinder practical real-time use; existing cache methods reuse features across adjacent timesteps.Earlier U-Net cache designs depend on U-Net architecture, while newer methods directly apply caching to transformers.Source: §1 · p. 1
  • Some tokens have adjacent-timestep feature distances nearly 2.5× the mean, while identical noise on different tokens causes final-layer propagation errors differing by more than 10×.These differences motivate prioritizing tokens with greater redundancy and lower caching error.Source: §1 · pp. 1–2
  • ToCa selects cache-suitable tokens using four scores and supports different caching ratios across layer types and depths, without additional computation costs.Its scores address temporal redundancy and error propagation.Source: §1 · p. 2
  • 2.36× acceleration on OpenSora, without training, outperforms halving timesteps by 1.56 on VBench.The experiments cover text-to-image, text-to-video, and class-to-image generation on PixArt-α, OpenSora, and DiT.Source: §1 · p. 2

2 RELATED WORK

Related work traces diffusion models' shift from U-Net to transformer architectures and reviews acceleration methods that reduce sampling steps or denoising-network computation. Feature caching avoids additional training, but prior approaches have architectural or token-uniformity constraints.

2 RELATED WORK

  • Diffusion models iteratively denoise an initial noise input, and Diffusion Transformers replace U-Net backbones with Transformers to scale across generative tasks.PixArt-α applies DiT to text-to-image generation, while Sora demonstrates its potential for video generation.Source: §2 · p. 3
  • Acceleration methods reduce sampling timesteps or speed denoising networks through approaches including solvers, distillation, token reduction, quantization, and pruning.Cache-based methods also reduce redundant computation without additional training.Source: §2 · p. 3
  • U-Net-specific cache methods do not apply to DiT, while DiT caching approaches use uniform cache solutions across tokens or layers and can degrade generation quality.Learning-to-Cache instead learns a cache strategy but requires training.Source: §2 · p. 3

3 METHODOLOGY

ToCa accelerates diffusion transformers by selectively reusing cached token features, rather than treating all tokens and layers alike. Its selection scores and layer-specific caching ratios account for token sensitivity, error propagation, and layer depth and type.

3.1 PRELIMINARY

  • Diffusion-transformer denoising repeatedly processes patch tokens through attention, MLP, and conditional cross-attention layers.The denoising network runs across timesteps, making its repeated computation a target for caching.Source: §3.1 · p. 4

3.2 NAIVE FEATURE CACHING FOR DIFFUSION TRANSFORMERS

  • Naive caching computes and stores every layer's features at the first timestep, then reuses them for the next N − 1 timesteps, but reuse error can grow with N.After each N-timestep period, the cache is initialized again.Source: §3.2 · p. 4

3.3 TOKEN-WISE FEATURE CACHING

  • ToCa targets differences in token-level caching impact with a finer-grained strategy for selecting which tokens to cache.The framework also adapts caching ratios across layers.Source: §3.3 · pp. 4–5

3.4 OVERALL FRAMEWORK

  • After computing all tokens to initialize the cache, ToCa scores tokens at subsequent timesteps, reuses cached features for selected tokens, computes the rest, and updates cache entries for computed tokens.The cache ratio specifies the fraction of tokens in each layer whose computation is skipped.Source: Figure 3 · p. 4, §3.4 · p. 5

3.5 TOKEN SELECTION

  • ToCa's score combines four criteria: influence on other tokens, cross-attention entropy as a measure of control influence, recent cache frequency, and spatial uniformity.Higher influence and entropy make a token less suitable for caching; frequently cached tokens are discouraged, while the spatial rule discourages locally concentrated caching.Source: §3.5 · pp. 5–6

3.6 DIFFERENT CACHE RATIOS IN DIFFERENT LAYERS

  • ToCa adjusts layer caching ratios by depth and type, reflecting differences in temporal redundancy and in how caching errors affect the final output.It assigns higher ratios to deeper than shallow layers; self-attention layers cache all tokens, while MLP and cross-attention ratios track their shares of total computation cost.Source: §3.6 · pp. 6–7

4 EXPERIMENT

The experiments evaluate ToCa across PixArt-α, OpenSora, and DiT using task-specific configurations and metrics. Reported results include lower FID and sFID than FORA on class-conditional generation, with additional text-to-image results, ablations, and token visualizations.

4.1 EXPERIMENT SETTINGS

  • The evaluation uses COCO-2017 captions for text-to-image, VBench prompts for text-to-video, and ImageNet classes for class-conditional generation.The reported evaluation sizes are 30,000 images, 4,750 videos, and 50,000 images, respectively; configurations vary by model and layer type.Source: §4.1 · pp. 7–8

4.2 RESULTS ON TEXT-TO-IMAGE GENERATION

  • 9.37 FID increments are reported when directly halving sampling steps, while ToCa has only a very marginal quality decrease at similar acceleration ratios.The comparison is for PixArt-α text-to-image generation; Table 1 covers MS-COCO2017 and PartiPrompts with 20 DPM++ steps by default.Source: Table 1 · p. 8, §4.2 · p. 9

4.3 RESULTS ON TEXT-TO-VIDEO GENERATION

  • Figure 8 presents VBench metrics and acceleration ratios for ToCa and other text-to-video methods.The supplied figure caption does not state the comparative outcomes.Source: Figure 8 · p. 9

4.4 RESULTS ON CLASS-CONDITIONAL IMAGE GENERATION

  • 0.39 lower sFID and 0.22 lower FID than FORA at a similar acceleration ratio are reported for class-conditional generation.Table 3 covers ImageNet class-to-image generation with DiT-XL/2.Source: §4.4 · p. 9, Table 3 · p. 10
  • Different layer-type and depth caching ratios reduce FID by 0.28 and 0.35, while scores s3 and s4 reduce FID by 0.02 and 0.33, respectively.The attention-score ablation finds s1 tends toward lower FID and s2 toward higher CLIP Score.Source: §4.4 · p. 10, Table 4 · p. 10, Table 5 · p. 10
  • With cache-frequency score s3, background tokens are cached less excessively, and Figure 9(c) shows over 4.5k fewer cached tokens.The visualization compares token caching with and without s3; darker pixels indicate more frequent caching.Source: §4.4 · p. 10, Figure 9 · p. 10

5 CONCLUSION

ToCa reuses features from previous timesteps by selectively skipping computation for tokens, based on token-specific temporal redundancy and error propagation. The paper reports better acceleration than prior caching methods across image and video generation.

5 CONCLUSION

  • ToCa adaptively skips computation for selected tokens by reusing their features from previous timesteps, using token-level differences in temporal redundancy and error propagation.The conclusion describes this as a token-wise feature caching method.Source: §5 · p. 10

A.1 ENGINEERING DETAILS

The section introduces engineering techniques.

A.1 ENGINEERING DETAILS

  • The section introduces engineering techniques for the method.Source: §A.1 · p. 13

A.1.1 STEP-WISE CACHING SCHEDULING

The scheduling details adjust caching ratio and forced activation cycle over generation timesteps. The stated schedules shift computation from earlier to later steps and use smaller activation cycles during later detail generation.

A.1.1 STEP-WISE CACHING SCHEDULING

  • The caching ratio is determined by multiplying R by layer-depth, layer-type, and timestep factors.The section describes adjusting R across layer depths and types as well as across timesteps.Source: §A.1.1 · p. 13
  • The timestep multiplier rt adjusts caching ratio R, shifting computation from earlier steps toward later steps to improve generated-image quality.The multiplier is defined as rt = 0.5+λt(0.5−t/T).Source: §A.1.1 · p. 13
  • The schedule uses a larger forced activation cycle N early and a smaller N during later detail generation to enhance detail quality.Nt is defined using expected theoretical speedup N0 and slope hyperparameter wt.Source: §A.1.1 · p. 13

A.1.2 PARTIALLY COMPUTATION ON SELF-ATTENTION

The method avoids partial computation in self-attention during non-forced activation, reflecting concern about error propagation. A sensitivity-study scheme can instead trade computation between self-attention and MLP layers.

A.1.2 PARTIALLY COMPUTATION ON SELF-ATTENTION

  • During non-forced activation, self-attention uses rtype = 0, avoiding partial computation in that module.The authors motivate this choice by the rapid propagation and accumulation of errors in self-attention.Source: §A.1.2 · p. 13
  • The sensitivity-study allocation sets self-attention rtype = 1 − 0.4λtype and MLP rtype = 1 + 0.6λtype.The factors 0.4 and 0.6 reflect the approximate computational ratio between the modules in DiT.Source: §A.1.2 · p. 14

A.1.3 TOKEN SELECTION FOR CFG AND NON-CFG

Token-selection consistency between classifier-free guidance (cfg) and non-cfg parts depends on conditioning type. Text-conditional models are sensitive to independent selection, whereas class-conditional DiT is not.

A.1.3 TOKEN SELECTION FOR CFG AND NON-CFG

  • For text-to-image and text-to-video models, independent cfg and non-cfg token selection substantially degrades performance, so the schemes should match.This applies to models such as PixArt-α and OpenSora.Source: §A.1.3 · p. 14
  • In class-conditional DiT, independent versus identical cfg and non-cfg selection produces only minor performance differences.The authors attribute the contrast to text cross-attention injecting conditioning unevenly into cfg and non-cfg parts, unlike class conditioning.Source: §A.1.3 · p. 14

A.2 MORE IMPLEMENTATION DETAILS ON EXPERIMENTAL SETTINGS

The experiments evaluate ToCa across class-to-image, text-to-image, and text-to-video generation, with model-specific settings and quality measures. They report benchmark data and computational setup for reproducibility.

A.2 MORE IMPLEMENTATION DETAILS ON EXPERIMENTAL SETTINGS

  • DiT uses ImageNet class prompts for 50,000 256 × 256 images; PixArt-α uses 30,000 COCO captions, and OpenSora evaluates 4,750 videos from VBench prompts.Evaluation includes FID-50k for DiT, FID-30k and CLIP Score for PixArt-α, and 16 VBench aspects for OpenSora.Source: §A.2 · p. 14
  • PixArt-α uses an average caching ratio R = 70%, while OpenSora fixes MLP caching at Rmlp = 85%.PixArt-α settings include N = 2 and rl = 0.3; OpenSora uses forced-activation cycles of 3 for temporal attention, spatial attention, and MLP, and 6 for cross-attention.Source: §A.2 · p. 14
  • DiT uses an average caching ratio R = 93%, with forced activation cycle N = 3 and N = 2 during sampling steps 50–100.Experiments ran on six A800 GPUs with 80 GB memory each.Source: §A.2 · pp. 14–15
  • Figure 10 provides visualization examples comparing different acceleration methods on DiT.The caption identifies the comparison but does not state an outcome.Source: Figure 10 · p. 15

A.3 SENSITIVITY STUDY

Across the supplementary studies, ToCa reports nearly lossless FLUX acceleration and favorable DDIM comparisons, alongside parameter sensitivity, computation-cost analyses, and token- and layer-level observations.

A.3 SENSITIVITY STUDY

  • λl = 0.06, λtype = 2.5, λt = 0.03, and λ3 = 0.25 were optimal; grid size 2 outperformed 4, and fixing β = 2 at steps 50–100 improved quality.The study also reports stable performance when transferring hyperparameters across DiT models of different sizes.Source: §A.3 · p. 15

A.4.1 MAIN COMPUTATIONS

  • The self-attention computation includes O(N^2D) terms, while the MLP cost includes terms proportional to N × D1 × D2 and N × D2.The appendix gives separate cost expressions for self-attention, cross-attention, and MLP layers.Source: §A.4.1 · p. 16

A.4.2 COMPUTATION COSTS FROM TOKEN SELECTION

  • Token-selection costs are small in a practical case: Cache Frequency scoring accounts for 0.044%, Uniform Spatial Distribution scoring for 0.15%, and sorting for 0.18% of main computation.The text-selection timing is typically less than 1% of main computation; at R = 90%, O(ND2) terms fall to 10% and O(N^2D) terms to 1%, with the latter ignored in practice when computation shifts to MLP.Source: §A.4.2 · p. 17

A.5 IMPLEMENTED RESULTS ON CLASS-CONDITIONAL IMAGE GENERATION

  • ToCa achieved 1.32 lower sFID and 0.28 lower FID than FORA at approximately 2.7× acceleration, and 0.21 lower FID than reduced-step sampling at approximately 2.5×.The authors select N = 3, R = 93% as the recommended DDIM scheme; Table 6 is identified as the class-to-image comparison.Source: §A.5 · p. 17, Table 6 · p. 17

A.6 RESULTS ON HIGHER-RESOLUTION AND MORE ADVANCED TEXT-TO-IMAGE MODELS

  • At 1.5× acceleration, ToCa's Image Reward scores were almost identical to unaccelerated FLUX.1-dev and FLUX.1-schnell, while all tested FORA configurations degraded quality.The comparison uses 50-step FLUX.1-dev and 4-step FLUX.1-schnell; Figure 13 shows FLUX.1-schnell examples.Source: §A.6 · p. 19, Table 7 · p. 19, Figure 13 · p. 20

A.7 DETAILS FOR DISTRIBUTION FIGURES

  • Deeper layers show poorer temporal redundancy, but errors introduced there have a smaller impact on the output at the same timestep.The appendix also reports varying temporal redundancy across layers and varying error accumulation and propagation across tokens.Source: §A.7 · pp. 19, 21
Loading 2410.05317v4…