dKV-Cache: The Cache for Diffusion Language Models

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.

한 줄 요약

dKV-Cache는 확산 언어 모델(DLM)의 추론 속도를 2~10배 가속화하는 지연형 키-밸류 캐시 메커니즘을 제안한다.

핵심 기여도

dKV-Cache-Decode는 거의 손실 없는 가속화를 제공하며, 긴 시퀀스에서도 성능 향상 (LLaDA, Dream 모델 기준).
dKV-Cache-Greedy는 캐시 수명을 줄여 2차 시간 복잡도로 더 빠른 추론을 달성하지만, 일부 성능 저하 발생.
기존 DLM에서 훈련 없이도 적용 가능한 캐시 기법 제시.
7B 규모 DLM에서 2~10× 추론 가속화를 달성.

핵심 아이디어

DLM은 비자율적 아키텍처와 양방향 어텐션으로 인해 기존 KV-Cache를 사용할 수 없었다. 그러나 토큰의 표현 동역학을 분석한 결과, 키-밸류 상태는 이미 디코딩된 토큰에 대해서는 재사용 가능하다는 점을 발견했다. 이에 따라, 디코딩된 토큰의 키-밸류를 **지연된 조건부 캐시** 방식으로 저장하는 dKV-Cache를 제안했다. 특히, 토큰이 디코딩되기 전까지는 캐시를 지연시켜 메모리 사용량을 줄이고, 디코딩 단계에서만 선택적으로 재사용함으로써 효율성을 극대화했다.

기술적 접근법

**dKV-Cache-Decode**: 토큰별 디코딩 시점에 키-밸류를 캐시하여 거의 손실 없는 가속화.
**dKV-Cache-Greedy**: 캐시 수명을 줄이고, 현재 디코딩 대상 토큰과 지연 토큰, 윈도우 토큰만을 대상으로 캐시하여 2차 시간 복잡도 달성.
**지연 캐시 전략**: 키-밸류는 디코딩 완료 후 1단계 지연되어 저장됨.
**복잡도 개선**: DLM의 원래 복잡도 $ \mathcal{O}(L^3) $ → $ \mathcal{O}(L^2) $로 감소.

주요 결과

**LLaDA, Dream 모델**에서 **2~10× 추론 가속화**.
**코드 생성, 수학 문제 해결, 일반 언어 이해** 벤치마크에서 성능 저하 거의 없이 가속화.
dKV-Cache-Decode는 긴 시퀀스에서도 성능 향상 (기존 DLM이 문맥 정보를 과소 활용한다는 시사).

의의 및 한계

DLM의 추론 효율성을 획기적으로 향상시켜 AR 모델과의 격차를 줄임.
훈련 없이도 적용 가능한 캐시 기법으로, 실용적 활용 가능성이 높음.
한계로는 **시스템 수준 최적화**(메모리 관리, 하드웨어 인식 실행 등)와의 통합이 부족하며, 알고리즘적 개선에만 집중한 점이 언급됨.

실용적 활용

대규모 DLM을 활용한 **코드 생성, 수학 문제 풀이, 대화형 응용**에서 추론 속도 향상에 활용 가능.
**LLM 서비스 제공자**가 DLM을 실시간 응답 시스템에 적용할 때 유용한 기술.
**저비용 추론 인프라** 구축에 기여할 수 있는 기술.