Trois pièces modernes, et comment mesurer une idée

Douzième article de la série Smallem. Précédent : Rendre le GPU rapide.

Le modèle des articles précédents est un GPT-2 : une architecture de 2019. Les modèles d'aujourd'hui gardent son squelette, mais presque toutes ses pièces ont été remplacées. Cet article en change trois (la normalisation, le feed-forward, les positions), et les suivants s'attaquent aux idées propres à la famille DeepSeek.

Avant de changer quoi que ce soit, il faut régler une question qui n'a rien à voir avec les réseaux de neurones : comment sait-on qu'un changement est une amélioration ?

Le problème : le bruit

Entraînons deux fois le même modèle, sur les mêmes données, avec le même code. La seule différence est la graine du hasard, qui décide des poids de départ et de l'ordre des lots. Les deux pertes finales ne sont pas égales.

L'écart est de 0,0007. C'est le bruit de la mesure : la part du résultat qui ne dit rien du modèle. Toute différence de cet ordre entre deux architectures ne prouve rien. On retiendra un seuil prudent de 0,002.

Sans ce chiffre, on ne peut rien conclure. Une idée qui « gagne » 0,001 n'a rien gagné, et on ne le sait que si l'on a mesuré le bruit d'abord.

Le protocole

Les règles qui suivent s'appliquent à cet article et aux trois suivants.

• Une référence fixe. Le modèle de 30 millions de paramètres de l'article 10 : 8 couches, dimension 512, contexte de 512 tokens, vocabulaire de 8 192 sous-mots.
• Un budget fixe. 6 000 pas de 16 séquences, soit 49 millions de tokens et 16 minutes sur la carte. C'est court : le modèle n'a pas fini d'apprendre. Mais c'est le prix pour pouvoir essayer une vingtaine de variantes.
• Une mesure fixe. À la fin, la perte est calculée sur 200 lots de validation, toujours les mêmes, quelle que soit la graine.
• Un changement à la fois, à nombre de paramètres comparable.
• Les contrôles habituels pour chaque pièce nouvelle : gradient vérifié par différences finies, CPU et GPU comparés, test du futur, sur-apprentissage d'un lot.

Un budget court a une conséquence qu'il faut garder en tête : on mesure ce qui apprend vite, pas forcément ce qui finit mieux. Les conclusions valent pour ce budget et cette taille.

Une architecture configurable

Pour comparer des variantes, il faut pouvoir les construire sans dupliquer le modèle. La configuration porte donc une description des différents éléments :

pub struct Architecture {
    pub norm: NormKind,                 // Layer ou Rms
    pub feed_forward: FeedForwardKind,  // Gelu ou SwiGlu
    pub positions: PositionKind,        // Learned ou Rotary
    pub attention: AttentionKind,       // voir l'article suivant
}

impl Architecture {
    pub const GPT2: Self = /* LayerNorm, GELU, positions apprises */;
    pub const MODERN: Self = /* RMSNorm, SwiGLU, RoPE */;
}

Un bloc ne contient plus une LayerNorm mais une Norm, une énumération à deux variantes :

pub enum Norm<B: Backend> {
    Layer(LayerNorm<B>),
    Rms(RmsNorm<B>),
}

L'enveloppe est transparente pour le fichier de poids : le gain d'une normalisation s'appelle blocks.0.ln1.gain quelle que soit la variante. Le GPT-2 de l'article 9 se charge donc toujours.

L'architecture voyage avec le modèle, dans l'en-tête du fichier de poids, sous la forme d'un texte JSON :

{"vocab":8192,"context":512,"layers":8,"heads":8,"embed":512,
 "architecture":{"norm":"rms","feed_forward":"swiglu","positions":"rope"}}

Un fichier dit ainsi lui-même comment reconstruire le modèle qu'il contient. Et sur la ligne de commande, --arch modern choisit les trois pièces, --norm rms une seule.

RMSNorm : une normalisation allégée

La LayerNorm de l'article 8 fait deux choses à chaque vecteur : elle lui retire sa moyenne, puis le divise par son écart-type. Un gain et un biais appris lui redonnent ensuite une échelle.

RMSNorm (Zhang et Sennrich, 2019) ne garde que la division. Chaque vecteur est divisé par la racine de la moyenne de ses carrés (root mean square), puis multiplié par un gain. Pas de moyenne retirée, pas de biais.

rms(x) = x / √(moyenne(x²) + ε)
sortie = rms(x) × gain

L'idée de ses auteurs est que c'est la mise à l'échelle qui stabilise l'entraînement, pas le recentrage. On économise une réduction et une soustraction.

Comme normalize à l'article 11, c'est une opération fusionnée, rms_normalize, avec son kernel : un warp par ligne pour la somme des carrés, puis un thread par élément. Et comme toute opération, elle a une définition de référence écrite avec les primitives, contre laquelle le kernel est testé.

SwiGLU : une porte dans le feed-forward

Le feed-forward de GPT-2 élargit le vecteur, applique GELU, puis le ramène à sa taille :

sortie = gelu(x · W₁) · W₂

SwiGLU (Shazeer, 2020) calcule deux projections de l'entrée et les multiplie l'une par l'autre :

sortie = (silu(x · W_porte) × (x · W_valeur)) · W_proj

La première projection, passée par la fonction SiLU, sert de porte à la seconde : elle décide, composante par composante, ce qui passe. silu(x) = x × sigmoïde(x) ressemble beaucoup à GELU : presque nulle pour les valeurs négatives, presque l'identité pour les grandes valeurs positives.

pub fn forward(&self, tape: &Tape<B>, x: Var) -> Var {
    let gate = tape.silu(self.gate.forward(tape, x));
    self.proj
        .forward(tape, tape.mul(gate, self.up.forward(tape, x)))
}

Il y a maintenant trois matrices au lieu de deux. Pour que la comparaison reste honnête, la largeur cachée passe de 4C à 8C/3 : trois matrices de largeur 8C/3 ont autant de paramètres que deux de largeur 4C. Pour C = 512, cela donne 1 368 après arrondi.

RoPE : la position par une rotation

C'est le changement le plus intéressant, et, on le verra, le seul qui compte.

Ce que faisait GPT-2

L'attention ne connaît pas l'ordre des tokens : pour elle, une séquence est un ensemble. GPT-2 corrige cela en ajoutant à chaque token un vecteur appris qui dépend de sa place : un vecteur pour la position 0, un autre pour la position 1, et ainsi de suite jusqu'à 511.

Cette solution a deux défauts. Le modèle apprend des positions absolues, alors que ce qui compte le plus souvent est l'écart entre deux tokens. Prenons une règle simple qu'une tête d'attention peut apprendre : regarder le token qui précède immédiatement, par exemple pour savoir qu'après « le » vient souvent un nom. Cette situation se présente partout dans le texte : « le chat » peut occuper les positions 9 et 10 d'une séquence, « le jardin » les positions 299 et 300 d'une autre. Avec des positions absolues, ce n'est pas la même relation : dans un cas, la position 10 regarde la position 9 ; dans l'autre, la position 300 regarde la position 299, et rien ne relie ces vecteurs entre eux. Ce que le modèle a appris à la position 10, il doit le réapprendre à la position 300. Et il ne peut rien faire au-delà de la dernière position de sa table.

L'idée

Ce qui compte dans une phrase, c'est la distance entre deux mots, pas leur rang. RoPE (rotary position embedding, Su et al., 2021) fait entrer la position là où elle sert : au moment où une requête est comparée à une clé.

Prenons deux composantes d'une clé et voyons-les comme un point du plan. RoPE fait tourner ce point autour de l'origine, d'un angle proportionnel à la position du token : θ pour la position 1, 2θ pour la position 2, mθ pour la position m. Il fait de même avec la requête.

Le produit scalaire de deux vecteurs du plan ne dépend que de leurs longueurs et de l'angle entre eux. Si la requête a tourné de mθ et la clé de nθ, l'angle entre elles a changé de (m - n)θ. Le score ne dépend donc que de m - n : la distance.

requête à la position m, tournée de mθ
clé à la position n, tournée de nθ
score = fonction de (requête, clé, m - n)

Aucun paramètre n'est appris. Les positions absolues ont disparu : décaler toute la séquence de cent positions ne change aucun score.

Plusieurs vitesses

Un vecteur de 64 composantes forme 32 paires, et chaque paire tourne à sa propre vitesse :

angle de la paire i à la position m = m × 10 000^(-2i/D)

La première paire tourne d'un radian par position : elle distingue deux tokens voisins, mais elle a fait un tour complet après six positions. La dernière tourne plusieurs milliers de fois moins vite : elle distingue le début et la fin d'un long texte. C'est le principe des aiguilles d'une horloge.

Le code

Les cosinus et les sinus sont calculés sur l'hôte, une fois par passe, et envoyés à la carte. L'opération elle-même est élémentaire :

(x₀, x₁)  devient  (x₀·cos - x₁·sin,  x₀·sin + x₁·cos)

Dans l'attention, deux lignes suffisent :

if let Some(rope) = rope {
    q = rope.rotate(tape, q);
    k = rope.rotate(tape, k);
}

Les valeurs ne sont pas tournées : la position sert à choisir qui écouter, pas à modifier ce qu'on entend.

La rétropropagation est tout aussi simple. Une rotation ne change pas les longueurs, donc son gradient est la rotation inverse du gradient de la sortie. Les angles sont des données et ne reçoivent pas de gradient.

Un test qui dit ce que RoPE promet

#[test]
fn positions_matter_with_rope() {
    // Sans information de position, permuter deux tokens du passé ne
    // changerait pas le logit de la dernière position.
    let last = |tokens: &[u32]| logits_of(&model, tokens, 1)[3 * TINY.vocab..].to_vec();
    assert_ne!(last(&[3, 7, 1, 9]), last(&[7, 3, 1, 9]));
}

Un autre test vérifie la propriété inverse sur l'attention : décaler toutes les positions du même nombre ne change rien.

Une conséquence pour la génération

L'article 11 a introduit le cache de génération, qui garde les clés et les valeurs des positions déjà traitées. Avec des positions apprises, ce cache a une limite : quand le texte dépasse le contexte, il faut oublier le plus vieux token, et tous les autres changent de position. Leurs vecteurs changent donc, et le cache ne vaut plus rien.

Avec RoPE, oublier le plus vieux token ne change rien aux distances entre ceux qui restent. Le cache continue de servir : il suffit de retirer sa première ligne. Le modèle RoPE génère 160 tokens par seconde sur un texte de 700 tokens, contre 144 pour le modèle à positions apprises, qui recalcule tout à chaque token une fois le contexte dépassé.

Les résultats

Chaque pièce est mesurée seule, puis les trois ensemble.

L'écart est pris par rapport à la moyenne des deux références.

RoPE est le seul changement qui compte. Il gagne 0,043 nat par octet, vingt fois le seuil de bruit, pour le même débit et 262 144 paramètres de moins : la table des positions a disparu.

RMSNorm et SwiGLU sont dans le bruit. Leurs écarts, +0,0018 et -0,0014, sont indiscernables de zéro. RMSNorm est 3 % plus rapide ; SwiGLU est 4 % plus lent, sa troisième matrice ajoutant un produit.

Les trois ensemble ne font pas mieux que RoPE seul.

Ce que ce résultat veut dire

Il serait faux d'en conclure que RMSNorm et SwiGLU ne servent à rien. Ils ont été adoptés par des équipes qui entraînent des modèles mille fois plus gros, bien plus longtemps. Ce que la mesure dit est plus étroit : à 30 millions de paramètres et 49 millions de tokens, avec une seule graine par variante, leur effet est plus petit que le bruit.

C'est aussi ce qui rend la mesure du bruit indispensable. Sans elle, le tableau se lirait « SwiGLU améliore, RMSNorm dégrade », et l'on bâtirait la suite sur deux conclusions fausses.

L'architecture modern devient la référence des articles suivants. Non parce que ses trois pièces ont toutes fait leurs preuves ici, mais parce que c'est le socle des modèles DeepSeek, et que la prochaine idée suppose RoPE.

Ce qu'il faut retenir

• Deux entraînements identiques ne donnent pas la même perte. On mesure cet écart avant de comparer quoi que ce soit.
• Un changement à la fois, à budget et à paramètres égaux.
• RMSNorm : diviser par la racine de la moyenne des carrés, sans recentrer.
• SwiGLU : deux projections dont l'une sert de porte à l'autre.
• RoPE : faire tourner requêtes et clés selon leur position, pour que le score ne dépende que de la distance. Aucun paramètre, et un cache de génération qui peut glisser.
• Un résultat nul est un résultat, à condition de savoir ce qu'il ne dit pas.

La suite s'attaque au cache lui-même : L'attention latente.