SeerAttention: Обучение внутреннему разреженному вниманию в ваших LLMs

SeerAttention: Learning Intrinsic Sparse Attention in Your LLMs

October 17, 2024

Авторы: Yizhao Gao, Zhichen Zeng, Dayou Du, Shijie Cao, Hayden Kwok-Hay So, Ting Cao, Fan Yang, Mao Yang

cs.AI

Аннотация

Внимание является основой современных больших языковых моделей (LLM). Однако его квадратичная сложность ограничивает эффективность и масштабируемость LLM, особенно для тех, у которых длинное окно контекста. Многообещающим подходом к преодолению этого ограничения является использование разреженности внимания. Однако существующие решения на основе разреженности в основном полагаются на заранее определенные шаблоны или эвристику для приближения разреженности. Эта практика недостаточно учитывает динамическую природу разреженности внимания в задачах на языке. В данной статье утверждается, что разреженность внимания должна быть изучена, а не заранее определена. Для этого мы разрабатываем механизм внимания SeerAttention, который дополняет обычное внимание обучаемым затвором, который адаптивно выбирает значимые блоки на карте внимания и считает остальные блоки разреженными. Такая разреженность на уровне блоков эффективно балансирует точность и ускорение. Для обеспечения эффективного обучения сети управления, мы разрабатываем настраиваемую реализацию FlashAttention, которая извлекает истину на уровне блоков карты внимания с минимальными накладными расходами. SeerAttention применим не только на этапе пост-обучения, но также превосходит в долгосрочной настройке на контекст. Наши результаты показывают, что на этапе пост-обучения SeerAttention значительно превосходит современные статические или эвристические методы разреженности внимания, а также более гибок и адаптивен к различным длинам контекста и коэффициентам разреженности. Примененный к долгосрочной настройке с YaRN, SeerAttention может достичь замечательного коэффициента разреженности в 90% при длине контекста 32k с минимальной потерей перплексии, обеспечивая ускорение в 5,67 раз по сравнению с FlashAttention-2.

English

Attention is the cornerstone of modern Large Language Models (LLMs). Yet its quadratic complexity limits the efficiency and scalability of LLMs, especially for those with a long-context window. A promising approach addressing this limitation is to leverage the sparsity in attention. However, existing sparsity-based solutions predominantly rely on predefined patterns or heuristics to approximate sparsity. This practice falls short to fully capture the dynamic nature of attention sparsity in language-based tasks. This paper argues that attention sparsity should be learned rather than predefined. To this end, we design SeerAttention, a new Attention mechanism that augments the conventional attention with a learnable gate that adaptively selects significant blocks in an attention map and deems the rest blocks sparse. Such block-level sparsity effectively balances accuracy and speedup. To enable efficient learning of the gating network, we develop a customized FlashAttention implementation that extracts the block-level ground truth of attention map with minimum overhead. SeerAttention not only applies to post-training, but also excels in long-context fine-tuning. Our results show that at post-training stages, SeerAttention significantly outperforms state-of-the-art static or heuristic-based sparse attention methods, while also being more versatile and flexible to adapt to varying context lengths and sparsity ratios. When applied to long-context fine-tuning with YaRN, SeerAttention can achieve a remarkable 90% sparsity ratio at a 32k context length with minimal perplexity loss, offering a 5.67x speedup over FlashAttention-2.