LLM 디코딩 탐색 알고리즘 정리: 그리디·빔서치·A*와 최적 탐색이 어려운 이유 (CMU 강의)

핵심 메시지

• 초기 언어모델은 최고 확률 출력을 찾으면 같은 말을 반복하거나 빈 출력을 내놓아, 일부러 샘플링으로 추론을 '망가뜨려' 쓰는 시기가 있었다.
• 모델 품질이 올라오면서 그리디 탐색이 온도 1 샘플링보다 나은 결과를 내는 상황이 되었고, 강연자는 이제 다시 점수가 높은 출력을 효율적으로 찾는 방법을 재고할 때라고 본다.
• 빔서치는 길이를 먼저, 그다음 점수를 보고 확장하는 근사 알고리즘이라 최적해를 보장하지 않는다. 반면 완전 탐색과 A*는 최적을 보장한다.
• A*의 최적성은 휴리스틱이 남은 비용을 결코 과대평가하지 않는다는 허용성 조건에서 나온다. 과소평가는 탐색을 늘릴 뿐 실패를 부르지 않는다.
• 트랜스포머 언어모델에서는 어휘 크기만큼 그래프가 지수적으로 커지고 허용 가능한 휴리스틱을 만들기도 어려워, 가설 재결합과 미래 비용 예측 같은 우회로가 필요하다.

쉽게 이해하기

강의는 기술적인 내용에 앞서 왜 지금 이 주제를 다루는지부터 설명한다. GPT-1이나 GPT-2 시절 모델은 최고 확률 경로를 찾으면 같은 문구를 끝없이 반복하거나 사실상 쓸모없는 출력을 내놓았다. 그래서 당시에는 분포에서 직접 뽑지 않는 샘플링 전략들로 추론을 일부러 비틀어 쓰는 것이 정석이었다. GPT-4급에 이르러 탐색이 제대로 작동하기 시작했고, 최근 추론 모델들은 아예 온도 1 샘플링만 허용하는 경우도 있다. 강연자는 그리디 탐색이 온도 1 샘플링보다 나은 결과를 내는 지금이야말로 높은 점수의 출력을 효율적으로 찾는 방법을 다시 볼 시점이라고 말한다.

설명은 탐색 공간을 그래프로 두는 데서 시작한다. 상태 집합, 어휘에 해당하는 알파벳, 전이, 시작 상태와 종료 상태, 그리고 전이에 붙는 실수 가중치로 이루어진 가중 유한상태오토마타 형식이다. 가중치는 확률로 둘 수도 있고 로그 확률로 둘 수도 있는데, 탐색을 보통 최단 경로 문제로 풀기 때문에 0에서 시작해 값을 더해 가는 음의 로그 확률이 편하다. 점수가 낮을수록 좋은 경로가 된다.

같은 예제 그래프 위에서 알고리즘들을 비교하면 차이가 뚜렷해진다. 그리디 탐색은 우선순위 큐 크기가 1이어서 매 단계 가장 좋은 간선만 따라가고, 네 번 만에 끝나지만 국소적으로 나쁜 결정을 되돌릴 수 없다. 빔서치는 큐 크기를 빔 폭만큼 유지하며 길이가 짧은 가설을 먼저 확장하고 같은 길이에서는 점수로 정렬한다. 여섯 번 만에 그리디와 다른(그리고 더 좋은) 경로를 찾아내지만, 여전히 최적해를 보장하지는 않는 근사 알고리즘이다. 참고로 빔 폭이 1인 빔서치가 곧 그리디 탐색이다.

최적을 보장하려면 가지치기를 하지 않는 완전 탐색이 필요하다. 항상 점수가 가장 좋은 경로를 확장하므로 종료 상태에 처음 도달하는 경로가 곧 최적 경로가 되지만, 같은 예제에서 아홉 단계가 걸렸다. 어휘가 크거나 길이가 길어지면 탐색이 시작 지점 근처만 맴돌게 된다.

A*는 이 비효율을 줄이면서 최적성을 지킨다. 우선순위를 정하는 값 f를 지금까지의 실제 비용 g와 목표까지 남은 비용의 추정치 h의 합으로 두는 방식이다. 강의는 문장 앞부분만 주고 '남은 부분에 확률을 얼마나 더 써야 할 것 같으냐'고 묻는 예로 h의 직관을 설명한다. 뻔한 마무리만 남았다면 남은 비용은 거의 0이고, 답을 계산해야 하는 문제라면 훨씬 크다. h가 참 비용을 절대 과대평가하지 않는 허용 가능한 휴리스틱이면 A*는 최적을 보장하며, 예제에서는 여덟 단계로 끝나 완전 탐색보다 빨랐다. 과소평가는 확장을 늘릴 뿐 실패로 이어지지 않는다는 점, 그리고 허용성은 충분조건이지 필요조건은 아니라는 점이 질의응답에서 정리된다.

그렇다면 왜 실제 언어모델에서는 A*를 쓰지 않을까. 첫째, 토큰을 하나 생성할 때마다 그래프가 어휘 크기만큼 갈라져 지수적으로 커진다. 둘째, 허용 가능한 휴리스틱을 만들기가 매우 어렵다. 강의가 든 예가 인상적이다. 인사말을 하고 미국 헌법을 낭독하라고 시키면 모델은 어느 순간 암기한 구간에 진입해 수천 토큰을 거의 확률 1로 뱉는다. 남은 비용을 안전하게 과소평가하는 함수를 설계하기가 그만큼 까다롭다는 뜻이다.

우회로는 두 가지다. 하나는 가설 재결합으로, 비슷한 상태를 묶어 하나만 남겨 탐색 공간을 줄인다. 최근 n개 단어가 같으면 묶는 n-그램 기반 클러스터링은 구현과 캐싱이 쉽지만 이전 맥락이 크게 다른 가설까지 합쳐 버릴 수 있고, 다음 토큰 분포 사이의 KL 발산이나 코사인 유사도로 묶는 방법도 있다. 다른 하나는 미래 비용 예측이다. 허용성은 포기하되 부분 시퀀스를 끝내는 데 드는 비용을 학습으로 예측해 점수에 반영하는 것으로, 본 모델과 함께 보조 손실로 학습하거나 완성된 시퀀스로 별도 예측기를 학습할 수 있다. 강화학습의 가치 함수와 성격이 비슷하며, 2020년 기계번역 논문에서 강한 베이스라인 대비 개선이 보고됐다.

마지막으로 소개된 최선 우선 빔서치 논문은 강연자가 방법론보다 정리 논문으로서 더 좋아한다고 밝힌 연구다. 빔서치가 빔 안의 모든 가설을 무조건 확장하는 낭비를 줄이기 위해 A*처럼 점수 기반 우선순위를 쓰되 빔 제약은 유지하고, 확장하면 점수가 나빠지기만 한다는 성질을 이용해 빔에서 탈락이 확정된 가설을 미리 쳐내고 완전한 가설이 나오면 조기 종료한다. 그 결과 동일한 결과를 유지하면서 표준 빔서치 대비 최대 10배 속도 향상을 보고했다.

주요 인사이트

• 디코딩 전략의 역사는 모델 품질의 함수였다. 모델이 나쁠 때는 탐색을 망가뜨리는 것이 최선이었고, 모델이 좋아지자 제대로 탐색하는 쪽이 다시 이득이 되었다.
• 빔서치와 최선 우선 탐색, A*, 너비 우선 탐색의 차이는 비교자, 최대 큐 크기, 휴리스틱 세 가지 파라미터로 설명된다. 빔서치는 길이를 먼저 보고 최선 우선 탐색은 점수를 먼저 보며, 큐 크기가 무한하고 휴리스틱이 붙으면 A*가 된다.
• 허용 가능한 휴리스틱은 좋은 탐색의 충분조건이지 필요조건은 아니다. 모든 노드에 같은 값을 더하는 식으로 일관되게 과대평가하면 순위가 바뀌지 않아 여전히 잘 동작한다.
• 미래 비용 예측기를 큰 모델 대신 작은 모델로 학습하는 선택지는 실용적인 의미가 크다. 4천억 파라미터급 모델을 직접 학습하지 않고도 탐색을 개선할 수 있기 때문이다.
• 가설 재결합의 기준을 무엇으로 삼느냐는 여전히 열린 문제다. 표면 형태가 잠시 비슷했다가 나중에 갈라지는 경우를 다룬 최신 연구는 별로 없다는 점이 강의에서 언급된다.

자주 묻는 질문

관련 AI 소식