TCNCA: Temporeel Convolutienetwerk met Gegroepeerde Aandacht voor Schaalbare Sequentieverwerking

TCNCA: Temporal Convolution Network with Chunked Attention for Scalable Sequence Processing

December 9, 2023

Auteurs: Aleksandar Terzic, Michael Hersche, Geethan Karunaratne, Luca Benini, Abu Sebastian, Abbas Rahimi

cs.AI

Samenvatting

MEGA is een recente transformer-gebaseerde architectuur die gebruikmaakt van een lineaire recurrent operator waarvan de parallelle berekening, gebaseerd op de FFT, schaalt als O(LlogL), waarbij L de sequentielengte is. Wij bouwen voort op hun aanpak door de lineaire recurrent te vervangen door een speciaal tijdconvolutienetwerk dat een groter receptief veld mogelijk maakt met ondiepere netwerken en de rekencomplexiteit reduceert tot O(L). Het resulterende model wordt TCNCA genoemd, een Tijdconvolutienetwerk met Chunked Attention. We evalueren TCNCA op EnWik8-taalmodelering, long-range-arena (LRA) sequentieclassificatie, evenals een synthetisch redeneerbenchmark voor associatief ophalen. Op EnWik8 presteert TCNCA beter dan MEGA, met een lager verlies en 1,37x/1,24x snellere voorwaartse/achterwaartse passes tijdens de training. De gedilateerde convoluties die in TCNCA worden gebruikt, zijn consistent en aanzienlijk snellere operaties dan de FFT-gebaseerde parallelle recurrent in GPU's, waardoor ze een schaalbare kandidaat zijn voor het verwerken van zeer grote sequentielengtes: ze zijn tot 7,07x/2,86x sneller in de voorwaartse/achterwaartse pass voor sequenties tot 131k. Verder behaalt TCNCA op LRA gemiddeld een 1,28x versnelling tijdens inferentie met een vergelijkbare nauwkeurigheid als MEGA. Op associatief ophalen vinden we dat zelfs een vereenvoudigde versie van TCNCA, zonder overmatige multiplicatieve en additieve interacties, superieur of competitief blijft ten opzichte van MEGA over een reeks sequentielengtes en vocabulaire groottes.

English

MEGA is a recent transformer-based architecture, which utilizes a linear recurrent operator whose parallel computation, based on the FFT, scales as O(LlogL), with L being the sequence length. We build upon their approach by replacing the linear recurrence with a special temporal convolutional network which permits larger receptive field size with shallower networks, and reduces the computational complexity to O(L). The resulting model is called TCNCA, a Temporal Convolutional Network with Chunked Attention. We evaluate TCNCA on EnWik8 language modeling, long-range-arena (LRA) sequence classification, as well as a synthetic reasoning benchmark associative recall. On EnWik8, TCNCA outperforms MEGA, reaching a lower loss with 1.37times/1.24times faster forward/backward pass during training. The dilated convolutions used in TCNCA are consistently and significantly faster operations than the FFT-based parallelized recurrence in GPUs, making them a scalable candidate for handling very large sequence lengths: they are up to 7.07times/2.86times faster in the forward/backward pass for sequences up to 131k. Further on LRA, TCNCA achieves, on average, 1.28times speed-up during inference with similar accuracy to what MEGA achieves. On associative recall, we find that even a simplified version of TCNCA, without excessive multiplicative and additive interactions, remains superior or competitive to MEGA on a range of sequence lengths and vocabulary sizes.