JetSpec 병렬 트리 드래프팅으로 speculative decoding의 스케일링 한계 극복

Autoregressive LLM의 디코딩은 본질적으로 직렬적이어서 긴 생성 작업에서 지연이 주요 병목이다. 스펙큘레이티브 디코딩은 초안 생성과 병렬 검증으로 이 지연을 줄이지만 초안 예산을 늘릴 때는 acceptance rate과 초안 비용의 동시 개선이 필요하다. JetSpec은 병렬 예측의 효율성과 분기별(in-branch) 인과적 조건화를 동시에 확보하여 더 큰 초안 예산을 실제 accepted prefix 길이와 엔드투엔드 속도 향상으로 변환했다.

Autoregressive 디코더는 토큰을 한 번에 직렬적으로 생성하므로 동작상 latency가 쌓이며 긴 생성에서 병목이 발생한다. 스펙큘레이티브 디코딩은 경량 drafter가 여러 토큰을 초안으로 만들고 대상 모델이 병렬로 검증하여 한 번의 반복에서 여러 토큰을 전진할 수 있게 한다. 하지만 초안 수를 늘릴수록 acceptance rate α가 유지되지 않거나 드래프팅 비용 c가 증가하면 이득이 사라진다.

기존 병렬 드래프팅은 한 번에 많은 위치를 예측하여 비용을 낮추지만 각 위치가 분기별 조상 토큰에 조건화되지 않아 서로 모순되는 토큰 조합이 만들어질 수 있다. 반대로 autoregressive drafters는 분기별 조건화를 유지하지만 트리 깊이가 커질수록 순차 드래프팅 비용이 증가한다. JetSpec의 핵심은 tree-causal attention 마스크를 적용한 causal parallel draft head로, 모든 활성 트리 노드의 로짓을 병렬로 계산하면서 각 노드가 자신의 조상만 참조하도록 제한하여 branch-wise autoregressive 분해를 근사한다.

이 접근은 드래프팅 비용을 낮춘 채로 초안 예산을 늘렸을 때 실제로 더 긴 accepted prefixes로 이어지게 만든다. 실험에서 JetSpec은 높은 예산 구간에서 DFlash 등 branch-agnostic 방법 대비 더 높은 수용 길이와 종단 간 속도 향상을 얻었으며 MATH-500에서 최대 9.64×, Qwen3-30B-A3B에서도 유사한 수준의 향상이 관찰되었다. 따라서 병렬성과 인과성을 동시 만족시키면 초안 예산을 효율적으로 활용할 수 있다.

전체 접근은 frozen target 모델의 은닉 상태를 융합하여 경량 드래프트 헤드가 하나의 전방 패스로 트리의 모든 활성 노드 로짓을 계산하도록 하는 것이다. 드래프트 헤드는 입력 접두사와 조상 토큰만을 볼 수 있게 하는 tree-causal attention 마스크를 적용한다. 마스크는 각 노드의 쿼리가 자신과 조상 토큰의 키-값에만 0을 부여하고 다른 노드에는 -∞를 부여하여 softmax 후 해당 항목이 무시되게 만든다.

트리 구축은 최대 깊이 N, 분기 폭 W, 노드 예산 B를 기준으로 best-first 탐색을 수행한다. 드래프트 로짓으로 각 위치의 상위-W 토큰을 취하고 각 경로는 누적 draft log-probability로 점수를 매긴다. 그 후 최상위 확장 가능한 노드를 팝하여 자식들을 생성하고 점수와 함께 재삽입하여 예산이 소진될 때까지 반복함으로써 후보 트리 𝒯(x)를 얻는다.

검증 단계에서는 대상 모델이 트리의 모든 노드를 병렬로 검증하며 각 분기에서 사후적으로 acceptance rule을 적용한다. 비확률적 그리디 모드에서는 대상 모델의 다음 토큰 예측과 일치하면 승인되며 논리적 비유의 비그리디 모드에서는 rejection sampling 기반의 α_t = min(1, p/q) 규칙이 사용되어 필요 시 corrective token이 삽입된다. 학습 단계에서는 대상 모델의 온도 조정 logits를 교사로 삼아 forward KL을 최소화하는 distillation objective로 드래프트 헤드를 학습했다.

구현 측면에서는 Triton 최적화된 디코딩 커널과 SM90 기반의 paged FlashAttention 확장을 포함한 커스텀 tree-attention 경로를 사용하여 vLLM에 통합했다. 서빙 실험에서는 배치 크기, 요청률, GPU 종류(H100, B200)에 따른 예산 효용을 측정하여 실전 적용 시 예산 선택의 로드 의존성을 분석했다.

JetSpec은 다양한 수학, 코딩, 대화 벤치마크에서 기존 head-based 및 block-parallel baselines을 일관되게 능가했다. 저예산(예: 16 토큰) 구간에서는 DFlash와 유사한 성능을 보였으나 고예산(예: 256 토큰)에서는 JetSpec이 MATH-500에서 최대 9.64×, 대화 작업에서 최대 4.58× 엔드투엔드 속도 향상을 기록했다. Dense 및 MoE Qwen3 계열에서 모두 개선이 관찰되어 방법의 일반화 가능성이 확인되었다.

학습 목적의 비교에서 forward KL 증류는 reverse KL보다 우수했고 hard-label SFT과는 유사하거나 약간 우수한 결과를 냈다. γ 하이퍼파라미터(위치 가중치 감쇠)와 아키텍처별 민감도 실험에서 causal head는 γ 변화에 대해 더 안정적이었고 diffusion-style head는 특정 γ에서 성능 붕괴를 보였다. 이 차이는 branch-agnostic 분포가 상호 모순되는 토큰 조합을 더 높게 랭크하는 현상에서 비롯되었고, 사례별 분석에서 causal head가 더 많은 토큰을 수용하여 평균 accepted length를 증가시켰다.

서빙 실험에서는 낮은 부하에서 큰 트리 예산이 TPS를 크게 높였고 예를 들어 batch size 1에서 예산을 16에서 128로 늘리면 TPS가 443.3에서 968.2로 증가했다. 반대로 높은 부하나 큰 배치에서는 검증 오버헤드와 메모리 트래픽으로 인해 예산 증가에 따른 이득이 포화되었다. 따라서 실무에서는 서빙 부하에 맞춘 예산 선택이 필요하며 본 연구는 정적 정책을 사용했고 동적 예산 스케줄링은 후속 과제로 남겼다.

전체 아키텍처는 frozen target model에서 fused hidden states h_x^o를 추출하고 이를 입력으로 causal parallel draft head가 병렬 로짓을 예측하는 구조이다. tree-causal attention 마스크는 각 노드 v에 대해 조상 노드와 자신만 접근하도록 M_v,u를 0 또는 -∞로 설정하여 softmax 기반 attention이 조상 토큰에만 의존하게 만든다. 이로써 병렬 계산으로도 각 분기의 조건부 분해 q(π(v)|x)=∏u q(y_u|x,h_x^o,π{<u})에 근사한다.

학습은 anchor position을 샘플링하여 길이 N 블록의 활성 위치들에 대해 목표 모델의 온도 조정 logits를 교사로 사용하고 forward KL D_{KL}(p~||q~)를 최소화하는 방식으로 수행되었다. 온도 T_KD와 활성 위치 가중치 w_m을 도입하여 손실을 정규화하고, γ 파라미터를 통해 anchor로부터 멀어질수록 손실 기여를 지수적으로 감소시키는 실험을 수행했다. reverse KL은 mode-seeking 성향으로 인해 tree drafting에서 수용 길이를 줄이는 부작용이 관찰되어 본 구현에서는 forward KL을 기본으로 채택했다.

트리 확장 알고리즘은 각 깊이에서 상위-W 후보를 고르고 누적 드래프트 로그확률로 경로 점수를 계산한 뒤 best-first 방식으로 노드를 팝하여 자식 노드를 삽입하는 반복 과정을 거친다. 검증 단계에서는 대상 모델이 트리의 모든 노드를 병렬로 평가하고 비그리디 모드에서는 α_t = min(1, p/q) 규칙을 통해 rejection sampling 기반의 승인 판정을 수행한다. 서빙 통합은 Triton 및 custom SM90 paged FlashAttention 커널을 사용하여 tree attention과 tree-tail masking을 효율적으로 구현했고 vLLM과의 통합을 통해 실제 request-rate 하에서 throughput 변화를 측정했다.

논문은 트리 예산의 최적 선택이 서빙 부하에 강하게 의존함을 명시적으로 보고했고 동적 예산 스케줄링은 구현 대상에서 제외되었다. 학습 데이터의 효율성 측면에서 대상 모델이 생성한 regenerated continuations가 최고 성능을 냈으나 이 재생성 과정은 계산 비용이 크기 때문에 대규모 재현성 확보에 추가 비용이 발생한다. 또한 JetSpec의 효용은 드래프트 헤드가 대상 모델의 은닉 표현을 사용할 수 있는 head-based 설정에 의존하므로 별도의 drafter 배포 전략이 필요한 환경에서는 적용 제약이 존재한다.