LoopFormer: Elastic-Depth Looped Transformers for Latent Reasoning via Shortcut Modulation

1. Problématique

Les Transformers en boucle (ou looped Transformers) sont une classe de modèles efficaces qui partagent les mêmes paramètres à travers plusieurs itérations (boucles) pour réaliser un raisonnement latent. Bien que ces modèles aient démontré des performances prometteuses sur des tâches algorithmiques et de raisonnement, ils souffrent d'une limitation fondamentale : ils sont généralement entraînés et évalués avec un nombre fixe d'itérations.

Cette rigidité pose deux problèmes majeurs :

1. Manque d'adaptabilité au budget de calcul : Une fois entraîné pour un nombre spécifique de boucles (ex: 8 itérations), le modèle ne peut pas être utilisé efficacement avec un budget différent (ex: 4 ou 12 itérations) sans re-entraînement.
2. Effondrement des représentations : Lorsque l'on tente d'utiliser un modèle entraîné avec un nombre fixe de boucles à une profondeur différente (plus courte ou plus longue), les représentations internes tendent à s'effondrer ou à stagner. Les itérations supplémentaires n'améliorent pas le raisonnement mais convergent vers des états dégénérés, rendant l'approche peu utile pour un calcul élastique.

L'objectif est donc de concevoir un modèle capable d'un raisonnement élastique : un seul modèle capable de produire des représentations de haute qualité et de raffiner son raisonnement dynamiquement en fonction d'un budget de calcul spécifié par l'utilisateur, sans re-entraînement.

2. Méthodologie : LoopFormer

Les auteurs proposent LoopFormer, un Transformer en boucle conçu pour gérer des trajectoires de longueur variable grâce à une modulation par raccourcis (shortcut modulation) et un entraînement sur des trajectoires de longueurs variables.

A. Architecture et Conditionnement

Contrairement aux approches précédentes qui conditionnent simplement sur l'index de la boucle, LoopFormer conditionne chaque itération $i$ sur deux variables continues :

• Le temps normalisé ($t_{i-1}$) : La position cumulative dans la trajectoire (de 0 à 1).
• La taille de l'étape ($\Delta_i$) : La granularité du pas de temps actuel.

Ces deux scalaires sont encodés via des embeddings sinusoïdaux et projetés par de petits MLP pour générer un signal de conditionnement $c_i$. Ce signal module les couches du bloc Transformer via un mécanisme de type AdaLN (Adaptive Layer Normalization) :

• Il ajuste les échelles (scales) des couches de normalisation RMSNorm.
• Il contrôle les portes (gates) appliquées aux résidus des modules MHSA (Self-Attention) et FFN (Feed-Forward Network).

Cela permet au modèle de comprendre où il se trouve dans la trajectoire et la résolution de l'étape actuelle, assurant une cohérence entre des trajectoires grossières (peu d'étapes) et fines (beaucoup d'étapes).

B. Entraînement : Cohérence des Raccourcis (Shortcut-Consistency)

Pour apprendre à généraliser à n'importe quel budget, LoopFormer utilise une stratégie d'entraînement innovante inspirée des modèles de diffusion et des équations différentielles neuronales :

1. Échantillonnage de trajectoires : À chaque batch, le modèle est exposé à deux types de trajectoires :
   • Une trajectoire complète de longueur $L$ (le budget maximal).
   • Une trajectoire "raccourcie" de longueur $S$ (où $S < L$), avec un schéma d'étapes $\Delta_S$ échantillonné aléatoirement.
2. Objectif de perte : La fonction de perte combine trois termes :
   • $L_L$ : Perte de prédiction du token suivant sur la trajectoire complète.
   • $L_S$ : Perte de prédiction du token suivant sur la trajectoire raccourcie.
   • $L_{cons}$ : Perte de cohérence (avec arrêt du gradient). Elle force les logits de la trajectoire raccourcie à correspondre à ceux de la trajectoire complète (en utilisant les logits de la trajectoire complète comme cible fixe).

Cet objectif de cohérence agit comme une forme d'auto-distillation : il apprend au modèle que même avec peu d'étapes, la représentation finale doit être informative et alignée avec celle d'une exécution plus longue.

C. Inférence Élastique

À l'inférence, l'utilisateur spécifie un budget de calcul $M$ (nombre de boucles) et un schéma d'étapes $\Delta_M$. Le modèle exécute simplement $M$ itérations en utilisant les paramètres appris. Grâce au conditionnement temporel et à la perte de cohérence, le modèle ne s'effondre pas : les représentations s'améliorent progressivement à mesure que le budget augmente.

3. Contributions Clés

1. Formulation de modèles en boucle à modulation de raccourcis : Une architecture qui conditionne chaque passage sur le temps normalisé et la taille de l'étape, permettant une gestion fine de la granularité du calcul.
2. Protocole d'entraînement à cohérence de raccourcis : Une méthode permettant d'entraîner un modèle unique sur une famille de trajectoires, rendant possible l'inférence sous budget contrôlé sans re-entraînement.
3. Analyse des dynamiques de représentation : Les auteurs utilisent des métriques géométriques et informationnelles (anisotropie, courbure, entropie, similarité CKA) pour démontrer que, contrairement aux modèles à sortie précoce (early exit) qui stagnent, LoopFormer maintient une évolution dynamique et non dégénérée des états cachés.
4. Performances robustes : Démonstration que LoopFormer maintient des performances élevées en modélisation du langage et en raisonnement zéro-shot, même sous des contraintes de calcul agressives, tout en surpassant les modèles en boucle baselines.

4. Résultats Expérimentaux

Les expériences ont été menées sur des modèles de ~1 milliard de paramètres, comparés à des Transformers non en boucle (iso-FLOP) et à d'autres variantes en boucle (TMLT, modèles à sortie précoce).

• Perplexité et Raisonnement : LoopFormer réduit considérablement l'écart de perplexité par rapport aux modèles non en boucle de référence. Sur des tâches de raisonnement zéro-shot (COPA, HellaSwag, ARC, etc.), il surpasse systématiquement les autres modèles en boucle, en particulier aux budgets réduits.
• Élasticité de la profondeur :
  • Aux budgets réduits (ex: 6 ou 12 itérations au lieu de 24), les modèles baselines (comme Naive-Loop-EE) voient leurs performances chuter drastiquement.
  • LoopFormer, en revanche, maintient une courbe de performance lisse : la qualité de la prédiction s'améliore continuellement à mesure que le nombre de boucles augmente, sans effondrement.
• Analyse des représentations :
  • Les modèles baselines montrent une stagnation : la similarité CKA entre les étapes reste élevée et l'entropie des représentations diminue rapidement, indiquant un effondrement.
  • LoopFormer montre une évolution progressive : la courbure et l'entropie augmentent au milieu de la trajectoire avant de se stabiliser, prouvant que le modèle continue de raffiner l'information à chaque étape.
• Choix de la trajectoire : L'étude montre que pour un budget fixe, l'allocation des étapes compte : il est préférable d'utiliser des étapes plus grossières au début et plus fines vers la fin pour optimiser la perplexité et le raisonnement.

5. Signification et Perspectives

Ce travail ouvre une nouvelle voie pour les modèles de langage à budget contrôlé.

• Efficacité : Il permet d'adapter dynamiquement la puissance de calcul d'un modèle à la difficulté d'une tâche ou aux contraintes matérielles, sans sacrifier la qualité.
• Raisonnement Latent : Il valide l'hypothèse que les architectures en boucle possèdent un biais inductif vers le raisonnement latent, à condition que les trajectoires soient entraînées pour être cohérentes à différentes résolutions.
• Limites et Futur : L'approche actuelle utilise un budget global au niveau de la séquence. Les travaux futurs pourraient explorer des politiques de budget adaptatives au niveau du token ou de l'instance, ainsi qu'une analyse théorique plus profonde de l'espace des représentations.

En résumé, LoopFormer transforme les Transformers en boucle d'architectures rigides en systèmes élastiques capables de "penser" plus ou moins longtemps selon les besoins, tout en maintenant une cohérence interne robuste.