dKV-Cache: La caché para modelos de lenguaje de difusión

dKV-Cache: The Cache for Diffusion Language Models

May 21, 2025

Autores: Xinyin Ma, Runpeng Yu, Gongfan Fang, Xinchao Wang

cs.AI

Resumen

Los modelos de lenguaje basados en difusión (DLMs, por sus siglas en inglés) han sido considerados como un competidor prometedor frente a los modelos de lenguaje autoregresivos. Sin embargo, los modelos de lenguaje basados en difusión han estado limitados durante mucho tiempo por una inferencia lenta. Un desafío central es que su arquitectura no autoregresiva y su atención bidireccional impiden el uso de la caché de clave-valor que acelera la decodificación. Abordamos este cuello de botella proponiendo un mecanismo similar a la caché KV, denominado KV-Cache retrasado, para el proceso de eliminación de ruido en los DLMs. Nuestro enfoque se basa en la observación de que diferentes tokens tienen dinámicas de representación distintas a lo largo del proceso de difusión. En consecuencia, proponemos una estrategia de almacenamiento en caché retrasada y condicionada para los estados de clave y valor. Diseñamos dos variantes complementarias para almacenar en caché clave y valor paso a paso: (1) dKV-Cache-Decode, que proporciona una aceleración casi sin pérdidas e incluso mejora el rendimiento en secuencias largas, sugiriendo que los DLMs existentes podrían estar subutilizando la información contextual durante la inferencia. (2) dKV-Cache-Greedy, que implementa un almacenamiento en caché más agresivo con una vida útil reducida, logrando mayores aceleraciones con complejidad temporal cuadrática a costa de cierta degradación en el rendimiento. En última instancia, dKV-Cache logra una aceleración de 2 a 10 veces en la inferencia, reduciendo significativamente la brecha entre los modelos autoregresivos (ARs) y los DLMs. Evaluamos nuestro dKV-Cache en varios benchmarks, demostrando aceleraciones en tareas de comprensión general del lenguaje, matemáticas y generación de código. Los experimentos muestran que la caché también puede utilizarse en DLMs, incluso de manera libre de entrenamiento a partir de los DLMs actuales.

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.