arXiv:2505.20353cs.LGcs.AI2025-05被引 26

用可学习线性近似加速扩散Transformer,减少冗余计算

FastCache: Fast Caching for Diffusion Transformer Through Learnable Linear Approximation

  • 通过空间感知选Token和跨时间步缓存激活值,减少重复计算
  • 在多个DiT模型上实现延迟与内存显著降低,生成质量优于现有方法
  • 适合追求高效生成的图像合成与扩散模型部署场景

扩散Transformer(DiT)是强大的生成模型,但因其迭代结构和深层Transformer堆叠而计算成本高昂。为缓解此效率问题,我们提出FastCache,一种基于隐藏状态的缓存与压缩框架,通过利用模型内部表示中的冗余性来加速DiT推理。FastCache采用双策略:(1) 基于隐藏状态显著性的空间感知令牌选择机制,自适应过滤冗余令牌;(2) 变换器级缓存,在时间步间重用潜在激活值,当变化低于预设阈值时生效。这两个模块协同工作,在保持生成保真度的同时减少不必要的计算,通过可学习线性近似实现。理论分析表明,在基于假设检验的决策规则下,近似误差有界。在多个DiT变体上的实证评估显示,该方法大幅降低延迟与内存占用,且在FID和t-FID指标上优于现有缓存方法。为进一步提升速度,我们还引入基于k-NN密度的令牌合并模块,用于合并冗余令牌。代码已开源。

原文摘要 · Abstract (English)

Diffusion Transformers (DiT) are powerful generative models but remain computationally intensive due to their iterative structure and deep transformer stacks. To alleviate this inefficiency, we propose \textbf{FastCache}, a hidden-state-level caching and compression framework that accelerates DiT inference by exploiting redundancy within the model's internal representations. FastCache introduces a dual strategy: (1) a spatial-aware token selection mechanism that adaptively filters redundant tokens based on hidden-state saliency, and (2) a transformer-level cache that reuses latent activations across timesteps when changes fall below a predefined threshold. These modules work jointly to reduce unnecessary computation while preserving generation fidelity through learnable linear approximation. Theoretical analysis shows that FastCache maintains bounded approximation error under a hypothesis-testing-based decision rule. Empirical evaluations across multiple DiT variants demonstrate substantial reductions in latency and memory usage, achieving the best generation quality among existing cache methods, as measured by FID and t-FID. To further improve the speedup of FastCache, we also introduce a token merging module that merges redundant tokens based on k-NN density. Code is available at \href{https://github.com/NoakLiu/FastCache-xDiT}{https://github.com/NoakLiu/FastCache-xDiT}.

扩散模型缓存优化Transformer加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。