Looking back at speculative decoding

https://research.google/blog/looking-back-at-speculative-decoding/

Looking back at speculative decoding

 

요약

추측 디코딩은 대규모 언어 모델(LLM)의 추론 속도를 대폭 향상시키면서도 동일한 품질의 출력을 보장하는 혁신적인 알고리즘입니다. 이 방법은 추측 실행(speculative execution)에서 영감을 받아, 작은 모델을 사용해 빠르게 예측한 결과를 대규모 모델과 병렬로 처리하여 병렬성을 극대화합니다.

 

이 기술은 GPT와 같은 대규모 트랜스포머 기반 모델에서 번역, 요약, 이미지 및 음성 생성 등 다양한 작업에서 2~3배 이상의 성능 향상을 제공하며, 분산 환경에서도 그 효과가 입증되었습니다. 또한, 작은 모델을 활용해 쉬운 토큰을 빠르게 처리하고, 어려운 토큰은 큰 모델로 처리하는 방식으로 효율성을 극대화합니다.

 

추측 디코딩은 현재 Google Search를 포함한 다양한 대규모 AI 제품에 적용되어, 추론 속도를 크게 개선하면서도 에너지 소비를 줄이는 데 기여하고 있습니다. AI의 확산과 함께 효율적인 추론에 대한 수요가 증가하는 가운데, 이 기술은 산업 전반에서 혁신적인 활용 사례와 결합되며 널리 채택되고 있습니다.

 

추측 디코딩은 AI 추론의 미래를 열어갈 중요한 기술로, 앞으로 더욱 다양한 응용과 새로운 접근법이 등장할 것으로 기대됩니다.

 

------------------

 

추측 디코딩(Speculative Decoding)은 품질을 저하시키지 않으면서 LLM의 추론 속도를 높이고 비용을 절감하는 효과적인 기술로 입증되었습니다. 또한, 다양한 최적화 기술의 패러다임으로서도 효과적임이 입증되었습니다.

 

대규모 언어 모델(LLMs)은 최근 인공지능(AI)의 급격한 발전의 중심에 있습니다. 이러한 모델들은 혁신적이지만, 크기 때문에 사용자 중심 제품에서는 추론(즉, 출력 생성)이 느려져 사용자 경험이 저하될 수 있는 문제가 있습니다.

 

2022년, 우리는 "추측 디코딩을 통한 트랜스포머의 빠른 추론(Fast Inference from Transformers via Speculative Decoding)"이라는 논문을 발표하며, 추론 시간을 대폭 줄일 수 있는 기술인 **추측 디코딩(Speculative Decoding)**을 소개했습니다. 이 알고리즘은 자기회귀 모델에서 여러 토큰을 병렬로 계산하여 생성 속도를 높이며, 출력 품질에는 영향을 미치지 않습니다. 사실, 이 방법은 동일한 출력 분포를 보장합니다. 동일한 하드웨어에서 더 빠르게 결과를 생성할 수 있다는 것은 동일한 트래픽을 처리하는 데 필요한 기계의 수가 줄어듦을 의미하며, 이는 모델 서비스를 위한 에너지 비용 절감으로 이어집니다.

 

오늘날 우리는 이 방법이 어떻게 작동하는지 다시 살펴보고, 추측 실행(speculative execution)에서 영감을 받은 점을 탐구하며, 새로운 통찰과 발전을 조명하려 합니다. 초기 논문에서는 번역 및 요약 작업에 이 접근 방식을 적용하여 약 2배에서 3배의 성능 향상을 보여주었습니다. 이후, 우리는 여러 Google 제품에 추측 디코딩을 적용하여 추론 속도를 크게 개선하면서도 응답 품질을 유지했습니다. 또한, 산업 전반에서 추측 디코딩이 채택되었으며, 이를 활용한 효과적이고 통찰력 있는 응용 사례와 아이디어를 많이 확인할 수 있었습니다.

 

https://storage.googleapis.com/gweb-research2023-media/media/SpeculativeDecoding-0-AIO.mp4

 

추측 디코딩(Speculative Decoding)은 Google 검색에서 AI 개요(AI Overviews)가 이전보다 더 빠르게 결과를 생성하면서도 동일한 응답 품질을 유지할 수 있도록 지원합니다. AI 개요와 같은 대규모 제품들은 지속적으로 새로운 기술로 최적화되고 있지만, 추측 디코딩은 여전히 최적화의 중요한 요소로 자리 잡고 있습니다.

 

배경

대규모 언어 모델(LLM)은 보통 하나의 토큰씩 출력을 생성합니다. 여기서 토큰은 단어 또는 단어의 일부를 의미합니다. 예를 들어, 일반적인 토크나이저를 사용하면 **"One small step for man, one giant leap for mankind"**라는 문장은 12개의 토큰으로 구성됩니다. 즉, 이 문장을 생성하려면 LLM이 12번 실행되어야 합니다. 이러한 실행 한 번을 디코딩 단계(decoding step)라고 합니다.

 

LLM이 클수록 더 높은 성능을 발휘할 수 있지만, 동시에 느려지는 경향이 있습니다. 예를 들어, 각 디코딩 단계마다 모델 전체 가중치를 읽어야 하기 때문입니다. 이는 단어 하나를 생성하는 데에도 테라바이트(TB) 단위의 데이터를 읽어야 할 수 있음을 의미합니다.

 

보통 LLM은 문서 요약이나 대화형 응답과 같이 많은 단어를 생성하는 것이 목표이며, 각 토큰은 이전에 생성된 토큰에 의존합니다. 따라서 토큰을 하나씩 생성해야 하며, 모델의 가중치를 반복적으로 읽어야 합니다. 이 과정에서 추측 디코딩 방법은 다음 두 가지 핵심 관찰에 의해 동기부여되었습니다.

 

관찰 1: 일부 토큰은 다른 토큰보다 생성이 더 쉽다

모든 토큰이 동일하지는 않습니다. 어떤 토큰은 생성이 더 쉽고, 어떤 토큰은 더 어렵습니다. 다음 텍스트를 예로 들어보겠습니다:

 

What is the square root of 7? The square root of 7 is 2.646.

 

강조된 토큰 **"7"**의 생성은 상대적으로 쉽습니다. 예를 들어, 이전에 생성된 **"square root of"**라는 토큰 패턴을 따라 다음 토큰을 복사하면 됩니다. 반면, **"2.646"**과 같은 토큰을 생성하는 것은 더 어렵습니다. 모델은 답을 계산하거나 기억해야 할 수도 있기 때문입니다.

 

이 관찰은 대규모 모델이 어려운 상황(예: "2.646")에서 더 좋은 성능을 보이는 반면, 많은 쉬운 경우(예: "7")에는 작은 모델이 대규모 모델의 적절한 근삿값을 제공할 수 있음을 시사합니다.

 

관찰 2: LLM 추론의 병목은 보통 메모리 대역폭이다

TPU와 GPU 같은 머신러닝 하드웨어는 초당 수백조 번의 연산이 가능한 고도로 병렬화된 장치입니다. 그러나 메모리 대역폭은 초당 수조 바이트 수준으로, 연산 성능에 비해 두 단계 정도 낮은 수준입니다. 이는 현대 하드웨어를 사용할 때, 메모리에서 데이터를 읽어들일 때마다 수백 번의 연산을 수행할 수 있다는 것을 의미합니다.

 

반면, 현대 LLM의 기반이 되는 Transformer 아키텍처는 추론 중 메모리에서 읽어들인 데이터에 대해 상대적으로 적은 연산을 수행합니다. 즉, 현대 하드웨어에서 LLM의 출력을 생성할 때는 여전히 많은 여유 계산 리소스가 남아 있다는 것을 나타냅니다.

 

 

추측 실행(Speculative Execution)

토큰이 순차적으로 계산되는 동안 추가적인 병렬 계산 리소스가 사용 가능하다는 가정에 기반하여, 우리의 방법은 여러 토큰을 병렬로 계산함으로써 병렬성을 증가시키는 것을 목표로 합니다. 이 접근법은 **추측 실행(Speculative Execution)**에서 영감을 받았으며, 이는 작업이 실제로 필요한지 확인하기 전에 또는 병렬로 작업을 수행하여 병렬성을 증가시키는 최적화 기술입니다. 현대 파이프라인 CPU에서의 분기 예측(branch prediction)은 추측 실행의 대표적인 사례입니다.

 

추측 실행이 효과적이기 위해서는 실제로 필요할 가능성이 높은 작업을 제안할 수 있는 효율적인 메커니즘이 필요합니다. 이를 보다 일반화하면, 다음과 같은 추측 실행의 추상적 설정을 고려할 수 있습니다. 여기서 f(X)와 g(Y)는 시간이 오래 걸리는 연산이라고 가정합니다:

Y = f(X)  
Z = g(Y)

위의 설정에서 f(X)는 Y를 계산하며, Y는 g(Y)의 입력이 됩니다. 이러한 연산은 보통 순차적으로 이루어집니다. 하지만 추측 실행은 빠른 근사 함수 f*(X)를 사용하면, 첫 번째 느린 연산 f(X)와 동시에 g(f*(X))를 병렬로 평가할 수 있다고 제안합니다. 이후 f(X)가 완료되었을 때 f*(X)의 출력이 Y와 동일한 경우, 병렬성을 높이는 데 성공한 것입니다. 만약 출력이 다르면, g(f*(X))의 계산을 폐기하고 순차적으로 g(Y)를 계산하면 됩니다. f*(X)가 더 정확할수록 병렬성이 높아질 가능성이 커지며, 어떤 경우에도 동일한 출력을 보장할 수 있습니다.

 

추측 샘플링(Speculative Sampling)

우리는 추측 실행을 확률론적 설정으로 일반화한 **추측 샘플링(Speculative Sampling)**을 제안했습니다. 여기서는 작업이 일정 확률로 실행되어야 한다는 점이 차별점입니다. 설정은 다음과 같습니다:

Y ~ f(X)  
Z = g(Y)

f(X)는 확률 분포를 출력하며, 이 분포에서 Y를 샘플링하여 g 함수의 입력으로 사용합니다. 빠른 근사 f*(X)를 통해, f(X)와 병렬로 g(sample(f*(X)))를 실행할 수 있습니다. 여기서 표준 추측 실행처럼 샘플이 일치하지 않을 경우 계산을 폐기할 수도 있지만, 이는 비효율적입니다. 예를 들어, f(X)와 f*(X)가 항상 1에서 100까지의 균등 확률 분포를 출력한다고 가정하면, 추측 실행은 샘플이 일치하는 경우만 계산을 수용하므로 100번 중 1번만 성공합니다. 이는 비효율적이며, f와 f*가 동일한 함수라면 샘플을 항상 수용할 수 있어야 합니다.

 

대신, 추측 샘플링은 f(X)와 f*(X)에 따라 f*의 추측을 확률적으로 수용 또는 폐기하는 방법을 제공하며, 최적성과 동일한 출력 분포를 보장합니다. 이 방법은 강화 학습이나 물리 시뮬레이션과 같은 다른 설정에서도 유용할 수 있습니다(자세한 내용은 논문 참조).

 

추측 디코딩(Speculative Decoding)

LLM은 단일 토큰을 생성하는 것이 아니라, 다음 토큰의 확률 분포를 출력합니다. 예를 들어, 텍스트 "The most well known movie director is" 다음에 "Steven"이라는 토큰이 70% 확률로, "Quentin"이라는 토큰이 30% 확률로 생성될 수 있습니다. 이러한 확률 분포 샘플링은 추측 실행을 직접 적용하기에는 비효율적입니다.

 

추측 디코딩은 추측 샘플링을 활용하여 이 문제를 해결합니다. 이 방법을 사용하면 낮은 비용으로 동일한 확률 분포에서 생성된 샘플을 보장받을 수 있습니다. 특히 탐욕적 디코딩(greedy decoding)의 경우, 가장 높은 확률의 단일 토큰을 항상 샘플링하기 때문에, 추측 실행을 효과적으로 적용할 수 있음이 이전 연구에서 입증되었습니다.

 

추측 디코딩은 자기회귀 모델(예: 트랜스포머) 추론에서 추측 샘플링을 적용하는 방법입니다. 이 경우, f(X)와 g(Y)는 동일한 함수로, 입력으로 시퀀스를 받아 한 토큰을 추가한 시퀀스의 확률 분포를 출력합니다. 추측 디코딩은 동일한 확률 분포를 유지하면서도 한 토큰과 그 뒤의 토큰들을 병렬로 효율적으로 계산할 수 있습니다(논문에서는 두 개 이상의 토큰 생성 병렬화도 다루고 있습니다).

 

추측 디코딩을 적용하려면 디코딩 함수의 빠른 근사 버전이 필요합니다. 관찰 1은 작은 모델이 더 쉬운 토큰에서 잘 작동할 가능성을 시사합니다. 실제로, 기존의 소형 모델이나 간단한 휴리스틱을 사용한 사례에서 좋은 결과를 얻었습니다. 예를 들어, 번역 작업에서 매개변수 110억 개의 T5-XXL 모델에 추측 디코딩을 적용하면서, 매개변수 6천만 개의 T5-Small 모델을 추측 메커니즘으로 사용했을 때 약 3배의 속도 향상을 달성했습니다.

 

https://storage.googleapis.com/gweb-research2023-media/media/SpeculativeDecoding-1-Illustration.mp4

 

애니메이션 설명
표준 디코딩과 추측 디코딩 알고리즘을 비교하는 애니메이션입니다. 텍스트는 GPT와 유사한 대규모 트랜스포머 디코더에 의해 생성됩니다. 추측 디코딩의 경우, 훨씬 작은 모델이 추측 메커니즘으로 사용됩니다. 채택된 추측(accepted guesses)은 초록색, 거부된 제안(rejected suggestions)은 빨간색으로 표시됩니다.

 

더 효율적인 AI를 향하여

산업 전반에서 추측 디코딩이 채택되었으며, 이를 통해 놀라운 성능 향상이 보고되었습니다. 이러한 폭넓은 채택은 추측 디코딩 패러다임을 활용한 통찰력 있고 효과적인 기법들 덕분에 가속화되었습니다. 이들은 종종 다른 혁신적인 방법과 결합되어 사용되었습니다. 예를 들어:

• 분산 환경에서 대규모 모델에 대한 효율성 증명
• 단일 추측 대신 여러 초안 추측(draft guesses) 사용
• 목표 모델(target model)의 지식을 초안 모델(draft model)에 증류(distilling)
• 초안 모델이 목표 모델의 일부를 사용하는 방식
• 초안과 목표 모델에 단일 모델 사용
• 초안 토큰 전체를 함께 검증

또한, 이 방법은 이미지와 음성 생성과 같은 도메인에서도 성공적으로 적용되었습니다.

LLM 사용이 증가함에 따라 추론을 더 효율적으로 만드는 필요성이 점점 더 중요해지고 있습니다. 우리는 추측 디코딩 및 기타 기존 패러다임을 활용한 추가 아이디어와 완전히 새로운 접근법들이 등장하기를 기대하고 있습니다.

 

감사의 말

이 연구는 아사프 아하로니(Asaf Aharoni), 아비나탄 하시딤(Avinatan Hassidim), **대니 바인스타인(Danny Vainstein)**과의 긴밀한 협력의 결과입니다.
추가적으로, 리뷰, 도움, 통찰력 있는 논의, 귀중한 피드백 및 지원을 제공한 다음 분들께 감사의 말씀을 전합니다:

• YaGuang Li, Blake Hechtman, Tao Wang, Toby Boyd, Nathan Lintz, Phil Chen, Nir Shabat, Jayant Madhavan, Aliaksei Severyn, Jakub Adamek, Jonathan Mallinson, Zhifeng Chen, Yoel Drori, Mariano Schain, Charlie Chen, Noam Velan, Nitish Kulkarni, Sidharth Mudgal, Sasha Goldshtein, Nadav Sherman, Pilar Manchon, Fernando Pereira, Eyal Segalis, Eyal Molad, Dani Valevski, Daniel Lumen, Valerie Nygaard, Steve Baker, Srinivasan (Cheenu) Venkatachary, Hema Budaraju, Ziteng Sun, Ananda Theertha Suresh, Elizabeth Hamon Reid, Jeff Dean, Prabhakar Raghavan, James Manyika
• Google Research, Google DeepMind, Google Search 팀들.

모든 협력과 지원에 깊이 감사드립니다.