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

왜 중요한가

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

핵심 기여

병렬 트리 드래프팅을 위한 causal parallel draft head 도입

JetSpec은 대상 모델의 frozen hidden states를 융합하여 단일 전방 패스로 트리의 모든 활성 노드에 대한 로짓을 예측하는 causal parallel draft head를 도입했다. 이 드래프트 헤드는 각 노드가 자신의 조상 토큰에만 조건화되도록 tree-causal attention 마스크를 적용하여 branch-wise autoregressive 분해와 정합성을 확보한다. 결과적으로 드래프팅 비용을 낮추면서도 분기별 수용률을 크게 개선하여 더 큰 예산을 효율적으로 활용했다.

예산 기반 트리 확장 알고리즘과 누적 로그확률 스코어링

JetSpec은 최대 깊이, 분기 폭, 전체 노드 예산을 입력으로 받아 best-first 방식으로 노드를 확장하는 트리 구축 과정을 설계했다. 각 경로에 대해 드래프트 로그확률을 누적 점수로 사용하여 우선순위 탐색을 수행하고 예산이 소진되거나 확장 불가 노드가 남을 때까지 반복한다. 이 방식은 branch-aware한 후보 트리를 생성하여 대상 모델의 병렬 검증 단계에서 더 긴 접두사가 수용되도록 만들었다.

forward KL 기반의 soft-label distillation 학습 규범 채택

드래프트 헤드는 대상 모델의 온도 조정된 logits를 교사로 사용하여 forward KL을 최소화하는 방식으로 학습되었다. 실험에서 forward KL은 reverse KL보다 tree drafting 성능 측면에서 우수했고 hard-label SFT보다 약간 더 나은 결과를 보였다. 학습 데이터로는 대상 모델이 재생성한 continuations와 코퍼스 기반 블록이 활용되며 활성 위치 가중치로 손실을 정규화했다.

실제 서빙 엔진(vLLM) 통합 및 로드 의존적 예산 판단 실험

JetSpec은 vLLM에 통합되어 트리 검증과 페이징 FlashAttention 커널 확장을 적용한 실제 서빙 환경에서 평가되었다. 배치 크기와 요청률에 따라 최적 트리 예산이 달라짐이 관찰되었고, 낮은 부하에서는 큰 예산이 TPS를 크게 향상시켰다. 이 결과는 실무에서 예산 선택이 서빙 부하에 민감하므로 동적 예산 스케줄링이 필요함을 시사한다.

핵심 아이디어 이해하기

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 배포 전략이 필요한 환경에서는 적용 제약이 존재한다.

실무 활용

JetSpec은 대상 모델의 은닉 상태를 재사용하는 경량 드래프트 헤드를 학습하여 실시간 추론에서 디코딩 반복 횟수를 줄이는 방식으로 서빙 레이턴시를 감소시킨다. 실제 vLLM 통합 결과 낮은 부하 및 중간 부하 환경에서 의미 있는 TPS 향상을 보였으며, 고부하 상황에서는 예산 조정으로 효율을 유지해야 한다. 연구자는 제공된 코드 저장소를 통해 실험적 통합과 예산 튜닝을 재현할 수 있다.

대화형 서비스에서 긴 문장이나 다단계 추론이 빈번한 워크로드의 응답 레이턴시를 줄이는 용도
수학·코딩 벤치마크처럼 출력 길이가 길고 검증이 중요한 생성 작업에서 추론 비용 대비 처리량을 높이는 용도
실제 서빙 환경에서 GPU 여유 자원을 활용해 반복 수를 줄이고 배치 처리 효율을 개선하는 용도

코드 공개 여부: 공개

코드 저장소 보기

키워드

Speculative Decoding(스펙큘레이티브 디코딩)JetSpec(JetSpec)Draft Head(드래프트 헤드)Tree-Causal Attention(트리 인과적 어텐션)vLLM(vLLM)

용어 해설

스펙큘레이티브 디코딩(Speculative Decoding): — 스펙큘레이티브 디코딩은 경량의 drafter가 다수의 토큰을 초안으로 생성하고 대상 모델이 병렬로 검증하여 일치하는 가장 긴 접두사를 채택하는 추론 기법이다. 초안의 수와 초안 품질, 그리고 초안 생성 비용이 전체 속도 향상에 직접적으로 영향을 주며, 수식으로 속도향상 = (1-α^{N+1})/((1-α)(Nc+1)) 형태로 표현된다. 본 논문에서는 초안 예산을 늘릴 때 acceptance rate α와 토큰당 비용 c의 동시 개선이 핵심 병목이라 규정한다.
드래프트 헤드(Draft Head): — 드래프트 헤드는 목표 모델의 중간 은닉 상태를 재사용하는 경량 예측 헤드로 한 번의 전방 패스로 여러 초안 토큰의 로짓을 생성한다. 별도 drafter 모델을 배포하지 않으면서도 낮은 per-token 비용을 달성하고 target-model의 신호를 반영하기 위해 hidden state를 KV에 주입하거나 feature fusion을 수행한다. JetSpec은 이러한 드래프트 헤드에 branch-wise causal conditioning을 더해 트리 드래프팅과 정합성을 확보한다.
트리 인과적 어텐션(Tree-Causal Attention): — 트리 인과적 어텐션은 트리의 각 노드가 자신과 조상 노드의 토큰만 볼 수 있고 형제나 자손은 볼 수 없게 하는 마스킹 규칙이다. 이 마스크는 각 분기(branch)의 확률 분포가 해당 분기의 선행 토큰들에 조건화되도록 하여 병렬 계산으로도 autoregressive factorization과 유사한 분해를 유지한다. JetSpec은 이 마스크를 사용하여 모든 트리 깊이의 로짓을 병렬로 계산하면서 branch-wise causality를 보존했다.
Forward KL 증류(Forward KL Distillation): — Forward KL 증류는 대상 모델의 온도 조정된 확률 분포 p~를 교사로 삼아 드래프트 분포 q~가 그 확률 질량을 포괄하도록 KL(p~ || q~)를 최소화하는 목적이다. 본 논문에서는 forward KL이 mode-seeking 성향의 reverse KL보다 트리 드래프팅에서 수용 길이(accepted length)와 end-to-end speedup 관점에서 우수함이 관찰되었다. 학습 손실은 온도 T_KD와 활성 위치 마스크로 정규화된 평균 forward KL로 구성된다.

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

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

용어 해설

관련 토론

댓글

관련 피드

왜 중요한가

핵심 기여

핵심 아이디어 이해하기

방법론

주요 결과

기술 상세

한계점

실무 활용

키워드

관련 토론

댓글

관련 피드