PRESTO: Prefix-Aligned Tree Drafting for Diffusion Speculative Decoding
이 논문은 디퓨전 마진널(diffusion marginals)과 자기회귀적 검증(autoregressive verification) 사이의 불일치를 해결하기 위해 접두사 정렬 스코어링(prefix-aligned scoring)과 우선순위 기반 트리 탐색(priority-based tree search)을 구현함으로써 디퓨전 기반 투기적 디코딩을 향상시키는 원칙적인 프레임워크인 PRESTO를 소개하며, 이를 통해 엔드 투 엔드 처리량을 크게 개선한다.
원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기
당신이 이야기의 다음 단어를 예측하려고 노력하고 있다고 상상해 보세요. 오랫동안 가장 똑똑한 컴퓨터들(대규모 언어 모델이라고 불리는)은 책을 소리 내어 읽는 사람처럼 한 단어를 읽을 때마다 멈춰서 다음에 올 것을 생각하며, 한 번에 한 단어씩 이 작업을 수행했습니다. 이는 정확하지만 느립니다. 최근 과학자들은 "확산(diffusion)" 모델을 사용하여 이러한 이야기를 쓰는 새로운 방법을 발견했습니다. 이것은 마치 조각가가 한 번에 아주 작은 조각을 깎아내는 것이 아니라, 대리석 블록에서 한꺼번에 전체 조각상을 깎아내는 것과 같습니다. 이를 통해 컴퓨터는 수많은 단어를 동시에 추측할 수 있으며, 이는 믿을 수 없을 정도로 빠릅니다.
하지만 함정이 있습니다. 여러 단어를 한꺼번에 추측하면 몇 개를 틀릴 수도 있습니다. 이를 해결하기 위해 "추측적 디코딩(speculative decoding)"이라는 영리한 기술이 있습니다. 이것은 마치 빠르고 유능한 주니어 조수가 다음 몇 단어를 추측하고, 그러면 느리지만 매우 똑똑한 보스가 그 추측이 맞는지 확인하는 것과 같습니다. 만약 보스가 동의한다면, 그 단어 묶음 전체를 즉시 수락하여 시간을 엄청나ally 절약합니다. 문제는 주니어 조수(확산 모델)가 개별 단어를 추측하는 데는 뛰어나지만, 그 단어들이 특정 순서로 어떻게 어우러지는지는 항상 알지 못한다는 점입니다. 이는 마치 조수가 케이크를 만들기 위한 개별 재료를 고르는 데는 능숙하지만, 어떤 재료의 조합이 실제로 맛있는 맛을 낼지는 잘 모르는 것과 같습니다.
여기서 새로운 논문이 등장합니다. 정 웽(Zheng Wang)과 동료들이 이끄는 연구진은 현재 이러한 빠른 조수들을 사용하는 방식이 많은 속도를 낭비하고 있다는 사실을 깨달았습니다. 그들은 확산 모델이 방대한 종류의 가능한 단어 조합을 생성할 수 있음에도 불구하고, 현재 방식은 단 하나의 경로만을 확인하고 있다는 것을 발견했습니다. 이는 마치 하나의 복도를 따라 걸어가며 끝에 있는 문이 열려 있기를 바라는 것과 같습니다. 저자들은 PRESTO(Prefix-Aligned Tree Drafting)라고 불리는 새로운 시스템을 제안합니다. 하나의 복도를 걷는 대신, PRESTO는 가능성의 나무(tree)를 구축하여 한 번에 많은 경로를 탐색합니다. 하지만 여기서 마법이 일어납니다. PRESTO는 조수의 자신감을 측정하는 근본적인 결함을 해결합니다. 조수의 원래 자신감은 "접두사 결여(prefix-blind)" 상태로, 즉 앞에 어떤 단어가 왔는지 상관하지 않습니다. PRESTO는 "접두사 정렬(prefix-aligned)" 점수를 추가하여, 탐색을 위해 선택된 경로들이 보스에게 수락될 가능성이 가장 높은 경로가 되도록 나침반 역할을 합니다.
그 결과, 훨씬 더 빠른 시스템이 탄생했습니다. 테스트에서 PRESTO는 컴퓨터가 한 번의 추측 단계에서 더 많은 단어를 수락하도록 도왔습니다. 기존의 가장 우수한 설정 중 일부에서는 전체 과정을 1.5배 더 빠르게 만들었습니다. 다른 설정에서는 1.12배의 속도 향상을 제공했습니다. 이 논문은 추측 과정을 직선적인 경로가 아닌 나무 형태의 가지치기 모험으로 다룸으로써, 확산의 속도와 신중한 검증의 정확성이라는 두 마리 토끼를 모두 잡을 수 있다고 제안합니다.
문제점: "단일 경로"의 함정
왜 PRESTO가 필요한지 이해하기 위해, 당신이 친구와 함께 "매드 립스(Mad Libs)" 게임을 하고 있고 친구가 빈칸에 들어갈 단어를 맞히려고 노력하고 있다고 상상해 보세요. 여기서 친구는 확산 모델입니다. 그들은 빈 공간을 보고 "여기에 '고양이'가 들어갈 것 같아!" 또는 "아니면 '강아지'일까?" 혹은 "혹은 '로켓'일까?"라고 말하는 데 매우 뛰어납니다. 그들은 이 모든 옵션을 동시에 외칠 수 있습니다.
하지만 이 친구를 사용하는 현재 방식은 매우 경직되어 있습니다. 방식은 친구의 최상위 추측을 가져와 적어 놓은 뒤, 그 추서가 맞는지 "보스"(대상 모델)에게 묻습니다. 만약 보스가 "아니오"라고 하면, 그 전체가 버려지고 처음부터 다시 시작해야 합니다. 만약 보스가 "예"라고 하면, 다음 단어로 넘어가서 반복합니다. 이것을 **선형적 드래프팅(linear drafting)**이라고 합니다. 이는 숲을 지나가면서 바로 눈앞에 있는 길만을 바라보는 것과 같습니다.
저자들은 이 접근 방식이 비효und 효율적이라는 것을 관찰했습니다. 확산 모델은 여러 옵션을 동시에 생성하기 때문에, 엄청난 "조합 공간(combinatorial space)"이 존재합니다. 이는 천 개의 서로 다른 산책로가 있는 지도를 가지고 있지만, 오직 하나의 길만 걸을 수 있는 것과 같습니다. 논문은 단 하나의 경로에만 집착함으로써 시스템이 많은 유효한 경로를 놓치고 있음을 보여줍니다. 실제로 GSM8K와 같은 수학 문제에서 현재 방식은 평균적으로 약 6.5개의 단어를 수락했지만, 연구진은 만약 모든 최선의 경로를 확인할 수 있었다면 거의 10개의 단어를 수락할 수 있었을 것이라고 계산했습니다. 이는 엄청난 격차입니다!
불일치: "눈먼" 나침반
연구진은 단순히 더 많은 경로를 확인하는 것(나무를 구축하는 것)이 왜 기존 방식으로는 완벽하게 작동하지 않는지 그 구체적인 이유를 파헤쳤습니다. 그들은 "근본적인 불일치"를 식별했습니다.
표준 AI(자기 회귀 모델)의 세계에서, 단어에 대한 자신감 점수는 그 앞에 어떤 단어가 왔느냐에 크게 의존합니다. 문장이 "고양이가 매트 위에..."라면, 모델은 "매트"가 매우 가능성 높은 다음 단어임을 알지만, "피자"는 그렇지 않다는 것을 압니다. 이것이 **접두사 정렬(prefix-aligned)**입니다.
하지만 확산 모델은 다르게 작동합니다. 그들은 각 위치에 대해 독립적인 "주변(marginal)" 확률을 생성합니다. 이는 모델이 "5번 위치에 '고양이'가 올 확률은 80%야"라고 말하는 것과 같으며, 4번 위치가 "고양이가"였는지 "빠른 갈색"이었는지는 신경 쓰지 않습니다. 이것이 **접로 결여(prefix-blind)**입니다.
이러한 눈먼 점수를 사용하여 추측의 나무를 만들려고 하면 순위 결정 문제가 발생합니다. 첫 번째 단어에는 훌륭해 보이지만 두 번째 단어에는 끔찍해 보이는 경로를 선택할 수도 있는데, 이는 모델이 첫 번째 단어가 문맥을 어떻게 변화시켰는지 깨닫지 못했기 때문입니다. 이는 현재 있는 거리만을 기준으로 방향을 알려주고, 방금 좌회전을 해서 일방통행로에 들어섰다는 사실은 무시하는 GPS와 같습니다. 논문은 이러한 눈먼 점수를 사용하여 나무를 구축하는 것이 "신뢰할 수 없는 경로 순위(unreliable path ranking)"를 초래하며, 즉 시스템이 잘못된 가지를 탐색하고 시간을 낭비하게 만든다고 주장합니다.
해결책: PRESTO
PRESTO(Prefix-Aligned Scoring and priority-based Tree search for diffusion Speculative decOding)는 확산 모델의 점수에 "교정"을 추가함으로써 이 문제를 해결합니다.
- 접두사 정렬 점수 산출 (Prefix-Aligned Scoring): 저자들은 확산 모델의 강력한 "주변(marginal)" 신호(단독으로 단어가 나타날 확률)와 "접두사 조건부(prefix-conditioned)" 신호(이전 단어들이 주어졌을 때의 확률)를 결합해야 한다는 것을 깨달았습니다. 그들은 단순한 n-gram 모델(단어 조합을 살펴보는 가벼운 도구)로부터 유도된 교정 계수를 확산 확률에 곱하는 새로운 점수 공식을 만들었습니다. 이를 통해 이야기의 흐름을 존중하는 점수를 생성합니다.
- 우선순위 기반 트리 탐색 (Priority-Based Tree Search): 단순히 최상위 경로 하나를 고르는 대신, PRESTO는 나무를 구축합니다. 이들은 새로운 교정 점수를 사용하여 어떤 가지를 키울지 결정합니다. 이들은 보스에게 수락될 가능성이 가장 높은 경로를 우선시합니다. 이는 마치 하이커가 단순히 직진하는 대신, 설령 시작할 때 가장 눈에 띄는 길은 아닐지라도 정상에 도달할 가능성이 가장 높은 경로를 지도를 보고 선택하는 것과 같습니다.
연구진은 이 나무를 성장시키는 두 가지 방법을 테스트했습니다: 빔 서치(Beam Search)(각 단계에서 고정된 수의 상위 경로를 유지함)와 최선 우선 탐색(Best-First Search)(현재까지 발견된 단일 최상 경로를 항상 확장함). 그들은 자신들의 특정 설정에서 빔 서치가 더 복잡한 최선 우선 탐색만큼 잘 작동한다는 것을 발견했고, 따라서 더 단순하고 효율적인 옵션을 선택했습니다.
결과: 더 빠르고 더 똑똑하게
저자들은 GSM8K, Math500(수학 문제), HumanEval, LiveCodeBench(코딩 챌린지), 그리고 채팅 대화 등 다양한 작업에 PRESTO를 적용하여 테스트했습니다. 그들은 두 가지 유형의 시스템을 사용했습니다:
- 전용 확산 드래프터 (Dedicated Diffusion Drafters): 더 큰 자기 회귀 모델(예: dFlash)을 위해 단어를 추측하는 작고 빠른 확산 모델.
- 자기 추측 확산 LLM (Self-Speculative Diffusion LLMs): 스스로 추측하고 스스로 검증하는 단일 확산 모델(예: Nemotron-Labs-Diffusion).
결과는 전반적으로 일관되었습니다. PRESTO는 평균 수락 길이(Average Acceptance Length), 즉 보스가 한 번에 수락하는 단어 수를 지속적으로 증가시켰습니다.
- dFlash 시스템(Qwen3-8B 사용)에서 평균 수락 길이는 약 6.6단어에서 9.6단어로 뛰었습니다. 이는 전체 과정에서 1.5배의 속도 향상으로 이어졌습니다.
- Nemotron-Labs-Diffusion 시스템에서 수락 길이는 8.8에서 9.9단어로 증가하여 1.12배의 속도 향상을 가져왔습니다.
가장 인상적인 점은 PRESTO가 시스템이 "확률적(stochastic)"인 상황(무작위성이 있는 상황, 보통 예측하기 더 어려움)에서도 작동한다는 것을 논문은 보여주었다는 것입니다. 이러한 경우 속도 향상이 더욱 두드러졌으며, 일부 벤치마크에서는 처리량이 거의 두 배에 달했습니다.
저자들은 또한 이 새로운 방법의 "비용"을 확인했습니다. 그들은 나무를 구축하고 새로운 점수를 계산하는 데 드는 추가 작업이 매우 미미하다는 것을 발견했습니다(전체 시간의 4% 미만). 전체 시간의 90% 이상은 여전히 보스 모델에 의한 실제 검증에 소비됩니다. 이는 PRESTO가 시스템을 느리게 만드는 오버헤드 없이 매우 효율적인 업그레이드임을 의미합니다.
PRESTO가 아닌 것
이 논문이 주장하지 않는 바를 명시하는 것도 중요합니다. 저자들은 단순히 (그들의 접두사 정렬 점수 없이) 나이브한 트리 구조를 적용하는 것은 최적이 아니라고 명시적으로 밝힙니다. 만약 확산 모델의 가공되지 않은 점수를 그대로 사용하여 나무를 만든다면, 점수의 "눈먼" 특성 때문에 완전한 이점을 얻을 수 없습니다. PRESTO는 바로 그 점수 불일치를 해결하는 것에 관한 것입니다.
또한, 이 논문은 확산 모델의 문제를 완전히 해결했다고 주장하지 않습니다. 그들은 자신들의 방법이 (n-gram 모델과 같이) 교정 작업을 수행하기 위해 "다루기 쉬운 접두사 정렬 신호"에 의존하고 있음을 인정합니다. 그들은 향후 연구에서 더 풍부한 신호를 탐구할 수 있을 것이라고 제안하지만, 현재로서는 그들의 단순한 교정만으로도 거대한 이득을 보기에 충분합니다.
이것이 중요한 이유
AI를 더 빠르고 효율적으로 만들기 위한 경쟁에서, 속도는 매우 중요합니다. 추측적 디코딩은 가장 정확한 모델을 사용하면서도 그에 따르는 시간적 손실을 지불하지 않게 해주기 때문에 뜨거운 주제였습니다. 그러나 현재의 방법들은 확산 모델을 선형적인 기계처럼 취급하여, 그들이 가진 다중 경로 생성 능력을 간과했다는 점에서 한계가 있었습니다.
PRESTO는 확산 모델을 다중 경로 탐험가로서 대함으로써 게임의 판도를 바꿉니다. 점수를 접두사와 정렬함으로써, PRESTO는 확산 모델의 병렬 생성 잠재력을 온전히 끌어냅니다. 그 결과, 이 시스템은 단순히 약간 더 빠른 것이 아니라 훨씬 더 효율적이며, 텍-스트를 생성하고, 수학 문제를 풀고, 코드를 작성하는 속도를 이전에는 불가능하다고 생각했던 수준으로 높여줍니다. 저자들이 말했듯, 그들은 "단일 경로"의 걷기를 "나무 기반의 원정"으로 바꾸었으며, 이를 통해 모든 발걸음이 정답을 향한 발걸음이 되도록 보장했습니다.
연구 분야의 논문에 파묻히고 계신가요?
연구 키워드에 맞는 최신 논문의 일일 다이제스트를 받아보세요 — 기술 요약 포함, 당신의 언어로.