Speculative Decoding: 작은 모델이 먼저 쓰면 큰 모델도 빨라진다

채팅창에 답이 한 글자씩 나타나는 모습을 보고 있으면 질문이 생긴다. GPU는 수많은 계산을 동시에 하는데, 왜 문장은 꼭 한 토큰씩 기다려야 할까. 앞 토큰이 정해져야 다음 토큰의 분포를 계산할 수 있기 때문이다. Speculative Decoding은 이 의존성을 무시하지 않는다. 대신 값싼 모델이 먼저 짐작한 미래를 비싼 모델이 한꺼번에 검토하게 한다. 여기까지는 그럴듯한 속도 개선이다. 더 놀라운 것은 제대로 구현하면 최종 생성 분포까지 그대로 유지할 수 있다는 점이다.

초안을 쓰는 일과 승인하는 일

목표 모델을 P, 작은 초안 모델을 Q라고 부르자. Q는 현재 문맥 뒤에 붙을 토큰 몇 개를 자기회귀적으로 생성한다. P는 그 후보 토큰들을 입력으로 받아 각 위치에서의 다음 토큰 분포를 병렬로 계산한다. 후보가 이미 주어졌으므로, P가 후보를 평가하는 동안 각 위치의 입력은 확정되어 있다.

이 계산은 미래를 확실히 안다는 뜻이 아니다. 초안 앞부분이 채택되었을 때만 그 뒤의 계산이 유효하다. 첫 번째 거절이 나오면 그 뒤 후보들은 버린다. 앞에서부터 채택 가능한 접두부를 남기고, 거절된 위치의 토큰을 보정해 다시 진행한다. 모두 채택되면 목표 모델이 계산해 둔 분포에서 추가 토큰 하나를 얻을 수 있다. 원문 알고리즘

내가 이 방식을 설명한다면 검토자와 초안 작성자의 업무를 떠올리겠다. 검토자가 빈 종이를 받아 매번 한 단어를 정해 주는 대신, 작성자가 준비한 몇 단어를 한 번에 살핀다. 다만 실제 알고리즘에서 검토는 의미가 적절한지 판단하는 자연어 평가가 아니라 토큰 확률을 사용한 계산이다. 이 차이를 놓치면 '큰 모델이 작은 모델의 답을 믿는다'는 부정확한 설명으로 흘러가기 쉽다.

같은 분포를 지키는 작은 장치

한 위치에서 후보 토큰 x의 목표 확률을 p(x), 초안 확률을 q(x)라고 하자. Q에서 뽑힌 x는 min(1, p(x)/q(x))의 확률로 채택한다. 거절되면 max(0, p−q)를 정규화한 잔여 분포에서 새 토큰을 뽑는다. 작은 모델이 지나치게 많이 제안하는 토큰은 일부 거절하고, 덜 제안하는 토큰은 잔여 분포로 보충하는 구조다.

손으로 확인할 수 있는 예를 만들어 보자. 토큰 A와 B만 있고, P는 A를 60%, B를 40%로 생성한다고 하자. Q는 A를 80%, B를 20%로 생성한다. Q가 A를 제안하면 채택 확률은 0.6/0.8, 즉 75%다. Q가 B를 제안하면 목표 확률이 더 크므로 항상 채택한다.

따라서 전체 시행 중 A가 곧바로 채택되는 비율은 0.8×0.75=0.6이다. B가 바로 채택되는 비율은 0.2다. 남은 0.2는 A 제안이 거절된 경우다. 잔여량은 A에서 0, B에서 0.2이므로, 이 경우 B를 뽑는다. 최종 확률은 A 0.6, B 0.4가 된다. 작은 모델의 편향이 최종 출력에 그대로 섞이지 않는 이유가 숫자 두 개로 드러난다. 저자들의 원리 설명

이 보정이 빠지면 이야기가 달라진다. 큰 모델이 보기에 '충분히 괜찮은' 후보를 임의의 임계값으로 받아들이는 것은 목표 분포를 정확히 보존하는 알고리즘과 같지 않다. 온도나 top-p 같은 샘플링 설정을 사용한다면, 알고리즘이 다루는 p와 q가 그 설정이 반영된 분포인지도 확인해야 한다.

품질을 내려서 얻은 속도가 아니다

분포 보존은 매우 강한 성질이지만, 그 의미를 과장해서도 안 된다. 목표 모델이 원래 가진 오류나 환각은 그대로 남는다. 작은 모델의 도움으로 더 정확한 답을 만드는 기법이 아니라, 같은 목표 모델의 샘플링을 더 빨리 수행하는 기법이다. 확률적 생성에서 같은 분포라는 말은 같은 랜덤 시드를 넣으면 반드시 같은 문자열이 나온다는 말과도 구별된다. 구현에 따라 난수 사용 순서가 달라질 수 있다.

또한 '샘플 하나를 빨리 얻는다'와 '총 계산을 줄인다'가 항상 같은 것은 아니다. 초안 생성에도 비용이 들고, 거절된 후보를 확인한 계산은 버려진다. 논문이 주목한 상황은 비싼 모델의 디코딩이 메모리 대역폭 등의 제약을 받고, 여러 위치를 함께 평가할 연산 여지가 있는 상황이다. 남는 병렬 계산 능력을 사용해 기다림을 줄이는 셈이다.

나는 이것을 효율의 기준이 실행 목적에 따라 달라지는 사례로 읽었다. 대화 응답은 사용자가 한 번의 답을 얼마나 오래 기다리는지가 중요하고, 대규모 배치 작업은 같은 장치에서 하루에 얼마나 처리하는지가 중요하다. 같은 알고리즘도 어느 지표를 최적화하는지에 따라 유용성이 달라질 수 있다. 이것은 원문 성능 주장과 별개로 실무에서 고려할 관점이다.

초안 모델은 작기만 하면 될까

작은 모델이 빠르면 좋지만, 목표 모델과 자주 의견이 달라지면 대부분의 초안이 거절된다. 반대로 초안 모델을 크게 만들어 두 분포를 가깝게 하면 검토 전에 기다리는 시간이 늘 수 있다. 중요한 것은 모델 크기 하나가 아니라 초안의 비용과 채택률의 조합이다.

비유적으로, 아주 빠르지만 매번 첫 단어부터 틀리는 작성자는 도움이 적다. 매우 정확해도 검토자만큼 오래 걸리는 작성자 역시 이점이 작다. 목표 모델과 비슷한 토큰 분포를 싸게 제안할 수 있는 조합이 필요하다. 과제에 따라 익숙한 문구나 반복 입력을 잘 복사하는 간단한 제안 방식도 후보가 될 수 있다. 논문은 초안 생성 방법을 반드시 별도의 거대한 신경망으로 한정하지 않는다.

한 번에 몇 토큰을 제안할지도 비슷한 선택이다. 더 길게 제안하면 전부 채택될 때 이득이 커지지만, 초안 비용과 거절 뒤의 낭비도 커진다. 채택률은 고정된 물성처럼 보지 않는 편이 좋다. 내가 테스트한다면 한국어 대화, 코드 완성, 문서 요약을 나누어 보겠다. 목표와 초안이 어느 구간에서 잘 맞는지 작업마다 달라질 수 있기 때문이다.

논문 결과에서 가져갈 것

원문은 T5-XXL의 번역과 요약을 당시 T5X 구현과 비교해 약 2~3배의 실행 시간 개선을 보고한다. 이 숫자는 그 모델, 장치, 구현, 과제에서 얻은 결과다. 모든 최신 모델이나 이미 큰 배치를 처리하는 서버에 같은 배수가 나오리라는 보장은 없다. ICML 2023 논문

실제 도입에서는 목표 모델 단독 생성과 같은 샘플링 조건을 맞추고, 초안 생성 시간과 검증 시간을 모두 포함해야 한다. 추가 모델이 차지하는 메모리, 캐시 관리, 배치의 변화도 측정에 들어간다. 짧은 답에서 체감되는 이득과 긴 답에서의 평균 토큰 지연을 나누어 보고 싶다. 알고리즘의 분포 보존 조건이 지켜지는지와 제품의 속도 이득이 있는지는 각각 확인할 문제다.

이 논문이 흥미로운 이유는 미래 토큰의 의존성을 없애지 않고도 병렬화할 수 있다는 데 있다. 모든 미래를 먼저 맞힐 필요는 없다. 맞힌 접두부만 안전하게 재사용하고, 틀린 부분은 정확히 보정할 수 있다면 추측도 엄밀한 계산의 일부가 된다.

읽은 논문: Yaniv Leviathan, Matan Kalman, Yossi Matias, Fast Inference from Transformers via Speculative Decoding, ICML 2023, 사전 공개 2022. 논문 · 저자 해설