AMS-KV: Adaptive KV Caching in Multi-Scale Visual Autoregressive Transformers

TL;DR

This paper introduces AMS-KV, a novel adaptive key-value caching strategy for multi-scale visual autoregressive transformers that significantly reduces memory usage and latency while maintaining high generation quality.

Contribution

The paper proposes AMS-KV, a scale-adaptive KV caching policy tailored for next-scale prediction in VAR models, addressing memory growth and efficiency challenges.

Findings

Reduces KV cache usage by up to 84.83%.

Lowers self-attention latency by 60.48%.

Enables stable scaling to larger batch sizes.

Abstract

Visual autoregressive modeling (VAR) via next-scale prediction has emerged as a scalable image generation paradigm. While Key and Value (KV) caching in large language models (LLMs) has been extensively studied, next-scale prediction presents unique challenges, and KV caching design for next-scale based VAR transformers remains largely unexplored. A major bottleneck is the excessive KV memory growth with the increasing number of scales-severely limiting scalability. Our systematic investigation reveals that: (1) Attending to tokens from local scales significantly contributes to generation quality (2) Allocating a small amount of memory for the coarsest scales, termed as condensed scales, stabilizes multi-scale image generation (3) Strong KV similarity across finer scales is predominantly observed in cache-efficient layers, whereas cache-demanding layers exhibit weaker inter-scale similarity. Based on the observations, we introduce AMS-KV, a scale-adaptive KV caching policy for next-scale prediction in VAR models. AMS-KV prioritizes storing KVs from condensed and local scales, preserving the most relevant tokens to maintain generation quality. It further optimizes KV cache utilization and computational efficiency identifying cache-demanding layers through inter-scale similarity analysis. Compared to the vanilla next-scale prediction-based VAR models, AMS-KV reduces KV cache usage by up to 84.83% and self-attention latency by 60.48%. Moreover, when the baseline VAR-d30 model encounters out-of-memory failures at a batch size of 128, AMS-KV enables stable scaling to a batch size of 256 with improved throughput.