Résumé technique : Revoir la convention de forme des modèles de langage Transformer

Énoncé du problème

Les modèles de langage Transformer denses ont historiquement adhéré à une convention architecturale rigide : le réseau de neurones à propagation avant (FFN - Feed-Forward Network) utilise une forme de perceptron multicouche (MLP) « étroit–large–étroit ». Dans cette conception, la dimension du modèle ($d_{model}$) est étendue vers une dimension intermédiaire plus large ($d_h$) avec un ratio d'expansion typiquement compris entre 2 et 4, avant d'être projetée à nouveau. Cette structure domine l'allocation des paramètres, comptant souvent pour la majorité des paramètres et des opérations de calcul en virgule flottante du modèle.

Les auteurs se demandent si cette forme spécifique est essentielle ou s'il ne s'agit que d'une instanciation pratique. Des travaux théoriques récents suggèrent que les MLP résiduels de type « large–étroit–large » (sablier ou hourglass) peuvent servir d'approximateurs de fonctions optimaux, offrant potentiellement une expressivité supérieure avec moins de paramètres. De plus, des preuves empiriques dans les modèles de vision génératifs et les Transformers de type Mixture-of-Experts (MoE) ont montré que les structures en sablier peuvent surpasser les conceptions standards. Le problème central abordé est de savoir si le FFN conventionnel « étroit–large–étroit » est une contrainte nécessaire ou si un FFN en sablier plus profond et plus étroit peut offrir une allocation plus efficace de la capacité entre les modules d'attention et de FFN.

Méthodologie

Le papier propose un variant de Transformer qui remplace le FFN conventionnel par un FFN en Sablier (Hourglass FFN). La méthodologie implique :

1. Redesign Architectural :

   • Le FFN en Sablier remplace la couche d'expansion large unique par une pile de $K$ sous-MLP en forme de sablier connectés par des voies résiduelles.
   • Chaque sous-MLP ($MLP_i$) suit une structure de compression-expansion : il projette l'entrée de $d_{model}$ vers une dimension de goulot d'étranglement $d_h$ (où $d_h < d_{model}$), applique une activation non linéaire (SwiGLU), et projette à nouveau vers $d_{model}$.
   • La sortie de la couche est le résultat du passage de l'entrée à travers $K$ de ces sous-blocs résiduels.

2. Stratégie de Redistribution des Paramètres :

   • En réduisant la dimension intermédiaire $d_h$ (en rendant le FFN « plus étroit »), les auteurs économisent des paramètres qui sont typiquement alloués à la largeur du FFN.
   • Ces paramètres économisés sont réalloués soit à :
     • L'augmentation de la dimension du modèle ($d_{model}$) et de la capacité d'attention.
     • L'augmentation de la profondeur de la pile FFN ($K$).
     • L'augmentation du nombre de couches Transformer ($L$).

3. Configuration Expérimentale :

   • Baselines : Des comparaisons sont effectuées contre des Transformers de style LLaMA et l'architecture OLMo-2 de pointe.
   • Échelles : Les expériences sont menées sur des tailles de modèles allant de 113M à 1B de paramètres.
   • Contraintes : Les expériences sont réalisées sous des budgets de paramètres totaux identiques (à 0,001 % près) et des tailles d'attention fixes pour isoler les effets de la forme du FFN.
   • Espace de recherche : Les auteurs effectuent une recherche architecturale faisant varier $K$ (profondeur de la pile sablier), $L$ (nombre de couches Transformer), $d_{model}$ et le ratio $d_h/d_{model}$.

Principales Contributions

Le papier apporte quatre contributions majeures :

1. Remise en question du standard : Il conteste la forme de FFN « étroit–large–étroit » par défaut, proposant un FFN en Sablier « large–étroit–large » efficace qui maintient l'expressivité grâce à des sous-MLP empilés et connectés de manière résiduelle.
2. Réallocation des paramètres : Il démontre une stratégie où la redirection des paramètres de la largeur du FFN vers la largeur de l'attention ou la profondeur interne du FFN ($K > 1$) produit des performances supérieures, particulièrement aux petites échelles (113M).
3. Insights Architecturaux : Il identifie un compromis largeur-profondeur en forme de U, trouvant des ratios optimaux pour $d_{model}/L$ et $d_h/d_{model}$, et établit que les structures de FFN en sablier plus profondes sont robustes.
4. Validation de la Scalabilité : Il valide le FFN en Sablier de 113M à 1B de paramètres, montrant qu'il est une alternative compétitive aux baselines standards, avec des configurations spécifiques surpassant les modèles conventionnels à plus petite échelle et les égalant à plus grande échelle.

Résultats

La validation empirique produit les conclusions suivantes :

• Performance à Attention Fixe : À l'échelle 113M avec des paramètres d'attention fixes, les FFN en Sablier atteignent une perte de validation et une perplexité (PPL) plus faibles que les FFN conventionnels. Par exemple, à $K=8$, la PPL est passée de 36,441 (conventionnel) à 36,179.
• Configurations Optimales :
  • Profondeur vs Largeur : Un ratio $d_h/d_{model} \approx 0,4$ avec $K=4$ couches a produit la perte de validation la plus faible.
  • Décalage de Paramètres : À 113M, déplacer les paramètres pour rendre le module d'attention dominant (ex: 66M attention vs 47M FFN) a donné la meilleure performance (PPL 35,101), suggérant que le traitement contextuel est plus précieux que la récupération factuelle statique pour les petits modèles.
  • Compromis Largeur-Profondeur : Une courbe distincte en forme de U a été observée pour le ratio $d_{model}/L$, avec un point optimal idéal entre 100 et 250.
• Scalabilité :
  • Petites à Moyennes Échelles (113M – 906M) : Les variantes en Sablier ont systématiquement surpassé les baselines conventionnelles en termes de perte de validation, de perplexité et de tâches de raisonnement en aval (ex: Arc Easy, HellaSwag, PIQA).
  • Grande Échelle (1B) : À l'échelle 1B de paramètres, la variante en Sablier a atteint une performance comparable à la forte baseline conventionnelle (Val PPL 20,082 vs 20,002) et a surpassé la configuration OLMo-2 (20,466).
• Études d'Ablation : L'augmentation de la profondeur du sablier $K$ a systématiquement amélioré les performances, bien qu'avec des rendements décroissants au-delà de $K=6$. Réduire $d_h/d_{model}$ en dessous de 0,4 a entraîné une dégradation significative des performances, indiquant une limite inférieure pour la compression du FFN.

Signification et Revendications

Le papier affirme que la forme MLP « étroit–large–étroit » de longue date dans les Transformers n'est pas essentielle et peut être redondante. En adoptant un FFN en Sablier, les auteurs démontrent que :

1. Efficacité : Les modèles peuvent atteindre des performances compétitives ou supérieures avec une distribution de paramètres plus efficace.
2. Réallocation des Ressources : La structure en Sablier permet un changement fondamental où le mécanisme d'attention peut devenir le principal consommateur de paramètres, défiant le budget traditionnellement dominé par le FFN.
3. Réévaluation du Design : Les résultats incitent à repenser l'équilibre entre l'attention et le FFN, suggérant que pour les modèles de langage modernes efficaces et expressifs, une conception combinant des FFN plus profonds et plus étroits avec des mécanismes d'attention plus larges peut être une stratégie supérieure.

Les auteurs restent modestes quant à la scalabilité vers des modèles de dizaines ou centaines de milliards de paramètres, notant que leurs résultats ont été extrapolés de l'échelle 113M et que la « malédiction de la profondeur » (problèmes de propagation du signal) nécessite des investigations supplémentaires dans des architectures plus profondes. Ils notent également que l'interaction entre les FFN en Sablier et les mécanismes d'attention avancés (ex: Group-Query Attention) reste inexplorée.