Computer Vision | AI Research | Arm

HeatKV

Head-tuned KV-cache Compression for Visual Autoregressive Modeling

Jonathan Cederlund1,2 Axel Berg1 William Isaksson1 Durmus Alp Emre Acar1 Chuteng Zhou1 Pontus Giselsson2

1 Arm    2 Department of Automatic Control, Lund University

Paper arXiv Code

Overview

Abstract

Visual Autoregressive (VAR) models can generate high-quality images with low latency, but their KV caches can consume gigabytes of accelerator memory for every generated image. This limits batch size and makes high-resolution image generation difficult on memory-constrained hardware.

HeatKV introduces a fine-grained compression policy that adapts cache allocation for each attention head and image scale. A small offline calibration set is used to calculate the Scale-dependent Cumulative Attention Score (S-CAS), which ranks how strongly each head depends on previously generated scales.

From these rankings, HeatKV builds a static pruning schedule for a chosen memory budget. It removes low-importance head-scale pairs while retaining the cached information that contributes most to future generation, preserving fidelity, prompt alignment, and perceptual quality under aggressive compression.

HeatKV workflow showing calibration images, offline S-CAS computation, and inference-time pruning of low-importance head-scale cache entries.
Figure 1. HeatKV uses calibration images to rank head-scale importance with S-CAS, then applies the resulting static pruning schedule during inference.

How it works

Method

HeatKV turns stable attention patterns into a predictable cache-allocation policy that meets a user-selected memory budget throughout image generation.

1

Collect calibration prompts

A small offline prompt set is used to observe attention patterns as the VAR model generates token maps at increasing scales.

2

Measure scale attention

Attention probabilities are aggregated for each layer, head, source scale, and future generation scale.

3

Rank with S-CAS

S-CAS ranks head-scale pairs by their contribution to future predictions, capturing more detail than binary head or layer classes.

4

Allocate the budget

The target memory budget determines how many low-ranked head-scale pairs must be removed while early sink scales are retained.

5

Prune during inference

A static schedule and greedy early-pruning policy keep cache use within budget as new scales are generated, without prompt-dependent allocation.

Evaluation

Results

2× higher KV-cache compression ratio at comparable image quality

On Infinity-2B, HeatKV at a 10% cache budget outperforms the strongest reported baseline at a 20% budget across FID, LPIPS, and PSNR.

  • Experiments evaluate Infinity-2B and Infinity-8B at 20%, 10%, and 4% of the original cache budget.
  • Testing covers fidelity on MS-COCO 2017, prompt alignment on GenEval, and perceptual quality with HPSv2.1.
  • S-CAS rankings remain stable with a small calibration set, supporting an offline schedule that can be prepared before inference.
Selected Infinity-2B results from Table 1 of the paper. Lower FID and LPIPS are better; higher PSNR is better.
Method Budget KV cache FID ↓ LPIPS ↓ PSNR ↑
Full cache 100% 38,550 MB — — —
ScaleKV 20% 7,800 MB 1.94 0.079 24.93
HeatKV 10% 3,900 MB 1.67 0.073 25.40
HeatKV 4% 1,590 MB 2.70 0.127 22.14

See the paper for full results across all baselines, budgets, models, and evaluation metrics.

Four paired image-generation examples comparing Infinity-8B with a full 42 GB KV cache against HeatKV with a 2.1 GB KV cache.
Figure 2. Infinity-8B examples using a full 42 GB KV cache (top) and HeatKV with a 2.1 GB cache (bottom), a 20× reduction in cache allocation.

Citation

BibTeX

If you find this work useful, please cite the HeatKV paper:

@misc{cederlund2026heatkvheadtunedkvcachecompression,
  title={HeatKV: Head-tuned KV-cache Compression for Visual Autoregressive Modeling},
  author={Jonathan Cederlund and Axel Berg and William Isaksson and
          Durmus Alp Emre Acar and Chuteng Zhou and Pontus Giselsson},
  year={2026},
  eprint={2605.14877},
  archivePrefix={arXiv},
  primaryClass={cs.CV},
  url={https://arxiv.org/abs/2605.14877}
}