ATTENTION IS ALL YOU NEED FOR KV CACHE IN DIFFUSION LLMS 논문리뷰

Elastic-Cache: diffusion LLM을 위한 적응형 KV 캐시 관리 

 

확산 거대 언어 모델(DLMs)은 반복적인 디노이징(denoising) 절차로 인해 추론 과정(inference)에서 엄청난 계산량

표준 구현 방식에서는 매 디노이징 단계와 모든 레이어에서 모든 토큰에 대한 쿼리(Q), 키(K), 값(V)을 다시 계산

그러나 KV 상태는 대부분의 단계에서, 특히 얕은 레이어(shallow layers)에서는 거의 변하지 않기 때문에, 이러한 방식은 상당한 **중복성(redundancy)**을 유발하며 지연 시간(latency)을 증가

 

이러한 DLM 디코딩의 중복성을 체크하고, 예측 품질을 극대화하면서 wall-clock latency을 최소화하기 위해 KV 캐시를 언제, 어떻게 적응적으로 다시 계산할지 결정

#--------------------------------------------------------------

 

2. Elastic-Cache의 기반이 되는 경험적 관찰

저자들은 낭비적인 재계산을 방지하고 비용을 절감하기 위한 세 가지 주요 경험적 관찰을 제시

1. MASK 토큰의 블록 캐싱: 멀리 떨어진 MASK 토큰은 현재 토큰을 언마스킹하는 데 미미한 영향을 미치며, 주로 길이 편향 사전 정보(length-bias prior) 역할

 

->활성 예측 창(active prediction window) 밖의 MASK 토큰의 KV는 블록 단위로 캐시하여 불필요한 작업을 피할 수 있습니다.

 

2. 깊이에 따른 KV 드리프트 증가: KV 드리프트(drift)—즉, 캐시된 키와 값의 단계별 변화—는 레이어 깊이가 증가함에 따라 커지는 경향

 

 이는 얕은 레이어(local lexical structure 인코딩)는 빨리 안정화되는 반면, 깊은 레이어(global, semantic dependencies 조정)는 계속 변화

 

이 관찰은 깊은 레이어부터 선택적으로 새로 고침(refresh)을 시작하는 것이 효율적임을 뜻(이론적으로 **레이어별 KV 드리프트 단조성(Layer-Wise KV Drift Monotonicity)**으로 formalized).

3. 가장 많이 어텐션 받은 토큰의 안정성:

 

주어진 단계에서 **가장 많이 어텐션 받은 토큰(most-attended token)**은 일반적으로 가장 작은 KV drift

--> 이는 캐시 변경에 대한 conservative lower bound을 제공

 이 토큰의 드리프트를 모니터링하면 오버헤드가 낮은 신뢰할 수 있는 전역 새로 고침 트리거를 제공

이 속성은 어텐션 집중 및 드리프트(Attention Concentration and Drift) 이론으로 

 

#--------------------------------------------------------------

 

3. Elastic-Cache 프레임워크 (훈련 불필요)

이러한 관찰을 바탕으로, 저자들은 훈련이 필요 없는(training-free), 아키텍처 비종속적인(architecture-agnostic) 전략인 Elastic-Cache를 제안합니다. Elastic-Cache는 Attention-Aware KV Cache Update와 Layer-Aware KV Cache Update를 결합합니다.

• Attention-Aware (언제 새로 고칠 것인가):

    ◦ 가장 많이 어텐션 받은 토큰에 대해 경량의 **드리프트 통계(drift statistic)**를 계산합니다.

    ◦ 이 통계가 임계값(γ)을 초과하면, 새로 고침(refresh)을 트리거합니다.

    ◦ 드리프트는 현재 단계(t)와 이전 단계(t−1) 사이의 가장 많이 어텐션 받은 토큰의 어텐션 가중치(attention weights) 변화를 코사인 유사도(σt,l​)로 측정하여 결정됩니다.

• Layer-Aware (어디를 새로 고칠 것인가):

    ◦ 새로 고침이 트리거되면, 선택된 레이어(\ell^{*})부터 마지막 레이어(L)까지만 KV 캐시를 다시 계산합니다 (ℓ≥ℓ∗).

    ◦ 얕은 레이어는 캐시를 재사용하고, 활성 예측 창 밖의 MASK 토큰도 블록 캐싱을 통해 재사용됩니다.

• 슬라이딩 윈도우 디코딩 (Sliding Window Decoding): 효율성을 높이기 위해, 모델은 가장 왼쪽에 가깝고 크기 β를 갖는 마스킹된 위치(Mtβ​)에 대해서만 예측을 수행하는 슬라이딩 윈도우 디코딩을 사용합니다. 이는 멀리 떨어진 MASK 토큰의 캐시 손실을 최소화하는 데 도움이 됩니다.

 

#--------------------------------------------------------------

 

4. 실험 결과 및 효율성

Elastic-Cache는 LLaDA-Instruct, LLaDA-1.5, LLaDA-V 등의 모델과 GSM8K, HumanEval 등 다양한 작업에서 일관된 속도 향상을 입증했습니다.

• 주요 성능: Elastic-Cache는 표준 기준선(baseline) 대비 최대 45.1배의 상당한 처리량(throughput) 향상을 달성했습니다.

    ◦ LLaDA-1.5 모델의 GSM8K (512 토큰)에서 45.1배 속도 향상을 기록했습니다.

    ◦ HumanEval에서 4.8배 속도 향상을 달성했습니다.

    ◦ GSM8K에서 기존 신뢰도 기반 접근 방식(Fast-dLLM) 대비 6.8배 높은 처리량(throughput)을 달성했습니다.

• 정확도 유지: 캐시를 적응적으로 업데이트함으로써, Elastic-Cache는 생성 품질의 손실을 무시할 수 있는 수준으로 유지하거나, GSM8K 및 MBPP 같은 일부 작업에서는 Fast-dLLM보다 더 높은 정확도를 유지했습니다.

• 확장성: Elastic-Cache는 생성 길이(generation length)가 길어질수록 처리량이 증가하는 독특한 확장 속성을 보였습니다. 이는 고정 크기의 슬라이딩 윈도우와 자동 캐시 업데이트 덕분입니다. 또한 모델의 정확도가 높을수록(즉, 예측이 정확할수록) Elastic-Cache의 효율성(처리량)이 더 높아지는 경향이 있음을 관찰했습니다.