dKV-Cache: O Cache para Modelos de Linguagem de Difusão

dKV-Cache: The Cache for Diffusion Language Models

May 21, 2025

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

cs.AI

Resumo

Os Modelos de Linguagem de Difusão (DLMs, na sigla em inglês) têm sido vistos como um concorrente promissor para os modelos de linguagem autoregressivos. No entanto, os modelos de linguagem de difusão têm sido limitados há muito tempo pela inferência lenta. Um desafio central é que sua arquitetura não autoregressiva e atenção bidirecional impedem o uso do cache de chave-valor (KV-cache) que acelera a decodificação. Nós abordamos esse gargalo propondo um mecanismo semelhante ao KV-cache, chamado de KV-Cache Atrasado (delayed KV-Cache), para o processo de remoção de ruído dos DLMs. Nossa abordagem é motivada pela observação de que diferentes tokens têm dinâmicas de representação distintas ao longo do processo de difusão. Consequentemente, propomos uma estratégia de cache atrasado e condicionado para os estados de chave e valor. Projetamos duas variantes complementares para armazenar em cache as chaves e valores passo a passo: (1) dKV-Cache-Decode, que oferece aceleração quase sem perdas e até melhora o desempenho em sequências longas, sugerindo que os DLMs existentes podem subutilizar informações contextuais durante a inferência. (2) dKV-Cache-Greedy, que possui um cache agressivo com tempo de vida reduzido, alcançando maiores ganhos de velocidade com complexidade de tempo quadrática, ao custo de alguma degradação de desempenho. O dKV-Cache, em última análise, alcança uma aceleração de 2 a 10 vezes na inferência, reduzindo significativamente a lacuna entre os modelos autoregressivos (ARs) e os DLMs. Avaliamos nosso dKV-Cache em vários benchmarks, proporcionando aceleração em tarefas de compreensão geral de linguagem, matemática e geração de código. Os experimentos demonstram que o cache também pode ser utilizado em DLMs, mesmo de forma livre de treinamento a partir dos DLMs atuais.

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.