Le Scalable-Softmax est supérieur pour l'attention.

Scalable-Softmax Is Superior for Attention

January 31, 2025

Auteurs: Ken M. Nakanishi

cs.AI

Résumé

L'élément maximal du vecteur produit par la fonction Softmax tend vers zéro à mesure que la taille du vecteur d'entrée augmente. Les modèles de langage basés sur les Transformers s'appuient sur le Softmax pour calculer les scores d'attention, ce qui entraîne une distribution d'attention qui s'aplatit à mesure que la taille du contexte augmente. Cela réduit la capacité du modèle à prioriser efficacement les informations clés et limite potentiellement sa capacité de généralisation de longueur. Pour résoudre ce problème, nous proposons Scalable-Softmax (SSMax), qui remplace le Softmax dans les scénarios où la taille du vecteur d'entrée varie. SSMax peut être intégré de manière transparente dans les architectures existantes basées sur les Transformers. Les résultats expérimentaux en modélisation de langage montrent que les modèles utilisant SSMax permettent non seulement d'obtenir une réduction plus rapide de la perte lors de la pré-entraînement, mais améliorent également de manière significative les performances dans les contextes longs et la récupération des informations clés. De plus, une analyse des scores d'attention révèle que SSMax permet au modèle de se concentrer sur les informations clés même dans les contextes longs. De plus, bien que les modèles utilisant SSMax dès le début de la pré-entraînement obtiennent une meilleure généralisation de longueur, ceux qui ont déjà commencé la pré-entraînement peuvent encore acquérir cette capacité en remplaçant le Softmax dans les couches d'attention par SSMax, soit pendant, soit après la pré-entraînement.

English

The maximum element of the vector output by the Softmax function approaches zero as the input vector size increases. Transformer-based language models rely on Softmax to compute attention scores, causing the attention distribution to flatten as the context size grows. This reduces the model's ability to prioritize key information effectively and potentially limits its length generalization. To address this problem, we propose Scalable-Softmax (SSMax), which replaces Softmax in scenarios where the input vector size varies. SSMax can be seamlessly integrated into existing Transformer-based architectures. Experimental results in language modeling show that models using SSMax not only achieve faster loss reduction during pretraining but also significantly improve performance in long contexts and key information retrieval. Furthermore, an analysis of attention scores reveals that SSMax enables the model to focus attention on key information even in long contexts. Additionally, although models that use SSMax from the beginning of pretraining achieve better length generalization, those that have already started pretraining can still gain some of this ability by replacing Softmax in the attention layers with SSMax, either during or after pretraining.