Decodificación asistida

La decodificación asistida acelera la generación de texto al permitir que un modelo auxiliar proponga tokens candidatos antes de que el modelo principal los confirme. El modelo principal verifica los tokens candidatos en una sola pasada hacia adelante. El modelo auxiliar es rápido y económico, y puede reemplazar a decenas de pasadas hacia adelante más costosas del modelo principal.

Esta guía cubre los métodos de decodificación asistida en Transformers.

Decodificación especulativa

Decodificación especulativa utiliza un modelo asistente más pequeño para redactar tokens candidatos. El modelo principal verifica estos tokens en una sola pasada. Los tokens validados entran en la salida final y los tokens rechazados activan el muestreo estándar. La generación es más rápida porque el modelo principal realiza menos pasadas hacia adelante costosas.

El método funciona mejor cuando el modelo asistente es significativamente más pequeño que el modelo principal y utiliza el mismo tokenizer. La decodificación especulativa admite la búsqueda codiciosa y el muestreo, pero no las entradas por lotes.

Pasa assistant_model a generate(). Establece do_sample=True para re-muestrear si la validación de tokens falla.

búsqueda voraz

sampling

Decodificación mediante búsqueda de indicaciones

El decodificado mediante búsqueda en el prompt no requiere un modelo asistente. Encuentra n-gramas superpuestos en el prompt para proponer tokens candidatos. Si no existe coincidencia, recurre a la decodificación autoregresiva normal. Esto es adecuado para tareas basadas en la entrada, como la resumición y la traducción, porque los tokens candidatos suelen reflejar patrones locales del texto de origen.

Pase prompt_lookup_num_tokens a generate(). Esto define cuántos tokens el algoritmo intenta copiar de una parte anterior del prompt cuando detecta un patrón repetido.

decodificación codiciosa

sampling

Decodificación especulativa auto

La decodificación especulativa autocontenida utiliza las capas intermedias de un modelo como asistente para proponer tokens candidatos. Si la propuesta coincide, el modelo finaliza de forma anticipada y las capas restantes verifican o corrigen los tokens.

Dado que todo es un solo modelo, los pesos y las cachés se comparten, lo que aumenta la velocidad sin un sobrecosto adicional de memoria. Esta técnica solo funciona para modelos entrenados para admitir logits de salida temprana desde capas intermedias.

Pasa assistant_early_exit a generate() para establecer la capa de salida.

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)

Capa de predicción de múltiples tokens (MTP)

La capa de predicción de múltiples tokens (MTP) genera tokens candidatos mediante capas de predicción adicionales que provienen del propio checkpoint del modelo principal. Las capas MTP reutilizan los embeddings y la cabeza de salida del modelo principal. Cada capa MTP genera un token, por lo que el número de candidatos por paso coincide con la cantidad de capas MTP del checkpoint. El modelo principal verifica los borradores en una sola pasada hacia adelante.

MTP solo funciona con puntos de control entrenados con capas MTP, como DeepSeek-V3 y GLM-4.5. Estos puntos de control incluyen pesos MTP y establecen num_mtp_layers en su configuración, que aparece como num_nextn_predict_layers en config.json.

Pase use_mtp=True a generate(), lo que genera un error si el modelo no tiene capas MTP. Al igual que la decodificación especulativa, MTP admite búsqueda codiciosa y muestreo, pero no entradas por lotes.

Aquí, MTP se refiere a las capas de predicción incluidas en el checkpoint objetivo, como en DeepSeek-V3 y GLM-4.5. Gemma 4 utiliza un checkpoint de asistente separado para la decodificación MTP. Pásalo como assistant_model (ver Drafter de predicción de múltiples tokens).

El pipeline text-generation y generate() aceptan ambos 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, puede inferir de la siguiente manera.

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 predicción de múltiples tokens

Un redactor MTP es un modelo asistente ligero que se carga desde un punto de control (checkpoint) separado. Propone tokens candidatos para que el modelo objetivo los verifique. El redactor utiliza las incrustaciones de entrada del modelo objetivo y atiende a su caché KV. Comienza con los estados ocultos del modelo objetivo y luego redacta tokens de forma autoregresiva. Para cada token siguiente, reintroduce su estado oculto proyectado en el redactor. El modelo objetivo verifica los tokens redactados en una sola pasada hacia adelante, tal como en la decodificación especulativa.

A diferencia de las capas MTP, que forman parte del checkpoint objetivo y comparten su cabeza de salida, un redactor MTP es un modelo asistente independiente. Utiliza los embeddings de entrada y la caché KV del modelo objetivo para proponer 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)

Decodificación asistida universal

El decodificado asistido universal (UAD) hace posible el decodificado especulativo incluso cuando los modelos principal y asistente tienen tokenizadores diferentes. Permite emparejar cualquier modelo asistente pequeño con el modelo principal. Los tokens candidatos se re-codifican y el algoritmo calcula la subsecuencia común más larga para que la continuación permanezca alineada.

Pase tokenizer, assistant_tokenizer y assistant_model a generate() para habilitar 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.'

Verificación de conjuntos estáticos

La decodificación especulativa estándar es sin pérdida: garantiza que la distribución de salida coincida exactamente con la del modelo objetivo. La verificación de conjunto estática relaja esta condición al verificar contra una mezcla de las distribuciones del modelo objetivo y del borrador, lo que aumenta la tasa de aceptación (inferencia más rápida) a costa de que la distribución de salida sea una mezcla en lugar de la distribución objetivo exacta. Un valor inicial recomendado es 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)

Esto requiere logits de candidatos del modelo asistente y no se admite con la decodificación mediante búsqueda de indicio.

Recursos

Lea la publicación de blog Assisted Generation: a new direction toward low-latency text generation para obtener más contexto sobre la latencia de la generación de texto y la generación asistida.

Update on GitHub