🤖 machine learning

DARTree: Speculative Diffusion Decoding with Autoregressive Draft Trees

DARTree는 자기회귀적 교정을 선형 체인에서 고정 너비 후보 트리로 확장하여, 토큰 수용을 극대화하기 위해 AR 헤드 추론을 순차적 연산으로부터 분리함으로써 상태 최고 수준의 무손실 속도 향상을 달성하는 훈련이 필요 없는 투기적 디코딩 방법이다.

원저자: Tianyi Li, Yaxin Luo, Xinyi Shang, Zhiqiang Shen

게시일 2026-08-14

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Tianyi Li, Yaxin Luo, Xinyi Shang, Zhiqiang Shen

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신은 매우 똑똑하지만 믿을 수 없을 정도로 느린 로봇 친구와 함께 이야기를 쓰려고 한다고 상상해 보세요. 이 로봇은 세상을 이해하고 멋진 문장을 쓸 수 있는 천재이지만, 엄격한 규칙이 하나 있습니다. 바로 한 번에 단 한 단어씩만 써야 한다는 것입니다. 다음 단어를 쓰기 전, 로봇은 지금까지 쓴 모든 내용을 다시 생각하고, 그 후 가장 적절한 단 하나를 신중하게 골라야 합니다. 이는 마치 요리사가 수프에 넣을 재료를 한 번에 하나씩 맛본 후에 다음에 무엇을 넣을지 결정하는 것과 같습니다. 이 방식은 수프의 맛을 완벽하게 만들어 주지만, 큰 요리를 만드는 데는 엄청난 시간이 걸립니다. 인공지능의 세계에서 이 '한 번에 한 단어씩' 생성하는 과정은 **자기회귀 생성(autoregressive generation)**이라고 불리며, 이것이 강력한 AI 챗봇들이 때때로 느리게 느껴지는 주요 원인입니다.

품질을 해치지 않으면서 속도를 높이기 위해, 과학자들은 **추측 디코딩(Speculative Decoding)**이라는 기술을 발명했습니다. 이것은 느린 로봇을 대신해 다음 몇 단어를 추측하는 빠르고 에너지가 넘치는 인턴을 고용하는 것과 같습니다. 인턴은 문장 전체를 외치고, 느린 로봇은 인턴이 맞았는지 빠르게 확인합니다. 만약 인턴이 올바르게 추측했다면, 로봇은 문장 전체를 즉시 받아들이고 다음 단계로 넘어갑니다. 만약 인턴이 실수를 했다면, 로봇은 그 단어 하나만 수정하고 다시 시작합니다. 마법은 인턴이 여러 단어를 연속으로 올바르게 추측할 때 일어납니다. 이때 느린 로봇은 어려운 생각 과정을 건너뛰고 텍스트 한 덩어리에 대해 그저 "그래, 맞아!"라고 말하며 넘어갈 수 있습니다.

최근 연구자들은 더 빠른 인턴을 만들기 위해 **확산 모델(Diffusion Model)**이라는 다른 종류의 뇌를 사용했습니다. 이 인턴은 단어를 하나씩 추측하는 대신, 마치 화가가 단 한 번의 붓질로 캔버스를 채우듯 다음 문장 전체를 한꺼번에 상상합니다. 이것은 매우 빠르지만 결함이 있습니다. 인턴이 문장 전체를 한꺼번에 추측하기 때문에, 첫 번째 단어가 두 번째 단어에 어떻게 영향을 미치는지, 혹은 두 번째가 세 번째에 어떻게 영향을 미치는지 제대로 알지 못한다는 점입니다. 이는 마치 중간 장면을 보지 않고 영화의 결말을 추측하는 것과 같습니다. 이를 해결하기 위해 다른 연구자들이 "교정" 단계를 추가했지만, 그 방식은 여전히 느리고 투박하여 로봇이 인턴의 작업을 단어별로 확인하게 만들었고, 결국 속도를 높이려던 목적을 무색하게 만들었습니다.

여기서 MBZUAI의 VILA 랩에서 나온 새로운 논문인 DARTree가 등장합니다. 연구자들은 기존의 인턴 작업 확인 방식이 도서관을 정리할 때 책 한 권을 집어 들고, 선반을 확인하고, 다시 제자리에 둔 뒤, 다음 책을 집어 드는 식의 너무 많은 움직임이 필요한 방식이라는 것을 깨달았습니다. 대신 DARTree는 다양한 가능성의 "나무"를 구축하는 새로운 방법을 제안합니다. 인턴이 단 하나의 단어 경로만을 추측하는 것이 아니라, 다양한 이야기의 가능성을 담은 울창한 나무를 그린다고 상상해 보세요. 그러면 느린 로봇은 이 전체 나무를 한꺼번에 보되, 특별한 방식을 사용합니다. 바로 나무의 "가지"들을 하나씩이 아니라 큰 묶음(배치) 단위로 확인하는 것입니다.

핵심 혁신은 "추측"과 "확인"을 분리하는 데 있습니다. 먼저, 가능한 많은 이야기 경로를 담은 넓고 일시적인 나무를 한꺼번에 구축합니다. 그다음, 유망해 보이지 않는 가지들을 쳐내는 스마트한 가지치기 도구를 사용하여, 느린 로봇에게 보여줄 가장 좋은 나무만을 남깁니다. 이 과정에서 많은 이야기 경로를 배치 단위로 확인함으로써, 모든 것을 지연시켰던 기존의 단계별 확인 과정을 피할 수 있었습니다. 논문에 따르면 이 방법은 큰 성공을 거두었습니다. 수학 문제, 코딩 작업, 채팅 대화 등 다양한 테스트에서 DARTree는 한 번의 확인 단계당 최대 12.97 토큰(단어 또는 단어의 일부)을 수용할 수 있었습니다. 이는 이전 방식들에 비해 엄청난 도약이며, 최고의 경쟁자 중 하나인 DFlash보다 98.6% 더 뛰어났고, Domino보다 27.9% 더 우수했습니다.

그 결과, 이 시스템은 놀라울 정도로 빠르면서도 완벽하게 정확합니다. 연구자들은 이 새로운 방식이 표준적인 작성 방식보다 9.73배 더 빠르면서도 품질을 떨어뜨리거나 가짜 정보를 만들어내지 않는다는 것을 측정했습니다. 이들은 다양한 유형의 AI 모델로 테스트를 진행했으며, AI가 매우 엄격하고 논리적일 때(수학처럼)나 창의적이고 자유로울 때(채팅처럼) 모두 잘 작동한다는 것을 발견했습니다. 이 논문은 여러 경로를 병렬로 확인한 뒤 최종적인 절단을 수행하는 이 "나무" 접근 방식이 이러한 똑똑한 로봇들의 속도를 높이는 최선의 방법이라고 주장합니다. 이는 구조만 제대로 갖춰져 있다면 속도와 지능 사이에서 하나를 포기할 필요 없이 둘 다 가질 수 있음을 증명합니다.

기술 요약: DARTree

문제 정의

현대의 자기회귀(autoregressive, AR) 대규모 언어 모델(LLM)은 본질적으로 순차적이며, 한 단계에 하나의 토큰만을 생성하는 동시에 거의 모든 모델 파라미터에 접근해야 합니다. 이는 I/O 바운드 병목 현상을 야기하여 추론 속도를 제한하고 현대 하드웨어의 병렬 연산 능력을 낭비하게 만듭니다.

추측적 디코딩(Speculative decoding)은 경량화된 "드래프터(drafter)"를 사용하여 여러 개의 미래 토큰을 제안하고, 이를 타겟 모델이 병렬로 검증함으로써 이 문제를 해결합니다. 그러나 기존 방식들은 제안 지연 시간(proposal latency)과 인과적 조건화(causal conditioning) 사이에서 트레이드오프를 겪습니다:

자기회귀형 드래프터(Autoregressive Drafters): 제안을 순차적으로 생성합니다. 강력한 인과적 의존성을 유지하지만, 드래프트 길이에 따라 지연 시간이 선형적으로 증가합니다.
확산 기반 블록 드래프터(Diffusion-Based Block Drafters): 단 한 번의 순전파(forward pass)로 전체 토큰 블록을 예측하여 제안 지연 시간을 획기적으로 줄입니다. 그러나 이들은 검증된 접두사(prefix)에만 조건화된 위치별 주변 분포(position-wise marginal distributions)를 예측하며, 동일한 드래프트 블록 내의 실현된 토큰들을 무시합니다. 이는 타겟 모델의 분포와 "인과적 불일치(causal mismatch)"를 생성합니다.
인과적 교정(Causal Correction): 경량화된 교정 헤드(correction heads)는 선택된 토큰들을 단일 드래프트 체인을 따라 전파하여 블록 내 의존성을 복구할 수 있습니다. 하지만 이를 드래프트 트리 검증(타겟 모델에 여러 후보 연속성을 노출함)과 결려할 경우, 새로운 효율성 병목 현상이 발생합니다. 특정 노드의 자식 노드에 대한 교정된 점수가 해당 노드에 도달하기 위해 취해진 특정 경로에 의존하기 때문에, 표준 최선 우선 탐색 트리 구축(max-heap 사용)은 교정 헤드가 순차적으로 작동하도록 강제합니다. 즉, 하나의 노드를 꺼내고, 교정하고, 자식 노드를 푸시한 후에야 다음 노드를 확장할 수 있습니다. 이러한 교정 추론과 힙(heap) 연산의 인터리빙(interleaving)은 배칭(batching)을 방해하여 트리 구축을 심각한 지연의 원인으로 만듭니다.

방법론: DARTree

저자들은 사전 학습된 인과 교정 블록 병렬 드래프터를 단일 체인에서 추측적 트리로 확장하는 학습 불필요(training-free) 추측적 디코딩 방법인 DARTree를 소개합니다. 핵심 혁신은 경로 의존적인 인과 교정을 순차적인 힙 연산으로부터 분리하는 것입니다.

1. 깊이별 병렬 AR 확장 (Depth-Wise Parallel AR Expansion)

DARTree는 트리를 노드 단위로 확장(순차적 교정을 강제함)하는 대신, 병렬 배치 방식으로 **깊이별(depth-by-depth)**로 트리를 확장합니다:

블록 병렬 드래프팅(Block-Parallel Drafting): 확산 드래프터가 먼저 전체 드래프트 블록에 대한 공유 표현(shared representations)을 생성합니다.
후보 선택(Candidate Selection): 각 깊이 $d$ 에서, 시스템은 베이스 로짓(base logits)으로부터 상위- $K$ 후보 토큰을 선택합니다.
배치 교정(Batched Correction): 현재 깊이의 모든 활성 브랜치(active branches)가 동시에 처리됩니다. 사전 학습된 교정 헤드는 각 브랜치의 특정 접두사 상태를 사용하여 모든 후보를 단일 배치로 평가합니다.
스코어링 및 유지(Scoring and Retention): 각 후보 확장은 교정된 누적 확률을 사용하여 점수가 매겨집니다. 시스템은 다음 깊이를 위해 가장 높은 점수를 받은 상위- $W$ 브랜치를 유지하며, 이들의 교정 상태를 전달합니다.
슈퍼트리 구축(Supertree Construction): 이 과정은 모든 드래프트 깊이에 대해 반복되어, 모든 깊이가 규칙적이고 텐서화된 워크로드를 갖는 "고정 폭 후보 슈퍼트리(fixed-width candidate supertree)"를 구축합니다.

2. 지연된 최선 우선 프루닝 (Deferred Best-First Pruning)

전체 슈퍼트리가 구축되고 모든 노드의 점수가 매겨진 후, DARTree는 지연된 프루닝(deferred pruning) 전략을 적용합니다:

전역 선택(Global Selection): 힙 연산을 교정 과정과 인터리빙하는 대신, DARTree는 슈퍼트리에 구현된 모든 접두사에 대해 단일 전역 Top- $B$ 선택을 수행합니다.
이론적 동등성: 저자들은 비양수 깊이 보너스( $\beta \le 0$ )를 적용하여 더 깊은 접두사에 페널티를 부여할 경우, 전역 Top- $B$ 연산에 의해 선택되는 노드 집합이 순차적 최선 우선 max-heap에 의해 선택되는 집합과 동일함을 증명합니다(Lemma 1). 이를 통해 시스템은 순차적 병목 현상을 완전히 우회할 수 있습니다.
검증: 결과로 생성된 압축된 트리는 트리 어텐션(tree attention)을 사용하여 타겟 모델의 단일 순전파로 전달됩니다.

주요 기여

순차적 병목 현상 식별: 본 논문은 경로 조건부 인과 교정과 노드 단위 최선 우선 트리 구축을 결합하는 것이 효율적인 배칭을 방 препят하는 지연 병목 현상을 생성함을 식별했습니다.
학습 불필요 깊이별 배치 구축: DARTree는 교정 헤드의 추가 학습 없이도 각 깊이에서 여러 브랜치에 걸쳐 인과 교정을 적용하는 방법을 도입하여, 손실 없는 속도 향상을 달aff합니다.
지연된 프루닝 전략: 저자들은 순차적인 힙 연산을 슈퍼트리가 완전히 스코어링된 후의 단일 전역 Top- $B$ 선택으로 대체하는 단순화된 최선 우선 프루닝 전략을 제안하며, 특정 스코어링 조건 하에서 최적의 트리 구조를 유지합니다.

실험 결과

저자들은 Qwen3-4B 및 Qwen3-8B를 대상으로 7개의 벤치마크(Math: GSM8K, MATH-500, AIME25; Code: HumanEval, MBPP; Chat: MT-Bench, Alpaca)에서 온도 $T=0$ 및 $T=1$ 에 대해 DARTree를 평가했습니다.

성능: DARTree는 네 가지 모델-온도 구성 모두에서 가장 높은 평균 수락 길이( $\tau$ $τ$ )와 속도 향상을 달성했습니다.
- 최적의 설정(Qwen3-4B, $T=0$ , GSM8K)에서 DARTree는 검증 라운드당 12.97개 토큰을 수락하였으며, 이는 DFlash 대비 98.6%, Domino 대비 27.9% 증가한 수치입니다.
- 이는 로컬로 측정된 자기회귀 디코딩 대비 9.73배의 손실 없는 속도 향상을 가져왔습니다.
비교:
- DARTree(pruned 버전)는 이전의 SOTA 확산 기반 방식인 Domino보다 수락 길이 측면에서 최대 46.8%, 속도 향상 측면에서 최대 40.1% 더 우수한 성능을 보였습니다.
- 또한, (교정 없이) 주변 예측을 사용하는 DDTree 및 (트리 구축 전 단일 체인에만 교정을 적용하는) Domino-chain + DDTree보다 뛰어난 성능을 보였습니다.
어블레이션(Ablation):
- 지연 시간(Latency): 힙을 사용하는 순차적 교정 방식은 라운드당 약 70ms가 소요된 반면, DARTree는 약 28ms를 유지하여 깊이별 배치 방식의 효율성을 입증했습니다.
- 일반화 가능성: DARTree는 DSpark 드래프터(마르코프 교정 헤드 사용)에 적용했을 때도 성능을 개선하는 데 성공하여, 인과적 교정이 적용된 병렬 드래프터에 대한 범용적인 방법임을 확인했습니다.
- 하이퍼파라미터: 레이어 너비( $W$ )를 늘리면 수락 길이가 일정 수준까지 향상되었으나( $W=12$ 이후 점진적 수익 감소), 검증 예산( $B$ )은 $B=64-128$ 부근에서 정점의 속도 향상을 보였습니다.

의의

본 논문은 추측적 디코딩에서 인과적 충실도(경로 특화 교정 필요)와 계산 효율성(병렬 배칭 필요) 사이의 근본적인 긴장 관계를 DARTree가 해결한다고 주장합니다. 트리를 노드 단위가 아닌 깊이 단위로 재구성함으로써, DARTree는 전통적인 힙 기반 탐색의 순차적 지연 페널티 없이도 고성능의 인과 교정 확산 드래프터를 사용할 수 있게 합니다. 이 방법은 기존의 사전 학습된 교정 헤드에 적용할 수 있는 학습 불료(training-free) 강화 기법으로 제시되며, 다양한 도메인과 디코딩 온도에서 상당하고 일관된 추론 속도 향상을 제공합니다.