Замкнутое внимание на слоты для эффективного моделирования последовательностей за линейное время

Gated Slot Attention for Efficient Linear-Time Sequence Modeling

September 11, 2024

Авторы: Yu Zhang, Songlin Yang, Ruijie Zhu, Yue Zhang, Leyang Cui, Yiqiao Wang, Bolun Wang, Freda Shi, Bailin Wang, Wei Bi, Peng Zhou, Guohong Fu

cs.AI

Аннотация

Линейные трансформеры внимания и их варианты с воротными элементами, известные своей способностью к параллельному обучению и эффективному рекуррентному выводу, все еще уступают в задачах, требующих высокой памяти, по сравнению с традиционными трансформерами и требуют значительных ресурсов для обучения с нуля. В данной статье представлено Воротное Слотовое Внимание (GSA), которое улучшает внимание с Ограниченным-память-Контролем (ABC), интегрируя механизм ворот, вдохновленный Воротным Линейным Вниманием (GLA). По сути, GSA состоит из двухслойного GLA, связанных через софтмакс, используя контекстно-осознанное чтение памяти и адаптивное забывание для увеличения объема памяти, сохраняя компактный размер рекуррентного состояния. Этот дизайн значительно улучшает как эффективность обучения, так и вывода благодаря аппаратно-эффективному алгоритму обучения GLA и уменьшенному размеру состояния. Кроме того, сохранение операции софтмакс особенно полезно в сценариях "донастройки предварительно обученных трансформеров к RNN" (T2R), уменьшая необходимость в обширном обучении с нуля. Обширные эксперименты подтверждают превосходное качество работы GSA в сценариях, требующих воспоминания в контексте, а также в настройках T2R.

English

Linear attention Transformers and their gated variants, celebrated for enabling parallel training and efficient recurrent inference, still fall short in recall-intensive tasks compared to traditional Transformers and demand significant resources for training from scratch. This paper introduces Gated Slot Attention (GSA), which enhances Attention with Bounded-memory-Control (ABC) by incorporating a gating mechanism inspired by Gated Linear Attention (GLA). Essentially, GSA comprises a two-layer GLA linked via softmax, utilizing context-aware memory reading and adaptive forgetting to improve memory capacity while maintaining compact recurrent state size. This design greatly enhances both training and inference efficiency through GLA's hardware-efficient training algorithm and reduced state size. Additionally, retaining the softmax operation is particularly beneficial in "finetuning pretrained Transformers to RNNs" (T2R) settings, reducing the need for extensive training from scratch. Extensive experiments confirm GSA's superior performance in scenarios requiring in-context recall and in T2R settings.