Segmented Recurrent Transformer: An Efficient Sequence-to-Sequence Model (2023)

Abstract: Transformers have shown dominant performance across a range of domains including language and vision. However, their computational cost grows quadratically with the sequence length, making their usage prohibitive for resource-constrained applications. To counter this, our approach is to divide the whole sequence into segments and apply attention to the individual segments. We propose a segmented recurrent transformer (SRformer) that combines segmented (local) attention with recurrent attention. The loss caused by reducing the attention window length is compensated by aggregating information across segments with recurrent attention. SRformer leverages Recurrent Accumulate-and-Fire (RAF) neurons’ inherent memory to update the cumulative product of keys and values. The segmented attention and lightweight RAF neurons ensure the efficiency of the proposed transformer. Such an approach leads to models with sequential processing capability at a lower computation/memory cost. We apply the proposed method to T5 and BART transformers. The modified models are tested on summarization datasets including CNN-dailymail, XSUM, ArXiv, and MediaSUM. Notably, using segmented inputs of varied sizes, the proposed model achieves 6-22% higher ROUGE1 scores than a segmented transformer and outperforms other recurrent transformer approaches. Furthermore, compared to full attention, the proposed model reduces the computational complexity of cross attention by around 40%.

## introduction since the inception of transformers @xcite , they have gradually become the de-facto model for natural language processing (nlp) as well as computer vision (cv) tasks @xcite . more recently, generative models have also achieved unprecedented results by employing transformers in language and vision for understanding and generation @xcite . one of the key recipes behind the success of transformers is their ability to leverage long-range connections baked in the all-to-all attention mechanism. this sort of attention eases optimization and enables learning of long-term dependency, resulting in improved performance over traditional recurrent neural networks (rnns). however, transformers incur huge memory and computing requirements, which may not be feasible for resource-constrained applications. moreover, the computational complexity of transformers increases quadratically with input length, making scaling-up transformers a big challenge. therefore, improving the efficiency of transformers is a key research direction @xcite . specifically, for sequence-to-sequence language tasks such as translation and summarization, the decoder part of the transformer inevitably generates words in steps @xcite . these steps are mostly correlated and hence, some form of recurrence is suited to harness the underlying information efficiently. without any explicit memory, the vanilla attention processes the same set of previous inputs repeatedly. such redundant computation can be eliminated by incorporating recurrent units with small memory overhead. as a result, the model can efficiently recapitulate relevant information from past states while computation is performed only on the new inputs. to that end, there have been efforts to combine transformers with the recurrent mechanism. in transformerxl @xcite , the long-term dependency is modeled by adding recurrence to transformers. however, this method simply concatenates the hidden states from the last segment with the current state for attention. it does not involve any explicit modification to the attention block itself to better suit the recurrence in order to bridge consecutive segments. though this is simple to implement, it may not be the best possible way to merge the attention mechanism with recurrence. moreover, most existing recurrent transformers are only designed for self attention, whereas, for summarization tasks, cross attention uses much longer inputs and incurs a higher cost. hence, there is scope to design a custom transformer architecture that can effectively combine recurrence with the embedded attention. in order to enhance the computational efficiency of transformers, we divide the input sequence into segments to compute segmented attention. next, to aggregate global information over segments, we add recurrent attention. we propose recurrent accumulate and fire (raf) neurons for recurrent attention. it has an implicit memory to accumulate information over segments and only propagates values upon crossing a learnable threshold. by accumulating the partial product of keys and values from other segments, recurrent attention compensates for the loss caused by segmentation. furthermore, raf neurons use a leak and a threshold to forget and select information. compared to rnns or lstms, a raf neuron has fewer parameters and thus does not add large computational overhead. based on the proposal of segmented attention and recurrent attention, we introduce srformer, as shown in fig. 1 . srformer is a sequence-tosequence model consisting of an encoder and an auto-regressive decoder. it replaces the decoder's cross-attention with segmented recurrent attention, while self-attention can still attend to all locations of the generated sequence. during training, it selects the corresponding segment for each input and learns recurrent attention between segments. because recurrent attention processes a segment as a whole, it enables partially parallel computation. during inference, it accesses past information from the memory and only updates the segmented recurrent attention with the current segment. hence, it is more efficient than full attention. in addition, the number of timesteps for backpropagation depends on the number of segments instead of the length of inputs. the reduced propagation length ensures efficiency and alleviates the vanishing gradient problem. moreover, we provide a matrix interpretation to demonstrate the harmonious alliance between recurrence and attention obtained through the proposed model. ation score for the proposed model is 6 -22% higher than a segmented transformer without recurrent blocks. interestingly, the performance of srformer is resilient to extremely small segment sizes (rouge1 score drops by only 0.25 as segment size is reduced from 64 to 8). on the other hand, for the same segment size reduction, the baseline transformer suffers from severe performance degradation (rouge1 score drops by 16.04). furthermore, the computation reduces by ∼ 40% in comparison with regular cross attention having access to the full sequence. this result clearly demonstrates the efficacy of segmented recurrent attention for sequential data processing, specifically for limited compute budgets. ## related works several approaches have been reported in the literature to improve the efficiency of transformers @xcite . usage of low-rank kernels @xcite , fixed or learnable patterns @xcite @xcite , block-wise processing @xcite , recurrent transformers @xcite , and global attention @xcite @xcite have been explored recently. our work falls under the category of transformers with recurrence. in this aspect, we closely examine the following two models. transformerxl @xcite concatenates the hidden states from the previous segment with the current one. then each hidden state depends on all previous ones and becomes built-in memory. however, the length of hidden states gets doubled resulting in additional complexity. the linear transformer reduces attention to a linear-time, constant memory recurrent neural network @xcite . it is similar to our recurrent attention without raf neurons and setting segment size to one. however, linear transformer usually suffers from performance degradation compared to regular attention with softmax. @xcite identified that the devil in linear attention is the scaling of attention matrices and used normalization to stabilize gradients. in our approach, we use raf neurons and normalization to make a better approximation to softmax attention. other types of transformers deal with long sequences by combining local attention with global attention, such as etc @xcite , big bird @xcite , longformer @xcite , and longt5 @xcite . each paper proposed a different design for global attention. etc pretrains global tokens with contrastive predictive coding. big bird adds random sparse attention patterns to global attention. longformer applies global attention on pre-selected locations and makes it symmetric: all locations attend to global tokens and vice versa. longt5 simply uses the sum of each block as global tokens. our work substitutes recurrent attention for global attention. we will compare our results with longt5 since it is the most recent work. there are two concurrent works about recurrent transformers. the first one is block recurrent transformer @xcite . they design recurrent cells to operate on blocks of tokens. different from our encoder-decoder model, they add cross attention to a decoder-only model, which attends to recurrent states in parallel with self attention. the second work is on recurrent memory transformer @xcite . the authors segment inputs similar to our work, then put read and write memory at different ends of the input to attention. however, they only add memory to the input, while we propose a novel attention mechanism that is a drop-in replacement for regular attention. notably, most of these related works focus on the self attention block of the transformer. however, we notice that the cross attention between the encoder and decoder is the bottleneck of computation in summarization tasks. hence, we focus on improving the efficiency of cross attention for sequence-to-sequence models. additionally, we use small segment sizes to improve efficiency, while other works use large sizes to improve scalability, which is another distinction between the proposed work and prior art. ## experimental results we test our proposed model and method on t5 sequence-to-sequence transformers @xcite and bart model @xcite . since we want to design a model for resource-constrained environments, we use t5-small, which contains 6 transformer layers and 60 million parameters, and bart-base which contains 6 layers and 140 million parameters. our implementation is built based on huggingface transformer platform, which is released under the apache-2.0 license, and we load the pretrained models before training @xcite . the implementation is based on pytorch python package @xcite . we run all experiments on nvidia a40 gpus. each experiment is run on a single core to compare gpu power and memory usage. we train the models for 25 epochs as the loss and evaluation scores converge at that point. the initial learning rate is 5e -5 . the optimizer is adamw and the learning rate scheduler is linear with default settings. we use rouge scores to evaluate srformer and compare it with other transformers. rouge-n is defined as the n-gram recall between model predictions and a set of ground-truth summaries, and rouge-l measures the longest common subsequence between them @xcite . note, a higher rouge score denotes better summarization performance. ## conclusion in this paper, we propose an efficient sequence-tosequence transformer using segmented recurrent attention. since cross-attention is the computation bottleneck for text summarization, replacing full attention with segmented recurrent attention can largely reduce its computation complexity. the proposed raf neurons help approximate full attention at a low overhead. during inference, they utilize the sequential process to accumulate information in their internal memory and bridge information across segments effectively. the efficacy of the proposed model is demonstrated through higher rouge scores on multiple summarization datasets compared to other segmented transformers. our proposed architecture is especially suitable for summarization and can be used for other sequence-to-sequence tasks such as translation and question-answering.