LLM 디코딩 전략

들어가며

이 글은 2024-10-29에 작성된 Decoding Strategies in Large Language Models (by Maxime Labonne)을 읽고 번역 및 요약한 글입니다.

이 글은 다음의 내용을 다룹니다.

• 디코딩 전략(greedy, beam search, top-k sampling, nucleus sampling)의 동작 방식
• 주요 파라미터(temperature, num_beams, top_k, top_p) 조정 방법

이 글에 사용된 코드는 Google Colab에서 확인할 수 있습니다.

---

배경

대규모 언어 모델(LLM)이 텍스트를 생성할 때, 사람들은 흔히 모델이 직접 단어를 "생각"해서 문장을 만든다고 오해합니다. 하지만 실제 과정은 훨씬 더 수학적입니다. 모델은 입력된 텍스트 시퀀스를 바탕으로, 어휘집(vocabulary)에 있는 모든 가능한 다음 토큰에 대해 점수, 즉 로짓(logit)을 계산합니다. 이 로짓 값은 소프트맥스(softmax) 함수를 통해 확률 분포로 변환됩니다.

예를 들어, GPT-2 모델에 I have a dream이라는 프롬프트를 입력하면, 모델은 다음에 올 토큰의 확률을 계산합니다. 가장 가능성 있는 토큰은 that일 수도 있고, 다른 토큰일 수도 있습니다. 원문 블로그의 예시에서는 다음 토큰으로 of가 올 확률, 즉 $P(\text{of} | \text{"I have a dream"}) = 17\%$ 라고 설명합니다.

이렇게 생성된 확률 분포로부터 실제로 다음 토큰 하나를 '선택'하는 방법이 바로 디코딩 전략(Decoding Strategy)입니다. 어떤 전략을 선택하느냐에 따라 생성되는 텍스트의 품질, 창의성, 일관성이 크게 달라집니다. 모델은 한 번에 한 토큰씩 텍스트를 생성하며, 이전에 생성된 토큰을 다시 입력의 일부로 사용하는 자동 회귀(auto-regressive) 방식을 따릅니다.

이제부터 주요 디코딩 전략들을 하나씩 살펴보겠습니다.

1. Greedy Search (탐욕 검색)

Greedy Search는 가장 직관적이고 간단한 디코딩 전략입니다. 각 단계(timestep)에서 가장 높은 확률을 가진 토큰을 선택하고, 다음 단계로 넘어갑니다. 이 과정을 원하는 길이의 시퀀스가 생성될 때까지 반복합니다.

동작 방식

1. I have a dream 입력에 대해 모델이 계산한 확률 분포에서 가장 확률이 높은 토큰을 선택합니다. (예: of)
2. 새로운 시퀀스 I have a dream of를 모델의 다음 입력으로 사용합니다.
3. 다시 가장 확률이 높은 다음 토큰을 선택합니다. (예: being)
4. 이 과정을 반복하여 최종적으로 I have a dream of being a doctor.와 같은 문장을 생성할 수 있습니다.

장점

• 속도와 효율성: 계산적으로 매우 간단하여 가장 빠르게 결과를 생성할 수 있습니다.
• 결정론적(Deterministic): 동일한 입력에 대해서는 항상 동일한 출력을 보장합니다.

단점

• 근시안적 접근: 각 단계에서는 최적의 선택이지만, 전체적으로는 최적의 시퀀스를 놓칠 수 있습니다. 예를 들어, being ($P=9.68\%$)과 doctor ($P=2.86\%$)의 확률이 상대적으로 낮았음에도 불구하고, 첫 단계에서 of를 선택했기 때문에 이 경로에 갇히게 됩니다. 더 높은 전체 확률을 가진 다른 시퀀스가 존재할 수 있지만 Greedy Search는 이를 고려하지 않습니다.
• 반복적인 결과: 고품질의 자연스러운 텍스트를 생성하는 데 실패하는 경우가 많으며, 특정 구문이 반복되는 현상이 발생하기 쉽습니다.

이러한 명백한 단점 때문에 Greedy Search는 실제 애플리케이션에서는 거의 사용되지 않지만, 다른 고급 전략들을 이해하기 위한 좋은 출발점입니다.

2. Beam Search (빔 검색)

Beam Search는 Greedy Search의 근시안적인 단점을 보완하는 전략입니다. 각 단계에서 단 하나의 최선이 아닌, num_beams라는 파라미터로 지정된 개수만큼의 가장 가능성 있는 시퀀스(가설)를 유지하고 확장해 나갑니다.

동작 방식

num_beams=2라고 가정해 보겠습니다.
1. 첫 번째 단계에서 가장 확률이 높은 2개의 토큰(예: of, that)을 선택하여 두 개의 독립적인 가설(I have a dream of, I have a dream that)을 생성합니다.
2. 두 번째 단계에서는, 이 두 가설 각각에 대해 다음에 올 모든 토큰의 확률을 계산합니다. 그러면 수많은 가능한 두 토큰 시퀀스가 생성됩니다.
3. 모든 가능한 시퀀스들의 누적 점수(일반적으로 로그 확률의 합)를 계산하고, 가장 높은 점수를 가진 상위 num_beams(2개)의 시퀀스만 유지하고 나머지는 폐기합니다.
4. 이 과정을 목표 길이에 도달할 때까지 반복한 후, 최종적으로 가장 높은 누적 점수를 가진 시퀀스를 결과로 반환합니다.

원문의 예시에서는 Beam Search를 적용했을 때 I have a dream... 이라는 더 일관성 있는 (혹은 안전한) 결과가 생성되었습니다.

장점

• 더 높은 품질: Greedy Search보다 훨씬 더 전체적인 맥락에서 높은 확률을 갖는 시퀀스를 탐색하므로, 더 일관되고 품질 좋은 문장을 생성합니다.
• 특정 과제에 유리: 번역이나 요약처럼 정답의 범위가 비교적 명확하고 높은 정확도가 요구되는 작업에서 뛰어난 성능을 보입니다.

단점

• 계산 비용: num_beams의 크기에 비례하여 계산량과 메모리 사용량이 증가합니다.
• 창의성 및 다양성 부족: 여전히 가장 가능성 있는 결과에 집중하기 때문에, 인간처럼 창의적이거나 예상치 못한 텍스트를 생성하기는 어렵습니다.
• 반복 문제: Greedy Search보다는 덜하지만, 여전히 일반적인 문구나 단어를 반복하는 경향이 있습니다. 이를 완화하기 위해 n-gram penalty와 같은 추가 기법이 필요합니다.

3. Sampling (샘플링)

Greedy Search와 Beam Search는 확률이 가장 높은 결과를 따라가는 결정론적 방법입니다. 하지만 인간의 언어는 항상 예측 가능하지 않으며, 때로는 덜 예상되는 단어가 문장을 더 풍부하게 만듭니다. 이러한 무작위성(stochasticity)과 창의성을 모델에 주입하기 위해 Sampling 기법이 사용됩니다.

가장 기본적인 샘플링은 모델이 출력한 확률 분포에 따라 다음 토큰을 무작위로 선택하는 것입니다. 하지만 확률이 매우 낮은 부적절한 단어까지 선택될 위험이 있어, 이를 제어하기 위한 여러 전략이 등장했습니다.

4. Temperature (온도)

Sampling의 무작위성을 조절하는 핵심 파라미터가 바로 Temperature입니다. Temperature는 소프트맥스 함수가 적용되기 전의 로짓 값에 적용되어 확률 분포를 변형시킵니다.

로짓 벡터 $L = (l_1, l_2, ..., l_V)$에 대해, Temperature $T$가 적용된 새로운 확률 $p_i$는 다음과 같이 계산됩니다.
$p_i = \frac{\exp(l_i / T)}{\sum_{j} \exp(l_j / T)}$

• $T < 1$ (예: 0.7): 로짓 값의 차이가 증폭되어 확률 분포가 더 뾰족해집니다(sharper). 모델은 확률이 높은 토큰을 선택할 가능성이 더 커져, 더 안정적이고 예측 가능한 텍스트를 생성합니다. Greedy Search와 유사해집니다.
• $T = 1$: 원래 모델의 확률 분포를 그대로 사용합니다.
• $T > 1$ (예: 1.5): 로짓 값의 차이가 줄어들어 확률 분포가 더 평평해집니다(flatter). 확률이 낮은 토큰도 선택될 가능성이 생겨, 더 다양하고 창의적이며 예상치 못한 텍스트가 생성됩니다. 하지만 너무 높으면 문맥과 무관한 단어가 등장할 수 있습니다.

5. Top-K Sampling

Top-K Sampling은 샘플링의 품질을 높이기 위한 간단하면서도 효과적인 방법입니다. 전체 어휘집에서 샘플링하는 대신, 가장 확률이 높은 상위 K개의 토큰으로 후보를 제한하고, 그 안에서만 확률에 비례하여 샘플링을 진행합니다.

동작 방식

1. top_k=50으로 설정하면, 모델이 예측한 모든 토큰 중 가장 확률이 높은 50개를 추립니다.
2. 이 50개의 토큰에 대해서만 확률 분포를 재정규화(re-normalize)합니다.
3. 재정규화된 분포에 따라 다음 토큰을 샘플링합니다.

원문의 예시에서는 Top-K 샘플링을 통해 I have a dream job and I want to와 같이 더 자연스러운 문장이 생성되었습니다.

장점

• 품질과 다양성의 균형: 확률이 매우 낮은 엉뚱한 토큰이 선택되는 것을 방지하여 문장의 일관성을 유지하면서도, 적절한 수준의 창의성을 부여합니다.

단점

• 고정된 K의 한계: 확률 분포의 형태에 따라 유연하게 대처하지 못합니다. 예를 들어, 모델이 다음에 올 단어를 매우 확신하는 경우(분포가 뾰족함) K가 50이면 불필요하게 많은 후보를 고려하게 되고, 반대로 여러 가능성이 있는 경우(분포가 평평함) K가 너무 작으면 좋은 후보를 놓칠 수 있습니다.

6. Nucleus Sampling (Top-P Sampling)

Nucleus Sampling (또는 Top-P Sampling)은 Top-K의 한계를 극복하기 위해 제안된 더욱 정교한 전략입니다. 후보 토큰의 개수(K)를 고정하는 대신, 후보들의 누적 확률 합(P)을 기준으로 동적으로 후보 집합의 크기를 조절합니다.

동작 방식

1. top_p=0.9로 설정하면, 확률이 높은 순서대로 토큰을 정렬합니다.
2. 누적 확률이 0.9를 초과하는 지점까지의 토큰들로 후보 집합(nucleus)을 구성합니다.
3. 이 후보 집합 내에서만 확률을 재정규화하여 샘플링을 수행합니다.

이 방식의 핵심은 확률 분포의 모양에 따라 후보 집합의 크기가 달라진다는 점입니다.

• 분포가 뾰족할 때 (모델의 확신이 높을 때): 소수의 토큰만으로도 누적 확률 $P$에 도달하므로 후보 집합이 작아집니다.
• 분포가 평평할 때 (모델의 확신이 낮을 때): 더 많은 토큰이 있어야 누적 확률 $P$에 도달하므로 후보 집합이 커집니다.

원문의 예시에서는 Nucleus Sampling을 통해 I have a dream. I'm going to라는, 의미적으로 더 발전된 문장이 생성되었습니다. 이 전략은 유연성 덕분에 현재 많은 LLM에서 기본적으로 사용될 만큼 높은 성능을 보여줍니다.

결론 및 요약

전략	동작 방식	장점	단점	주요 파라미터
Greedy Search	각 단계에서 가장 확률이 높은 토큰($P_{max}$)을 선택	빠르고 간단함	근시안적, 반복적, 품질 저하	-
Beam Search	num_beams개의 가장 가능성 있는 시퀀스를 유지하며 탐색	Greedy Search보다 고품질, 정확성	계산 비용 높음, 창의성 부족	num_beams
Top-K Sampling	확률 상위 K개의 토큰 중에서 샘플링	품질과 다양성의 균형	확률 분포 변화에 둔감	top_k, temperature
Nucleus Sampling	누적 확률 P를 넘는 최소한의 토큰 집합(Nucleus)에서 샘플링	확률 분포에 따라 동적으로 후보 조절, 고품질 및 다양성	Top-K보다 약간 복잡	top_p, temperature

최적의 디코딩 전략은 해결하려는 과제에 따라 달라집니다.

• 정확성과 사실 기반 응답이 중요할 때 (번역, 요약, 질의응답): num_beams를 적절히 설정한 Beam Search나, temperature를 낮춘 샘플링 방식이 유용할 수 있습니다.
• 창의성과 다양성이 중요할 때 (스토리 생성, 창의적 글쓰기, 챗봇 대화): temperature를 약간 높이고 Nucleus Sampling (Top-P)을 사용하는 것이 가장 일반적이고 효과적인 조합입니다.

이러한 디코딩 전략과 파라미터들의 작동 원리를 이해하고 조절함으로써,
우리는 LLM이 생성하는 결과물의 방향성과 품질을 효과적으로 제어할 수 있습니다.