Le Prism Transformer introduit un programme de têtes progressif qui varie le nombre de têtes d'attention entre les couches, en commençant par moins de têtes plus larges dans les premières couches et en augmentant monotoniquement avec la profondeur. Cette approche défie l'allocation uniforme standard en répondant aux besoins structurels distincts des premières et dernières couches sans ajouter de surcharge architecturale.
- Les premières couches utilisent des sous-espaces par tête larges (dh=256) pour une capture riche de motifs locaux, tandis que les dernières couches utilisent de nombreuses têtes étroites pour une décomposition spécialisée.
- Les matrices de poids conservent des formes standard dmodel×dmodel, maintenant le nombre de paramètres neutre.
- Le total des FLOPs reste mathématiquement invariant par rapport au nombre de têtes, assurant la neutralité du calcul.
- Les dimensions de tête de puissance de 2 (dh ∈ {256, 128}) préservent l'alignement des Tensor Core pour la neutralité du débit.
- Les résultats montrent une perte de validation inférieure à chaque échelle par rapport aux bases uniformes avec le même nombre de tokens/seconde et le même temps d'exécution.
- Le modèle atteint des gains ou une parité sur les benchmarks incluant PIQA, WinoGrande, HellaSwag et ARC-Easy.
L'analyse de la distance d'attention par couche confirme que le gain est structurel, car les premières couches Prism s'attendent plus localement avant de passer à l'intégration globale. L'implémentation nécessite uniquement un changement d'une ligne par couche d'attention pour rendre le nombre de têtes dépendant de la couche.