윈도우 디퓨전 토큰 프루닝과 캐싱으로 확산 언어 모델 추론 가속화

초록

본 논문은 사전 학습된 확산 언어 모델(DLM)의 추론 단계에서 전체 시퀀스에 대한 반복적인 연산을 제거하기 위해, 토큰을 활성, 버퍼, 원거리 세 그룹으로 나누고 윈도우 기반으로 연산을 제한하는 방법을 제안한다. 활성 토큰은 실시간으로 계산하고, 버퍼 토큰은 KV 상태를 캐시·주기적으로 갱신하며, 원거리 토큰은 윈도우 밖에서 완전히 프루닝한다. 실험 결과, 동일한 연산 예산 하에 최대 99배의 속도 향상을 달성하면서도 생성 품질 저하가 미미함을 보였다.

상세 분석

본 연구는 확산 언어 모델(DLM)의 추론 과정이 실제로는 매우 높은 구조적 국소성을 가진다는 세 가지 핵심 관찰을 기반으로 한다. 첫째, 각 디노이징 단계에서 업데이트가 일어나는 토큰은 전체 미디코드 영역의 앞쪽(prefix)에 집중되는 경향이 강해, 먼 위치의 토큰은 거의 변하지 않는다. 이는 “활성 토큰”이 주로 프리픽스에 존재한다는 의미이며, 원거리 토큰을 매 단계 재계산하는 것은 큰 연산 낭비가 된다. 둘째, 활성 토큰이 필요로 하는 컨텍스트는 제한된 길이의 프리픽스만으로 충분히 정보를 제공한다. 실험적으로 남은 미디코드 토큰을 길이 W 로 잘라서 KL 발산을 측정했을 때, W 가 작아도 활성 토큰의 예측 분포와 전체 시퀀스 기준이 거의 동일함을 확인했다. 또한, 비활성 토큰의 KV 상태를 이전 단계에서 캐시하고 재사용하면 정확도 손실이 거의 없으며, 캐시 없이 단순 트렁케이션만 할 경우보다 더 작은 KL을 보였다. 셋째, 디코딩된 토큰의 표현은 시간적으로 두 단계로 구분된다. 새로 디코딩된 토큰은 디노이징 직후 급격히 변하는 “포스트 디코드 트랜지언트”를 겪지만, 일정 시간 이상 지난 토큰은 V값이 높은 코사인 유사도를 유지하며 거의 정착한다. 이러한 안정성은 KV 캐시를 장기간 유지할 수 있음을 의미한다. 위 세 관찰을 종합하면, 전체 시퀀스에 대한 매 단계 전면 연산은 불필요하고, 윈도우 안에서만 활성·버퍼 토큰을 대상으로 연산하고, 원거리 토큰은 프루닝하거나 캐시만 유지하면 된다. 따라서 저자는 (i) 활성 토큰을 실시간으로 계산, (ii) 버퍼 토큰은 KV 캐시를 주기적으로 갱신, (iii) 원거리 토큰은 윈도우 밖에서 완전히 제외하는 “윈도우‑디퓨전” 프레임워크를 설계했다. 윈도우는 디노이징 진행에 따라 오른쪽으로 슬라이딩하며, 각 단계마다 활성·버퍼 영역이 재조정된다. 실험에서는 LLaDA와 Dream 두 대규모 사전 학습 모델에 적용했으며, 동일한 연산 예산(플롭수) 기준으로 기존 방식 대비 2.3×6.6×의 기본 가속을 달성하고, 어댑티브 길이 추론과 결합했을 때 최대 99배의 속도 향상을 기록했다. 품질 측면에서는 BLEU, ROUGE, 코드 정확도 등 다양한 벤치마크에서 0.10.3% 수준의 미세한 성능 저하만 관측되었다. 이 결과는 사전 학습된 DLM을 재학습 없이도 효율적으로 가속화할 수 있음을 입증한다.