TCNCA: Temporal CNN with Chunked Attention for Efficient Training on Long Sequences
摘要
MEGA is a recent Transformer-based neural network utilizing a linear recurrent architecture whose computation can be parallelized, a useful property during training as well as sequence encoding tasks. The parallel computation is based on the fast Fourier transform and scales as O(LlogL), with L being the sequence length. We replace the linear recurrence in MEGA with a temporal convolutional network (TCN) which permits a large receptive field size with few TCN layers, and reduces the computational complexity to O(L). We call the resulting model TCNCA, a Temporal Convolutional Network with Chunked Attention. We evaluate TCNCA on the tasks of EnWik8 language modelling, associative recall, a synthetic reasoning benchmark, and long-range-arena (LRA) sequence classification, and observe consistent improvements compared to MEGA in terms of both task-specific metrics and runtimes. More specifically, on EnWik8, TCNCA outperforms MEGA-chunk by a 0.01 BPC loss with a \(1.22\times\) / \(1.28\times\) faster forward/backward pass. On LRA, TCNCA outperforms MEGA-chunk by \(0.8\%\) on average with a \(1.42\times\) / \(1.16\times\) forward/backward pass speed-up. We further demonstrate the efficacy of our approach by comparing the runtimes of our approach and MEGA over a wide range of sequence lengths and embedding dimensions.