DBLAST: Dependent Block Drafting for Stochastic Speculative Decoding

기술 요약: DBlast – 확률적 투기적 디코딩을 위한 의존적 블록 초안 작성 (Dependent Block Drafting)

문제 정의

투기적 디코딩(Speculative decoding)은 경량화된 드래프터(drafter)가 여러 개의 미래 토큰을 제안하고, 타겟 모델이 이를 병렬로 검증하게 함으로써 대규모 언어 모델(LLM) 추론을 가속화한다. 최근의 블록 기반 드래프터(예: DFlash)는 단 한 번의 패스로 여러 위치를 예측할 수 있지만, 일반적으로 **독립적 블록 샘플링(independent block sampling)**에 의존한다. 이 방식은 드래프트 블록 내의 토큰들이 프리픽스(prefix)가 주어졌을 때 조건부로 독립적이라고 가정한다.

이러한 가정은 탐욕적(greedy) 또는 저엔트로피 디코딩에서는 잘 들어맞지만, 대화, 창의적 글쓰기, 강화 학습 탐색과 같은 개방형 작업에 필수적인 확률적(stochastic, 비탐욕적) 디코딩 환경에서는 취약해진다. 확률적 설정에서는 타겟 분포가 의도적으로 다양하며, 여러 가지 연속적인 흐름이 타당할 수 있다. 본 논문은 근본적인 불일치를 식별한다. 즉, 드래프터는 주변(marginal) 타겟 분포를 맞추기 위해 토큰을 독립적으로 샘플링하지만, 타겟 모델은 이전에 수락된 토큰들에 조건화되어 각 토큰을 순차적으로 검증한다는 점이다. 결과적으로, 각 위치에서의 주변 분포는 정확하더라도 블록 자체가 일관된 조건부 궤적(conditional trajectory)을 형성하지 못할 수 있으며, 이는 타겟 샘플링 분포의 엔트로피가 높아짐에 따라 수락된 드래프트 길이의 상당한 저하를 초래한다.

방법론

저자들은 병렬 예측의 효율성을 희생하지 않으면서 블록 확산 드래프터(block diffusion drafters)에 의존성을 주입하는 프레임워크인 DBlast(Dependent Block Drafting)를 제안한다. 이 방법은 두 가지 핵심 구성 요소로 이루어진다:

1. 잠재 혼합(Latent Mixture)을 통한 의존적 블록 샘플링:
   제안 분포를 위치별로 분해하는 대신($q(x_{1:b}|y) = \prod q_i(x_i|y)$), DBlast는 블록 내의 상관관계를 유도하기 위해 범주형 잠재 변수(categorical latent variable) $z$("브랜치" 또는 "모드")를 도입한다. 결합 분포는 저계수 정형(low-rank canonical polyadic, CP) 혼합으로 모델링된다:
   $$q_{dep}(x_{1:b} | y) = \sum_{z=1}^{K} q(z | y) \prod_{i=1}^{b} q_i(x_i | y, z)$$

   • 아키텍처: 이 방법은 트랜스포머 은닉 상태(hidden states)를 $K$개의 범주별 오프셋(offset)으로 매핑하는 경량 "범주 은닉 확장기(category hidden expander)"와 $q(z|y)$를 예측하는 "범도 우선순위 헤드(category-prior head)"를 추가하여 DFlash 스타일의 드래프터를 수정한다.
   • 추론: 추론 과정 중에는 잠재 범주를 샘플링하고(선택적으로 온도 $Z_T$ 적용), 그에 해당하는 브랜치를 단 한 번의 병렬 패스로 탐욕적으로 디코딩한다. 이는 블록 드래프팅의 단일 패스 효율성을 유지하면서도, 모델이 일관된 대안적 블록 수준의 연속성을 표현할 수 있게 한다.

2. 수락 지향적 학습 목적 함수(Acceptance-Oriented Training Objective):
   표준 학습은 타겟 블록의 음의 로그 가능도(NLL)를 최소화하는데, 이는 제안의 가능성(likelihood)을 최적화할 뿐 순차적인 검증 과정을 무시한다. DBlast는 **기대 수락 길이(Expected Accepted Length, AL)**에 기반한 대리 손실(surrogate loss)을 도입한다.

   • 이 목적 함수는 접두사 수락 이벤트들로 분해된, 수락된 프리픽스 길이의 기댓값으로부터 유도된다.
   • 드래프터와 타겟 분포 사이의 중첩이 적을 때(학습 초기) 발생하는 중요도 가중치(importance weighting)의 불안정성을 처리하기 위해, 저자들은 **임계값 절단 로그 도메인 대리 손실(threshold-truncated log-domain surrogate)**을 제안한다. 이 손실은 중요도 비율이 임계값 $\tau$를 초과하는 유지된 프리픽스들에 대해 로그 확률비와 조건부 수락 길이의 로그를 합산한다.
   • 이 접근 방식은 드래프터의 종속성 모델링이 블록 가능도만으로 충분하지 않을 때를 보완하며, 검증기의 순차적 수락 동작을 직접적으로 겨냥한다.

주요 기여

1. 불일치의 진단: 저자들은 독립적 블록 샘플링이 타겟 디코딩이 덜 결정론적일수록 실패한다는 것을 경험적으로 입증한다. 이들은 독립적 제안이 순차적 검증에 필요한 일관성이 부족하기 때문에, 고엔트로피 환경에서 수락된 길이가 현저히 저하됨을 보여준다.
2. 의존적 블록 샘플링: 비탐욕적 투기적 샘플링을 위해 CP 스타일의 잠재 혼합 제안을 서브워드 블록 확산 드래프터에 처음으로 적용하였다. 이 방법은 병렬 토큰 예측을 유지하면서 블록 위치 간의 상관관계를 유도한다.
3. 수락 지향적 손실: 기대 수락 길리에 의해 동기 부여된 새로운 학습 대리 함수를 제안한다. 이 목적 함수는 임계값 절단을 통해, 드래프터를 순차적 검증 과정과 일치시키는 안정적인 학습 신호를 제공한다.

실험 결과

실험은 Qwen3-4B 및 Q3-8B 타겟 모델을 사용하여 네 가지 벤치마크(GSM8K(수학), MT-Bench(채팅), HumanEval(코딩), 창의적 글쓰기)에서 수행되었다. 평가는 고결정성(Temperature 0.7)부터 저결정성(Temperature 1.5)까지의 세 가지 타겟 샘플링 환경을 다루었다.

• 성능 향상: DBlast는 독립적 블록 샘플링(DFlash) 및 NLL로 학습된 베이스라인보다 일관되게 우수한 성능을 보였다.
  • 고엔트로피 설정(Target Temp 1.5)에서, DBlast는 독립적 베이스라인 대비 Qwen3-8B의 수락 길이에서 평균 12.1%의 이득을 달了.
  • 이득은 저결정성 영역에서 가장 두드러졌으며, 이는 의존성 모델링이 여러 가지 연속 가능성이 존재할 때 가장 가치 있음을 확인시켜 준다.
• 절제 연구(Ablation Studies):
  • 잠재 범주($K$): $K$를 늘리면 성능이 향상되었으나, $K=4$ 이후로는 수확 체감 현상이 관찰되었다.
  • 학습 목적 함수: 수락 지향적(AL) 손실은 독립적($K=1$) 드래프터에서도 NLL 학습보다 일관되게 우수한 성능을 보였으나, 의존성과 AL을 결합했을 때 최상의 결과를 얻었다.
  • 추론 온도: 범주 샘플링 온도($Z_T$)를 조정함으로써 드래프터가 타겟 분포의 확률성을 더 잘 맞추도록 할 수 있었고, 이는 수락 길이를 더욱 향상시켰다.

의의 및 주장

본 논문은 DBlast가 투기적 디코딩의 핵심적인 간극, 즉 현재의 블록 드래프터가 확률적 검증을 효율적으로 처리하지 못하는 문제를 해결한다고 주장한다. 저자들은 블록 내부의 의존성 모델링과 수락 정렬 학습이 상호 보완적임을 입증함으로써, 투기적 디코딩의 매력인 병렬 효율성을 희생하지 않고도 다양한 비탐욕적 디코딩 시나리오에서 높은 수락률을 달성할 수 있음을 보여준다.

저자들은 자신들의 대리 손실이 공식적인 불편 추정량(unbiased estimator)이나 하한(lower bound)은 아니지만, 실용적이고 효과적인 학습 신호 역할을 한다고 겸허히 언급한다. 또한 학습 시 사용되는 소프트 잠재 혼합과 추론 시 사용되는 그리디 브랜치 디코딩 사이의 관계는 공식적인 특성화가 남아 있는 개방된 영역이지만, 실험 결과는 효과적인 전이를 보여준다고 밝혔다. 이 연구는 개방형 생성을 위한 미래의 투기적 디코딩 시스템이 단순한 독립적 토큰 예측을 넘어 미래 시퀀스의 구조적 일관성을 포착해야 함을 시사한다.