Dupliquer des couches pendant l'entraînement fait passer GSM8K de 12% à 19% sur un modèle de 3B

Les chercheurs Ferdinand Kapl, Kaitlin Maile, Stefan Bauer et Johannes von Oswald ont comparé trois façons de construire un modèle de langage allant jusqu'à 3 milliards de paramètres. La première est un réseau dense classique. La deuxième est un réseau qui gagne progressivement en profondeur pendant l'entraînement, par duplication des couches intermédiaires (méthode LIDAS).

La troisième est un réseau « bouclé », où les mêmes couches sont traversées plusieurs fois en une seule passe. L'étude porte sur l'in-context learning (ICL) : la capacité d'un modèle à comprendre une tâche à partir de quelques exemples dans le prompt, sans réentraînement. Résultat : l'extension de la profondeur et le bouclage apprennent l'ICL nettement mieux que le modèle dense.

À l'échelle 3B, LIDAS atteint environ 19% de correspondance exacte sur GSM8K avec 8 exemples dans le prompt, contre environ 12% pour le modèle dense.

Deux façons de réutiliser la profondeur

Dans un modèle ordinaire, le nombre de couches est fixé d'avance et chaque couche est entraînée à partir de zéro. LIDAS procède autrement : l'entraînement démarre avec un réseau peu profond, puis les couches intermédiaires sont dupliquées, et le réseau grandit en cours de route. Au départ, les copies sont presque identiques, puis elles divergent. La phase initiale se déroule sur un petit modèle, elle coûte donc moins cher. Selon les auteurs, à 3B, LIDAS dépense 26% de FLOPs (opérations en virgule flottante) en moins à l'entraînement, tout en dépassant le modèle dense.

Le bouclage est plus simple. On prend un petit ensemble de couches, et les données le traversent plusieurs fois au cours d'une seule passe avant. Le modèle compte moins de paramètres uniques, mais sa profondeur effective est plus grande.

Les expériences ont porté sur des modèles de 360M à 3B de paramètres. Le code repose sur la bibliothèque nanotron de Hugging Face. Les auteurs ne publient ni leur propre code ni les poids.

Où les modèles prennent l'avantage

GSM8K est un ensemble de problèmes de mathématiques de niveau scolaire. On place huit exemples résolus dans le prompt et on vérifie si la réponse du modèle correspond à la référence. Résultats des modèles 3B :

• LIDAS : environ 19% de correspondances exactes ;
• modèle dense : environ 12%.

Le second test est RegBench : de nouveaux langages formels que le modèle n'a jamais vus. À partir de quelques exemples, il doit saisir les règles et prédire le symbole suivant. La qualité se mesure par la distance de variation totale, c'est-à-dire l'écart entre les probabilités prédites et les vraies. Plus elle est faible, mieux c'est :

• extension de la profondeur : ~0.119 ;
• bouclage : ~0.125 ;
• modèle dense : ~0.137.

C'est bien de l'ICL pur : le test ne laisse rien à retrouver dans les poids, il ne reste qu'à lire le contexte.

Têtes d'induction multipliées et prix du gain

Les auteurs expliquent ce résultat par les têtes d'induction. Ce sont des têtes d'attention qui repèrent dans le contexte une occurrence antérieure du token courant et suggèrent ce qui la suivait. Elles sont considérées comme l'un des principaux mécanismes de l'ICL.

Dans le modèle 3B avec LIDAS, 10 des 12 têtes d'induction les plus fortes descendent de deux ancêtres issus du réseau initial à quatre couches. Le motif trouvé dans le petit modèle se copie sur toute la profondeur au fil de la croissance. Il ne semble pas être ensuite évincé par les connaissances que le réseau mémorise dans ses poids.

Une autre vérification consiste à désactiver ces têtes. Les auteurs ont coupé les 20% de têtes les plus fortes et mesuré la baisse du F1 moyen sur 19 tâches à quelques exemples :

• modèle dense : baisse d'environ 0.08 ;
• LIDAS : environ 0.18 ;
• modèle bouclé : environ 0.28.

Les modèles étendus et bouclés s'appuient donc davantage sur ces motifs, et toute leur capacité d'apprentissage par l'exemple y est concentrée.

Ce que les ingénieurs doivent retenir

Pour ceux qui préentraînent des modèles, la conclusion est pratique. Si l'on veut économiser du calcul tout en obtenant un bon ICL, l'extension progressive de la profondeur paraît une option raisonnable. À 3B, elle a donné à la fois de meilleurs résultats et 26% de FLOPs économisés.

Le bouclage coûte plus cher. Le modèle compte moins de paramètres uniques, et il extrait moins bien les connaissances stockées dans les poids. Le gain en ICL se paie : le modèle lit mieux le contexte, mais retient moins bien les faits.

Il y a aussi une réserve sur la fiabilité. Chaque configuration de 360M à 3B n'a été entraînée qu'avec une seule graine (seed), c'est-à-dire un seul état initial du générateur de nombres aléatoires. La dispersion entre les exécutions n'a pas été mesurée. Sept points de pourcentage sur GSM8K paraissent convaincants, mais sans répétitions, impossible de dire s'ils sont robustes. Une taille jusqu'à 3B reste en outre modeste pour les standards actuels, et l'extrapolation des conclusions à des dizaines de milliards de paramètres doit être vérifiée séparément.

Le résultat n'en est pas moins inattendu. Le modèle dense, entraîné de la manière habituelle, se révèle moins bon précisément là où on le croyait fort : l'apprentissage à partir d'exemples dans le contexte. L'explication des auteurs, la copie précoce de motifs d'attention, peut être testée à d'autres tailles.