Gegatterte Slot-Aufmerksamkeit für effizientes sequenzielles Modellieren in linearer Zeit

Gated Slot Attention for Efficient Linear-Time Sequence Modeling

September 11, 2024

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

Zusammenfassung

Lineare Aufmerksamkeits-Transformer und ihre gesteuerten Varianten, die dafür bekannt sind, paralleles Training und effiziente wiederkehrende Inferenz zu ermöglichen, bleiben im Vergleich zu traditionellen Transformern bei auf Rückruf basierenden Aufgaben hinter diesen zurück und erfordern erhebliche Ressourcen für das Training von Grund auf. Dieses Papier stellt das Gesteuerte Slot-Aufmerksamkeit (GSA) vor, das die Aufmerksamkeit mit begrenzter Speichersteuerung (ABC) verbessert, indem es einen Steuermechanismus integriert, der von der Gesteuerten Linearen Aufmerksamkeit (GLA) inspiriert ist. Im Wesentlichen besteht GSA aus einem zweischichtigen GLA, das über Softmax verbunden ist, wobei kontextbewusstes Speicherlesen und adaptives Vergessen genutzt werden, um die Speicherkapazität zu verbessern und gleichzeitig die kompakte wiederkehrende Zustandsgröße beizubehalten. Dieses Design verbessert sowohl die Effizienz des Trainings als auch der Inferenz erheblich durch den hardwareeffizienten Trainingsalgorithmus von GLA und die reduzierte Zustandsgröße. Darüber hinaus ist die Beibehaltung der Softmax-Operation besonders vorteilhaft in "Feinabstimmung vortrainierter Transformer auf RNNs" (T2R)-Einstellungen, was den Bedarf an umfangreichem Training von Grund auf reduziert. Umfangreiche Experimente bestätigen die überlegene Leistung von GSA in Szenarien, die In-Context-Abruf erfordern, sowie in T2R-Einstellungen.

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.