텍스트, 오디오, 이미지 입력과 최대 256K의 긴 컨텍스트 창을 지원하는 Gemma 4 가 출시되었습니다. 자세히 알아보기

Google uses AI technology to translate content into your preferred language. AI translations can contain errors.

Hugging Face 트랜스포머를 사용하는 Gemma 4 다중 토큰 예측 (MTP)

Gemma 4 모델의 추론 속도를 개선하기 위해 기본 라인업과 함께 새로운 자동 회귀 '초안 작성기' 모델 시리즈가 출시되었습니다. 기본 Gemma 4 모델('타겟' 모델이라고 함)에만 의존하는 대신, 초안 모델은 타겟 모델이 하나의 토큰을 처리하는 데 걸리는 시간에 여러 토큰을 자기 회귀적으로 예측합니다. 이 기법을 추측 디코딩이라고도 합니다.

이제 초안 작성자가 여러 초안 토큰을 예측한 후 타겟 모델은 제안된 초안 토큰만 확인하면 됩니다. 검증은 병렬로 실행되므로 추론 속도가 크게 빨라집니다. 이렇게 하면 타겟 모델이 각 토큰에 대해 실행해야 하는 정방향 패스 수가 줄어듭니다. Google의 초안 작성기는 확인을 위해 토큰 시퀀스를 생성하므로 이를 다중 토큰 예측 (MTP) 헤드라고 합니다.

png

Gemma 4 제품군용으로 출시된 초안 모델은 작으며, 타겟 모델 활성화 및 KV-cache를 사용하여 더 나은 예측을 얻는 등 초안 토큰의 품질을 개선하고 추론 속도를 더욱 높이기 위한 여러 개선사항이 도입되었습니다.

이러한 개선사항은 유사한 품질을 보장하면서 디코딩 속도를 크게 높여 이러한 체크포인트가 짧은 지연 시간 및 온디바이스 애플리케이션에 적합합니다.

Python 패키지 설치

Gemma 4 및 Gemma 4 어시스턴트 모델을 실행하는 데 필요한 Hugging Face 라이브러리를 설치합니다.

# Install PyTorch & other libraries
pip install torch accelerate

# Install the transformers library
pip install transformers

모델 로드

각 타겟 모델 (Gemma 4 모델의 기본 모델 중 하나)에는 추론 속도를 높이는 데 도움이 되는 어시스턴트가 있습니다. 따라서 두 모델을 로드합니다.

타겟 (예: google/gemma-4-E2B-it): 전체 Gemma 4 타겟 모델
Drafter (예: google/gemma-4-E2B-it-assistant): 후보 토큰을 제안하는 경량 4계층 MTP drafter

모델이 더 큰 모델이 예측할 토큰을 선택하는 데 도움을 주기 때문에 초안 작성기는 어시스턴트라고도 합니다.

transformers 라이브러리를 사용하여 다음 코드 예와 같이 AutoProcessor 및 AutoModelForCausalLM 클래스를 사용하여 processor 및 model의 인스턴스를 만듭니다.

TARGET_MODEL_ID = "google/gemma-4-E2B-it" # @param ["google/gemma-4-E2B-it","google/gemma-4-E4B-it", "google/gemma-4-31B-it", "google/gemma-4-26B-A4B-it"]
ASSISTANT_MODEL_ID = TARGET_MODEL_ID + "-assistant"

import torch
from transformers import AutoProcessor, AutoModelForCausalLM

# Target Model
processor = AutoProcessor.from_pretrained(TARGET_MODEL_ID)
target_model = AutoModelForCausalLM.from_pretrained(
    TARGET_MODEL_ID,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

# Assistant Model (the drafter)
assistant_model = AutoModelForCausalLM.from_pretrained(
    ASSISTANT_MODEL_ID,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

[transformers] `torch_dtype` is deprecated! Use `dtype` instead!
Loading weights:   0%|          | 0/1951 [00:00<?, ?it/s]
Loading weights:   0%|          | 0/50 [00:00<?, ?it/s]

어시스턴트가 탑재된 Gemma 4

transformers에서 어시스턴트를 사용하는 것은 다행히 매우 간단하며 어시스턴트 모델을 model.generate 함수에 전달해야 합니다.

# Process inputs with the `target_model`
messages = [
    {
        "role": "user",
        "content": "Explain the concepts of speculative decoding and MTP in 3 sentences."
    }
]
input_text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = processor(text=input_text, return_tensors="pt").to(target_model.device)

# `assistant_model=assistant_model` is all you need to enable MTP!
outputs = target_model.generate(
    **inputs,
    assistant_model=assistant_model,
    max_new_tokens=256,
    do_sample=False,
)

# Decode the response into text
response = processor.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
print(response)

**Speculative decoding** is a technique where a smaller, faster language model (the "draft model") generates several candidate tokens, which are then quickly verified by a larger, more accurate model to produce a final, high-quality output much faster than decoding the large model alone. **MTP (Multi-Task Prediction)** involves training a single model to perform multiple related tasks simultaneously, allowing it to leverage shared knowledge across different objectives. Together, these methods aim to significantly accelerate the inference speed of large language models while maintaining or improving output quality.

내부적으로 프로세스는 다음과 같습니다.

초안 작성자는 자동 회귀적으로 생성된 N개의 토큰을 제안합니다.
타겟 모델은 하나의 순방향 패스에서 모든 N 토큰을 검증합니다.
확률이 높은 초안 토큰이 허용됨
확률이 낮은 초안 토큰은 거부됨
타겟 모델은 정방향 전달을 실행하므로 승인되거나 거부된 초안 토큰 수와 관계없이 항상 토큰 1개를 자체적으로 생성합니다.

초안 토큰

초안 작성자는 타겟 모델이 확인할 수 있도록 원하는 만큼의 토큰을 생성할 수 있습니다. 하지만 타겟 모델은 여전히 특정 토큰을 거부할 수 있습니다. 이 경우 이후의 모든 토큰은 무시됩니다.

png

따라서 초안 토큰 수에 다양한 값을 사용할 때의 절충안을 알아두는 것이 중요합니다.

더 많은 초안 토큰

토큰을 많이 작성하면 (예: 15개) 일부 토큰이 수락되지 않을 가능성이 높습니다. 따라서 낭비되는 컴퓨팅 리소스가 더 많을 수 있습니다. 반대로 수락률이 높으면 추론 속도가 빨라지는 경향이 있습니다.

png

더 적은 초안 토큰

토큰을 적게 작성하면 초기 프롬프트와 위치가 더 가까운 토큰이 더 정확하므로 수락률이 높아지는 경향이 있습니다. 하지만 일부 토큰만 초안이 작성되므로 더 빠른 초안 작성 모델에서 얻을 수 있는 속도 향상이 줄어듭니다.

png

다행히도 transformers에서 사용 사례에 가장 적합한 값을 실험할 필요가 없습니다. num_assistant_tokens_schedule을 '휴리스틱'으로 설정하면 런타임에 초안 토큰 수가 자동으로 조정되기 때문입니다.

모든 토큰이 허용됨: 프롬프트에 대한 초안 작성자의 정확도가 높으므로 초안 작성할 토큰 수를 2개 늘립니다. 드래프트된 토큰 수를 늘리면 해당 토큰도 수락되는 경우 속도가 빨라질 수 있습니다.
거부된 토큰 - 거부된 토큰이 있으면 초안을 작성할 토큰 수를 1개 줄입니다. 토큰 수를 줄이면 타겟 모델이 대부분의 토큰을 계속 거부하더라도 낭비되는 초안이 너무 많지 않습니다.

마찬가지로 다음과 같이 drafter에서 num_assistant_tokens를 업데이트하여 초안 토큰 수를 업데이트할 수 있습니다.

# Update how many draft tokens are generated at the start of inference
assistant_model.generation_config.num_assistant_tokens = 4

# Update how the number of draft tokens are updated ("heuristic" for a dynamic schedule and "constant" for a constant schedule)
assistant_model.generation_config.num_assistant_tokens_schedule = "heuristic"

# Run with MTP
outputs = target_model.generate(
    **inputs,
    assistant_model=assistant_model,
    max_new_tokens=256,
    do_sample=False,
)

# Decode the response into text
response = processor.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
print(response)

**Speculative decoding** is a technique where a smaller, faster language model (the "draft model") generates several candidate tokens, which are then verified by a larger, more accurate model to quickly produce a high-quality output. **MTP (Multi-Task Prediction)** involves training a single model to perform multiple related tasks simultaneously, allowing it to leverage shared knowledge across different objectives. Together, these methods aim to significantly speed up the inference process of large language models while maintaining or improving output quality.