SpecBlock: 동적 트리 드래프팅을 이용한 블록 반복형 추측 디코딩

LLM 추론 속도를 높이는 추측 디코딩 기술에서 기존의 순차적 방식(EAGLE-3)과 병렬 방식(Medusa)의 단점을 동시에 해결했습니다. 한 번의 연산으로 여러 토큰 간의 의존성을 유지하며 예측하는 블록 단위 접근법을 통해 연산 비용은 절반으로 줄이면서도 정확도는 높게 유지합니다.

기존의 추측 디코딩은 두 가지 극단적인 방식이 존재했다. EAGLE-3 같은 자기회귀 방식은 토큰을 하나씩 순차적으로 예측하여 정확도는 높지만 연산 횟수가 많아 속도가 느리다. 반면 Medusa 같은 병렬 방식은 여러 토큰을 한 번에 예측해 빠르지만, 토큰 간의 관계를 무시하고 독립적으로 예측하기 때문에 검증 단계에서 거부될 확률이 높다.

SpecBlock은 이 문제를 해결하기 위해 '블록'이라는 개념을 도입했다. Transformer의 Attention 메커니즘이 시퀀스 전체를 한 번에 처리할 수 있다는 점에 착안하여, K개의 토큰 위치를 하나의 블록으로 묶어 동시에 계산한다. 이때 단순히 병렬로 나열하는 것이 아니라, 레이어와 레이어 사이에서 이전 토큰의 정보를 옆으로 밀어 넣어주는 'Layer-wise Shift'를 통해 병렬 연산 안에서도 순차적인 의존성을 모방한다.

결과적으로 한 번의 연산으로 여러 개의 연관된 토큰 후보를 만들어낼 수 있게 된다. 이는 마치 문장의 다음 단어들을 예측할 때, 앞 단어가 무엇인지 고려하면서도 여러 단어를 한 호흡에 내뱉는 것과 유사한 원리다. 이를 통해 연산 효율성과 예측 정확도라는 두 마리 토끼를 모두 잡았다.

SpecBlock은 타겟 모델의 중간 레이어 hidden state를 입력받아 K개의 연속된 위치를 예측한다. 입력 단계에서 Prefix Broadcast를 통해 타겟 모델의 문맥 정보를 모든 예측 위치에 공유하며, 각 위치마다 고유한 learnable query q를 할당하여 위치 정보를 구분한다.

핵심인 Within-block dependence는 Layer-wise Shift로 구현된다. ℓ번째 레이어의 k번째 위치 상태 h(ℓ,k)와 k-1번째 위치 상태 h(ℓ,k-1)를 결합(concatenation)한 후 선형 투영(linear projection)을 거쳐 ℓ+1번째 레이어의 입력으로 전달한다. [이전 레이어의 현재 위치 상태와 이전 위치 상태 입력 → 선형 행렬 연산 → 다음 레이어의 입력값 생성] 과정을 통해 병렬 연산 중에도 위치 간 정보 교류가 발생한다.

트리 구성은 Rank Head를 통해 동적으로 이루어진다. Rank Head는 각 위치의 hidden state와 소프트맥스 분포의 통계적 요약을 입력받아 4개의 버킷(b0~b3) 중 하나로 분류한다. [hidden state와 분포 요약값 입력 → 분류기 연산 → 버킷 레이블 출력] 결과에 따라 해당 위치에서 몇 개의 후보(branching width)를 생성할지 결정하며, 불확실성이 높은 위치에서 더 많은 가지를 뻗도록 설계됐다.

Llama-3.1-8B 모델을 타겟으로 실험한 결과, SpecBlock은 EAGLE-3 대비 평균 8~13%의 추가 속도 향상을 달성했다. 특히 드래프팅에 소요되는 연산 비용(TD%)은 EAGLE-3가 전체 추론 시간의 31%를 차지하는 반면, SpecBlock은 16% 수준으로 약 절반에 불과했다.

서빙 타임 적응(Cost-aware adaptation) 기법을 적용했을 때 성능 향상은 더욱 두드러졌다. 도메인이 다른 데이터셋에서 성능이 저하되는 기존 모델들과 달리, 실시간 피드백을 통한 가중치 업데이트를 통해 EAGLE-3 대비 최대 19%까지 속도 우위를 확장했다. Ablation study에서는 Layer-wise Shift를 제거했을 때 속도가 3.21x에서 2.95x로 급감하여, 블록 내 의존성 유지가 성능의 핵심임을 입증했다.

SpecBlock 아키텍처는 타겟 모델의 low, mid, top 레이어 hidden state를 추출하여 컨디셔닝 벡터 ct를 생성한다. 이를 임베딩된 마지막 토큰 및 위치 쿼리와 융합하여 드래프터의 입력으로 사용한다. 드래프터는 타겟 모델과 동일한 차원의 L개 Transformer 레이어로 구성된다.

학습 시에는 Valid-prefix mask를 사용하여 이전 토큰이 틀렸을 경우 이후 위치의 손실(loss)을 계산에서 제외하는 커리큘럼 학습을 수행한다. 이는 추론 시 검증 모델이 첫 번째 오류 지점에서 생성을 중단하는 특성을 반영한 것이다. 또한 Cross-block training을 통해 타겟 모델의 정보 없이 드래프터 자신의 이전 블록 상태만으로 다음 블록을 예측하는 상황에 대비한다.

온라인 적응을 위한 Bandit 알고리즘은 'Skip', 'Head-only update', 'Full update'의 세 가지 액션 중 하나를 선택한다. 업데이트 신호 s는 거부된 토큰들의 확률 오차 합으로 계산되며, EWMA(지수 이동 평균)를 통해 예상되는 처리량 이득이 비용보다 클 때만 역전파(backpropagation)를 수행하여 서빙 중 연산 낭비를 방지한다.

Rank Head의 4개 버킷 분류 정확도가 완벽하지 않아 일부 위치에서 검증 예산이 비효율적으로 할당될 수 있다. 또한 블록 너비 K는 학습 시 고정되므로 추론 시 하드웨어 사양에 맞춰 동적으로 변경하기 어렵다는 한계가 있다.