针对图像生成的内存瓶颈,按注意力头特性智能分配缓存空间。
Head-Aware Key-Value Compression for Efficient Autoregressive Image Generation

- 根据注意力头的局部或全局关注特性,动态分配缓存预算。
- 在不降低图像质量前提下,显著减少内存占用并提升生成速度。
- 无需额外训练,可通用适配不同输入与模型架构。
自回归视觉生成虽表现优异,但因需缓存已生成的视觉标记而导致内存高、吞吐低。近期研究发现,仅保留少量缓存标记即可维持高质量图像,同时大幅降低内存消耗与提升吞吐。然而,现有方法对各注意力头均分配固定预算,忽略了头间异质性,导致资源分配不佳。本文观察到不同层的注意力头具有各异的关注模式:部分专注局部邻域,另一些则捕捉更广上下文依赖。基于此,提出名为HeadKV的头感知键值(KV)缓存压缩框架,将较小预算分配给局部偏好头,较大预算分配给具有广泛注意力的头。关键挑战在于识别每个头的类型以指导压缩。进一步发现,同一层内,各头在不同标记位置的行为具有一致性——早期与后期行为一致。据此,可在生成初期确定头类型,并持续用于整个生成过程的缓存压缩。该方法无需额外训练或数据集统计,且能无缝跨输入泛化。此外,设计分层标记淘汰策略,有效保留长程信息。大量实验验证了其在多个自回归图像生成模型上的有效性。
原文摘要 · Abstract (English)
Autoregressive (AR) visual generation has achieved remarkable performance but suffers from high memory usage and low throughput, as it requires caching previously generated visual tokens. Recent research has shown that retaining only a few lines of cache tokens can maintain high-quality images while significantly reducing memory usage and improving throughput. However, these methods allocate a fixed budget to each attention head, overlooking the heterogeneity among attention heads, leading to suboptimal memory allocation. In this paper, we observe that attention heads across different layers exhibit diverse attention patterns, where some heads focus on local neighborhoods while others capture broader contextual dependencies. Based on this insight, we propose a novel head-aware key-value (KV) cache compression framework for autoregressive image generation, called HeadKV, which assigns smaller budgets to locality-biased heads and larger budgets to heads with broader attention. A key challenge lies in identifying the type of each attention head to guide cache compression. We further observe that, within the same layer, each head exhibits consistent attention patterns across token positions, \emph{i.e.}, a head's behavior for early tokens remains consistent with that for later tokens. This insight suggests that head types can be identified during the early stage and reused for KV compression throughout generation. Its advantage is that it requires no additional training or dataset-level statistics and generalizes seamlessly across different inputs. Moreover, we design a Stratified Token Eviction strategy to effectively preserve long-range information. Extensive experiments demonstrate its effectiveness across multiple autoregressive image generation models.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。