dKV-Cache: De Cache voor Diffusie Taalmodellen

dKV-Cache: The Cache for Diffusion Language Models

May 21, 2025

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

cs.AI

Samenvatting

Diffusie Taalmodellen (DLMs) worden gezien als een veelbelovende concurrent voor autoregressieve taalmodellen. Echter, diffusie taalmodellen zijn lange tijd beperkt geweest door trage inferentie. Een kernuitdaging is dat hun niet-autoregressieve architectuur en bidirectionele aandacht het key-value cache mechanisme uitsluiten dat decodering versnelt. Wij pakken dit knelpunt aan door een KV-cache-achtig mechanisme, de vertraagde KV-Cache, voor te stellen voor het denoisingsproces van DLMs. Onze aanpak is gemotiveerd door de observatie dat verschillende tokens verschillende representatiedynamieken hebben gedurende het diffusieproces. Dienovereenkomstig stellen wij een vertraagde en geconditioneerde cachingstrategie voor voor key- en value-states. We ontwerpen twee complementaire varianten om key en value stap-voor-stap te cachen: (1) dKV-Cache-Decode, dat een bijna verliesloze versnelling biedt, en zelfs de prestaties op lange sequenties verbetert, wat suggereert dat bestaande DLMs contextuele informatie tijdens inferentie mogelijk onderbenutten. (2) dKV-Cache-Greedy, dat een agressieve caching heeft met een verkorte levensduur, waardoor hogere snelheidsverbeteringen worden bereikt met kwadratische tijdcomplexiteit ten koste van enige prestatievermindering. dKV-Cache bereikt uiteindelijk een versnelling van 2-10x in inferentie, waardoor de kloof tussen ARs en DLMs aanzienlijk wordt verkleind. We evalueren onze dKV-Cache op verschillende benchmarks, waarbij versnelling wordt geleverd over benchmarks voor algemeen taalbegrip, wiskunde en codegeneratie. Experimenten tonen aan dat cache ook kan worden gebruikt in DLMs, zelfs op een trainingsvrije manier vanuit huidige 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.