DARTree: 오토리그래시브 초안 트리를 활용한 추론 기반 디코딩

DARTree: Speculative Diffusion Decoding with Autoregressive Draft Trees

추론 기반 디코딩은 여러 개의 초안 토큰을 병렬로 검증하여 오토리그래시브 언어 모델의 속도를 향상시키는 기술입니다. 확산 모델 기반의 초안 생성기는 전체 토큰 블록을 예측하며 지연 시간을 더욱 줄이지만, 기존 방식과 달리 각 초안 경로에서 선택된 토큰에 대한 조건부 분포를 활용하지 않고 위치별 분포만 사용합니다. 기존의 순환 수정(recurrent correction)은 단일 초안 체인 내의 인과적 정보를 활용하는 반면, 확산 모델 기반의 트리 구조는 후보 범위를 넓히지만 이러한 수정 정보를 개별 분기 경로에 반영하지 않습니다. 본 연구에서는 사전 훈련된 AR 수정 모듈을 체인에서 트리 구조로 확장하는, 학습이 필요 없는 추론 기반 디코딩 방법인 DARTree를 소개합니다. DARTree는 먼저 각 깊이에서 모든 노드를 확장하고 평가하여 고정 폭의 후보 트리를 구성한 다음, 최적-우선 탐색(best-first pruning)을 사용하여 검증할 트리를 선택함으로써 AR 수정 모듈의 추론 연산을 순차적인 힙 연산과 분리합니다. 7개의 수학, 코드 및 채팅 벤치마크에서 DARTree는 모든 모델-온도 구성에서 가장 높은 평균 채택 길이와 속도 향상을 달성했으며, 동일한 설정에서 DFlash보다 최대 98.6%, Domino보다 27.9% 더 많은 토큰(최대 12.97개)을 검증하고, 로컬 측정된 오토리그래시브 디코딩에 비해 최대 9.73배의 손실 없는 속도 향상을 보였습니다.

Speculative decoding losslessly accelerates autoregressive language models by verifying multiple draft tokens in parallel. Diffusion-based drafters further reduce proposal latency by predicting an entire token block in parallel, but their position-wise distributions are marginal rather than conditioned on tokens selected along each draft path. Existing recurrent correction incorporates causal information along a single draft chain, whereas diffusion-based tree construction broadens candidate coverage without carrying this correction along individual branches. We introduce DARTree, a training-free speculative decoding method that extends a pretrained AR correction head from chains to trees. DARTree first constructs a fixed-width candidate tree by expanding and scoring all nodes at each depth in a single batch, and then only applies best-first pruning to select the verification tree, decoupling AR-head inference from sequential heap operations. Across seven math, code, and chat benchmarks, DARTree achieves the highest average acceptance length and speedup in all four model--temperature configurations, accepting up to 12.97 tokens per verification round, 98.6\% more than DFlash and 27.9\% more than Domino in the same setting, and reaching up to 9.73$\times$ lossless speedup over locally measured autoregressive decoding.