2. Beam Search와 그 한계

필기에 내가 적어둔 메모 중에 2³개...? (X) 라는 게 있었다. beam search가 매 스텝 후보를 k개씩 펼치니까 스텝이 지날수록 $k^2,k^3,\dots$으로 지수적으로 폭발하는 거 아닌가 싶었던 흔적이다. 결론부터 말하면 그렇지 않다. 그게 왜 그런지 따라가보면서 beam search를 정리해봤다.

score#

beam search는 매 타임스텝마다 k개의 후보 토큰을 들고 간다. 이때 k를 beam size라고 부른다. 문장 후보 $y_1,\dots ,y_t$의 score는 이렇게 정의한다.

$$score(y_1,\dots ,y_t)=\log P_{LM}(y_1,\dots ,y_t\mid \mathbf{x})=\sum _{i=1}^t\log P_{LM}(y_i\mid y_1,\dots ,y_{i-1},\mathbf{x})$$

앞서 분해한 곱 형태에 log를 붙인 것이다. 0~1 확률에 log를 취했으니 score는 모두 음수이고, 높을수록(0에 가까울수록) 더 적절한 문장으로 본다. 토큰을 생성할 때마다 score가 가장 높은 k개를 추적한다.

k²개로 펼쳤다가 k개만 남긴다#

내가 헷갈렸던 지점이 정확히 여기다. k개 후보 각각에 대해 다음 토큰 후보를 다시 k개씩 뽑으면, 한 스텝에 $k^2$개의 후보가 생긴다. 하지만 그걸 전부 들고 가는 게 아니다. $k^2$개 중 score가 가장 높은 k개만 남기고 나머지 가지는 쳐낸다. 그래서 다음 스텝도 다시 k개에서 출발한다. 후보 수가 매 스텝 k로 고정되니 $k^2,k^3$으로 폭발하지 않는다.

k=2로 예를 들어 따라가봤다. START에서 확률이 높은 두 토큰 He(-0.7), I(-0.9)를 고른다. 각각에서 다음 토큰을 둘씩 뻗으면

• He → hit (-1.7), struck (-2.9)
• I → was (-1.6), got (-1.8)

여기서 score는 누적합이라 He의 -0.7, I의 -0.9가 각각 더해진 값이다. 이 네 후보 중 상위 2개는 was(-1.6)와 hit(-1.7)이고, struck과 got은 버린다. 다시 이 둘에서 각각 두 개씩 뻗어 4개를 만들고 또 상위 2개만 남긴다. 매 타임스텝에서 $k^2$개의 후보 중 k개만 선택하는 것이 핵심이다.

종료 조건#

beam search는 갈래가 여러 개라 종료 기준이 따로 필요하다. 내가 정리해둔 세 가지는 이렇다.

1. EOS를 만나 조기 종료된 문장은 따로 보관해두고, beam size를 유지하며 나머지 갈래로 계속 서치한다.
2. EOS가 안 나와도 미리 정해둔 최대 문장 길이에 도달하면 완성으로 처리한다.
3. 보관된 문장이 특정 개수에 도달하면 decoding을 종료한다.

최종적으로는 보관된 문장 중 score가 가장 높은 것을 선택한다.

길이 정규화#

그런데 여기서 한 가지 문제가 생긴다. score는 음수가 누적되는 구조다. 문장이 길수록 음수 항이 더 많이 더해지니 score가 낮아질 확률이 높다. 그러면 짧은 문장이 부당하게 유리해진다. 그래서 최종 score를 문장 길이 t로 나눠 평균을 낸다.

$$score(y_1,\dots ,y_t)=\frac{1}{t}\sum _{i=1}^t\log P_{LM}(y_i\mid y_1,\dots ,y_{i-1},\mathbf{x})$$

길이로 나눠 평균을 취하면 길이 요소가 제거되어, 길다는 이유만으로 손해 보는 일이 사라진다.

beam search도 결국 반복에 약하다#

beam search를 끝까지 밀어붙여도, 결국 확률이 높은 토큰만 고른다는 본질은 greedy와 같다. 그래서 아무리 돌려도 그럴듯하지만 비슷비슷한, 반복적인 문장만 나오는 경향이 있다. 필기에도 맨날 확률 높은 것만 반복함, 반복에 매우 취약이라고 적어뒀다.

확률 곡선으로 비교하면 차이가 분명하다. beam search가 생성한 문장은 매 시점 확률이 높게 평평하게 유지되는 반면, 실제 인간의 언어는 확률이 높은 단어와 낮은 단어를 끊임없이 오르내린다. 인간은 늘 확률 최댓값 단어만 말하지 않는다. 좋은 문장과 확률 최대화가 같지 않다는 뜻이다.

이 간극을 메우려고 확률분포 자체를 건드리거나 일부러 뽑기를 도입하는 것이 sampling 기반 디코딩이다.