Taipan: Эффективные и выразительные модели языков состояний с селективным вниманием

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

October 24, 2024

Авторы: 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

Аннотация

Эффективное моделирование языка с длинным контекстом остается значительной проблемой в обработке естественного языка (NLP). В то время как трансформеры доминируют в языковых задачах, они испытывают трудности с длинными последовательностями из-за квадратичной вычислительной сложности во время обучения и линейного увеличения затрат памяти во время вывода. Недавние модели пространства состояний (SSM), такие как Mamba, предлагают альтернативы с постоянным использованием памяти, но они показывают худшие результаты в задачах, требующих обширного поиска в контексте. Мы представляем Taipan, новую гибридную архитектуру, которая объединяет Mamba-2 с Слоями Селективного Внимания (SAL). Эти SAL идентифицируют токены, требующие взаимодействия на большие расстояния, удаляют менее важные признаки, а затем улучшают их представления с помощью модуля внимания. Этот подход сбалансированно сочетает эффективность Mamba с производительностью, схожей с трансформером, в задачах, требующих больших объемов памяти. Ограничивая бюджет внимания, Taipan расширяет точные прогнозы на длины контекста до 1 миллиона токенов, сохраняя при этом вычислительную эффективность. Наши эксперименты демонстрируют превосходную производительность Taipan на различных масштабах и задачах, предлагая многообещающее решение для эффективного моделирования языка с длинным контекстом.

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.