이론적인 관점에서 볼 때 LLM이 내놓을 수 있는 최적의 답변은 항상 가장 높은 확률과 연관된 토큰을 선택할 때 얻어진다. 이 방식은 LLM의 출력을 결정론적으로 만들며, 이는 여러 애플리케이션에서는 바람직하지 않은 특성이다. 이런 이유로 LLM의 창의성과 맥락에 대한 충실도 사이의 균형을 맞추기 위해 최근 몇 년간 다양한 샘플링 알고리즘이 제안되었다.

오늘날 가장 널리 쓰이는 방법 중 하나이자 사실상 기본값처럼 사용되는 것이 top-p다. 이는 확률 합이 “p”에 도달할 때까지 상위 점수 토큰들을 모은 뒤 가중 랜덤 샘플링을 수행하는 뉴클리어스 샘플링의 한 형태이다.

이 블로그 포스트에서는 뉴클리어스 샘플링이 최선의 접근법이 아닐 수 있다고 생각하는 이유를 살펴보고, 뉴클리어스 샘플링의 문제점을 피하기 위한 간단하고 이해하기 쉬운 대안을 제시하려 한다. 이 알고리즘은 아직 작업 중인 단계이지만, 지금 공개함으로써 논의와 다양한 시도를 촉진할 수 있기를 바란다.

로짓 속에 금이 있다

LLM의 로짓은 LLM 내부 동작 중 몇 안 되는 완전히 이해 가능한 부분임에도, 그 특성을 연구하거나 더 발전된 샘플링 방법을 탐구하거나, 불확실성과 환각 가능성을 감지해 사용자에게 알리는 작업에 대한 관심은 전반적으로 매우 적은 편이다. 연속된 토큰들의 확률 분포를 시각화하는 것은 통찰을 얻기 위한 간단하고 실용적인 방법이다.

이미지에서는 상위 32개 후보 토큰이 확률에 따라 색상으로 표시된 모습(흰색 = 0, 파란색 = 1)과 선택된 토큰, 그리고 선택된 토큰의 순위(가장 높은 확률 = 0, 그 다음 = 1, 이런 식으로)를 볼 수 있다.

위 예시에서 Mistral 베이스 모델은 움베르토 에코의 출생일과 사망일을 알고 있기 때문에, 전체 확률의 대부분을 가장 가능성 높은 토큰에 자신 있게 할당한다. 다른 경우에는 텍스트의 이어지는 내용을 표현하는 방법이 여러 가지이거나 특정 사실에 대해 확신하지 못하기 때문에 모델이 더 갈팡질팡한다. 학습 과정에서 생일을 학습하지 못한 인물의 날짜를 물으면 전혀 다른 토큰 확률 분포가 나타난다.

하지만 일반적으로 우리가 피하고자 하는 것은 차선의 토큰을 선택해 LLM 생성을 최소 퍼플렉서티 경로에서 벗어나게 하고 환각을 유발하는 것이다. 뉴클리어스 샘플링은 분포에 따라 “p”까지 토큰을 누적하는 과정에서 첫 번째 선택보다 현저히 약한 토큰을 포함할 수 있기 때문에 이를 제대로 막지 못한다. 예를 들어 위 이미지의 “writer” 토큰을 생각해 보자. 움베르토 에코는 실제로 작가였다. writer에 해당하는 토큰은 아주 높은 점수를 받지는 못했지만 여전히 두 번째 후보보다 훨씬 높은 점수를 유지한다. 그럼에도 p를 0.3으로 설정하면 0.01 같은 낮은 값을 가진 두 번째 토큰까지 누적되어 몇 퍼센트의 확률로 선택될 수 있으며, 이는 생성을 잘못된 경로로 이끌 위험이 있다.

반면 다음 토큰인 “who”의 경우에는 점수가 비교적 비슷한 여러 선택지가 있어 다양한 생성을 위해 활용될 수 있다.

눈사태 효과 활용하기

여기서 핵심적인 관찰은 다양성을 위해 너무 약한 토큰을 선택하는 것은 그다지 좋은 거래가 아니라는 것이다. 출력을 다양화하기 위해 양호한 후보가 몇 개 있는 생성 시점만을 활용하더라도 눈사태 효과가 우리를 돕는다. 입력 컨텍스트가 달라지면 LLM의 출력이 교란되고, 그 결과 텍스트의 다른 버전을 생성할 가능성이 더 높아지는 효과가 나타나기 때문이다.

First Token Cutoff (FTC) 알고리즘

면책 조항: 여기에 설명된 알고리즘은 어떠한 과학적 검증도 거치지 않았다. 나는 “최악 토큰을 제한하는” 특성을 가진 다양한 샘플링 알고리즘을 며칠간 실험해 봤고, 이 알고리즘이 적용 가능성, 결과, 이해 용이성 사이에서 가장 균형이 잘 잡힌 것으로 보인다.

여기에 설명된 알고리즘은 비공식적으로 다음과 같이 말할 수 있다.

• LLM이 특정 후보에 강하게 편향되어 있을 때는 그 토큰을 선택한다.
• 여러 개의 유효한 후보가 있을 때는 대안을 생성한다.
• 선택될 수 있는 가장 낮은 품질의 토큰은 일정 수준으로 제한되어야 한다.

Tail Free Sampling 같은 이전 연구에서도 LLM이 내놓은 소수의 고품질 토큰 집합 내에서 선택이 이루어져야 한다고 지적한 바 있다. 다만 TFS에서는 미분을 수행해 토큰 품질이 급격히 떨어지기 전까지 완만한 곡선을 유지하는 구간에 해당하는 클러스터를 선택함으로써 해당 집합을 식별한다.

반면 여기에서 제안하는 알고리즘에서는 최고 점수 토큰인 T0를 기준으로 더 제한적이고 이해하기 쉬운 컷오프를 따르도록 하고자 하며, 다른 모든 토큰과 비교해 T0에 특별한 의미를 부여한다. 즉, 해당 토큰을 생성할 때 LLM이 갖는 확신도(기본적으로 퍼플렉서티의 대리 지표)이다. 따라서 이 알고리즘은 T0와 비교해 일정 비율 이상으로 나쁜 모든 토큰을 배제한다.

0에서 1 사이의 값을 가질 수 있는 컷오프 비율을 “co”라고 부른다. 실현 가능한 “co”의 예로는 0.5가 있다.

알고리즘은 다음과 같이 동작한다.

1. 로짓에 대해 softmax()를 계산한다.
2. 토큰을 확률 순으로 정렬한다.
3. 최적 토큰의 확률인 T0가 주어졌을 때, 나머지 모든 토큰에 대한 비율을 r = 1 - (T[i] / T0)로 계산한다.
4. r <= co인 토큰만 선택한다.
5. 선택된 토큰들 중에서 가중 랜덤 선택을 수행한다.

이 방식에서는 토큰이 부드럽게 단조 감소하는 값을 가질 수 있음에도 불구하고, 가능성 집합에 포함할 수 있는 토큰에 대해 엄격한 상한이 존재한다는 점에 유의하자. 고득점 클러스터를 식별하려는 다른 방법들에서는 그렇지 않다.

실제 예시

뉴클리어스 top-p 샘플링이 실제에서는 치명적으로 실패하는 것처럼 보이지 않는 이유 중 하나는, 종종 첫 번째 토큰의 확률이 매우 높아 퍼플렉서티가 낮고, 우연히 저품질 토큰을 모아서 선택하지 않는 경우에도 생성이 합리적인 경로를 계속 유지하기 때문이다.

다음과 같은 확률을 가진 토큰들이 연속될 때는 문제가 더 심각해진다.

0.25, 0.14, 0.01

p=0.4인 경우 세 번째 저품질 토큰까지 포함하게 되어 약 3% 확률로 선택될 수 있다.

이제 co 값을 0.5(토큰이 첫 번째 토큰보다 최대 50%까지 나빠질 수 있음)로 한 First Token Cutoff를 고려해 보자.

두 번째 토큰의 r 값은 다음과 같다.

r[t1] = 1-(0.14/0.25) = 0.44 # 0.44 <= 0.5, 이 토큰은 허용됨

r[t2] = 1-(0.01/0.25) = 0.96 # 0.96 > 0.5, 이 토큰은 거부됨

출력 예시

프롬프트 “Sorted sets are”에 대해 co=0.7로 설정한 Mistral 베이스 모델(인스트럭트 없음)의 출력이다. 아래 출력은 품질을 고려해 골라낸 것이 아닌, 성공적인 세 가지 출력이다.

1. Sorted sets are a powerful data structure in Redis. They can be used to store sorted lists, to store unique values, to store scores for ranking, and to store a sorted list of sorted sets.
2. Sorted sets are a very powerful data structure. They allow you to store data in a way that makes it easy to find the highest or lowest values in the set, and they also allow you to sort the data. This can be useful for many different tasks, such as ranking users by their score, or finding the most popular items in a database.
3. Sorted sets are a very powerful data structure that can be used to solve many different problems. The most common use case is to store a list of unique elements, each of which has an associated value. For example, you could use a sorted set to store the names of all the people in your family, with their ages as the associated values.

이해 가능성이 중요한 이유

샘플링 파라미터는 LLM의 최종 사용자나 API 사용자가 반드시 튜닝해야 하는 몇 안 되는 요소 중 하나다. 시행착오가 필요한 경우가 많지만, 현실 세계에 대한 직관적인 설명이 바로 떠오르는 하나의 튜닝 가능한 파라미터를 갖는 것은 큰 도움이 된다. 게다가 "co"는 선형 파라미터이므로, temperature 같은 파라미터나 선형이면서도 로짓의 분포 형태에 크게 의존하는 top_p의 "p" 같은 파라미터에 비해 추론하기가 특히 간단하다.

향후 과제

아직 연구의 초기 단계에 있으므로, 이 알고리즘을 더 연구하고 평가할 예정이다. 무엇보다도 샘플링 알고리즘에 대한 관심이 더 높아지고, top-p와 유사한 접근 방식에서 벗어나 앞으로 나아가려는 관심이 많아지기를 바란다.

로짓 분포로부터 수집할 수 있는 흥미로운 정보가 있을지도 모른다. 예를 들어 선형 프로브(linear probe)가 LLM의 은닉층이 언제 사실적 정보를 다루고 있는지를 학습할 수 있을 가능성이 있다. 이는 토큰 퍼플렉서티와 함께 LLM 출력의 일부가 틀렸을 가능성이 높음을 사용자에게 보여주는 데 사용될 수 있다. 일반적으로 토큰 확률 분포를 시각화하는 것은 매우 유익하며, 어떤 면에서는 LLM이 어떻게 동작하고 각 단계에서 어떤 후보들이 있는지를 직접 손으로 만져보는 듯한 경험을 제공한다.

레퍼런스 구현

logits = mx.softmax(logits)
np_logits = np.array(logits) # MX -> NumPy
np_logits = np_logits.flatten()
sorted_indices = np.argsort(np_logits)
sorted_indices = sorted_indices[::-1]

co = 0.7
j = 1
t0 = np_logits[sorted_indices[0]]
while 1 - (np_logits[sorted_indices[j]] / t0) < co and j < len(np_logits):
    j += 1
accepted_logits = []
for i in range(0,j):
    accepted_logits.append(float(np_logits[sorted_indices[j]]))
accepted_logits = mx.array(accepted_logits)

idx = mx.random.categorical(accepted_logits)
idx = int(np.array(idx)) # Convert zero-dim array to scalar
token_id = sorted_indices[idx]

크레딧

이러한 실험은 Apple의 MLX 라이브러리와 그 라이브러리를 끊임없이 개발해 온 멋지고 영리한 개발자들 덕분에 정말 간단하게 수행할 수 있었다. MLX는 마땅히 그래야 하듯 극도로 접근하기 쉽다. 결국 LLM 자체는 불가해하지만, 추론 과정 자체는 단순한 프로세스이기 때문이다.