TCNCA : Réseau de Convolution Temporelle avec Attention Segmentée pour le Traitement Évolutif de Séquences

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

Résumé

MEGA est une architecture récente basée sur des transformateurs, qui utilise un opérateur récurrent linéaire dont le calcul parallèle, basé sur la FFT, s'échelonne en O(LlogL), où L représente la longueur de la séquence. Nous nous appuyons sur leur approche en remplaçant la récurrence linéaire par un réseau convolutif temporel spécial qui permet une taille de champ réceptif plus grande avec des réseaux moins profonds, et réduit la complexité computationnelle à O(L). Le modèle résultant est appelé TCNCA, un Réseau Convolutif Temporel avec Attention Segmentée. Nous évaluons TCNCA sur la modélisation de langage EnWik8, la classification de séquences long-range-arena (LRA), ainsi qu'un benchmark de raisonnement synthétique de rappel associatif. Sur EnWik8, TCNCA surpasse MEGA, atteignant une perte plus faible avec un passage avant/arrière 1,37 fois/1,24 fois plus rapide pendant l'entraînement. Les convolutions dilatées utilisées dans TCNCA sont systématiquement et significativement plus rapides que la récurrence parallélisée basée sur la FFT sur les GPU, en faisant un candidat scalable pour traiter des séquences de très grande longueur : elles sont jusqu'à 7,07 fois/2,86 fois plus rapides dans le passage avant/arrière pour des séquences allant jusqu'à 131k. De plus, sur LRA, TCNCA obtient, en moyenne, une accélération de 1,28 fois pendant l'inférence avec une précision similaire à celle de MEGA. Sur le rappel associatif, nous constatons qu'une version simplifiée de TCNCA, sans interactions multiplicatives et additives excessives, reste supérieure ou compétitive par rapport à MEGA sur une gamme de longueurs de séquences et de tailles de vocabulaire.

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.