通过头级注意力分析,实现视觉自回归模型的高效缓存压缩。
HeatKV: Head-tuned KV-cache Compression for Visual Autoregressive Modeling

- 按注意力强度动态分配每头缓存资源,实现精细化压缩。
- 在2倍压缩率下保持图像质量与人类感知评分不下降。
- 适合需要低延迟高吞吐的生成式视觉模型部署场景。
视觉自回归(VAR)模型近期在保持低延迟的同时展现出出色的图像生成质量。然而,其KV缓存存在严重的内存瓶颈,单张图像生成常需数GB内存。本文提出HeatKV,一种新型压缩方法,根据各注意力头对先前生成尺度的关注程度,自适应调整缓存分配。利用小型离线校准集,按注意力得分对注意力头进行排序,并基于此构建针对特定内存预算的静态剪枝策略。应用于Infinity-2B模型时,HeatKV相较现有方法实现两倍的缓存压缩比,同时保持或提升图像保真度、提示对齐度及人类感知评分。该方法在VAR模型的KV缓存压缩上达到新SOTA,验证了细粒度头级缓存分配的有效性。代码与校准脚本已开源于https://github.com/arm-research/heatkv。
原文摘要 · Abstract (English)
Visual Autoregressive (VAR) models have recently demonstrated impressive image generation quality while maintaining low latency. However, they suffer from severe KV-cache memory constraints, often requiring gigabytes of memory per generated image. We introduce HeatKV, a novel compression method that adapts cache allocation in each head based on its attention to previously generated scales. Using a small offline calibration set, the attention heads are ranked according to their attention scores over prior scales. Based on this ranking, we construct a static pruning schedule tailored to a given memory budget. Applied to the Infinity-2B model, HeatKV achieves $2 \times$ higher compression ratio in memory allocation for KV cache compared to existing methods, while maintaining similar or better image fidelity, prompt alignment and human perception score. Our method achieves a new state-of-the-art (SOTA) for VAR model KV-cache compression, showcasing the effectiveness of fine-grained, head-specific cache allocation. Code and calibration script available at https://github.com/arm-research/heatkv.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。