빔 서치 실전 — num_beams 매개변수

빔 서치 실전 — num_beams 매개변수

Transformers의 generate() 에서 빔 서치는 num_beams 를 켜는 것만으로 시작돼요. HF 블로그는 이 디코딩 기법을 코드로 바로 체감하게 해줘요.

가장 간단한 예제

from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")

outputs = model.generate(**inputs, num_beams=5)
tokenizer.batch_decode(outputs, skip_special_tokens=True)

num_beams=5 면 5개의 후보 시퀀스를 동시에 이끌어 가장 적합한 최종 시퀀스를 골라요.

변형들

  • 다양한 빔 서치(diverse beam search): num_beam_groupsdiversity_penalty 로 시퀀스 간 유사도를 벌점화해 더 다양한 후보를 만들어요.
  • 빔 + 샘플링: num_beams>1do_sample=True 를 주면 각 빔에서 샘플링해요.

언제 쓸까

전역에 정답이 정해진 입력 기반 작업에서 greedy보다 훨씬 안정적으로 고품질 결과를 내요. 다만 추론 비용이 빔 수에 비례해 커지므로 지연은 감안해야 해요.

더 알아보기