첫 토큰 컷오프 LLM 샘플링
이론적으로 보면 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”라고 부릅니다. 예를 들어 0.5는 충분히 쓸 만한 co 값입니다.
알고리즘은 다음과 같이 동작합니다.
- 로짓에 대해 softmax()를 계산합니다.
- 토큰을 확률 순으로 정렬합니다.
- 가장 좋은 토큰의 확률인 T0가 주어졌을 때, 나머지 모든 토큰에 대해 비율을 r = 1 - (T[i] / T0)로 계산합니다.
- r <= co인 토큰만 선택합니다.
- 선택된 토큰들 사이에서 가중 무작위 선택을 수행합니다.
이 방식을 사용하면 토큰 값들이 매끄럽게 단조 감소하는 경우에도 가능 집합에 포함할 수 있는 토큰에 엄격한 상한이 생긴다는 점에 유의해야 합니다. 고득점 클러스터를 식별하려는 다른 방법들에서는 그렇지 않습니다.
실전 예시
실제로 뉴클리어스 top-p 샘플링이 치명적으로 실패하는 것처럼 보이지 않는 이유 중 하나는 첫 번째 토큰의 확률이 매우 높은 경우가 많기 때문입니다. 퍼플렉서티가 낮을 때는 물론, 우연히 저품질 토큰을 모았다가 선택하지 않는 경우에도 생성은 합리적인 경로를 계속 따라갑니다.
다음과 같이 연속된 토큰들의 확률이 있을 때 문제가 더 심각해집니다.
0.25, 0.14, 0.01
p=0.4인 경우 세 번째 저품질 토큰까지 포함하게 되어 약 3% 확률로 선택될 수 있습니다.
이제 co 값을 0.5로 한 First Token Cutoff를 생각해 보겠습니다(토큰은 첫 번째 토큰보다 최대 50%까지만 나빠질 수 있습니다).
두 번째 토큰의 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 베이스 모델(no instruct)의 출력입니다. 다음 세 출력은 품질을 고려해 골라낸 것이 아닌, 성공적인 출력 세 개입니다.
- 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.
- 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.
- 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와 유사한 방식에서 벗어나 앞으로 나아가려는 관심이 많아지기를 바랍니다.
로짓 분포에서 수집할 수 있는 흥미로운 정보가 있을지도 모릅니다. 예를 들어 선형 프로브가 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 자체는 불가해하지만, 추론 자체는 단순한 과정이니까요.