E NCODING RECURRENCE INTO T RANSFORMERS
A BSTRACT
This paper novelly breaks down with ignorable loss an RNN layer into a sequence of simple RNNs, each of which can be further rewritten into a lightweight positional encoding matrix of a self-attention, named the Recurrence Encoding Matrix (REM).
Thus, recurrent dynamics introduced by the RNN layer can be encapsulated into the positional encodings of a multihead self-attention, and this makes it possible to seamlessly incorporate these recurrent dynamics into a Transformer, leading to a new module, Self-Attention with Recurrence (RSA).
The proposed module can leverage the recurrent inductive bias of REMs to achieve a better sample efficiency than its corresponding baseline Transformer, while the self-attention is used to model the remaining non-recurrent signals.
The relative proportions of these two components are controlled by a data-driven gated mechanism, and the effectiveness of RSA modules are demonstrated by four sequential learning tasks.
I NTRODUCTION
Sequential data modeling is an important topic in machine learning, and the recurrent networks such as LSTM (Hochreiter & Schmidhuber, 1997) and GRU (Chung et al., 2014) have served as the benchmarks in this area over a long period of time.
The success mainly contributes to the variety of recurrent dynamics introduced by these models, referred to as the recurrent inductive bias.
More specifically, the dependence between any two inputs can be described by a parametric form, which heavily depends on their relative temporal locations.
However, the recurrent models are well known to suffer from two drawbacks.
The first one is the gradient vanishing problem (Hochreiter et al., 2001), i.e. the recurrent models have difficulty in depicting the possibly high correlation between distant inputs.
This problem cannot be solved fundamentally by the recurrent models themselves, although it can be alleviated to some extent, say by introducing long memory patterns (Zhao et al., 2020).
Secondly, the sequential nature renders these models difficult to be trained in parallel (Vaswani et al., 2017).
In practice, many techniques were proposed to improve the computational efficiency of recurrent models, while they all come with compromises (Luo et al., 2020; Lei et al., 2017).
In recent years, Transformers (Vaswani et al., 2017) have been revolutionizing the field of natural language processing by achieving the state-of-the-art performance on a wide range of tasks, such as language modeling (Kenton & Toutanova, 2019), machine translation (Dai et al., 2019) and text summarization (Liu & Lapata, 2019), etc.
They have also demonstrated great potential in other types of sequence learning problems, for instance, time series forecasting (Zhou et al., 2021; Li et al., 2019).
The success of Transformers is due to the fact that the similarity between any two tokens is well taken into account (Vaswani et al., 2017) , and hence they can model long range dependence effortlessly.
Moreover, contrary to the recurrent models, the self-attention mechanism in Transformers is feed-forward in nature, and thus can be computed in parallel on the GPU infrastructure (Vaswani et al., 2017).
However, the flexibility also leads to sample inefficiency in training a Transformer, i.e. much more samples will be needed to guarantee good generalization ability (d’Ascoli et al., 2021).
Moreover, the chronological orders are usually ignored by Transformers since they are time-invariant, and some additional efforts, in the form of positional encoding, will be required to further aggregate the temporal information (Shaw et al., 2018; Vaswani et al., 2017; Dai et al., 2019).
In short, both recurrent and Transformer models have the pros and cons in modeling sequential data.
On one hand, due to inductive bias, the recurrent models excel at capturing the recurrent patterns even with relatively small sample sizes; see Figure 1(a).
Meanwhile, sample size is the performance bottleneck for the Transformer models and, when there are sufficient samples, they are supposed to be able to depict any recurrent or non-recurrent patterns in the data; see Figure 1(b).
On the other hand, sequential data have recurrent patterns more or less, and Transformers may have an improved performance if the recurrent model can be involved to handle these patterns, especially when the sample size is relatively small.
Specifically, if the recurrent and non-recurrent components are separable, then one can apply a parsimonious recurrent model on the recurrent component and a
Transformer on the non-recurrent one.
As a result, the sample efficiency can be improved comparing to the Transformer-only baseline; see illustration in Figure 1(c).
There have been various attempts in the literature to combine the two models.
Some earlier works were to simply stack them together in a straightforward manner.
Chen et al.
(2018) mixed and matched a Transformer’s encoder with an recurrent-based decoder.
Hao et al.
(2019) introduced an additional recurrent encoder to a Transformer, while Wang et al.
(2019) stacked a recurrent layer prior to the multihead self-attention.
These proposals inherit both the aforementioned shortcomings of Transformer and recurrent models.
In particular, for a very long input sequence, the sequential operation in the recurrent layers become extremely expensive.
Recent efforts have been spent on integrating recurrence and self-attention systematically.
Feedback
Transformer (Fan et al., 2021) introduces the memory vectors to aggregate information across layers, and uses them to update the next token in a recursive manner.
However, the computationally expensive sequential operation limits its attractiveness.
Another line of research applies the recurrent operation only to aggregate the temporal information at a coarser scale, while the token-by-token dependence is learned by self-attention instead.
Transformer-XL (Dai et al., 2019) partitions the long inputs into segments and introduces a segment-level recurrence.
Meanwhile, Temporal Latent Bottleneck (TLB) (Didolkar et al., 2022) and Block-Recurrent Transformer (BRT) (Hutchins et al., 2022) further divide the segments into smaller chunks, and each chunk is summarized into a few state vectors.
A recurrent relation is then formed on the sequence of state vectors.
These hierarchical designs are useful to reduce the computational burden, while they overlook recurrent dynamics at a finer scale.
In an attempt to simplify the numerical calculation of RNNs, we found surprisingly that an RNN layer with linear activation can be broken down into a series of simple RNNs with scalar hidden coefficients.
Each simple RNN induces a distinct recurrent pattern, and their combination forms the recurrent dynamics of the RNN layer.
Hence the calculation time can be greatly reduced by training these simple RNNs in parallel.
On top of that, it can be equivalently rewritten into the positional encodings of a multihead self-attention (MHSA).
This spontaneously inspires a solution, the multihead Self-Attention with Recurrence (RSA), to combine self-attention with RNN into one single operation while maintaining parallel computation.
This solution enables our design to preserve the merits from both Transformer and recurrent models, while their respective shortcomings are avoided.
More importantly, it can be used to replace the self-attention of existing networks, such as
Transformer XL, TLB and BRT, to further explore recurrent dynamics at the finer scale.
Our paper makes three main contributions below.
With ignorable approximation loss, we demonstrate that an RNN layer with linear activation is equivalent to a multihead self-attention (MHSA); see Figure 2.
Specifically, each attention
head can be used to recover a type of recurrent dependence structure, and multiple heads working in sync replicate the complex temporal dynamics of the RNN layer.
In the above MHSA, recurrent dynamics of the RNN layer are encapsulated entirely into the positional encodings, which we name the Recurrence Encoding Matrices (REMs).
This makes it possible to add the REMs to any existing Transformer architecture that has the self-attention mechanism, leading to the Self-Attention with Recurrence (RSA) module; see Figure 1(d).
A gated mechanism is used to control the proportion of REMs
and the gate value may also be interpreted as the proportion of the recurrent signals in the data.
Our experiments on four sequential tasks demonstrate that the proposed RSA module can effectively enhance the baseline Transformers’ forecasting power.
Moreover, from the fitted gate values, we have an interesting finding that time series data have stronger recurrent signals than those of the regular languages, while the code or natural languages have much weaker recurrent signals
1. O THER RELATED WORKS
Relative positional encoding The proposed REMs can provide complementary positional encoding in the RSA module.
Essentially, its recurrent pattern conveys relative location information.
This property is inherited from the distance-awareness of the RNN.
Compared to the relative positional encoding (RPE) in the existing literature, it is more parsimonious and interpretable.
The fully learnable RPE proposed by Shaw et al.
(2018) requires O ( T 2 d ) parameters, where T is the number of tokens and d is the token dimension.
Later, different variants are introduced by Huang et al.
and Dai et al.
(2019), which reduce the number of parameters to O
( Td ) and O ( d ) , respectively.
In contrast, our REMs have parametric forms of merely one or two parameters.
Time series forecasting Recent literature sees a surge in deep learning models designed specially for time series forecasting tasks, which are not Transformer-based.
For instance, DLinear (Zeng et al., 2022) and Autoformer (Wu et al., 2021) decompose the series into trend and non-trend components.
Autoformer and SCINet (Liu et al., 2021) further exploit the temporal dependence structure of time series via the autocorrelation function and multi-level convolution operations, respectively.
FiLM (Zhou et al., 2022a) and FEDformer (Zhou et al., 2022b) transform the series to the frequency domain to capture meaningful signals while removing noises.
This paper, however, differs from these literature in both target and method.
We aim to add the recurrent inductive bias to general-purpose
Transformers, and hence its performance rely heavily on the baseline Transformer that it modifies.
R ELATIONSHIP BETWEEN RNN AND MULTIHEAD SELF - ATTENTION
This section demonstrates that, without loss of much generality, an RNN layer can be approximated by a series of simple RNNs with scalar (hidden) coefficients, which can be further represented in the form of a multihead self-attention.
2. B REAKING DOWN AN RNN LAYER
Consider an RNN layer with the input variables { x t ∈ R d in , 1 ≤ t ≤ T } , and it has the form of h t
=
g ( W h h t − 1 + W x x t + b ) , where g ( · ) is the activation function, h t ∈ R
d is the output or hidden variable with h 0 = 0 , b ∈ R d is the bias term, W h ∈ R d × d and W x ∈ R
d
×
d in are weights.
When the activation function is linear, i.e. g ( x ) = x , the RNN becomes
[Equation]
equivalently
[Equation]
where the bias term b is suppressed for simplicity.
Although it has a feedforward form, the RNN cannot be trained in parallel, and this is mainly caused by the power j of recurrent weights W h , where 1 ≤ j ≤ t − 1 .
This section makes an effort to block diagonalize W h such that the RNN at (1) can be broken down into a sequence of simple RNNs with scalar (hidden) coefficients.
Lemma 1 (Theorem 1 in Hartfiel (1995)) .
Real matrices with R distinct nonzero eigenvalues are dense in the set of all d × d real matrices with rank at most R , where 0 < R ≤ d .
Suppose that the weight matrix W h has rank R ≤ d .
By Lemma 1, without loss of much generality, we can assume that the nonzero eigenvalues of W h are all distinct.
Specifically, W h has r real nonzero eigenvalues λ 1 , . . .
, λ r , and s pairs of complex nonzero eigenvalues λ r +1 , . . .
, λ r +2 s , where ( λ r +2 k − 1 , λ r +2 k )
= ( γ k e iθ k , γ k e − iθ k ) for 1 ≤ k ≤ s , i is the imaginary unit, and R = r + 2 s .
As a result, we have the real Jordan form, W h =
BJB − 1 , where B ∈ R d × d is invertible and
J ∈ R
d × d is a block diagonal matrix.
It holds that W jh = BJ j B − 1 for all j ≥ 1 , and we can then break down the recurrence induced by W h into that of the p × p block matrices in J with p = 1 or 2.
Similar to (1), we define three types of RNNs with linear activation below,
[Equation]
and h C2 t ( γ, θ ) =
(cid:80) t − 1 j =1 γ j sin( jθ ) W C2 x x t − j , where the first one corresponds to the real eigenvalues, i.e. the 1 × 1 block matrices in J , while the last two correspond to the complex eigenvalues, i.e. the 2 × 2 block matrices in J .
Note that each of the three RNNs has the recurrent weights of λ or ( γ, θ ) , and its form with a nonlinear activation function is given in the Appendix.
Proposition 1.
Suppose that W h has rank R = r
+ 2 s ≤ d , and its eigenvalues are defined above.
Let h 0 ,t
= W x x t , and then the RNN with linear activation at (1) can be equivalently rewritten into
[Equation]
2. A N EQUIVALENT MHSA REPRESENTATION
Consider the RNN of { h R t } , and let X = ( x 1 , . .
, x T ) ′
∈ R T ×
d in be an input matrix consisting of T tokens with dimension d in , where the transpose of a matrix A is denoted by A ′ throughout this paper.
We first give the value matrix V by projecting X with a linear transformation, i.e. V = XW V with
W V = W R ′ x ∈ R d in × d , and the relative positional encoding matrix is set to
[Equation]
where f t ( λ ) =
λ t for 1 ≤ t ≤ T − 1 .
As a result, the first RNN at (2) can be represented into a self-attention (SA) form,
[Equation]
We call P Rmask the recurrence encoding matrix (REM) as it summarizes all the recurrence in { h R t } .
For the RNN of { h Ci t } , i = 1 or 2 , the REM is denoted by P Cimask ( γ, θ ) , which has the form of (3) with f t ( λ ) being replaced by f t ( γ, θ ) =
γ t cos( tθ ) , i = 1 or f t ( γ, θ ) =
γ t sin( tθ ) , i = 2 for 1 ≤ t ≤ T − 1 .
And the value matrix has the form of V = XW V with W V =
W Ci ′ x ∈ R d in × d .
Thus, these two RNNs at (2) can also be represented into SA forms,
[Equation]
The remaining term in Proposition 1, h 0 ,t depends on x t only, and there is no inter-dependence involved.
Mathematically, we can represent it into a SA with the identity relative positional encoding matrix.
Finally, all the query and key matrices Q and K are set to zero in the above reformulations.
Proposition 2.
If the conditions of Proposition 1 hold, then the RNN with linear activation at (1) can be represented into a multihead self-attention (MHSA) with r + 2 s + 1 heads, where the query and key matrices are zero, and relative positional encoding matrices are { P Rmask ( λ k ) , 1 ≤ k ≤ r } , { P C1 mask ( γ k , θ k ) , P C2 mask ( γ k , θ k ) , 1 ≤ k ≤ s } and an identity matrix, respectively.
The three simple RNNs at (2) provide different temporal decay patterns: h R t provides the regular exponential decay induced by the real eigenvalues λ k ’s, and { h C1 t , h C2 t } provide the cyclical damped cosine or sine decay induced by the pair of complex eigenvalues ( γ k e iθ k , γ k e − iθ k ) ’s. These temporal decay patterns are further summarized into the regular REM, P Rmask ( λ ) , and cyclical REMs,
P C1mask ( γ, θ ) and P C2mask ( γ, θ ) , respectively; see Figure 2 for the illustration.
From Proposition 2, the combination of these three types of patterns forms the recurrent dynamics of the RNN layer at (1).
For each head, the REM has one or two parameters, and W V can be regarded as one d × d learnable matrix.
This leads to a parameter complexity of O ( Rd 2 ) , and it is slightly larger than that of the
RNN at (1) since R is usually much smaller than d (Prabhavalkar et al., 2016).
Moreover, the MHSA representation in Proposition 2 gives us a chance to make use of parallel matrix calculation on the
GPU hardware; see Appendix D.3 for an illustration of the computational efficiency of REMs.
E NCODING RECURRENCE INTO SELF - ATTENTION
While the query and key matrices are set to zero in the MHSA representation at Proposition 2, they play a central role in a standard Transformer.
This motivates us to propose the Self-Attention with
Recurrence (RSA) module to seamlessly combine the strengths of RNNs and Transformers:
[Equation]
for each head, where P is a regular or cyclical REM, and σ ( µ ) ∈ [0 , 1] is a gate with σ being the sigmoid function and µ being the learnable gate-control parameter.
Figure 1(d) provides a graphical illustration of one RSA head.
Note that the REMs in Section 2 are all lower triangular matrices, which correspond to unidirectional
RNNs.
In the meanwhile, for non-causal sequential learning tasks (Graves et al., 2005; Zhou et al., 2016), bidirectional RNNs are usually applied, and accordingly we can define the unmasked versions of REMs.
Specifically, the regular REM is P Runmask ( λ ) =
P Rmask ( λ ) + [ P Rmask ( λ )] ′ , and the cyclical
REMs are P Ciunmask ( γ, θ ) =
P Cimask ( γ, θ ) + [ P Cimask ( γ, θ )] ′ with i = 1 and 2.
In practice, these
REMs will be explosive when | λ | > 1 or | γ | > 1 .
To avoid this problem, we further bound these two parameters by transformations in this paper, i.e. λ = tanh( η ) and γ = sigmoid ( ν ) , and these notations then become P R k ( η ) ,
P C1 k
( ν, θ ) and P C2 k
( ν, θ ) with k ∈ { mask , unmask } accordingly.
Sequential data usually have some recurrent patterns, which can be captured by REMs.
In the meanwhile, the remaining non-recurrent patterns can be modeled by the baseline Transformer via softmax ( QK ′ ) .
The allocation of weights between REMs and self-attention is adjusted by the learnable gate σ ( µ ) .
On one hand, comparing with the baseline Transformer, the inclusion of REMs will lead to a significant improvement of sample efficiency in modeling the recurrent patterns, and hence a higher accuracy in prediction can be expected from the proposed RSA.
On the other hand, comparing with the recurrent models, the non-recurrent patterns can be taken care of by the flexible
Transformers, and hence the representation power of the RSA module is as good as that of the baseline Transformer.
In addition, the REM can act as a complementary positional encoding scheme to further provide relative location information.
For the multihead RSA, the gate-control parameter µ only varies across layers, while the parameters controlling the matrix P vary across all heads and layers.
Initialization.
The gate-control parameter µ is initialized in the interval [ − 3 , 3] for all the layers.
For ( λ, γ, θ ) which determine the recurrent patterns, we initialize λ ’s at different heads to spread out between [ − 2 , − 1] ∪ [1 , 2] and ν ’s to spread out between [1 , 2] , and θ is initialized at π/ 4 , to encourage REMs to be non-zero and well-diversified.
Dilated REM variants.
The dilated REMs can be further obtained by considering the block matrix formed by the first T columns and first T rows of P ⊗ I d , where d is the dilating factor.
In fact, the dilated REMs can encapsulate the recurrence patterns of the dilated RNNs (Chang et al., 2017); see Proposition 4 in the Appendix.
They describe potentially periodic recurrence over a long-term timespan, and can significantly enrich the temporal dynamics of our RSA.
Hyperparameters k i and d .
In each multihead RSA layer with H heads, the number of the six types of REMs, namely one regular, two cyclical, one dilated regular and two dilated cyclical REMs, are denoted by k 1 − k 6 , respectively.
Since (cid:80) 6 i =1 k i
=
H , we can apply constrained hyperparameter search to optimize over the choices of k i ’s. For simplicity, we can set k 2 = k 3 and k 5 = k 6 , which indicates that the cyclical REMs come in pairs.
The search for d , on the other hand, can be guided by various data analytic tools.
For instance, d can be observed from the autocorrelation plots as the seasonal period length for time series data (Cryer & Chan, 2008); while for the language-related data, d can be heuristically deduced from the recurrence plots (Webber et al., 2016).
More discussions on REMs.
Each type of REMs is basically a masked or unmasked linear aggregation of tokens, and we may alternatively consider a more general Toeplitz, or even a fully learnable, matrix P at (4) such that all kinds of temporal decay patterns can be automatically captured.
Although more flexible than REMs, it will need O ( T ) or O ( T 2 ) additional parameters, where T is the sequence length, while each REM requires merely one or two parameters.
Note that the proposed RSA at (4) also includes a standard self-attention, which is flexible enough for the remaining nonlinearities that failed to be captured by REMs, and hence it may not be necessary to consider a more general yet less efficient structure for REMs.
see Appendix B for both theoretical and empirical evidences.
E XPERIMENTS
This section contains four sequential modeling tasks and, for each task, we modify some popular
Transformer baselines by adding the REMs to their attention weights via a gated mechanism as in (4).
The modified models are named with the prefix “RSA-”.
The trained gate σ ( µ ) controls the percentage of weights allocated to the REMs, and it hence provides a rough measure for recurrent signals in the data.
From Figure 3, we may argue that time series have the strongest recurrent signals, followed by the regular language (Yu, 1997) and finally the code or natural languages.
The gate values also vary across different layers of Transformer models, indicating that the proportion of recurrent signals may change as the data propagate through layers.
Moreover, the data-driven gated mechanism can help maintain the optimal allocation between REM and self-attention to achieve better sample efficiency and improved performance.
All experiments are conducted on Nvidia V100 32GB GPUs.
4. T IME SERIES FORECASTING
Time series data possess unique attributes, such as strong recurrent signals, that can improve forecasting accuracy.
Although Transformers are not currently the state-of-the-art model for this task, it is interesting to explore how the proposed RSA module could enhance their performance.
Our experiments are performed on two public benchmark datasets: the ETT 1 dataset is comprised of seven features related to the electric power long-term deployment, where {ETTh 1 , ETTh 2 } are recorded by the hour and ETTm 1 is recorded by 15-minute intervals; and the Weather 2 dataset contains twelve climate indicators collected every 1 hour over a 4-year period.
For the baseline models, we adopt Informer (Zhou et al., 2021),
LogSparse Transformer (Li et al., 2019) and Transformer-XL (Dai et al., 2019).
In particular, to accommodate the cases where the attention map is non-square, we extend the recurrent patterns and adjust the shape of our REMs; see details in Appendix D.2.
All hyperparameters in baseline models are set to the optimal setting in Zhou et al.
, and we also follow their train/val/test division and training schemes to conduct our experiments.
Table 1 summarizes the mean squared error (MSE) and mean absolute error (MAE) for the three baselines against their RSA counterparts, averaged across 5 repetitions.
Overall, the revised models show significant improvement in performance over their baselines.
Moreover, it is interesting to see that the REMs can act as effective substitutions for positional encoding.
For Transformer-XL in particular, although we remove their learnable relative positional encoding, superior performance is attained by RSA-XL while the total number of parameters is reduced.
Since more than 70% of weights are allocated to the REMs in Figure 3, it numerically verifies that the recurrent signals dominate the non-recurrent ones for time series data.
[Table]
ETT m 10.
W ea t h e r0.
[Table] [Table]
4. R EGULAR LANGUAGE LEARNING
Regular languages are intimately related to the linear recurrence sequences, such as Fibonacci numbers (Kotek & Makowsky, 2010).
Some works report that Transformers have difficulty in generalizing the rules of regular languages (Bhattamishra et al., 2020; Tran et al., 2018; Hahn, 2020).
Experiments are conducted to show that the proposed RSA module can leverage the recurrent inductive bias of the
REMs to improve Transformer’s sample efficiency, and subsequently its generalization ability.
For the experiments, we deliberately choose six types of regular languages on which the Transformer model has been shown to perform poorly, namely Parity, Tomita 3, 5 & 6 (Tomita, 1982), D 2 and
D 4 (Bhattamishra et al., 2020).
We use the decoder-only vanilla Transformer model as the baseline, and modify its self-attention layers into our RSA forms under four hyperparameter settings (detailed in Table 2).
As a comparison, another benchmark model (denoted by Transformer ⋆ ) is obtained by replacing the absolute sinusoidal positional encoding with the learnable relative positional encoding (Dai et al., 2019).
To evaluate their generalization ability, the test set is further separated into two bins, where bin 1 contains longer test samples than bin 0.
Following Bhattamishra et al.
(2020), we recast the prediction task into a classification one, where the prediction is deemed a success only when the output values at all positions are correct, and a failure otherwise.
From Table 2, it can be seen that the vanilla Transformer model fails to grasp the recurring rules of these languages.
With assistance of the learned relative positional encoding, they can learn better on D 2 , D 4 and Tomita 3, but not the others.
However, the RSA-Transformer not only achieves near 100% accuracy on D 2 , D 4 and Tomita 3, but also gains a significant boost in performance on the other three languages.
In particular, Figure 3 shows that approximately 60% of weights are allocated to the REMs, implying strong recurrent signals in the datasets.
We further provide some visualization of the REMs under Cases I – IV and their interpretaions in Appendix E.2.
[Table]
4. C ODE AND NATURAL LANGUAGE MODELING
Different from regular languages, the recurrence relationship in programming or natural languages is weaker and harder to interpret.
Nevertheless, they have a wide range of applications on which the
Transformers are serving as the baseline models.
Our experiments show that the baseline models can improve the performance with inclusion of a small percentage of REMs.
We first conduct a defect detection task based on the C language dataset provided by Zhou et al.
(2019), which is a binary classification task to evaluate whether a code is vulnerable to external attacks.
The pretrained CodeT5 (Wang et al., 2021) 3 is adapted to this downstream task.
Specifically, we use CodeT5-small as our baseline model with 12 layers and 8 heads, and directly add the REMs on top of the self-attention with fully learnable RPE via a gated mechanism.
The gate-control parameter µ is initialized at − 3 and we fine-tune both the baseline and our RSA models forepochs with early stopping.
It can be observed from the boxplot in Figure 3 that the gate values at different layers remain close to the initialization, with only 5% of weights allocated to the REMs.
Although the recurrent signals in the code language is apparently much weaker than those in time series or regular language, our RSA-CodeT5-small is still able to achieve a significant improvement against the CodeT5-small baseline as shown in Table 3(a).
The RSA-CodeT5-small model has only about 320 additional parameters over the baseline with 66 million parameters, which is equal to 5 × 10 − 4 % .
More surprisingly, it surpasses the current SOTA result of 65.78%, which is achieved by the CodeT5-base model with a total number of 220 million parameters (Wang et al., 2021).
Due to limited resources, we have not yet fine-tuned RSA-CodeT5-base model.
The natural language modeling task is conducted on three standard NLP datasets, namely the wordlevel WikiText-
(Merity et al., 2017),
Enwik8 (LLC., 2009) and Text8 (LLC., 2009).
We use two baseline models: Transformer-XL and Block Recurrent Trasnformer (BRT), and their corresponding
RSA models are formed by directly adding the REMs on top of some of the self-attention heads via a gated mechanism.
For Wikitext-103, all transformers are 16-layer with 12 heads.
In the meantime, for both Enwik8 and Text8 there are 14 layers and 8 heads.
The details of model structures can be referred to Appendix E.3.
We follow the training schemes of Transformer-XL, and the results are reported in Table 3(b).
Overall, it can be seen that RSA models can achieve better performance than their corresponding baselines.
The gated values for RSA-XL on the three datasets are presented in Figure 3, and we can see that REMs contribute more substantially for Text8 and Enwik8 than
WikiText-103, which may imply stronger recurrent signals for these two datasets.
4. A BLATION STUDIES
This section conducts three experiments to analyze how RSA scales with respect to sample size, data size and sequence length, and its performance is compared against the baseline models.
Model scaling The first ablation studies the performance of RSA-XL on Text8 dataset when the number of parameters increases.
Baseline models include Transformer-XL, and two other benchmarks suggested by a reviewer, which replace the parsimonious matrix function P in RSA (see equation (4)) with (1) a Toeplitz matrix, or (2) a fully learnable matrix.
They are incorporated into Transformer-XL baseline and are hence referred to as TL-XL and FL-XL, respectively.
To have comparable parameter counts, the number of layers in FL-XL varies from 8 to 14, whereas that in others varies from 10 to 16.
We adopt the original training schemes of Transformer-XL.
From the results presented in Figure 4 (left), it can be observed that RSA-XL is the most parameter efficient across all model scales.
See
Appendix E.4 for similar results comparing BRT and RSA-BRT on model scaling performance.
Data scaling The second ablation study compares the sample efficiency of RSA-XL against
Transformer-XL.
Specifically, we use different proportions of training data from Enwik8 to train the models, and maintain the same evaluation and test sets.
Results are presented in middle panel of Figure 4, and we further visualize the gate values σ ( ν ) allocated to REMs at every layer in the right panel.
It shows that, as available training samples shrink, learnable gate will automatically increase the weights of the parsimonious REMs, and this subsequently helps RSA-XL to maintain a consistently higher sample efficiency than the baseline XL model.
Sequence length scaling In this ablation study, we adopt a synthetic copy task akin to the one in
Nguyen et al.
(2021) or Katharopoulos et al.
It requires the model to copy a sequence of symbols with different maximum lengths.
For every sample, a sequence of ten symbols is randomly generated and repeated multiple times, with a reserved symbol separating each replication.
Four models are used: BRT, TLB, RSA-BRT and RSA-TLB, and we adopt a small training sample size of 6400 as well as a larger one of 12800.
The implementation is detailed in the appendix.
Figure 5 first depicts the accuracy with respect to the sequence length and then shows the cross entropy loss with respect to the number of gradient steps.
It can be observed that RSA-BRT and RSA-TLB converge faster than their baselines, respectively, and achieve perfect accuracy across all sequence lengths.
C ONCLUSION AND DISCUSSION
By formulating RNN into an MHSA form, we novelly propose an RSA module to incorporate the recurrent dynamics of an RNN into a Transformer.
Specifically, the lightweight REMs are combined with the self-attention weights via a gated mechanism, maintaining a parallel and efficient computation on the GPU infrastructure.
Experiments on four sequential learning tasks show that the proposed RSA module can boost the sample efficiency of the baseline Transformer, supported by significantly improved performance.
In addition, with the fitted gate values to measure the recurrent signal strength in the data, it can be observed that the code or natural languages have much smaller proportion of recurrent patterns, which further explain the superiority of Transformer models in these two data domains, while the inclusion of REMs still can boost the baseline performance.
This paper can be extended along three directions below.
Firstly, while six typical examples of
REMs are introduced, we may also consider other types of REM designs, such as the non-distinct eigenvalues of the RNN weight matrix W h , and long-memory filters operated on top of the RNNs (Zhao et al., 2020).
Secondly, the proposed RSA module can be further applied to other types of sequential learning tasks such as video learning (Bulat et al., 2021) and skeleton-based human motion detection (Shu et al., 2021).
Finally, it is of interest to incorporate the REMs into models other than
Transformers, or even to apply it directly as a standalone recurrent unit.
R EPRODUCIBILITY STATEMENT
To ensure the reproducibility and completeness of this paper, we include the Appendix with five main sections.
Appendix A discusses the relationship between RSA and some important related works.
Appendix B provides a theoretical and empirical study on the performance gap between linear and nonlinear RNNs.
Appendix C contains the complete proofs for all propositions presented in this paper.
An efficient method to compute the REMs is detailed in Appendix D, together with an extension to non-square REMs and an empirical study on the time efficiency of the proposed method.
The experiments in the paper are reproducible with additional implementation details provided in
Appendix E. We also include the hyperparameter settings for all results contained in Tables 1 and
A CKNOWLEDGMENTS
We would like to thank the three anonymous reviewers for spending time and efforts and bringing in constructive questions and suggestions, which help us greatly to improve the quality of the paper.
We would like to also thank the Program Chairs and Area Chairs for handling this paper and providing the valuable and conprehensive comments.
A PPENDIX
This Appendix contains five sections.
The first section provides an illustrative plot for multiscale recurrence in various Transformer models.
The theoretical and empirical gaps between the linear and nonlinear RNNs are studied in the second section.
The third section presents Proposition 3, and then gives the detailed proofs for Propositions 1 – 3.
An efficient method to compute the REMs is detailed in Section 4, and we further extend it to non-square REMs and conduct an empirical study on the training time per batch for the RSA models against their baselines.
Section 5 provides additional implementation details for the experiments, additional ablation studies and visualization of the learned patterns.
A M ULTISCALE RECURRENCE
For sequential learning tasks, the recurrence relationship can be well observed at different levels of temporal granularity.
This feature has inspired new designs to add recurrence to Transformers from varying scales; see illustration in Figure 6.
Transformer-XL (XL) (Dai et al., 2019) partitions a long input sequence into segments, and places them into consecutive batches.
A segment-level recurrence is then introduced by using a cached hidden state to pass historical information iteratively to the next segment.
Temporal Latent Bottleneck (TLB) (Didolkar et al., 2022) further divides the segment within one batch into smaller chunks, and then adopts the state vectors to aggregate both high-level information across layers and temporal information across chunks.
The chunk-level recurrence helps create a slow information stream in the sequence to learn a more condensed representation.
Block-Recurrent Transformer (BRT) (Hutchins et al., 2022) also establishes recurrence across chunks (or blocks), while their recurrent states are layer-specific and updated with an LSTM-style gated design.
As a comparison, the proposed RSA follows the RNN to account for the recurrence between individual inputs.
In other words, it models the token-level recurrence, which is at the most finegrained scale.
Subsequently, it can be easily incorporated into the aforementioned coarser-grained designs, and may potentially bring benefits to their performance.
For illustration, we use XL and
BRT as our baseline models in Section 4.3.
B T HEORETICAL AND EMPIRICAL GAPS BETWEEN LINEAR AND NONLINEAR RNN S B. T HEORETICAL GAP
This subsection theoretically evaluates prediction errors when a linear RNN model is used to train the data generated by a nonlinear RNN.
For simplicity, we consider 1D case only, and many-to-many
RNNs are assumed.
Specifically, the nonlinear RNN model used to generate the data has the recursive form of g t = σ h ( u t ) with u t = w ∗ h g t − 1 + w ∗ x x t + b ∗ , (5) and σ h ( · ) is the activation function satisfying
[Equation]
Note that many commonly used activation functions, including Tanh and Sigmoid, satisfy the above condition.
We further consider an additive error ε t , i.e. y t =
g t + ε t , where ε t has mean zero and a finite variance denoted by γ .
For the generated data { y t } , we train a linear RNN model,
[Equation]
where the parameters θ = ( w h , w x , b ) .
Then the mean squared prediction error can be defined as
[Equation]
and its theoretical bound is provided in the following proposition.
Proposition 3.
Suppose that E ( u 2 t ) ≤ α and E ( u 2 t u 2 s ) ≤ β for all t, s ∈ Z .
If | w ∗ h | < 1 and the condition at (6) holds, then
[Equation]
where the first part is due to the misspecification of using the linear activation to approximate σ h ( · ) .
Proof of Proposition 3.
Let θ ∗ =
( w ∗ h , w ∗ x , b ∗ ) , and denote h t = h t ( θ ∗ ) for all t ∈ Z , i.e.
[Equation]
By the definition of e pred , it holds that
[Equation]
where the equality comes from E ( ε t )
= 0 and var ( ε t )
=
γ .
By applying second-order Taylor expansion at zero and from (5) and (6), we have
[Equation]
where
[Equation]
and (cid:101) u lies between u t and zero.
This, together with (7), leads to g t − h t
=
σ h (0) + w ∗ h ( g t − 1 − h t − 1 )
+ R t (0) .
Let δ t =
g t − h t , and it then holds that
[Equation]
where the second equality is obtained by applying the first equality recursively.
As a result, by the condition that | σ h (0) | < 1 , E ( u 2 t ) ≤ α and E ( u 2 t u 2 s ) ≤ β for all t, s ∈ Z , we can show that
[Equation]
where ξ = (cid:80) ∞ j =0 | w ∗ h | j .
If | w ∗ h | < 1 , we have ξ = (1 −| w ∗ h | ) − 1 .
This, together with (8), accomplishes the proof.
B. E MPIRICAL GAP
This subsection conducts a synthetic experiment to evaluate the performance of linear RNNs when there is nonlinearity in the data.
We first generate the data by using a two-layer nonlinear RNN model with the form of
[Equation]
and
[Equation]
with i = 1 and 2, where z (0) t = x t , z ( i ) t ∈ R 2 for 1 ≤ i ≤ 3 , σ h ( · ) is a nonlinear activation function, and 0 ≤ α ≤ 1 is the weight of nonlinearity.
An additive error is further assumed, i.e. y t =
z (2) t + ε t , where { x t } and { ε t } are independent and follow the standard multivariate normal distribution.
Three nonlinear functions are considered for σ h ( · ) , Tanh, Sigmoid and ReLU.
As α increases from 0 to 1, the data generating process gradually changes from a strictly linear RNN to a nonlinear one, i.e. α essentially controls the proportion of nonlinearity involved.
The sequence { y t , 1 ≤ t ≤ T } is fitted separately by a linear RNN, a nonlinear RNN with the corresponding activation, and a linear RNN combined with a self-attention, i.e. the proposed RSA.
Specifically, we generate a sequence of length 10000 and then divide it into 100 segments, each of length 100.
In each segment, we train with the first 99 observations and calculate the prediction error for the last observation.
The Adam optimizer is adopted for training, and the training procedure will be terminated when the training loss drops by less than 10 − 5 .
The mean squared prediction error (MSPE) averaged over 100 segments are denoted by e L pred , e NL pred and e RSA pred for the three models, respectively.
Using nonlinear RNNs as the benchmark, the MSPE ratio for the linear RNN or the
RSA is defined as
[Equation]
where i ∈ { L , RSA } .
Figure 7 presents the MSPE ratios for three types of activation functions.
It can be seen that, when α = 1 , nonlinear RNNs perform the best, while linear RNN suffers from misspecification error.
Alternatively, when α = 0 , the opposite can be observed.
Moreover, as α increases, i.e. there are more nonlinearity, it is expected that linear RNNs become less favorable, while the proposed RSA can remedy the problem to some extent.
Especially when α > 0 .
6 , the RSA consistently achieves better prediction performance than the pure linear RNN.
C P ROPOSITION 3 AND PROOFS FOR ALL P ROPOSITIONS C. P ROPOSITION 3 FOR DILATED RNN S
For some positive integer d , let ¯ P R mask be the block matrix formed by the first T columns and the first
T rows of P Rmask ⊗ I d .
And ¯ P C i mask can be defined similarly for i = 1 or 2 .
Consider a dilated RNN (Chang et al., 2017) with the dilating factor d .
It has the form h t
=
g
( W h h t − d + W x x t + b ) where g ( · ) is the activation function, h t ∈ R
d is the output or hidden variable with h 0 = 0 , b ∈ R d in is the bias term, W h ∈ R d × d and W x ∈ R
d
×
d in are weights.
When the activation function is linear, i.e. g ( x ) = x , the RNN becomes
[Equation]
where the bias term b is suppressed for simplicity.
We have the following proposition.
Proposition 4.
If the conditions of Proposition 1 hold for W h in (10), then the RNN with linear activation at (10) can be represented into a multihead self-attention (MHSA) with r + 2 s +heads, where the query and key matrices are zero, and relative positional encoding matrices are { ¯ P Rmask ( λ k ) , 1 ≤ k ≤ r } , { ¯ P C1mask ( γ k , θ k ) , ¯ P C2mask ( γ k , θ k ) , 1 ≤ k ≤ s } and an identity matrix, respectively.
Proof.
The proof follows directly from Propositions 1 and 2.
C. P ROOF FOR P ROPOSITION
Let W h be a d × d real matrix with distinct eigenvalues, and from Chapter 3 of Horn & Johnson (2012), we have the Jordan decomposition, W h
=
BJB − 1 , where B ∈ R d × d is invertible, and
J ∈ R d × d has a real Jordan form, J = diag { λ 1 , . . .
, λ r , C 1 , . . .
, C s , 0 } with
[Equation]
k
[Equation]
γ k ·
[Equation]
cos
[Equation]
Then,
[Equation]
for all j ≥
[Equation]
where G R k ’s, G C1 k ’s and G C2 k ’s are d × d real matrices determined jointly by B and B − 1 .
Let h 0 ,t
= W x x t
and then,
[Equation]
where h R
[Equation]
x x
[Equation]
− j or equivalently, h R t =
[Equation]
g ( · ) being the identity function and h C 1 t
[Equation]
t
[Equation]
1 j =1 γ j sin
[Equation]
jθ ) G C2 W x x t − j which can be obtained via the recursive relationships,
[Equation]
with G being G C1 or G C2 , respectively.
For a more general form, we can further assume that the activation function g ( · ) is nonlinear for the simple RNNs at (11) and (12).
C. P ROOF FOR P ROPOSITION
Using the SA form, we can represent the three types of the RNNs by ( h R 1 ( λ k ) , · · · , h R 1 ( λ k ))
=
SA R k ( X ) and ( h C
i 1 ( γ k , θ k ) , · · · , h C i 1 ( γ k , θ k ))
=
SA C ik ( X ) for i = 1 or 2 .
Therefore, the first term in Proposition 1 can be represented as
[Equation] [Equation]
where W o =
( I d , . .
, I d )
′ ∈ R rd
×
d with I d being d -dimensional identity matrix.
Similarly the
MHSA for second and third term in Proposition 1 is given by
[Equation]
where W o =
( I d , . .
, I d ) ′ ∈ R sd × d .
And we define the additional head as SA 0 ( X ) =
( h 0 , 1 , . . .
, h 0 ,T ) ′
=
[softmax( QK ′ )
+ I ] V with
W Q = W K = 0 and W V
= W ′ x .
Combine with (13) and (14), we have
[Equation]
where W o =
( I d , . .
, I d )
′ ∈ R
( r +2 s +1)
d × d .
D C OMPUTATION OF THE REM S D. E FFICIENT COMPUTATION OF THE REM S
For any matrix A , let cos ( A ) and sin ( A ) refer to applying the cosine and sine functions to each element of A .
To compute the REMs, we first construct the exponent matrix L that has a Toeplitz form:
[Equation]

Denote the square matrix whose elements are all equal to λ as λ .
Similarly, we define γ and θ for γ and θ , respectively.
Extending the conventional power notation, we let λ L = pow ( λ , L ) and γ L = pow ( γ , L ) , i.e., ( λ L ) i,
j = λ
L i,j i,j and ( γ L ) i,j = γ
L i,j i,j .
Then we can calculate the REMs by
P = (cid:26)
[Equation]
−
[Equation]
T for regular REM ,
D C
i − I T for cyclical REMs with 1 ≤ i ≤ 2 , where D R = λ
L , D C1 =
γ L ⊙ cos ( θ ⊙ L ) , D C2 = γ L ⊙ sin ( θ ⊙ L ) and ⊙ is the Hadamard product.
Such design allows us to circumvent redundant computations and achieve parallelization as much as possible.
To avoid extremely large or small values, we mask the entries that have power larger than a preset number K .
In practice, we set K = 200 .
Moreover, we can easily extend the above calculation to dilated REMs by substituting L with L d , where ( L d ) i,j = I { L i,j mod d =0 } L i,j /d and d is the dilating factor.
The remaining operations are unchanged.
D. F LEXIBLE ADAPTATION TO DIFFERENT ATTENTION SHAPES
In some cases, the query and key have different time spans with an overlapping period; examples including the cross-attention and all attention in the Transformer-XL.
We show in this section that
our
REMs can be adapted for such cases.
Typically, the historical information will be mostly stored in the query matrix, and the new information will be included in the key and the value matrices.
Assume that the query matrix contains the tokens from time 1 to time T , and the key matrix contains that from time T − K to time T + L , where K > 0 is the overlapping period length and L > 0 is the length of new information.
Then the REM is of size ( L + K + 1) × T and takes the form of    ···
T − K −
T − K ···
T −
T
T − K f T − K − 1 ( β ) f T − K − 2 ( β ) · · · f 1 ( β ) · · · f K − 1 ( β ) f K ( β )
T − K + f T − K − 2 ( β ) f T − K − 3 ( β ) · · · f 2 ( β ) f 1 ( β ) · · · f K − 2 ( β ) f K − 1 ( β ) .
T −f T − 2 ( β ) f T − 3 ( β ) · · · f K ( β ) f K − 1 ( β ) · · · f 1 ( β )
T f T − 1 ( β ) f T − 2 ( β ) · · · f K +1 ( β ) f K ( β ) · · · f 1 ( β ) T + f T ( β ) f T − 1 ( β ) · · · f K +2 ( β ) f K +1 ( β ) · · · f 2 ( β )
f 1 ( β ) .
T + L f T + L − 1 ( β ) f T + L − 2 ( β ) · · · f K + L
+
β ) f K + L ( β ) · · · f L +1 ( β ) f L ( β )   , which maintains the recurrent patterns while considering the shift in the relative time locations, where β = λ for regular REM or ( γ, θ ) for cyclical REMs.
And we can extend to the dilated REMs similarly.
D. T HE COMPUTATIONAL EFFICIENCY OF THE REM S FOR VARIOUS T RANSFORMER BASELINES
We first conduct an empirical study on the Weather dataset to evaluate the computational time of the
REMs for RSA-Informer, RSA-LogSparse and RSA-XL.
Specifically, we record the training time per batch for the RSA models, while simultaneously we exclude the REMs from the models and record the training time again.
The difference between the two gives the additional computation time due to the REMs, as shown by the green bar in Figure 8.
It can be seen that the calculations involving the
REMs account for less than 20% of the total training time in general.
In the next experiment, we further compare the computational time between the RSA models and their respective baseline models.
Note that the original Informer and LogSparse Transformer use absolute positional encoding (APE), while Transformer-XL adopts a learned relative positional encoding (RPE).
These positional encoding schemes are excluded from the RSA models.
The training time per batch for all the models are reported in Figure 9.
We can observe a small increase in computational time for RSA-Informer and RSA-LogSparse against their baselines, while that for RSA-XL remains nearly unchanged.
This implies that the REMs requires a bit more computation effort than the APEs, and the same computation effort as the RPE introduced by Transformer-XL.
E A DDITIONAL EXPERIMENTS AND RESULTS E. T IME SERIES FORECASTING
We apply the same time series datasets as in the Informer paper (Zhou et al., 2021), and the details are given below for completeness.
The first dataset is the Electricity Transformer Temperature (ETT).
It contains the records for seven electric power indicators over a two-year period from two separated counties in China.
Zhou et al.
(2021) created three separate datasets ETTh 1 and ETTh 2 for hourly level, and ETTm 1 for the fifteen-minute level.
The train/val/test is 12/4/4 months.
The second dataset consists of climatological data over a four-year period from 2010 to 2013 for various locations in
US.
The train/val/test is 28/10/10 months.
To ensure the quality of our baselines models, we make reference to the following two github repositories, namely Informer2020 4 and transformer-xl 5 .
Since
Transformer-XL requires a sequential order for training, testing and validation sets, we generate the training samples by sliding window starting from different points of the original time series following the practice in Salinas et al.
In this task, we let the REMs substitute the original positional encoding scheme in the RSA models.
Moreover, since the distilling operation in the original Informer design is likely to distort the relative
[Table]
ETT
[Table]
ETT m
[Table]
W
[Table] [Table]
positions, we exclude it from both our Informer baseline and the RSA-Informer.
The hyperparameters in RSA are fine tuned via grid search.
Table 5 reports the fine-tuned hyperparameter settings considered in Table 1 of the main paper.
To provide additional information about the standard deviation, we provide a more detailed version of
Table 1 of the main paper, given in Table 4.
Ablation on the hyperparameters To shed light on the importance of different recurrent dynamics, we provide an ablation study with RSA-Informer on ETTh 1 .
The model has a total of 8 heads, and we set k 2 = k 3 to some positive integer s and k 5 = k 6 to some positive integer ¯ s .
In particular, when ¯ s = 0 , we also take k 4 =
and hence there are not any dilated REMs, leading to k 1 = 8 − 2 s .
Otherwise, we take k 4 = 4 − 2¯ s and k 1 = 4 − 2 s .
It can be observed from Table 6 that the inclusion of dilation and the cyclical patterns can significantly improve the forecasting performance.
The concurrence of dilated and non-dilated RSA heads allows the model to capture both the short- and long-term recurrent dynamics.
Gate-control parameter learning To study how much proportion of the RSA output is accounted for by the REMs, Figure 10 plots the change in its relative importance i.e. σ ( µ ) during training for all the attention layers, starting from different initializations.
As the learning rate drops rapidly, all curves are flattened towards the end.
The relative positional attention is the most effective for the cross-attention layer, and the least effective for the self-attention layer in the decoder, whose recurrent patterns may be disrupted by zero paddings.
We observe a decrease in MSE as the initialization increases from − 1 to 1 , which provides an empirical support for initializing at µ = 1 .
[Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
ETT m 1(
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
W ea t h e r(0,0,0,0,4,4)
[Table]
pred window =
[Table]
E. R EGULAR LANGUAGE LEARNING
The six types of regular language datasets are obtained from the github repository TransformerFormal-Language .
6 We briefly describe the rules of these languages:
• The Parity language contains the numbers { 0 , 1 } with an even number of 1’s.
For example, the string 0110 belongs to the language but 0111 does not.
• The D n language is defined on the alphabets { a , b }.
The D 1 language can be written using the regular expression ( ab ) ⋆ , where “ ⋆ ” represents the Kleene star operation.
Then for n ≥ 1 , D n can be expressed by the recursive relationship D n = ( aD n − 1 b ) ⋆ .
We choose Dand D 4 , where D 2 =
( a ( ab ) ⋆ b ) and D 4 = ( a ( a ( a ( ab ) ⋆ b ) ⋆ b ) ⋆ b )
⋆ .
The Tomita languages also use the numbers { 0 , 1 } and include seven different grammars.
In our experiment, we choose Tomita 3, 5 and 6.
In Tomita 3, consecutive 1’s must be of an odd number and consecutive 0 ’s must be of an even number.
In Tomita 5, the total numbers of 1 ’s
and 0 ’s must both be even.
Finally, Tomita 6 requires the difference between the total number of 1 ’s and 0 ’s to be a multiple of 3.
In this experiment, we adopt a simple setting with 3 layers and 5 heads for all models.
The embedding size is set to 20.
During training, the Adam optimizer is initialized at 0.005 with a 0.5 × decay every 5 epochs.
And all models are run for 25 epochs without early stopping.
Both the training set and bin 0 contain samples with length in the range of [2 , 50] for Parity and Tomita, and [2 , 100] for D n , while bin 1 contains longer test samples with length in the range of [51 , 100] for Parity and Tomita, and [101 , 200] for D n .
The models are trained on 10K samples and tested on 2K samples per bin for Parity and Tomita, and trained on 5K samples with 1K test samples per bin for D n .
Following
Bhattamishra et al.
(2020), we train the models as below: starting from the first position of an input sequence, the model outputs a string of possible characters for the next position, and this continues until the sequence is exhausted.
The prediction task is recast into a classfication task, where the prediction is deemed a success if and only if the output values at each position are correct, and a failure otherwise.
From Table 2, we can see that the Transformer-only models may require much more training samples to learn these patterns.
Moreover, when the training samples are shorter than the test samples, they will have difficulty in extending the learned patterns to longer sequences, even with absolute or learned relative positional encoding.
Meanwhile, REMs can easily extend its learned patterns to any sequence length, which may account for the better generalization ability of RSA-Transformers.
To help interpret the results in Table 2, we provide some visualizations of the learned recurrent patterns on different languages in Figure 11.
Learning Parity requires the model to count the number of 1’s, and then decide whether it is an even number.
Its recurrent pattern at (a) mainly consists of alternating {− 1 , 1 } that leads to regular oscillations when encountering an even or odd number of 1’s.
Tomita 5 has an additional requirement that the number of 0 is also even, and hence its learning can be assisted by dilated patterns; see (b) for illustration.
Although Tomita 6 is more complicated, the regular cyclical patterns at (c) seems sufficient for it to learn well.
This gives us a glimpse of how the
RSA-Transformer manages to learn the regular languages.
E. C ODE AND NATURAL LANGUAGE MODELING
For RSA-CodeT5-small, our code is based on the original CodeT5, which can be referred to in the repository https://github.com/salesforce/CodeT5 .
For the RSA-XL models adopted for the natural language modeling, only some of the heads are replaced with the RSA heads, while the remaining ones are still the original XL heads.
In the RSA-BRT and baseline BRT models, the recurrent layer is put at the third layer counted from the top, i.e., the recurrent layer is at layer 12 for 14-layer models and at layer 14 for 16-layer models.
The nonrecurrent layers are of the same structure as RSA-XL and XL respectively.
In Table 7, we report the detailed hyperparameters settings for the RSA and their baseline models used in Table 3, where “n_RSA_heads" refers to the number of heads that are changed to the RSA heads.
[Table]
E. A BLATION STUDY
This subsection contains one additional ablation experiment to compare the parameter efficiency of
Block Recurrent Transformer (BRT) against its RSA variant i.e. RSA-BRT, and implementation details of the sequence length ablation study.
We conduct an additional experiment to compare the scaling performance of RSA-BRT and BRT with respect to the number of parameters on Enwik8 dataset.
Specifically, the number of layers in both models varies in [8 , 10 , 12 , 14] , and the training schemes for NLP task in Section 4 is adopted.
From the results presented in Figure 12, it can be observed that RSA-BRT consistently achieves better performance than its baseline.
Next, we provide implementation details for sequence length scaling experiement.
For all four models (i.e. BRT, TLB, RSA-BRT, RSA-TLB), a 4 layer transformer with 8 attention heads are trained using a batch size of 16 and the Adam optimizer with a learning rate of 0.0001 which is reduced by half every 5 epochs for BRT and RSA-BRT.