Decodificação assistida

O decodificação assistida acelera a geração de texto ao permitir que um modelo auxiliar proponha tokens candidatos antes que o modelo principal os confirme. O modelo principal verifica os tokens candidatos em uma única passagem de frente. O modelo auxiliar é rápido e barato e pode substituir dezenas de passagens de frente mais caras do modelo principal.

Este guia aborda os métodos de decodificação assistida no Transformers.

Decodificação especulativa

O decodificação especulativa usa um modelo assistente menor para rascunhar tokens candidatos. O modelo principal verifica esses tokens em uma única passagem. Os tokens validados entram na saída final e os tokens rejeitados disparam a amostragem padrão. A geração é mais rápida porque o modelo principal executa menos passes de frente caros.

O método funciona melhor quando o modelo assistente é significativamente menor que o modelo principal e usa o mesmo tokenizer. O speculative decoding suporta busca gulosa e amostragem, mas não entradas em lote.

Passe assistant_model para generate(). Defina do_sample=True para reamostrar caso a validação do token falhe.

busca gulosa

amostragem

Decodificação por busca de prompt

O prompt lookup decoding não requer um modelo assistente. Ele localiza n-grams sobrepostos no prompt para propor tokens candidatos. Se não houver correspondência, ele recorre ao decodificação autoregressiva normal. Isso se adequa a tarefas baseadas na entrada, como sumarização e tradução, pois os tokens candidatos frequentemente espelham padrões locais no texto de origem.

Passe prompt_lookup_num_tokens para generate(). Isso define quantos tokens o algoritmo tenta copiar de uma parte anterior do prompt quando detecta um padrão repetido.

decodificação gulosa

amostragem

Decodificação especulativa autônoma

O auto-decodificação especulativa usa as camadas intermediárias de um modelo como assistente para propor tokens candidatos. Se a proposta corresponder, o modelo encerra precocemente e as camadas restantes verificam ou corrigem os tokens.

Como tudo é um único modelo, os pesos e os caches são compartilhados, o que aumenta a velocidade sem sobrecarga adicional de memória. Essa técnica só funciona para modelos treinados para suportar logits de saída antecipada a partir de camadas intermediárias.

Passe assistant_early_exit para generate() para definir a camada de saída.

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)

Camada de previsão de múltiplos tokens (MTP)

A camada de previsão de múltiplos tokens (MTP) rascunha tokens candidatos com camadas de previsão adicionais que provêm do próprio checkpoint do modelo principal. As camadas MTP reutilizam os embeddings e a cabeça de saída do modelo principal. Cada camada MTP rascunha um token, de modo que o número de candidatos por etapa corresponde à quantidade de camadas MTP no checkpoint. O modelo principal verifica os rascunhos em uma única passagem de frente.

O MTP funciona apenas com checkpoints treinados com camadas MTP, como DeepSeek-V3 e GLM-4.5. Esses checkpoints contêm pesos MTP e definem num_mtp_layers em sua configuração, listado como num_nextn_predict_layers em config.json.

Passe use_mtp=True para generate(), que gera um erro se o modelo não tiver camadas MTP. Assim como o speculative decoding, o MTP suporta busca gulosa e amostragem, mas não suporta entradas em lote.

Aqui, MTP se refere às camadas de previsão incluídas no checkpoint de destino, como em DeepSeek-V3 e GLM-4.5. O Gemma 4 usa um checkpoint de assistente separado para decodificação MTP. Passe-o como assistant_model (veja Drafter de previsão de múltiplos tokens).

O pipeline text-generation e o generate() aceitam 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]

Alternativamente, você pode inferir da seguinte forma.

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)

Drafter de previsão de múltiplos tokens

Um redator MTP é um modelo assistente leve carregado a partir de um checkpoint separado. Ele propõe tokens candidatos para que o modelo de destino os verifique. O redator usa os embeddings de entrada do modelo de destino e atende ao seu cache de KV. Ele parte dos estados ocultos do modelo de destino e, em seguida, redige tokens de forma autoregressiva. Para cada próximo token, ele alimenta seu estado oculto projetado de volta ao redator. O modelo de destino verifica os tokens redigidos em uma única passagem de frente, como no decodificação especulativa.

Diferentemente das camadas MTP, que fazem parte do checkpoint de destino e compartilham seu cabeçalho de saída, um redator MTP é um modelo assistente separado. Ele usa os embeddings de entrada e o cache KV do modelo de destino para propor tokens.

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)

Decodificação assistida universal

O decodificação assistida universal (UAD) torna a decodificação especulativa possível mesmo quando os modelos principal e assistente possuem tokenizers diferentes. Ele permite combinar qualquer modelo assistente pequeno com o modelo principal. Os tokens candidatos são re-codificados e o algoritmo calcula a subsequência comum mais longa para que a continuação permaneça alinhada.

Passe tokenizer, assistant_tokenizer e assistant_model para generate() para habilitar o UAD.

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.'

Verificação de ensemble estático

O decodificação especulativa padrão é sem perdas — ela garante que a distribuição de saída corresponda exatamente à do modelo de destino. A verificação de ensemble estática relaxa isso ao verificar contra uma mistura das distribuições de destino e de rascunho, aumentando a taxa de aceitação (inferência mais rápida) ao custo de tornar a distribuição de saída uma mistura, em vez da distribuição de destino exata. Um valor inicial recomendado é 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)

Isso requer logits candidatas do modelo assistente e não é suportado com prompt lookup decoding.

Recursos

Leia o post no blog Assisted Generation: a new direction toward low-latency text generation para mais contexto sobre a latência de geração de texto e a geração assistida.

Update on GitHub