TCNCA: Temporales Faltungsnetzwerk mit segmentierter Aufmerksamkeit für skalierbare Sequenzverarbeitung

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

December 9, 2023

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

cs.AI

Zusammenfassung

MEGA ist eine kürzlich entwickelte Transformer-basierte Architektur, die einen linearen rekurrenten Operator verwendet, dessen parallele Berechnung auf der FFT basiert und sich mit O(LlogL) skaliert, wobei L die Sequenzlänge ist. Wir bauen auf diesem Ansatz auf, indem wir die lineare Rekurrenz durch ein spezielles temporales Faltungsnetzwerk ersetzen, das größere rezeptive Feldgrößen mit flacheren Netzwerken ermöglicht und die Rechenkomplexität auf O(L) reduziert. Das resultierende Modell wird TCNCA genannt, ein Temporales Faltungsnetzwerk mit Chunked Attention. Wir evaluieren TCNCA auf EnWik8-Sprachmodellierung, Long-Range-Arena (LRA)-Sequenzklassifizierung sowie einem synthetischen Reasoning-Benchmark für assoziatives Abrufen. Auf EnWik8 übertrifft TCNCA MEGA, erreicht einen niedrigeren Verlust und ist 1,37-mal/1,24-mal schneller im Vorwärts-/Rückwärtsdurchlauf während des Trainings. Die in TCNCA verwendeten dilatierten Faltungen sind durchweg und deutlich schnellere Operationen als die FFT-basierte parallelisierte Rekurrenz auf GPUs, was sie zu einem skalierbaren Kandidaten für die Verarbeitung sehr großer Sequenzlängen macht: Sie sind bis zu 7,07-mal/2,86-mal schneller im Vorwärts-/Rückwärtsdurchlauf für Sequenzen bis zu 131k. Weiterhin erreicht TCNCA auf LRA im Durchschnitt eine 1,28-mal schnellere Inferenz bei ähnlicher Genauigkeit wie MEGA. Beim assoziativen Abrufen stellen wir fest, dass sogar eine vereinfachte Version von TCNCA, ohne übermäßige multiplikative und additive Interaktionen, MEGA bei einer Reihe von Sequenzlängen und Vokabulargrößen übertrifft oder konkurrenzfähig bleibt.

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.