[논문 리뷰] SpecTr: Fast Speculative Decoding via Optimal Transport

Ziteng Sun, Ananda Theertha Suresh외 4명|arXiv (Cornell University)|2023. 10. 23.

Natural Language Processing Techniques인용 수 4

한 줄 요약

SpecTr는 최적 운반 이론을 사용하여 고품질의 드래フト 토큰을 선택하는 새로운 사전 추론 방법을 제안한다. 이는 품질 저하가 없이 더 빠른 자동회귀 생성을 가능하게 한다. 최신 언어 모델에서 기준 자동회귀 추론 대비 2.13배의 월클록 속도 향상을 달성했으며, 표준 사전 추론 대비 1.37배 향상되었다.

ABSTRACT

Autoregressive sampling from large language models has led to state-of-the-art results in several natural language tasks. However, autoregressive sampling generates tokens one at a time making it slow, and even prohibitive in certain tasks. One way to speed up sampling is $ extit{speculative decoding}$: use a small model to sample a $ extit{draft}$ (block or sequence of tokens), and then score all tokens in the draft by the large language model in parallel. A subset of the tokens in the draft are accepted (and the rest rejected) based on a statistical method to guarantee that the final output follows the distribution of the large model. In this work, we provide a principled understanding of speculative decoding through the lens of optimal transport (OT) with $ extit{membership cost}$. This framework can be viewed as an extension of the well-known $ extit{maximal-coupling}$ problem. This new formulation enables us to generalize the speculative decoding method to allow for a set of $k$ candidates at the token-level, which leads to an improved optimal membership cost. We show that the optimal draft selection algorithm (transport plan) can be computed via linear programming, whose best-known runtime is exponential in $k$. We then propose a valid draft selection algorithm whose acceptance probability is $(1-1/e)$-optimal multiplicatively. Moreover, it can be computed in time almost linear with size of domain of a single token. Using this $new draft selection$ algorithm, we develop a new autoregressive sampling algorithm called $ extit{SpecTr}$, which provides speedup in decoding while ensuring that there is no quality degradation in the decoded output. We experimentally demonstrate that for state-of-the-art large language models, the proposed approach achieves a wall clock speedup of 2.13X, a further 1.37X speedup over speculative decoding on standard benchmarks.

연구 동기 및 목표

• 멤버십 비용을 고려한 최적 운반 이론을 활용한 사전 추론의 원리적인 프레임워크를 개발하기 위해.
• 드래프트 토큰을 k개로 일반화하여 토큰 수준에서의 사전 추론을 가능하게 하여 드래프트 품질과 수락 확률을 향상시키기 위해.
• 각 토큰당 거의 선형 시간 내에 실행되는 (1−1/e)-최적 수용 확률을 가지는 드래프트 선택 알고리즘을 설계하기 위해.
• 최종 출력 분포가 대규모 언어 모델의 분포와 일치하도록 보장하여 품질 저하가 없음을 보장하기 위해.
• 최신 모델에서의 실증적 검증을 통해 실제 추론 환경에서 뚜렷한 속도 향상을 입증하기 위해.

제안 방법

• 멤버십 비용을 고려한 최적 운반 문제로 사전 추론를 공식화하여 최대 결합 프레임워크를 확장한다.
• 수용 확률을 최대화하기 위해 운반 계획을 최적화하는 k-토큰 드래프트 선택 전략을 도입한다.
• 선형 프로그래밍을 통해 최적의 드래프트 선택을 해결하지만, (1−1/e)-최적성을 확보하는 실용적 근사치를 제안한다.
• 각 토큰당 어휘 크기와 거의 선형적인 시간 내에 드래프트 선택을 계산할 수 있는 빠른 알고리즘을 설계한다.
• 대규모 모델이 드래프트 토큰을 병렬로 평가할 수 있도록 허용하는 새로운 자동회귀 샘플링 알고리즘인 SpecTr에 드래프트 선택을 통합한다.
• 통계적 수락 메커니즘을 사용하여 최종 출력이 대규모 모델의 진짜 분포를 따르도록 보장한다.

실험 결과

연구 질문

• RQ1 최적 운반 이론을 사용하여 사전 추론의 원리적인 프레임워크를 제공할 수 있는가?
• RQ2 사전 추론을 k-토큰 드래프트를 허용하는 방식으로 일반화하여 수용 확률을 향상시킬 수 있는가?
• RQ3 최적 운반 공식화 하에서 드래프트 선택 품질의 이론적 한계는 무엇인가?
• RQ4 각 토큰당 거의 선형 시간 내에 근사 최적의 드래프트 선택 알고리즘을 효율적으로 계산할 수 있는가?
• RQ5 제안된 방법은 출력 품질 저하 없이 실질적으로 빠른 속도 향상을 달성하는가?

주요 결과

• LM1B 벤치마크에서 대규모 모델로 PALM-2-Bison을 사용할 때 SpecTr는 기준 자동회귀 추론 대비 2.13배의 월클록 속도 향상을 달성했다.
• L=8 및 K=8 조건에서 SpecTr는 4.0 블록 효율성을 기록했으며, K=1인 표준 사전 추론 대비 뚜렷한 우수성을 보였다. 표준 사전 추론은 블록 효율성이 2.9에 불과했다.
• K=1 및 L=8 조건에서 SpecTr는 표준 사전 추론 대비 1.37배의 추가 속도 향상을 기록했다.
• 제안된 드래프트 선택 알고리즘은 (1−1/e)-최적 수용 확률을 달성했으며, 이러한 방법들 중에서 알려진 최고의 곱셈 보장이다.
• 실증 결과는 SpecTr의 출력 분포가 대규모 모델의 분포와 일치함을 확인했으며, 이는 품질 저하가 없음을 보장한다.
• 이론적 성과는 구현 오버헤드와 상당한 소규모 모델 추론 시간이 존재하더라도 실질적으로 대부분 유지된다.