1. 문제 정의 (Problem Statement)

확산 기반 대규모 언어 모델 (dLLMs) 은 기존 autoregressive (AR) 모델과 달리 양방향 어텐션 (bidirectional attention) 을 사용하여 전체 시퀀스를 병렬적으로 생성합니다. 그러나 이 아키텍처는 다음과 같은 근본적인 한계로 인해 추론 비용이 매우 높습니다.

• KV 캐싱의 불가능성: AR 모델에서는 새로운 토큰이 추가될 때 이전 토큰의 Key-Value (KV) 쌍이 변하지 않아 KV 캐싱이 손실 없이 가능합니다. 반면, dLLM 은 한 토큰의 마스크가 해제 (unmask) 되면 전체 시퀀스의 표현이 변경되므로, 매 디노이징 (denoising) 단계마다 전체 시퀀스에 대한 전체 포워드 패스 (full forward pass) 를 수행해야 합니다.
• 기존 방법의 한계: 기존 근사 KV 캐싱 방법들 (Fast-dLLM, d2Cache, Elastic-Cache 등) 은 캐시를 선택적으로 업데이트하여 비용을 줄이려 시도했습니다. 하지만 이러한 방법들은 문맥 길이 (context length) 나 모델 깊이 (model depth) 에 비례하는 높은 결정 오버헤드 (decision overhead) 를 가지거나, 정확도와 처리량 간의 트레이드오프가 심했습니다.

2. 방법론 (Methodology)

저자들은 EntropyCache 를 제안하며, 이는 학습이 필요 없는 (training-free) KV 캐싱 전략입니다. 이 방법은 두 가지 실증적 관찰 (Empirical Observations) 에 기반합니다.

핵심 관찰 (Key Observations)

1. 해독된 토큰의 엔트로피는 KV 캐시 드리프트를 예측합니다:
   • 이전 디노이징 단계에서 새로 해독된 토큰 분포의 최대 엔트로피 (Maximum Entropy) 가 높을수록, 다음 단계에서 KV 캐시의 드리프트 (변화) 가 큽니다.
   • 즉, 높은 엔트로피는 모델이 불확실한 토큰을 선택했음을 의미하며, 이는 내부 표현의 큰 변화를 유발하므로 캐시를 재계산 (recompute) 해야 함을 나타내는 저렴한 신호가 됩니다.
2. 특성 변동성 (Feature Volatility) 은 여러 단계 지속됩니다:
   • 기존 동적 캐싱 방법들은 토큰이 해독된 직후의 한 단계만 재계산하는 경향이 있었습니다. 하지만 분석 결과, 일부 토큰은 해독된 후에도 여러 디노이징 단계에 걸쳐 불안정하게 유지됩니다.
   • 따라서 직전 단계뿐만 아니라 최근에 해독된 k 개의 토큰까지 재계산 범위로 포함해야 정확도를 유지할 수 있습니다.

EntropyCache 의 작동 원리

각 디노이징 단계 $t$ 에서 다음과 같은 3 단계 프로세스를 수행합니다.

1. 포워드 패스 (Forward Pass):
   • 이전 단계의 결정에 따라 전체 시퀀스를 계산하거나, 캐시된 KV 와 함께 현재 마스크 토큰과 최근 $k$ 개 토큰만 부분적으로 계산합니다.
2. 해독 및 엔트로피 평가 (Decode & Entropy Evaluation):
   • 모델 로짓에서 새로운 토큰을 해독하고, 해당 토큰 분포의 최대 엔트로피 ($E_{t+1}$) 를 계산합니다.
   • 계산 비용은 어휘 크기 $V$ 에만 비례 ($O(V)$) 하며, 문맥 길이 $L$ 이나 모델 크기와 무관합니다.
3. 스킵 결정 및 최근 토큰 선택 (Skip Decision & Selection):
   • 임계값 ($\tau$) 비교: $E_{t+1} > \tau$ 이면 불확실성이 크다고 판단하여 다음 단계에서 전체 포워드 패스를 수행합니다.
   • 부분 재계산: $E_{t+1} \le \tau$ 이면 캐시를 재사용하되, 현재 마스크 토큰과 최근 $k$ 개 해독 토큰에 대해서만 KV 를 재계산합니다.
   • 최근 토큰 선택: 해독 히스토리 ($\mathcal{H}$) 를 유지하여 가장 최근에 해독된 $k$ 개의 토큰을 선택합니다. 이는 고정된 토큰 수 ($k$) 를 기반으로 하므로, 병렬 해독 시 한 번에 많은 토큰이 생성되더라도 계산 비용을 일정하게 유지합니다.

3. 주요 기여 (Key Contributions)

1. 경량화된 캐시 노후화 신호 (Lightweight Staleness Proxy):
   • 고비용의 어텐션 롤아웃 (attention rollout) 이나 은닉 상태 비교 대신, 최대 엔트로피라는 단일 스칼라 값을 사용하여 KV 캐시 드리프트를 예측합니다.
   • 이 결정 오버헤드는 문맥 길이, 모델 깊이, 은닉 차원과 무관한 상수 시간 ($O(V)$) 으로만 발생합니다.
2. 다단계 특성 변동성 대응:
   • 토큰 해독 후에도 특성 변동성이 지속된다는 사실을 발견하고, 이를 해결하기 위해 최근 $k$ 개 토큰을 재계산하는 로컬 리컴퓨팅 전략을 도입했습니다.
3. 성능 및 효율성 입증:
   • LLaDA-8B-Instruct 와 Dream-7B-Instruct 모델에서 기존 방법들을 압도하는 속도와 정확도를 달성했습니다.

4. 실험 결과 (Results)

실험은 NVIDIA RTX 3090 GPU 에서 수행되었으며, GSM8K, MATH500, MBPP, HumanEval 등 표준 벤치마크와 CoT (Chain-of-Thought) 벤치마크에서 평가되었습니다.

• 처리량 (Throughput) 향상:
  • 표준 벤치마크: LLaDA-8B 에서 15.2 배, Dream-7B 에서 26.4 배의 속도 향상 (Speedup) 을 달성했습니다.
  • CoT 벤치마크: 긴 문맥을 요구하는 작업에서 더 큰 효과를 보였으며, LLaDA 에서 22.4 배, Dream 에서 24.1 배의 속도 향상을 기록했습니다.
• 정확도 (Accuracy):
  • EntropyCache 는 대부분의 벤치마크에서 베이스라인 (전체 재계산) 과 경쟁력 있는 정확도를 유지하거나, 다른 캐싱 방법들보다 더 높은 정확도를 기록했습니다. 특히 HumanEval 과 같은 코드 생성 작업에서 다른 방법들이 정확도가 떨어지는 반면, EntropyCache 는 베이스라인 수준의 정확도를 유지했습니다.
• 오버헤드 분석:
  • EntropyCache 의 결정 로직은 전체 추론 시간의 0.5% 만 차지합니다. (비교: d2Cache 9.2%, Elastic-Cache 14.5%). 이는 문맥 길이가 길어질수록 상대적인 효율성이 더욱 커짐을 의미합니다.

5. 의의 및 결론 (Significance)

EntropyCache 는 확산 기반 언어 모델 (dLLMs) 의 추론 병목 현상을 해결하는 획기적인 접근법입니다.

• 상수 오버헤드: 기존 방법들이 문맥 길이나 모델 크기에 따라 오버헤드가 증가하는 문제를 해결하여, 대규모 모델과 긴 문맥에서도 일관된 효율성을 제공합니다.
• 정확도 - 속도 트레이드오프 최적화: 단순한 캐시 재사용이 아니라, 모델의 불확실성 (엔트로피) 을 기반으로 "언제" 재계산할지 지능적으로 판단함으로써, 속도 향상과 정확도 유지 사이의 최적 균형을 찾았습니다.
• 실용성: 학습이 필요 없으며 (training-free), 기존 dLLM 아키텍처에 쉽게 통합할 수 있어 실제 배포 환경에서 dLLM 의 에너지 효율성과 응답 속도를 크게 개선할 수 있는 잠재력을 가집니다.

결론적으로, 이 연구는 dLLM 의 실용화를 가로막는 계산 비용 문제를 해결하고, 확산 모델이 생성형 AI 분야에서 autoregressive 모델과 경쟁할 수 있는 강력한 대안이 될 수 있음을 입증했습니다.