Эффективное использование памяти в визуальном авторегрессионном моделировании с масштабно-осознанным сжатием кэша ключей и значений

Memory-Efficient Visual Autoregressive Modeling with Scale-Aware KV Cache Compression

May 26, 2025

Авторы: Kunjun Li, Zigeng Chen, Cheng-Yen Yang, Jenq-Neng Hwang

cs.AI

Аннотация

Визуальное авторегрессивное моделирование (VAR) привлекло значительное внимание благодаря своему инновационному подходу к прогнозированию на следующем масштабе, что приводит к существенному улучшению эффективности, масштабируемости и обобщения в условиях нулевого сэмплинга. Однако методология "от грубого к точному", присущая VAR, вызывает экспоненциальный рост кэша ключей и значений (KV) во время вывода, что приводит к значительному потреблению памяти и вычислительной избыточности. Для устранения этих узких мест мы представляем ScaleKV — новый фреймворк для сжатия кэша KV, разработанный специально для архитектур VAR. ScaleKV основывается на двух ключевых наблюдениях: различной потребности в кэше между слоями трансформера и различиях в паттернах внимания на разных масштабах. На основе этих инсайтов ScaleKV классифицирует слои трансформера на две функциональные группы: черновики и уточнители. Черновики демонстрируют рассеянное внимание на нескольких масштабах, что требует большего объема кэша. Напротив, уточнители фокусируют внимание на текущей карте токенов для обработки локальных деталей, что, в свою очередь, значительно снижает потребность в кэше. ScaleKV оптимизирует многомасштабный конвейер вывода, идентифицируя черновики и уточнители для каждого масштаба, что позволяет осуществлять дифференцированное управление кэшем, адаптированное под каждый масштаб. Оценка на семействе современных моделей VAR для генерации изображений из текста, Infinity, демонстрирует, что наш подход эффективно сокращает требуемую память для кэша KV до 10%, сохраняя при этом точность на уровне пикселей.

English

Visual Autoregressive (VAR) modeling has garnered significant attention for its innovative next-scale prediction approach, which yields substantial improvements in efficiency, scalability, and zero-shot generalization. Nevertheless, the coarse-to-fine methodology inherent in VAR results in exponential growth of the KV cache during inference, causing considerable memory consumption and computational redundancy. To address these bottlenecks, we introduce ScaleKV, a novel KV cache compression framework tailored for VAR architectures. ScaleKV leverages two critical observations: varying cache demands across transformer layers and distinct attention patterns at different scales. Based on these insights, ScaleKV categorizes transformer layers into two functional groups: drafters and refiners. Drafters exhibit dispersed attention across multiple scales, thereby requiring greater cache capacity. Conversely, refiners focus attention on the current token map to process local details, consequently necessitating substantially reduced cache capacity. ScaleKV optimizes the multi-scale inference pipeline by identifying scale-specific drafters and refiners, facilitating differentiated cache management tailored to each scale. Evaluation on the state-of-the-art text-to-image VAR model family, Infinity, demonstrates that our approach effectively reduces the required KV cache memory to 10% while preserving pixel-level fidelity.