Draft Less, Retrieve More: Speculative Decoding을 위한 하이브리드 트리 구성

동적 깊이 prune은 드래프트 비용을 낮추지만 MAT를 떨어뜨려 Pareto frontier를 굳건히 만든다. Graft는 pruning으로 해방된 예산을 retrieval으로 보충해 동일 예산에서 후보 범위를 확장하고, lossless한 검증 경로를 유지하며 짧은 컨텍스트와 긴 컨텍스트 양쪽에서 실용적 속도 향상을 달성한다. 대형 모델에서의 평균 속도 up to 5.41×를 기록하고, Qwen3-235B에서 EAGLE-3 대비 평균 속도 향상을 최대 21.8% 달성했다. 또한 LLaMA3.1-8B의 장-context에서 3.22×의 평균 속도를 달성하고 Qwen3-14B에서 EAGLE3-64K 대비 16.6%를 넘었다. DFlash와의 확장 가능성에 대한 초기 탐색도 제시한다.

단계별 요지 1) 시작점은 SD에서의 단일 혹은 다수의 드래프트 토큰 체인인 트리 구조이다. Dense 트리에 비해 Dynamic-depth pruning은 계산을 줄이지만, 구조적으로 Dense Tree의 후보 경로 전체를 활용하지 못하게 되고 MAT가 상한에 묶인다. 2) 해결 원리는 prune으로 남겨진 슬롯을 retrieval으로 채워 budget를 재분배하는 것이다. GPU 상의 adjacency matrix를 통해 root 토큰에서 retrieval이 병렬로 시작되며, pruning checkpoint에서 released budget에 맞춰 stage-specific retrieval 템플릿이 호출된다. 3) 결과적으로 pruning으로 얻은 예산은 retrieval에 의해 보완되고, 최종 검증 경로의 예산은 변경 없이 유지된다. 이로써 단일 fixed-budget에서 더 넓은 후보 공간과 높은 MAT를 모두 달성할 수 있다.

첫째, Dynamic-depth pruning으로 budget를 해방한다. 깊이 d에서 c_d를 계산해 prune 여부를 결정하고, 각 pruning 스테이지 s마다 예산 분할을 결정한다(Kmax = K_draft_s + K_ret_s). 둘째, GPU-resident adjacency matrix M를 이용해 stage-specific retrieval 템플릿 G_ret_s를 root에서부터 매칭된 prefix로 채운다. 이때 retrieval은 루트 주변의 후보를 선점해 전체 예산을 확장하며, g_s는 prune 여부를 gating한다. 셋째, retained draft nodes와 retrieval 노드를 하나의 hybrid 트리로 합친 뒤, 표준 target 모델의 verification 경로에서 검증한다. 네째, 온라인 업데이트를 통해 verified tree의 logits를 이용해 M를 지속적으로 갱신한다. 다섯째, long-context에서 프롬프트 evidence를 확장하고 YaRN-adapted EAGLE-64K를 도입해 retrieval의 초기 신뢰도를 높인다.

주요 벤치마크에서 Graft는 Short-context에서 1.83×~5.41×의 속도향상을 기록했다. Vicuna-13B에서 평균 속도 향상은 5.41×에 달했고, large-scale Qwen3-235B에서 EAGLE-3 대비 평균 향상은 최대 21.8%였다. LLaMA3-1B-8B의 long-context에서 평균 속도향상은 3.22×이며, Qwen3-14B에서 EAGLE3-64K 대비 16.6%를 넘었다. Long-context에서의 평균 속도향상은 3.22×로, Qwen3-235B에서도 2.09×의 평균 속도향상을 달성했다. Ablation 연구에서 w/o retrieval은 LLaMA3-8B에서 평균 3.74×, Qwen3-8B에서 2.20×로 하락했고, w/o prune 시에는 3.62×, 1.95×로 추가 하락했다. Warm-up five 라운드 도입 시 HumanEval/GSM8K/MT-Bench에서 각각 2.46×/2.86×/2.46×의 속도향상을 보였다. DFLASH(16) 대비 Graft-DFLASH은 Qwen3-8B에서 평균 3.71×로 상승했고, DFLASH 대비 9.1%의 상대 향상을 보였다. 고병렬(batch) 서빙에서도 Graft는 EAGLE3 대비 throughput 및 MAT를 개선했다(예: HumanEval에서 319.54 tokens/s vs 303.97 tokens/s; MAT 3.22 vs 3.10).

Graft의 아키텍처는 (i) budget-releasing dynamic-depth pruning, (ii) GPU-resident retrieval adjacency matrix, (iii) stage-adaptive retrieval templates, (iv) hybrid verification 및 온라인 업데이트의 네 축으로 구성된다. Pruning은 depth별 confidence c_d를 기반으로 𝜏_d 임계치를 사용해 d0, d1, d5에서 후보를 제거한다. 예산 분해는 Kmax = K_draft_s + K_ret_s로 표기되며, s가 작아질수록 더 많은 budget이 retrieval에 할당된다. Retrieval 템플릿은 base-tree Envelope와 호환되도록 stage별 prefix를 미리 준비하고, pruning stage에 맞춘 prefix를 M에서 행렬 조회로 채운다. Retreival과 retained draft를 합쳐 T_s = T_draft_s ∪ G_ret_s로 구성하고, |T_s| = Kmax를 유지한다. 검증은 기존의 tree-attention 경로를 따라 하나의 target forward pass에서 이루어지며, retrieval 노드 역시 proposal로 취급되며 최종 분포는 lossless하게 보장된다. online 업데이트는 verified tree에서 얻은 logits를 이용해 M를 지속적으로 업데이트하며, long-context에서는 YaRN-adapted EAGLE-64K를 도입해 retrieval의 초기 신뢰성을 높인다.