Taipan : Modèles de langage d'espace d'états efficaces et expressifs avec Attention Sélective

Taipan: Efficient and Expressive State Space Language Models with Selective Attention

October 24, 2024

Auteurs: Chien Van Nguyen, Huy Huu Nguyen, Thang M. Pham, Ruiyi Zhang, Hanieh Deilamsalehy, Puneet Mathur, Ryan A. Rossi, Trung Bui, Viet Dac Lai, Franck Dernoncourt, Thien Huu Nguyen

cs.AI

Résumé

La modélisation efficace du langage sur de longs contextes reste un défi majeur en Traitement Automatique du Langage Naturel (TALN). Bien que les Transformers dominent les tâches linguistiques, ils rencontrent des difficultés avec les séquences longues en raison de la complexité computationnelle quadratique lors de l'entraînement et des coûts de mémoire qui augmentent linéairement lors de l'inférence. Les Modèles d'Espace d'État (SSM) récents tels que Mamba offrent des alternatives avec une utilisation de mémoire constante, mais ils sont moins performants dans les tâches nécessitant une récupération approfondie en contexte. Nous présentons Taipan, une nouvelle architecture hybride qui combine Mamba-2 avec des Couches d'Attention Sélective (CAS). Ces CAS identifient les jetons nécessitant des interactions à longue portée, éliminent les caractéristiques moins importantes, puis augmentent leurs représentations en utilisant le module d'attention. Cette approche équilibre l'efficacité de Mamba avec des performances similaires à celles des Transformers dans les tâches intensives en mémoire. En limitant le budget d'attention, Taipan étend les prédictions précises à des longueurs de contexte allant jusqu'à 1 million de jetons tout en préservant l'efficacité computationnelle. Nos expériences démontrent la performance supérieure de Taipan à différentes échelles et dans diverses tâches, offrant une solution prometteuse pour une modélisation efficace du langage sur de longs contextes.

English

Efficient long-context language modeling remains a significant challenge in Natural Language Processing (NLP). While Transformers dominate language tasks, they struggle with long sequences due to quadratic computational complexity in training and linearly scaling memory costs during inference. Recent State Space Models (SSMs) such as Mamba offer alternatives with constant memory usage, but they underperform in tasks requiring extensive in-context retrieval. We introduce Taipan, a novel hybrid architecture that combines Mamba-2 with Selective Attention Layers (SALs). These SALs identify tokens requiring long-range interactions, remove less important features, and then augment their representations using the attention module. This approach balances Mamba's efficiency with Transformer-like performance in memory-intensive tasks. By constraining the attention budget, Taipan extends accurate predictions to context lengths of up to 1 million tokens while preserving computational efficiency. Our experiments demonstrate Taipan's superior performance across various scales and tasks, offering a promising solution for efficient long-context language modeling.