arXiv:2511.16047cs.CV2025-11AAAI被引 3

提出自适应分尺度KV缓存,显著降低图像生成模型显存占用

AMS-KV: Adaptive KV Caching in Multi-Scale Visual Autoregressive Transformers

  • 按尺度重要性动态分配缓存,优先保留关键局部与粗粒度特征
  • 显存使用减少84.83%,注意力计算延迟降低60.48%
  • 适合大规模图像生成场景,尤其提升高并发下的稳定性

基于下一尺度预测的视觉自回归建模(VAR)已成为可扩展的图像生成范式。尽管大语言模型中的键值(KV)缓存已得到广泛研究,但针对基于下一尺度的VAR变压器,其缓存设计仍鲜有探索。主要瓶颈在于随着尺度数量增加,KV内存呈指数增长,严重制约可扩展性。系统性分析发现:(1) 关注局部尺度的令牌显著提升生成质量;(2) 为最粗粒度尺度(即浓缩尺度)分配少量内存可稳定多尺度图像生成;(3) 在缓存高效层中,精细尺度间存在强KV相似性,而高需求层则表现出较弱的跨尺度相似性。基于上述观察,我们提出AMS-KV——一种面向下一尺度预测的自适应分尺度KV缓存策略。该方法优先存储来自浓缩尺度和局部尺度的KVs,以保留最关键信息,维持生成质量;同时通过跨尺度相似性分析识别高缓存需求层,优化缓存利用率与计算效率。相比基线模型,AMS-KV将KV缓存使用量降低最高达84.83%,自注意力延迟降低60.48%。当基线模型在批量大小128时发生显存溢出,AMS-KV可稳定支持批量大小256,且吞吐量提升。

原文摘要 · Abstract (English)

Visual autoregressive modeling (VAR) via next-scale prediction has emerged as a scalable image generation paradigm. While Key and Value (KV) caching in large language models (LLMs) has been extensively studied, next-scale prediction presents unique challenges, and KV caching design for next-scale based VAR transformers remains largely unexplored. A major bottleneck is the excessive KV memory growth with the increasing number of scales-severely limiting scalability. Our systematic investigation reveals that: (1) Attending to tokens from local scales significantly contributes to generation quality (2) Allocating a small amount of memory for the coarsest scales, termed as condensed scales, stabilizes multi-scale image generation (3) Strong KV similarity across finer scales is predominantly observed in cache-efficient layers, whereas cache-demanding layers exhibit weaker inter-scale similarity. Based on the observations, we introduce AMS-KV, a scale-adaptive KV caching policy for next-scale prediction in VAR models. AMS-KV prioritizes storing KVs from condensed and local scales, preserving the most relevant tokens to maintain generation quality. It further optimizes KV cache utilization and computational efficiency identifying cache-demanding layers through inter-scale similarity analysis. Compared to the vanilla next-scale prediction-based VAR models, AMS-KV reduces KV cache usage by up to 84.83% and self-attention latency by 60.48%. Moreover, when the baseline VAR-d30 model encounters out-of-memory failures at a batch size of 128, AMS-KV enables stable scaling to a batch size of 256 with improved throughput.

图像生成自回归缓存优化

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