Résumé Technique : Réexamen de la Paramétrisation des Couches de Transformer par Minimisation Causale de l'Énergie

Énoncé du Problème

Les architectures de Transformer, spécifiquement la combinaison de l'Attention Multi-Têtes (MHA) et des Perceptrons Multicouches (MLP) à portes, servent de blocs de construction fondamentaux pour les modèles de langage de grande taille (LLM) modernes. Cependant, les choix spécifiques concernant la paramétrisation des couches — tels que le couplage des poids, les structures de projection et les règles de mise à jour — sont largement motivés par des études empiriques plutôt que par des principes computationnels explicites. Bien que les modèles basés sur l'énergie (EBM) aient précédemment offert une perspective où le calcul est vu comme une recherche d'états de basse énergie (par exemple, dans les réseaux de Hopfield), les travaux existants n'ont pas pleinement connecté ces interprétations basées sur l'énergie aux paramétrisations spécifiques des couches standard de Transformer, ni systématiquement exploré comment cette perspective peut guider la conception de variantes de couches plus efficaces ou plus performantes.

Méthodologie : Minimisation Causale de l'Énergie (CEM)

Les auteurs introduisent la Minimisation Causale de l'Énergie (CEM), un cadre qui reformule les couches de Transformer comme des étapes d'optimisation sur des fonctions d'énergie conditionnelles. La prémisse centrale est qu'une couche mappe une séquence d'entrée $h_{1:J}$ vers une sortie $h'_{1:J}$ en traitant chaque position $i$ comme une variable d'optimisation $x_i$ (initialisée à $h_i$) qui est mise à jour via une procédure $\mathcal{A}$ pour minimiser une énergie conditionnelle $\epsilon(x_i | h_{1:i})$ dépendante de l'historique causal.

Le papier dérive les composants standard de Transformer à partir de ce cadre et propose des extensions :

1. Dérivation de l'Attention à Poids Couplés :
   En définissant une énergie d'interaction basée sur une forme log-sum-exp (similaire aux réseaux de Hopfield modernes), les auteurs montrent qu'une seule étape de descente de gradient retrouve la mise à jour standard MHA. Crucialement, cette dérivation nécessite naturellement des contraintes spécifiques de partage de poids :

   • La projection Clé ($W^K$) est couplée à la projection Valeur ($W^V$).
   • La projection Requête ($W^Q$) est couplée à la projection Sortie ($W^O$).
   • Le paramètre de température $\tau$ est lié à la dimension de la tête.

2. Dérivation des MLP à Portes à Poids Couplés :
   En définissant une énergie élément par élément qui attribue de l'énergie indépendamment au vecteur de caractéristiques de chaque token, une mise à jour par gradient retrouve la structure des MLP à portes. Cette interprétation implique que les matrices de projection « up » et « down » dans le MLP devraient être couplées (spécifiquement, la transposée de la projection down est égale à la projection up).

3. Extensions Proposées :
   S'appuyant sur la perspective d'optimisation par l'énergie, les auteurs proposent trois améliorations architecturales spécifiques :

   • Paramétrisation Diagonale-plus-Rang-Faible (D+LR) : Au lieu de projections purement de rang faible pour la matrice d'interaction $A_k$, ils proposent $A_k = \text{diag}(d_k) + W^{Q\top}_k W^K_k$. Cela capture les interactions que les matrices de rang faible manquent tout en restant computationnellement efficace. Ils proposent également de partager le terme diagonal entre les têtes pour réduire les paramètres.
   • Préconditionneurs Légers Apprentis : Pour améliorer la convergence au-delà d'une seule étape de gradient, ils introduisent des matrices de préconditionnement $P$ avec une structure D+LR ($P = \text{diag}(\text{softplus}(d)) + UV^\top + VU^\top$). Ces matrices agissent comme des proxies appris pour l'information de courbure, redimensionnant les gradients sans le coût des méthodes d'ordre complet.
   • Mises à Jour Récursives au Sein de la Couche : Plutôt qu'une seule mise à jour, le cadre permet $T$ étapes récursives au sein d'une seule couche ($x^{(t+1)}_i = x^{(t)}_i - \eta \nabla \epsilon$). Cela permet à la couche de minimiser sa fonction d'énergie plus profondément sans ajouter de nouveaux paramètres, distinct des Transformers « en boucle » qui réutilisent des blocs entiers.

Contributions Clés

• Unification Théorique : Le papier établit que les couches standard de Transformer avec des schémas spécifiques de partage de poids (Q/O et K/V couplés dans l'attention ; up/down couplés dans les MLP) sont mathématiquement équivalentes à des étapes de gradient uniques sur des fonctions d'énergie d'interaction et élément par élément spécifiques.
• Expansion de l'Espace de Conception : La perspective CEM identifie un nouvel espace de conception pour les couches de Transformer, incluant les interactions D+LR, les préconditionneurs légers et les mises à jour récursives, qui sont motivés par la théorie de l'optimisation plutôt que par des heuristiques.
• Validation Empirique : Les auteurs évaluent les couches dérivées de CEM à une échelle modérée (centaines de millions de paramètres) sur des tâches de modélisation de langage, démontrant que ces paramétrisations contraintes peuvent égaler ou approcher les bases de référence Transformer standard.

Résultats Expérimentaux

Les expériences ont été menées sur le jeu de données SlimPajama en utilisant des tailles de modèle allant d'environ 86M à 162M paramètres, entraînées avec des budgets de tokens optimaux en calcul (mise à l'échelle Chinchilla).

• Efficacité des Paramètres : Les modèles CEM avec poids couplés (utilisant environ la moitié des paramètres MHA et deux tiers des paramètres MLP des bases de référence Llama) ont atteint une perplexité de test compétitive. Pour l'attention CEM, la dégradation des performances était minimale malgré la réduction significative des paramètres.
• Impact des Extensions :
  • Récursion : L'augmentation du nombre d'étapes récursives de $T=1$ à $T=2$ a constamment amélioré les performances pour l'attention CEM et les MLP. Notamment, les variantes d'attention CEM récursives ont surpassé les bases de référence Llama tout en utilisant moins de paramètres.
  • Préconditionneurs : Les préconditionneurs appris n'ont apporté que des gains marginaux par rapport à la récursion, suggérant que la structure récursive est le principal moteur d'amélioration.
  • Termes Diagonaux : L'inclusion d'un terme diagonal dans la matrice d'interaction (spécifiquement un diagonal partagé entre les têtes) était essentielle pour de bonnes performances, surpassant les variantes purement de rang faible ou sans diagonal.
• Stabilité : Les Transformers composés entièrement de couches dérivées de CEM (avec $T=2$ et préconditionneurs D+LR) se sont entraînés de manière stable de bout en bout, atteignant une perplexité comparable aux bases de référence Llama avec moins de paramètres.
• Données Synthétiques : Les expériences sur des données de processus gaussien synthétiques ont confirmé qu'une récursion plus profonde ($T=4, T=8$) produit de nouvelles améliorations en RMSE, bien que les gains dans les configurations complètes de Transformer aient été limités à $T=2$ dans des configurations d'entraînement standard.

Importance et Revendications

Le papier revendique que CEM fournit une lentille théorique utile pour comprendre la paramétrisation des couches de Transformer, connectant efficacement les architectures de Transformer aux modèles basés sur l'énergie. L'importance réside dans :

1. Conception Principielle : Elle fait passer la conception des couches de l'essai-erreur empirique à un cadre fondé sur la minimisation de l'énergie et la dynamique d'optimisation.
2. Efficacité : Elle démontre que le partage de poids et les mises à jour récursives, dérivés de ce cadre, permettent des modèles économes en paramètres qui maintiennent des performances compétitives.
3. Voies Futures : Les auteurs suggèrent que cette perspective motive une exploration plus poussée vers :
   • Kernels Personnalisés : L'optimisation du temps d'exécution pour le partage de poids structuré et les mises à jour récursives (par exemple, fusionner les opérations).
   • Mise à l'Échelle au Moment du Test : Utiliser la récursion au sein de la couche comme un axe supplémentaire pour mettre à l'échelle le calcul lors de l'inférence.
   • Co-conception Matérielle : Concevoir des paramétrisations de couches spécifiquement pour de nouvelles architectures matérielles.

Les auteurs restent modestes, notant que leur évaluation est limitée à l'échelle des centaines de millions de paramètres et que des études à plus grande échelle et des évaluations en aval sont nécessaires pour évaluer pleinement la transférabilité de ces bénéfices aux modèles de l'état de l'art. Ils reconnaissent également que les implémentations actuelles ne sont pas encore optimisées pour les performances d'exécution.