dKV-Cache: La Cache per i Modelli Linguistici di Diffusione

dKV-Cache: The Cache for Diffusion Language Models

May 21, 2025

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

cs.AI

Abstract

I modelli linguistici basati su diffusione (Diffusion Language Models, DLMs) sono stati considerati un promettente concorrente dei modelli linguistici autoregressivi. Tuttavia, i DLMs sono stati a lungo limitati da un'inferenza lenta. Una sfida fondamentale è che la loro architettura non autoregressiva e l'attenzione bidirezionale impediscono l'uso della cache chiave-valore che accelera il processo di decodifica. Affrontiamo questo collo di bottiglia proponendo un meccanismo simile alla KV-cache, chiamato delayed KV-Cache, per il processo di denoising dei DLMs. Il nostro approccio è motivato dall'osservazione che token diversi hanno dinamiche di rappresentazione distinte durante il processo di diffusione. Di conseguenza, proponiamo una strategia di caching ritardato e condizionato per gli stati chiave e valore. Progettiamo due varianti complementari per memorizzare chiavi e valori passo dopo passo: (1) dKV-Cache-Decode, che fornisce un'accelerazione quasi senza perdite e migliora persino le prestazioni su sequenze lunghe, suggerendo che i DLMs esistenti potrebbero non sfruttare appieno le informazioni contestuali durante l'inferenza. (2) dKV-Cache-Greedy, che utilizza un caching aggressivo con una durata ridotta, ottenendo accelerazioni maggiori con una complessità temporale quadratica al costo di un certo degrado delle prestazioni. In definitiva, dKV-Cache raggiunge un'accelerazione da 2 a 10 volte nell'inferenza, riducendo significativamente il divario tra i modelli autoregressivi (ARs) e i DLMs. Valutiamo il nostro dKV-Cache su diversi benchmark, dimostrando un'accelerazione in ambiti come la comprensione linguistica generale, la matematica e la generazione di codice. Gli esperimenti mostrano che la cache può essere utilizzata anche nei DLMs, persino in modo training-free a partire dai DLMs attuali.

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.