기술 요약: Polestar – 효율적인 Diffusion LLM 추론을 위한 드리프트 인지형 캐시 교정 및 토큰 확정

1. 문제 정의

디퓨전 대규모 언어 모델(dLLM)은 반복적인 디노이징(denoising)을 통해 병렬 토큰 디코딩을 가능하게 함으로써 자기회귀적 LLM(AR-LLM)의 유망한 대안으로 주목받고 있다. 그러나 이들의 추론 효율성은 현재 두 가지 상호 연관된 과제에 의해 제약을 받고 있다:

1. KV-캐시 비효율성: AR-LLM과 달리, dLLM은 양방향 어텐션(bidirectional attention)을 사용하므로 매 디코딩 단계마다 컨텍스트가 업데이트됨에 따라 토큰 표현이 변화(드리프트, drift)한다. 이는 표준 Key-Value(KV) 캐싱을 비효율적으로 만들며, 캐싱된 표현이 빠르게 노후화되게 한다. 기존 솔루션들은 블록 내 표현이 정적이라고 가정하거나(정확도 저하 초래), 시퀀스 전반의 토큰 표현의 공동 진화를 포착하지 못하는 거친 프록시 신호(예: 가장 많이 어텐션되는 토큰)에 의존한다.
2. 병렬성 vs 품질 간의 트레이드오프: 디코딩 병렬성(한 번의 포워드 패스에서 여러 토큰 디코딩)을 높이는 것은 일반적으로 정적인 확정 임계값(static confidence thresholds)을 요구한다. 이러한 고정된 임계값은 다양한 디코딩 단계와 입력에 적응하지 못하여, 조기 확정(품질 저하) 또는 보수적 확정(처리량 제한)을 초래한다.

저자들은 두 가지 과제가 모두 **토큰 표현 드리프트(token representation drift)**라는 공통된 현상에서 기인한다는 점에 주목했다. 토큰이 디코딩됨에 따라, 양방향 어텐션을 통한 컨텍스트 통합으로 인해 토кеn의 은닉 상태(hidden state) 표현이 진화한다. 이 드리프트는 단순한 오류 아티팩트가 아니라 dLLM 추론의 본질적인 특성이다.

2. 방법론: Polestar 프레임워크

Polestar는 토큰 표현 드리프트를 KV-캐시 노후화와 토큰 확정 문제를 동시에 해결하기 위한 통합된 신호로 활용하는 학습이 필요 없는(training-free) 추론 프레임워크이다. 이는 세 가지 핵심 구성 요소로 이루어진다:

A. Polestar-Cache (드리프트 인지형 KV 리프레시)

전체 KV 캐시를 새로고침하거나 정적인 간격에 의존하는 대신, Polestar-Cache는 드리프트 크기를 기반으로 노후화된 위치를 식별한다.

• 드리프트 측정: 드리프트는 연속된 단계 사이의 토큰 어텐션 분포 간의 **KL-발산(KL-divergence)**을 사용하여 정량화된다 ($D^{(t)}_i = KL(A^{(t)}_i \| A^{(t-1)}_i)$). 이는 부분 집합에 대한 코사인 유사성보다 컨텍스트 변화를 더 정확하게 포착한다.
• 로컬 윈도우 및 클러스터링: 오버헤드를 최소화하기 위해, 시스템은 현재 디코딩 블록 주변의 로컬 윈도우(접두사 및 접미사 블록)를 유지한다. 시스템은 **구형 K-평균(Spherical K-Means)**을 사용하여 이 윈도우 내의 은닉 상태를 클러스터링하고, NVFP4를 사용하여 잔차(residual)를 양자화함으로써 압축된 프록시 표현을 생성한다.
• 희소 업데이트(Sparse Updates): 각 단계에서 시스템은 클러스터 중심점(centroids)을 사용하여 "프록시 어텐션"을 계산한다. 시스템은 드리프트가 심한 상위-$k$ 클러스터를 식별하고, 이 위치들에 대해서만 **희소 포워드 패스(sparse forward pass)**를 수행하여 "업데이트 패킷"을 생성한다. 이 패킷은 후속 레이어의 해당 KV-캐시 엔트리와 은닉 상태를 새로고침하여 전체 시퀀스 재계산을 방지한다.
• 스케줄: 업데이트는 디코딩된 컨텍스트가 상당한 불일치를 유발할 때(예: $\tau_{upd}$ 토큰 디코딩 후) 선택적으로 트리거되거나, 비로컬 위치에 대해 주기적인 전체 시퀀스 패스를 통해 실행된다.

B. Polestar-Commit (드리프트 인지형 토큰 확정)

Polestar-Commit은 "급격한 드리프트 이벤트(sharp drift events)"를 감지하여 토큰의 표현이 수렴했는지 확인하며, 이를 통해 정적 신뢰도 임계값보다 더 빠르고 신뢰할 수 있는 확정을 가능하게 한다.

• 현재 블록 확정: 시스템은 역사적 평균으로부터의 드리프트 편차인 **드리프트 델타($\Delta^{(t)}_j$)**를 모니터링한다. 급격한 양(+)의 델타는 빠른 컨텍스트 통합을 나타낸다. 토큰은 신뢰도 조건부 동적 임계값($\tau^{(t)}_{d,j} = \alpha(\tau_s - c^{(t)}_j)^2$)을 초과할 경우 확정된다. 이를 통해 신뢰도가 낮은 토큰이라도 강력한 수렴 신호를 보이면 확정될 수 있다.
• 접미사 블록 확정: 시스템은 마지막 레이어의 업데이트 패킷 또한 분석한다. 만약 접미사 윈도우의 클러스터 중심점이 상위 드리프트 클러스터로 선택된다면, 이는 현재 블록으로부터 강력한 컨텍스트 업데이트가 일어나고 있음을 의미한다. 접미사 토큰은 높은 신뢰도 임계값($\ge \tau_s$)을 충족할 때만 확정되어, 안전성을 확보하면서도 병렬성을 높인다.

C. 시스템 최적화

높은 처리량을 보장하기 위해 Polestar는 다음을 채택한다:

• CPU 오프로딩: 양자화된 은닉 상태 클러스터는 고정된(pinned) CPU 메모리로 오프로드되며, 중심점(centroids)만 GPU에 남는다. 상위-$k$ 클러스터는 비동기적으로 가져온다.
• 지연 시간 은닉(Latency Hiding): CUDA 스트림을 사용하여 보조 작업(클러스터링, 양자화, 프록시 어텐션, 업데이트 패킷 생성)을 메인 dLLM 포워드 패스와 중첩시켜, 드리프트 계산의 오버헤드를 효과적으로 숨긴다.

3. 주요 기여

1. 통합 신호: 토큰 표현 드리프트를 단순한 캐싱 에러가 아닌, KV-캐시 노후화와 병렬 확정을 동시에 해결하기 위한 원칙적인 신호로 식별하였다.
2. Polestar-Cache: 레이어별 드리프트 감지와 클러스터 기반 프록시를 사용하여 최소한의 오버헤드로 캐시 유효성을 유지하는 새로운 희소 리프레시 스킴을 제안하였다.
3. Polestar-Commit: 급격한 드리프트 이벤트를 감지하여 신뢰할 수 있는 조기 토큰 확정을 가능하게 함으로써, 정적 신뢰도 임계값을 낮출 때 발생하는 정확도 붕괴 없이 높은 디코딩 병렬성을 구현하였다.
4. 시스템 공동 설계(System Co-Design): 드리프트 모니터링의 런타임 비용을 최소화하기 위해 CPU 오프로딩과 지연 시간 은닉을 활용한 최적화된 실행 흐름을 사용하였다.

4. 실험 결과

저자들은 수학(GSM8K, MATH), 코딩(HumanEval, MBPP), 병렬화 스트레스 테스트(ParallelBench) 벤치마크에서 여러 dLLM 제품군(LLaDA-8B, LLaDA-1.5, Dream-7B)을 대상으로 Polestar를 평가하였다.

• 정확도-처리량 파레토 프런티어(Pareto Frontier): Polestar는 기존 베이스라인 대비 최대 10.73%의 정확도 향상을 달면서 동시에 처리량을 높여 새로운 SOTA를 달성했다.
• 처리량 이득: 베이스라인 모델 대비 최대 **3.7배 높은 처리량(TPS)**을 달성했으며, 특정 구성(예: GSM8K-512)에서는 베이스라인 대비 최대 20.5배까지 향상되었다.
• 디코딩 병렬성: Polestar는 **포워드 패스당 3.67개 토큰(TPF)**이라는 높은 디코딩 병렬성을 달려, 정확도를 유지하거나 개선하면서도 Fast-dLLM(2.72 TPF) 및 Elastic-Cache(2.93 TPF)와 같은 베이스라인을 크게 앞질렀다.
• 절제 연구(Ablation Studies): 실험을 통해 KL-발산 기반의 드리프트 측정이 코사인 유사성이나 불확실성 기반 선택보다 우수함을 확인했다. 또한, 드리프트 인지형 확정과 접미사 블록 확정의 결 조합이 정확도 저하 없이 높은 TPF를 달축하는 데 필수적임을 입증했다.

5. 의의

본 논문은 Polestar가 dLLM의 효율성 패러다임을 근본적으로 전환한다고 주장한다. 토큰 표현이 양방향 어텐션으로 인해 본질적으로 드리프트한다는 점을 인식함으로써, 이 프레임워크는 이 "드리프트"를 부채가 아닌 제어 신호로 변환한다. 이를 통해 다음이 가능하다:

• 효율적인 캐시 재사용: 이전 방식에서 나타난 정확도 페널티 없이 dLLM에서 KV 캐시를 재사용할 수 있게 한다.
• 신뢰할 수 있는 병렬성: 디코딩 병렬성을 경직된 정적 신뢰도 임계값으로부터 분리하여, 모델이 실제 표현의 수렴에 기반해 토큰을 확정할 수 있도록 한다.

결과적으로 Polestar는 학습이 필요 없는 추론 최적화가 정확도와 처리량을 동시에 극대화할 수 있음을 보여주며, 디퓨전 기반 언어 모델의 실질적인 배포를 위한 새로운 기준을 세운다.