Accelerating Large Language Model Decoding with Speculative Sampling

https://arxiv.org/abs/2302.01318

글 작성일 : 2024.01

내용 정리 + 간단한 설명(필자가 이해한 내용) 으로 구성함

아래 논문을 참고하면 좋음

Fast Inference from Transformers via Speculative Decoding

Introduction

Transformer model의 sampling의 경우 대부분 memory bandwidth에 bound되어 있다. 따라서 하나의 토큰을 만드는 시간은 모델 파라미터 수와 transformer 메모리 크기에 비례한다. 거기에 더해 auto-regressive한 transformer 구조의 특성상, 새로운 토큰을 생성하기 위해선 foward path를 한 번 거처야 하기 때문에 더욱 많은 시간이 소요된다.

그래서 논문의 저자들은 latency를 줄이기 위해 Sepculative sampling(SpS)을 제안했다. 이는 다음과 같은 단계로 구성된다.

  1. 작은 Draft model이 auto regressive하게 k개의 토큰을 만든다. ( draft model을 k번 돌림)
  2. 크고 강력한(powerful) Target model이 draft model의 output을 scoring 한다.
  3. Rejection sampling scheme에 따라 draft model output중 일부를 채택(accept)하고 target model의 distribution을 회복(recover)한다.

이 같은 sampling 기법은 Target model의 변형이나 확률 분포의 변형 없이 latency를 크게 향상시킬 수 있다.

Chinchilla를 사용한 실험에서 SpS가 일반적인 샘플링 기법보다 2-2.5배 speed up을 기록했다. 그리고 이는 memory bandwith 한계로 인한 이론적인 auto-regressive sampling의 한계보다 더 높은 기록이다.

Related Work