기술 요약: xPress: 확산형 드래프터(Diffusion Drafters)를 위한 병렬 정교화(Parallel Refinement)

1. 문제 정의

추측적 디코딩(Speculative Decoding, SD)은 경량 드래프터가 제안한 토큰을 더 큰 타겟 모델이 검증하게 함으로써 자기회귀적(Autoregressive, AR) 생성을 가속화합니다. EAGLE과 같은 AR 드래프터는 효과적이지만, 긴 블록을 생성할 때 순차적인 오버헤드 문제를 겪습니다. 블록-확산 드래프터(Block-diffusion drafters, 예: dFlash)는 단 한 번의 병렬 순방향 패스(forward pass)로 전체 토큰 블록을 생성함으로써 드래프팅 레이턴시를 획기적으로 줄여 이 문제를 해결합니다.

그러나 확산 드래프터는 근본적인 정확도 한계를 가지고 있습니다: 이들은 결합 분포(joint distributions)가 아닌 주변 분포(marginal distributions)를 기반으로 토큰을 예측합니다. 단일 패스 확산 과정에서 토큰 $k$는 토큰 $1$부터$k-1$까지의 실현된 값들에 조건화되지 않은 채 샘플링됩니다. 이러한 블록 내 인과관계(intra-block causality)의 결여는 개별 토큰은 국소적으로 확률이 높지만, 결합되었을 때는 부조화스러운(예: 주어-동사 불일치) 시퀀스를 초래합니다. 검증 과정에서 타겟 모델은 이러한 부조화스러운 시퀀스를 조기에 거부하며, 이는 수용 길이($\tau$)를 심각하게 제한하고 병렬 드래프팅의 잠재적인 속도 향상을 상쇄합니다.

PRESTO나 DDTree와 같은 트리 기반 드래프팅 방식은 후보 공간을 트리 형태로 확장하여 이를 완화하려 시도합니다. 그러나 이러한 방법들은 계산 비용이 많이 들고 서빙 구조가 복잡하며, 근본적인 인과적 부조화를 해결하기보다는 단순히 재현율(recall)을 목표로 삼아 문제를 우회하는 경향이 있습니다.

2. 방법론: xPress

저자들은 느린 토큰 단위의 자기회귀 루프로 되돌아가지 않으면서도, 확산 드래프터의 누락된 인과적 의존성을 복구하도록 설계된 경량 인과 정교화 도구인 xPress를 제안합니다.

2.1 핵심 개념

xPress는 확산 드래프터의 초기 블록 제안을 받아 병렬 정교화(parallel refinement)를 통해 이를 조정하는 "인과적 정교화 도구(causal refiner)"로 작동합니다. xPress는 토큰을 순차적으로 다시 생성하는 대신, 확산 드래프터의 로짓(logits)에 인과적 정보를 주입하여 주변 분포를 교정합니다.

2.2 아키텍처

정교화 도구는 각 위치 $k$에서 드래프터의 기본 로짓 $s_k$에 학습된 로짓 바이어스(logit bias, $\delta_k$)를 더하는 경량 모듈입니다. 아키텍처는 다음 네 단계로 구성됩니다:

1. Fuse (융합): 세 가지 입력을 저차원 공간($r=256$)으로 투영합니다:
   • 드래프터의 위치별 은닉 상태($h_k$).
   • 평균 풀링(mean-pooling)을 통해 얻은 블록 전역 요약값($g$).
   • 이전 토큰 추측값($x_{k-1}$)의 임베딩.
2. Mix (혼합): 블록 위치 전체에 걸쳐 엄격한 하삼각 선형 믹서(strictly lower-triangular linear mixer, $L$)를 적용합니다. 이를 통해 위치 $k$는 단 한 번의 연산으로 모든 선행 위치($j \le k$)로부터 정보를 집계할 수 있으며, 순차적 계산 없이 인과 구조를 강제합니다.
3. Trans (변환): 저차원 공간에서의 잔차 MLP(residual MLP)입니다.
4. Read (읽기): 결과를 어휘 사전 크기로 다시 투영하여 바이어스 $\delta_k$를 생성합니다.

최종 정교화된 로짓은 $\ell_k = s_k + \delta_k$입니다.

2.3 Jacobi 디코딩을 통한 병렬 정교화

직관적인 인과 정교화 도구는 순차적인 좌-우(left-to-right) 패스를 요구할 것이며, 이는 xPress가 피하고자 하는 레이턴시를 다시 도입하게 됩니다. 이를 해결하기 위해 xPress는 Jacobi 디코딩을 사용합니다:

• 초기화: 블록은 드래프터의 초기 그리디(greedy) 예측값으로 시딩(seeding)됩니다.
• 반복: 모든 위치는 현재 블록의 상태를 다음 반복을 위한 접두사(prefix)로 사용하여 동시에 병렬적으로 업데이트됩니다.
• 수렴: 이 과정은 $K$번의 반복을 거칩니다. 이론적으로 블록 크기가 $B$일 때, 순차적 고정점(fixed point)은 최대 $B$번의 반복 내에 도달합니다. 경험적으로 저자들은 드래프터의 시드가 이미 고품질의 추측값이기 때문에 $K \approx 6$번의 반복만으로도 전체 순차 디코딩의 수용 길이에 도달할 수 있음을 발견했습니다.

2.4 학습 전략

xPress는 합성 손실 함수(composite loss function)를 사용하여 확산 드래프터와 공동 학습됩니다:

1. 교차 엔트로피 (Teacher Forcing): 정답 접두사가 주어졌을 때 정답 토큰이 나타날 확률을 극대화합니다.
2. 전변동 (Total Variation, TV) 손실: 정교화된 분포와 타겟 모델의 분포 사이의 거리를 최소화하여 수용률(acceptance rate)을 직접 최 optimization 합니다.
3. 일관성 (Consistency) 손실: 드래프터 자신의 예측값을 입력으로 사용하는 두 번째 순방향 패스에 적용되는 보조 손실입니다. 이는 학습 분포를 추론 시의 Jacobi 과정과 일치시켜, 정교화 도구가 골드(gold) 접두사에 과적합되는 것을 방지합니다.
4. 드래프터 앵커 (Drafter Anchor): 기본 드래프터의 출력에 적용되는 손실 항으로, 기저의 확산 모델이 높은 품질을 유지하도록 보장합니다.

3. 주요 기여

• 인과적 의존성 주입: xPress는 드래프터의 은닉 상태에 작용하는 최소한의 메커니즘(저차원 투영 및 삼각 믹서)을 사용하여 확산 드래프터에 인과적 정보를 주입하며, 이를 통해 파라미터 증가를 최소화합니다.
• 병렬 정교화: Jacobi 디코딩을 활용하여 인과적 의존성을 병렬로 해결합니다. xPress는 $B$번의 종속적인 자기회귀 단계를 단 $K \ll B$번의 병렬 반복만으로 구현하여, 병렬 드래프팅의 속도 이점을 보존합니다.
• 일관된 엔드-투-엔드 이득: xPress는 검증 과정을 수정하거나 더 큰 드래프터를 요구하지 않고도 수용 길이와 처리량(throughput)을 모두 유의미하게 개선합니다.

4. 실험 결과

실험은 Qwen3-8B를 사용하여 블록-16 dFlash 드래프터를 대상으로 7개의 벤치마크(Math: GSM8K, MATH-500, AIME25; Code: HumanEval, MBPP, LiveCodeBench; Chat: MT-Bench)에서 수행되었습니다.

• 수용 길이 (Acceptance Length): xPress는 기존 dFlash 확산 드래프터에 비해 평균 수용 길이를 약 30% 증가시킵니다.
  • 그리디 디코딩 ($T=0$) 시, 이득은 18%에서 56% 사이입니다.
  • 샘플링 ($T=1$) 시, 이득은 19%에서 58% 사이입니다.
• 처리량 속도 향상 (Throughput Speedup):
  • 그리디 ($T=0$): xPress는 dFlash 드래프터 대비 평균 1.3배의 엔드-투-엔드 속도 향상을 달야하며, GSM8K와 같은 경우 최대 1.7배까지 도달합니다.
  • 샘플링 ($T=1$): 속도 향상은 평균 1.46배로 개선되며, 최대 1.6배에 달합니다.
• 순차적 정교화 도구와의 비교: Markov head(DSpark의 순차적 인과 정교화 도구)와 비교했을 때, xPress는 수용 길이를 대등하게 맞추거나 능가하면서도 훨씬 빠르게 실행됩니다.
  • $K=4$ Jacobi 반복에서, xPress는 Markov head의 정확도와 일치하면서도 훨씬 빠른 속도(1.6배 빠름: 블록당 379 µs vs 601 µs)로 실행됩니다.
  • Markov head는 고정된 15단계 직렬 디코딩이 필요하지만, xPress는 더 적은 병렬 단계 내에 수렴합니다.
• 확장성 (Scalability): 대규모 배치 서빙 시나리오(vLLM, 배치 크기 16 및 32)에서도 xPress는 이점을 유지합니다. 배치가 커짐에 따라 (검증 작업이 연산 집약적이 되므로) AR 베이스라인 대비 절대적인 속도 향상은 줄어들지만, xPress는 지속적으로 일반 dFlash 드래프터보다 1.2–1.4배 높은 처리량을 보여주며 부하 상황에서도 정교화 도구의 효용성을 입증합니다.

5. 의의 및 주장

본 논문은 xPress가 병렬 확산 드래프팅과 순차적 검증 사이의 "내재적 불일치"를 해결하기 위해, 탐색 공간을 확장하는 대신 주변 분포 자체를 수선한다고 주장합니다.

• 효율성: xPress는 드래프팅 오버헤드 대비 거의 제로에 가까운 비용으로 인과적 정보를 주입할 수 있음을 보여줍니다. 정교화 비용은 전체 드래프팅 시간의 극히 일부이며, 병렬적 특성 덕분에 순차적 교정으로 인한 레이턴시 페널티를 피할 수 있습니다.
• 직교성 (Orthogonality): 이 방법은 다른 개선 사항들과 직교합니다. 즉, 더 큰 드래프터나 트리 기반 검증 방법과 결래하여 사용할 수 있으며, 이러한 방법들이 의존하는 "백본" 로짓을 개선합니다.
• 일반화 가능성: 저자들은 상호 의존적인 예측들을 몇 번의 병렬 반복을 통해 조정하는 이 철학이, 모델이 상호 의존적인 출력을 동시에 생성하는 확산 드래프팅 외의 영역에도 적용될 수 있다고 제안합니다.

결론적으로 xPress는 단일 패스 확산 드래프팅에서 상실된 인과성을 성공적으로 복구하여, 병렬 효율성을 희생하지 않으면서도 더 긴 수용 체인을 가능하게 하고 더 높은 처리량을 달 수 있음을 입증했습니다.