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.

错误:搜索内容不能为空,请输入英文关键词
错误:关键词超出字数限制,请精简
高级检索

TCNCA: Temporal CNN with Chunked Attention for Efficient Training on Long Sequences

  • Aleksandar Terzić,
  • Michael Hersche,
  • Geethan Karunaratne,
  • Abu Sebastian,
  • Abbas Rahimi

摘要

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.