추측 디코딩으로 LLM 추론을 2~3배 가속하기 — 정확도 손실 없이 속도를 높이는 방법

핵심 메시지

• 거대 언어 모델은 토큰 하나를 만들 때마다 전체 모델을 한 번씩 통과해야 해서 느리다.
• 추측 디코딩은 작고 빠른 '초안 모델'이 여러 토큰을 먼저 제안하고, 크고 정확한 '목표 모델'이 이를 한 번의 연산으로 병렬 검증하는 방식이다.
• 거부 샘플링이라는 수학적 장치 덕분에 속도를 높이면서도 큰 모델만 썼을 때와 정확도가 동일하다.
• 딥마인드는 2~2.5배, 구글은 2~3.4배 속도 향상을 보고했고, 한 번에 제안하는 토큰 수 K는 3~7 사이를 권장한다.
• 최악의 경우에도 큰 모델 연산 한 번마다 최소 1개 토큰은 생성되므로 기존 방식보다 느려지지 않는다.

쉽게 이해하기

트랜스포머 기반 언어 모델은 토큰을 하나씩 차례로 생성한다. 토큰 하나를 만들 때마다 모델 전체를 한 번씩 통과해야 하는데, 1000억 개 파라미터 모델이라면 토큰 하나에 약 1000억 번의 연산이 든다. 문제는 'University of' 다음에 올 쉬운 단어처럼 굳이 거대한 모델을 쓰지 않아도 되는 토큰까지 똑같이 비싼 연산을 반복한다는 점이다.

추측 디코딩의 첫 번째 아이디어는 쉬운 토큰은 훨씬 작은 모델에 맡기자는 것이다. 두 번째 아이디어는 트랜스포머의 구조적 특성에서 나온다. 트랜스포머는 토큰 하나를 생성하는 것과 비슷한 비용으로, 여러 개의 토큰을 한 번의 연산(forward pass)으로 한꺼번에 검증할 수 있다.

알고리즘은 이렇게 작동한다. 먼저 작은 '초안 모델'을 자기회귀적으로 K번 돌려 K개의 토큰을 만든다. 그다음 큰 '목표 모델'을 단 한 번 병렬로 돌려 그 토큰들에 대한 확률 분포를 동시에 얻는다. 이때 토큰은 항상 초안 모델에서 샘플링하고, 목표 모델은 검증용 확률 분포만 만든다.

핵심은 '거부 샘플링'으로 어떤 토큰을 받아들일지 정하는 단계다. 토큰을 앞에서부터 검사해, 목표 모델의 확률 Q가 초안 모델의 확률 P보다 크면 수용하고, 작으면 Q/P 확률로 수용 여부를 정한다. 한 토큰이 거부되면 그 뒤의 토큰은 모두 버리고, 거부된 자리에서는 (Q-P)의 양수 부분 분포로 토큰 하나를 다시 뽑는다. 이 절차를 거치면 최종적으로 뽑히는 토큰의 분포가 목표 모델 분포 Q와 정확히 같아져 정확도 손실이 없다.

실험 결과 딥마인드 논문은 K를 3~4로 두고 자기회귀 방식 대비 약 2~2.5배, 구글 논문은 K를 3~7로 두고 2~3.4배의 속도 향상을 보고했다. 두 결과 모두 정확도는 그대로였다. K를 키우면 더 빨라지지만 수확체감이 나타난다.

주요 인사이트

• 구글과 딥마인드가 사실상 같은 아이디어를 거의 동시에 독립적으로 발표했고, 두 알고리즘에 큰 차이가 없다.
• 큰 모델은 토큰을 '생성'하는 데 쓰이지 않고 '검증'에만 쓰여, 한 번의 연산으로 여러 토큰을 한꺼번에 확인한다.
• 모든 토큰은 초안 모델에서 샘플링되지만, 거부 샘플링 보정 덕분에 최종 분포는 목표 모델과 동일하게 맞춰진다.
• 최악의 경우에도 토큰이 0개 생성되는 일이 없어 알고리즘이 멈추거나 기존보다 느려지지 않는다.
• 한 번에 제안하는 토큰 수 K를 키우면 속도는 오르지만 어느 지점부터 수확체감이 나타난다.

자주 묻는 질문

관련 AI 소식