Segmented Recurrent Transformer: An Efficient Sequence-to-Sequence Model
Authors: Yinghan Long
Sayeed Shafayet Chowdhury
Kaushik Roy
Venue: EMNLP 2023 Findings (2023)
Keywords: abstractive summarization
transformers
language models

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%.

Segmented Recurrent Transformer: An Efficient Sequence-to-Seqeenc
Model
Yinghan Long Purdue University 3ong273@purdue.coo edu
Sayeed Shafayet Chowdhury
Kaushik Roy Purdue University
Purdue University chowdh23@purduee edu
kaushikepurdue.cek edu
Abstract
vision for understanding and generation (Saharia et al., 2022; Rombach al., 2022). One of the key recipes behind the success of transformers is their ability to leverage long-range connections baked the all-to-all attention mechanism. This sort of at- tention eases optimization and enables learning of long-term dependency, resulting in improved per- formance traditional recurrent neural networks (RNNs). However, transformers incur huge mem- ory and computing requirements, which may not be feasible for resource-constrainee, applications. Morcover, the computational complexity of trans- formers increases quadratically with input length, making scaling-up up transformers a big challenge. Therefore, improving the efficiency of transform- is key research direction (Tay et al., 2022). Specifically, for sequence-to-sequenee language tasks such as translation and summarization, the decoder part of the transformer inevitably generates words in steps (Guo et al., 2022; Bao et al., 2020). These steps are mostly correlated and hence, some form of recurrence is suited to harness the under- lying 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 current units with small memory overhead. As re- sult, the model can efficiently recapitulate relevant information from past states while computation is performed only on the new inputs.
Transformers have shown dominant perfor- mance across range of domains including language and vision. However, their compu- lational cost grows quadratically with the se- quence length, making their usage prohibitive for ussorcc-conssrringe
applications. To counter this, our approach to divide whole sequence into segments and apply attention to the individual segments. We propose a seg- mented recurrent transformer (SRformer) that combines segmented (local) attention with re- current attention. The loss caused by reducing the attention window length is compensated by aggregating information across segments with recurrent attention. SRformer leverages Re- current Accumulate-ane and-Fire (RAF) neurons' inherent memory to update the cumulative prod- uct of keys and values. The segmented atten- tion and lightweight RAF neurons ensure the efficiency of the proposed transformer. Such an approach leads models with sequential pro- cessing capability at lower computation/mem- ory cost. apply proposed method to and BART transformers. The modified models are tested on summarization datasets including CNN-dailymail, XSUM, ArXiv, and Media- SUM. Notably, using segmented inputs of var- ied sizes, the proposed model achieves 6 22% higher ROUGEI scores than segmented trans- former and outperforms other recurrent trans- former approaches. Furthermore, compared to full attention, the proposed model reduces the computational complexity of cross attention by around 40%.
To that end, there have been efforts to combine transformers with the recurrent mechanism. TransformerXL (Dai et al., 2019), the long-term dependency is modeled by adding recurrence to transformers. However, this method simply con- catenates the hidden states from the last segment with the current state for attention. It does not involve any explicit modification the attention block itself better suit the recurrence in order bridge consecutive segments. Though this simple to implement, it may not be the best possible way
1 Introduction
Since the inception of transformers (Vaswani et al., 2017), they have gradually become the de-facto model for Natural Language Processing (NLP) as well as Computer Vision (CV) tasks (Brown al., 2020; Dosovitskiy et al., 2021). More recently, gen- crative models have also achieved unprecedented results by employing transformers in language and

---

merge the attention mechanism with recurrence. Moreover, most existing recurrent transformers are only designed for self attention, whereas, for sum- marization tasks, cross attention uses much longer inputs and incurs a higher cost. Hence, there is scope to design custom transformer architecture that can effectively combine recurrence with the embedded attention.
Generated
Decoder layer
Feed Forward Network
Segmented Recurrent
Cross Attention
In order to enhance the computational efficiency 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 accu- mulate and fire (RAF) neurons for recurrent atten- tion. It has an implicit memory to accumulate infor- mation over segments and only propagates values upon crossing a learnable threshold. By accumu- lating the partial product of keys and values from other segments, recurrent attention compensates for the loss caused by segmentation. Furthermore, RAF neurons use leak and a threshold to for- get and select information. Compared to RNNs or LSTMs, RAF neuron has fewer parameters and thus does not add large computational overhead.
K
Q
Attention
Transformer
K
Q
Original text
Text
Figure 1: Segmented Recurrent Transformer (SR- former). Cross attention is replaced by segmented recur- rent attention.
ation score for the proposed model is 6 22% higher than a segmented transformer without current blocks. Interestingly, the performance of SRformer is resilient to extremely small segment sizes (ROUGEI score drops by only 0.25 as seg- ment size reduced from 64 to 8). On the other hand, for the same segment size reduction, the base- line transformer suffers from severe performance degradation (ROUGEI score drops by 16.04). Fur- thermore, the computation reduces by 40% in comparison with regular cross attention having ac- cess to the full sequence. This result clearly demon- strates the efficacy of segmented recurrent attention for sequential data processing, specifically for lim- ited compute budgets.
Based on the proposal of segmented attention and recurrent attention, we introduce SRformer, as shown 1. SRformer a sequence-to- sequence model consisting of an encoder and an auto-regressive decoder. It replaces the decoder's 's cross-attention with segmented recurrent attention, while self-attention can still attend to all locations of the generated sequence. During training, it se- lects the corresponding segment for each input and learns recurrent attention between segments. Be- cause 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 recur- rent attention with the current segment. Hence, it is more efficient than full attention. In addition, the number of timesteps for backpropagation de- pends on the number of segments instead of the length of inputs. The reduced propagation length ensures efficiency and alleviates the vanishing gra- dient problem. Moreover, we provide matrix interpretation to demonstrate the harmonious al- liance between recurrence and attention obtained through the proposed model.
In the next section, we introduce the related works. In Section 3, provide details of SR- former. Finally, the experimental results dis- cussed Section
2 Related Works
Several approaches have been reported in the litera- ture to improve the efficiency of transformers (Tay et al., 2022). Usage of low-rank kernels (Wang et al., 2020), fixed or learnable patterns (Child al., 2019; Ho et al., 2019; Roy et al., 2021; Kitaev et al., 2020), block-wise processing (Lefaudeux et al., 2022; Dao al., 2022), recurrent transform- ers (Dai et al., 2019; Katharopoulos et al., 2020), and global attention (Ainslie et al., 2020; Zaheer
SRformer achieves better computational com- plexity and accuracy trade-off. When tested on the -NN-dailymail dataset, the ROUGE! evalu-

---

al., 2020; Beltagy et al., 2020; Guo et al., 2021) have been explored recently.
propose novel attention mechanism that is drop-in replacement for regular attention.
Our work falls under the category of transform- ers with recurrence. In this aspect, we closely ex- amine the following two models. TransformerXL. (Dai et al., 2019) concatenates the hidden states from the previous segment with the current one. Then each hidden state depends 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 linear-time, constant mem- ory recurrent neural network (Katharopoulos et al., 2020). It is similar to our recurrent attention with- out RAF neurons and setting segment size to one. However, linear transformer usually suffers from performance degradation compared regular at- tention with softmax. Qin al. (2022) identified that the devil in linear attention is the scaling of at- tention matrices and used normalization to stabilize gradients. In our approach, we use RAF neurons and normalization to make a better approximation to softmax attention.
Notably, most of these related works focus on the self attention block of the transformer. How- ever, we notice that the cross attention between the encoder and decoder is the bottleneck of com- putation in summarization tasks. Hence, we fo- cus improving the efficiency of cross attention for requence-to-sequeee
models. Additionally, small segment sizes improve efficiency, while other works use large sizes to improve scal- ability, which another distinction between the proposed work and prior art.
3 Proposed method
3.1 Segmented Attention
Suppose that Q is the query matrix, K is the key matrix, and V is the value matrix. The cross at- tention weight A and the output attention O are computed by
Other types of transformers deal with long se- quences by combining local attention with global attention, such as ETC (Ainslie et al., 2020), Big Bird (Zaheer et al., 2020), Longformer (Beltagy et al., 2020), and longT5 (Guo et al., 2021). Each paper proposed a different design for global atten- tion. ETC pretrains global tokens with Contrastive Predictive Coding. Big Bird adds random sparse attention patterns 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.
qx A axxd Q Qxaxk KT
(1)
oxd = softmax(A) skd
(2)
We show the dimension of each matrix at its upper- left superscript. 1 is the length of encoded features and is the length of decoder embeddings. d is the inner dimension of each attention head. For summarization tasks, q is usually much smaller than k. Note that the batch dimension and the number of heads are ignored for simplicity.
If we split encoded features into segments of size s, then the number of segments would be m k We define the i-th segment of keys and values s] and Vs, Vli * (i + 1) * s). To find the corresponding segment for a query at time t, we compute the segment index by = min(int(t * m//), 11mmm11). -1).
Ks: K[i
There are two concurrent works about recurrent transformers. The first one is block recurrent trans- former (Hutchins et al., 2022). They design recur- rent cells to operate on blocks of tokens. Different from our encoder-decodee model, they add cross attention to decoder-only model, which attends to recurrent states in parallel with self attention. The second work is on recurrent memory transformer (Bulatov et al., 2022). The authors segment in- puts 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
Let the query at time be denoted as Q.- Then the full attention weight A[t] and output O[t] at time are computed using the current query Q. with full keys and values. However, in the case of seg- mented cross attention, we use only the correspond- ing segment of keys and values compute the at- tention for the query at time f, instead of the full matrix. The segmented attention weight (As(t) and output (Os(t)) are computed as in Eqn. and

---

Eqn. 6.
Oul softmax (Q+ K Vic..
O[1 = Cs 0s|1 + c Cs Onl c
1xk 1[d 1xd Oj * dxck KT
(3)
(9)
1x8 As [t] _1xd Qi .xx KT
(4)
O[t] = -0s + +1-
(10)
[xd softmax(A[1) * kxd V
(5)
where and es are divisors of softmax. They can be replaced by weighing factor a. Please see the appendix for detailed analyses.
1s = (A)(t)) exd Vss
Unlike directly projecting K to a lower rank of size d x 8 for all timesteps, each query is multiplied with different segment Ks; and Vs.. Hence every segment used at least once during the process and there are no extra weights or information loss caused by projection. The computation complexity is reduced to (qqds) from O(gdk), where << k. If s is equal to the summarization ratio k/q, then the computation complexity is reduced to (kd) from d/ss and no longer grows quadratically with the length. Segmented attention described above largely improves efficiency but at the loss of the ability to access all information, which causes accuracy degradation. Therefore, next, we focus on how to compensate for the loss and achieve a better trade-off between efficiency and accuracy using recurrent units between segments.
In the proposed approach, approximate the output time by computing the segmented at- tention Os of the local segment using Eqn. 6, then adding the approximated Or for other segments. The weighing factor a can be a trainable parameter. Note, the simplified sum shown in Eqn. 11 still works because of the automatic adjustment of other parameters.
O = 0s 0
(11)
To efficiently approximate the compensation part OR, we propose recurrent atlention. This is impor- tant because if we directly compute it, the total computational complexity becomes the same non-segmented model. Inspired by linear atten- tion proposed in Katharopoulos et al. (2020), we multiply keys and values first, then multiply with queries. This validated by the associativity prop- erty of matrix multiplication. The switch of order allows storing the product of keys and values as a state of RNN. The softmax operation applied on the attention weights is approximated as-
3.2 Approximating full attention with segmented recurrent attention
The segmented attention described in the last sub- section simply replaces K and V with the current segments Ks; and Vss, but it neglects a large part of the keys and values. Therefore, we need to com- pensate for the loss to better approximate the full attention.
ORl = softmax (2+ * K%&)
(12)
Qt * RAF K R& * VRR)
(13)
norm (K)
Suppose we divide the inputs into two parts: the current segment St and the remaining segments R, By linearity of matrix operations, if we switch the order of columns in K and that of corresponding rows in V, it does not affect the matrix multiplica- tion result. Hence, we can partition matrix K into [Ks: KRi] and V into [Vs; VRil where K Ra and VR: are the remaining parts of the matrices ex- cluding the current segment. The output O of full attention is computed by
The RAF in this equation is a recurrent unit de- signed to accumulate sequential information. It is explained in detail in subsection 3.3. In Eqn. 13, division by norm K) is used to simulate the normalization factor of softmax.
We simplify the computation of multiplying K and Va by the multiplication rule of the matrix, follows. The product is denoted by P.
Vs: VIR
K V = [Ks
Km:]
(14)
K + V Kst Vs: + Ri VRi
(15)
= Q x KT
= =Q= ++ [Kst KRd]
(7)
P. K VR K V Ks: * Vss
(16)
Vsi
0 = softmax(Q [Kss KRi))
(8)
VRi
Thus, can easily reduce the computational com- plexity of Eqn. 16 from Old (k s) * d) (d s * d). For P; of different segments, we only need to compute K + V once. Fig. shows the computation process of recurrent cross attention.
By expanding Eqn. 8, full attention at time can be rewritten as the weighted sum of segmented attention [t and the attention of remaining parts

---

output
ReLU
mem;
RAF
mem4+1
threshold
Leak
soft-reset
mem++1
A
Softmax
Linear
bs:
A
input
Recurrent Accumulation-anddFFin neuron. The leaky factor and threshold are trained parameters.
Figure 3:
Project
Ks Vsi
P.
soft reset by subtracting the value of the threshold. Then will not fire again until new information is given. Because KR and VRi different steps have overlaps, reset important to forget redun- dant information. The design RAF results in a sparse and low-complexity bio-inspired recurrence mechanism. Unlike an LSTM, a RAF neuron does not apply weights on the hidden memory, but only a Iearnable Icaky factor, threshold, and soft reset as forget and output gates. Hence, requires fewer weights and lesser computation compared an LSTM.
Figure Segmented Recurrent Attention in cross attention block. Ks; and Ve, are segmented keys and values. H, the hidden input, which is projected to Q: by linear layer. by: is the bias T5 model.
3.3 Recurrent Accumulation-and-FP neuron
RAF neurons accumulate the partial products of keys and values, help approximate softmax, and filter out unnecessary information. It is inspired by Integrate-and-F neurons from biological brains and Spiking Neural Networks (Lee et al., 2016;
Wozniak et al., 2020; Ponghiran and Roy, 2021; Chowdhury et al., 2022). The architecture of an RAF neuron is shown in Fig. 3. We also provide Pytorch-like pseudo code for RAF in Appendix A.2.
Algorithm One step of cross attention blocks in SRformer
Input: Query Q Keys K and Values V projected from encoded features; Timestep 4; previous index prev;
Each RAF neuron uses a membrane potential (memory) to efficiently aggregate sequential infor- mation. At each timestep, the memory is updated by adding the projected input with the previous memory, where memory gets decayed by learn- able leak. The input to RAF neurons is the partial product K VR: containing information of the se- quence except the current segment. It is projected by a linear layer. Since Eqn. 13 does not use soft- max, the linear layer is necessary to adjust the scale of inputs. Next, we check if its memory crosses the threshold of information propagation to the next layer. If yes, the information in its memory gets passed the next processing layer, else, the out- put is zero. This activated by applying ReLU on the thresholded memory. The learned thresh- old works like a bias to activate only values over it. If the memory has activated firing, it will be
Output: Output O,
K₃ = split(K)); Vú split(V)
min(inttt * m /g) m 1)
A(| = matmul [i)
A[t] softmax(A[())
if 0 then
KV = matmul(K,1
memory = zero tensor
end
Oş = matmul (A(t), Vi(i)) if prev! i then
p= KV [Kt
P. mem RAF( P, mem)
end
On matmul(Q:, (ynommm ())
O, Oş + Or prev =

---

Comp. complexity
Example (K)
Memory complexity
Example (K)
Transformer ((gkd)
gkd 8389
O kd)
197
Segmented
((gsd)
qsd 524
O + kd)
74
SRformer
Olgsd + kd)) qsd + kdf 4718
O(gs + kd + md²)
139
Table 1: Computation and memory complexity of cross attention. For example, in one of our experiments on the CNN-dailymail dataset, q 128; A 1024; d 64; s 64. To illustrate theoretical computation and memory cost, we show the product of example sizes in thousands.
3.4 Segmented Recurrent Transformer
in replacement for attention, it can be used in all encoder-decoder transformers, such T5 and BART. Please note that the usage of positional bias and masks varies across transformers and needs modification. We provide details in Appendix. A.3.
SRformer aims at utilizing the sequential genera- tion process to capture the relation between loca- tions of a long sequence, without paying attention to all locations or propagating from start to end. After reducing the length of keys and values to re- duce computational complexity, it uses Eqn. 13 to approximate the attention mechanism.
3.5 Complexity Analysis of cross attention
Algorithm shows the process of one step of cross-attention blocks in SRformer. First, keys and values are split into segments. Then it computes the attention score A[t] for the current query with the corresponding segment of keys Kst- Because the length of queries is usually greater than the number of segments, some queries might use the same segment. Only when = 0, which means is the beginning of a new sequence, it initializes the memory of RAF and compute the full product of KV. After applying softmax, attention scores are multiplied with segmented values to get the segmented attention output Os. It is important to compute the partial product of keys and values only once for each index. addition, the computa- tion of P can be parallelized to improve efficiency, Then is given to the RAF layer for accumulat- ing with the memory through the recurrent process. The output of RAF is multiplied by the query and then normalized. Finally, we add the segmented attention output Os and the complementary OR-
The total complexity of a srquence-to-sqquenee transformer is composed of the cost of the en- coder and that of the decoder. Here, we only com- pare the complexity of cross-attention since the other components are not affected by the proposed method. In table. 1, show both computation and memory (space) complexity analyses of reg- ular transformer, segmented one, and the pro- posed SRformer. Please note that the batch size and the number heads are not considered. To compute the segmented qkv attention, the complex- ity is (ggsd.. The bottleneck for computation is to multiply Ks: * Vs: for m times, which results in O (msd)) complexity. Because the number of seg- ments m k/s, becomes Okkd)). After that, need to multiply Q with it, which costs Oggd²). but as 9 << k. this term can be ignored. The re- current unit also costs much less computation than the bottleneck, so can also be ignored. Since the segment size and the inner dimension of each head are much smaller than k in summarization datasets, the number of computations is fewer than regular transformer, shown in table 1.
Our approach builds connections between seg- ments by using the partial product of keys and val- ues as memory. Since the RAF neuron is applied between segments, the puckpropagation-thruuge time in our model is more efficient than using re- current layers on each input. As discussed before, recurrent attention makes up for the loss caused by segmented attention. Together, they approximate a complete attention mechanism at a lower cost. Although some existing works proposed global at- tention for a similar purpose, global attention based selected locations cannot make up for the loss. Since segmented recurrent attention is a drop-
The memory complexity depends on the sizes of input matrices and intermediate matrices. The size of key and value matrices are k x d. The intermedi- ate matrix is the attention weight A, whose size is x kt By segmentation, its size is reduced to g x s. For recurrent attention, compute the product of keys and values, P. for segments. Hence, there are m matrices of size d d. In table. that given the example sizes, the theoretical reduc- tion of computation is 43%, and that of memory complexity is 29%.
show

---

4 Experimental Results
Table 2 lists the detailed results including ROUGEI, ROUGE2, ROUGEL, and ROUGEL- Sum scores. From the table, we can see that a model with a higher ROUGEI score usually gives higher other scores, and hence can be considered as better. The sizes of keys and values for cross at- tention are 1024. With the default segment size set to 64, the sizes of keys and values are reduced by 16. As consequence, both T5 and BART experi- ence large performance degradation. Then we show ROUGEI scores of SRformer in bold numbers. SR- former benefits from its recurrent attention mech- anism and achieves 3 10% higher scores than the segmented transformers. Although its ROUGE scores are slightly lower than the baseline model, the computational complexity is largely reduced, and will show it also cuts down GPU power usage and memory access time in section 4.4.2.
We test our proposed model and method on T5 sequence-O-ssqueenee transformers (Raffel et al.. 2022) and BART model (Lewis et al., 2020). Since we want to design model for resouree-constraaned environments, we use T5-small, which contains 6 transformer layers and 60 million parameters, and BART-base which contains 6 layers and 140 mil- lion parameters. Our implementation built based on Huggingface transformer platform, which is re- leased under the Apache-2.0 License, and we load the pretrained models before training (Wolf et al., 2019). The implementation is based on Py python package (Paszke et al., 2019).
We run all experiments on Nvidia A40 GPUs. Each experiment is run on single core to compare GPU power and memory usage. We train the mod- els for 25 epochs as the loss and evaluation scores converge at that point. The initial learning rate is 5e The optimizer is AdamW and the learning rate scheduler linear with default settings.
45.00
40.00
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 set of ground-truth summaries, and ROUGE-L measures the longest common sub- sequence between them (Lin, 2004). Note, higher ROUGE score denotes better summarization per- formance.
35.00
SRformer
Rouge1
30.00
LongT5
Transformer
25.00
20.00
Segmented
T5-small
15.00
16
4.1 Results on summarization datasets
Segment Size
The summarization datasets we use evaluate the model include CNl-dailymail (Nallapati et al., 2016), XSUM (Narayan et al., 2018), ArXiv Sum- marization (Cohan et al., 2018), and MediaSum (Zhu et al., 2021). CNN-dailymail dataset con- tains 31 1k anonymized English news articles and their summaries, among which 286,817 training pairs, 13,368 are validation pairs, and 1 11,487 are test pairs (Hermann et al., 2015; See al., 2017). XSUM is an extreme summarization dataset, and thus its summaries only have 23 words on average (Narayan et al., 2018). The numbers of documents for training, validation, and test sets of XSUM dataset are 204,045/1 1,332/1 1,334. ArXiv dataset has 215 thousand documents, including 5% for val- idation and another 5% for test (Cohan et al., 2018) 8). MediaSUM is a media interview dataset consisting of 463.6K dialogue transcripts and summaries, and validation and test set each randomly select 10K instances (Zhu et al., 2021).
Figure 4: Rougel scores CNN-dailymail dataset for different segment sizes across various models.
4.2 Effeet of segment sizes
Figure 4 shows ROUGEI scores on CNN-dailymail dataset with regard to different segment sizes 8. 16, and 64. The number of segments should not exceed the number of generation steps, otherwise, some segments would never be used. Hence, the segment size should not be smaller than 8 for the CNN-dailyme dataset.
The blue line in Fig. 4 shows that for a segmented transformer without recurrence, the ROUGEI score drops significantly by 16% when reducing segment size from 64 to 8. The yellow line shows the results of applying the method from TransformerXI. (Dai et al., 2019) to cross atten- tion blocks of the T5-small model. We add hid- den memory for keys and values and concatenate

---

SRformer
Segmented ROUGE1
Baseline
Dataset
Model
ROUGE2
ROUGEL
ROUGELsum
ROUGEI
ROUGE1
T5
16.79
27.29
36.81
39.39
33.31
41.60
CNN-DM
BART
19.80
29.44
40.40
43.19
33.87
44.54
T5
11.72
26.88
26.88
34.48
27.03
35.35
XSUM
BART
16.38
31.54
31.54
39.03
36.22
41.33
T5
11.50
24.53
31.86
36.04
32.64
37.79
ArXiv
BART
14.97
25.29
37.86
42.99
36,43
43.95
T5 12.89
25.72
26.29
28.59
24.43
29.25
MediaSum
BART 14.70
24.75
28.46
32.36
28.40
33.88
Table Results of SRformer Summarization datasets compared baseline models and their segmented version. The baseline models are T5-small and BART-base.
the hidden memory with the current segment to compute cross attention. With that setting, Trans- fomerXL. performs better than segmented trans- former, but its performance still decreases with the segment size. The green line shows the results of longT5 (Guo et al., 2021). When the number of segments grows as the segment size decreases, it can use more tokens for computing global attention. Hence, its performance inversely proportional to the segment size. For example, when segment size is uses 8 local and 128 global tokens, and its performance is 3% better than using 64 local and 16 global ones. However, the computational over- head of global attention also grows quadratically with the number of tokens. In contrast, SRformer is able to keep a stable performance regardless of the size reduction. Its ROUGEI scores 6 - 22% higher than those of segmented T5 model, and it gives a larger boost in performance compared to TransformerXL. and IongT5. Even when other models use a segment 8 times larger, they are still not able to compete with our model. The results show that SRformer can indeed better approximate full attention.
64 and the model is modified from T5-small. The dataset used for ablation study CNN-dailymail. We varied the following components one by one while keeping the rest of the model fixed:
(I) Remove segmented attention. We modify the model to only use recurrent attention. We observe noticeable drop in performance, however, it is much better than only using segmented attention. This shows that recurrent attention made crucial contribution to the overall performance. (II) Remove normalization. Without dividing norm(K), ROUGE 1 score drops by 3.
(III) Without RAF. Skipping RAF will remove re- current computation and corresponding weights. However, without RAF, it is not able to approxi- mate full attention well. The performance decrease shows the importance of RAF.
(IV) Replacing KR VR with K₃ V.. In this case. can skip computing K V. RAF accumulates information as timesteps increase. However, at the beginning of generation, it only has information from the first few segments. Hence the perfor- mance is not good.
In summary, our ablation study shows that all terms in Eqn. 13 are important components.
Setting
ROUGE1
Setting
ROUGE1
SRformer
39.39
4.4 Evaluation of improvement in efficiency
(I)
36.86
(II)
36.77
4.4.1
Theoretical improvement
(III)
35.84
(IV)
29.72
In 1able. 1, we have shown that SRformer reduces the computation and memory complexity of cross attention. In our experiment, as the length of the generated summary q equals 128 and the length of the input article equals 1024, the theoretical reduction of computation 43%, and that of mem- ory complexity 29%. In Fig. 5, we show how the improvement would change when the length the generated summary scales up. The costs are
Table 3: Ablation study using CaNl-ailymail dataset
4.3
Ablation Study
To better understand the contribution of each com- ponent the performance of SRformer, we con- ducted an ablation study by removing or modifying terms in Eqn. 11 or 13. The segment size is fixed to

---

Computation cost VS text length
Memory cost VS text length
6.01E+08
1.01E107
9.10E106
5.01E+08
8. 8.10E
7.10E106
4.01E:08
Transformer
6.10E+06
Transformer
3.01E+08
SRformer
3.10E+06
SRformer
4, 4.10E+06
2.01E+08
3.10E+06
2. 2.10E+06
1.014008
1.10E:06
1.00E:06
1.00E+05
128
256
512
1024
128
256
512
1024
(a) Computation cost
(b) Memory cost
Figure 5: Theoretical computational cost and memory cost of cross attention when the length of summary (g)
increases from 128 to 1024.
computed based the thcoretical complexity in 1. Suppose that the length of the input article is eight times the length of the summary (A: 8* regular transformer
Model
Mem
Time
Power
BART
49
91
99
TransformerXL
82
70
89
The cost of cross attention will increase quadratically with the text length. The cost cross attention in SRformer not dependent where is small segment size.
LongT5
73
80
90
SRformer
60
60
80
Table Percentages (%) of GPU usage in terms of memory, memory accessing time, and power usage
on q k but Therefore, it increases linearly and the reduction would be 93% compared to transformer when 9 1024. Although existing datasets for abstrac- tive summarization have relatively short summary lengths, large language models may be required to generate long summary for a very long document at a fixed ratio. In that case, SRformer can provide a large improvement in efficiency.
is larger than vanilla BART store weights and gradients, our method less memory than Trans- formerXL and LongT5. In conclusion, by allevi- ating the cross attention bottleneck, the proposed transformer becomes more efficient.
Conclusion
4.4.2 GPU power and memory usage
We track GPU memory, memory accessing time, and power using Weights& Biases tools (Biewald, 2020). The GPU usage of SRformer modified from BART and other models are listed in Table. 4. Our model consumes less GPU power and spends less time accessing memory, compared lar transformer, TransformerXL (Dai al., 2019), and LongT5 (Guo et al., 2021). BART uses 99% GPU power, which almost reaches the limitation of GPU. The percentage of time GPU spent accessing memory is 91%. It means GPU spent most of the time fetching data from memory. However, with segmented recurrent attention, SRformer does not need to all encoded features each step. Hence it saves considerable amount of computation and memory access. only uses 81% % GPU power and 60% GPU memory accessing time. Although GPU memory usage of all recurrent transformers
In this paper, we propose an efficient sequence-to sequence transformer using segmented recurrent attention. Since cross-attention is the computation bottleneck for text summarization. replacing full attention with segmented recurrent attention largely reduce its computation complexity. The proposed RAF neurons help approximate full tention at a low overhead. During inference, they utilize the sequential process to accumulate infor- mation in their internal memory and bridge infor- mation across segments effectively. The efficacy of the proposed model is demonstrated through higher ROUGE scores multiple summarization datasets compared to other segmented transform- ers. Our proposed architecture is especially suit- able for summarization and can be used for other sequence-to-sqqueee to-sequence tasks such as translation and dression-answerigg
regu-

---

Limitations
Winter, Chris Hesse, Mark Chen, Eric Sigler, Ma- teusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. 2020, Language models are few-shot learners. In Ad- vances in Neural Information Processing Systems, volume 33, pages 1877-1901. Curran Associates, Inc.
This work focuses on designing an efficient encoder-decoder model for summarization tasks. The proposed method can not be directly applied to an encoder-only or decoder-only model, although it is possible design modified version for those models. In the case that the length of encoded features is not much larger than the length of de- coder output, it may not be able to show significant reduction in computation and memory cost.
Aydar Bulatov, Yuri Kuratov, and Mikhail Burtsev. 2022. Recurrent memory transformer. In Advances in Neu- ral Information Processing Systems.
Rewon Child, Scott Gray, Alec Radford, and Nya Sutskever. 2019. Generating long sequences with sparse transformers. CoRR, abs/1904. 10509.
Ethics Statement
This work complies with the ACL Ethics Policy. It contributes a new model the deep learning and computational linguistics community which can be used in many applications. We will make the codes open-source for other researchers.
Sayeed Shafayet Chowdhury, Nitin Rathi, and Kaushik Roy. 2022. One timestep is all you need: Train- ing spiking neural networks with ultra low latency. abs/21 10.05929.
Arman Cohan, Franck Dernoncourt, Doo Soon Kim. Trung Bui, Scokhwan Kim, Walter Chang, and Nazli Goharian. 2018. A discourse-awee aware attention model for abstractive summarization of long documents. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computa- tional Linguistics: Humun Language Technologies, Volume (Short Papers), pages 615-621, New Or- leans, Louisiana. Association for Computational Lin- guistics.
A Acknowledgement
This work was supported by the Center for Code- sign of Cognitive Systems (CoCoSys), one of the seven centers in JUMP 2.0, a Semiconductor Re- search Corporation (SRC) program sponsored by DARPA.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Car- bonell, Quoe Le, and Ruslan Salakhutdinov. 2019. Transformer-X XL: Attentive language models beyond fixed- length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pages 2978-2988, Florence, Italy. Asso- ciation for Computational Linguistics.
References
Joshua Ainslie, Santiago Ontañón, Chris Alberti, Philip Pham, Anirudh Ravula, and Sumit Sanghai. 2020. ETC: encoding long and structured data in transform- ers. CoRR, his/2004.08483.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. 2022. Flash Attento:::: Fast and memory-efficient exact attention with 1O-awareness. In Advances in Neural Information Processing Sys- tems.
Hangbo Bao, Li Dong, Furu Wei, Wenhui Wang, Nan Yang. Xiaodong Liu, Yu Wang, Songhao Piao, Jian- feng Gao, Ming Zhou, and Iisiao-Wuen Hon. 2020. Unilmv2: Pseudo-masked language models for uni- fied language model pre-training. In Proceedings of the 37th International Conference on Machine Learn- ing, ICMI.20. JMLR.org.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner. Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. 2021. An image is worth 16x16 words: Transformers for image recognition at scale. abs/2010.11929.
Iz Beltagy, Matthew E. Peters, and Arman Cohan, 2020. Longformer: The long-document transformer. CoRR, abs/2004.05150..
Lukas Biewald. 2020. weights and biases. Software available from wandb.com.
Experiment tracking with
Mandy Guo, Joshua Ainslie, David Uthus, Santiago On- tanon, Jianmo Ni, Yun-Hsuan Sung, and Yinfei Yang. 2022. Long Efficient text-to-text transformer for long sequences. In Findings of the Asxociation for Computational Linguistics: NAACL 2022, pages 724- 736, Seattle, United States. Association for Compu- tational Linguistics.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Knueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens
Mandy Guo, Joshua Ainslie, David C. Uthus, Santi- ago Ontañón, Jianmo Ni, Yun-Hsuan Sung, and Yin- Yang. 2021. Longt5: Efficient eex--to-text trans- former for long sequences. CoRR, abs/21 12.07916.
Phttps://wwa awwweb, org/portal/content) cel-code-ethics

---

Karl Moritz Hermann, Tomás Kociský, Edward Grefen- stette, Lasse Espeholt, Will Kay, Mustafa Suleyman. and Phil Blunsom. 2015. Teaching machines to read and comprehend. In NIPS, pages 1693-1701.
Conference on Empirical Methods in Natural Lan- guage Processing, Brussels, Belgium.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lever, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Te- jani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. 2019. PyTorch: An Imperative Style, High-Performance Deep Learn- ing Library. In Advances in Neural Information Pro- cessing Systems 32, pages 8024 8035, Curran Asso- ciates, Inc.
Jonathan Ho. Nal Kalchbrenner, Dirk Weissenbor, and Tim Salimans. 2019. Axial attention in multidimen- sional transformers. CoRR, abs/1912. 1800.
DeLesley S. Hutchins, Imanol Schlag, Yuhuai Wu, Ethan Dyer, and Behnam Neyshabur. 2022. Block- recurrent transformers. ArXiv, abs/2203.07852.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pap- pas, and François Fleuret. 2020. Transformers are rnns: Fast autoregressive transformers with lin- ear attention. In Proceedings of the 37th Interna- rional Conference on Machine Learning, ICML'20. JoLR.org.
Wachirawit Ponghiran and Kaushik Roy. 2021. Spiking neural networks with improved inherent recurrence dynamics for sequential learning. In AAAI Confer- ence on Artificial Intelligence.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. 2020. Reformer: The efficient transformer. In Inter- national Conference on Learning Representations.
Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong. Nick Barnes, and Yiran Zhong. 2022. The devil in linear transformer. In Proceedings of the 2022 Conference on Empirical Methods in Nat- ural Language Processing, pages 7025-7041, Abu Dhabi, United Arab Emirates. Association for Com- putational Linguistics.
Jun Haeng Lee, Tobi Delbruck, and Michael Pfeiffer. 2016. Training deep spiking neural networks using backpropagation. Frontiers in Neuroxcience, 10.
Benjamin Lefaudeux, Francisco Massa, Diana Liskovich, Wenhan Xiong, Vittorio Caggiano, Sean Naren, Min Xu, Jieru Hu, Marta Tin- tore, Susan Zhang, Patrick Labatut, and Daniel Haziza. 2022. xformers: A modular and hack- able transformer modelling library. https: / thg.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. 2022. Exploring the limits of transfer learning with unified text-to-text trans- former. J. Mach. Learn. Res., 21(1).
Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy. Veselin Stoyanov, and Luke Zettlemoyer. 2020. BART: Denoising sequence-to-sequencc pre-training for natural language generation, translation, and com- prehension. Proceedings of the 58th Annual Meet- ing of the Association for Computational Linguistics, pages 7871-7880, Online. Association for Computa- tional Linguistics.
Robin Rombach, Andreas Blattmann, Dominík Lorenz, Patrick Esser, and Björn Ommer. 2022. High- resolution image synthesis with latent diffusion mod- els. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). pages 10684-10695.
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. 2021. Efficient content-based sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53- 68.
Fei-Fei Li. 2022. Cs231n: Convolutional neural net- works for visual recognition.
Chin-Yew Lin. 2004. ROUGE: A package for auto- matic evaluation of summaries. In Text Summariza- tion Branches Out, pages 74-81, Barcelona, Spain. Association for Computational Linguistics.
Chitwan Saharia. William Chan, Saurabh Saxena. Lala Li, Jay Whang, Emily Denton, Seyed Kam- yar Scyed Ghasemipour, Raphael Gontjjo-opes, Bureu Karagol Ayan, Tim Salimans, Jonathan Ho, David J. Fleet, and Mohammad Norouzi. 2022. Pho- torealistic text-to-image diffusion models with deep language understanding. In Advances Neural In- formation Processing Systems.
Ramesh Nallapati, Bowen Zhou, Cicero dos Santos, Çağlar Gulgehre. and Bing Xiang, 2016. Abstrac- tive text summarization using pepueece-to-sequeeee RNNs and beyond. In Proceedings of the 20th SIGNLL Conference on Computational Natural Lan- guage Learning, pages 280-290, Berlin, Germany. Association for Computational I Linguistics.
A Abigail See, Peter J. Liu, and Christopher D. Manning, 2017. Get to the point: Summarization with pointer- generator networks. In Proceedings of the 55th An- nual Meeting of the Association for Computational Linguistics ( (Volume I: Long Papers), pages 1073- 1083, Vancouver. Canada, Association for Computa- tional Linguistics.
Shashi Narayan, Shay B. Cohen, and Mirella Lapata. 2018. Don't give me the details, just the summary! Topic-aware convolutional neural networks for ex- treme summarization. In Proceedings of the 2018

---

Tay, Mostafa Dehghani, Dara Bahri, and Donald Met- zler. 2022. Efficient transformers: A survey. ACM Comput, Surv, 55(6).
A
Appendix
A.1 Expanding attention equations
Denote the current i-th segment as S = [i 8, (i + 1)s} The indices of remaining segments are de- noted as Rg { [1, & S.). Then the full attention, segmented attention, and remaining attention at timestep / will be
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, ukasz Kaiser, and Illia Polosukhin. 2017. Attention is all you need. In Advances in Neural Information Pro- cessing Systems, volume 30. Curran Associates, Inc.
Sinong Wang. Belinda Li, Madian Khabsa, Han Fang. and Hao Ma. 2020. Linformer: Self-attention with linear complexity. CoRR, abs/2006.04768.
Olt] softmax (altmaxx(rrra K" K"W
Os([] - softmax (Qr K K(;)Vs Vss
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pier- ric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, and Jamie Brew. 2019. Huggingface's 's transformers: State-of-the-an natural language processing. CoRR, abs/1910.03771.
On(4 softmax Vns
For example, the softmax operation applied on the remaining attention OR can be expanded as
On softmax (Q: K/Rs Vis
(17)
Stanislaw Wozniak, Angeliki Pantazi, Thomas Bohnst- ingl, and Evangelos Eleftheriou. 2020. Deep learning incorporating biologically inspired neural dynamics and in-memory computing. Nature Machine Inteli- gence, 236..-336.
= Φ(Qi * Kk) Vis
(18)
ΣJER, Φ(Q1 KJ)
ΣjER, (Ø(Q, * KJ)VV,
(19)
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang. and Amr Ahmed. 2020. Big bird: Transformers for longer sequences. CoRR, abs/2007.14065
ΣjER, p(Q Kÿ)
In the common implementation of softmax atten- tion, (xX) exp(X) an exponential function. For the numerical stability of the computation, an- other choice is 6(X) = exp(X - log(max; X.)) (Li, 2022). However, Katharopoulos et al. (2020) proposed to replace it with a kernel function such that ((XY) d( ())())) Then we have
Chenguang Zhu, Yang Liu, Jie Mci, and Michael Zeng. 2021. MediaSum: A large-scale media interview dataset for dialogue summarization. In Proceedings of the 2021 Conference of the North American Chap- ter of the Association for Computational Linguistics: Human Language Technologies, pages 5927-5934, Online, Association for Computational Linguistics.
0[1] ((Q,) (((K)))))
(20)
ΣØ= oQQ KT)
1sl 8l (QQ)) Ejes, (0 (o(K)) KF));)
(21)
EjEs: b(Q; KT)
OR( ((2)) EjeR, (0 K3))))
(22)
Ejek. Φ(Q) KF)
Denote the divisor in the above equations using Ca Cr, and c.
C. Σ s(Q+ K})
jES,
Σ
((Q+= K!)
Gr
jERl
(Qp Ky) = c +
c =
Then we can write the full attention as a weighted summation, where the weights depend on c and c,-

---

We can use factor represent the weights.
BART adds positional bias the inputs of the first layer, so there is no need to take care of that sepa- rately. However, the encoder of BART generates features of various sizes and adds padding to each batch. To ignore the padding, an encoder mask is used. Because softmax is removed in recurrent attention. adding negative infinity to masked loca- tions is not suitable. Instead, we filter out abnormal values of K V and P with clamp function before computing recurrent attention.
Olt! _ 0st + _ [n[t] Olt] x0 + (1-)))
practice, we find that simple addition without weighting still works. Our recurrent attention fur- ther approximates OR below.
Qi RAF(K) * VR)
Onl! ~
norm (
To apply SRformer to other types of transformers, we recommend checking the positional bias and masks and adjusting accordingly.
Moveover, we estimate the error caused by segmented attention.
O -Os
G * p(2 * K"
c* (Q: K
Cs
c* #(Q++ K R
R Cp * Φ1Qt * KTJV
(+)
(0R 0)
Cr (OR 0s)
Cr
A.2 Pseudo code for RAF neurons
class RAF(nn, Module):
def init (dim) ):
self linear = Linear (dim, dim)
self. activation = ReLU
self. = parameter (1.0,
[guuiregrad=True)
thre = parameter (0. eequire_grad=te
def forward (input. mem. t):
x self linear (input)
mem = leak * mem + x
y = mem/ self thre - 1.0
mem = mem - self. thre +( y>0) out = ReLU(y)
return out
A.3 Positional bias and masks
Different transformers such T5 and BART use positional bias and masks differently. Hence, we need to take this into consideration when imple- menting segmented recurrent attention.
For example, T5 adds positional bias to the atten- tion weights A. For segmented attention, we need to segment bias as shown in the following equa- tions.
bs: bli x (1 + 1) Osit] softmax (As( + bsg) Vsi


Reviews:
[{"review_id": "d2HDwKvG4a", "reviewer": "Reviewer_f7Z1", "summary": "", "questions": "- In Fig.2 , why do we still need calculate $(QK^T)V$, can we calculate all attention in the linear form shown in eq 13?", "limitations": "", "rating": 3, "confidence": 4, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the feedback and for finding our work well-motivated, novel and effective. \n\n**Comparison to existing works**: Regarding the comparison to existing methods, we have compared the performance of the proposed SRformer to Transformer-XL and LongT5 in Fig. 4 (submitted manuscript). Only our method can achieve results comparable to full attention when small segment sizes are used, and its performance is not affected by the reduction of segment size. We cited RMT in section 2, but since it was a concurrent work, we had not compared to it. Additionally, we provide comparison with another previous work, Longformer [1]. Our proposed method with segment size set to 64 achieves 36.04 and 42.99 ROUGE 1 scores on ArXiv dataset when applying to T5 or BART, whereas Longformer with segment size set to 1024 achieves 35.21 using BART and shows a clear trend of decreasing with the segment size (Figure 3 of [1]). We have recorded the actual memory footprints and GPU usage of training different models as shown in Table B. Our method uses the lowest percentages of GPU power and memory accessing time, providing 29%/31% improvement over the baseline. Although GPU memory usage of all recurrent transformers is larger than vanilla BART to store weights and gradients, our method uses less memory than TransformerXL and LongT5. \n\n**Theoretical improvement:** As mentioned, the proposed method improves the computational complexity from O(qkd) to O(qsd + kd^2). In many cases, the length of summary, q, depends on the length of article, k. Then the improvement is from quadratic to linear. We would like to provide Table A, which is an extension for Table 1 in our paper, to show theoretical computational reduction increases proportionally with the lengths. When q=128 and k=1024, the overall cost is reduced by 43.8%. When they both increase by 4 times, the improvement will be 85.9%.\n\n**Table A: Number of computation of cross attention VS length**\n|\tq/k\t    |    128/1024\t|        256/2048\t   |     512/4096|\n|---------------|-------------------|--------------------|--------------|\n|Transformer (full attention)|\t8.39E+06 |\t3.36E+07 |\t1.34E+08|\n|SRformer\t    |    4.72E+06 |\t9.44E+06 |\t1.89E+07|\n|reduction\t     |   43.8%\t|\t71.9%\t|\t85.9%|\n\n**Table B: Memory and power comparisons**\n|\t\t\t   |     GPU Mem(%)\t| GPU mem access time(%)\t|Power usage(%)\t |\n|---|---|---|---|\n|BART\t(Baseline)\t|\t49\t\t|\t91\t\t\t|\t99|\n|TransformerXL\t\t|\t82\t\t|\t70\t\t\t|\t89|\n|LongT5\t\t\t\t |       73\t\t|\t80\t\t\t|\t90|\n|SRformer - BART\t\t  |      60\t\t|\t60\t\t\t|\t80|\n|SRformer - T5\t\t|\t        62\t|\t\t60\t\t|\t\t81|\n\n**Block-wise attention and RNN-based attention**: Although we combine block-wise attention and recurrent based attention, we analyze the loss caused by block-wise attention and re-design the recurrent attention to specifically compensate for the lost information. As far as we know, the existing work did not show the complementary relation between local attention and global attention mathematically. Hence, we believe our work gives a novel view of how local attention and global/recurrent attention collaborate.\n\n**Attention in the encoder**: We keep the vanilla attention in the encoder because it runs only once during generation, while the decoder takes steps to generate new tokens. Hence it is acceptable for it to use vanilla attention. Allowing the encoder to access all tokens would make sure the summary captures all information. It is possible to apply the segmentation and recurrence approach to the encoder as well, which will improve efficiency. However, it leads to further information loss before the decoder part resulting in performance degradation, so it is a trade-off.\n\n**Decoder-only models**: It is possible to integrate the proposed method into the decoder-only structures. However, decoder-only models would keep using the original article as inputs at each generation step while concatenating the generated words to its end. Then it becomes difficult to apply segmented attention to the article and summary separately. On the other hand, encoder-decoder models separately process encoded articles and generated summary and explicitly learn input-output relationships. Hence, we think it is necessary to design an efficient encoder-decoder model for summarization task, and our method is more suitable for such a model.\n\n**Response to the question**: In section 4.3 ablation study, we have shown that if we only use recurrent attention in Eq. 13, the ROUGE score would drop by 3. In addition, we have tried linear attention, proposed in paper [2], but it gives ROUGE scores lower than 10. We think the reason might be providing too little information at a time (segment size=1) and a lack of exponential softmax. Hence, we think segmented attention with softmax and recurrent attention are both important and keep them in our design.\n\nWe thank the reviewer for referring  to “The Devil in Linear Transformer” [3] as it is highly related to our work. We have noticed that unbounded values could impact the convergence, and our solution is to apply normalization on recurrent (linear) attention and keep using softmax in segmented attention. We will add this paper [3] to our related works.\n\n[1] Iz Beltagy, Matthew E. Peters, and Arman Cohan. 2020. Longformer: The long-document transformer. CoRR, abs/2004.05150.\n\n[2] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. 2020. Transformers are rnns: Fast autoregressive transformers with linear attention. In Proceedings of the 37th International Conference on Machine Learning, ICML’20.\n\n[3] Zhen Qin and XiaoDong Han and Weixuan Sun and Dongxu Li and Lingpeng Kong and Nick Barnes and Yiran Zhong, The Devil in Linear Transformer, 2022, abs/2210.10340.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "- The paper is well written and easy to follow.\n- The motivation is clear, the segment attention can reduce the computation and the recurrent can fuse the global contexts.\n- The design of RFA is novel which is a new method to include global information into segments.\n- The experiments show the effectiveness of this model on the summarization task.", "reasons_to_reject": "- The actual memory footprints and speed comparison tables are needed to validate the efficiency. This method should be compared to existing methods such as Transformer-XL, RMT or techniques such as KV-cache since they share similar methodology.\n- The novelty of the propose method is limited, it only incremental improvements over the existing block-wise attention and RNN-based attention.\n- The summarization task commonly takes long-sequence document as input, why the encoder still uses vanilla attention as it might has large computational overhead? \n- With the advent of decoder-only structures, I cannot see the reasons for designing a specific model for summarization task as the decoder-only model can also do this task, Why don't simply integrate this method into the decoder-only structure?", "questions_to_authors": "- In Fig.2 , why do we still need calculate $(QK^T)V$, can we calculate all attention in the linear form shown in eq 13?", "paper_topic_and_main_contributions": "- This paper proposes a new structure for sequence-to-sequence task\n- It proposes a segmented attention to reduce the computation complexity of vanilla attention\n- It proposes a recurrent module to accumulate the information across the segments.\n- The method is validated on summarization benchmarks", "excitement": 3, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Missing_References": "- 'The Devil in Linear Transformer'.  This paper also adopts segmented attention and uses linear attention to fuse information"}, "extra_text": {}}, {"review_id": "NMHJ5HH4cn", "reviewer": "Reviewer_yFAe", "summary": "", "questions": "1. The model seems not that computationally efficient compared with other models based on the analysis, though it achieves seemingly better performance, and more experimental results need to be explored based on fair comparison.\n2. Do you consider comparing with more previous methods?\n3. In this paper, only T5 and BART are used. Do you try testing on larger models now that reducing cost for large models also counts?", "limitations": "", "rating": 3, "confidence": 3, "soundness": 3, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the comments. \n\n**Regarding the computational reduction**, we improve the computational complexity from O(qkd) to O(qsd + kd^2). In many cases, the length of summary, q, depends on the length of article, k. Then the improvement is from quadratic to linear. The computational reduction increases proportionally with the lengths, as shown in Table A, which is an extension for Table 1 in our paper. When q=128 and k=1024, the overall cost is reduced by 43.8%. When they both increase by 4 times, the improvement will be 85.9%. Moreover, most existing works only reported the computation reduction of attention and ignored the other layers such as feed forward layers in overall cost. We have recorded the actual memory footprints and GPU usage of training different models as shown in Table B. Our method uses the lowest percentages of GPU power usage and memory accessing time, providing 29%/31% improvement over the baseline. Although GPU memory usage of all recurrent transformers is larger than vanilla BART to store weights and gradients, our method uses less memory than TransformerXL and LongT5. Considering these, we believe our improvement is significant. Besides, we have provided experiments on four different datasets in Table 2 in our paper and an ablation study in section 4.3 to validate the effectiveness of the method. \n\n**Table A: Number of computation of cross attention VS length**\n|\tq/k\t    |    128/1024\t|        256/2048\t   |     512/4096|\n|---------------|-------------------|--------------------|--------------|\n|Transformer (full attention)|\t8.39E+06 |\t3.36E+07 |\t1.34E+08|\n|SRformer\t    |    4.72E+06 |\t9.44E+06 |\t1.89E+07|\n|reduction\t     |   43.8%\t|\t71.9%\t|\t85.9%|\n\n**Table B: Memory and power comparisons**\n|\t\t\t   |     GPU Mem(%)\t| GPU mem access time(%)\t|Power usage(%)\t |\n|---|---|---|---|\n|BART\t(Baseline)\t|\t49\t\t|\t91\t\t\t|\t99|\n|TransformerXL\t\t|\t82\t\t|\t70\t\t\t|\t89|\n|LongT5\t\t\t\t |       73\t\t|\t80\t\t\t|\t90|\n|SRformer - BART\t\t  |      60\t\t|\t60\t\t\t|\t80|\n|SRformer - T5\t\t|\t        62\t|\t\t60\t\t|\t\t81|\n\n\n**Regarding the comparison to existing methods**, we have compared the ROUGE1 scores of the proposed SRformer to Transformer-XL and LongT5 in Fig. 4 (submitted paper). Compared to these existing baselines, only SRfromer achieves comparable results to full attention when small segment sizes are used, and its performance is only slightly affected by the reduction of segment size. Furthermore, we provide a new comparison with another existing work, Longformer [1]. As shown in Table 2 (submitted manuscript), our proposed method with segment size set to 64 achieves 36.04 and 42.99 ROUGE 1 scores on ArXiv dataset when applying to T5 or BART, whereas Longformer with segment size set to 1024 achieves 35.21 using BART and shows a clear trend of decreasing with the segment size (Figure 3 of [1]).\n\n**Regarding testing on larger models**, although larger models have more layers and more heads, the architectures of transformers are similar, and the attention mechanism remains the same. Hence, we expect the proposed method to be valid. The percentage of computational reduction is dependent on the lengths q and k, therefore, it will show similar trends in larger models. Due to resource constraints in an academic setting, it is difficult for us to experiment with training these LLMs. But we thank the reviewer for the suggestion. \n\n[1] Iz Beltagy, Matthew E. Peters, and Arman Cohan. 2020. Longformer: The long-document transformer. CoRR, abs/2004.05150.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "1. SRformer achieves better computational complexity and accuracy trade-off through aggregating information across segments with recurrent attention.\n2. Experimental results show that the method achieves better performance.", "reasons_to_reject": "1. Computational reduction is not that surprising compared to other methods.\n2. It seems that more fair experiments are needed to demonstrate the effectiveness of the method.", "questions_to_authors": "1. The model seems not that computationally efficient compared with other models based on the analysis, though it achieves seemingly better performance, and more experimental results need to be explored based on fair comparison.\n2. Do you consider comparing with more previous methods?\n3. In this paper, only T5 and BART are used. Do you try testing on larger models now that reducing cost for large models also counts?", "paper_topic_and_main_contributions": "This paper proposed a segmented recurrent transformer that combines segmented attention with recurrent attention. This method achieved higher scores than other approaches and reduced the computational complexity of cross attention.", "excitement": 3, "reproducibility": 3, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "IYumMxhbut", "reviewer": "Reviewer_qsxJ", "summary": "", "questions": "> line 358 \"Unlike RNNs, RAF does not apply weights on the hidden memory but only uses a learnable leak and threshold. Hence, it requires fewer weights and lesser computation compared to an RNN or LSTM.\"\n\nDoes following the LSTM style not improve performance?", "limitations": "", "rating": 3, "confidence": 3, "soundness": 4, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer’s positive feedback. \n\n**RAF Neuron**: In SRformer, we propose Recurrent-accumulate-and-fire (RAF) neurons for implementing recurrent attention. The temporal dynamics of a RAF neuron can be expressed as\n\n    U[t]=λ*U[t-1]+linear(X[t]) (i),\n    Y[t]=U[t]/V_th -1.0 (ii)\n    if U[t]>V_th,O[t]=Y[t],U[t]=U[t]-V_th (iiiA)\n    else,O[t]=0, (iiiB)\n\nwhere U is the accumulated membrane potential, λ is the learnable leak, V_th is the learnable threshold parameter, X is the input, Y is just a scaled and shifter version of U and O is the output. This RAF neuron is conceptually inspired by Leaky-Integrate-and-Fire(LIF) neurons of biological brains [1][2], which can efficiently aggregate sequential information. At each timestep, the membrane potential U[t] is updated by adding the incoming input X[t] with the previous information U[t-1], where U[t-1] gets decayed by a learnable leak λ. Next, we check if the accumulated U[t] crosses the threshold (V_th) of information propagation to the next layer. If yes, that accumulated information Y[t] gets passed to the next processing layer (and V_t is subtracted from the U[t] as a reset of potential), else the output is 0 (and U[t] is retained). This results in a sparse and low-complexity bio-inspired recurrence mechanism. Note, the difference between LIF and RAF is that LIF integrate binary input spikes and propagates discrete output spikes, whereas RAF accumulates the partial products of keys and values and uses ReLU activation. The pseudo-code for RAF neurons is provided in the submitted manuscript (appendix section A.2) and the schematic of their working mechanism is given in Fig. 3. \n\nA limitation of RAF neurons is that modeling very long sequences and very complicated dependencies can be tricky using them. However, as our results demonstrate, they can provide highly efficient solutions as a low-cost recurrent block for processing segments and achieve results comparable to the costly baselines on popular summarization benchmarks.\n\n**RAF versus LSTM/RNN**: We have tried replacing RAF with RNN and the results show that RNN does not significantly improve the performance (ROUGE1 ±0.2), although it uses extra weights on hidden states. Hence, we believe the recurrent attention is important for making up the loss caused by segmentation, and a simple recurrent unit can do the work. We have not tried LSTM because it takes four times more parameters/computation than RNN. Similar results are reported in [2], where LSTMs provide only marginal performance improvements compared to accumulate-and-fire type activations, however the parameter count becomes 2X, with increased compute. As a result, the RAF can be regarded as a lightweight and efficient recurrent unit in lieu of LSTM cells, where the leak ($\\lambda$) of RAF acts as a simple forget mechanism and the thresholding operation ($V_{th}$) leads to sparse and efficient information propagation.\n\nThe effect of the RAF neurons is elucidated in Table 3, case (III) of the submitted paper, where the ROUGE1 score drops by 4.0 when RAF is not used. Furthermore, we have tested how the design of RAF would affect the model performance. For example, we have tried to remove the linear layer, replace ReLU with binary spike activation used in LIF, or replace the linear layer with element-wise multiplication with learnable parameters, they all perform worse than the current design of RAF (ROUGE1 drop to 32.62%, 32.63%, and 28.52%, respectively, on CNN-dailymail). Due to the page limitation, we only showed the final design of RAF that gives the best results. We will add these ablation studies and an in-depth analysis of RAF in the appendix to the camera-ready version.\n\nFor the typo regarding the format of the rouges in Table 2, we will capitalize ROUGE scores to be uniform.\n\n**Additional comparison tables**\n\nRegarding the computational reduction, we improve the computational complexity from O(qkd) to O(qsd + kd^2). In many cases, the length of summary, q, depends on the length of article, k. Then the improvement is from quadratic to linear. The computational reduction increases proportionally with the lengths, as shown in Table A, which is an extension for Table 1 in our paper. When q=128 and k=1024, the overall cost is reduced by 43.8%. When they both increase by 4 times, the improvement will be 85.9%. Moreover, most existing works only reported the computation reduction of attention and ignored the other layers such as feed forward layers in overall cost. \n\nWe have recorded the actual memory footprints and GPU usage of training different models as shown in Table B. Our method uses the lowest percentages of GPU power usage and memory accessing time, providing 29%/31% improvement over the baseline. Although GPU memory usage of all recurrent transformers is larger than vanilla BART to store weights and gradients, our method uses less memory than TransformerXL and LongT5. Considering these, we believe our improvement is significant. Besides, we have provided experiments on four different datasets in Table 2 in our paper and an ablation study in section 4.3 to validate the effectiveness of the method. \n\n**Table A: Number of computation of cross attention VS length**\n|     length\tq/k\t    |    128/1024\t|        256/2048\t   |     512/4096|\n|---------------|-------------------|--------------------|--------------|\n|Transformer (full attention)|\t8.39E+06 |\t3.36E+07 |\t1.34E+08|\n|SRformer\t    |    4.72E+06 |\t9.44E+06 |\t1.89E+07|\n|reduction\t     |   43.8%\t|\t71.9%\t|\t85.9%|\n\n**Table B: Memory and power comparisons**\n|\t\t\t   |     GPU Mem(%)\t| GPU mem access time(%)\t|Power usage(%)\t |\n|---|---|---|---|\n|BART\t(Baseline)\t|\t49\t\t|\t91\t\t\t|\t99|\n|TransformerXL\t\t|\t82\t\t|\t70\t\t\t|\t89|\n|LongT5\t\t\t\t |       73\t\t|\t80\t\t\t|\t90|\n|SRformer - BART\t\t  |      60\t\t|\t60\t\t\t|\t80|\n|SRformer - T5\t\t|\t        62\t|\t\t60\t\t|\t\t81|\n\n\n[1] Lee, J. H., Delbruck, T., & Pfeiffer, M. (2016). Training deep spiking neural networks using backpropagation. Frontiers in neuroscience, 10, 508.\n\n [2] Ponghiran, W., & Roy, K. (2022, June). Spiking neural networks with improved inherent recurrence dynamics for sequential learning. In Proceedings of the AAAI Conference on Artificial Intelligence (Vol. 36, No. 7, pp. 8001-8008).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "\nSRformer achieves a superior trade-off between computational complexity and accuracy. It reduces computation by approximately 40% when compared to regular cross-attention models that have access to the full sequence. This outcome demonstrates the effectiveness of segmented recurrent attention for processing sequential data, particularly when dealing with extremely limited computational resources.", "reasons_to_reject": "Although RAF neurons are introduced as a key element of SRformer, the paper lacks an in-depth analysis of their working mechanisms and limitations. Providing more insights into how RAF neurons affect model performance would improve the overall contribution of the paper.", "questions_to_authors": "> line 358 \"Unlike RNNs, RAF does not apply weights on the hidden memory but only uses a learnable leak and threshold. Hence, it requires fewer weights and lesser computation compared to an RNN or LSTM.\"\n\nDoes following the LSTM style not improve performance?", "paper_topic_and_main_contributions": "Transformers have exhibited exceptional performance in various domains, including language and vision. However, their computational cost increases quadratically with sequence length, which makes them unsuitable for resource-constrained applications. To address this issue, the paper proposes an alternative approach by dividing the entire sequence into segments and employing local attention mechanisms on each segment.\n\nSegmented Recurrent Transformer (SRformer), combines segmented (local) attention with recurrent attention. By using recurrent attention, the loss incurred by reducing the attention window length is compensated and information across segments is aggregated effectively. The SRformer utilizes Recurrent Accumulate-and-Fire (RAF) neurons, which possess inherent memory, to update the cumulative product of keys and values. The combination of segmented attention and lightweight RAF neurons ensures the efficiency of the proposed transformer. ", "excitement": 3, "reproducibility": 4, "ethical_concerns": "No", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"Typos_Grammar_Style_and_Presentation_Improvements": "The format of the rouges in Table 2 does not appear to be uniform."}, "extra_text": {}}]

Author rebuttal:

Meta-review: Reviewers agree that the work is moderately sound and moderately exciting. Reviewers like the efficiency and strong experimental results of the novel way to use attention during generation. While some flaws are highlighted there is consensus that the overall method and evaluation is sound. I recommend acceptance to the main conference or findings.
Decision: Accept-Findings