Staircase Attention for Recurrent Processing of Sequences
Authors: Da JU
Stephen Roller
Sainbayar Sukhbaatar
Jason E Weston
Venue: NeurIPS 2022 Accept (2022)

Abstract: Attention mechanisms have become a standard tool for sequence modeling tasks, in particular by stacking self-attention layers over the entire input sequence as in the Transformer architecture. In this work we introduce a novel attention procedure called staircase attention that, unlike self-attention, operates across the sequence (in time) recurrently processing the input by adding another step of processing. A step in the staircase comprises of backward tokens (encoding the sequence so far seen) and forward tokens (ingesting a new part of the sequence). Thus our model can trade off performance and compute, by increasing the amount of recurrence through time and depth. Staircase attention is shown to be able to solve tasks that involve tracking that conventional Transformers cannot, due to this recurrence. Further, it is shown to provide improved modeling power for the same size model (number of parameters) compared to self-attentive Transformers on large language modeling and dialogue tasks, yielding significant perplexity gains.

Staircase Attention for Recurrent Processing of Sequences
Da Ju Meta AI dajuêfb.com com
Stephen Roller Meta AI roller@fb.cm com
Sainbayar Sukhbaatar Meta AI sainbarefb.. com
Jason Weston Meta Al jase@meta, com
Abstract
Attention mechanisms have become a standard tool for sequence modeling tasks, in particular by stacking self-attention layers over the entire input sequence as in the Transformer architecture. In this work we introduce a novel attention procedure called staircase attention that, unlike self-attention, operates across the sequence (in time) recurrently processing the input by adding another step of processing. A step in the staircase comprises of backward tokens (encoding the sequence so far seen) and forward tokens (ingesting a new part of the sequence). Thus our model trade off performance and compute, by increasing the amount of recurrence through time and depth. Staircase attention is shown to be able to solve tasks that involve tracking that conventional Transformers cannot, due this recurrence. Further, it is shown to provide improved modeling power for the same size model (number of parameters) compared self-attentive Transformers on large language modeling and dialogue tasks, yielding significant perplexity gains.
Introduction
Early breakthrough work in neural language modeling considered a fixed context size of tokens that are embedded with a lookup table, followed by nonlinearities and final softmax produce a probability distribution for the next output token in sequence [5] Such models were replaced, that were able to consider arbitrary context length via the ability to store state in their memory using recurrent steps through the data, in contrast the fixed length constraint of earlier models. Moreover, the repeated application of the recurrent network across the sequence also made the models considerably deeper: a given representation a function of large number of nonlinearities due to previous state. This allows such models to track state store long-term memories, and potentially solve highly nonlinear sequential tasks. Today, with the advent of attention- based models [2] and particular Transformers (361, fixed length inputs that eschew recurrence
pre- Trenrrassfrrer, with recurrent models such as RNNs and LSTMs I22
16. 28
back the norm, thanks mainly due to deep
stacks of nonlincarities on those fixed inputs that are also well suited to modern hardware, leading the authors of Vaswani et al. (36] to claim that non-recurrent attention is "all you need." However, some of the advantages just mentioned of earlier models tracking state, and solving highly nonlinear sequential tasks have to some degree been lost [
In this work, we introduce novel recurrent model that utilizes a novel attention procedure called staircase attention. We show that our new model, which utilize both sequence aligned recurrence (in time) and recurrence depth can bring advantages to modern models, particular in terms of lower language modeling perplexities given the same number of parameters, and for solving nonlinear state-tres tracking tasks. Staircase attention, like self-attention, processes tokens in parallel for speed, but unlike self-attention, operates across the sequence (in time) recurrently processing the input by
keep track of an evolving state given sequence of changes to it
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

transformer core
indentity
input layer
output layer
time
layers
C tokens
N chunks
(a) Transformer core
(b) Staircase
(c) Cached Staircase
Visualization of our Staircase models. (a) The building block of our models is Transformer core with layers that processes tokens (equals chunk) in parallel. (b) In the Staircase model, Transformer cores are stacked diagonally, each step sees one new input chunk. column corresponds to processing
Figure 1:
single chunk and is composed of cores sandwiched between input and output layers. All N chunks in a row are processed in parallel and are accessible by self-attention. (c) the Cached Staircase, some final cores are replaced by an identity function to reduce computation. The output from the last core cached and included within the attention span of later chunks, so the context size remains the same.
adding another step of processing. A step (processing block) in the staircase consists of backward tokens (encoding the sequence so far seen) and forward tokens (ingesting new part of the sequence). Thus, on each time step the block moves forward in time, retaining memory comprised of multiple vectors stored in the backward tokens (the recurrent tokens). The blocks utilize the same model weights for each step, hence giving recurrence in depth.
Compared to Transformers, Staircase models can retain a recurrent memory in time, and repeated application of it in the recurrent network over the sequence also makes the model considerably deeper for the same number of parameters (but not necessarily the same amount of compute).
We show on two tasks requiring state- trackin that Staircase models can perform successfully where Transformers fail. We then show on two language modeling and dialogue modeling task for the same number of parameters, significantly lower perplexities can be obtained compared to standard Transformers for certain kinds of Staircase models. We thus analyze our models and show that one can control the amount of recurrence and depth which can trade off compute for performance, which has practical implications depending on available memory and compute architecture constraints. The Staircase models perform well on both state-tracking tasks and language modeling tasks, providing good performance across the board. Our code will be made publicly available.
Related Work
2, non-recurrent models have superseded in popularity recurrent 16. which were for time dominant in NLP applications, particularly when involving sequence generation. The first models use stacked layers of attention over input word embeddings and position encodings, as replacement to recurrence in time, were end-to-end memory networks 34] Those models were shown on the task of language modeling to perform well compared to LSTMs, but in experiments still shared some weights across depth, which we refer to as recurrence depth (also referred to as a "recurrent attention mechanism" Vaswani al. (6)). Transformers [36] removed the constraint of sharing any weights among layers at all. and showed this improves performance (at the cost of using more parameters). Transformers additionally contributed other notable improvements such multi-head, self-attention and residual blocks. Such models do not have any recurrence at all, and are the current of-te-of-the-ar architecture choice for many tasks.
Since attention was introduced
models like RNNs and LSTMs 2
Since then, several variants Transformers have arisen that attempt to incorporate recurrence again by sharing some weights. The Universal Transformer proposes an extreme variant of applying the same layer (with shared weights) repeatedly [T]. Similarly, ALBERT 22 also shares the same weights across all layers for pretraining/finetuiing setting where the BERT model [[] is effectively

---

compressed; they also consider sharing only the self-attention or only the ferd-forward weights. We note also that several works, while not sharing parameters of layers, have studied the ordering of the sublayers of Transformers, in particular Sandwich (29] and Macaron 25
Transformers.
Some works have also attempted incorporate sequence-aligned recurrence Transformers. Chen et al. [ incorporate LSTM layers into the Transformer, and Hao et al. TSI blend a non-recurrent and recurrent model (e.g., an RNN) together with gating function. Transformer-X XL p employs
segment-level recurrence mechanism to effectively cache and speed up computations in long-context sequence tasks. We note that number of recent architectures have also focused on allowing long- context in Transformers, although typically without employing recurrence ? 211 4. Linearizing self-attention for efficiency [19 31 uses state that is updated in recurrent way, but the model still remains feedforward. Finally, the Feedback Transformer (13), perhaps the most similar work to ours, incorporates step-wise recurrence in the Transformer, with a step size of one and a fixed cached memory in the past. It achieves good results but has relatively high computational cost due to its architecture not fully exploiting parallelism.
In this work, we compare architectures with the number of model parameters fixed, and explore increasing recurrence and/or compute given that fixed budget. An orthogonal topic of study is to fix the compute budget instead, but do not fix the amount of parameters, e.g. research into large, sparse (typically non-recurrent) models that may require to be spread over cluster [4] We focus on the former here, but learnings from each direction should be complementary.
3 Method
3.1 Background
In this paper, we consider decoder-only Transformers [1 that are applied to sequential tasks like language modeling. In this setting, a Transformer model takes as input a sequence (r1, T2,,,,,,...... ...,,..... &T) of 7' tokens and outputs sequence of the same size
V1 1/2 yr = TRANSFORMER(11) 2,,.,,,,,
(1)
we separate out the input embedding h, = fm (z+) and the final output module yt = (he), we are left with the Transformer core as shown Figure I
h hy...., hy TC(b,, b by....,
(2)
which consists of L layers that compute final hidden states for each token. Each layer is composed of self-attention and feedforward sublayers. In the self-attention sublayer, causal masking is applied to
prevent tokens from attending to any future token, and we use relative position embeddings [32]. See
Vaswani et al. 36 for more details about the sublayer architecture of Transformers.
3.2 Staircase Model
We now describe our model that utilize staircase attention. Later, we also introduce Cached Staircase version that is more computationally efficient. Their graphical representation may be found in
Figure I
Unlike Transformers, a Staircase model ingests input tokens in smaller chunks, as shown Figure I.. Inside Staircase model lies Transformer steps. With each recurrent step, a Staircase model moves C' tokens forward in time, which we call the forward step size. In addition to those C' forward tokens, the model simultaneously also processes NC - C tokens that come from the previous step, which we call backward tokens. We refer to the total number of tokens NC that are being processed parallel the step size.
that processes each input token in N > recurrent
Let us denote a chunk C input tokens as HD [h4c+1...... hcc+c} for brevity. Here hy is the embedding vector of the token It., At each step, the Staircase model processes N chunks in parallel via the Transformer core
(3) Among the input chunks, only HOpN contains new token embeddings (i.e., forward tokens) while the remaining N chunks come from the previous recurrent step (i.e., backward tokens). The
H.N HNN 7442 H TC(HN-1 H
H.Lx).
3

---

superseript nt of H," denotes the number of computational passes through the Transformer core. After N passes through the core module, the output for particular token I, is computed with
y/ = Saun (ha) for all h, H..
As you can see, an input token gets processed by the same core module N times. This makes it possible to control the amount of computation by varying the number of recurrent steps N without changing the number of parameters of the model.
Feeding states computed by the previous step into the next step computation makes Staircase models recurrent in time like RNNs because each recurrent step moves forward C' tokens. There are two advantages this type of recurrence. First, the number of non-linear computations between an input token and output token Ut++ linearly increases with their distance т. In contrast, Transformers are strictly a feedforward model that has fixed number of computation steps. The second advantage is that information can pass to future steps without any limits, whereas standard Transformers are limited by their token truncation length. These two advantages make recurrent models capable of maintaining an internal state, but more importantly of updating this state continuously through time. However, unlike RNNs, Staircase models take advantage of the attention mechanism the same way as Transformers, and store state across multiple vectors: the NC C backward tokens. Like Transformers, they thus take advantage of parallelism.
We use the same techniques as Transformer-LL [ for processing very long or unbounded sequences. First, each token will attend fixed number of previous tokens S rather than the whole sequence. This reduces the computational complexity of the self-attention from O(T²) to (TTS) assuming < T. Next, we split the input sequence into smaller manageable segments and let the model process them sequentially. To avoid the segment boundaries from obstructing attention across segments, the hidden states computed in the previous segments are fixed and kept in memory. Then, this fixed-memory is made available in the self-attention sublayer for subsequent segments token can attend to a token in the previous segment. See Dai et al. 9 for more details about this mechanism.
3.2.1 Cached Staircase Model
In Staircase models, the self-attention sublayer processes NC tokens at time. This means how far token can directly attend to is limited by this context size NC. However, the hyperparameter N also controls the number of recurrent computations, and one might want to decouple these two factors to control context size versus recurrence.
Here we propose simple solution for increasing the context size while keeping the recurrent computation constant. We do this by introducing a new hyperparameter M N and put hidden states in cache after M recurrent steps. Figure Tp shows case where N = 4 and M = 1. Once hidden state is in the cache, stays the same, requiring no additional computation
H H for M.
This means the number of recurrent computations particular input is limited to M. However, hidden states stay in the cache for the remaining N -M steps so other tokens still can attend to them. This achieved by including cached hidden states only when computing keys and values in the self-attention sublayer of the Transformer core. As a result, the self-attention sublayer will have NC keys and values, but only MC queries, reducing its computational complexity from CN to O MC As cached hidden states are excluded from the feedforward sublayer altogether, the computational complexity there changes from O(NC)) (MC). Thus, the context size NC can be increased by picking larger N, but the amount of computation can be reduced by choosing smaller M. For example, for M = 1, we can see that the reduction computation N fold.
3.3 Relation to Existing Models
Transformer The standard Transformer corresponds to a Staircase model with a large chunk size and no recurrence. While it is efficient at processing tokens in parallel, it has no ability of retaining and recomputing state across sequences, other than by fitting those tokens into the current processing block.

---

Table I: Results summary all tasks. We compare four architectures where fix the number learnable parameters to be the same for models on the same task. * is from [[]]
Random Walk (error %)
Algorithm (error %)
Enwik8 (test bpc)
Reddit (lest ppl)
BASE (valid ppl)
Model
LSTM
0.6
1.0"
1.38
Transformer- XL
9
84.1
48.7
1.15
26.2
28.0
Feedback Transformer R3
0.1
0.2
1.12
25.5
26.6
Our models
Staircase
0.1
0.2
1.14
22.6
23.0
Cached Staircase
0.1
1.2
1.13
26.1
27.8
Random Walk
Algorithm
Enwik8
Pushshift.io Reddit
Borasd Fodlack
Toan-SL - Foilbokk
RodBack
Stancie
Oatbed Statcese
Cachel Stare on
Oabeel Shatowe
Can Steman
24-
Bocrrer
Bn
Recaavaa
Figure 2: The performance of various models on the four tasks with different numbers recurrent steps. The number of parameters is the same for all the models.
Feedback Transformer The Feedback Transformer |K is equivalent to Cached Staircase model with chunk size of C (i.e., a forward step of a single token), and caching after M = step, i.e. all tokens in the past part of a fixed cached memory. In contrast, larger chunk sizes allow general Staircase models to exploit parallelism and be far more efficient than the Feedback Transformer, while increasing M can give more expressive power. We compare this model in our experiments.
Recurrent Neural Networks RNNs 12 that store recurrent state in single vector and ingest tokens one at time can be compared to a Staircase model with single backward token and single forward token, i.e. chunk size of C = 1 and N = 2. Staircase models exploit parallelism similar to Transformers while maintaining several chunks recurrent (per token) features more expressively track state than conventional RNNs.
Memory Networks MemNets as implemented Sukhbaatar et al.
34 employ recurrence in the stacked layers of attention and computation for the current token, but only compute input embeddings hā = fin for previous tokens, and can thus be seen as a Cached Staircase with a chunk size of C 1 and caching at all previous steps, M
Transformer-XI. Transformer-XI uses Bldding-window attention where each token attends to fixed-number of previous tokens. Like Cached Staircase, it also has caching mechanism which eases computation when dealing with earlier chunks of tokens. The Staircase model has a similar sliding attention window, but its unit is a chunk of tokens instead of one. The more important difference is that Staircase models take the last state of earlier chunks and process that state further in recurrent way; Transformer-X0 on the other hand extends each layer of the Transformer's attention mechanism to using old cached states at each layer, i.e. does not build further computations top of the old state. We use this baseline in our experiments.
propose to tie all the layer weights in Transformer. In contrast Staircase models repeat multiple layers with different weights, thus allowing large number of parameters without prohibitive computational cost.
Universal Transformer Universal Transformers 10
5

---

Experiments
We types of tasks to models and compare its variations, along with Transformer- XL [9] and Feedback Transformer [3] baselines. First, we have two artificial state tracking tasks specifically designed test the model's ability to keep track of evolving changes. Next, we use real-world language modeling tasks. See Appendixx A for further details of our experimental setup for training, including all hyperparameter choices. We also make the staircase implementation public available on GitHub
4.1 Tasks
Random Walk At each time step, agent in small grid takes random action that turns the agent, or moves it forward. A model has to predict the agent's location given these actions. This seemingly simple task requires recurrency and has been shown to make feedforward models like Transformers fail. We borrow this task from Fan et al. [3], but increase the length from 100 to 400 make more challenging. See Fan al. KK] for more details about this task.
This task consists of simple algorithmic operations that need to be executed sequentially, Some operations depend on the current variable values, which makes it necessary to keep track of variable values and update them necessary. Thus, it also requires recurrency, and like the Random Walk task has been shown to make Transformers fail. We use the 3 variable version of the task from Fan et al.
Algorithm
Enwik8 Enwik8 is a character-level language modeling task [26] that consists of 100M tokens taken from Wikipedia articles. The challenging part of this data that it requires long-term reasoning 35 because tokens characters instead of words.
We use variant of Reddit discussions, which has also been used in several 20, 33 Following Humeau et al. [17. we use a previously existing Reddit dataset extracted and obtained by a third party and made available on pushshift.io R training to generate comment conditioned on the full thread leading up to the comment, spanning 1.5B training examples from Reddit obtained from Pushshift through July 2019. See Roller al. details. We concatenate the dataset view it language modeling task.
Pushshift.io Reddit
existing studies 337 22
30 for more
BASE Data We use the language modeling dataset from Lewis et 233 which consists of approximately 100B tokens, combining the corpora used in Liu et al. [24] that consists of Wikipedia, BookCorpus, CC-News, OpenWebTes and Stories, along with the English subset of the CC100 corpus
4.2 Results
Our results all of the tasks summarized in Table For each task, all the models have the same number of parameters and use the same Transformer core architecture implementation. For Random Walk and Algorithm tasks, trained each model with multiple seeds and chose the best seed as measured by their validation performance. The specific configuration of each model can be found in Table 5 in Appendix A We see clear wins from our models on all the tasks, and the following subsections will analyze these results in detail.
4.2.1 Staircase models have strong performance state tracking tasks
The Random Walk and Algorithm tasks are specifically designed to model's capability of tracking states: to store given information internally and update it with new information coming at each time step. In Table we report results from running multiple training seeds, and selecting the one with best performance on the validation set. In Figure 2 we show detailed results when varying the recurrent computation steps, reporting the mean and standard deviations amongst the seeds.
Irie et al. 18] shown that a recurrent LSTM does well on the Algorithm task with an error rate of 1 1.. Our experiment shows LSTM also solve the Random Walk task. The powerful

---

26
Trins-X XL.
Caiched
Universul
Foedback
Test (ppl)
Staircase
24
Trans XL.
Staircuse
Staircase
22
250
500
1,000
2,000
Training time (ms/batch)
Figure 3: Trade-off between training time and performance (lower ppl is better) on Pushshift.io Reddit. The circle radius is proportional to the number of parameters. The Staircase models can improve performance for longer training time without increasing the model size.
Table 2: Detailed performance on Pushshift.io Reddit. We compare our models with varying recurrent steps baselines with the same number of parameters, twice as deep (2x) baseline, and Universal Transformers with either the same layer size or much larger layer size, but same total parameters.
Model
Num. of Recurrent steps
Step size
Forward size
Valid. (ppl)
Test (ppl)
Train batch
params
time (ms)
Transformer-Xt XI.
117M
26.5
26.2
178
Transformer-X0 deep
218M
23.7
23.4
Feedback Transformer KK
102M
512
25.8
25.5
3260
Cached Staircase
117M
1
256
128
26.4
26.1
246
Staircase
117M
2
256
128
25.0
24.8
297
Universal Transformer
29M
2
39.9
39.5
51
Staircase
117M
4
256
64
23.7
23.4
580
Universal Transformer
29M
4
34,9
34.5
88
Staircase
117M
8
256
32
22.9
22.6
1147
Universal Transformer
29M
8
32.3
32,0
163
Universal Transformer
120M
8
25.9
25.6
766
Universal Transformer
29M
16
31.5
31.1
328
Transformer-X XL baseline performs poorly here due to its lack recurrent mechanism, confirming the results from Fan et al. [13]. The self-attention access hidden state far away in the past, but updating that hidden state with a new piece of information brings it up one layer higher. Thus, in Transformer with I. layers, a particular hidden state be updated only L times before it reaches the final output layer, and becomes unavailable for future computations. This limited computation depth a problem in the Random Walk task, for example, because the model needs to internally store the agent's 's location and update it with actions taken at every time step for hundreds of steps.
The Staircase model successfully solves both tasks, recurrence through time, the computation depth is only restricted by the input sequence length itself. More concretely, each recurrent step update the output from the previous step and feed it to the next step, making it possible to maintain and update an internal state without limit.
with only two recurrent steps. Thanks to its
The Cached Staircase model also performs reasonably well on those tasks. While we only ran this model with M computation step, it still recurrent time which more critical for these tasks than increased computation.

---

Table 3: Staircase model performance on Pushshift.io Reddit with different forward sizes. The step size is also changed to keep the recurrent steps constant for each section.
Model
Recurrent steps
Step size
Forward size
Valid. (ppl)
Batch
time (ms)
Cache
288
32
26.4
489
Cache St.
320
64
26.4
318
Cache St.
384
128
26.4
246
Staircase
128
64
25.2
310
Staircase
2
256
128
25.0
297
Staircase
128
32
23.8
605
Staircase
4 256
64
23.7
580
The Feedback Transformer solves both tasks, which is not surprising as it is particular case of Cached Staircase model with forward step C However, such fine-grained steps make it slow to train in practice because of the reduced parallelism, will see in the analysis in the next section.
4.2.2
Staircase models outperform Transformers for the same number of parameters on language modeling tasks
Table shows results on the three language modeling tasks, Enwik8, Pushshift.io Reddit and BASE Data. We show performance versus recurrence plots for the first two tasks in particular in Figure 2 also show more detailed performance numbers on the Pushshift.io Reddit task in Table In all three tasks, one general trend is that more recurrent steps improve the performance significantly. On the
We
Pushshift.io Reddit task, see --4 perplexity point improvement over the Transformer-X XL baseline without adding any new parameters when using recurrent steps, and a ~~5 point improvement on BASE Data. Making twice as deep Transformer-XI. (marked with "2x" in Table 2) improves the baseline at the expense of having twice as many parameters than the Staircase models, but still ~ perplexity point worse, showing the power of recurrent models. Our models provides new way of improving model performance without increasing the number of parameters that is generally applicable.
On Enwik8, we saw smaller improvement with Staircase model over a Transformer-X This could be due to the long context requirement of the character-level data of Enwik8. The Staircase model tries to compress past context into fixed number of hidden states, equal NC- C' backward tokens to be precise. The Cached-Staircase model works better on this dataset, and only 0. bpp behind the Feedback Transformer despite being much faster to train. also trained an LSTM model of baseline, but its performance was far worse at 1.38 test bpc, showing simple recurrence alone is not sufficient.
similar size on Enwik8
also shows the time it takes for training on single batch for each model. Models with more recurrent steps take longer run as they have perform more computations per token, but are still tractable and much faster than the Feedback Transformer, can be seen in Figure 3) The Feedback Transformer does not perform more computations, but it is slow because processes one token at time and also generally performs worse in our language model experiments. In contrast, the Staircase model fast because it processes NC tokens parallel despite being recurrent in time. While Cached Staircase model did not bring much performance improvement on this task, we can also see does not increase the training time, and that is because does not add any new computation when M 1. The memory usage during training increases with more recurrent steps, but inference time memory usage will stay the same because we only need keep the latest layer computation in memory.
Table 2
4.2.3 Staircase model's forward size and step size control its performance
The forward step chunk size C' and overall staircase step size NC are hyperparameters in Staircase models, where the effective number of recurrent steps is determined by those choices of N in the Staircase model, or truncated to only M steps due to caching in the Cached Staircase model. In Table 3 we compare different values of step size and forward size on the Pushshift. Reddit task for those models with differing numbers of recurrent steps. We see that, in general, the models are robust

---

to different choices of those values. Larger forward step sizes are preferable in terms of computational efficiency because they allow more parallelism, but if they are too large some performance in terms of perplexity is lost. We see evidence of this in Table annd] Z Appendix B where both Staircase and Cached Staircase models perform poorly its forward size increases.
Conclusion
In this work, we proposed Staircase attention, which re-introduce recurrence back into the family of Transformer-based models across both time and depth. We show that our Staircase model is able to solve tasks which require tracking of state that conventional Transformers cannot via its recurrence in time. It also delivers more modeling power per parameter than conventional Transformers via its recurrence in depth, thus also giving improved performance language modeling tasks for the same number of parameters, which is especially important in regimes which memory rather than compute bound. The Cached Staircase variant trades off depth-recurrenn for efficiency, but still maintains rimerrecurrnncy and do well the state-tracking tasks well character-level language modeling task. Future work should continue to investigate how recurrence can be built into sequence models, without memory component systems will always be limited only short-term reactive tasks with limited input. The approaches detailed here are one way forward and should be studied in further applications.
References
[1] Al-Rfou, R., Choe, D., Constant, N., Guo, M., and Jones, L. Character-level language modeling with d self-attention. In Pruceedings the 33rd AAAI Conference on Artificial Intelligence, 2019.
[2] Bahdanau, D., Cho, K., and Bengio, Y. Neural machine translation by jointly learning to align and translate. In ICLR, 2015.
[3] Baumgartner, J., Zannettou, S., Keegan, B., Squire, M., and Blackbum, J. The pushshift reddit dataset. arXiv preprint arXi:20001088435, 2020.
[4] Beltagy, I,, Peters, M. E.,, and Cohan, A. Longformer: The long-document transformer. arXiv preprint 25xv:2004.005550, 2020.
[5] Bengio, Y., Ducharme, R., Vincent, P., and Janvin, C. A neural probabilistic language model. The journal of machine learning research, 3:1 137-1155, 2003.
[6] Chen, M. X., Firat, O., Bapna, A., Johnson, M., Macherey, W., Foster, G., Jones, L., Parmar, N., Schuster, M., Chen, Z., et al. The best of both worlds: Combining recent advances in neural machine translation. arXiv preprint arXiv: 1804. 2018.
[7] Child, R., Gray, S., Radford, A., and Sutskever, I. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904. 2019.
[8] Conneau, A., Khandelwal, K., Goyal, N., Chaudhary, V.. Wenzek, G., Guzmán, F., Grave, E., Ott, M., Zettlemoyer, L., and Stoyanov, V. Unsupervised cross- moss-linguul representation learning at scale. arXiv preprint arXiv:1911.02111 2019.
[9] Dai, Z., Yang. Z., Yang, Y., Carbonell, J. G.. Le, Q. V., and Salakhutdinov, Transformer- xl: Attentive language models beyond fixed-length context. In ACL. (1), pp. 2978-2988. Association for Computational Linguistics, 2019.
[10] Dehghani, M., Gouws, S., Vinyals, O., Uszkoreit, J., and Kaiser, E.. Universal transformers. arXiv preprint arXiv: 1807.0389, 2018.
[11] Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 04,000 2018
[12] Elman, J. Finding structure in time. Cogn. Sci.,
1990.

---

[13] Fan, A., Lavril, T., Grave, E., Joulin, A., and Sukhbaatar, S. Addressing some limitations of transformers with feedback memory. arXiv preprint arXiv:2002.09402, 2020.
[14] Fedus, W., Zoph, B., and Shazeer, N. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. arXiv preprint arXiv:2101.0396/ 2021.
[15] Hao, J., Wang, X., Yang. B., Wang. L.,, Zhang, J., and Tu, Z. Modeling recurrence for transformer. arXiv preprint arXiv: 1904.0309,, 2019.
[16] Hochreiter, S. and Schmidhuber, J. Long short-term memory. Neural computation, 9(8):
1785-1780, 1997.
[17] Humeau, S., Shuster, K., Lachaux, M., and Weston, J. Poly-encoders: Architectures and pre-training strategies for fast and accurate multi-sentence scoring. In Proceedings of the International Conference on Learning Representations, 2019.
[18] Irie, K., Schlag, I., Csord' as, R., and Schmidhuber, J. Going beyond linear transformers with recurrent fast weight programmers. ArXiv, abs/2106.06293 2021.
[19] Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are rnns: Fast autoregres- sive transformers with linear attention. ArXiv, abs/2006.162366 2020.
[20] Keskar, N. S., McCann, B., Varshney, L. R., Xiong, C., and Socher, R. CTRL: A conditional transformer language model for controllable generation. arXiv preprint arXiv: 1909.05858, 2019.
[21] Kitaev, N., Kaiser, L., and Levskaya, A. Reformer: The efficient transformer. In International Conference on Learning Representations, 2019.
[22] Lan, 7.,, Chen, M., Goodman, S., Gimpel, K., Sharma, P., and Soricut, R. Albert: A lite bert for self-supervised learning of language representations. arXiv preprint arXiv: /909. 1942, 2019.
[23] Lewis, M., Bhosale, S., Dettmers, T., Goyal, N., and Zettlemoyer, L. Base layers: Simplifying training of large, sparse models. arXiv preprint arXiv:2003.16716 2021.
[24] Liu, Y., Out, M., Goyal, N., Du, Joshi, M., Chen, D., Levy, O., Lewis, M., Zettlemoyer, L., and Stoyanov, V. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: 11692 2019.
[25] Lu, Y., Li, Z., He, D., Sun, Z., Dong, B., Qin, T., Wang, L., and Liu, T.-Y. Understanding and improving transformer from multi-particle dynamic system point of view. arXiv preprint arXiv:/906.00 02762, 2019.
[26] Mahoney, M. Large text compression benchmark. URL: http:///ww. mattmahoney. nethtext/text. html, 2011.
[27] Mazaré, P.-.,, Humeau, S., Raison, M., and Bordes, A. Training millions of personalized dialogue agents. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 1775-2779, Brussels, Belgium, October-Novemble 2018. Association for Computational Linguistics.
[28] Mikolov, T., Karafiát, M., Burget, L., Černocký, J., and Khudanpur, Recurrent neural network based language model. In Eleventh annual conference of the international speech communication association, 2010.
[29] Press, O., Smith, N. A., and Levy, O. Improving transformer models by reordering their sublayers. arXiv preprint arXiv: 191 1911.91814,,.4.. 2019.
[30] Roller, S., Dinan, E., Goyal, N., Ju. D., Williamson, M., Liu, Y., Xu, J.. Ott, M., Shus- ter. K., Smith, E. M., et al. Recipes for building an open-domain chatbot. arXiv preprint arXiv:2004.1366 2020.
[31] Schlag, I., Irie, K., and Schmidhuber, Linear transformers are secretly fast weight program- mers. In ICML, 2021.
10

---

[32] Shaw, P., Uszkoreit, J.. and Vaswani, A. Self-attention with relative position representations. In NAACL-HLT (2), 2018.
[33] Shuster, K., Ju, D., Roller, S., Dinan, E., Boureau, Y.-L., and Weston, J. The dialogue dodecathlon: Open-domain knowledge and image grounded conversational agents, 2019.
[34] Sukhbaatar, S., Szlam, A. D., Weston, J., and Fergus, R. End-to-end memory networks. In NIPS, 2015.
[35] Sukhbaatar, S., Grave, É., Bojanowski, P., and Joulin, A. Adaptive attention span in transformers. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 331-335, 2019.
[36] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.. Gomez, A. N.. Kaiser, L., and Polosukhin, Attention is you need. In Advances in neural information processing systems, pp. (1.55999666 6008, 2017.
[37] Yang, Y., Yuan, S., Cer, D., Kong, S.-y., Constant, N., Pilar, P., Ge, H., Sung, Y.-.., Strope, B., and Kurzweil, R. Learning semantic textual similarity from conversations. In Proceedings of The Third Workshop on Representation Learning for NLP, pp. 164-174, Melbourne, Australia, July 2018. Association for Computational Linguistics.


Reviews:
[{"review_id": "L20BvTUnz5", "reviewer": "Reviewer_Cujd", "summary": "The paper builds upon the Feedback Transformer (Fan et al., 2020) to introduce recurrence into the Transformer architecture. While a regular transformer processes C tokens in parallel, the proposed approach - staircase attention - processes C new tokens at each time step and NC-C tokens from prior time steps. Thus, the transformer core moves ahead in time and resembles a staircase. The output for a given token is generated after N passes. The approach is compute-heavy compared to a regular transformer but uses the same number of parameters. To save computation, the paper proposes a variant that caches states to decrease recurrence while keeping the context size fixed. The paper shows that staircase attention outperforms transformer variants (transformer XL) on state tracking tasks (Random walk and Algorithm) as well as language modeling (Enwik8, Reddit, BASE). The staircase attention obtains a similar performance as the Feedback transformer on the state tracking tasks and improves upon it on some language models tasks (Reddit/BASE). \n\n", "questions": "* Although the cached version of the staircase attention is similar to the feedback transformer, they do not perform similarly on Algorithm and BASE. It would be useful to provide more analysis to understand why this is the case.\n* Why does validation error increase when going from 6 to 8 recurrence steps for the 'Algorithm' task?\n* Enwik8: \"The Cached-Staircase model works better on this dataset, and only 0.1bpc behind the Feedback Transformer despite being much faster to train. \". It'd be good to understand why Cached-Staircase performs better than standard staircase for this task.\n* Eqn 2: Please mention that TC stands for 'transformer core' when this acronym is used for the first time.\n\nTypos:\n* Figure 1: 'indentity' -> 'identity'\n* Sec 2: 'non-recurrent models have superseded in popularity recurrent' -> 'non-recurrent models have superseded recurrent' or 'non-recurrent models have become more popular than recurrent'\n\n", "limitations": "Yes, the authors have addressed limitations. They have not discussed potential negative societal impact of their work.", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* Novel idea to introduce recurrence into a transformer architecture\n* Architecture allows a control between amount of recurrence and context size\n* Proposes a cached version which controls the amount of recurrence and provides speeds at the cost of performance\n* Demonstrates improvements over Feedback transformer in terms of training speed\n\nWeaknesses:\n* It would be useful to add results from standard transformer and LSTM for all state-tracking/LM tasks in Table 1. Such a comparison would be useful to all researchers.\n* It would be useful to provide some more analysis (see below).\n* Given that staircase attention is compute heavy, it would be useful to see a comparison of the inference speed between all models for say one of the tasks (e.g. Reddit). This would enable researchers to understand the trade-off between performance gains vs added compute. \n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "EAlZyjAe_2J", "reviewer": "Reviewer_FFvF", "summary": "This paper introduces recurrences into transformers with staircase attention. Most existing transformer models process the entire input sequence through multiple layers. Staircase attention departs from this. It segments the input sequence into multiple chunks. As the model processes the input vertically (i.e., feeding into upper layers), it “slides forward” along the input, receiving a new chunk of input while dropping the oldest one. The neural architecture used to process the input chunks is the same as the transformer. To improve the model’s efficiency and flexibility, a cached variant is proposed, disabling the gradient updates of certain (older) input chunks, thereby reducing the computational cost. The proposed models are evaluated on multiple sequence modeling benchmarks, including variants of language modeling tasks and others. Results show that staircase attention achieves a better tradeoff between modeling sizes (in terms of the number of parameters) and performance and that the recurrence it introduces is crucial in certain tasks.\n\nThe paper is interesting and novel to the best of my knowledge, and it is clearly written. Most of the claims are reasonable, except that some important experiment details are missing. This is the reason why my initial review is a bit reserved. I’m happy to revise it if the authors can clarify some details (see below in the question section).\n", "questions": "- Could the authors clarify the connections to sliding-window attention?\n- For the “random walk” and “algorithm” tasks, which representation is fed into the classification layer?\n- Have the authors compared to a staircase baseline ablating the recurrences? For example, such a baseline for a “4-layer staircase transformer with N=8” would be, a canonical 32-layer transformer model, but ties the parameters of layers 1–8, 9–16, 17–24, and 25–32. This baseline has more computational overhead, but is important to establish the claims on the benefits of recurrences.  \n- In the language modeling experiments, say we have a staircase model with C=1. Is it the case that the model always has the access to N-1 preceding tokens when making a prediction? This relates to the “context size” in [7], and has a substantial impact on transformer language models’ performance. It would be great if the authors can clarify and make sure the models are comparable.\n\n[7] https://arxiv.org/abs/1809.10853", "limitations": "The paper did not explicitly discuss the limitations or the potential negative societal impact.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- Staircase attention is interesting and novel to my knowledge.\n- The experiments cover a diverse set of sequence modeling tasks, and have a detailed analysis.\n- Well-written.\n\nWeaknesses:\n- Details on the context sizes need to be clarified (see below)\n- The claim on the benefits of recurrences could be better supported (see below).\n- Missing discussion of several recent recurrences <-> transformer works [1,2,3,4,5,6].\n- Efficiency discussion should be complemented by a discussion on memory overhead, time overhead of the entire training procedure, inference efficiency, etc.\n\n[1] https://arxiv.org/abs/2006.16236\n[2] https://arxiv.org/abs/2103.02143\n[3] https://arxiv.org/abs/2102.11174\n[4] https://arxiv.org/abs/2106.06295\n[5] https://aclanthology.org/2021.emnlp-main.602/\n[6] https://arxiv.org/abs/2110.02488", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "7mH6gZigCno", "reviewer": "Reviewer_ZfSh", "summary": "In this paper, authors propose a new attention algorithm to operate the input sequence recurrently. Each time, the staircase attention process the seen and unseen tokens at the same time. They conduct several experiments on various tasks including ones which can be done by the standard Transformer. For other tasks, their method try to obtain a better efficiency vs performance tradeoff.", "questions": "Please refer to the above section.", "limitations": "Yes.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n1. Their idea is easy to follow and their motivation is clear.\n2. They conduct a list of experiments to show the performance of their models and compare to other models.\n3. Comparison to relevant models is sufficient and clear.\n\nWeakness:\n1. My major concern is the tradeoff between the standard staircase attention and cased one. It seems that staircase attention is able to obtain much better performance than cached attention but with low efficiency. For the cached one, it can not achieve great performance compared to other models such as feedback transformer. Can you find a better tradeoff between cache and original staircase attention.\n2. Will you release the code? It seems not straightforward to implement the proposed method.\n3. Is that possible to try various recurrent steps for cached staircase (or do I miss something here)? Also ablation study between M and N for cached staircase attention.", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "445A9OD8zBw", "reviewer": "Reviewer_3i3K", "summary": "The paper falls into a set of efforts that incorporate recurrence mechanisms in Transformer-based models while keeping the model reasonably efficient. While Transformers have been outperforming recurrent models in a number of practical tasks, some prior works have highlighted limitations of Transformer's feedforward-like architecture on tasks that involve state-tracking (among others).\n \nSeveral prior works have tried to integrate recurrence in a Transformer-like architecture to mitigate such limitations but most of them end up being computationally inefficient compared to standard Transformers. One such approach was Feedback Transformer (FeedT) which this paper builds on. Although feedback Transformer reportedly performed better than Transformer across different tasks, it suffered from efficiency issues. This paper takes a step towards mitigating the efficiency issues while preserving some of the benefits of recurrence.\n \nThey propose a modified version of attention architecture namely Staircase attention which recurrently processes different chunks/blocks of tokens and processes each chunk in a parallel manner. The computation for each chunk depends on the state of some previous chunks which makes the model recurrent. Compared to FeedT which incrementally processes each token (like an RNN) this reduces the computational cost by a factor of the size of the chunks. Additionally, each chunk can go through multiple recurrent computations. They also introduce a cached version of the model where the number of recurrent computations for each chunk is limited to reduce computational costs. They compare their approach against several Transformer variants and report the performance on 3 language modeling and 2 synthetic algorithmic tasks. Their results indicate that Staircase models achieve near-perfect accuracy in synthetic tasks while performing well on language modeling tasks.", "questions": "(Q1) Is there any limitation of processing chunks compared to the standard recurrent model? Maybe when the chunk size is large? I do not think there would be any limitation in terms of expressiveness of the model until recurrence is removed but there could be differences in performance as the chunk size grows larger. \n \n \n(Q2) Table 2: Are the recurrent steps in universal transformers and the Staircase model comparable?\n \n(Q3) Is there any relation between the proposed model to the Block-recurrent Transformers method [1]? On a high level, there seems to be some similarity in terms of processing blocks of tokens in parallel and recurrently processing blocks.  \n \n \n[1] BLOCK-RECURRENT TRANSFORMERS. https://arxiv.org/pdf/2203.07852.pdf\n \n \n \nTypos:\n \nLine 55\nFigure 1: Indentity -> Identity?", "limitations": "Not clear where limitations are explicitly discussed in the paper.", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "\n \n### Strengths\n \n(S1) **Simplicity.** The proposed idea seems simple and does not seem to make tedious modifications to standard Transformers. From the results, it seems to achieve the advantages of both recurrence (strong performance on state tracking tasks) and attention (in terms of preserving efficiency) while obtaining competitive performance on practical language modeling tasks. \n \n(S2) **Good results.** The experimental setup and evaluation seem sound. In particular, relevant baselines such as universal Transformers and Transformers-XL have been included apart from direct ones such as feedback Transformers. The results indicating slightly better performance on LM tasks with a similar number of parameters and training time look promising. \n \n \n(S3) **Writing.** The paper is well written and the general narrative of the paper is easy to follow. The problem is well-motivated and the experiments/results are well-explained. However, there is some lack of clarity about the details of the model (more on that below). \n \n \n  \n### Weakness\n  \n \n\n(W1) **Evaluation.** The set of tasks seems quite narrow and it seems like a lot more could have been done to support the efficacy of the proposed approach. The method is evaluated on 2 synthetic tasks and 3 LM tasks. Although the results on those tasks look promising it does not seem to provide a broad picture of the effectiveness of the method. The improvements in LM tasks are also incremental which is not an issue if the method consistently performs competitively across more tasks. I believe a comprehensive analysis across different tasks and across different aspects of the model is important to judge its effectiveness.\n \n**[Possible Improvements]**\n\nTasks. Since the method builds on feedback Transformers I was expecting most tasks explored in that paper to be naturally a part of the experiments here. For instance, copy, reverse, and counting should be easy to implement and experiment with. There are formal language tasks where standard Transformers have also been found to struggle compared to RNNs. If the method can perform well on those while maintaining efficiency it could strengthen the case. On the practical side, long document datasets such as GitHub, arxiv, and PG19 datasets could be explored. Positive results on some of the above-mentioned tasks could significantly help strengthen the case for the proposed method.\n \nModel Analysis. The effect of varying parameters such as chunk size and recurrence across depth is unclear. For algorithmic tasks, it could be easier to study the effect of increasing the chunk size. I would guess as the chunk size becomes very large, the model would get closer to standard Transformers and its performance might decrease. Similarly, it could be helpful to highlight if and where recurrence across depth leads to any advantage.\n \n(W2) **Novelty.** The contribution in terms of algorithmic novelty is weak in my opinion. Given the previous approaches, it seems like an extension to reduce computational costs. The recurrence across depth seems different but it is not clear if it leads to improvements. In my opinion, it is not necessarily an issue if the method consistently achieves strong performance across several tasks.\n\n\nOther Comments.\n \n(W3) **Method description.** The main method (Staircase attention model) is not very clearly described. It seems to be described in an informal and conceptual manner. Although the method is clear up to a certain degree, I am not confident if I were to reimplement it based on my understanding then it would be the same as the author's version. \n[Suggestion] It would be helpful to the readers if a very detailed description of the model is provided in the appendix. \n \n\n\n ", "was_revised": false, "final_justification": "", "rebuttal": "", "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": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept