Hugging Face Transformers 라이브러리가 Greedy Search, Beam Search, Top-K/Top-P Sampling 등 3가지 텍스트 생성 디코딩 방식을 API로 제공하여 언어 모델의 출력 품질 제어 가능

Context

트랜스포머 기반 대규모 언어 모델의 등장으로 개방형 텍스트 생성 수요가 증가했으나, 단순한 그리디 탐색만으로는 높은 확률을 가진 단어 시퀀스를 놓치거나 반복적인 텍스트를 생성하는 문제가 발생했다. 또한 각 디코딩 방식의 장단점을 실무에서 쉽게 비교하고 적용할 방법이 명확하지 않았다.

Technical Solution

• Greedy Search 구현: 각 타임스텝 t에서 조건부 확률 P(w_t | w_1:t-1)이 최대인 단어를 선택하는 방식을 model.generate() API로 제공
• Beam Search 구현: num_beams 파라미터로 각 타임스텝에서 상위 N개 가설을 유지하고 최종적으로 전체 확률이 가장 높은 가설을 선택하는 방식 제공
• Top-K Sampling 구현: 확률이 높은 상위 K개 토큰만 샘플링 후보로 제한하여 낮은 확률 토큰 제거
• Top-P Sampling 구현: 누적 확률이 P에 도달할 때까지만 토큰을 샘플링 후보로 포함시켜 동적으로 후보 토큰 개수 조절
• Transformers 라이브러리의 통일된 API: 모든 디코딩 방식을 max_new_tokens, num_beams, top_k, top_p 등 파라미터로 호출 가능하게 통합

Impact

Greedy Search 사용 시 반복 텍스트 생성 문제가 보고되었으나, Top-K와 Top-P Sampling이 Greedy Search와 Beam Search 대비 더 유창한 텍스트 생성 경향을 보임. 모델 학습 방식을 조정하면 Beam Search가 Top-P Sampling 대비 인간 평가 기준 더 유창한 텍스트 생성 가능 (Welleck et al., 2019 & 2020 논문 기준).

Key Takeaway

Autoregressive 텍스트 생성에서 최적의 디코딩 방식은 사용 사례마다 상이하므로, 통일된 API를 통해 여러 방식을 쉽게 실험하고 비교하는 것이 중요하다. Transformers 라이브러리는 GPT-2, TensorFlow, JAX를 포함한 다양한 모델과 프레임워크에서 동일한 API로 이 실험을 가능하게 함으로써 선택지 없는 개발자를 해방시켰다.