Décodage assisté

Le décodage assisté accélère la génération de texte en permettant à un modèle auxiliaire de proposer des jetons candidats avant que le modèle principal ne s’y engage. Le modèle principal vérifie les jetons candidats en un seul passage avant. L’auxiliaire est rapide et peu coûteux et peut remplacer des dizaines de passages avant plus onéreux du modèle principal.

Ce guide présente les méthodes de décodage assisté dans Transformers.

Décodage spéculatif

Speculative decoding utilise un modèle assistant plus petit pour rédiger des jetons candidats. Le modèle principal vérifie ces jetons en un seul passage. Les jetons validés entrent dans la sortie finale et les jetons rejetés déclenchent un échantillonnage standard. La génération est plus rapide car le modèle principal effectue moins de passages avant coûteux.

La méthode fonctionne de manière optimale lorsque le modèle assistant est nettement plus petit que le modèle principal et utilise le même tokenizer. Le décodage spéculatif prend en charge la recherche gloutonne et l'échantillonnage, mais pas les entrées par lots.

Passez assistant_model à generate(). Définissez do_sample=True pour rééchantillonner si la validation des jetons échoue.

recherche gloutonne

échantillonnage

Décodage par recherche de prompt

Le décodage par recherche de prompt n’a pas besoin d’un modèle assistant. Il trouve des n-grammes superposés dans le prompt pour proposer des jetons candidats. En l’absence de correspondance, il bascule sur un décodage autorégressif standard. Cette approche convient aux tâches ancrées dans l’entrée, comme la résumée et la traduction, car les jetons candidats reflètent souvent les motifs locaux du texte source.

Passez prompt_lookup_num_tokens à generate(). Cela définit le nombre de jetons que l'algorithme tente de copier depuis le début du prompt lorsqu'il détecte un motif répété.

décodage glouton

échantillonnage

Décodage spéculatif auto

Le décodage spéculatif auto-référentiel utilise les couches intermédiaires d'un modèle comme assistant pour proposer des jetons candidats. Si la proposition correspond, le modèle sort en avance et les couches restantes vérifient ou corrigent les jetons.

Puisqu’il s’agit d’un seul et même modèle, les poids et les caches sont partagés, ce qui augmente la vitesse sans surcoût mémoire supplémentaire. Cette technique ne fonctionne que pour les modèles entraînés pour prendre en charge les logits de sortie anticipée à partir des couches intermédiaires.

Passez assistant_early_exit à generate() pour définir la couche de sortie.

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)

Couche de prédiction multi-tokens (MTP)

La couche de prédiction multi-tokens (MTP) génère des tokens candidats à l'aide de couches de prédiction supplémentaires issues du checkpoint du modèle principal. Les couches MTP réutilisent les embeddings et la tête de sortie du modèle principal. Chaque couche MTP génère un token, de sorte que le nombre de candidats par étape correspond au nombre de couches MTP du checkpoint. Le modèle principal vérifie les brouillons en un seul passage avant.

MTP ne fonctionne qu’avec des points de contrôle entraînés avec des couches MTP, tels que DeepSeek-V3 et GLM-4.5. Ces points de contrôle contiennent les poids MTP et définissent num_mtp_layers dans leur configuration, répertorié sous le nom num_nextn_predict_layers dans config.json.

Passez use_mtp=True à generate(), ce qui lève une erreur si le modèle ne possède pas de couches MTP. Comme le décodage spéculatif, MTP prend en charge la recherche gloutonne et l’échantillonnage, mais pas les entrées par lots.

Ici, MTP désigne les couches de prédiction incluses dans le checkpoint cible, comme dans DeepSeek-V3 et GLM-4.5. Gemma 4 utilise un checkpoint assistant distinct pour le décodage MTP. Transmettez-le en tant que assistant_model (voir Drafter de prédiction multi-tokens).

Le pipeline text-generation et generate() acceptent tous deux 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]

Vous pouvez également inférer comme suit.

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 prédiction multi-tokens

Un rédacteur MTP est un modèle assistant léger chargé à partir d’un point de contrôle distinct. Il propose des jetons candidats que le modèle cible doit vérifier. Le rédacteur utilise les plongements d’entrée du modèle cible et porte son attention sur son cache KV. Il part des états cachés du modèle cible, puis rédige des jetons de manière autorégressive. Pour chaque jeton suivant, il réinjecte son état caché projeté dans le rédacteur. Le modèle cible vérifie les jetons rédigés en un seul passage avant, comme dans le décodage spéculatif.

Contrairement aux couches MTP, qui font partie du point de contrôle cible et partagent sa tête de sortie, un rédacteur MTP est un modèle assistant distinct. Il utilise les plongées d’entrée et le cache KV de la cible pour proposer des jetons.

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)

Décodage assisté universel

Le décodage assisté universel (UAD) rend le décodage spéculatif possible même lorsque les modèles principal et assistant utilisent des tokenizers différents. Il permet d’associer n’importe quel petit modèle assistant au modèle principal. Les jetons candidats sont réencodés et l’algorithme calcule la plus longue sous-séquence commune afin que la continuation reste alignée.

Passez tokenizer, assistant_tokenizer et assistant_model à generate() pour activer 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.'

Vérification statique d'ensembles

Le décodage spéculatif standard est sans perte — il garantit que la distribution de sortie correspond exactement à celle du modèle cible. La vérification d'ensemble statique assouplit cette contrainte en vérifiant par rapport à un mélange des distributions cible et de brouillon, ce qui augmente le taux d'acceptation (inférence plus rapide) au prix de rendre la distribution de sortie un mélange plutôt que la distribution cible exacte. Une valeur de départ recommandée est 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)

Cela nécessite des logits candidats du modèle assistant et n'est pas pris en charge avec le décodage par recherche de prompt.

Ressources

Lisez le billet de blog Assisted Generation: a new direction toward low-latency text generation pour plus de contexte sur la latence de la génération de texte et la génération assistée.

Update on GitHub