보조 디코딩

Assisted decoding은 보조 모델이 메인 모델이 토큰을 확정하기 전에 후보 토큰을 제안하도록 하여 텍스트 생성 속도를 높입니다. 메인 모델은 하나의 forward pass에서 후보 토큰을 검증합니다. 보조 모델은 빠르고 비용이 적게 들며, 메인 모델의 수십 번의 더 비용이 높은 forward pass를 대체할 수 있습니다.

이 가이드는 Transformers의 보조 디코딩(assisted decoding) 방법을 다룹니다.

추측적 디코딩

Speculative decoding은 더 작은 어시스턴트 모델을 사용하여 후보 토큰을 초안합니다. 메인 모델은 이 토큰들을 한 번의 패스에서 확인합니다. 검증된 토큰은 최종 출력에 포함되고, 거부된 토큰은 표준 샘플링을 트리거합니다. 메인 모델이 비용이 높은 포워드 패스를 더 적게 실행하기 때문에 생성 속도가 빨라집니다.

이 메서드는 어시스턴트 모델이 메인 모델보다 훨씬 작고 동일한 토크나이저를 사용하는 경우에 가장 잘 작동합니다. Speculative decoding은 greedy search와 sampling을 지원하지만 배치된 입력은 지원하지 않습니다.

assistant_model을 generate()에 전달합니다. 토큰 검증이 실패할 경우 다시 샘플링하도록 do_sample=True를 설정합니다.

탐욕스러운 탐색

샘플링

프롬프트 조회 디코딩

Prompt lookup decoding은 어시스턴트 모델이 필요하지 않습니다. 프롬프트에서 겹치는 n-gram을 찾아 후보 토큰을 제안합니다. 매칭이 없으면 일반적인 자기회귀 디코딩으로 전환됩니다. 이는 요약 및 번역과 같이 입력에 기반한 작업에 적합하며, 후보 토큰이 원문 텍스트의 로컬 패턴을 자주 반영하기 때문입니다.

prompt_lookup_num_tokens를 generate()에 전달합니다. 이는 반복 패턴이 감지될 때 알고리즘이 프롬프트의 앞부분에서 복사하려고 시도하는 토큰 수를 설정합니다.

탐욕적 디코딩

샘플링

자기 추측 디코딩

Self-speculative decoding은 모델의 중간 레이어를 어시스턴트로 사용하여 후보 토큰을 제안합니다. 제안이 일치하면 모델은 조기 종료되고, 나머지 레이어가 토큰을 검증하거나 수정합니다.

모든 것이 하나의 모델이므로 가중치와 캐시가 공유되어 추가적인 메모리 오버헤드 없이 속도가 향상됩니다. 이 기법은 중간 레이어에서 조기 종료(early-exit) 로짓을 지원하도록 학습된 모델에서만 작동합니다.

assistant_early_exit를 generate()에 전달하여 종료 레이어를 설정합니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("facebook/layerskip-llama3.2-1B")
model = AutoModelForCausalLM.from_pretrained("facebook/layerskip-llama3.2-1B", dtype="auto")
inputs = tokenizer("Hugging Face is an open-source company", return_tensors="pt")

outputs = model.generate(**inputs, assistant_early_exit=4, do_sample=False, max_new_tokens=20)
tokenizer.batch_decode(outputs, skip_special_tokens=True)

멀티 토큰 예측(MTP) 레이어

Multi-token prediction(MTP) 레이어는 메인 모델의 자체 체크포인트에서 가져온 추가 예측 레이어를 사용하여 후보 토큰을 생성합니다. MTP 레이어는 메인 모델의 임베딩과 출력 헤드를 재사용합니다. 각 MTP 레이어는 토큰 하나를 생성하므로, 단계별 후보 토큰 수는 체크포인트의 MTP 레이어 수와 일치합니다. 메인 모델은 하나의 포워드 패스에서 생성된 후보를 검증합니다.

MTP는 DeepSeek-V3와 GLM-4.5와 같이 MTP 레이어로 학습된 체크포인트에서만 작동합니다. 이러한 체크포인트에는 MTP 가중치가 포함되어 있으며, config.json에서 num_nextn_predict_layers로 나열되는 num_mtp_layers가 설정되어 있습니다.

use_mtp=True를 generate()에 전달합니다. 모델에 MTP 레이어가 없으면 오류가 발생합니다. 스펙ulative 디코딩과 마찬가지로 MTP는 그리디 검색과 샘플링을 지원하지만 배치 입력은 지원하지 않습니다.

여기서 MTP는 DeepSeek-V3와 GLM-4.5와 같이 대상 체크포인트에 포함된 예측 레이어를 의미합니다. Gemma 4는 MTP 디코딩을 위해 별도의 어시스턴트 체크포인트를 사용합니다. 이를 assistant_model로 전달하세요(Multi-token prediction drafter 참고).

text-generation 파이프라인과 generate()는 모두 use_mtp=True를 지원합니다.

from transformers import pipeline

pipe = pipeline(
    "text-generation",
    model="zai-org/GLM-4.5-Air",
    device_map="auto",
)

outputs = pipe(
    "Hugging Face is an open-source company",
    use_mtp=True,
    do_sample=False,
    max_new_tokens=50,
    add_special_tokens=True,
    return_full_text=True,
)

[output["generated_text"] for output in outputs]

또는 다음과 같이 추론할 수 있습니다.

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("zai-org/GLM-4.5-Air")
model = AutoModelForCausalLM.from_pretrained("zai-org/GLM-4.5-Air", device_map="auto")
inputs = tokenizer("Hugging Face is an open-source company", return_tensors="pt").to(model.device)

outputs = model.generate(**inputs, use_mtp=True, do_sample=False, max_new_tokens=50)
tokenizer.batch_decode(outputs, skip_special_tokens=True)

멀티 토큰 예측 드래프터

MTP drafter는 별도의 checkpoint에서 로드되는 경량 어시스턴트 모델입니다. 대상 모델이 검증할 후보 토큰을 제안합니다. drafter는 대상 모델의 입력 임베딩을 사용하고, 대상 모델의 KV 캐시에 주의를 기울입니다. 대상 모델의 히든 상태(hidden states)에서 시작하여 토큰을 자기회귀적으로 작성합니다. 각 다음 토큰에 대해, 자신의 투영된 히든 상태를 drafter로 다시 입력합니다. 대상 모델은 스펙ulative decoding과 마찬가지로 한 번의 포워드 패스에서 작성된 토큰을 검증합니다.

MTP 레이어는 대상 체크포인트의 일부이며 출력 헤드를 공유하지만, MTP 드래프터는 별도의 어시스턴트 모델입니다. 대상의 입력 임베딩과 KV 캐시를 사용하여 토큰을 제안합니다.

from transformers import AutoProcessor, AutoModelForCausalLM

processor = AutoProcessor.from_pretrained("google/gemma-4-E4B-it")
target_model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-4-E4B-it",
    dtype="auto",
    device_map="auto",

)

assistant_model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-4-E4B-it-assistant",
    dtype="auto",
    device_map="auto",
)

messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Write a short joke about saving RAM."},
]

inputs = processor.apply_chat_template(
    messages,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
    add_generation_prompt=True,
    enable_thinking=False
).to(target_model.device)
input_len = inputs["input_ids"].shape[-1]

outputs = target_model.generate(
    **inputs,
    assistant_model=assistant_model,
    max_new_tokens=256,
)
response = processor.decode(outputs[0][input_len:], skip_special_tokens=False)

processor.parse_response(response)

범용 보조 디코딩

Universal assisted decoding(UAD)은 메인 모델과 어시스턴트 모델의 토크나이저가 서로 다른 경우에도 speculative decoding을 가능하게 합니다. 이를 통해 메인 모델과 임의의 소형 어시스턴트 모델을 페어링할 수 있습니다. 후보 토큰은 다시 인코딩되며, 알고리즘은 최장 공통 부분 수열(longest common subsequence)을 계산하여 계속 부분이 정렬되도록 유지합니다.

UAD를 활성화하려면 tokenizer, assistant_tokenizer, assistant_model을 generate()에 전달하세요.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

assistant_tokenizer = AutoTokenizer.from_pretrained("double7/vicuna-68m")
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b")
model = AutoModelForCausalLM.from_pretrained("google/gemma-2-9b", dtype="auto")
assistant_model = AutoModelForCausalLM.from_pretrained("double7/vicuna-68m", dtype="auto")
inputs = tokenizer("Hugging Face is an open-source company", return_tensors="pt")

outputs = model.generate(**inputs, assistant_model=assistant_model, tokenizer=tokenizer, assistant_tokenizer=assistant_tokenizer)
tokenizer.batch_decode(outputs, skip_special_tokens=True)
'Hugging Face is an open-source company that is dedicated to creating a better world through technology.'

정적 앙상블 검증

표준 speculative decoding은 손실 없음(lossless)입니다. 출력 분포가 대상 모델과 정확히 일치하도록 보장합니다. 정적 앙상블 검증(Static ensemble verification)은 대상 분포와 드래프트 분포의 혼합물에 대해 검증함으로써 이 보장을 완화합니다. 이를 통해 수용률이 높아져 추론 속도가 빨라지지만, 출력 분포가 정확한 대상 분포가 아닌 혼합물이 됩니다. 권장 시작 값은 0.7입니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM-1.7B")
model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM-1.7B", dtype="auto")
assistant_model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM-135M", dtype="auto")
inputs = tokenizer("Hugging Face is an open-source company", return_tensors="pt")

outputs = model.generate(**inputs, assistant_model=assistant_model, assistant_ensemble_weight=0.7, do_sample=True)
tokenizer.batch_decode(outputs, skip_special_tokens=True)

이 기능은 어시스턴트 모델의 후보 logits가 필요하며, 프롬프트 룩업 디코딩에서는 지원되지 않습니다.

리소스

텍스트 생성 지연 시간과 보조 생성에 대한 더 자세한 배경은 Assisted Generation: a new direction toward low-latency text generation 블로그 포스트를 참고하세요.

Update on GitHub