PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning-to-Cache: Accelerating Diffusion Transformer via Layer Caching

Xinyin Ma, Gongfan Fang, Michael Bi Mi, Xinchao Wang

arXiv:2406.01733v2cs.LGcs.CV

TL;DR

Diffusion transformers deliver strong generative results but are slow because every denoising step evaluates a large model. L2C learns differentiable, timestep-variant layer-routing policies that reuse prior computations without updating model parameters. It substantially reduces computation, preserves image quality in reported settings, and outperforms comparable samplers and cache-based methods, while remaining model-dependent and capped at 2× acceleration.

  • Problem

    Diffusion transformers have high deployment cost and slow inference, motivating methods that reduce sampling steps or per-step computation.

  • Method

    L2C treats transformer layers as caching units and learns differentiable timestep-variant routers that interpolate between cached and full computations without changing diffusion-model parameters.

  • Results

    L2C cacheable ratios reach 93.68% for U-ViT-H/2 and 47.43% for DiT-XL/2 with ΔFID < 0.01, while outperforming DDIM, DPM-Solver, and prior cache-based methods at comparable speed.

  • Takeaways & Limitations

    Layer caching can remove substantial diffusion-transformer computation while retaining high image quality and improving on comparable acceleration baselines.

  • Takeaways & Limitations

    Lossless caching does not uniformly exist across models, and the two-step schedule caps acceleration at 2×.

Abstract

from arXiv · show

Diffusion Transformers have recently demonstrated unprecedented generative capabilities for various tasks. The encouraging results, however, come with the cost of slow inference, since each denoising step requires inference on a transformer model with a large scale of parameters. In this study, we make an interesting and somehow surprising observation: the computation of a large proportion of layers in the diffusion transformer, through introducing a caching mechanism, can be readily removed even without updating the model parameters. In the case of U-ViT-H/2, for example, we may remove up to 93.68% of the computation in the cache steps (46.84% for all steps), with less than 0.01 drop in FID. To achieve this, we introduce a novel scheme, named Learning-to-Cache (L2C), that learns to conduct caching in a dynamic manner for diffusion transformers. Specifically, by leveraging the identical structure of layers in transformers and the sequential nature of diffusion, we explore redundant computations between timesteps by treating each layer as the fundamental unit for caching. To address the challenge of the exponential search space in deep models for identifying layers to cache and remove, we propose a novel differentiable optimization objective. An input-invariant yet timestep-variant router is then optimized, which can finally produce a static computation graph. Experimental results show that L2C largely outperforms samplers such as DDIM and DPM-Solver, alongside prior cache-based methods at the same inference speed. Code is available at https://github.com/horseee/learning-to-cache

1 Introduction

Diffusion transformers offer strong generative performance but incur high inference costs. L2C learns which layers to cache between timesteps, achieving substantial computation reduction with minimal quality loss and outperforming comparable baselines.

  • Diffusion transformers face high deployment costs and slow inference because each timestep evaluates a large transformer model.Sampling cost grows with both the number of timesteps and model size per timestep.
  • 93.68% of U-ViT-H/2 layers and 47.43% of DiT-XL/2 layers are cacheable with ΔFID < 0.01.The reported cacheable ratios apply to cache steps on ImageNet.
  • L2C converts layer selection into a differentiable optimization problem to learn caching patterns for diffusion transformers.The approach addresses the exponentially growing search space of layer-cache configurations.
  • L2C significantly outperforms fewer-step samplers and prior cache-based methods at comparable inference speed.The introduction reports this comparison as a central empirical result.

2 Related Work

Prior diffusion acceleration methods reduce sampling steps, reduce per-step computation, or reuse computations across denoising steps. Existing cache methods mainly exploit consecutive-step feature similarity and U-Net computation structure.

  • Diffusion transformers provide an alternative to U-Net backbones and support applications including video, speech, and 3D generation.DiT emphasizes transformer scalability, while U-ViT uses long skip connections.
  • Diffusion acceleration methods primarily reduce the number of sampling steps or the inference cost per step.Examples include trajectory distillation, solver approximations, model compression, and low-precision formats.
  • Cache-based diffusion methods reuse feature maps or computations across consecutive steps, especially by exploiting U-Net structure.These approaches are motivated by high similarity between high-level features in consecutive denoising steps.

3 Method

L2C interpolates between full computation and cached layer outputs, then learns timestep- and layer-specific routing coefficients that minimize approximation error under a computation budget. During inference, these coefficients are thresholded to activate or disable layers while leaving diffusion-model parameters unchanged.

  • 3.2 Approximating ϵθ(·)with a lightweight substitute: L2C learns a lightweight substitute by interpolating between the previous-step prediction and the full current-step prediction.The interpolation is designed to approximate the full model output while reducing inference cost.
  • 3.3 Caching the Layer: A Feasible Choice for the Interpolation I: A zero layer coefficient reuses the previous timestep’s output and skips that layer’s non-residual computation, while a nonzero coefficient triggers normal computation.The coefficients therefore act as routers selecting activated or disabled layers.
  • 3.3 Caching the Layer: A Feasible Choice for the Interpolation I: Approximation error varies across timesteps, layers, and block types, motivating timestep-variant coefficients β across the model depth.The reported differences include higher later-step error in DiT and distinct behavior between attention and feedforward layers.
  • 3.4 Learning to Cache: The objective minimizes the discrepancy between the interpolated and full predictions while constraining total computation cost.The optimization seeks as many zero coefficients as possible with minimal approximation error.
  • 3.4 Learning to Cache: Caching every two steps combines one full model inference with one cheaper step, limiting the maximum acceleration ratio to 2×.The same two-step caching schedule is used in the training and inference algorithms.
  • 3.4 Learning to Cache: β remains continuous during optimization and is thresholded into binary routing decisions at inference, while diffusion-model parameters remain unchanged.The training procedure optimizes only β, using sigmoid-constrained coefficients and a threshold for discretization.

4 Experiments

Experiments evaluate L2C on DiT and U-ViT using ImageNet generation, showing strong quality–latency tradeoffs, substantial cacheability, and advantages over samplers and prior cache-based methods.

  • Main Results: FID 3.46 with 20-step DDIM on DiT-XL/2 nearly matches unaccelerated quality and improves over DDIM’s 4.68 at comparable evaluation.The method uses more denoising steps with lower average time per step, while fewer-step samplers use fewer but more expensive steps.
  • Main Results: L2C consistently improves quality across acceleration ratios and significantly outperforms DDIM and DPM-Solver in the FID–latency tradeoff.Heuristic layer-selection strategies degrade image quality when layer dependencies matter, whereas L2C accounts for those dependencies.
  • Main Results: Up to 94% of U-ViT layers can be discarded during cache steps with less than 0.01 FID degradation, while DiT supports a considerably lower cacheable proportion.Full inference occurs at step s, while cached layers are used at step m; cacheable ratios also differ between FFN and MHSA components.
  • Main Results: L2C achieves better quality than other cache-based methods on U-ViT under the comparison setting.Prior cache-based methods are strongly coupled to U-Net structure and therefore are not directly applicable to DiT.
  • Analysis: U-ViT’s learned router makes the middle layers nearly cacheable while retaining computation at both model ends, whereas DiT shows no equally clear pattern.Across models, later stages tend to retain more computation than earlier stages.
  • Analysis: Higher acceleration ratios favor larger thresholds, while lower ratios favor smaller thresholds, indicating that layer importance does not follow a strict sequential ranking.The threshold study varies λ across six points on a curve.

5 Limitation

The method depends on the structure of trained diffusion models and has a capped acceleration ratio, limiting how uniformly and extensively lossless caching can be applied.

  • Lossless caching does not uniformly exist across all models because the method depends strongly on trained diffusion-model structure.The authors report a slight FID drop for DiT-XL/2 at 512 resolution despite outperforming the baseline.
  • 2× is the maximum acceleration under the current two-step schedule, because each pair contains one full-model inference and one cheaper step.

6 Conclusion

The paper proposes L2C for faster diffusion-transformer inference by learning a differentiable layer-caching strategy. Experiments report faster inference with high-fidelity images and improvements over DDIM, DPM-Solver, and cache-based baselines.

  • L2C interpolates between an inexpensive but suboptimal model and an expensive optimal model, then discretizes a trained continuous router for inference.
  • L2C largely outperforms DDIM, DPM-Solver, and other cache-based methods in experiments.
  • The learned solution provides faster inference while producing high-fidelity images for some diffusion-transformer models.

A Proof

The proof establishes two equivalent ways to obtain x_t: directly updating from x_s or computing an intermediate x_m before reaching x_t.

  • Two approaches yield equivalent results for obtaining x_t from x_s.
  • The direct approach updates x_t from x_s using the solution at time t.
  • The intermediate approach first computes x_m from x_s, then computes x_t from x_m while preserving ϵ_θ(x_m, m) = ϵ_θ(x_s, s).
  • The proof separately considers the solution for x_m from x_s before deriving x_t under the matched-noise condition.

A.2 Layer interpolation and Interpolation I

The layer interpolation is designed to connect the models at timesteps s and m through a differentiable function. Setting all layer interpolation parameters to either endpoint recovers the corresponding model output.

  • The interpolation is defined to connect ϵ_θ(x_s, s) and ϵ_θ(x_m, m) through layer-wise interpolation.
  • The interpolation must satisfy endpoint interpolation, continuity, and differentiability; the proof focuses on establishing the endpoint condition.
  • Setting α_i and β_i to 0 across layers recovers ϵ_θ(x_s, s), one endpoint of the interpolation.
  • Setting α_i and β_i to 1 across layers recovers ϵ_θ(x_m, m), the other interpolation endpoint.

B Additional Experiments

The appendix compares DPM-Solver with and without shifted cache steps and explains why shifting is necessary for DPM-Solver-2.

  • Table 6 compares DPM-Solver performance with and without shifted cache steps while caching all layers.
  • Shifting cache steps from [2,4,6,8,10,...] to [3,5,7,9,11,...] is necessary when using DPM-Solver-2.DPM-Solver-2 computes the first-order derivative by subtracting the output at timestep i from that at timestep i + 1.
  • The shift accounts for DPM-Solver-2's dependence on model outputs at adjacent timesteps.

B.2 Layer Dropout v.s. Layer Cache

Layer caching preserves the original pipeline more effectively than layer dropout, while random dropping severely damages image quality.

  • B.2 Layer Dropout v.s. Layer Cache: Layer caching significantly outperforms layer dropout and maintains pixel-wise consistency with the original pipeline.
  • B.2 Layer Dropout v.s. Layer Cache: Around 10% layer dropout causes substantial performance degradation.
  • B.2 Layer Dropout v.s. Layer Cache: Algorithmically selected dropped layers can preserve correct semantics, whereas randomly dropped layers severely compromise acceptable image generation.
  • B.2 Layer Dropout v.s. Layer Cache: At approximately 1.18×–1.19× speedup, L2C is compared with DDIM, learned layer dropping, and random layer dropping.

B.3 Effect of the hyper-parameter λ and θ

The router is relatively insensitive to most optimization hyperparameters, while λ and θ govern the latency–quality trade-off and reproduce reported settings.

  • B.3 Effect of the hyper-parameter λ and θ: Only the training λ and inference threshold θ materially affect the learned router's behavior; learning rate, training epoch, and optimizer hyperparameters do not.
  • B.3 Effect of the hyper-parameter λ and θ: Table 8 lists λ and θ settings used to reproduce the results in Table 1.
  • B.3 Effect of the hyper-parameter λ and θ: Table 9 reports performance across different λ values with θ fixed at 0.1.
  • B.3 Effect of the hyper-parameter λ and θ: λ controls router sparsity and trades inference latency against generated-image fidelity.
  • B.3 Effect of the hyper-parameter λ and θ: The appropriate λ is setting-dependent rather than universal, partly because DiT and U-ViT differ in implementation.

NeurIPS Paper Checklist

The checklist records the paper's claims, experimental reproducibility information, statistical-reporting disclosure, compute-resource information, ethics statement, and limitations discussion.

  • NeurIPS Paper Checklist: The paper states that its claims are clearly presented and that Section 5 discusses limitations.
  • NeurIPS Paper Checklist: The checklist marks experimental reproducibility as satisfied based on information in Section 4.1 and related materials.
  • NeurIPS Paper Checklist: The checklist references Section 4.1 and implementations for experimental-setting and compute-resource information.
  • NeurIPS Paper Checklist: The paper states that it does not report error bars or other statistical-significance information.
  • NeurIPS Paper Checklist: The authors state that their research conforms to the NeurIPS Code of Ethics and discuss relevant matters in Appendix C.
Loading 2406.01733v2…