[논문 리뷰] Fast Inference from Transformers via Speculative Decoding

https://arxiv.org/abs/2211.17192
Yaniv Leviathan, Matan Kalman, Yossi MatiasICML
2022. 11. 30.

 

논문 요약

• 대규모 autoregressive 모델, 특히 Transformer로부터의 추론은 K개의 토큰을 디코딩하기 위해 K번의 직렬 모델 실행이 필요하므로 느립니다.
• 추론 속도를 높이기 위한 기존 방법들은 종종 아키텍처 변경, 재훈련 또는 출력 분포 변경을 요구하여 동일한 결과를 유지하지 못합니다.
• 이 연구는 병렬 컴퓨테이션을 활용하여 모델 아키텍처, 훈련 절차 또는 출력 분포를 수정하지 않고 autoregressive model inference를 가속화하는 것을 목표로 합니다.

 

LLM 추론의 병목 현상

https://huggingface.co/blog/tngtech/llm-performance-prefill-decode-concurrent-requests

Transformer의 Inference 동작 방식은 Prefill과 Decoding, 2가지로 나눌 수 있음

1) Prefill 단계는 초기 입력 시퀀스인 input prompt 전체를 한 번에 모델에 넣고, 이에 대한 attention 계산을 수행한다. 그리고, KV Cache 기능을 사용한다면, 이를 생성

2) Decoding 단계는 Prefill 단계 이후, 하나씩 토큰을 생성해서 output을 만들어내는 단계 이 단계는 autoregressive 하기 때문에, 이전에 생성한 토큰들에 의존해서 한 토큰씩 반복적으로 수행

 

1. 모델 크기 증가: 시간이 지남에 따라 모델의 파라미터 크기가 기하급수적으로 커지고 있으며, 1조 개 이상의 파라미터를 가진 모델도 등장하고 있음
2. 트랜스포머 구조의 한계: 현재 표준화된 트랜스포머 아키텍처는 오토리그레시브(Autoregressive) 모델 구조의 한계를 그대로 가짐
3. 순차적 토큰 처리: 다음 토큰의 출력이 다음 입력으로 들어가 다음 출력이 생성되는 구조로 인해, 토큰 처리 과정 자체가 순차적으로 진행되어야 하는 제약이 있음
4. 비효율성: 출력 길이의 지속적인 증가와 맞물려 수십억 파라미터 모델이 수만 번 포워드(순차 처리)되어야 하므로 매우 비효율적임
5. 부분적 개선: KV 캐싱 등 여러 방법론이 적용되어 구동 자체가 토큰 길이만큼 걸리는 것은 아니지만, 근본적인 순차 처리 문제는 남아있음

Speculative Decoding의 기본 아이디어

https://www.youtube.com/watch?v=9MWt1-DjYlg

1. 목표: 모델이 커지고 계산량이 한정되어 있을 때, 출력을 더 빨리 달성하기 위한 방법론을 제시함.
2. 핵심 방법론: 모델 두 개를 활용하는 재밌는 방법론을 주장함
3. Speculative Execution 개념 차용: 컴퓨터 과학에서 통용되는 개념인 Speculative Execution에서 아이디어를 얻음. 이는 각 작업이 실제로 필요한지 검증한 뒤 검증된 작업만 실행하여 자원이나 시간을 아끼는 방식임
4. Speculative Decoding의 정의: 소형 모델 A와 거대 모델 B를 동시에 사용하여, A의 출력을 B가 검증하는 방식으로 사용함
5. 결과 활용: 검증을 통과하면 A의 출력을 그대로 사용함

1. 토큰 출력 비교: 초록색 토큰은 작은 모델 A가 이번 주기에 뱉어낸 것이고, 파란색 토큰은 타겟 모델 B (거대 모델)가 뱉어낸 출력임
2. 거절된 토큰: 붉은색 토큰은 거대 모델 B가 거절하여 취소된 토큰임
3. 검증 주기: 일정한 주기 γ 또는 초록색 토큰의 한 번에 나올 길이에 따라 검증 프로세스가 진행됨을 이해하는 것이 중요함

확률 기반 검증 원리

위 알고리즘을 쉽게 설명하자면

1. 모델 정의: 가벼운 모델 A (Draft Model)와 무거운 모델 B (Target Model)가 존재함. 현재 논의는 하나의 출력 토큰 X_T와 ~X_T에 대한 내용임

1. 확률 정의: A가 X_T 토큰에 가지는 확률을 Q(X), B가 동일한 단어에 가지는 확률을 P(X)라고 정의함.
2. 승인/거절 기준:
   1. P(X) > Q(X)일 경우: 정답 모델 B의 확신도가 더 높으므로 승인함
   2. 그렇지 않은 경우: Q(X)와 P(X)의 확률로 승인함 (실제로는 거절에 가까움)
3. 시행 시점: 이 과정은 앞에서 언급한 검증 주기 r(gamma)마다 시행됨

Speculative Sampling의 조작 및 기대값

https://www.youtube.com/watch?v=v5al_cwvkJQ

1. 추가된 조작: Speculative Decoding의 Speculative Sampling은 Rejection Sampling과 거의 동일하나, 하나의 간단한 조작이 추가됨
2. 조작 내용: Q에서 뽑은 샘플이 T 내부라면 승인하고, 아니라면 거절하며 맞는 분포로 이동한다는 조작이 추가됨
3. 기대값: 이 조작 덕분에 최종 기대값은 목표하는 분포 P와 거의 동일한 분포를 기대할 수 있음

Batched Decoding

1. 의문점: N개의 토큰을 비교하려면 N번의 출력도 필요한데, 순차 처리의 한계를 그대로 가지는 것이 아닌가 하는 의문이 생길 수 있음. (오히려 드래프트 모델만 N번 더 실행하는 것이 아닌가?)
2. 병렬 처리 가능성:
   • 오토리그레시브 프로세스에서, 이미 알고 있는 입력 프리픽스(A_1, A_2, …)를 모두 가지고 있다면, 빅모델 B에 들어갈 입력(프리픽스)을 모두 가지고 있다는 뜻임
   • 따라서 이 작업은 병렬적으로 처리가 가능하며 출력 또한 한 번에 나오게 됨
   • 시퀀셜 실행이 아니게 된다는 것이 중요함
3. 가속 시점: 병렬화가 가능하므로 B_1부터 B_r까지 한 번에 배치 디코딩으로 처리 가능하며, 이 지점이 가장 큰 가속이 일어나는 시점임
4. 오류 처리: 만약 앞부분이 틀렸다면 뒤는 모두 버리고 해당 부분부터 다시 시작함

진짜 빠른걸까?

외부에서 봤을 때 Draft Model의 도입이 오히려 모델 추론을 2번 하게 되는 셈이기 때문에 더 느려지는거 아닌가?

논문에서는 이를 Wall-time Improvement 파트에서 설명하고 있음

 

결론부터 말하자면, 병럴 처리 자원이 충분하다는 전제 하에,

• Main Model의 Decoding 횟수는 감소하고,
• 도입 되는 Draft Model 추론이 굉장히 빠르기 때문이다.

peculative Decoding에서 성능에 영향을 줄 수 있는 파라미터는 2가지가 있음

1. α 값은 Draft Model이 Main Model을 얼마나 잘 근사하는가를 나타낸다. 당연히 Draft Model의 근사 분포가 Main Model과 거의 유사하다면, 발생 할 다음 토큰의 확률이 거의 Main Model과 유사해지기 때문에, 거절 되는 토큰의 수가 줄어들 것이다.
2. γ 값은 Draft Model이 생성하는 후보 토큰의 수를 의미한다. Draft Model이 한 번에 고려 할 수 있는 후보 토큰을 많이 만들수록 Main Model은 기대 할 수 있는 토큰의 수도 증가한다.

이 파라미터를 토대로 논문에서 유도한 Speculative Decoding을 적용 했을 때의 Main Model의 호출 횟수의 배율을 아래와 같은 수식으로 작성 할 수 있음

결국 Draft Model이 Main Model을 더 잘 근사하고, 한 번에 많은 후보 토큰을 생산 할수록 Main Model의 호출은 저 수식 값의 배율만큼 호출 횟수가 줄어듬

 

 

가속 효과 및 하이퍼파라미터

1. 가속 시각화: 거대 모델만 사용할 때보다, 정해진 길이 γ 마다 큰 모델을 병렬적으로 호출하여 검증하는 방식이 γ 길이에 따라 직관적으로 계산 시간이 줄어드는 것을 확인할 수 있음.
2. 하이퍼파라미터 α: 검증할 길이 r와 드래프트 모델이 얼마나 승인되는가를 고려하여 정해지는 승인율의 기대값을 α라고 부름.
3. 1. α에 영향을 미치는 요소:
   • 검증을 수행할 길이 γ
   • 작은 모델이 큰 모델과 비교할 때 얼마나 작은가 (시간 차이)
   • 작은 모델 자체의 성능
   • 작은 모델과 큰 모델의 아키텍처 유사도.
4. 알고리즘 효율성 지표: α가 Speculative Decoding 알고리즘 효율성을 나타내는 좋은 지표가 됨

모델 크기에 따른 실험 결과

1. 모델 선택: 논문에서 작은 모델(드래프트 모델)은 기본적으로 타겟 모델에 비해 100배 이상 작은 모델을 선택함 (동일 T5 모델의 라지/베이스/스몰 기준).
2. 드래프트 모델 크기 vs 가속량: 드래프트 모델 크기가 작아질수록 가속량은 늘어나지만, 승인율이 줄어들어 검증한 곳을 다시 만들어야 하는 상황이 자주 연출됨.
3. 승인율과 가속:
   • 가벼운 모델은 α값(승인율)이 높지 않았음.
   • 큰 모델은 승인율 자체가 높아서 이미 지나온 것을 다시 하지 않아도 되므로 가속이 되었음.
4. 트레이드오프: 이 경우 트레이드오프가 존재함을 알 수 있음.
5. Temperature 영향: Temperature 설정에 따라 출력 확률 분포의 소프트맥스 처리가 달라지며, Temperature가 0일 때는 큰 모델과 작은 모델이 유사한 값을 내어 가속이 조금 더 됨.

   

참고자료

[1] https://medium.com/@TitanML/in-the-fast-lane-speculative-decoding-10x-larger-model-no-extra-cost-f33ea39d065a

[2] https://velog.io/@with1015/ICML23-Fast-Inference-from-Transformers-via-Speculative-Decoding

[3] https://www.youtube.com/watch?v=9MWt1-DjYlg

[4] https://www.youtube.com/watch?v=v5al_cwvkJQ