기술 요약: TokenTiming

문제 정의

Speculative Decoding(SD) 은 작은 "드래프트" 모델을 사용하여 토큰을 제안하고, 이를 더 큰 "타겟" 모델이 검증함으로써 대규모 언어 모델 (LLM) 추론을 가속화하는 두드러진 기술입니다. 그러나 현재 SD 방법들은 근본적인 제약에 직면해 있습니다: 드래프트 모델과 타겟 모델은 동일한 어휘 ($V_d = V_t$) 를 공유해야 합니다. 이 요구 사항은 사용 가능한 드래프트 모델의 선택을 심각하게 제한하며, 종종 새로운 드래프트 모델을 처음부터 훈련하거나, 작은 변형 모델조차 효율적인 가속화에 너무 큰 특정 모델 계열을 사용하도록 강요합니다.

"이질적인 어휘"(여기서 $V_d \neq V_t$) 를 해결하려는 기존 시도들은 제한된 성공을 보였습니다. Token-Level Intersection(TLI) 와 같은 방법들은 어휘의 교집합에 의존하여 어휘 밖 토큰에 대한 확률 질량을 폐기함으로써 디코딩 과정의 무손실성을 위반합니다. String-level Exact Match(SLEM) 와 같은 다른 접근법들은 확률적 샘플링을 수행할 수 없습니다. 결과적으로, 재훈련 없이 시중의 모든 드래프트 모델을 활용할 수 있는 보편적이고 무손실적인 SD 프레임워크가 부재합니다.

방법론: TokenTiming

저자들은 타겟 또는 드래프트 모델을 수정하지 않고 이질적인 어휘 간 무손실 Speculative Decoding 을 가능하게 하는 동적 정렬 방법인 TokenTiming을 제안합니다. 핵심 혁신은 시계열 분석의 Dynamic Time Warping(DTW) 에서 차용한 Dynamic Token Warping(DTW) 알고리즘입니다.

TokenTiming 파이프라인은 다음과 같이 작동합니다:

1. 재 토큰화 (문자 수준 브릿지):
   각 디코딩 단계에서 드래프트 모델은 토큰 시퀀스 $D = (d_1, \dots, d_k)$를 생성합니다. TokenTiming 은 이러한 토큰들을 타겟 어휘와 직접 비교하는 대신, 드래프트 토크나이저를 사용하여 $D$를 중간 문자열 표현 $S$로 디코딩합니다. 그런 다음 이 문자열을 타겟 토크나이저로 다시 인코딩하여 "프록시 타겟 토큰 시퀀스" $T = (t_1, \dots, t_m)$를 생성합니다. 이 단계는 서로 다른 토큰화 세분성 사이의 간극을 메우기 위해 기본 문자 스트림이 동일하도록 보장합니다.

2. Dynamic Token Warping(DTW) 정렬:
   토크나이저 차이로 인해 드래프트 토큰 수 ($k$) 와 프록시 타겟 토큰 수 ($m$) 가 종종 다르기 때문에 단순한 1 대 1 매핑은 불가능합니다. TokenTiming 은 $D$와 $T$ 사이의 최적의 다대다 정렬 경로 $\pi^*$를 구축하기 위해 DTW 를 사용합니다.

   • 거리 척도: 알고리즘은 토큰 간의 비유사성을 측정하기 위해 Levenshtein 거리 (편집 거리) 를 사용합니다.
   • 제약 조건: 계산 효율성을 보장하기 위해 Sakoe-Chiba Band 제약이 적용되어 검색 공간을 폭 $w$의 대각선 밴드로 제한합니다. 이로 인해 복잡도가 $O(k \cdot m)$에서 $O(w \cdot \max(k, m))$로 감소합니다.
   • 출력: 결과는 드래프트 토큰을 프록시 타겟 토큰에 정렬하는 매핑 $\pi^*$이며, 하나의 드래프트 토큰이 여러 타겟 토큰에 매핑되는 경우 (예: "Scaling" $\to$ "Scale", "ing") 또는 그 반대의 경우를 처리합니다.

3. 확률 전달 및 검증:
   정렬 경로 $\pi^*$를 사용하여 드래프트 모델의 확률 분포가 타겟 어휘로 전달됩니다.

   • 다대일: 여러 드래프트 토큰이 단일 타겟 토큰에 매핑되는 경우, 최종 드래프트 토큰의 확률이 타겟 토큰에 할당됩니다.
   • 일대다: 하나의 드래프트 토큰이 여러 타겟 토큰에 매핑되는 경우, 드래프트의 확률이 모든 해당 타겟 토큰에 복사됩니다.
     그런 다음 타겟 모델은 프록시 시퀀스 $T$를 병렬로 검증합니다. 토큰들은 표준 Speculative Sampling 기준인 $r < \min(1, q(t_j)/p(t_j))$에 따라 수용됩니다. 여기서 $q$는 타겟 확률이고 $p$는 전달된 드래프트 확률입니다. 이는 최종 출력 분포가 타겟 모델의 분포와 동일하게 유지됨을 보장하여 무손실을 보장합니다.

주요 기여

• 보편적 호환성: TokenTiming 은 공유 어휘 제약을 제거하여 아키텍처나 토크나이저에 관계없이 시중의 모든 드래프트 모델을 재훈련이나 모델 수정 없이 타겟 모델에 플러그인할 수 있게 합니다.
• 무손실 가속화: DTW 를 통해 완전한 토큰 매핑을 구축함으로써, 이 방법은 Speculative Sampling 이 무손실로 유지되도록 보장하며, 어휘가 불일치할 때조차 타겟 모델의 출력 분포 무결성을 유지합니다.
• 동적 정렬: 이 방법은 각 디코딩 단계에서 실시간으로 작동하며, DTW 메커니즘을 통해 가변적인 토큰 수와 세분성을 동적으로 처리합니다.

실험 결과

저자들은 7B 에서 70B 파라미터 타겟과 68M 에서 350M 파라미터 드래프트 모델을 포함한 다양한 모델 쌍을 사용하여 요약, 번역, 코드 생성, 수학, QA 등 다양한 작업 세트를 평가했습니다.

• 가속화 성능: TokenTiming 은 다양한 작업에서 autoregressive 기준 대비 최대 1.57 배의 속도 향상을 달성했습니다. 예를 들어, Qwen3-0.6B 드래프트를 사용한 Qwen3-32B 타겟에서 1.57 배의 속도 향상을 달성하여 TLI 기준 (1.21 배) 을 크게 능가했습니다.
• 최첨단 (SOTA) 과의 비교: 7B 및 33B 타겟 모델에서 TokenTiming 은 EAGLE-3 와 같은 동질적 어휘 기반 최첨단 (SOTA) 방법과의 성능 격차를 좁혔습니다. 구체적으로, 33B 타겟에서 TokenTiming 은 2.27 배의 속도 향상을 달성하여 Medusa(1.71 배) 와 EAGLE-1(2.21 배) 을 능가하면서도 이질적인 드래프트를 사용할 수 있는 유연성을 유지했습니다.
• 작업 강건성: 이 방법은 평가된 다섯 가지 작업 카테고리 전반에 걸쳐 TLI 를 일관되게 능가했습니다. 수학 및 요약 작업에서 속도 향상은 특히 두드러졌습니다 (예: 수학 작업에서 2.53 배).
• 오버헤드: DTW 정렬 및 재 토큰화로 인한 계산 오버헤드는 무시할 수 있으며, 디코딩 주기당 차단 시간이 약 663µs 이므로 총 실행 시간의 0.1% 에서 0.5% 만 차지합니다.

중요성 및 주장

이 논문은 TokenTiming 이 LLM 가속화를 위한 다용도적이고 실용적인 도구로 Speculative Decoding 을 만드는 중요한 단계라고 주장합니다. 어휘 불일치라는 근본적인 병목 현상을 제거함으로써, 이는 "드래프트 모델 선택을 위한 보편적 접근법"을 가능하게 합니다. 이를 통해 실무자들은 새로운 드래프트 헤드를 훈련하거나 특정 타겟 어휘를 위해 모델을 재훈련하는 비용과 비유연성 없이 기존에 존재하는 작고 이질적인 모델들을 드래프트로 활용할 수 있습니다. 이 연구는 "플러그 앤 플레이" 호환성을 통해 고속 추론을 달성할 수 있음을 보여주며, SD 의 이론적 효율성과 다양한 모델 생태계의 실용적 제약 사이의 간극을 메웠습니다.