💻 computer science

Head-Aware Key-Value Compression for Efficient Autoregressive Image Generation

This paper proposes HeadKV, a novel framework for autoregressive image generation that improves memory efficiency and throughput by dynamically assigning variable key-value cache budgets to attention heads based on their distinct locality or global attention patterns, which are identified early and reused throughout the generation process without requiring additional training.

Original authors: Guotao Liang, Baoquan Zhang, Zhiyuan Wen, Yunming Ye

Published 2026-05-21

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Guotao Liang, Baoquan Zhang, Zhiyuan Wen, Yunming Ye

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 massive, intricate mural, but you have a very small desk to work on. Every time you add a new brushstroke, you have to keep a reference of every single previous stroke you made so you don't lose the picture's context. Eventually, your desk gets so cluttered with these references that you can't move your hands anymore, and painting slows to a crawl.

This is exactly what happens in Autoregressive Image Generation. These AI models create images one tiny piece (token) at a time. To make sure the new piece fits with the old ones, the computer has to remember (cache) every single piece it has already generated. For high-resolution images, this "memory desk" becomes huge, eating up all the computer's resources and making generation painfully slow.

The Old Way: One Size Fits All

Previously, researchers tried to clear off the desk by throwing away some old references. But they did it the same way for every part of the painting process. They assumed every "brain cell" (attention head) in the AI needed the same amount of history.

The paper's authors realized this was like giving a map of the whole city to someone who only needs to know the next street corner, while giving a tiny street sign to someone trying to navigate the entire country. It's inefficient.

The New Discovery: Two Types of "Brain Cells"

By looking closely at how these AI models think, the authors found that the AI's "brain cells" actually fall into two distinct personality types:

The "Local" Neighbors: Some brain cells only care about what's happening right next to them. They are like a person looking at a single brick in a wall; they only need to see the bricks immediately touching it to do their job.
The "Global" Architects: Other brain cells care about the big picture. They are like an architect looking at the whole building; they need to remember details from the other side of the room to ensure the roof doesn't collapse.

The Solution: HeadKV (The Smart Desk Manager)

The authors created a new system called HeadKV. Instead of treating all brain cells the same, HeadKV acts like a smart desk manager that assigns different amounts of space based on who is working:

For the "Local" Neighbors: HeadKV says, "You only need to see the last few items." It keeps a tiny, sliding window of recent history and throws the rest away immediately. This saves a massive amount of space.
For the "Global" Architects: HeadKV says, "You need to remember the big picture." It keeps a larger history but uses a special trick called Stratified Token Eviction.

The "Stratified Token Eviction" Trick:
Imagine you are cleaning a bookshelf. If you just pick the "most important" books, you might accidentally throw away all the books from the 1990s because the 2020s books are louder and more popular. But you still need the 1990s books for context!

HeadKV splits the history into two piles: "Nearby" and "Far Away." It makes sure to keep a few important books from both piles. This ensures the AI doesn't forget the distant, crucial details (like the shape of a whole bird or the color of the sky) just because it's focused on the immediate details (like the texture of a feather).

How They Do It Without Training

Usually, teaching a computer to know which brain cell is which requires years of extra training. HeadKV is clever: it figures this out on the fly.

Think of it like meeting a new person. You don't need to know their whole life story to know if they are a "local" or "global" thinker. You just watch them for the first few minutes. If they only talk about what's happening right now, they are a "local" thinker. If they start referencing things from far away, they are "global."

HeadKV watches the AI for a tiny moment at the start of generating an image, decides which brain cells are which, and then applies the right memory rules for the rest of the process. It requires no extra training and works on any image the AI tries to make.

The Result

By using this "smart desk" approach, the authors showed that they could:

Shrink the memory usage: They could keep only 1/4, 1/6, or even 1/8 of the original memory needed.
Speed things up: The AI generates images much faster because it isn't drowning in unnecessary data.
Keep the quality: The images still look just as good as the ones made with the full, cluttered memory.

In short, HeadKV stops the AI from hoarding useless information, letting it paint beautiful pictures faster and with less effort.

Technical Summary: Head-Aware Key-Value Compression for Efficient Autoregressive Image Generation

Problem Statement

Autoregressive (AR) visual generation models, which represent images as sequences of visual tokens, have achieved high-resolution and photorealistic results by scaling model size and token counts. However, this scaling introduces substantial memory overhead and low inference throughput due to the Key-Value (KV) cache required during autoregressive decoding. While recent works in natural language processing (NLP) have explored KV cache management (e.g., token selection, eviction, and quantization), directly applying these techniques to the visual domain is challenging due to the inherent spatial structure of images.

Existing visual-specific methods, such as LineAR, attempt to reduce cache usage by retaining only a subset of tokens based on the observation that tokens primarily attend to spatially adjacent neighbors. However, these approaches allocate a fixed, uniform budget to every attention head. This strategy overlooks the heterogeneity among attention heads, leading to suboptimal memory allocation. Specifically, some heads focus on local neighborhoods (requiring small budgets), while others capture long-range dependencies like object structures and background regions (requiring larger budgets). Uniform allocation fails to adapt to these distinct patterns, resulting in either wasted memory or degraded generation quality.

Methodology: HeadKV

The authors propose HeadKV, a novel, training-free framework that performs head-aware KV cache compression. The core insight is that attention heads in AR image generation models exhibit two distinct patterns—Local and Global—and that a head's behavior is consistent across token positions within the same layer for a given input.

The HeadKV pipeline consists of three key steps:

Cache Prefilling: The model encodes conditional inputs and generates an initial sequence of visual tokens without compression. This stage accumulates sufficient context to estimate head-wise attention distributions.
Head Grouping (Classification): Once the cache reaches a predefined length, the framework classifies each attention head as either "Local" or "Global" using a threshold-based attention accumulation criterion.
- For each head, attention scores are accumulated backward from the current token position until a threshold ( $\tau = 0.9$ ) is reached.
- If the number of tokens required to reach the threshold is smaller than a predefined local window size, the head is classified as Local. Otherwise, it is classified as Global.
- This classification is performed once during the early generation stage and reused for the entire decoding process, requiring no additional training or dataset-level statistics.
Head-Aware Cache Compression: Different compression strategies are applied based on the head type:
- Local Heads: These heads primarily attend to nearby context. The framework assigns a compact budget, retaining only the most recent tokens via a sliding-window mechanism. Historical tokens are omitted entirely.
- Global Heads: These heads capture long-range dependencies. To preserve informative distant tokens without incurring linear memory growth, the authors design a Stratified Token Eviction (STE) strategy.
  - Stratified Token Eviction: Candidate historical tokens are partitioned into two disjoint subsets: Near-Range and Long-Range tokens based on their distance from the current position.
  - Importance scores are computed independently for each subset using past query features.
  - Tokens with the lowest scores are evicted from each subset (e.g., $M/2$ from near-range and $M/2$ from long-range). This prevents locally dominant tokens from overshadowing semantically relevant distant tokens, ensuring a balanced preservation of both local details and global structure.

Key Contributions

Analysis of Attention Heterogeneity: The paper provides an in-depth analysis of AR image generation models, identifying two distinct attention patterns (Local and Global) and revealing that head behaviors are invariant across token positions within a layer for a fixed input.
HeadKV Framework: A novel, training-free compression framework that allocates asymmetric budgets across attention heads. It assigns compact budgets to local heads and larger budgets to global heads, guided by an early-stage classification mechanism.
Stratified Token Eviction (STE): A specific strategy for global heads that partitions tokens by distance to ensure the retention of informative long-range dependencies, addressing the limitation of standard importance scoring which often favors local tokens.
Empirical Validation: Extensive experiments demonstrate that HeadKV significantly reduces KV cache memory usage and improves inference throughput across multiple models (Janus-Pro, Lumina-mGPT, LlamaGen-XL) while maintaining high image quality.

Results

The authors evaluated HeadKV on four autoregressive image generation backbones: Janus-Pro-1B, Janus-Pro-7B, Lumina-mGPT-768, and LlamaGen-XL.

Image Quality: HeadKV achieves performance comparable to the full-cache baseline on benchmarks such as GenEval and DPG. For instance, on Janus-Pro-7B with a compression ratio ( $\rho$ ) of 1/8, HeadKV maintains a GenEval score of 0.99 (vs. 0.99 for Full) and a DPG score of 83.78 (vs. 84.20 for Full).
Comparison with LineAR: Compared to the uniform allocation method LineAR, HeadKV shows less performance degradation as compression ratios increase (e.g., at $\rho=1/8$ ). On LlamaGen-XL, HeadKV achieves a better FID-50K score (2.75) compared to LineAR (2.95) under extreme compression.
Efficiency: HeadKV significantly reduces memory usage and improves throughput.
- On LlamaGen-XL ( $\rho=1/6$ ), HeadKV reduces memory by 63.60% and achieves a 3.00× speedup compared to the full setting, outperforming LineAR (2.93×).
- On Janus-Pro-7B ( $\rho=1/6$ ), it achieves a 3.68× speedup with a 52.08% memory saving.
- On Lumina-mGPT-768, while latency is slightly higher than LineAR due to implementation overheads in grouping, it still achieves a 1.35× speedup with significant memory savings.

Significance and Claims

The paper claims that HeadKV offers a practical solution to the memory bottleneck in autoregressive image generation without requiring additional training or external supervision. By recognizing and exploiting the inherent heterogeneity of attention heads, the method moves beyond the "one-size-fits-all" cache allocation of previous works.

The authors emphasize that their approach is generalizable across different inputs and model architectures because it relies on intrinsic attention patterns observed during the early stages of generation. The Stratified Token Eviction strategy is highlighted as a critical component for maintaining the quality of generated images under tight memory budgets, specifically by preserving long-range structural information that is often lost in standard eviction methods. The work suggests that adaptive, head-specific resource allocation is a more efficient paradigm for visual generation than uniform token retention.