dKV-Cache: Der Cache für Diffusions-Sprachmodelle

dKV-Cache: The Cache for Diffusion Language Models

May 21, 2025

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

cs.AI

Zusammenfassung

Diffusion Language Models (DLMs) gelten als vielversprechende Konkurrenten zu autoregressiven Sprachmodellen. Allerdings waren Diffusion Language Models lange Zeit durch langsame Inferenz eingeschränkt. Eine zentrale Herausforderung besteht darin, dass ihre nicht-autoregressive Architektur und bidirektionale Aufmerksamkeit den Key-Value-Cache ausschließen, der die Dekodierung beschleunigt. Wir adressieren diesen Engpass durch einen KV-Cache-ähnlichen Mechanismus, den verzögerten KV-Cache, für den Denoising-Prozess von DLMs. Unser Ansatz wird durch die Beobachtung motiviert, dass verschiedene Tokens unterschiedliche Repräsentationsdynamiken während des Diffusionsprozesses aufweisen. Entsprechend schlagen wir eine verzögerte und bedingte Caching-Strategie für Key- und Value-Zustände vor. Wir entwerfen zwei komplementäre Varianten, um Key und Value schrittweise zu cachen: (1) dKV-Cache-Decode, das eine nahezu verlustfreie Beschleunigung bietet und sogar die Leistung bei langen Sequenzen verbessert, was darauf hindeutet, dass bestehende DLMs kontextuelle Informationen während der Inferenz möglicherweise nicht ausreichend nutzen. (2) dKV-Cache-Greedy, das ein aggressives Caching mit reduzierter Lebensdauer aufweist und höhere Beschleunigungen mit quadratischer Zeitkomplexität auf Kosten eines gewissen Leistungsverlusts erreicht. dKV-Cache erzielt schließlich eine 2- bis 10-fache Beschleunigung in der Inferenz und verringert damit die Lücke zwischen ARs und DLMs erheblich. Wir evaluieren unseren dKV-Cache anhand mehrerer Benchmarks und erzielen Beschleunigungen in allgemeinem Sprachverständnis, mathematischen und Code-Generierungs-Benchmarks. Experimente zeigen, dass der Cache auch in DLMs verwendet werden kann, sogar in trainingsfreier Weise mit aktuellen 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.