💻 computer science

HeatKV: Head-tuned KV-cache Compression for Visual Autoregressive Modeling

HeatKV is a novel KV-cache compression method for Visual Autoregressive models that employs a head-tuned, static pruning schedule based on offline attention ranking to achieve a 2x higher memory compression ratio while maintaining or improving image generation quality.

Original authors: Jonathan Cederlund, Axel Berg, Durmus Alp Emre Acar, Chuteng Zhou, Pontus Giselsson

Published 2026-05-15

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jonathan Cederlund, Axel Berg, Durmus Alp Emre Acar, Chuteng Zhou, Pontus Giselsson

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to paint a masterpiece, but you are working in a tiny room with very limited shelf space. As you paint, you need to keep looking back at your earlier brushstrokes to make sure the new ones fit perfectly. In the world of AI image generation, these "earlier brushstrokes" are called KV-caches (Key-Value caches). They are the AI's short-term memory of what it has already generated.

The problem is that for modern AI models (specifically Visual Autoregressive or VAR models), this memory grows incredibly fast. Generating a single high-quality image can require a shelf so huge (gigabytes of memory) that it forces the AI to run on expensive, specialized hardware, limiting how many people can use it at once.

Enter HeatKV, a new method invented by researchers at Lund University and Arm. Think of HeatKV as a smart, space-saving organizer for that AI's memory shelf.

Here is how it works, using simple analogies:

1. The Problem: The "One-Size-Fits-All" Shelf

Previous methods tried to save space by treating all parts of the AI's brain the same. Imagine a librarian who decides, "Okay, we only have space for 50% of the books, so we'll throw away half the books from every section of the library."

The Issue: This is wasteful. Some sections (like the "History" section) might be rarely used, while others (like "Cooking") are checked constantly. Throwing away 50% of the "Cooking" books hurts the library's ability to function, even if you save space.

2. The Solution: HeatKV's "Personalized" Organization

HeatKV changes the rules. Instead of treating every section of the library the same, it looks at exactly which books are being read most often and keeps those on the shelf, while tossing the ones nobody touches.

The "Head" Concept: The AI model has many "attention heads" (think of them as different specialized librarians). Some librarians care deeply about the early stages of the painting (the rough sketch), while others care about the final details.
The "Scale" Concept: VAR models build images in layers, starting small and getting bigger (like zooming in).
The HeatKV Magic: HeatKV analyzes how much each librarian cares about each specific layer of the painting. It creates a custom "memory budget" for every single librarian.
- Librarian A might need to remember the first 3 layers but can forget the 4th.
- Librarian B might need to remember the 2nd and 5th layers but can forget the rest.

3. How It Decides What to Throw Away

Before the AI starts generating an image for a user, HeatKV does a quick "rehearsal" using a small set of practice images (called a calibration set).

It watches how the librarians pay attention to the different layers.
It ranks them: "This layer is super important for this librarian; that layer is boring."
Based on this ranking, it creates a static schedule (a fixed plan) for what to keep and what to delete to fit within a specific memory limit (e.g., "We only have space for 10% of the total memory").

4. The "Greedy" Cleanup

As the AI generates the image, the memory fills up. HeatKV uses a clever "greedy" strategy to stay under the limit:

It doesn't wait until the shelf is full to start cleaning.
It constantly checks: "If we add this new piece of memory, will we break the budget?"
If yes, it immediately removes the least important piece of memory (the one the librarian cares about the least) to make room. It does this so precisely that it never accidentally runs out of space.

The Results: More Images, Same Quality

The researchers tested this on a massive AI model called Infinity-2B.

The Old Way: To generate an image, the model needed about 42 GB of memory.
The HeatKV Way: It achieved the same high-quality results using only 2.1 GB of memory.
The Win: This is a 20x compression. They managed to squeeze the memory usage down to 10% (or even 4%) of the original size without the AI "forgetting" how to draw good pictures. The images looked just as sharp, and the AI followed instructions just as well as the full-memory version.

Why This Matters

The paper claims that HeatKV allows these powerful image generators to run on hardware with much less memory. This means:

Cheaper Hardware: You might not need the most expensive, giant servers to run these models.
More Users: A single server could handle many more people generating images at the same time because each person takes up less "shelf space."

In short, HeatKV is like a master organizer that knows exactly which memories are vital and which can be let go, allowing the AI to paint beautiful pictures in a much smaller room.

Technical Summary: HeatKV - Head-tuned KV-cache Compression for Visual Autoregressive Modeling

Problem Statement

Visual Autoregressive (VAR) models have emerged as a competitive alternative to diffusion models for image generation, offering fast inference speeds and favorable scaling properties by generating images as a sequence of multi-scale token maps. However, these models face severe memory constraints due to the nature of their KV-cache. Unlike Large Language Models (LLMs) which typically generate one token per step, VAR models generate entire token maps at increasing resolutions. Consequently, the number of tokens in the KV-cache grows super-linearly with the generation scale.

State-of-the-art VAR models, such as Infinity-2B and Infinity-8B, can require up to 10 GB of KV-cache memory to generate a single 1024 × 1024 image. This high memory footprint limits the maximum batch size and the number of concurrent users a service provider can support on specialized hardware accelerators. Existing compression methods for VARs, such as ScaleKV and HACK, often rely on coarse-grained strategies, such as classifying attention heads or layers into binary categories (e.g., "high-budget" vs. "low-budget"). The authors argue that these binary classifications are overly restrictive, as attention patterns within the same layer often vary significantly between different heads, leading to wasted memory or suboptimal pruning.

Methodology: HeatKV

The authors propose HeatKV (Head-tuned KV-cache Compression), a novel strategy that adapts cache allocation at the granularity of individual attention heads and specific image scales. The method consists of three primary components:

1. Head-Scale Importance (Offline Calibration)

Unlike previous methods that use variance-aware metrics, HeatKV ranks attention heads based on their Scale-dependent Cumulative Attention Score (S-CAS).

Mechanism: Using a small offline calibration dataset, the method aggregates attention probabilities from tokens at future scales back to tokens at prior scales.
Metric: For a given layer $\ell$ , head $h$ , and source scale $k$ , the S-CAS measures the average attention mass from all future scales ( $\tau > k$ ) to the tokens at scale $k$ .
Outcome: This generates a specific ranking of "head-scale" pairs for each source scale, indicating how critical a specific head's cache of a specific prior scale is for future token generation. This allows the system to identify that a specific head might need to retain Scale 1 but can discard Scale 2, a level of granularity binary methods cannot achieve.

2. Budget Allocation

HeatKV calculates the number of heads ( $N_k$ ) that must be pruned after generating each scale $k$ to ensure the total cache size remains within a user-defined memory budget $b$ .

Constraint: The method adheres to the "Streaming-LLM" principle of retaining the first $s$ scales as an "attention sink" (essential for stability) and never caching the current scale being generated.
Calculation: The algorithm determines the minimal number of heads to prune to satisfy the budget constraint, allowing for a dynamic reduction of the cache as the generation progresses.

3. Removal Order Optimization (Greedy Early Pruning)

A critical challenge in VARs is that cache size grows super-linearly within a single generation step (across transformer layers). Standard methods that prune only between scales can cause transient memory over-allocations that exceed hardware limits.

Solution: HeatKV introduces a Greedy Early Pruning algorithm (Algorithm 1). It enforces the budget constraint after every transformer layer, not just between scales.
Strategy: If the cache size exceeds the budget after a specific layer, the algorithm greedily selects additional head-scales to prune from that layer or earlier layers. It prioritizes pruning heads with lower S-CAS ranks and, in the case of ties, prioritizes pruning from later layers or later source scales to maximize memory relief.
Optimality: The authors prove (Proposition 1) that this greedy approach yields a minimum-cardinality set of early pruned heads, ensuring the budget is never violated while maximizing the utilization of the available memory.

Key Contributions

Fine-Grained Observation: The authors demonstrate that attention patterns in VAR models vary significantly between different scales within the same attention head but remain relatively stable across different input prompts. This motivates a shift from binary head classification to a continuous, scale-dependent budget allocation.
S-CAS Metric: The introduction of the Scale-dependent Cumulative Attention Score allows for the construction of a static pruning schedule tailored to a specific memory budget, enabling each head to discard non-essential scales while retaining critical ones.
Optimized Pruning Algorithm: The development of a greedy eviction algorithm that enforces budget constraints at the layer level (rather than just the scale level) prevents transient memory overflows, which is crucial for the super-linear growth characteristic of VARs.
State-of-the-Art Performance: The method achieves a new state-of-the-art for VAR model KV-cache compression, demonstrating that fine-grained allocation significantly outperforms coarse-grained approaches.

Experimental Results

The authors evaluated HeatKV on the Infinity-2B and Infinity-8B models using the MS-COCO 2017 validation set. They compared HeatKV against baselines including StreamingLLM, SnapKV, HACK, and ScaleKV across memory budgets of 4%, 10%, and 20%.

Fidelity: HeatKV consistently outperformed all baselines. Notably, on the Infinity-2B model, HeatKV at a 10% budget achieved better FID, LPIPS, and PSNR scores than ScaleKV at a 20% budget. At a 4% budget, HeatKV maintained an FID of 2.70 (Infinity-2B), significantly lower than the next best method (ScaleKV at 4.92).
Prompt Alignment: Using GenEval, HeatKV maintained strong prompt alignment even at low budgets, with scores comparable to the full-cache baseline.
Perceptual Quality: Human Preference Score (HPSv2.1) results indicated that HeatKV preserved high perceptual quality, outperforming other methods at 4% budget on Infinity-2B.
Efficiency: HeatKV achieved a 2× higher compression ratio in memory allocation compared to existing methods while maintaining similar or better image fidelity. The method also enabled larger batch sizes, effectively reducing generation latency compared to full-cache inference.

Significance and Claims

The paper claims that HeatKV represents a significant advancement in making VAR models viable for deployment on hardware with limited memory capacity, such as consumer-grade GPUs, or for serving a larger number of concurrent users on existing infrastructure.

The authors emphasize that their method does not require changes to the model architecture or training process; it relies on a small offline calibration set (as few as 1–10 prompts) to generate a static pruning schedule. The results suggest that fine-grained, head-specific cache allocation is essential for preserving the semantic structure and generation quality of VAR models under aggressive memory constraints, outperforming methods that rely on coarse binary classifications of attention heads or layers.

The paper concludes by noting that while HeatKV optimizes memory, future work could explore combining it with other optimization techniques, such as attention sparsity exploitation or model quantization, to further enhance efficiency.