Transformer Beam Search 란?

빔 서치(Beam Search)는 디코더가 여러 개의 후보 토큰 시퀀스를 유지하면서 최적의 시퀀스를 선택하는 기법입니다. 이를 통해 단순히 최대 확률을 갖는 단어 하나만 선택하는 대신, 여러 개의 후보를 고려하여 더 나은 텍스트를 생성할 수 있습니다. 아래는 빔 서치의 사용 방법을 간단히 설명한 것입니다:

빔 크기(Beam Size) 설정:
빔 서치에서는 각 타임 스텝마다 빔 크기라고 하는 후보 시퀀스의 개수를 설정합니다. 빔 크기가 클수록 더 많은 후보가 유지되므로 더 정확한 결과를 얻을 가능성이 높아집니다. 하지만 계산 비용도 증가하므로 적절한 빔 크기를 선택해야 합니다.

초기화:
빔 서치를 시작하기 위해 시작 토큰(예: "[CLS]")을 입력으로 사용하여 디코더를 초기화합니다. 이 시작 토큰을 통해 디코더는 텍스트 생성을 시작하게 됩니다.

디코더 실행:
디코더는 시작 토큰을 입력으로 받고, 각 타임 스텝마다 확률이 높은 상위 빔 크기만큼의 단어를 예측합니다. 이러한 예측된 단어들을 각각의 빔에 연결하여 후보 시퀀스를 유지합니다.

끝 토큰 검사:
디코더는 끝 토큰(예: "[SEP]" 또는 "." 등)을 예측하면 해당 후보 시퀀스를 종료합니다. 빔 서치는 최소 길이 기준을 충족하면서 확률이 가장 높은 시퀀스를 선택합니다.

반복:
빔 크기만큼의 후보 시퀀스를 유지한 상태로 디코더를 다음 타임 스텝으로 진행합니다. 디코더가 끝 토큰을 예측할 때까지 반복합니다.

결과 선택:
끝 토큰을 예측한 후, 빔 크기 내에서 가장 높은 확률을 갖는 시퀀스를 최종 결과로 선택합니다.

빔 서치를 구현할 때는 보통 우선순위 큐(Priority Queue)를 사용하여 빔의 상위 후보들을 효율적으로 관리합니다. 또한, 빔 크기가 클수록 계산 비용이 증가하므로 적절한 빔 크기를 선택하는 것이 중요합니다. 빔 서치를 사용하면 모델의 생성 결과를 향상시키는데 도움이 되며, 언어 모델링 등의 다양한 자연어 처리 작업에 활용됩니다.