Gegrendelde Sleuf Aandacht voor Efficiënte Lineaire-Tijd Sequentiemodellering

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

Samenvatting

Lineaire aandachttransformatoren en hun afgeleiden met poorten, die bekend staan om het mogelijk maken van parallelle training en efficiënte terugkerende inferentie, schieten nog steeds tekort bij taken die veel oproepen vereisen in vergelijking met traditionele transformatoren en vereisen aanzienlijke middelen voor training vanaf nul. Dit artikel introduceert Gated Slot Aandacht (GSA), dat Aandacht verbetert met Begrensde-geheugencontrole (ABC) door een poortmechanisme te incorporeren geïnspireerd door Gated Lineaire Aandacht (GLA). In essentie bestaat GSA uit een tweelaagse GLA verbonden via softmax, waarbij contextbewust geheugen lezen en adaptief vergeten worden gebruikt om de geheugencapaciteit te verbeteren terwijl de compacte terugkerende toestandsgrootte behouden blijft. Deze ontwerp verbetert aanzienlijk zowel de training als de inferentie-efficiëntie door het hardware-efficiënte trainingsalgoritme van GLA en de verkleinde toestandsgrootte. Bovendien is het behouden van de softmax-operatie met name gunstig in "fijnafstemming van vooraf getrainde transformatoren naar RNNs" (T2R) instellingen, waardoor de noodzaak voor uitgebreide training vanaf nul wordt verminderd. Uitgebreide experimenten bevestigen de superieure prestaties van GSA in scenario's die in-context oproep vereisen en in T2R-instellingen.

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.