플래시 스펙터티브 디코딩을 위한 블록 디퓨전

초록

DFlash는 경량 블록 디퓨전 모델을 초안기로 활용해 대형 LLM의 추론을 가속화한다. 목표 모델의 은닉 특성을 조건으로 주입함으로써 초안 품질을 크게 향상시키고, 한 번의 전방 패스로 토큰 블록을 병렬 생성한다. 실험 결과 Qwen‑3‑8B 기준 6배 이상의 손실 없는 가속과 기존 최고 성능인 EAGLE‑3 대비 2.5배 높은 속도 향상을 달성했다.

상세 분석

본 논문은 LLM 추론의 핵심 병목인 순차적 토큰 생성 문제를 해결하기 위해 ‘스펙터티브 디코딩’이라는 프레임워크에 블록 디퓨전 모델을 결합한 DFlash를 제안한다. 기존 스펙터티브 디코딩(EAGLE‑3 등)은 경량 초안 모델을 사용하지만, 초안 생성 자체가 여전히 자동회귀 방식이므로 토큰 수 γ에 비례해 초안 생성 비용 T_draft가 증가한다. 이는 초안 모델의 용량을 얕게 만들게 하여 초안 품질이 제한되고, 결과적으로 수용 길이 τ가 포화돼 실질적인 속도 향상이 2‑3배 수준에 머문다.

DFlash는 이 한계를 블록 디퓨전 모델이 제공하는 ‘동시 다중 토큰 생성’ 특성을 활용해 극복한다. 블록 디퓨전은 마스크된 토큰 블록을 한 번의 포워드 패스로 복원하므로 T_draft ≈ t_parallel이며, γ가 커져도 초안 생성 지연이 크게 증가하지 않는다. 따라서 더 깊고 표현력이 풍부한 초안 모델을 설계해도 실시간 지연에 큰 영향을 주지 않는다.

핵심 혁신은 ‘타깃 컨텍스트 피처’를 초안 모델에 주입하는 방식이다. 대형 목표 모델(M_t)은 프리‑필 단계에서 최초 토큰을 생성하면서 여러 레이어(浅層~深層)의 은닉 상태를 추출한다. 이 은닉 상태를 층별로 결합·프로젝션해 압축된 컨텍스트 벡터를 만든 뒤, 초안 모델의 KV 캐시에 직접 삽입한다. 이렇게 하면 초안 모델은 토큰 임베딩뿐 아니라 목표 모델이 이미 인코딩한 장기 의존성과 작업‑특화 정보를 지속적으로 활용할 수 있다. 기존 방법이 초안 입력에만 피처를 결합하고 깊은 레이어로 갈수록 희석되는 것과 달리, DFlash는 KV‑Injection을 통해 모든 초안 레이어에 동일한 고품질 컨텍스트를 제공한다.

학습 단계에서도 특수한 설계가 적용된다. 블록 구성은 ‘앵커 토큰 + 마스크 토큰’ 방식으로, 앵커는 목표 모델이 검증 단계에서 제공하는 보너스 토큰과 일치한다. 이는 추론 시 초안이 항상 목표 모델의 최신 토큰을 조건으로 삼는 상황을 정확히 모방한다. 또한, 블록 내 마스크 토큰을 무작위로 선택해 다양한 컨텍스트에 대한 일반화를 촉진한다. 이러한 학습 전략은 수용 길이 τ를 크게 늘리고, 초안 품질을 향상시켜 손실 없는 가속을 가능하게 한다.

실험 결과는 두드러진 성능 향상을 보여준다. Qwen‑3‑8B를 기준으로 16‑token 블록을 5‑layer 초안 모델로 처리했을 때, 초안 생성 지연이 기존 자동회귀 초안(1‑layer)보다 3‑5배 낮았으며, 전체 추론 속도는 6.1×까지 상승했다. 다양한 벤치마크(GSM8K, HumanEval, MBPP 등)에서도 평균 2.5× 이상의 속도 향상을 기록했으며, 품질 지표(정확도, 코딩 성공률)는 목표 모델과 거의 동일한 수준을 유지했다. 메모리 측면에서도 1‑2 GB 정도의 추가 비용만 발생해 실제 서비스 환경(SGLang)에서도 적용 가능함을 입증했다.

요약하면, DFlash는 (1) 블록 디퓨전 기반의 병렬 초안 생성, (2) 목표 모델 은닉 피처의 KV‑Injection을 통한 강력한 조건부 생성, (3) 스펙터티브 디코딩 프레임워크와의 자연스러운 결합이라는 세 축을 통해 기존 자동회귀 초안 방식의 한계를 극복하고, 손실 없는 고속 LLM 추론을 실현한다.