확률적 디코딩에서 동적 트리 추측 디코딩의 원-핫 붕괴를 푸는 이중 확률 설계

무엇인가

추측 디코딩(speculative decoding)은 가벼운 드래프트 모델이 여러 토큰을 병렬로 예측하고 타깃 모델이 한 번의 포워드 패스로 검증하는 방식으로 LLM 추론 지연을 줄인다. 최근에는 후보 시퀀스를 트리로 묶는 트리 기반 추측 디코딩이 주류가 됐고, EAGLE-2/3 같은 동적 트리 방법은 각 스텝에서 문맥에 맞춰 토폴로지를 적응적으로 만든다. 문제는 이 파이프라인이 본질적으로 결정적이라는 점이다. 드래프트 분포가 계산되면 top-K 확장과 전역 프루닝으로 트리 구조와 후보가 완전히 결정되고 무작위성이 개입할 여지가 없다. T=0 그리디 디코딩에서는 이 방식이 잘 작동하지만, T>0 확률적 디코딩에서는 검증 단계에서 드래프트 확률이 1.0으로 취급되면서 분포가 원-핫으로 붕괴하고 꼬리 분포 탐색이 사라져 수용률이 크게 떨어진다. 논문은 이를 딜레마로 정식화한다. 동적 트리 방법은 문맥 인지 토폴로지를 지키는 대신 확률적 샘플링을 포기하고, 정적 트리 방법은 확률적 샘플링을 지키는 대신 문맥 무관 구조를 택한다. 근본 원인은 같은 확률 분포가 트리 구성과 토큰 검증이라는 상충하는 두 작업에 동시에 쓰인다는 결합(coupling)이다. 토큰의 생존 여부가 자기 확률 값에 의해 결정되면, 생존을 조건으로 한 조건부 분포가 원래 샘플링 분포에서 벗어나 무손실성이 깨진다.

어떻게 동작하나

RheoSampling의 핵심은 이 결합을 끊는 이중 정체성(dual-identity) 설계다. 각 확장 단계의 후보 풀 K개 슬롯을 세 부분으로 나눈다. 먼저 드래프트 확률이 가장 높은 top-m 토큰을 결정적으로 뽑고(lead 토큰), 그다음 top-m 너머의 잔차 분포 q~에서 대표 토큰 x_s를 하나 샘플링하며, 남은 K-m-1 슬롯은 나머지 어휘에서 순위가 높은 fill 토큰으로 채운다. 결정적 고확률 토큰들과 분포 꼬리를 탐색하는 확률적 탐침 하나가 섞인 후보 집합이 만들어지는 셈이다. 결정적 부분은 x_s를 트리 구성 단계에서 대리 확률 q_proxy(x_s) = min{q(x_m), z}로 다룬다. 여기서 z = 1 - Σ_{i=1}^{m} q(x_i)는 잔차 질량이고 q(x_m)은 m번째 결정적 확률이다. 이 값은 x_s를 후보 풀 안의 모든 fill 토큰보다 앞세우므로, 전역 프루닝을 통과할지 여부가 실제로 샘플링된 토큰의 정체성과 독립이 된다. m=0일 때는 대리 확률을 q(x_1)+ε(ε>0)로 두어 샘플 토큰을 첫 슬롯에 고정한다. 반면 검증 단계에서는 대리 확률이 아니라 실제 샘플링 확률 q~(x_s)로 토큰을 평가한다. 트리는 통제된 추정치 위에서 세우고, 검증은 실제 샘플링 분포를 존중해 무손실성을 유지하는 구조다.

무엇과 다른가

m은 이 논문이 rheostat(가변 저항)라고 부르는 조절 손잡이다. m이 작으면 x_s에 더 큰 대리 확률이 배정되어 최종 트리에 살아남을 확률이 올라가지만 결정적 백본이 줄고, m이 크면 트리 토폴로지 품질은 좋아지지만 대리 확률이 낮아져 샘플 노드가 리랭킹에서 잘릴 위험이 커진다. 검증은 Optimal Transport(OT) 기반의 RheoVerification으로 수행하며, 샘플 토큰을 먼저 배치하는 stochastic-first 전략을 쓴다. 대어휘에서 샘플링 비용을 줄이기 위해 sparse draft 분포 전략도 도입한다. 소프트맥스 전에 로짓을 top-128까지만 남기고 나머지를 -∞로 잘라 희소 분포를 만들며, 검증도 드래프팅과 같은 절단 분포를 쓰는 한 무손실성이 유지된다. top-128이 드래프트 분포 확률 질량의 대부분을 담기 때문에 수용률도 유지된다고 저자들은 밝힌다.

어떻게 쓰나

이론적 보장의 핵심은 무손실성 정리다. 동적 트리에서는 샘플 토큰 Y ~ q~가 자기 값뿐 아니라 이후 fill 토큰들의 정체성까지 결정하고, 이것이 다시 깊은 노드의 경로 점수와 전역 프루닝 생존 여부를 바꾼다. 그래서 프루닝된 트리 T_R은 토폴로지와 각 노드의 토큰 채움 모두가 확률적인, 결합된 확률 변수가 된다. 논문은 등가류(equivalence-class) 분석으로 이 확률적 트리 공간을 다룰 수 있는 클래스로 압축해, 임의의 프루닝 예산 R ≥ 1과 임의 토큰 시퀀스에 대해 E_{T_R}[Pr(Rheo(T_R) = Seq)] = p(Seq)임을 보인다(Theorem 1). 이는 확률적 동적 트리에 대한 첫 엄밀한 무손실성 증명이라고 저자들은 주장한다. 대리 확률이 q~와 독립인 것만으로는 충분하지 않고 q_proxy ≥ q(x_{m+1})라는 순위 하한을 만족해야 한다는 것도 Lemma 1로 제시한다. 어휘 {A,B,C}에 드래프트 분포 {0.5, 0.3, 0.2}, K=3, m=1인 예에서 Rheo-proxy를 min{0.5, 0.5}=0.5로 두면 샘플 토큰은 어떤 토큰이 뽑히든 항상 rank 2에 고정되어 생존이 Y와 독립이지만, adhoc-proxy를 0.25로 두면 순위가 fill 토큰에 따라 흔들리고 생존을 조건으로 Y가 확률 1로 B에 강제되어(q~(B)=0.6이어서) 검증이 잘못된 기저 확률을 쓰게 되며 무손실성이 깨진다. 단일 레이어 수용률의 닫힌 형태도 Theorem 3로 제시된다.

전제와 한계

실험은 6개 벤치마크에서 수행했다. Alpaca, GSM8K, HumanEval, MT-bench, Natural Questions, CNN/DailyMail이며 각 데이터셋 80문항으로 지시 따르기, 수학 추론, 코드 생성, 멀티턴 대화, 질의응답, 요약을 아우른다. 타깃 모델은 Llama-3.1-8B-Instruct, Vicuna-13B-v1.3, DeepSeek-R1-Distill-Llama-8B이고, 드래프트 모델은 공식 EAGLE-3 체크포인트를 추가 파인튜닝 없이 사용했다. 지표는 사이클당 평균 수용 토큰 수인 평균 수용 길이 τ와 자기회귀 디코딩 대비 실제 wall-clock 속도 향상이다. 트리 크기 60, 드래프트 깊이 8, NVIDIA A6000 한 장, 시드 3회, 기본 온도 T=1.0이다. 결과적으로 RheoSampling은 모든 구성에서 vanilla Top-K 베이스라인 대비 τ를 일관되게 개선했고, 개선폭은 Vicuna-13B에서 +0.14, DeepSeek-R1-Distill-8B에서 +0.22였다. Llama-3.1-8B에서는 τ가 5.04에서 5.25로 4.2% 올랐고 속도 향상은 2.84×에서 2.93×로 3.2% 증가했다. 속도 개선폭이 τ 개선폭보다 약간 작은 것은 확률적 탐침의 샘플링·검증 오버헤드 때문이며, 그럼에도 모든 모델에서 순 속도 향상은 양수였다.

절제 실험에서는 m=1이 대부분 설정에서 최고 또는 근접 성능을 냈다. m=0과 m=2도 Top-K 베이스라인을 상회했지만 m≥3은 대리 확률이 줄어 샘플 토큰이 리랭킹에서 잘리고 이득이 미미해졌다. m=0은 잔차 분포가 전체 드래프트 분포로 퇴화해 OT 기반 확률 재배치 이점을 충분히 쓰지 못하고, 샘플 토큰이 원래 top-1 후보와 겹치면 트리 구조가 Top-K와 같아지지만 아니면 크게 흔들리는 all-or-nothing 성격을 보인다. 온도와의 상호작용도 분석한다. 트리 구성(확장·리랭킹)은 온도 스케일링 없이 원래 드래프트 확률로 수행하고 온도는 샘플링과 검증 단계에만 영향을 준다. T→0에서 m=0은 샘플링과 구조가 모두 Top-K로 붕괴하는 이중 퇴화를, m=1은 샘플링만 붕괴하고 트리 구조는 Top-K와 다르게 남는 단일 퇴화를 보이며, 이 구조적 차이가 T→0에서 m=1과 Top-K/m=0 사이의 수용률 격차를 만든다.

실무 관점에서 이 논문은 T>0로 샘플링하는 LLM 서빙에서 EAGLE-3 계열 동적 트리 추측 디코딩을 쓸 때의 수용률 손실을 줄이는 방법을 제시한다. 적용 시 확인할 점은 몇 가지다. 검증 단계가 드래프팅과 동일한 절단 분포(top-128)를 쓰는지, 대리 확률이 순위 하한 q_proxy ≥ q(x_{m+1})을 만족하는지, m 값을 워크로드별로 튜닝했는지다. 저자들이 밝힌 전제와 한계도 분명하다. 확률적 탐침의 샘플링·검증 오버헤드 때문에 속도 이득이 τ 이득보다 작고, m≥3에서는 이득이 사실상 사라지며, m=0은 구조가 확률적으로 크게 흔들릴 수 있다. 또한 실험은 EAGLE-3 체크포인트와 T=1.0 기본 설정, 트리 크기 60·깊이 8이라는 특정 구성에 기반하므로 다른 드래프트 모델이나 트리 예산에서는 재검증이 필요하다.