Deviner, puis vérifier

Quatrième article de la série Smallem : servir un modèle. Précédent : Huit bits, puis quatre.

Le deuxième article a montré une chose étrange : une passe du modèle sur huit tokens coûte presque autant que sur un seul, 1,19 ms contre 1,02. Il s'en servait pour générer plusieurs séquences à la fois. Mais une seule séquence peut-elle en profiter ?

À l'entraînement, le modèle lit une séquence entière en une passe, et chaque position y prédit la suivante. En génération, on ne peut pas faire pareil : pour lire le token n+1, il faut d'abord l'avoir tiré, donc avoir fini la passe du token n. Sauf si quelqu'un devine les tokens à venir. Le modèle n'a alors plus qu'à vérifier, en une passe, toutes les devinettes à la fois.

C'est le décodage spéculatif (Leviathan, Kalman et Matias, 2023 ; Chen et al., 2023).

Vérifier sans changer le tirage

En mode glouton, la règle est évidente. Un proposeur bon marché devine k tokens. Le modèle les lit en une passe, qui donne sa distribution à chacune des k + 1 positions. On garde les propositions tant qu'elles sont son premier choix. Au premier désaccord, on prend le sien. Une passe donne ainsi entre 1 et k + 1 tokens, et exactement ceux que le modèle aurait choisis seul.

En tirage, garder les propositions « qui auraient pu être tirées » ne suffit pas : ce serait un autre tirage, biaisé vers ce que le proposeur aime. Le proposeur a tiré x selon sa distribution q. Le modèle visé donne la sienne, p. On accepte x avec la probabilité

min(1, p(x) / q(x))

Un token que le proposeur sous-estime est toujours gardé. Un token qu'il surestime est gardé dans la proportion où le modèle visé le tirerait. Au premier refus, on tire un token de la distribution corrigée, max(0, p − q), renormalisée : elle rend exactement la probabilité que les refus ont retirée. Si tout est accepté, la passe donne en prime un token de plus, tiré de p à la dernière position.

La probabilité de sortir x à une position est alors la somme de deux chemins : x est proposé et accepté, ou une proposition est refusée et x est tiré de la correction. La somme vaut p(x), quel que soit q. Le texte produit suit exactement la distribution du modèle visé : le proposeur ne change que la vitesse.

let (p, q) = (target[i][token], draft[i][token]);
let keep: f64 = rng.random();
if keep * q >= p {
    // Refus : la distribution corrigée, `max(0, p − q)`.
    …
}

Une telle propriété se vérifie par la statistique. Un test tire 200 000 paires de tokens sur un vocabulaire de cinq, avec un proposeur et un modèle visé qui ne s'accordent pas, et dont le second token dépend du premier. Il compare les 25 fréquences à celles du modèle visé seul par un test du khi-deux : au seuil de 0,1 %, avec 24 degrés de liberté, le khi-deux doit rester sous 51,2. Il y reste. Le même test, appliqué à la distribution du proposeur, dépasse 1 000 : il sait reconnaître une distribution fausse.

Revenir en arrière

Les propositions refusées ont traversé le modèle visé, et le proposeur lui-même a lu les siennes. Leurs clés et leurs valeurs sont dans les caches. Il faut les oublier.

Avec le cache du premier article, c'est une ligne : une séquence ne fait que pointer vers des lignes de tables, et revenir en arrière, c'est raccourcir sa liste. Les pages devenues vides retournent à la réserve.

self.target.keep(start + verdict.accepted);
self.draft.keep(self.draft.seen.len().min(start + verdict.accepted));

Deux proposeurs

Le plan en prévoyait deux.

Un petit modèle. Deux, en fait, entraînés pour l'occasion avec le même tokenizer : la configuration large de la première série, largeur 384, avec 2 ou 4 couches. Neuf minutes d'entraînement à eux deux. Ils finissent à 1,055 et 0,995 nat par octet, contre 0,882 pour le modèle visé. Ils génèrent seuls 3 264 et 1 992 tokens par seconde : un token leur coûte 0,31 et 0,50 passe du modèle visé.

Le module de prédiction multi-token de l'article 15 de la première série. À la position i, il reçoit l'état caché h_i du modèle et le token suivant, et prédit celui d'après. À l'entraînement, il ne servait qu'à enrichir le signal. En génération, il avait été jeté. DeepSeek-V3 s'en sert justement comme proposeur, et c'est l'occasion de lui donner un usage. Chaque passe du modèle lit le dernier token accepté et la proposition qui le suit : elle vérifie la proposition, et rend les états cachés d'où le module tire la suivante. Le proposeur ne coûte qu'un bloc et la sortie, et il voit tout ce que le modèle a calculé. Son bloc a sa propre attention, donc son propre cache, d'un étage.

Le modèle à module de la première série n'est pas le modèle visé des autres mesures : c'est celui de l'article 15, entraîné 16 minutes. Il se compare à lui-même.

Les mesures

Huit amorces de 32 tokens tirées du texte de validation, 200 tokens chacune. En mode glouton d'abord, où les huit textes sont, à chaque ligne, identiques à ceux du modèle seul :

Le modèle seul fait 1 006 tokens par seconde ; avec son module, 1 545.

Le coût d'un tour l'explique. Avec un proposeur qui coûte c passes par token, un tour coûte k × c plus une passe de vérification, à peine plus chère qu'une passe d'un token. Il rapporte en moyenne (1 − α^(k+1)) / (1 − α) tokens si chaque proposition est acceptée avec la probabilité α, indépendamment des autres. Le petit proposeur, avec k = 2 : 2 × 0,31 + 1,05 ≈ 1,67 passe pour 2,17 tokens, soit ×1,30, contre ×1,34 mesuré.

Le proposeur à 4 couches devine mieux, mais coûte trop : pour chaque token gagné, il en dépense la moitié d'un. Et au-delà de 3 ou 4 propositions, les dernières sont presque toujours refusées : on paie leur calcul pour rien.

Le module, lui, a tout pour lui. Il coûte un huitième du modèle, et devine juste 94 fois sur 100 : il lit les états cachés du modèle lui-même, qui contiennent déjà, l'entraînement l'y a poussé, de quoi prévoir deux tokens.

En tirage, à température 1 :

Tout baisse. En tirage, deux modèles qui hésitent entre plusieurs tokens tombent moins souvent d'accord qu'en prenant chacun son premier choix. Et l'hôte travaille plus : chaque vérification construit des distributions complètes de 8 192 tokens, des deux côtés. Les textes ne sont plus identiques à ceux du modèle seul, puisque les tirages diffèrent ; c'est le test du khi-deux qui garantit qu'ils suivent la même loi.

Ce qu'il faut retenir

• Vérifier k devinettes coûte une passe, parce qu'une passe sur quelques tokens coûte à peine plus que sur un seul. Le décodage spéculatif transforme ce temps fixe en tokens.
• La règle min(1, p/q), avec sa correction max(0, p − q), ne change pas la distribution du texte. Un test du khi-deux le vérifie.
• Le meilleur proposeur n'est pas celui qui devine le mieux, mais celui qui devine assez bien pour pas cher. Notre proposeur à 4 couches perd contre celui à 2.
• Le module de prédiction multi-token, inutile en génération jusqu'ici, fait le meilleur proposeur : ×1,53 en mode glouton, ×1,30 en tirage.

Article suivant : Servir plusieurs requêtes.