Apprendre à répondre

Deuxième article de la série Smallem, après le pré-entraînement : apprendre à répondre. Précédent : Un format pour converser.

Le modèle de départ ne sait rien faire de nos quatre tâches : 0 % partout, même avec des exemples dans la consigne. Le remède le plus direct est de lui montrer des conversations, et de continuer à l'entraîner dessus. C'est le fine-tuning supervisé : le même entraînement que la première série, sur d'autres données, avec une seule vraie différence.

Une perte masquée

Pendant le pré-entraînement, le modèle apprend à prédire le token suivant partout dans le texte. Ici, il ne l'apprend que là où ce token est écrit par l'assistant : il lit la question, mais n'apprend qu'à écrire la réponse.

pub fn masked_loss(&self, tape, tokens, targets, scored, batch) -> Var {
    let rows: Vec<u32> = (0..tokens.len() as u32).filter(|&i| scored[i as usize]).collect();
    let logits = self.logits(tape, tokens, batch);
    let flat = tape.reshape(logits, &[tokens.len(), self.config.vocab]);
    let kept: Vec<u32> = rows.iter().map(|&row| targets[row as usize]).collect();
    tape.cross_entropy(tape.gather_rows(flat, &rows), &kept)
}

Prenons <|utilisateur|> What is 12 + 30? <|assistant|> 42 <|fin|>. Après What, le modèle devrait prédire is : cette prédiction n'est pas notée, puisque is fait partie de la question. Après <|assistant|>, il doit prédire 42 : celle-là l'est. Le modèle n'apprend donc pas à poser des questions, seulement à y répondre.

La question n'est pas notée, mais elle est lue. Pour prédire 42, le modèle doit regarder 12 + 30 : si la question disait 12 + 31, la bonne réponse changerait. Un test vérifie ces deux propriétés.

Ranger les conversations

Nos conversations font une trentaine de tokens. Il faut les ranger dans des lignes de longueur fixe, deux façons sont possibles :

Une conversation par ligne, complétée par des tokens de remplissage jamais notés. Des lignes de 64 tokens, 32 par lot : la moitié du calcul est perdue sur le remplissage. Mais l'attention étant causale, ces tokens, placés après la conversation, ne changent rien à son calcul.

Des conversations mises bout à bout, autant qu'il en tient dans une ligne de 512 tokens : presque rien de perdu. En échange, chaque conversation voit celles qui la précèdent dans sa ligne. Les isoler demanderait un masque d'attention par document, que nos kernels n'ont pas : la variante n'est pas faite.

Les réglages

Huit mille conversations d'entraînement, deux mille par famille, et quatre cents de validation, tirées comme elles mais jamais vues. 1 500 pas de 32 lignes de 64 tokens, soit six passages sur les données. AdamW sans décroissance des poids, un taux de 10⁻⁴ au pic, dix fois sous celui du pré-entraînement, puis une descente en cosinus.

Toutes les 250 pas : la perte sur les conversations de validation, l'exactitude sur chaque famille, et la perte sur FineWeb-Edu.

Ce que le modèle apprend

Exactitudes à la fin de l'entraînement, perte en nats par octet.

Deux tâches sont apprises. La lecture dès les 250 premiers pas : il suffit de recopier un mot de la question, et le modèle sait recopier. Le tri, plus lentement, jusqu'à 86 %. Et le modèle s'arrête de lui-même : sans <|fin|>, sa réponse se prolongerait de texte, et ne serait jamais exacte.

Deux autres ne le sont pas du tout. L'addition reste sous 2 %, la recopie à l'envers à 0 %. Les deux butent sur le tokenizer. Un nombre de trois chiffres est souvent un seul token, ou deux, sans rapport visible avec ses chiffres : additionner, c'est alors connaître la table d'addition de milliers de tokens. Et retourner un mot demande de savoir de quelles lettres est fait chaque token, ce que le modèle n'a jamais eu besoin d'apprendre. Huit mille exemples n'y suffisent pas à cette taille.

Oublier

La dernière colonne est la mauvaise nouvelle. Le modèle de départ prédisait FineWeb-Edu à 0,883 nat par octet, après le fine-tuning de référence, à 1,478. Il a perdu une bonne partie de ce qu'il savait du texte, pour apprendre à trier des nombres. C'est l'oubli catastrophique, et il est rapide : l'essentiel est fait dès les 250 premiers pas.

Il y a deux remèdes que l'on peut apporter. Un taux plus faible abîme moins : 1,075. Et surtout, continuer à montrer du texte : un quart des lignes de chaque lot sont prises dans FineWeb-Edu, et notées partout. L'oubli tombe à +0,051. Avec la moitié des lignes, et deux fois plus de pas pour garder autant de conversations, à +0,041. Le prix est sur le tri, appris plus lentement : 71 et 74,5 % au lieu de 86.

On souhaitait un oubli de 0,02 au plus. Avec le taux de référence, même la moitié du lot en texte n'y suffit pas : le modèle perd 0,04 dès les premiers pas, quand il apprend le format, et ne le rattrape pas. Il faut les deux remèdes ensemble, la moitié du lot en texte et un taux de 3 × 10⁻⁵ : l'oubli tombe à +0,011, sous le seuil. Le tri s'arrête alors à 57 %.

Entre savoir trier et ne rien oublier, il faut choisir. Pour la suite de la série, nous prenons ce dernier modèle : il tient le seuil d'oubli, et laisse de la marge aux méthodes qui suivent.

Surapprendre

Avec des données finies, le modèle finit par les apprendre par coeur. Deux essais le montrent.

Les conversations mises bout à bout voient deux fois plus de conversations par pas, onze passages en tout. Leur perte de validation est au plus bas au pas 500, à 0,477, puis remonte jusqu'à 0,617 : le modèle apprend les conversations d'entraînement elles-mêmes, et plus ce qui les rend justes.

Avec cent conversations par famille au lieu de deux mille, c'est flagrant :

La perte d'entraînement tombe à zéro : les 400 conversations sont sues par coeur. Celle de validation ne fait que monter. C'est pourquoi le programme garde le modèle dont la perte de validation est la plus basse, et non le dernier. Et c'est la première fois du projet : pendant toute la première série, avec 850 millions de tokens, aucun modèle n'a vu deux fois le même texte.

Ce qu'il faut retenir

• Le fine-tuning supervisé, c'est l'entraînement ordinaire, noté sur les seuls tokens de l'assistant.
• Un modèle de 29 millions de paramètres apprend en quelques centaines de pas ce qui se recopie ou se compare (lecture, tri), pas ce qui demande de voir dans les tokens (addition, recopie à l'envers).
• L'oubli est rapide et fort. Mêler du texte ordinaire aux conversations le divise par dix.
• Avec des données finies, la perte de validation finit par remonter : il faut garder le bon modèle, pas le dernier.

Article suivant : Ajuster sans tout réentraîner.