Cost-Aware Diffusion Draft Trees for Speculative Decoding

기술 요약: 추측적 디코딩을 위한 비용 인식 확산 드래프트 트리 (Cost-Aware Diffusion Draft Trees for Speculative Decoding)

문제 정의
추측적 디코딩(Speculative decoding)은 경량화된 드래프터(drafter)가 후보 토큰들을 제안하고, 타겟 모델이 이를 병렬로 검증함으로써 자기회귀적 언어 모델 추론을 가속화한다. 트리 기반 검증은 여러 후보 연속형(continuations)을 동시에 검증할 수 있게 해주지만, 기존의 드래프트 트리 구축 방식들은 최적화 목표 설정에서 근본적인 불일치 문제를 겪고 있다.

DDTree와 같은 기존 방식들은 고정된 노드 예산(node budget) 하에서 **기대 수락 길이(expected acceptance length)**를 최대화하도록 드래프트 트리를 구성한다. 그러나 저자들은 수락 길이가 예산에 대해 비감소(non-decreasing)한다는 점을 관찰했다. 즉, 더 큰 트리가 항상 더 긴 수락 길이를 제공하므로, 최적의 예산 크기를 선택할 원칙적인 근거가 부족하다. 또한 최적의 예산은 정적이지 않다. 드래프터의 확신도와 현재 컨텍스트 길이에 따라 매 디코딩 라운드마다 크게 달라진다. 기존 방식들은 전체 노드 예산을 오프라인에서 고정된 값으로 선택하는 데 의존하며, 이는 커버리지(coverage)와 검증 비용 사이의 트레이드오프에 대한 라운드별 가변성에 대응하지 못한다.

방법론: CaDDTree
본 논문은 수락 길이가 아닌 **토큰 처리량(token throughput, 단위 시간당 생성되는 기대 토큰 수)**을 직접 최적화하는 프레임워크인 CaDDTree(Cost-aware Diffusion Draft Tree)를 소개한다. 이 방법은 단 한 번의 순전파(forward pass)로 각 위치별 주변 분포(per-position marginal distributions)를 생성하는 블록 확산 드래프터(예: DFlash)를 기반으로 작동한다.

1. 처리량 목표: 저자들은 라운드 비용을 $C(T; \ell) = C_d + C_v(|T|; \ell)$로 모델링한다. 여기서 $C_d$는 드래프팅 비용이며, $C_v$는 트리 크기 $|T|$와 컨텍스트 길이 $\ell$에 의존하는 검증 비용이다. 목표는 기대 수락 토큰 수($1 + \Phi(T)$)를 총 비용으로 나눈 비율을 최대화하는 것이다.
2. 분해: 저자들은 모든 유효한 트리에 대해 처리량을 최대화하는 것이 노드 예산 $n$에 대한 일차원 탐색으로 귀결됨을 증명한다. 임의의 고정된 $n$에 대해, 최적의 트리 구조는 "그리디 베스트 퍼스트(greedy best-first)" 트리(가장 높은 접두사 확률을 가진 노드를 선택)이며, 이는 이미 DDTree에서 사용되는 구성 방식이다.
3. 단봉성(Unimodality) 및 그리디 중단: 본 논문의 핵심적인 이론적 기여는 검증 비용 함수에 완만한 볼록성(convexity) 조건이 있을 때, 대리 처리량 목표(surrogate throughput objective)가 트리 크기에 대해 **단봉성(unimodal)**을 띤다는 것을 증명한 것이다. 이 성질 덕분에 효율적인 그리디 알고리즘이 가능하다. 즉, 알고리즘은 트리에 노드를 점진적으로 추가하다가 한계 이득(marginal gain)이 양수가 되지 않는 즉시 중단한다.
4. 적응형 예산: 기존 연구와 달리, CaDDTree는 오프라인 그리드 탐색을 통해 고정된 예산을 요구하지 않는다. 대신, 현재의 위치별 분포와 추정된 검증 비용(컨텍스트 길이에 따라 변함)을 바탕으로 매 디코딩 라운드마다 예산 $n^*$을 적응적으로 선택한다.
5. 비용 프로파일링: 검증 비용 함수 $C_v(n; \ell)$은 다양한 컨텍스트 길이와 노드 수에 걸쳐 측정된 지연 시간을 이용한 볼록 회귀(convex regression)를 통해 오프라인에서 한 번만 추정된다. 이를 통해 비용 모델이 볼록하고 단조성을 유지하도록 하여 단봉성 보장을 충족시킨다.

주요 기여

• 목표의 전환: 수락 길이 최대화라는 잘못 정렬된 기준 대신, 직접적인 처리량 최적화를 드래프트 트리 구성의 주요 기준으로 제안한다.
• 이론적 보장: 처리량 목표가 볼록한 검증 비용 하에서 단봉성을 가짐을 증명하여, 추가 학습 없이도 증명 가능한 최적의 그리디 중단 규칙을 가능하게 한다.
• 적응형 메커니즘: 매 라운드마다 최적의 트리 크기를 동적으로 선택하는 방법을 도입하여, 수동 또는 오프라인 예산 튜닝의 필요성을 제거한다.
• 구현: 블록 확산 드래프터를 사용하여 프레임워크를 구현하였으며, 이들의 단일 패스 주변 분포를 활용하여 트리를 효율적으로 구축한다.

실험 결과
본 방법은 추론(MATH-500, GSM8K, AIME 2025), 코딩(HumanEval, MBPP, LiveCodeBench), 지시 이행(MT-Bench, Alpaca)을 아우르는 8개 벤치마크에서 Qwen3-4B 및 Qwen3-8B를 대상으로 평가되었다.

• 성능: CaDDTree는 거의 모든 작업에서 오라클 예산 선택(grid search를 통해 발견된 최적의 고정 예산)을 적용한 DDTree와 대등하거나 이를 능가하는 성능을 보인다.
• 효율성: 자가회귀(autoregressive) 베이스라인 및 고정 예산 DDTree와 비교했을 때 더 낮은 토큰당 생성 시간(더 높은 처리량)을 달고한다. 예를 들어, Qwen3-8B 기반의 MATH-500 작업에서 CaDDTree는 4.60 ms의 토큰당 시간을 기록하여, 자가회귀 베이스라인의 5.83 ms보다 빠른 성능을 보였다.
• 적응성: 분석 결과, 최적의 예산은 라운드마다 크게 변하는 것으로 나타났다. MATH-500과 같은 추론 작업에서는 드래프터의 확신도에 따라 선택된 예산이 고정된 오라클 예산과 크게 차이 나며 넓게 변동(표준 편차 $\approx 42$)하였다. 지시 이행 작업에서는 예산이 더 안정적이었으나 여전히 고정된 그리드 지점과는 구별되었다.
• 강건성: 절제 연구(Ablation studies)에 따르면, 본 방법은 비용 파라미터 추정치의 섭동(perturbation)에 강건하며, 이는 정밀한 오프라인 캘리브레이션이 엄격하게 필수적이지 않음을 시사한다.

의의
본 논문은 CaDDTree가 트리 커버리지와 검증 지연 시간 사이의 긴장을 해결하는, 원칙적이고 비용 인식적인 추측적 디코딩 접근법을 제공한다고 주장한다. 처리량 목표의 단봉성을 증명함으로써, 저자들은 추가 학습 없이도 각 디코딩 단계의 난이도에 자동으로 적응하는 효율적인 알고리즘을 가능하게 했다. 이는 휴리스틱한 예산 선택의 필요성을 제거하며, 동적이고 비용 인식적인 트리 구성이 다양한 작업과 모델 크기에 걸쳐 정적이고 오라클로 튜닝된 베이스라인을 일관되게 능가할 수 있음을 입증한다.