Taipan: Modelos de Linguagem de Espaço de Estado Eficientes e Expressivos com Atenção Seletiva

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

Resumo

A modelagem eficiente de linguagem de longo contexto continua sendo um desafio significativo no Processamento de Linguagem Natural (PLN). Enquanto os Transformers dominam as tarefas de linguagem, eles enfrentam dificuldades com sequências longas devido à complexidade computacional quadrática no treinamento e ao aumento linear dos custos de memória durante a inferência. Modelos recentes de Espaço de Estado (SSMs), como o Mamba, oferecem alternativas com uso de memória constante, mas têm desempenho inferior em tarefas que exigem extensa recuperação no contexto. Apresentamos o Taipan, uma arquitetura híbrida inovadora que combina o Mamba-2 com Camadas de Atenção Seletiva (SALs). Essas SALs identificam tokens que requerem interações de longo alcance, removem características menos importantes e depois aumentam suas representações usando o módulo de atenção. Esse método equilibra a eficiência do Mamba com o desempenho semelhante ao Transformer em tarefas intensivas em memória. Ao limitar o orçamento de atenção, o Taipan estende previsões precisas para comprimentos de contexto de até 1 milhão de tokens, preservando a eficiência computacional. Nossos experimentos demonstram o desempenho superior do Taipan em várias escalas e tarefas, oferecendo uma solução promissora para a modelagem eficiente de linguagem de longo 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.