Простые линейные модели внимания языка уравновешивают компромисс между recall (полнотой) и throughput (производительностью).

Simple linear attention language models balance the recall-throughput tradeoff

February 28, 2024

Авторы: Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, Christopher Ré

cs.AI

Аннотация

Недавние исследования показали, что языковые модели на основе внимания отличаются по способности к воспроизведению, к способности привязывать генерации к токенам, ранее увиденным в контексте. Однако эффективность моделей на основе внимания ограничивается в процессе вывода из-за агрессивного потребления памяти KV-кэшем. В данной работе мы исследуем, можно ли улучшить эффективность языковой модели (например, путем снижения потребления памяти) без ущерба для воспроизведения. Применяя эксперименты и теорию к широкому набору архитектур, мы выявляем ключевой компромисс между размером состояния модели и способностью к воспроизведению. Мы показываем, что эффективные альтернативы вниманию (например, H3, Mamba, RWKV) поддерживают фиксированный рекуррентный стейт, но испытывают трудности с воспроизведением. Мы предлагаем BASED - простую архитектуру, объединяющую линейное и скользящее оконное внимание. Изменяя размер окна BASED и размер признаков линейного внимания, мы можем настраивать размер состояния и двигаться по кривой компромисса между воспроизведением и памятью, восстанавливая полное качество внимания с одной стороны и небольшой размер состояния альтернатив вниманию - с другой. Мы обучаем языковые модели до 1.3 млрд параметров и показываем, что BASED соответствует самым мощным субквадратичным моделям (например, Mamba) по перплексии и превосходит их на реальных задачах с высоким воспроизведением на 6.22 пункта точности. Реализации линейного внимания часто менее эффективны, чем оптимизированные стандартные реализации внимания. Для того чтобы сделать BASED конкурентоспособным, мы разрабатываем алгоритмы, учитывающие ввод-вывод, которые обеспечивают в 24 раз большую производительность при генерации текста, чем FlashAttention-2, при генерации 1024 токенов с использованием моделей с 1.3 млрд параметров. Код для этой работы предоставлен по ссылке: https://github.com/HazyResearch/based.

English

Recent work has shown that attention-based language models excel at recall, the ability to ground generations in tokens previously seen in context. However, the efficiency of attention-based models is bottle-necked during inference by the KV-cache's aggressive memory consumption. In this work, we explore whether we can improve language model efficiency (e.g. by reducing memory consumption) without compromising on recall. By applying experiments and theory to a broad set of architectures, we identify a key tradeoff between a model's state size and recall ability. We show that efficient alternatives to attention (e.g. H3, Mamba, RWKV) maintain a fixed-size recurrent state, but struggle at recall. We propose BASED a simple architecture combining linear and sliding window attention. By varying BASED window size and linear attention feature dimension, we can dial the state size and traverse the pareto frontier of the recall-memory tradeoff curve, recovering the full quality of attention on one end and the small state size of attention-alternatives on the other. We train language models up to 1.3b parameters and show that BASED matches the strongest sub-quadratic models (e.g. Mamba) in perplexity and outperforms them on real-world recall-intensive tasks by 6.22 accuracy points. Implementations of linear attention are often less efficient than optimized standard attention implementations. To make BASED competitive, we develop IO-aware algorithms that enable 24x higher throughput on language generation than FlashAttention-2, when generating 1024 tokens using 1.3b parameter models. Code for this work is provided at: https://github.com/HazyResearch/based.