Transformers à largeur variable

Résumé

L'augmentation de la taille des modèles, notamment en termes de profondeur et de largeur, a permis des progrès significatifs dans les modèles de langage basés sur les transformeurs. Cependant, la plupart des architectures conservent une largeur constante à travers toutes les couches, allouant un budget fixe de paramètres et de calcul de manière uniforme, bien que différentes couches puissent jouer des rôles computationnels distincts. Dans ce travail, nous étudions empiriquement l'allocation non uniforme des capacités sur la profondeur du réseau en proposant une architecture > <former en forme de sablier. Cette conception maintient des couches plus larges au début et à la fin tout en rétrécissant les couches intermédiaires, en utilisant un mécanisme de redimensionnement résiduel sans paramètre. Sur des modèles de langage à simple décodeur allant de 200 millions à 2 milliards de paramètres (denses) et 3 milliards de paramètres (MoE), notre > <former surpasse systématiquement les bases de référence uniformes appariées en paramètres en termes de perte de modélisation du langage. En réduisant la largeur moyenne des couches, cette architecture nécessite également moins de FLOPs globaux (réduction de 22 % sous des courbes d'échelle ajustées à perte équivalente) et une empreinte mémoire et un coût d'entrée-sortie de cache KV plus faibles (réduction de 15 %). Dans l'analyse, nous montrons que cette structure en goulot d'étranglement produit des représentations qualitativement différentes dans les flux résiduels. Dans l'ensemble, nos résultats démontrent qu'une allocation non uniforme de la largeur peut conduire à une mise à l'échelle plus optimale en termes de ressources des modèles de langage.

English

Scaling model size, specifically depth and width, has driven significant progress in transformer-based language models. However, most architectures maintain a constant width across all layers, allocating a fixed parameter and computation budget evenly despite different layers potentially playing distinct computational roles. In this work, we empirically investigate nonuniform capacity allocation across network depth by proposing a times-shaped > <former architecture. This design maintains wider early and late layers while narrowing the middle layers, utilizing a parameter-free residual resizing mechanism. Across decoder-only language models ranging from 200M to 2B parameters (dense) and 3B parameters (MoE), our > <former consistently outperforms parameter-matched uniform baselines on language modeling loss. By reducing the average layer width, this architecture also requires fewer overall FLOPs (22% reduction under fitted loss-matched scaling curves) and smaller KV cache memory and I/O cost (15% reduction). In analysis, we show that this bottleneck structure results in qualitatively different representations in residual streams. Overall, our results demonstrate that nonuniform width allocation can result in more resource-optimal scaling of language models.