Attention à fentes contrôlées pour une modélisation efficace des séquences en temps linéaire

Gated Slot Attention for Efficient Linear-Time Sequence Modeling

September 11, 2024

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

Résumé

Les Transformers d'attention linéaire et leurs variantes à portes, célèbres pour permettre un entraînement parallèle et une inférence récurrente efficace, restent en deçà dans les tâches intensives en rappel par rapport aux Transformers traditionnels et exigent des ressources significatives pour être entraînés à partir de zéro. Ce document présente l'Attention à Fentes à Portes (GSA), qui améliore l'Attention avec Contrôle de Mémoire Bornée (ABC) en incorporant un mécanisme de portes inspiré par l'Attention Linéaire à Portes (GLA). Essentiellement, le GSA comprend un GLA à deux couches lié via softmax, utilisant une lecture de mémoire consciente du contexte et un oubli adaptatif pour améliorer la capacité de mémoire tout en maintenant une taille d'état récurrent compacte. Cette conception améliore considérablement à la fois l'efficacité de l'entraînement et de l'inférence grâce à l'algorithme d'entraînement efficace en matériel du GLA et à la réduction de la taille de l'état. De plus, le maintien de l'opération softmax est particulièrement bénéfique dans les paramètres de "réglage fin des Transformers pré-entraînés aux RNN" (T2R), réduisant le besoin d'un entraînement approfondi à partir de zéro. Des expériences approfondies confirment la performance supérieure du GSA dans les scénarios nécessitant un rappel contextuel et dans les paramètres 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.