À durée égale

Seizième article de la série Smallem. Précédent : Prédire plus loin, optimiser autrement.

Les quatre articles précédents ont comparé des idées à nombre de pas égal : chaque variante voyait les mêmes 49 millions de tokens. C'est la comparaison la plus simple, et elle a un défaut que l'article 15 a montré avec Muon : une idée peut gagner par pas et perdre par minute.

Sur une seule carte, c'est la minute qui compte. Cet article reprend donc tout à durée égale, puis essaie quelques idées publiées après DeepSeek-V4, en les jugeant sur ce seul critère.

Une courbe de référence

Pour comparer à durée égale, il faut savoir ce que vaut le modèle de référence après n'importe quelle durée. Quatre entraînements d'AdamW, de plus en plus longs, donnent cette courbe :

Chaque ligne est un entraînement complet, avec son taux d'apprentissage qui redescend jusqu'au bout : on ne lit pas ces valeurs sur la courbe d'un seul entraînement long.

Une variante qui met 25 minutes à faire ses 6 000 pas ne se compare plus à 0,9592, mais à 0,9247.

La dette des articles précédents

Deux résultats attendaient cette comparaison.

Les deux perdent. Le mélange d'experts gagnait 0,019 à pas égaux ; il en perd 0,015 à durée égale. La prédiction multi-token gagnait 0,004 ; elle en perd 0,013.

Ce n'est pas un jugement sur ces idées. Le mélange d'experts promet un coût de calcul inchangé, et nos kernels ne tiennent pas cette promesse : un pas coûte 1,6 fois plus cher. Avec des kernels qui la tiendraient, la comparaison serait à refaire. Mais c'est le résultat sur notre carte, et il vaut mieux le connaître.

Avec celui de Muon à l'article 15, le bilan de la partie DeepSeek devient net : à durée égale, aucune des idées essayées ne bat le modèle de l'article 12 entraîné plus longtemps.

Muon : où part le temps

Muon était le cas le plus frustrant. Il apprend nettement plus par pas, et son pas coûte 2,3 fois plus cher. Avant d'essayer des variantes, il fallait savoir pourquoi.

Le profileur de l'article 11 répond en une mesure :

AdamW   162 ms par pas    219 produits matriciels     96 ms
Muon    386 ms par pas  1 899 produits matriciels    267 ms

L'orthogonalisation ajoute 1 680 produits matriciels à chaque pas : trente pour chacune des 56 matrices de poids du modèle. Ils sont petits, des carrés de 512 de côté, et c'est le problème. Chacun prend un dixième de milliseconde, soit un débit de moins de 3 TFLOP/s pour un kernel qui en atteint 19 sur de grandes matrices. Lancer un kernel a un coût fixe, et un petit produit ne remplit pas la carte.

C'est le même constat qu'à l'article 14 avec les experts : notre carte aime les gros calculs réguliers.

Empiler

Le remède vient de là. Les 32 matrices d'attention du modèle ont toutes la même forme. On peut les ranger dans un seul tenseur à trois dimensions et les orthogonaliser ensemble : le produit matriciel de l'article 5 traite déjà les lots.

// Une pile par forme : trente produits matriciels pour trente-deux
// matrices d'attention, au lieu de neuf cent soixante.
for group in by_shape.values() {
    let matrices: Vec<&B::Tensor> = group.iter().map(|&at| &aheads[at].1).collect();
    let stacked = stack(backend, &matrices);
    let result = orthogonalize(backend, &stacked, self.variant.newton_schulz());
    for (position, &at) in group.iter().enumerate() {
        directions[at] = Some(unstack(backend, &result, position));
    }
}

Le calcul est le même, nombre pour nombre. Seule la façon de le présenter à la carte change. Chaque matrice garde sa propre norme : la normalisation qui précède l'itération se fait matrice par matrice dans la pile.

Un test vérifie que trois matrices, dont deux de même forme, reçoivent exactement les mises à jour qu'elles auraient reçues seules.

Moins d'itérations

DeepSeek-V4 fait dix itérations de Newton-Schulz. L'article qui propose SAMuon (Spectral Allocation), dont il est question plus bas, en fait quatre, avec des coefficients différents à chaque itération, réglés pour que le résultat soit proche de 1 dès qu'une valeur singulière dépasse 2 % de la norme. Douze produits au lieu de trente.

Le surcoût de Muon passe de 138 % à 19 %. Aucun kernel n'a été écrit.

Trois variantes de Muon

Trois articles récents proposent d'améliorer Muon. Chacun tient en quelques lignes une fois Muon écrit.

Le Muon périodique par lignes (Scaling Muon for Diffusion Transformers) vise le coût. L'orthogonalisation n'a lieu qu'un pas sur trois. Les autres pas se contentent de ramener chaque ligne de la moyenne des gradients à la norme 1, ce qui ne coûte presque rien. Une matrice orthogonale a aussi des lignes de norme 1, mais en plus perpendiculaires entre elles : on ne garde que la moitié bon marché de la contrainte. Les pas par lignes sont plus courts, d'un facteur 0,15 retenu par l'article.

NorMuon (NorMuon) vise la qualité. Dix itérations, ou quatre, ne donnent pas une matrice parfaitement orthogonale, et certaines colonnes de la mise à jour restent plus grosses que d'autres. Or une colonne est un neurone. NorMuon divise la mise à jour de chaque neurone par la racine de la moyenne glissante de son carré : c'est l'idée d'Adam, appliquée par neurone plutôt que par nombre.

SAMuon-lite vise la qualité aussi, par une observation. Ses auteurs mesurent, direction par direction, le plus grand pas que la perte tolère. La première direction, celle de la plus grande valeur singulière, est instable et ne tolère qu'un petit pas. Toutes les autres en toléreraient un bien plus grand. Muon, qui donne le même pas à toutes, est donc bridé par la première.

SAMuon-lite multiplie toute la mise à jour par un facteur, puis ramène la première direction à sa taille d'origine :

O' = γ × O − (γ − 1) × u₁·v₁ᵀ

u₁ et v₁ sont les premiers vecteurs singuliers. On les trouve sans décomposition, par la puissance itérée : multiplier un vecteur par Mᵀ·M encore et encore l'aligne sur la direction que la matrice amplifie le plus. Dix itérations suffisent, et elles s'empilent comme le reste.

Les trois variantes se choisissent sur la ligne de commande, et se combinent :

--optimizer muon:short,period=3
--optimizer muon:short,norm
--optimizer muon:short,boost=5

Les résultats

Chaque variante fait les mêmes 6 000 pas que la référence. Sa durée change, et avec elle le point de la courbe d'AdamW auquel elle se compare. Ce point est interpolé entre les quatre entraînements du premier tableau.

La troisième ligne n'a pas été réentraînée jusqu'au bout : empiler ne change pas le calcul, donc sa perte est celle de l'article 15, et seul son débit a été remesuré, sur 500 pas.

Quatre lectures.

Muon devient rentable. Avec les matrices empilées et quatre itérations, il atteint 0,9256 en 19 minutes. AdamW a besoin de 25 minutes pour faire aussi bien. Le résultat de l'article 15 est renversé, pas une idée nouvelle sur l'optimisation, seulement une meilleure façon de présenter le même calcul à la carte.

Quatre itérations valent dix. 0,9256 contre 0,9251 : l'écart est quatre fois plus petit que le bruit. L'orthogonalisation n'a pas besoin d'être précise.

Le Muon périodique ne gagne rien ici. Il est le plus rapide, et il perd exactement ce qu'il économise. Deux pas sur trois sans orthogonalisation coûtent 0,027 de perte. Une fois l'orthogonalisation ramenée à 19 % du pas, il ne reste plus grand-chose à économiser en la sautant.

NorMuon et SAMuon-lite ne se distinguent pas de Muon. Leurs pertes sont à moins de 0,001 de celle de muon:short, et leurs pas sont un peu plus chers. À durée égale, ils sont donc légèrement derrière.

Ce dernier résultat demande des réserves, parce qu'il ne reproduit pas les articles :

• Une seule graine, et un taux d'apprentissage qui n'a été réglé pour aucune variante.
• SAMuon amplifie le bruit du gradient, et ses auteurs le disent : le gain demande de grands lots. Les leurs comptent 1 024 séquences, les nôtres 16.
• Notre NorMuon garde l'échelle de mise à jour de Muon, 0,18, là où l'article prend 0,2.
• Le facteur 0,15 du Muon périodique vient d'un article sur des modèles de diffusion.

La contre-épreuve

L'article 15 avait une ligne cruelle : AdamW entraîné 36 minutes atteignait 0,8997, et battait Muon entraîné aussi longtemps. Il fallait refaire cette mesure-là, sur un entraînement entier et pas sur une interpolation.

Muon l'emporte de 0,017, huit fois le bruit. C'est le meilleur modèle de la série, et c'est la première idée de cette partie qui bat la référence à la minute.

L'explication de l'article 15, « Muon suppose des lots énormes », était donc trop rapide. Le coût de l'orthogonalisation ne dépend pas du lot, c'est vrai, et Practical Efficiency of Muon for Pretraining mesure son avantage avec des lots de 128 000 tokens et plus. Mais ce coût était surtout celui de notre façon de l'écrire.

Traverser les couches deux fois

L'idée vient du modèle Nanbeige 4.2 (rapport technique) : faire passer les tokens deux fois dans la même pile de couches. Le modèle devient deux fois plus profond sans un seul poids de plus.

Le code tient en une boucle autour de celle qui existait :

let layers = self.blocks.len();
for pass in 0..self.config.architecture.passes {
    for (index, block) in self.blocks.iter().enumerate() {
        let past = cache
            .as_deref_mut()
            .map(|cache| &mut cache.layers[pass * layers + index]);
        x = block.forward(tape, x, rope, past, &mut lent);
    }
}

L'autograd n'a rien à apprendre. Un poids utilisé deux fois reçoit deux gradients, qui s'additionnent : c'est déjà ce qui arrive à la table d'embedding, utilisée à l'entrée et à la sortie depuis l'article 8.

Le seul piège est le cache de génération. Les poids sont partagés, mais pas les états : à la seconde traversée, une couche reçoit d'autres vecteurs qu'à la première, donc calcule d'autres clés. Il faut un étage de cache par couche et par traversée. Le rapport de Nanbeige dit avoir essayé de partager aussi les clés entre traversées, et y avoir perdu.

Un test vérifie le tout d'un coup : un modèle de 2 couches traversé deux fois donne exactement les sorties d'un modèle de 4 couches dont on a recopié les poids.

Trois modèles se comparent. Le deuxième a les poids du premier et le coût du troisième.

À pas égaux, la seconde traversée rapporte 0,008, quatre fois le bruit. C'est un tiers de ce que rapportent 16 vraies couches (0,023), pour aucun paramètre de plus.

À durée égale, les deux perdent, et nettement. Un pas coûte 1,9 fois plus cher dans les deux cas, et le modèle de 8 couches fait meilleur usage de ce temps en voyant presque deux fois plus de texte. La mémoire suit le même chemin : 16,7 Gio pour deux traversées, contre 9,4.

C'est la leçon de l'article 15 qui revient : dans notre régime, où le modèle n'a vu que quelques pour cent du corpus, rien ne vaut des tokens en plus. Une seconde traversée est une idée pour qui est limité par la taille du modèle, sur un téléphone par exemple. Nous sommes limités par le temps.

Deux réserves :

• Une seule graine.
• L'initialisation n'a pas été adaptée : les projections sont réglées pour 8 couches qui s'additionnent dans le flux résiduel, alors que deux traversées en font 16.

Partager les clés entre couches

L'article 13 réduisait le cache de génération en compressant ce que chaque couche y range. Gemma 4 fait plus simple : toutes les couches ne remplissent pas le cache. Les premières calculent leurs clés et leurs valeurs, comme d'habitude. Les suivantes réutilisent celles de la dernière couche qui en a calculé, et ne gardent à elles que leurs requêtes.

Une couche qui emprunte est une couche d'attention amputée de deux projections sur quatre :

pub struct BorrowedAttention<B: Backend> {
    heads: usize,
    pub query: Linear<B>,
    pub proj: Linear<B>,
}

Le calcul de l'attention lui-même ne change pas. Il a suffi de le découper en deux fonctions, l'une qui produit les clés et les valeurs, l'autre qui s'en sert, et de faire passer les clés d'un bloc au suivant.

L'option --own-keys N dit combien de couches calculent leurs clés. Avec 8 couches :

Partager les clés de la moitié des couches divise le cache par deux et coûte 0,004, deux fois le bruit. C'est six fois moins que l'attention latente. Et là où celle-ci ralentissait tout, le partage accélère : l'entraînement de 6 % (55 000 tokens par seconde contre 52 000) et la génération de 13 %, puisque quatre couches ont deux projections de moins à calculer.

À durée égale, la perte de 0,004 disparaît : le modèle finit ses 6 000 pas en un peu moins de 15 minutes, et la référence arrêtée au même moment en serait à peu près au même point. Ce « à peu près » est une extrapolation de la courbe en dessous de son premier point, pas une mesure.

Aller plus loin ne marche pas. Avec deux couches seulement, la perte remonte au niveau de GPT-2, et le cache n'est toujours que quatre fois plus petit, contre 6,4 fois pour l'attention latente. Les couches profondes peuvent se passer de leurs propres clés mais les premières, non.

Pour un cache deux fois plus petit, c'est donc la meilleure des deux méthodes essayées. Pour un cache six fois plus petit, l'attention latente reste la seule. Et chez nous, comme à l'article 13, 16 Mio de cache ne gênaient personne.

Ce qu'il faut retenir

• Une idée se compare à la référence entraînée aussi longtemps, pas pendant autant de pas. Il faut pour cela une courbe de référence, faite d'entraînements complets.
• À cette aune, le mélange d'experts, la prédiction multi-token et la double traversée perdent tous : dans notre régime, rien ne vaut des tokens en plus.
• Avant de juger une idée trop chère, profiler. Muon coûtait 2,3 fois le pas d'AdamW à cause de 1 680 petits produits matriciels, pas à cause de l'idée.
• Empiler des matrices de même forme transforme beaucoup de petits calculs en quelques gros. C'est ce que la carte sait faire.
• Une orthogonalisation approchée suffit : quatre itérations valent dix.
• Un résultat publié à grande échelle ne se retrouve pas toujours à la nôtre. NorMuon et SAMuon-lite n'ont rien donné ici, avec des lots soixante fois plus petits que les leurs.
• Les couches profondes peuvent emprunter les clés des premières. C'est la façon la moins chère de réduire le cache de moitié.

Reste un mur, que l'article suivant fait tomber : L'attention fusionnée.