Meilleur score ICML2024 ! Modifiez l'attention comme par magie, permettant aux petits modèles de se battre deux fois plus que les grands modèles.

Améliorez le mécanisme de base de Transformer pour se concentrer, afin que les petits modèles puissent être fabriqués deux fois plus grands !

Dans l'article ICML+2024 à score élevé, l'équipe de Caiyun Technology a construit le framework DCFormer pour remplacer le module d'attention des composants de base de Transformer (MHA) et a proposé l'attention multi-têtes combinable dynamiquement (DCMHA).

DCMHA supprime la liaison fixe de la boucle de sélection de recherche et de la boucle de transformation de la tête d'attention MHA, leur permettant d'être combinées dynamiquement en fonction de l'entrée, ce qui améliore fondamentalement la capacité d'expression du modèle.

La signification originale est que chaque couche a des têtes d'attention H fixes. Maintenant, on peut comprendre que chaque couche a des têtes d'attention H fixes. Maintenant, avec presque la même quantité de paramètres et de puissance de calcul, elle peut être combinée dynamiquement jusqu'à. Têtes d'attention HxH. Le contenu affiné peut exprimer plus clairement le sens du texte original, comme suit : Chaque couche du modèle d'origine contient un nombre fixe de têtes d'attention H. Nous pouvons désormais utiliser

DCMHA plug-and-play pour remplacer MHA dans n'importe quelle architecture Transformer afin d'obtenir une nouvelle architecture DCFormer polyvalente, efficace et évolutive.

Ce travail a été réalisé conjointement par des chercheurs de l'Université des postes et télécommunications de Pékin et la startup d'IA Caiyun Technology.

Le modèle DCPythia-6.9B construit par les chercheurs sur la base de DCFormer est meilleur que le modèle open source Pythia-12B en termes de perplexité avant la formation et d'évaluation des tâches en aval.

Le modèle DCFormer est comparable en performances à ceux des modèles Transformer qui nécessitent 1,7 à 2 fois plus de calculs.

Quelles sont les limites du module d'attention multi-têtes ?

La loi d'échelle des grands modèles nous dit qu'à mesure que la puissance de calcul augmente, le modèle devient plus grand et contient plus de données, et l'effet du modèle deviendra de mieux en mieux. Bien que personne ne puisse expliquer clairement quelle est la hauteur du plafond de cette route et si elle peut atteindre l’AGI, il s’agit effectivement de l’approche la plus courante à l’heure actuelle.

Mais en plus de cela, une autre question mérite également d'être réfléchie : la plupart des grands modèles actuels sont basés sur Transformer. Ils sont construits un par un avec des blocs Transformer comme des blocs de construction, le Transformer lui-même. et quelle est la marge d’amélioration ?

C'est la question fondamentale à laquelle il faut répondre dans la recherche sur la structure des modèles, et c'est également le point de départ des travaux DCFormer réalisés conjointement par Caiyun Technology et l'Université des postes et télécommunications de Pékin.

Dans le module d'attention multi-têtes de Transformer (MHA) , chaque tête d'attention fonctionne de manière complètement indépendante les unes des autres.

Cette conception a connu un grand succès dans la pratique en raison de sa simplicité et de sa facilité de mise en œuvre, mais elle entraîne également un faible classement de la matrice de score d'attention, ce qui affaiblit la capacité d'expression, et une redondance répétée des gaspillages de fonctions de la tête d'attention. paramètres et calculs, ressources et autres inconvénients. Sur cette base, certains travaux de recherche de ces dernières années ont tenté d'introduire une certaine forme d'interaction entre les têtes d'attention.

Selon la théorie des boucles de transformateur, dans MHA, le comportement de chaque tête d'attention est décrit par quatre matrices de poids W^Q, W^K, W^V, W^O (où WO ^{Obtenu par segmentant la matrice de projection de sortie de MHA)}.

Q^WK ^{est appelée la boucle QK} (ou boucle de sélection de recherche) , qui détermine sur quel (certains) token dans le contexte se concentrer à partir du jeton actuel, par exemple :

O^WV ^{est appelée la boucle OV} (ou boucle de transformation de projection), qui détermine quelles informations récupérer du jeton concerné (ou quels attributs projeter) pour écrivez dans le flux résiduel à la position actuelle, puis prédisez le jeton suivant. Par exemple :

Les chercheurs ont remarqué que la recherche (où obtenir) et la transformation (quoi obtenir) sont à l'origine deux choses indépendantes, qui doivent être spécifiées séparément et librement combinées selon les besoins (tout comme dans les requêtes SQL, la sélection les conditions après WHERE et la projection d'attribut après SELECT sont écrites séparément), MHA force à les "regrouper" dans QKOV avec une tête d'attention, ce qui limite la flexibilité et les capacités d'expression.

Par exemple, supposons qu'il existe un modèle avec les têtes d'attention A, B et C dont les boucles QK et OV peuvent compléter l'exemple ci-dessus =, puis remplacez-le par :

Vous devez croiser- combinez les boucles QK et OV des têtes d'attention existantes, le modèle peut "ne pas pouvoir se retourner" (vérifié par l'ensemble de tests synthétiques construit par le système du chercheur, le modèle

À quoi ressemble la combinaison dynamique d’une longue attention ?

Avec ceci comme point de départ, l'équipe de recherche de cet article a introduit l'opération de composition dans MHA :

Comme le montre la figure ci-dessous, DCMHA est obtenu :

△Figure 1. Le La structure globale du DCMHA

sera QW^{Q La matrice de score d'attention A}S ^{et la matrice de poids d'attention A}W ^{calculées par} et KW^K sont mappées linéairement sur la dimension num_heads pour obtenir une nouvelle matrice A' avant étant multiplié par VW^V, à travers différentes matrices de cartographie linéaire (carte de composition) , pour obtenir les effets de diverses combinaisons de têtes d'attention.

Par exemple, sur la figure 2(c), les boucles QK des têtes 3 et 7 sont combinées avec la boucle OV de la tête 1 pour former une « nouvelle » tête d'attention.

△ Figure 2. Fonctions de carte de composition typiques simplifiées de 8 têtes d'attention, les couleurs claires représentent de grandes valeurs

Afin de maximiser les capacités d'expression, les chercheurs espèrent que la matrice de cartographie est dynamiquement pilotée par l'entrée Générer , c'est-à-dire déterminer dynamiquement comment les têtes d'attention sont combinées.

Mais la matrice de mappage qu'ils veulent générer n'en est pas une, mais une telle matrice doit être générée pour chaque paire de requête Q_i à la position source et clé K_j à la position de destination dans la séquence, ce qui nécessite à la fois la surcharge de calcul et l’utilisation de la mémoire, ce sera difficile à accepter.

À cette fin, ils décomposent davantage la matrice de cartographie en la somme d'une matrice statique indépendante de l'entrée W_b, d'une matrice de bas rang w₁w₂ et d'une matrice diagonale Diag(w_g ), Ils sont respectivement responsables des combinaisons de base, des combinaisons dynamiques de manière limitée (c'est-à-dire le rang R entre les têtes d'attention, et du gating dynamique des têtes elles-mêmes (voir Figure 2 (d) et Figure 3 (b )) . Ces deux dernières matrices sont générées dynamiquement par la matrice Q et la matrice K.

Réduisez la complexité des calculs et des paramètres à un niveau presque négligeable sans sacrifier l'effet (voir l'analyse de la complexité dans l'article pour plus de détails). Combiné avec l'optimisation au niveau de la mise en œuvre de JAX et PyTorch, DCFormer peut former et déduire efficacement.

△Figure 3. Comment se déroule le calcul

de Compose ?

Extension d'échelle

Pour évaluer la qualité d'une architecture, l'indicateur principal sur lequel les chercheurs se concentrent est l'efficacité de la conversion de la puissance de calcul en intelligence (ou rapport de puissance de calcul de performance), c'est-à-dire les performances du modèle qui peuvent être atteintes. en investissant dans l'amélioration de la puissance de calcul unitaire - dépensez moins de puissance de calcul et obtenez de meilleurs modèles.

Cela peut être vu à partir des courbes de loi d'échelle sur les figures 4 et 5 (en coordonnées logarithmiques, la perte de chaque architecture de modèle peut être tracée sous la forme d'une ligne droite approximative à mesure que la puissance de calcul change. Plus la perte est faible, mieux c'est le modèle) , DCFormer peut obtenir l'effet du modèle Transformer avec 1,7 à 2 fois la puissance de calcul, c'est-à-dire que le taux de conversion intelligent de la puissance de calcul est augmenté de 1,7 à 2 fois.

△Figure 4. Effet d'expansion d'échelle de Transformer et DCFormer

△Figure 5. Effet d'expansion d'échelle de Pythia et DCPythia

Comment comprendre cette amélioration ?

Depuis la naissance de Transformer en 2017, du point de vue de l'amélioration des performances et du rapport puissance de calcul, GLU MLP et l'encodage de position de rotation RoPE sont deux des rares améliorations architecturales qui se sont révélées universellement efficaces et largement adoptées dans un grand nombre de pratiques.

L'architecture qui ajoute ces deux améliorations au Transformer original est également appelée Transformer++. Les modèles open source les plus puissants tels que Llama et Mistral utilisent tous cette architecture. Quelle que soit l'architecture Transformer ou Transformer++, des améliorations significatives peuvent être obtenues grâce à DCMHA.

À l'échelle du modèle 1,4B, l'amélioration de DCMHA est supérieure à la somme des deux améliorations de Transformer++, et l'évolutivité est meilleure (Comparaison de la ligne bleu-vert et de la ligne noire sur la figure 4, l'amélioration de DCMHA atténue à mesure que la puissance de calcul augmente plus lentement, et comparaison de la figure 4 et de la figure 5).

On peut dire que DCFormer élève les capacités de Transformer à un nouveau niveau.

Évaluation des tâches en aval

L'équipe de recherche a formé deux modèles, DCPythia-2.8B et DCPythia-6.9B, pour évaluer les tâches en aval de la PNL grand public et les a comparés avec le modèle open source Pythia de la même échelle (la formation utilise le même hyperparamètres comme paramètres Pythia) .

△Tableau 1. Performances de DCFormer et Pythia dans les tâches en aval

Comme le montre le tableau 1, DCPythia-2.8B et 6.9B ont non seulement des effectifs inférieurs sur l'ensemble de validation de pile, mais également des effectifs inférieurs. sur la plupart des tâches, il surpasse considérablement Pythia sur les tâches en aval. La précision moyenne de DCPythia6.9B sur les tâches ppl et en aval dépasse même Pythia-12B.

DCFormer++2.8B est encore amélioré par rapport à DCPythia-2.8B, vérifiant l'efficacité de la combinaison de l'architecture DCMHA et Lllama.

Vitesse de formation et d'inférence

Bien que l'introduction de DCMHA entraînera une surcharge de formation et d'inférence supplémentaire, il ressort du tableau 2 que la vitesse de formation de DCFormer++ est de 74,5 % à 89,2 % de celle de Transformer++ et que la vitesse d'inférence est de 81,1 %. -89,7 %, et à mesure que les paramètres du modèle augmentent, la charge de calcul supplémentaire diminuera progressivement.

△Tableau 2. Comparaison des vitesses d'entraînement et d'inférence entre Transformer++ et DCFormer++

La vitesse d'entraînement est comparée lorsque le pod TPU v3, la longueur de la séquence est de 2048, la taille du lot est de 1k, la vitesse d'inférence est lorsque le GPU A100 80G Pour l'évaluation, le la longueur d'entrée est de 1024 et la longueur générée est de 128.

Expérience d'ablation

Les résultats sont les suivants :

△Tableau 3. Expérience d'ablation de DCMHA

Du tableau 3, nous pouvons voir les points suivants :

• Bien que l'ajout de poids de combinaison statique puisse réduire les personnes, mais l'introduction de poids de combinaison dynamique peut réduire davantage le nombre de personnes, ce qui illustre la nécessité d'une combinaison dynamique.
• La combinaison dynamique de bas rang fonctionne mieux que le gate dynamique.
• Le ppl obtenu en utilisant uniquement une combinaison dynamique par requête ou par clé est très similaire, et l'écart avec DCFormer++ est très faible.
• Il est plus efficace de faire une combinaison de têtes d'attention après softmax qu'avant softmax, probablement parce que la probabilité après softmax peut affecter plus directement le résultat.
• Le rang du poids combiné dynamique n'a pas besoin d'être défini trop grand, ce qui illustre également le faible rang du poids combiné.

En outre, les chercheurs ont également réduit davantage les frais de formation et d'inférence en augmentant la proportion de couches d'attention locales et en utilisant uniquement une combinaison dynamique par requête. Voir le tableau 10 de l'article pour plus de détails.

De manière générale, l’équipe de recherche tire deux conclusions.

À propos des poids dynamiques : les travaux récents sur SSM et l'attention linéaire/RNN tels que Mamba, GLA, RWKV6, HGRN, etc. ont rattrapé Transformer++ en introduisant des poids dynamiques (dépendants des entrées), mais DCFormer utilise une combinaison dynamique de têtes d'attention. La méthode montre que lors de l'utilisation de softmax attention, l'effet de Transformer++ peut être considérablement amélioré en introduisant des poids dynamiques.

À propos de l'innovation en matière d'architecture de modèle : ce travail montre que s'il existe une « architecture de modèle idéale » avec une puissance de calcul extrême et une efficacité de transformation intelligente, bien que l'architecture Transformer actuelle soit déjà très puissante, il est probable qu'il reste encore un long chemin à parcourir. partir de cette architecture idéale, il reste encore une énorme marge d’amélioration. Par conséquent, outre le développement vigoureux de miracles en empilant la puissance de calcul et les données, l'innovation dans l'architecture des modèles présente également un grand potentiel.

L'équipe de recherche a également déclaré que Caiyun Technology sera la première à appliquer DCformer sur ses produits Caiyun Weather, Caiyun Xiaoyi et Caiyun Xiaomeng.

Pour plus de détails sur la recherche, veuillez vous référer à l'article original.

Lien papier ICML2024 : https://icml.cc/virtual/2024/poster/34047.
Lien de l'article Arxiv : https://arxiv.org/abs/2405.08553.
Lien du code : https://github.com/Caiyun-AI/DCFormer.