TCNCA: Red de Convolución Temporal con Atención Segmentada para el Procesamiento Escalable de Secuencias

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

December 9, 2023

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

cs.AI

Resumen

MEGA es una arquitectura reciente basada en transformadores que utiliza un operador recurrente lineal cuya computación paralela, basada en la FFT, escala como O(LlogL), donde L es la longitud de la secuencia. Nos basamos en su enfoque al reemplazar la recurrencia lineal con una red convolucional temporal especial que permite un mayor tamaño de campo receptivo con redes más superficiales y reduce la complejidad computacional a O(L). El modelo resultante se llama TCNCA, una Red Convolucional Temporal con Atención Segmentada. Evaluamos TCNCA en modelado de lenguaje EnWik8, clasificación de secuencias en el long-range-arena (LRA), así como en un benchmark sintético de razonamiento de recuerdo asociativo. En EnWik8, TCNCA supera a MEGA, alcanzando una pérdida menor con un paso hacia adelante/hacia atrás 1.37 veces/1.24 veces más rápido durante el entrenamiento. Las convoluciones dilatadas utilizadas en TCNCA son operaciones consistentemente y significativamente más rápidas que la recurrencia paralelizada basada en FFT en GPUs, lo que las convierte en un candidato escalable para manejar secuencias muy largas: son hasta 7.07 veces/2.86 veces más rápidas en el paso hacia adelante/hacia atrás para secuencias de hasta 131k. Además, en LRA, TCNCA logra, en promedio, una aceleración de 1.28 veces durante la inferencia con una precisión similar a la de MEGA. En el recuerdo asociativo, encontramos que incluso una versión simplificada de TCNCA, sin interacciones multiplicativas y aditivas excesivas, sigue siendo superior o competitiva frente a MEGA en un rango de longitudes de secuencia y tamaños de vocabulario.

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.