Ajuster sans tout réentraîner

Troisième article de la série Smallem, après le pré-entraînement : apprendre à répondre. Précédent : Apprendre à répondre.

Le fine-tuning de l'article précédent modifie tous les paramètres du modèle, et cela a deux prix. Le premier, nous l'avons mesuré : en apprenant les tâches, le modèle oublie une partie de ce qu'il savait. Même en continuant à lui montrer du texte, sa perte sur FineWeb-Edu passe de 0,8827 à 0,9338. Le second compte pour les grands modèles : il faut garder, pour chaque paramètre, les deux moments d'Adam, et enregistrer à la fin une copie complète du modèle, une par tâche.

LoRA (Hu et al., 2021) répond au second : chaque matrice est gelée, et seule une petite correction apprend, quelques pour cent des paramètres. À notre échelle, la mémoire ne manque pas. Mais un modèle qui ne peut changer qu'un peu pourrait aussi moins oublier : c'est ce que cet article mesure.

Une correction de rang faible

Une couche linéaire calcule x · W, avec W de [entrée, sortie]. LoRA gèle W et lui ajoute un adaptateur : le produit de deux matrices minces, A et B, seules à apprendre.

x · (W + A · B)       A : [entrée, r]     B : [r, sortie]

Avec r petit, la correction A · B est de rang r au plus. Pour une matrice 512 × 512 et r = 8, 8 192 paramètres au lieu de 262 144. L'idée de l'article est que la correction utile à une tâche nouvelle a peu de dimensions : pas besoin de pouvoir changer tout W.

B part de zéro : au départ, la correction est nulle, et le modèle est exactement celui d'avant. A part du hasard, sinon le gradient de B, qui passe par A, serait nul lui aussi, et rien ne bougerait.

pub fn forward(&self, tape: &Tape<B>, x: Var) -> Var {
    tape.matmul(tape.matmul(x, self.down.var(tape)), self.up.var(tape))
}

Le calcul est (x · A) · B, pas x · (A · B) : on ne fabrique jamais la matrice pleine.

Geler W est simple : un paramètre gelé entre sur la tape comme une constante. La rétropropagation ne calcule alors aucun gradient pour lui, et l'optimiseur, qui saute déjà les paramètres sans gradient, ne le touche pas. Le gradient continue pourtant de traverser W, vers les couches précédentes : leurs adaptateurs en ont besoin pour apprendre.

Le problème des tokens de rôle

Une difficulté propre à notre post-entraînement est apparue en écrivant. Les trois tokens de rôle sont nés identiques, à la moyenne des lignes de la table d'embedding. Si la table est gelée, ils le restent : en entrée, le modèle ne peut pas distinguer <|utilisateur|> de <|assistant|>, et en sortie, <|fin|> a toujours le même logit que les deux autres.

La table reçoit donc elle aussi son adaptateur, [V, r] · [r, C]. La table effective est E + A · B, recalculée à chaque passe, elle sert à l'entrée comme à la sortie. Pour r = 4, cela fait 35 000 paramètres de plus, et chaque token, de rôle ou non, peut recevoir sa propre correction.

Fusionner

Une fois l'entraînement fini, la correction s'ajoute à la matrice : W ← W + A · B. Le modèle redevient un modèle ordinaire, sans adaptateur, ni calcul de plus. Un test vérifie que le modèle fusionné donne les mêmes logits, à 10⁻¹² près, que le modèle avec ses adaptateurs, et qu'un modèle muni d'adaptateurs neufs donne exactement les logits de départ.

Les mesures

Mêmes données que le fine-tuning complet : 1 500 pas, un quart des lignes en FineWeb-Edu. LoRA apprend avec un taux plus fort, 10⁻³ au lieu de 10⁻⁴, comme l'article le recommande : les adaptateurs partent de zéro et doivent faire tout le chemin.

L'addition et la recopie à l'envers restent à 0 ou presque, comme avec le fine-tuning complet.

Avec 1,2 % des paramètres, LoRA apprend la lecture aussi bien, et le tri 10 points moins bien. Au rang 32, l'écart sur le tri disparaît presque. On pouvait espérer, sous 2 % des paramètres, une exactitude à quelques points du fine-tuning complet : c'est tenu pour la lecture, pas pour le tri.

La réponse est dans la dernière colonne. Le fine-tuning complet faisait perdre 0,051 nat par octet sur FineWeb-Edu, LoRA au rang 4, 0,004. Dix fois moins d'oubli. Une correction de rang 4 ne peut déplacer chaque couche que dans peu de directions : elle modifie moins le modèle initial que le fine-tuning complet, et cela limite l'oubli. Elle ne l'empêche pas pour autant. Ces quelques directions agissent sur tout ce que la couche calcule, et rien ne garantit qu'elles épargnent ce que le modèle savait. L'écart mesuré ici vaut pour ce modèle, ces tâches et cet entraînement, il peut se réduire avec un rang plus élevé, un entraînement plus long ou des données plus éloignées du pré-entraînement. Plus le rang monte, plus le modèle apprend, et plus il oublie : le rang règle cet échange.

La mémoire, enfin. Les moments d'Adam du fine-tuning complet occupent 235 Mo, ceux de LoRA au rang 4, 2,8 Mo. À 29 millions de paramètres, c'est négligeable à côté des activations, à 70 milliards, c'est ce qui décide si l'entraînement tient sur une carte. Les durées, elles, n'ont pas pu être comparées proprement : la carte était partagée avec d'autres entraînements pendant ces mesures.

Ce qu'il faut retenir

• LoRA gèle les matrices et apprend une correction A · B de rang faible, nulle au départ, fusionnable à la fin.
• Des tokens nouveaux, nés identiques, ont besoin que leur table d'embedding apprenne : elle reçoit, elle aussi, son adaptateur.
• Avec 1,2 % des paramètres, LoRA apprend presque aussi bien et, dans cette expérience, oublie dix fois moins : il modifie moins le modèle initial, ce qui limite l'oubli sans l'empêcher.

Article suivant : Apprendre d'une préférence.