TCNCA: Rede de Convolução Temporal com Atenção Segmentada para Processamento Escalável de Sequências

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

Resumo

O MEGA é uma arquitetura recente baseada em transformadores, que utiliza um operador recorrente linear cuja computação paralela, baseada na FFT, escala como O(LlogL), com L sendo o comprimento da sequência. Nós construímos sobre essa abordagem substituindo a recorrência linear por uma rede convolucional temporal especial que permite um tamanho de campo receptivo maior com redes mais rasas e reduz a complexidade computacional para O(L). O modelo resultante é chamado de TCNCA, uma Rede Convolucional Temporal com Atenção Segmentada. Avaliamos o TCNCA na modelagem de linguagem EnWik8, na classificação de sequências do long-range-arena (LRA), bem como em um benchmark sintético de raciocínio de recall associativo. No EnWik8, o TCNCA supera o MEGA, alcançando uma perda menor com passos forward/backward 1,37 vezes/1,24 vezes mais rápidos durante o treinamento. As convoluções dilatadas usadas no TCNCA são operações consistentemente e significativamente mais rápidas do que a recorrência paralelizada baseada em FFT em GPUs, tornando-as uma candidata escalável para lidar com comprimentos de sequência muito grandes: elas são até 7,07 vezes/2,86 vezes mais rápidas no passo forward/backward para sequências de até 131k. Além disso, no LRA, o TCNCA alcança, em média, uma aceleração de 1,28 vezes durante a inferência com uma precisão semelhante à alcançada pelo MEGA. No recall associativo, descobrimos que mesmo uma versão simplificada do TCNCA, sem interações multiplicativas e aditivas excessivas, permanece superior ou competitiva em relação ao MEGA em uma variedade de comprimentos de sequência e tamanhos de vocabulário.

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.