Attention Multi-Jetons

Résumé

L'attention douce est un mécanisme essentiel permettant aux LLM de localiser les parties pertinentes dans un contexte donné. Cependant, les poids d'attention individuels sont déterminés par la similarité d'un seul vecteur de jeton de requête et de clé. Cette "attention à un seul jeton" limite la quantité d'informations utilisées pour distinguer une partie pertinente du reste du contexte. Pour résoudre ce problème, nous proposons une nouvelle méthode d'attention, l'attention multi-jetons (MTA), qui permet aux LLM de conditionner leurs poids d'attention sur plusieurs vecteurs de requête et de clé simultanément. Cela est réalisé en appliquant des opérations de convolution sur les requêtes, les clés et les têtes, permettant ainsi aux requêtes et clés voisines d'influencer mutuellement leurs poids d'attention pour une attention plus précise. En conséquence, notre méthode peut localiser un contexte pertinent en utilisant des informations plus riches et nuancées, dépassant la capacité d'un seul vecteur. Grâce à des évaluations approfondies, nous démontrons que la MTA améliore les performances sur une gamme de benchmarks populaires. Notamment, elle surpasse les modèles de base Transformer sur les tâches standard de modélisation du langage, ainsi que sur les tâches nécessitant la recherche d'informations dans des contextes longs, où la capacité de notre méthode à exploiter des informations plus riches s'avère particulièrement bénéfique.

English

Soft attention is a critical mechanism powering LLMs to locate relevant parts within a given context. However, individual attention weights are determined by the similarity of only a single query and key token vector. This "single token attention" bottlenecks the amount of information used in distinguishing a relevant part from the rest of the context. To address this issue, we propose a new attention method, Multi-Token Attention (MTA), which allows LLMs to condition their attention weights on multiple query and key vectors simultaneously. This is achieved by applying convolution operations over queries, keys and heads, allowing nearby queries and keys to affect each other's attention weights for more precise attention. As a result, our method can locate relevant context using richer, more nuanced information that can exceed a single vector's capacity. Through extensive evaluations, we demonstrate that MTA achieves enhanced performance on a range of popular benchmarks. Notably, it outperforms Transformer baseline models on standard language modeling tasks, and on tasks that require searching for information within long contexts, where our method's ability to leverage richer information proves particularly beneficial.