dKV-Cache: Кэш для диффузионных языковых моделей

dKV-Cache: The Cache for Diffusion Language Models

May 21, 2025

Авторы: Xinyin Ma, Runpeng Yu, Gongfan Fang, Xinchao Wang

cs.AI

Аннотация

Диффузионные языковые модели (DLMs) рассматриваются как перспективная альтернатива авторегрессивным языковым моделям. Однако долгое время их применение ограничивалось медленным выводом. Основная проблема заключается в том, что их неавторегрессивная архитектура и двунаправленное внимание исключают использование кэша ключей и значений (KV-cache), который ускоряет декодирование. Мы устраняем это узкое место, предлагая механизм, аналогичный KV-cache, — отложенный KV-Cache (delayed KV-Cache) для процесса удаления шума в DLMs. Наш подход основан на наблюдении, что разные токены имеют различные динамики представления в ходе диффузионного процесса. Соответственно, мы предлагаем стратегию отложенного и условного кэширования состояний ключей и значений. Мы разработали два взаимодополняющих варианта для пошагового кэширования ключей и значений: (1) dKV-Cache-Decode, который обеспечивает практически без потерь ускорение и даже улучшает производительность на длинных последовательностях, что указывает на недостаточное использование контекстной информации существующими DLMs в процессе вывода. (2) dKV-Cache-Greedy, который использует агрессивное кэширование с сокращённым сроком жизни, достигая большего ускорения с квадратичной временной сложностью ценой некоторого снижения производительности. В итоге dKV-Cache обеспечивает ускорение вывода в 2–10 раз, значительно сокращая разрыв между авторегрессивными и диффузионными моделями. Мы оценили наш dKV-Cache на нескольких бенчмарках, демонстрируя ускорение в задачах общего понимания языка, математических задач и генерации кода. Эксперименты показывают, что кэш может быть успешно применён в DLMs даже без дополнительного обучения на основе текущих моделей.

English

Diffusion Language Models (DLMs) have been seen as a promising competitor for autoregressive language models. However, diffusion language models have long been constrained by slow inference. A core challenge is that their non-autoregressive architecture and bidirectional attention preclude the key-value cache that accelerates decoding. We address this bottleneck by proposing a KV-cache-like mechanism, delayed KV-Cache, for the denoising process of DLMs. Our approach is motivated by the observation that different tokens have distinct representation dynamics throughout the diffusion process. Accordingly, we propose a delayed and conditioned caching strategy for key and value states. We design two complementary variants to cache key and value step-by-step: (1) dKV-Cache-Decode, which provides almost lossless acceleration, and even improves performance on long sequences, suggesting that existing DLMs may under-utilise contextual information during inference. (2) dKV-Cache-Greedy, which has aggressive caching with reduced lifespan, achieving higher speed-ups with quadratic time complexity at the cost of some performance degradation. dKV-Cache, in final, achieves from 2-10x speedup in inference, largely narrowing the gap between ARs and DLMs. We evaluate our dKV-Cache on several benchmarks, delivering acceleration across general language understanding, mathematical, and code-generation benchmarks. Experiments demonstrate that cache can also be used in DLMs, even in a training-free manner from current DLMs.