Mécanisme d'Attention Innovant Réduction du Cache KV

Introduction

L'utilisation croissante des grands modèles de langage (LLM) et l'émergence de nouveaux paradigmes d'inférence ont mis en évidence le défi d'une inférence efficace à grande échelle. Un goulot d'étranglement important est le cache Key-Value (KV) au sein des mécanismes d'attention traditionnels. Ce cache s'étend linéairement avec la taille du lot et la longueur de la séquence, devenant un véritable "gouffre à mémoire" qui entrave la mise à l'échelle et l'expansion des LLM.

Bien que des variantes comme MQA, GQA et MLA aient émergé pour résoudre ce problème, elles ont souvent du mal à maintenir les performances sous des contraintes de mémoire strictes ou introduisent des complexités qui posent des défis d'ingénierie et des problèmes de compatibilité.

Multi-matrix Factorization Attention (MFA)

Un article récent de chercheurs de Stepes, de l'université de Tsinghua et d'autres institutions présente une nouvelle architecture de mécanisme d'attention : le Multi-matrix Factorization Attention (MFA) et sa variante MFA-Key-Reuse (MFA-KR). Ce mécanisme réduit considérablement le coût de l'inférence des modèles de langage tout en améliorant simultanément les performances.

MFA et MFA-KR surpassent non seulement MLA en termes de performances, mais égalent également les performances traditionnelles de MHA tout en réduisant l'utilisation du cache KV jusqu'à 93,7 %.

MFA est conçu pour la simplicité, la reproduction facile, la faible sensibilité aux hyperparamètres et la compatibilité avec diverses méthodes d'intégration de position.

Approche et Analyse de MFA

L'équipe de recherche a analysé la conception générale et la capacité des mécanismes d'attention, identifiant deux dimensions critiques liées à la capacité. Cette analyse a conduit au développement de nouvelles méthodes analytiques et de principes de conception.

Ils ont introduit le concept de Generalized Multi-Head Attention (GMHA) comme un cadre unificateur pour comprendre les différentes variantes de MHA.

L'équipe a également exploré le calcul et le stockage des key-values d'un point de vue inférence et a examiné la capacité du modèle d'un point de vue décomposition.

Le Fully Parameterized Bilinear Attention (FPBA) a été établi comme la limite supérieure théorique des performances. Ils ont constaté que MHA et ses variantes sont des décompositions de faible rang de FPBA.

Comparaison avec MQA et MLA

L'analyse s'est concentrée sur deux schémas d'amélioration représentatifs : le Multi-Query Attention (MQA) et le Multi-Head Latent Attention (MLA).

MQA utilise une stratégie de partage de paramètres plus agressive, où toutes les têtes d'attention partagent le même ensemble de paramètres key-value. Cela réduit l'utilisation de la mémoire, mais peut affecter l'expressivité du modèle.

MLA introduit un espace latent partagé pour la compression des paramètres, mais la puissance expressive réelle est limitée par la plus petite dimension, ce qui signifie que l'augmentation des dimensions intermédiaires n'améliore pas significativement les performances.

Innovations Clés de MFA

Le développement de MFA a été motivé par l'objectif de créer un mécanisme d'attention qui minimise la consommation de ressources tout en approchant les limites théoriques de performance. La conception de MFA intègre trois innovations clés :

Augmenter significativement le nombre et la dimension des têtes d'attention pour maximiser la capacité du modèle.
Employer une stratégie de décomposition de faible rang agressive pour maintenir l'efficacité des paramètres tout en augmentant le nombre et les dimensions des têtes d'attention.
Utiliser une conception de tête key-value unique pour maintenir une consommation de mémoire minimale, même avec une complexité de modèle accrue.

Mesure et Comparaison de la Capacité

Pour analyser davantage MFA et d'autres mécanismes d'attention, l'équipe a introduit deux métriques clés :

Total Effective Rank (TER) : Le produit du nombre de têtes d'attention et du rang de factorisation par tête (FRH).
Shared Latent Subspace Dimension (SLSD) : La dimension de l'espace caché partagé par toutes les têtes d'attention.

MFA atteint un SLSD et un TER plus élevés par rapport à MQA.

Par rapport à MLA, MFA atteint une taille de cache KV plus petite et un TER plus élevé avec des budgets de paramètres similaires, tout en maintenant un SLSD comparable.

Par rapport au MHA traditionnel, MFA a un TER plus élevé, même si son SLSD est plus petit.

Résultats Expérimentaux

Des expériences approfondies ont été menées pour évaluer les performances de la nouvelle architecture à plus grande échelle, en testant des modèles allant de 1 milliard à 7 milliards de paramètres et des données d'entraînement allant de 10 milliards à 1 trillion.

MFA a démontré des capacités de mise à l'échelle comparables au MHA traditionnel, en maintenant d'excellentes performances même à plus grande échelle.

Alors que MFA-KR présentait des performances légèrement inférieures, sa tendance à la mise à l'échelle s'alignait sur MHA. Les avantages en termes d'économie de mémoire de MFA et MFA-KR ont continué de s'étendre avec la taille du modèle, MFA réalisant une économie de mémoire de 87,5 % et MFA-KR réduisant l'utilisation de la mémoire à 6,25 % à la plus grande échelle.

Études d'Ablation

Les études d'ablation ont validé l'efficacité de MFA et MFA-KR. Leurs avantages en termes de performances ont également été confirmés dans diverses méthodes d'encodage positionnel courantes.

Perspectives

MFA offre des améliorations significatives avec une conception simple, en s'attaquant efficacement au goulot d'étranglement de la mémoire dans l'inférence LLM sans ajouter de complexité d'ingénierie supplémentaire. Il s'intègre de manière transparente dans l'écosystème Transformer existant, accélérant l'application des LLM dans divers scénarios.