Prédire plus loin, optimiser autrement

Quinzième article de la série Smallem. Précédent : Un mélange d'experts.

Les trois articles précédents ont changé des couches. Celui-ci change deux choses qui n'en sont pas : ce qu'on demande au modèle pendant l'entraînement, et la façon dont un gradient devient une mise à jour. Il se termine par l'inventaire de ce qui, dans les modèles DeepSeek les plus récents, n'a pas été transposé, et pourquoi.

La prédiction multi-token

L'idée

Un modèle de langage apprend à prédire le token suivant. À chaque position, il reçoit un seul signal : le bon token était celui-ci.

DeepSeek-V3 lui demande en plus de prédire le token d'après. L'idée a été étudiée d'abord par Gloeckle et al. (2024), avec des têtes parallèles, DeepSeek enchaîne les prédictions. Le signal d'entraînement est plus dense : deux réponses à trouver par position au lieu d'une. Et pour prévoir deux tokens d'avance, un état interne doit contenir plus qu'une simple continuation locale : le modèle est poussé à anticiper.

Cette seconde prédiction ne sert que pendant l'entraînement. Elle est confiée à un module d'appoint, qu'on retire ensuite.

Le module

À la position i, le modèle principal a lu les tokens jusqu'à t_i. Son dernier bloc produit un état caché h_i, d'où sort la prédiction de t_{i+1}.

Le module reprend cet état, y joint l'embedding du vrai token suivant, et prédit t_{i+2} :

x      = [rms(h_i) ; rms(embedding(t_{i+1}))] · W      de 2C à C
h'_i   = bloc(x)                                        un bloc de Transformer
logits = sortie du modèle principal, appliquée à h'_i

Trois choses sont à remarquer.

• Les deux vecteurs sont mis bout à bout : c'est la concaténation écrite pour l'attention latente à l'article 13, qui trouve ici son second usage.
• Le module ne possède qu'une projection et un bloc. La table d'embeddings et la sortie sont celles du modèle principal.
• Le gradient de la seconde prédiction remonte dans h_i, donc dans tout le modèle. C'est par là que le module agit : il ne sert pas à prédire, il sert à façonner les états cachés.

La perte d'entraînement est la somme des deux, la seconde pesant moins :

perte = perte principale + λ × perte du module

λ vaut 0,3 pendant les deux premiers tiers de l'entraînement, puis 0,1.

Dans le code

Le modèle expose désormais deux pertes.

/// Celle qui se compare d'un modèle à l'autre.
pub fn loss(&self, tape, tokens, targets, batch) -> Var

/// Celle qu'on minimise : la précédente, plus les termes d'appoint.
pub fn training_loss(&self, tape, tokens, targets, batch) -> Var

La distinction existait déjà en germe avec la pénalité d'équilibrage de l'article 14. Elle est essentielle ici : la perte de validation doit rester celle du modèle principal, sans quoi on comparerait des quantités différentes d'une variante à l'autre.

Le piège du module est un décalage d'indices. Les cibles d'un lot sont les tokens suivants : targets[i] est t_{i+1}. Le module reçoit donc targets[i] en entrée et doit prédire targets[i + 1]. La dernière position de chaque séquence n'a pas de cible aussi lointaine, et il faut la laisser de côté, sans aller chercher la sienne dans la séquence suivante du lot :

// Toutes les positions sauf la dernière de chaque séquence.
let rows: Vec<u32> = (0..batch)
    .flat_map(|sequence| (0..length - 1).map(move |t| (sequence * length + t) as u32))
    .collect();
let next: Vec<u32> = rows.iter().map(|&row| targets[row as usize]).collect();
let after: Vec<u32> = rows.iter().map(|&row| targets[row as usize + 1]).collect();

Un test garde cette frontière : la perte du module sur un lot de deux séquences doit être la moyenne de ses pertes sur chacune. Si une position débordait sur la séquence voisine, l'égalité casserait.

Retirer le module

Un fichier de poids entraîné avec le module contient ses paramètres. Pour générer, on construit le modèle sans lui :

let config = GptConfig::from_metadata(&file.metadata)?.for_inference();

Le chargement lit les poids par leur nom : ceux du module, que le modèle ne réclame pas, sont ignorés. Un test vérifie que le modèle allégé donne les mêmes logits que le modèle complet.

Le résultat

Même protocole : 6 000 pas, une graine, bruit d'environ 0,002 nat par octet. La perte comparée est celle du modèle principal seul.

Le gain existe, mais il est mince : un peu moins de deux fois le bruit, avec une seule graine. Il coûte 19 % de débit, puisque le module ajoute un neuvième bloc à chaque pas. La comparaison à temps d'entraînement égal n'a pas été faite, et rien ne dit que le gain y survivrait.

Le modèle qu'on garde, lui, n'a pas grossi : une fois le module retiré, il a ses 29,4 millions de paramètres et sa vitesse de génération d'origine.

DeepSeek rapporte un gain sur un modèle vingt mille fois plus gros. À notre échelle, le signal est favorable et faible.

Muon

Ce que fait AdamW, et ce qu'il ignore

AdamW, depuis l'article 7, règle le pas de chaque nombre séparément : il divise la moyenne récente de son gradient par la racine de la moyenne récente de son carré. Pour lui, une matrice de poids est un sac de nombres.

Or une matrice de poids est une transformation, et son gradient en est une aussi. Toute matrice se décompose en directions d'entrée, directions de sortie, et un facteur d'amplification pour chacune : ses valeurs singulières. Le gradient d'une matrice est typiquement dominé par quelques directions, celles où la perte varie le plus. Les autres existent, mais leur facteur est minuscule.

L'idée

Muon, proposé par Keller Jordan (2024) et utilisé par DeepSeek-V4, garde les directions du gradient et jette les facteurs. Il remplace la moyenne des gradients par la matrice orthogonale la plus proche : mêmes directions, mais toutes les valeurs singulières ramenées à 1.

M = μ·M + G                                         moyenne glissante du gradient
O = orthogonaliser(μ·M + G) × √max(lignes, colonnes) × γ
W = W × (1 - η·λ) - η·O

Toutes les directions avancent alors du même pas. Celles que le gradient signalait faiblement, et qu'AdamW aurait négligées, progressent autant que les autres.

Le facteur √max(lignes, colonnes) × γ est une mise à l'échelle. Une matrice orthogonale a des éléments de taille 1/√max(…) ; ce facteur ramène ceux de la mise à jour à la taille γ = 0,18, proche de celle d'un pas d'AdamW. Le même taux d'apprentissage sert alors aux deux optimiseurs.

Muon ne s'applique qu'aux matrices de poids. Les embeddings, les gains et les biais restent à AdamW.

Orthogonaliser sans décomposer

Calculer exactement la matrice orthogonale la plus proche demande une décomposition en valeurs singulières, coûteuse et mal adaptée à un GPU. Muon s'en passe avec une itération qui n'utilise que des produits matriciels :

X ← a·X + b·(X·Xᵀ)·X + c·(X·Xᵀ)²·X

Cette formule ne touche pas aux directions de X. Elle applique à chacune de ses valeurs singulières s le polynôme a·s + b·s³ + c·s⁵. Avec de bons coefficients, répéter ce polynôme pousse toute valeur entre 0 et 1 vers 1.

DeepSeek-V4 fait dix itérations en deux temps. Huit avec des coefficients choisis pour faire grandir très vite les petites valeurs, sans chercher la précision. Puis deux avec (2 ; -1,5 ; 0,5), un polynôme qui vaut 1 en 1 et y est plat : il finit le travail proprement.

for (a, b, c) in NEWTON_SCHULZ {
    let transposed = backend.transpose(&x, 0, 1);
    let gram = if wide {
        backend.matmul(&x, &transposed)
    } else {
        backend.matmul(&transposed, &x)
    };
    // `b·A + c·A²`, puis son produit avec `X` du bon côté.
    let polynomial = backend.add(
        &backend.mul(&gram, &scalar(b)),
        &backend.mul(&backend.matmul(&gram, &gram), &scalar(c)),
    );
    // …
    x = backend.add(&backend.mul(&x, &scalar(a)), &correction);
}

Aucun kernel nouveau : tout l'optimiseur est écrit avec les primitives de l'article 4. Le code choisit le plus petit des deux carrés, X·Xᵀ ou Xᵀ·X, selon que la matrice est large ou haute.

Les tests vérifient ce que la théorie promet : le résultat est orthogonal, une rotation étirée redevient la rotation, multiplier la matrice par mille ne change rien, et le premier pas vaut exactement le gradient orthogonalisé multiplié par le pas.

Un point de reprise qui ne connaît pas l'optimiseur

Muon a un état de plus à enregistrer : sa moyenne glissante. Jusqu'ici, le code des points de reprise connaissait AdamW par son nom et allait chercher ses deux moments. Un second optimiseur aurait demandé un second code.

Chaque optimiseur décrit maintenant son état lui-même : des compteurs, et des séries de tenseurs rangés par paramètre, chacune avec un préfixe.

pub struct OptimizerState<T> {
    pub counters: BTreeMap<String, String>,
    pub slots: Vec<(String, Vec<Option<T>>)>,   // "adamw.m.", "muon.m."…
}

Le module des fichiers écrit et relit cette description sans savoir de quel optimiseur elle vient. Le test de reprise de l'article 10, qui exige qu'un entraînement interrompu finisse sur exactement les mêmes poids, passe avec Muon comme avec AdamW.

Le résultat

Même modèle, même taux d'apprentissage, même graine : seul l'optimiseur change.

À nombre de pas égal, Muon gagne nettement : 0,034 nat par octet, dix-sept fois le bruit. C'est le plus grand écart de toute cette partie de la série, devant RoPE et devant le mélange d'experts. Chaque pas de Muon apprend plus que le pas d'AdamW correspondant.

À durée égale, AdamW gagne. Un pas de Muon coûte 2,3 fois plus cher. Avec le même temps de calcul, AdamW fait 13 600 pas et descend à 0,8997, soit 0,025 de mieux que Muon.

Les deux lignes sont vraies, et laquelle compte dépend de ce qui est rare.

Chez nous, c'est le temps. L'orthogonalisation fait une trentaine de produits matriciels pour chacune des matrices du modèle, à chaque pas. Son coût ne dépend pas de la taille du lot, alors que celui du reste du pas en dépend. Avec nos lots de 8 192 tokens, elle pèse plus lourd que le calcul du gradient lui-même.

À l'échelle de DeepSeek, un lot contient des millions de tokens : le même travail d'orthogonalisation devient négligeable devant le reste, et il ne reste que le gain par pas. Le résultat est donc cohérent avec leur choix, sans le reproduire.

Trois réserves :

• Une seule graine
• Le taux d'apprentissage est celui réglé pour AdamW, et n'a été cherché pour aucun des deux
• Notre orthogonalisation n'a pas été optimisée : elle refait ses transpositions à chaque itération. L'article suivant s'en occupe, et le verdict change.

Le troisième entraînement apporte une leçon de plus, qui dépasse Muon : entraîner 2,3 fois plus longtemps rapporte 0,06 nat par octet, plus que n'importe laquelle des idées de ces quatre articles prise seule.

Ce qui n'a pas été transposé

Trois idées des modèles DeepSeek récents n'ont pas été implémentées. Les raisons diffèrent.

L'attention clairsemée (V3.2) et l'attention compressée (V4). Comparer chaque requête à toutes les clés coûte un temps qui croît avec le carré de la longueur du texte. Pour un contexte d'un million de tokens, c'est intenable. DeepSeek-V3.2 place devant l'attention un indexeur léger qui ne retient que les quelques milliers de positions les plus prometteuses. V4 va plus loin : il compresse les clés par groupes de 4 dans certaines couches, et par groupes de 128 dans d'autres.

Le plan prévoyait une mesure avant toute décision : allonger le contexte et regarder ce que coûte l'attention ordinaire. RoPE le permet sans rien changer au modèle. Voici le temps d'un pas d'entraînement pour le même nombre de tokens, découpés en séquences de plus en plus longues :

Chaque doublement du contexte ajoute environ 35 ms par tranche de 512 positions. À 2 048, la part qui dépend de la longueur pèse la moitié du pas. Et à 4 096, l'entraînement ne tient plus dans les 24 Gio de la carte : un seul tenseur de scores occupe 1 Gio, et la tape en garde plusieurs par couche.

Il y aurait donc quelque chose à gagner à partir de 2 048 positions. L'indexeur n'a pourtant pas été écrit, pour trois raisons.

• L'ordre de grandeur n'y est pas. DeepSeek-V3.2 garde 2 048 positions sur 128 000. Sur un contexte de 2 048, garder 256 positions diviserait par 8 le nombre de scores, dans le meilleur des cas.
• Ce meilleur des cas suppose des kernels aussi bons que le produit matriciel. L'attention dense passe par notre kernel le plus travaillé. Une attention sur des clés choisies ligne par ligne demande des lectures dispersées, une sélection des meilleurs scores sur le GPU, et une rétropropagation où plusieurs requêtes écrivent dans la même clé. L'article 14 vient de montrer ce que coûtent des calculs irréguliers : 1,7 fois le temps pour le même nombre d'opérations. Un facteur 8 théorique n'y survivrait pas en entier.
• L'indexeur compare lui aussi toutes les paires. Il est plus léger, mais il ne lève pas le mur de mémoire rencontré à 4 096.

La mesure est faite, la suite ne l'est pas, et l'attention compressée de V4 non plus : elle répond au même problème.

Les connexions mHC (V4). Depuis l'article 8, chaque sous-couche ajoute sa sortie au flux résiduel : x = x + f(x). DeepSeek-V4 remplace ce flux unique par quatre flux parallèles, que des matrices apprises mélangent à chaque couche. Ces matrices sont contraintes à ne jamais amplifier le signal, par une vingtaine d'itérations de normalisation de leurs lignes et de leurs colonnes. Le but est la stabilité de modèles très profonds.

Notre modèle a 8 couches et son entraînement n'a jamais été instable. Quadrupler le flux résiduel quadruplerait la mémoire des activations pour traiter un problème que nous n'avons pas.

L'entraînement en basse précision. DeepSeek entraîne ses modèles avec des nombres de 8 bits, et même de 4 bits pour une partie des poids. Notre précision mixte de l'article 11 s'arrête à 16 bits, pour les seuls produits matriciels.

Ce qu'il faut retenir

• Un modèle peut avoir deux pertes : celle qu'on minimise et celle qu'on compare. Ne jamais les confondre.
• La prédiction multi-token ajoute un module d'appoint qui prédit un token plus loin, et qu'on retire pour générer.
• Muon regarde une matrice de poids comme une transformation : il garde les directions du gradient et égalise leurs tailles.
• Orthogonaliser une matrice ne demande que des produits matriciels.
• Un optimiseur doit décrire son propre état : le code des fichiers n'a pas à le connaître.
• Une technique se transpose quand on a le problème qu'elle résout.

L'article suivant reprend ces comparaisons là où celui-ci les laisse : À durée égale.