기술 요약: VIA-SD

1. 문제 정의

추측적 디코딩(Speculative Decoding, SD)은 대규모 언어 모델(LLM)의 추론 지연 시간을 줄이기 위해 널리 채택되는 기술이다. 이는 경량화된 드래프터(drafter) 모델이 후보 토큰을 생성하면, 더 큰 타겟 모델이 이를 병렬적으로 검증하는 2단계 "초안 작성-검증(draft-verify)" 패러다임으로 작동한다. 검증기가 후보를 수용하면 해당 토큰들이 방출되고, 거절되면 시스템은 검증기를 통한 전체 자기회귀(autoregressive) 디코딩 방식으로 회귀한다.

본 논문은 기존 SD 방식의 근본적인 한계를 지적한다. 기존 방식은 이진 결정 규칙(수용 또는 전체 재계산)에 의존한다. 저자들은 많은 토큰이 드래프터의 예측이 검증기와 유사하지만 완전히 신뢰할 수는 없는 "중간 지대(middle zone)"에 놓여 있다는 점을 관찰했다. 현재의 프레임워크에서 이러한 토큰들은 불필요한 계산 비용을 초래하며 검증기에 의한 전체 재계산을 유발한다. 또한, 기존의 계층적 접근 방식은 종종 중간 모델을 이진 결정을 내리는 용도로만 사용하여, 중간 모델을 능동적인 토큰 생성이나 교정에 활용하지 못한다.

핵심 연구 질문은 드래터와 검증기 사이의 분포적 격차를 늘리지 않으면서, 일련의 중간 검증 단계를 도입함으로써 중간 지대 토큰들을 처리할 수 있는가 하는 것이다.

2. 방법론: VIA-SD

저자들은 이진 수용-거절 메커니즘을 계층적 검증 프로세스로 대체하는 다계층 프레임워크인 **VIA-SD(Verification via Intra-Model Routing for Speculative Decoding)**를 제안한다.

이론적 기초: KL 기하학

저자들은 SD를 정보 이론적 관점에서 재구성한다. 손실 없는(lossless) SD에서 총 변동(Total Variation, TV) 거리가 거절률을 엄격하게 규제한다면(이는 중간 단계가 거절을 줄일 수 없음을 의미함), 쿨백-라이블러(Kullback-Leibler, KL) 발산은 더 유연한 기하학적 관점을 제공한다.

• 피타고라스 분해(Pythagorean Decomposition): 정보 기하학의 일반화된 피타고라스 정리에 따라, 드래프터 분포 $p$와 검증기 분포 $q$ 사이의 KL 발산은 중간 분포 $u$를 통해 다음과 같이 분해될 수 있다: $D_{KL}(p \| q) \ge D_{KL}(p \| u) + D_{KL}(u \| q)$.
• 시사점: 이는 $u$가 적절한 투영(projection)인 경우, 중간 분포 $u$를 도입함으로써 직접적인 $p \to q$ 경로보다 누적 발산(및 그에 따른 검증 비용)을 줄일 수 있음을 시사한다.

아키텍처: 슬림-검증기(The Slim-Verifier)

중간 분포 $u$를 구현하기 위해, VIA-SD는 **동적 모델 내부 라우팅(Dynamic Intra-Model Routing, DIMR)**을 통해 구축된 슬림-검증기를 도입한다.

• 모델 내부 라우팅(Intra-Model Routing): 별도의 중간 모델을 학습시키는 대신, 슬림-검증기는 전체 검증기 $q$로부터 파생된 라우팅된 하위 모델(routed submodel)이다. 이는 라우팅 마스크 $z$를 기반으로 트랜스포머 레이어를 선택적으로 건너뜀으로써 생성된다.
• 분포적 일관성(Distributional Consistency): 슬림-검증기는 전체 검증기와 임베딩 및 출력 헤드를 공유하므로, 동일한 크기의 독립적인 중간 모델보다 더 나은 분포적 정렬을 유지한다.
• DIMR 최적화: 최적의 라우팅 마스크는 KL 스타일의 비용 함수를 최소화하기 위해 무작위 탐색과 베이지안 최적화를 결 combination하여 오프라인에서 찾아지며, 이를 통해 슬림-검증기가 중간 분포의 안정적인 근사치가 되도록 보장한다.

VIA-SD 파이프라인

추론 과정은 드래프트된 토큰 블록에 대해 세 단계의 검증 과정을 거친다.

1. 직접 수용(Direct Acceptance): 높은 신뢰도를 가진 토큰(드래프터 $p$가 슬림-검증기 $q'$와 잘 일치하는 경우)은 즉시 수용된다.
2. 슬림-검증기 재생성(Slim-Verifier Regeneration): 중간 정도의 신뢰도를 가진 토큰( $p$가 가깝지만 완벽하지 않은 경우)은 슬림-검증기 $q'$에 의해 처리된다. $q'$가 확신을 가지면 토큰을 재작성하고, 그렇지 않으면 다음 단계로 진행한다.
3. 전체 모델 검증(Full-Model Verification): 낮은 신뢰도 또는 불확실한 토큰은 최종 검증 또는 재생성을 위해 전체 검증기 $q$로 에스컬레이션된다.

이러한 계층적 게이팅은 시스템이 계산 비용이 저렴한 슬림-검증기를 사용하여 "중간 지대" 토큰을 해결할 수 있게 하며, 가장 어려운 경우에만 전체 모델을 사용하도록 예약한다.

3. 주요 기여

1. 이론적 재구성: 본 논문은 TV 거리 기반의 이진 패러다임의 한계를 넘어, 중간 검증 단계 도입을 정당화하기 위한 KL 기반의 정보 이론적 관점을 도입한다.
2. 슬림-검증기 설계: 중간 모델의 메모리 오버헤드를 줄이기 위해, 중간 정도의 신뢰도를 가진 토큰을 타겟팅하여 전체 모델의 분포를 효율적으로 근사하는 동적으로 라우팅된 모델 내부(intra-model) 슬림-검증기를 설계한다.
3. VIA-SD 프레임워크: 실질적인 3단계 시스템(Drafter $\to$ Slim-Verifier $\to$ Full Verifier)으로 구체화된 다계층 프레임워크를 제안한다. 이는 기존 방식의 이진적 한계를 완화한다.
4. 호환성: 드래프터의 재학습이나 검증기의 훈련 절차 수정 없이도 기존 SD 프레임워크(예: EAGLE, PEARL)와 호환된다.

4. 실험 결과

저자들은 T5, Gemma2, LLaMA2, Qwen 등 다양한 모델 제품군과 다양한 크기 구성(예: 2B$\to$9B, 7B$\to$70B)을 사용하여 네 가지 대표적인 작업(요약, 번역, 추론, QA, 코딩)에서 VIA-SD를 평가했다.

• 거절률 감소: VIA-SD는 강력한 추측적 디코딩 베이스라인(예: Cascade SD, Faster Cascades)에 비해 거절률을 지속적으로 0.10–0.22 감소시켰다.
• 속도 향상: 이 방법은 최첨단 추측적 디코딩 베이스라인 대비 10–20%의 속도 향상을 제공했다. 드래프팅이 없는 그리디 디코딩(greedy decoding)과 비교했을 때 2.5–3배의 가속을 달ucceeded했다.
• 품질: 정확도(ROUGE, BLEU 및 작업별 점수로 측정)는 유지되거나 약간 향상되었으며, 이는 "손실적(lossy)"인 중간 라우팅 특성이 생성 품질을 저하시키지 않음을 입증한다.
• 효율성: 라우팅 마스크를 위한 오프라인 탐색 비용(DIMR)은 모델 쌍당 18~68분으로 적절한 수준이며 단 한 번만 발생하지만, 온라인 디코딩 속도 향상은 상당했다.
• 어블레이션(Ablation): 실험을 통해 라우팅된 슬림-검증기가 독립적인 중간 모델(메모리 오버헤드를 증가시킴) 및 무작위 레이어 스킵 방식(분포적 일관성을 유지하지 못함)보다 우수함을 확인했다.

5. 의의 및 주장

본 논문은 VIA-SD가 확장 가능하고 효율적인 LLM 추론을 위한 다계층 추측적 디코딩을 일반적인 패러다임으로 확립한다고 주장한다.

• 패러다임 전환: 이 연구는 분야를 이진 "수용/거절" 로직에서 검증 자원의 "단계적 할당"으로 이동시켜, 토큰 불확실성의 "중간 지대"를 효과적으로 활용한다.
• 실용성: 타겟 모델 자체에서 (레이어 스킵을 통해) 중간 검증기를 도출함으로써, 별도의 중간 모델을 로드하는 데 따르는 메모리 및 엔지니어링 오버헤드를 피하며, 이는 배포 시 매우 실용적이다.
• 확장성: 결과는 드래프터와 검증기 사이의 용량 차이가 커질수록 라우팅된 슬림-검증기를 활용한 계층적 검증의 이점이 더욱 뚜렷해짐을 시사하며, 대규모 LLM 서빙을 가속화하기 위한 견고한 솔루션을 제공한다.