dKV-Cache : Le cache pour les modèles de langage à diffusion

dKV-Cache: The Cache for Diffusion Language Models

May 21, 2025

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

cs.AI

Résumé

Les modèles de langage par diffusion (Diffusion Language Models, DLMs) sont apparus comme des concurrents prometteurs face aux modèles de langage autorégressifs. Cependant, les DLMs ont longtemps été limités par une inférence lente. Un défi majeur réside dans le fait que leur architecture non autorégressive et leur attention bidirectionnelle empêchent l'utilisation du cache clé-valeur (KV-Cache) qui accélère le décodage. Nous abordons ce goulot d'étranglement en proposant un mécanisme similaire au KV-Cache, appelé KV-Cache différé (delayed KV-Cache), pour le processus de débruitage des DLMs. Notre approche est motivée par l'observation que les différents tokens présentent des dynamiques de représentation distinctes tout au long du processus de diffusion. En conséquence, nous proposons une stratégie de mise en cache différée et conditionnée pour les états clé et valeur. Nous concevons deux variantes complémentaires pour mettre en cache les clés et valeurs étape par étape : (1) dKV-Cache-Decode, qui offre une accélération quasi sans perte et améliore même les performances sur les séquences longues, suggérant que les DLMs existants pourraient sous-utiliser l'information contextuelle lors de l'inférence. (2) dKV-Cache-Greedy, qui adopte une mise en cache agressive avec une durée de vie réduite, obtenant des gains de vitesse plus importants avec une complexité temporelle quadratique au prix d'une certaine dégradation des performances. En fin de compte, dKV-Cache permet une accélération de l'inférence de 2 à 10 fois, réduisant considérablement l'écart entre les modèles autorégressifs (ARs) et les DLMs. Nous évaluons notre dKV-Cache sur plusieurs benchmarks, démontrant une accélération dans des tâches de compréhension générale du langage, de raisonnement mathématique et de génération de code. Les expériences montrent que le cache peut également être utilisé dans les DLMs, même de manière sans entraînement à partir des DLMs actuels.

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.