通过缓存机制减少扩散Transformer的冗余计算,提升生成速度。
Token Caching for Diffusion Transformer Acceleration
- 基于分块架构,动态选择需缓存的令牌和块。
- 在多步推理中实现速度与质量的平衡,加速显著。
- 适合追求高效扩散模型推理的应用场景。
扩散Transformer因其卓越性能在生成建模中受到广泛关注,但其计算开销大,尤其注意力机制的二次复杂度和多步推理过程成为实际应用的主要瓶颈。为解决此问题,本文提出TokenCache,一种利用Transformer分块架构的新型加速方法,旨在减少冗余计算。该方法回答三个关键问题:(1) 哪些令牌应被剪枝并重用以消除冗余?(2) 哪些块应被选中进行高效缓存?(3) 在哪些时间步应用缓存策略可平衡速度与质量?针对这些问题,TokenCache引入缓存预测器,通过层次化策略实现:(1) 令牌剪枝:为每个令牌分配重要性分数以决定是否剪枝和重用;(2) 块选择:自适应分配各块的剪枝比例以确定缓存块;(3) 时间调度:决定在哪些时间步应用缓存策略。在多个模型上的实验表明,TokenCache在保持生成质量的同时显著提升推理速度。
原文摘要 · Abstract (English)
Diffusion transformers have gained substantial interest in diffusion generative modeling due to their outstanding performance. However, their computational demands, particularly the quadratic complexity of attention mechanisms and multi-step inference processes, present substantial bottlenecks that limit their practical applications. To address these challenges, we propose TokenCache, a novel acceleration method that leverages the token-based multi-block architecture of transformers to reduce redundant computations. TokenCache tackles three critical questions: (1) Which tokens should be pruned and reused by the caching mechanism to eliminate redundancy? (2) Which blocks should be targeted for efficient caching? (3) At which time steps should caching be applied to balance speed and quality? In response to these challenges, TokenCache introduces a Cache Predictor that hierarchically addresses these issues by (1) Token pruning: assigning importance scores to each token to determine which tokens to prune and reuse; (2) Block selection: allocating pruning ratio to each block to adaptively select blocks for caching; (3) Temporal Scheduling: deciding at which time steps to apply caching strategies. Experimental results across various models demonstrate that TokenCache achieves an effective trade-off between generation quality and inference speed for diffusion transformers.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。