Cactus: 제약된 수용 추론 샘플링을 통한 자기회귀 디코딩 가속화

초록

추측적 샘플링(SpS)은 작은 초안 모델을 활용하여 자기회귀적 대규모 언어 모델의 디코딩 처리량 가속화에 성공해왔습니다. SpS는 생성된 분포가 검증기 LLM의 분포와 정확히 일치하도록 엄격하게 강제합니다. 이는 검증기의 분포를 약간 변형하는(top-k나 temperature를 적용한 샘플링 등) 경우에도 허용될 수 있음을 고려할 때 불필요하게 제한적입니다. 일반적 수용 샘플링(TAS)은 엔트로피 기반 휴리스틱을 사용해 더 많은 토큰을 수용함으로써 이 문제를 완화합니다. 그러나 이 방식은 검증기 분포를 왜곡시켜, 검증기가 중요한 정보를 인코딩한 경우 출력 품질이 저하될 가능성이 있습니다. 본 연구에서는 제약 최적화 관점에서 추측적 샘플링 알고리즘을 정형화합니다. 이 정형화를 바탕으로 검증기 분포로부터의 제어된 divergence와 증가된 수용률을 보장하는 방법인 Cactus(제약 수용 추측적 샘플링)를 제안합니다. 다양한 벤치마크에서의 실험 결과는 우리 접근법의 효과성을 입증합니다.

English

Speculative sampling (SpS) has been successful in accelerating the decoding throughput of auto-regressive large language models by leveraging smaller draft models. SpS strictly enforces the generated distribution to match that of the verifier LLM. This is unnecessarily restrictive as slight variations of the verifier's distribution, such as sampling with top-k or temperature, would also be acceptable. Typical acceptance sampling (TAS) alleviates this issue by accepting more tokens using entropy-based heuristics. However, this approach distorts the verifier distribution, potentially degrading output quality when the verifier encodes critical information. In this work, we formalize the speculative sampling algorithm through the lens of constrained optimization. Based on this formulation, we propose Cactus (constrained acceptance speculative sampling), a method that guarantees controlled divergence from the verifier distribution and increasing acceptance rates. Empirical results across a wide range of benchmarks confirm the effectiveness of our approach.

Cactus: 제약된 수용 추론 샘플링을 통한 자기회귀 디코딩 가속화

Cactus: Accelerating Auto-Regressive Decoding with Constrained Acceptance Speculative Sampling

초록

Support