TreeFlash: 더 빠른 추측적 디코딩 (Speculative Decoding)을 위한 병렬 AR 근사 방식

추측적 디코딩 (Speculative Decoding)을 위한 원샷 블록 드래프터 (One-shot block drafters)는 단일 순전파 (forward pass) 과정에서 전체 초안 (draft)을 생성하며, 순차적인 토큰 생성 과정을 제거함으로써 강력한 처리량 (throughput)을 달성합니다. 그러나 이들은 각 초안 토큰을 오직 접두사 문맥 (prefix context)에만 의존하여 예측하며, 이전에 생성된 초안 토큰들과의 의존성은 고려하지 않습니다. 이러한 비자기회귀적 (non-autoregressive) 조건화는 초안의 깊이 (draft depth)가 깊어질수록 드래프터의 분포가 검증기 (verifier)의 실제 자기회귀적 (autoregressive) 분포로부터 벗어나게 만듭니다. 이 문제는 서로 다른 가지 (branches)들이 후속 토큰에 대해 동일한 주변 분포 (marginal distribution)를 공유하도록 강제되는 트리 기반 드래프팅 (tree-based drafting)에서 더욱 심각해집니다. 우리는 드래프터의 은닉 상태 (hidden state)와 이전 토큰에 조건화된 MLP 레이어를 통합하여 자기회귀적 분포를 근사함으로써 이 문제를 해결하는 TreeFlash를 제안합니다. TreeFlash는 2단계 근사 메커니즘 (two-stage approximation mechanism)을 채택하여 원샷 드래프터의 $\mathcal{O}(1)$ 디코딩 시간 복잡도를 유지합니다. TreeFlash는 다양한 작업과 모델에 대해 최첨단 (state-of-the-art) 성능을 달성하였으며, 주변 트리 드래프팅 (marginal tree drafting) 대비 블록 효율성 (block efficiency)은 $12%$ 높고, 속도 향상 (speedup)은 $9%$ 더 높게 개선되었습니다.