추측 디코딩으로 Whisper 추론 속도 2배 향상

TL;DR

Hugging Face는 추측 디코딩을 OpenAI의 Whisper 모델에 적용하면 추론 시간이 대략 2배 빨라진다고 보여줍니다(예: 작은 벤치마크에서 73 초에서 33 초로 감소). 전사 품질은 전혀 손실되지 않으며, 이 방법은 원본 모델과 정확히 동일한 토큰 출력을 보장합니다.

추측 디코딩이란?

Leviathan et al.(2022)이 제안한 추측 디코딩은 빠른 보조 모델과 더 큰 메인 모델을 짝지어 사용합니다. 보조 모델은 짧은 후보 토큰 시퀀스(예: 다섯 토큰)를 생성합니다. 메인 모델은 이 시퀀스를 한 번의 순전파로 검증하고, 첫 번째 불일치가 나타날 때까지의 토큰을 받아들이고 나머지는 버립니다. 수정된 접두사는 새로운 컨텍스트가 되어 보조 모델에 다시 전달되고, 이 과정이 반복됩니다. 보조 모델이 메인 모델보다 훨씬 빠르게 실행되기 때문에 전체 디코딩 속도가 향상되며, 검증 단계가 동일한 최종 출력을 보장합니다.

핵심 요구 사항:

  • 보조 모델은 메인 모델과 정확히 같은 토크나이저/어휘를 공유해야 합니다.
  • 메인 모델보다 최소 3배 빠르면서, “쉬운” 토큰의 대다수(≈70‑80 %)를 올바르게 예측해야 합니다.

기본 Whisper 추론 속도

저자들은 Whisper large‑v2를 73개 샘플(LibriSpeech 검증 세트, 총 ≈9 MB)으로 벤치마크했습니다:

  • float16 정밀도, FlashAttention(attn_implementation="sdpa"), 저‑CPU‑메모리 로딩 사용.
  • 샘플당 생성 시간을 측정.

결과: 총 72.99 초(샘플당 ≈1 초)이며, 단어 오류율(WER)은 3.5 %.

Whisper에 추측 디코딩 적용하기

보조 모델 선택

  • Distil‑Whisper distil‑large‑v2를 보조 모델로 선택했습니다. Whisper의 인코더는 그대로 유지하지만 디코더 레이어를 32개 중 2개만 사용해 6배 빠른 속도를 제공하면서, OOD 데이터에서 전체 모델 대비 1 % 이내의 WER 차이를 보입니다.
  • 인코더를 메인 모델과 공유할 수 있어 VRAM 오버헤드가 적습니다(≈8 % 추가).

구현 세부 사항

assistant_model = AutoModelForCausalLM.from_pretrained(
    "distil-whisper/distil-large-v2",
    torch_dtype=torch_dtype,
    low_cpu_mem_usage=True,
    use_safetensors=True,
    attn_implementation="sdpa",
).to(device)

# 보조 모델을 이용한 생성
outputs = model.generate(**inputs, assistant_model=assistant_model, **kwargs)

assistant_model 인자는 🤗 Transformers에서 보조 생성 전략을 활성화하며, 이는 추측 디코딩을 구현합니다.

속도 향상 결과 (영어)

  • 추측 디코딩 시간: 총 32.70 초(샘플당 ≈0.45 초).
  • 속도 향상: 기본 대비 2.2배 빠름.
  • WER: **3.5 %**로 동일, 정확히 같은 출력을 유지함을 확인.

같은 방법을 pipeline API에서도 generate_kwargs={"assistant_model": assistant_model}을 전달해 사용할 수 있습니다.

다국어 전사

Distil‑Whisper 체크포인트는 영어 전용이므로, 다국어 사용을 위해 저자들은 가장 작은 다국어 Whisper 체크포인트(tiny)를 보조 모델로 사용했습니다.

VoxPopuli의 네덜란드어 샘플 73개에 대한 벤치마크:

  • 기본 (large‑v2): 116.5 초, WER 12.8 %.
  • 추측 디코딩 (보조 = tiny): 62.1 초, WER 12.8 %.
  • 속도 향상: 1.9배.

language와 task 인자를 generate에 적절히 전달하면 전사와 번역 작업 모두에 적용할 수 있습니다.

효율성 극대화를 위한 전략

보조 모델 선택

  • 메인 모델 대비 ≥ 3배 빠르고, 토큰 일치율이 ≥ 70‑80 %인 모델을 목표로 합니다.
  • 특정 언어에 대해서는 큰 Whisper 모델(예: large‑v3)을 메인 모델로 fine‑tune하고, 동일 아키텍처를 distil해 빠른 보조 모델을 만들면 토큰 분포가 맞아떨어져 두 모델 모두 WER이 개선됩니다.

배치 크기 고려사항

  • 추측 디코딩은 배치 크기 = 1일 때 가장 큰 이득을 제공합니다.
  • 큰 배치에서는 배치 전체 후보가 메인 모델과 일치해야 하며, 불일치가 발생하면 조기에 버려져 속도 향상이 감소합니다.
  • 실험적으로 배치 크기 4까지는 속도 향상이 유지되지만, 그 이상에서는 오버헤드가 이점을 상쇄합니다(참조: Distil‑Whisper 논문, Sec. D.3).

실용적인 요점

  • 추측 디코딩은 기존 Whisper 파이프라인에 그냥 추가하면 되는 드롭‑인 대체 방식입니다: assistant_model 인자를 넣기만 하면 됩니다.
  • 무료 2배 추론 가속을 제공하면서 전사 품질이 변하지 않음을 보장합니다.
  • 모델에 구애받지 않는 접근법으로, 동일 토크나이저를 공유하는 충분히 빠른 보조 모델만 있다면 어떤 트랜스포머 기반 시퀀스 모델도 이점을 얻을 수 있습니다.

“추측 디코딩은 기존 Whisper 파이프라인에 완벽한 드롭‑인 대체이며, 정확도가 동일한 상태에서 무료 2배 속도 향상을 제공합니다.” — Sanchit Gandhi, Hugging Face 블로그 (2023‑12‑20)

감사의 글: 이 글은 Patrick von Platen, Pedro Cuenca, Joao Gante에게 피드백을 받아 작성되었으며, 🤗 Transformers의 보조‑생성 구현에 감사드립니다.

Sources