TCNCA: Rete Convoluzionale Temporale con Attenzione a Blocchi per l'Elaborazione Scalabile di Sequenze

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

December 9, 2023

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

cs.AI

Abstract

MEGA è una recente architettura basata su transformer, che utilizza un operatore ricorrente lineare la cui computazione parallela, basata sulla FFT, scala come O(LlogL), dove L è la lunghezza della sequenza. Costruiamo sul loro approccio sostituendo la ricorrenza lineare con una speciale rete convoluzionale temporale che consente una dimensione del campo ricettivo maggiore con reti più superficiali e riduce la complessità computazionale a O(L). Il modello risultante è chiamato TCNCA, una Rete Convoluzionale Temporale con Attenzione a Blocchi. Valutiamo TCNCA sulla modellazione del linguaggio EnWik8, sulla classificazione di sequenze long-range-arena (LRA), nonché su un benchmark sintetico di ragionamento associative recall. Su EnWik8, TCNCA supera MEGA, raggiungendo una perdita inferiore con un passaggio in avanti/indietro 1.37 volte/1.24 volte più veloce durante l'addestramento. Le convoluzioni dilatate utilizzate in TCNCA sono operazioni costantemente e significativamente più veloci rispetto alla ricorrenza parallelizzata basata su FFT nelle GPU, rendendole un candidato scalabile per gestire sequenze molto lunghe: sono fino a 7.07 volte/2.86 volte più veloci nel passaggio in avanti/indietro per sequenze fino a 131k. Inoltre, su LRA, TCNCA ottiene, in media, un'accelerazione di 1.28 volte durante l'inferenza con un'accuratezza simile a quella raggiunta da MEGA. Su associative recall, troviamo che anche una versione semplificata di TCNCA, senza eccessive interazioni moltiplicative e additive, rimane superiore o competitiva rispetto a MEGA su un intervallo di lunghezze di sequenza e dimensioni del vocabolario.

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.