Sparser работает быстрее, и меньше - это больше: эффективное разреженное внимание для трансформеров с большим диапазоном.

Sparser is Faster and Less is More: Efficient Sparse Attention for Long-Range Transformers

June 24, 2024

Авторы: Chao Lou, Zixia Jia, Zilong Zheng, Kewei Tu

cs.AI

Аннотация

Эффективное обработка длинных последовательностей в авторегрессионных трансформерах, особенно в рамках расширенного окна контекста, представляет существенные трудности из-за квадратичной вычислительной сложности и значительных требований к памяти KV, присущих механизмам самовнимания. В данной работе мы представляем SPARSEK Attention, новый разреженный механизм внимания, разработанный для преодоления этих вычислительных и памятных препятствий при сохранении производительности. Наш подход интегрирует сеть оценки и дифференцируемый оператор маскирования top-k, SPARSEK, для выбора постоянного числа пар KV для каждого запроса, тем самым обеспечивая оптимизацию на основе градиентов. В результате SPARSEK Attention предлагает линейную временную сложность и постоянный объем памяти во время генерации. Экспериментальные результаты показывают, что SPARSEK Attention превосходит предыдущие разреженные методы внимания и обеспечивает значительное увеличение скорости как во время обучения, так и вывода, особенно в языковом моделировании и задачах на уровне приложений. Более того, наш метод может быть легко интегрирован в предварительно обученные Большие Языковые Модели (LLM) с минимальной донастройкой, предлагая практическое решение для эффективного управления дальнодействующими зависимостями в различных приложениях.

English

Accommodating long sequences efficiently in autoregressive Transformers, especially within an extended context window, poses significant challenges due to the quadratic computational complexity and substantial KV memory requirements inherent in self-attention mechanisms. In this work, we introduce SPARSEK Attention, a novel sparse attention mechanism designed to overcome these computational and memory obstacles while maintaining performance. Our approach integrates a scoring network and a differentiable top-k mask operator, SPARSEK, to select a constant number of KV pairs for each query, thereby enabling gradient-based optimization. As a result, SPARSEK Attention offers linear time complexity and constant memory footprint during generation. Experimental results reveal that SPARSEK Attention outperforms previous sparse attention methods and provides significant speed improvements during both training and inference, particularly in language modeling and downstream tasks. Furthermore, our method can be seamlessly integrated into pre-trained Large Language Models (LLMs) with minimal fine-tuning, offering a practical solution for effectively managing long-range dependencies in diverse applications.