Hugging Face Transformers 4.45.0 동적 추측 디코딩

Hugging Face와 Intel Labs는 동적 추측 디코딩을 개발했으며, 이는 작업에 따라 텍스트 생성 속도를 최대 2.7배 가속합니다. 이 기술은 Transformers 4.45.0 릴리스부터 보조 생성(assisted generation)의 기본 운영 모드가 됩니다.

동적 추측 디코딩 작동 방식

동적 추측 디코딩은 표준 추측 디코딩을 개선하여 "speculation lookahead"(SL)를 최적화합니다—즉, 빠른 초안 모델이 더 큰 대상 모델이 병렬로 검증하기 전에 생성하는 토큰 수를 의미합니다.

이전 방법들은 정적인 SL이나 이전 반복의 수용률에 기반한 휴리스틱을 사용했지만, 동적 추측 디코딩은 보조 모델 자체의 신뢰도를 사용해 언제 중단할지 결정합니다. 구체적으로 시스템은 각 예측 토큰에 대한 로짓의 softmax를 모니터링합니다. 보조 모델의 신뢰도가 사전에 정의된 assistant_confidence_threshold 이하로 떨어지면 해당 반복에서 토큰 생성이 중단되고, 최대 num_assistant_tokens에 도달하지 않았더라도 시퀀스가 대상 모델로 전송되어 검증됩니다.

성능 벤치마크

RTX 4090에서 greedy decoding(temperature = 0)으로 수행한 벤치마크 결과, 동적 접근 방식은 다양한 모델 조합 및 작업에서 휴리스틱 기반 방법보다 일관되게 우수한 성능을 보였습니다:

대상 모델 초안(보조) 모델 작업 가속도 - 휴리스틱 가속도 - 동적
facebook/opt-6.7b facebook/opt-125m summarization 1.82x 2.71x
facebook/opt-6.7b facebook/opt-125m open-ended generation 1.23x 1.59x
Salesforce/codegen-6B-mono Salesforce/codegen-350M-mono code generation (python) 0.89x 1.09x
google/flan-t5-xl google/flan-t5-small summarization 1.18x 1.31x
meta-llama/Llama-3.1-8B meta-llama/Llama-3.2-1B summarization 1.00x 1.52x
meta-llama/Llama-3.1-8B meta-llama/Llama-3.2-1B open-ended generation 1.00x 1.18x
meta-llama/Llama-3.1-8B meta-llama/Llama-3.2-1B code generation (python) 1.09x 1.15x

이 벤치마크에서 도출된 주요 결과는 다음과 같습니다:

  • Llama 3.1/3.2 조합: 동적 접근 방식은 요약 작업에서 1.52배 가속을 달성했으며, 휴리스틱 접근 방식은 의미 있는 가속을 보이지 않았습니다.
  • 코드 생성: codegen-6B-mono의 경우, 휴리스틱 접근 방식은 오히려 속도가 감소(0.89x)했지만, 동적 접근 방식은 1.09배 가속을 제공했습니다.

Transformers 4.45.0 구현

동적 추측은 Transformers 4.45.0에서 보조 디코딩의 기본 모드입니다. generate 메서드에 assistant_model을 전달하면 사용할 수 있습니다:

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

prompt = "Alice and Bob"
checkpoint = "EleutherAI/pythia-1.4b-deduped"
assistant_checkpoint = "EleutherAI/pythia-160m-deduped"
device = "cuda" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
inputs = tokenizer(prompt, return_tensors="pt").to(device)

model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
assistant_model = AutoModelForCausalLM.from_pretrained(assistant_checkpoint).to(device)

outputs = model.generate(**inputs, assistant_model=assistant_model)

구성 매개변수

사용자는 assistant_model.generation_config에서 다음 매개변수를 조정하여 특정 데이터셋에 맞는 성능을 최적화할 수 있습니다:

  • assistant_confidence_threshold: 초안 모델이 생성을 중단하는 신뢰도 수준(로짓의 softmax).
  • num_assistant_tokens: 보조 모델이 한 반복에서 생성할 수 있는 최대 토큰 수.
  • num_assistant_tokens_schedule: 'dynamic'(기본값), 'heuristic', 'constant' 중 하나로 설정하여 서로 다른 추측 전략을 전환할 수 있습니다.

이론적 근거: 오라클 모델

동적 조정의 필요성을 검증하기 위해 연구진은 "오라클"을 사용했습니다. 이 오라클은 초안 모델과 대상 모델 사이에 불일치가 발생할 때까지 토큰을 생성함으로써 각 반복에 대한 최적의 추측 lookahead를 식별합니다. MBPP 데이터셋과 Alpaca 데이터셋을 분석한 결과, 최적 초안 토큰 수에 큰 변동성이 존재함을 확인했으며, 이는 정적인 lookahead 값이 추론 속도 최적화에 비효율적임을 증명합니다.

Sources