Un transformeur bouclé de 1B de paramètres revient 1,4 fois moins cher en FLOP

Les transformeurs bouclés sont devenus un sujet à la mode cette année : des articles paraissent presque chaque jour. L'idée est simple : au lieu d'ajouter de nouvelles couches, le modèle repasse plusieurs fois par les mêmes. Le nombre de paramètres reste identique, mais le modèle « réfléchit » plus en profondeur. Un nouveau travail de 149 pages, signé de chercheurs parmi lesquels Mike Lewis et Artidoro Pagnoni, répond enfin à une question pratique : comment boucler pour que cela soit rentable.

Où est le problème du bouclage

Imaginez une chaîne de 24 machines. On peut en acheter 24 de plus pour un traitement plus sophistiqué, mais c'est coûteux en mémoire et en paramètres. Ou bien on peut faire repasser la pièce par les mêmes machines pour un deuxième tour. La mémoire nécessaire reste la même, mais le calcul augmente. Cette approche a deux points faibles. Premièrement, on ne sait pas quoi faire tourner en boucle : toutes les couches ou une partie. Deuxièmement, lors de la génération de texte, chaque nouveau token attend la fin de toutes les boucles du précédent, et ce décodage est difficile à paralléliser.

Ce qu'ont découvert les auteurs

L'équipe a pré-entraîné des décodeurs denses jusqu'à 1B de paramètres et jusqu'à 32 boucles, et comparé deux modes. Le Full-Stack Looping répète toutes les couches. Le Middle-Stack Looping ne répète que celles du milieu, les premières et les dernières restant telles quelles.

La conclusion dépend de ce qui vous fait défaut. Si le nombre de paramètres est strictement limité, mieux vaut répéter toutes les couches. Si c'est le calcul (FLOP) qui est limité, les couches centrales l'emportent. Pour un ingénieur, cela se traduit immédiatement en règle de choix d'architecture.

La seconde partie du travail porte sur le décodage parallèle. Dans le schéma Parallel Loop Transformers, les états des boucles sont décalés vers les tokens voisins, de sorte que différentes boucles sont calculées simultanément. Les auteurs ont constaté qu'avec plusieurs boucles la qualité se dégrade, car l'embedding du token courant et l'état précédent sont mélangés dans une même somme. Leur correctif, Normalized Shifted PLT, normalise ces deux entrées séparément via RMSNorm avant de les combiner. Ils ont de plus ajouté des adaptateurs de faible rang partagés pour les boucles (Loop Matrix), en somme une correction de type LoRA, afin que chaque passage diffère légèrement du précédent.

Voici les chiffres. Sur le jeu de données DCLM, à 200M de paramètres, la version à quatre boucles ramène les bits par octet à ~1,060, contre ~1,087 pour le PLT d'origine et ~1,088 pour le modèle sans boucles. À 1B de paramètres, le modèle à deux boucles obtient ~0,86 contre ~0,90 avec 80 tokens par paramètre. La précision sur huit tâches atteint ~47,5 % contre ~46,1 % avec 160 tokens par paramètre. En entraînement long, les modèles à deux boucles nécessitent jusqu'à 1,4 fois moins de FLOP que les modèles sans boucles.

Que faire de tout cela

Si vous n'entraînez pas de modèles from scratch, il n'y a rien à faire pour l'instant. Mais si vous prévoyez un modèle compact maison pour un appareil ou un serveur à mémoire strictement limitée, le bouclage devient un vrai levier. La mémoire nécessaire reste la même, tandis que la qualité progresse. Le schéma de choix est simple : si ce sont les paramètres qui bloquent, prenez la boucle complète ; si ce sont les FLOP, bouclez le milieu. Si la vitesse de génération compte, regardez le décalage normalisé : c'est lui qui rend le décodage parallèle opérationnel.

Limites

Trois réserves importantes. Premièrement, les expériences se sont arrêtées à 1B de paramètres, et l'on ignore si les tendances se maintiennent sur des modèles de plusieurs dizaines de milliards. Deuxièmement, les auteurs n'ont pas mesuré le temps de génération réel. Une économie de FLOP ne signifie pas encore une réponse plus rapide : en pratique, la mémoire et les noyaux de calcul pèsent lourd. Troisièmement, le gain de précision est faible, environ 1,4 point de pourcentage, c'est donc plutôt une optimisation d'ingénierie qu'une percée.

En somme : ce travail est utile comme carte du terrain. Il montre que le bouclage n'est pas une technique unique mais un ensemble de choix aux compromis différents, et indique lequel retenir selon votre budget. Pour l'instant, cela concerne ceux qui entraînent eux-mêmes leurs modèles, mais la tendance est clairement de tirer le maximum de chaque paramètre.