Taipan: Modelos de Lenguaje de Espacio de Estados Eficientes y Expresivos con Atención Selectiva

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

October 24, 2024

Autores: 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

Resumen

La modelización eficiente de lenguaje de largo contexto sigue siendo un desafío significativo en el Procesamiento del Lenguaje Natural (PLN). Si bien los Transformadores dominan las tareas de lenguaje, tienen dificultades con secuencias largas debido a la complejidad computacional cuadrática en el entrenamiento y al aumento lineal de los costos de memoria durante la inferencia. Los Modelos de Espacio de Estado (MES) recientes como Mamba ofrecen alternativas con uso de memoria constante, pero tienen un rendimiento inferior en tareas que requieren una recuperación extensa en contexto. Presentamos Taipan, una arquitectura híbrida novedosa que combina Mamba-2 con Capas de Atención Selectiva (CAS). Estas CAS identifican tokens que requieren interacciones a larga distancia, eliminan características menos importantes y luego mejoran sus representaciones utilizando el módulo de atención. Este enfoque equilibra la eficiencia de Mamba con el rendimiento similar al de los Transformadores en tareas intensivas en memoria. Al limitar el presupuesto de atención, Taipan extiende predicciones precisas a longitudes de contexto de hasta 1 millón de tokens mientras conserva la eficiencia computacional. Nuestros experimentos demuestran el rendimiento superior de Taipan en diversas escalas y tareas, ofreciendo una solución prometedora para la modelización eficiente de lenguaje de largo contexto.

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.