Atenção por Espaços com Portões para Modelagem Eficiente de Sequências em Tempo Linear

Gated Slot Attention for Efficient Linear-Time Sequence Modeling

September 11, 2024

Autores: 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

Resumo

Os Transformers de atenção linear e suas variantes com portões, conhecidos por possibilitar treinamento paralelo e inferência recorrente eficiente, ainda ficam aquém em tarefas intensivas em recall em comparação com os Transformers tradicionais e exigem recursos significativos para treinamento do zero. Este artigo apresenta a Atenção por Slot com Portões (GSA), que aprimora a Atenção com Controle de Memória Limitada (ABC) incorporando um mecanismo de portões inspirado na Atenção Linear com Portões (GLA). Essencialmente, o GSA consiste em um GLA de duas camadas conectado via softmax, utilizando leitura de memória consciente do contexto e esquecimento adaptativo para melhorar a capacidade de memória mantendo um tamanho compacto de estado recorrente. Este design melhora significativamente a eficiência tanto no treinamento quanto na inferência através do algoritmo de treinamento eficiente em hardware do GLA e do tamanho reduzido do estado. Além disso, manter a operação softmax é particularmente benéfico em configurações de "ajuste fino de Transformers pré-treinados para RNNs" (T2R), reduzindo a necessidade de treinamento extensivo do zero. Experimentos extensos confirmam o desempenho superior do GSA em cenários que exigem recall contextual e em configurações 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.