TCNCA: Временная сверточная сеть с фрагментированным вниманием для масштабируемой обработки последовательностей

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

December 9, 2023

Авторы: Aleksandar Terzic, Michael Hersche, Geethan Karunaratne, Luca Benini, Abu Sebastian, Abbas Rahimi

cs.AI

Аннотация

MEGA — это недавно разработанная архитектура на основе трансформеров, которая использует линейный рекуррентный оператор, чье параллельное вычисление, основанное на быстром преобразовании Фурье (FFT), масштабируется как O(LlogL), где L — длина последовательности. Мы развиваем их подход, заменяя линейную рекуррентность специальной временной сверточной сетью, которая позволяет увеличить размер рецептивного поля при меньшей глубине сети и снижает вычислительную сложность до O(L). Полученная модель называется TCNCA (Temporal Convolutional Network with Chunked Attention). Мы оцениваем TCNCA на задаче языкового моделирования EnWik8, классификации последовательностей в рамках long-range-arena (LRA), а также на синтетическом бенчмарке ассоциативного воспроизведения. На EnWik8 TCNCA превосходит MEGA, достигая меньших потерь при ускорении прямого/обратного прохода в 1.37/1.24 раза во время обучения. Дилатированные свертки, используемые в TCNCA, стабильно и значительно быстрее операций, основанных на FFT, на GPU, что делает их масштабируемым решением для обработки очень длинных последовательностей: они до 7.07/2.86 раза быстрее в прямом/обратном проходе для последовательностей длиной до 131 тыс. элементов. На LRA TCNCA достигает в среднем ускорения в 1.28 раза при выводе с аналогичной точностью по сравнению с MEGA. На ассоциативном воспроизведении мы обнаружили, что даже упрощенная версия TCNCA, без избыточных мультипликативных и аддитивных взаимодействий, остается превосходящей или конкурентоспособной по сравнению с MEGA для различных длин последовательностей и размеров словаря.

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.