L ANGUAGE M ODELING USING LMU S AND I MPLICIT S ELF -A TTENTION
Anonymous authors
Paper under double-blind review
A BSTRACT
Recent studies have demonstrated that the performance of transformers on the task of language modeling obeys a power-law relationship with model size over six orders of magnitude.
While transformers exhibit impressive scaling, their performance hinges on processing large amounts of data, and their computational and memory requirements grow quadratically with sequence length.
Motivated by these considerations, we introduce a novel attention module called implicit self-attention and construct a Legendre Memory Unit based model hat exhibits an O ( n ) and O ( n ln n ) (or better) dependency for memory and computation respectively.
Over three orders of magnitude
, we show that for the same amount of training our model improves the loss over transformers about as much as transformers improve over LSTMs.
Additionally, we demonstrate that adding global self-attention complements our architecture and the augmented model improves performance even further.
I NTRODUCTION
Self-attention architectures such as the transformer (Vaswani et al., 2017) have been extremely successful in dealing with sequential data and have come to replace LSTMs and other RNN-based methods, especially in the domain of Natural Language Processing (Radford et al., 2018; 2019; Devlin et al., 2018).
Transformers facilitate parallelization within training examples, and this allows them to fully leverage hardware accelerators such as GPUs, making training on datasets as large as 750GB feasible (Raffel et al., 2019; Gao et al., 2020).
In addition to parallelizing training, selfattention architectures are much better at handling long-range dependencies relative to traditional
RNNs, and this allows them to take advantage of context much longer than the ∼ 100 - 1000 tokens typical of RNNs, like the LSTM (Voelker & Eliasmith, 2018; Kaplan et al., 2020).
Transformers are general-purpose architectures that can be applied to a wide variety of problems and modalities.
One of the drawbacks of such generality is the lack of a priori structure, which makes them heavily reliant on large quantities of data to achieve good results.
Another limiting factor is that self-attention involves the computation of the attention matrix, QK T , which is of shape n × n , with n being the sequence length.
Thus, transformer’s compute and memory requirements grow quadratically with respect to the sequence length.
In this work, we explore a way of addressing these limitations.
We base our approach on the nonparametric Linear Time-Invariant (LTI) component of the Legendre Memory Unit (Voelker et al., 2019).
This LTI system, which we refer to it here as the LMU, 1 projects a sliding window of the input sequence onto Legendre polynomials to provide a temporal representation and compression of the input signal.
Although the LTI system is an RNN, it has been shown to support both sequential and parallel processing of sequences (Chilkuri & Eliasmith, 2021).
Another crucial component of our model is a modiﬁed attention mechanism that operates only on the output of the LMU at each time step, and not across time steps.
The LMU state at each step captures information about the past tokens, and hence we call this attention mechanism implicit self-attention .
[Figure] Recent work (Kaplan et al., 2020; Gao et al., 2020) has explored the scaling properties of transformers in the context of autoregressive language modelling.
They show that the performance of transformers scales as a power-law with model size (excluding embedding parameters), dataset size and the amount of compute used for training, when not bottlenecked by the other two.
Inspired by these results, we validate our method by studying the scaling properties of our models.
To that end, we start by reviewing some necessary background in Section 2, and then present the architectural details of our model in Section 4.
In Section 5, we present our experiments with models containing up to 10 million parameters (or 1 million non-embedding parameters), which demonstrate a smooth power-law relationship with respect to the model size, and a loss that scales better than transformers.
When the loss is matched to that of transformers, 10x fewer tokens are required.
B ACKGROUND : T HE L EGENDRE M EMORY U NIT
The non-parametric LTI component of the LMU (Voelker & Eliasmith, 2018) is the focus of our study.
This LTI system is mathematically derived to project a sliding window of length θ of the input sequence onto q Legendre polynomials.
Thus, the two main hyper-parameters to choose when using it are θ and q .
Naturally, if we desire to capture the ﬁne-grained details of the input sequence, a large θ , which sets the length of the sliding window, should be accompanied by a large q , the number of Legendre polynomials used in approximating the input.
We present the state-update equations of the LTI component of the LMU below,
[Equation]
where the ¯ A = e A ∈ R q × q and ¯ B
=
A − 1 ( e A − I ) B ∈ R q × 1 matrices are frozen during training, with A and B deﬁned as follows:
[Equation]
Crucially, when needed, the LTI equation (1) above can be evaluated as a convolution, in parallel, as shown below (Chilkuri & Eliasmith, 2021):
[Equation]
or equivalently, deﬁning
[Equation]
Hidden State
LMU
Output
V K Q
Input, t
M
Hidden State
LMU
Output
V K Q
Input, t
M
Hidden State
LMU
Output
V K Q
Input, t
M
Hidden State
LMU
Output
Reshape and Project
Implicit Self-Attention
V K Q
Project and FFN FFN / Global Attention Input M (n, dq) (n, q’, d) (n, 3, q’, d) (n, d) (n, d) the above convolution equation can be written as an element-wise multiplication in the Fourier space as follows:
[Equation]
R ELATED W ORK
Our work falls under the broad category of combining convolution with self-attention.
Recent work has demonstrated that combining these two modules can be beneﬁcial for language modeling, speech and other NLP applications
(Yang et al., 2019; Wu et al., 2020; Gulati et al., 2020).
For instance, Wu et al.
(2020) introduce the Long-short Range Attention module that features a two-branch design, with the self-attention branch learning global interactions and the convolutional branch learning local interactions.
Gulati et al.
(2020) on the other hand propose a model that uses a single branch architecture, with the convolutional block connected directly to the attention block, and demonstrate improved performance on the task of speech recognition.
While our work is mathematically related to the models mentioned above, it differs from the previous approaches in three crucial ways: ﬁrst, we do not learn the convolutional weights, but instead work with the analytically deﬁned weights of the LMU; second, we introduce the novel implicit self-attention module that works on the hidden states of the LMU at each time-step; ﬁnally, we systematically study the scaling properties of our model in the context of language modeling.
Our model is also related to the many studies on reducing the quadratic computational and memory complexity of self-attention (Tay et al., 2020).
Out of the several ways of improving efﬁciency, our work shares some similarity with the sliding window attention approach (Beltagy et al., 2020;
Zaheer et al., 2020).
While these methods use a form of masking of the full self-attention matrix to constraint attention to the k -nearest tokens, our technique relies on the LMU to compute optimal compressed representations of sliding windows of input vectors, each of which is then used as input to the implicit attention module.
A RCHITECTURE
In this paper, we modify the architecture presented in Chilkuri & Eliasmith (2021) to better deal with the task of language modelling, especially when the sequences are long and high-dimensional.
Starting with the base LMU model, we describe the major components of our model below.
An illustration of our architecture is presented in Figure 2.
Memory Matrix Consider an input sequence { x 1 , x 2 , . . .
, x n } of length n where the individual elements are of dimension x i ∈ R d .
The most natural way of using an LMU-based model on such sequences is to set θ = n and use an appropriately large q .
The downside of using the LMU in such a manner, however, is that the hidden state of the LMU scales with input dimension and order: m ∈ R dq .
For example, in Section 5, we deal with n = 1024 and d that is as large as 204 .
Even when using a small q of 100, we may end up with hidden states that are as large as R 20 k , which is highly undesirable.
One way around this issue is to take inspiration from standard convolutional network architectures (CNNs) and work with a smaller sliding window, θ ≈ 10 , which in turn allows us to use a small
LMU order, q ≈ 5 , thus taming the hidden state dimension (see Chilkuri & Eliasmith (2021) for more details).
However, enforcing a small sliding window prompts the use of many stacked LMU layers in order to increase the ‘receptive ﬁeld’ (or the effective θ ) of the model, very similar to how
CNNs often use small kernels with many convolutional layers.
Unsurprisingly, such an approach results in very deep models, which can be problematic to train.
Here, we choose to follow the middle path, i.e, 0 (cid:28) q (cid:28) n , but instead of working directly with the potentially high-dimensional hidden state m , we disentangle the input dimensions from the order.
In other words, we perform our operations on the matrix
M ∈ R d × q and not on the vector m ∈ R dq .
This is beneﬁcial because while a fully-connected layer needs d 2 · q 2 parameters to process the m vector, processing the matrix M with the help of two fully connected layers, one for the rows and one for the columns, requires only d 2 + q 2 parameters.
Implicit Self-Attention
The main feature distinguishing our architecture from past work is the
LMU.
As mentioned above, the output of the LMU layer compresses past history at each timestep, which is captured by the M ∈ R q × d matrix.
Our modiﬁed self-attention acts on this matrix to combine temporal information.
As a result, self-attention does not act directly on the input sequence, but rather on a compressed version of the input, which is available at each moment in time and covers a window, determined by θ .
Ignoring the bias vectors, normalization layers, and skip-connections, we ﬁrst execute the following sets of operations simultaneously:
[Equation]
where L i ∈ R q (cid:48) × q , σ is a non-linearity such as gelu , and the matrices Q , K , V are all in R q (cid:48)
× q .
The primary motivation behind this transformation was to reduce the size of the matrix dimension, and we have found the setting q (cid:48)
=
q/ 10 to work well in our experiments, thus the attention matrices contain far fewer elements than n 2 .
For example, in our largest model we set q = 250 , resulting in q (cid:48)
= 25 .
Following the computation of Q , K , and V two additional computations result in a d -dimensional vector m :
[Equation]
where p ∈ R 1 × q (cid:48) .
In practice, equation (8) can be made far more efﬁcient computationally, especially during inference, by following the recipe outlined in Section A.1.
Feedforward Network We have also found it beneﬁcial to include a feedforward network (FFN) (Vaswani et al., 2017) before the LMU and after the implicit self-attention block.
The FFN component is deﬁned below:
[Equation]
where W 1 ∈ R d × d (cid:48) , W 2 ∈ R d (cid:48) × d , b 1 ∈ R d (cid:48) and b 2 ∈ R d .
[Table]
Additional
[Table]
Global Self-Attention We also explore the use of a global self-attention block in place of the FFN component before the LMU.
We ﬁnd that introducing this layer, while computationally expensive, further improves the cross-entropy score.
We believe that the improvement comes from the fact the
LMU and self-attention are complementary: the LMU’s implicit self-attention is good at prediction with limited context [Table] and the traditional self-attention captures long-range dependencies.
We wish to explore the use of efﬁcient self-attention blocks – which scale better than O ( n 2 ) – in the future.
Complexity As shown in Table 1, our architecture employing the parallel LMU along with implicit self-attention has memory requirements that are linear with respect to the sequence length [Table] and it has computational requirements that also grow as n ln n .
When we use the recurrent version of the LMU, the memory and compute requirements scale as O (1) and O ( n ) respectively.
Recurrent implementations, while not as efﬁcient on GPU architectures for large batch sizes, are ideally suited to edge applications, especially with efﬁcient hardware support.
Notably, if we add global attention to our model, then both compute and memory become quadratic, just like the original transformer.
We also list the number of ﬂoating point operations per-token for the (parallel) LMU model in
Table 2.
E XPERIMENTS
Dataset We train our models on the publicly available internet text dataset called OpenWebText (OWT2).
3 Similar to the WebText2 dataset (Radford et al., 2019), OWT2 was created using URLs extracted from Reddit submissions with a minimum score of 3 as a proxy for quality, and it consists of Reddit submissions from 2005 up until April 2020.
After additional ﬁltering, applying the pretrained GPT2 tokenizer containing 50257 tokens (Radford et al., 2019) results in approximatelybillion tokens in total.
We use a train/validation/test split of 96/3/1%.
Training Details We train our models in an autoregressive manner using the Adam optimizer with all the default settings.
We use sequences containing 1024 tokens, and in cases where the documents have fewer than 1024 tokens, we pack multiple documents into the same sequence, separated by the 
 token.
We use a learning rate schedule with a linear warmup and cosine decay to zero, while also reducing the learning rate on plateau.
We chose to train our models to [Figure] process a maximum of 13 billion tokens; at a batch size of 512 , this amounts to training forsteps.
Additionally, one of the important considerations when doing NLP experiments is the size of the embedding vectors, d .
In this work, in order to facilitate a fair comparison to the transformer models (Kaplan et al., 2020), we use the following rule to determine d :
[Equation]
where N represents the number of non-embedding (and trainable) parameters.
We also note that we attempted to match compute across all experiments to the corresponding transformer models.
Results Here we present the results of our experiments that use the LMU architecture described above with the non-embedding (and trainable) parameters ranging from 55k to 1M (i.e., from 2.5M to 10M, if we include all parameters).
The cross-entropy results are presented in Figure 3.
For the transformer models, we list the scores obtained by using the following power-law ﬁt,
[Equation]
where N refers to the number of non-embedding parameters and S is the total number of training steps (set to 25000 at batch size of 512, or equivalent).
The power-law is obtained from Kaplan et al.
(2020); the ﬁts were generated by training several transformer models, ranging in size from 768 to 1.5 billion non-embedding parameters, on OpenAI’s WebText2 dataset.
In addition, we compare against the power-law for LSTM models, also from Kaplan et al.
(2020), that use 10x more training steps than the transformer and LMU models:
[Equation]
Similar to the transformer and LSTM models, we notice that the performance of our models depends strongly on scale, with the loss of the LMU model exhibiting the following power-law relationship with respect to N :
[Equation]
[Figure]
The LMU model with global self-attention scales as follows:
[Equation]
It remains to be seen whether our models retain this performance advantage when N (cid:29) 10 6 .
The utility of adding global attention becomes clear when we observe the per-token loss plots of the three models.
In Figure 4 (right), we notice that although the LMU’s per-token loss is better overall than the transformer’s, it ﬂattens relatively early, around 100 tokens, suggesting that implicit attention alone does not capture long context.
The LMU and Attention model, on the other hand, continues improving with increasing context, similar to the transformer.
It is also interesting to note that we can compare LMUs and transformers by determining approximately how much training is required for a transformer to match LMU loss.
Figure 5 demonstrates that our models, trained on 13 billion tokens, have similar scaling to transformers trained on 130 billion tokens.
Consequently, the LMU architecture is 10x more data efﬁcient.
In addition, our LMU models with global attention continue to outperform transformer models trained on 10x more tokens (or with 10x more training steps) by a signiﬁcant margin.
D ISCUSSION
Semi-supervised learning has proven to be a very effective technique in Natural Language Processing.
General purpose language models pre-trained on a large corpus of text in an unsupervised manner and ﬁne-tuned on tasks such as sentiment analysis and question answering often outperform highly task-speciﬁc architectures that receive no pre-training.
The performance of models on the task of language modelling is thus a crucial metric that is indicative of the downstream performance of such models on a slew of tasks involving natural language.
While the performance of our models on the task of language modelling suggests an interesting trend, due to the scale of our experiments however, we do not consider this to be deﬁnitive evidence for the superiority of our LMU architecture.
As a result, a core objective for future research is to show that the observed trends hold over 6 orders of magnitude, as demonstrated by Kaplan et al.
(2020) for transformers.
Additionally, we would like to point out three things in relation to the use of OpenWebText2: 1)
We made sure to construct or ﬁlter the OpenWebText 2 dataset to resemble Open AI’s WebTextdataset as closely as possible. We trained transformer models on our version of the dataset and conﬁrmed that the ﬁnal performance is lower bounded – lower the better – by the power-law and scores reported in Kaplan et al. Our experiments are reinforced by the results listed in Gao et al.
and Kim (2021), where the authors use OpenWebText2 to study the scaling properties of GPT-2-like transformers.
[Figure]
Regarding the use of the original variant of the transformers as the point of comparison, we note that while efﬁcient transformers (Zaheer et al., 2020; Beltagy et al., 2020; Wang et al., 2020), thanks to their lower memory requirements, can be trained on batches of extremely long sequences and this proves advantageous for tasks that beneﬁt from long context, such as summarization and question answering.
While the efﬁcient transformers tend to outperform the original transformer in such settings, the original variant maintains better performance when working with typical sequence lengths.
Since we are dealing with sequences of length 1024, we reasoned that it would be most appropriate to compare against models employing full self-attention and not an approximate version.
In short, we are demonstrating advantages where efﬁcient transformers do not.
C ONCLUSION
In this work, we employ the Legendre Memory to construct a model that is well-suited to handling long sequences with high-dimensional elements.
We apply our architectures to model natural language in the inﬁnite data limit, demonstrating that: (1) like the established architectures such as transformers and LSTMS, our models also exhibit a power-law relationship between the crossentropy loss and model size; and (2) at the small-medium scale, our models have better scaling properties than other approaches.
A A PPENDIX A. R EDUCED - ORDER LMU
When implementing the LMU in these models, we use a parallelizable approach that computes the impulse responses of the LMU (which are essentially the Legendre polynomials), and convolve those with the input sequences (either using raw convolution or FFT-based convolution).
Speciﬁcally, given the q -dimensional impulse response H , we compute the LMU memory state M ∈ R d × q at the current time as
[Equation]
(13) where X ∈ R n × d is the time-series of previous inputs to the LMU, H ∈ R q × n is the LMU impulse response, and ∗ is the convolution operator.
We have found that our models are most expressive when using a value of q that is signiﬁcantly larger than q (cid:48) , as this allows the LMU to ”remember” the time history with high ﬁdelity, but only use the parts of the history that are most relevant.
Rather than explicitly computing the full LMU output M and then reducing this with the L i transformations as per equation (8), we propose applying the L i transformations directly to the impulse responses
[Equation]
14) and then applying these individually to directly compute Q , K , and V :
[Equation]
This is mathematically equivalent to the formulation expressed previously, but uses signiﬁcantly fewer operations per token, particularly when q is large or the ratio of q (cid:48) to q is small:
[Equation]
where L FFT ( q (cid:48) ) is given by Equation 19 with q → q (cid:48) .
A. LMU I MPLEMENTATION T RADE - OFFS
The LMU itself is a LTI dynamical system, with a number of options for implementation.
One implementation is to perform the update each timestep in state-space, using state-space matrices discretized using the zero-order hold (ZOH) method for high accuracy.
The operations required (per LMU layer and per token) are the multiplications by the ¯ A and ¯ B matrices (with number of elements q 2 and q , respectively):
[Equation]
Another option is to use an explicit Runge-Kutta method to update the LMU states.
By taking advantage of the unique structure of the A and B matrices (Equations 2 and 3), this implementation is able to reduce the complexity from O ( q 2 ) to O ( q ) , requiring the following approximate number of operations:
[Equation]
where r is the order of the Runge-Kutta method.
The disadvantage to this option is that it does not implement the exact same dynamics as the ideal system discretized with ZOH, and is less numerically stable particularly for higher values of q .
A disadvantage to both these options is that they must update LMU states sequentially, which is particularly ill-suited when using highly parallel hardware (e.g. GPU) with a long sequence of inputs available.
In this case, we can take the impulse response of the LMU system (discretized with
ZOH), and convolve it with an input in the FFT domain.
This implements the exact same dynamics as the ZOH state-space system, but with a complexity that is O ( q ) rather than O ( q 2 ) :
[Equation]
Here, C ( n ) is the number of FLOPs for a radix-2 Cooley-Tukey FFT implementation (Johnson &
Frigo, 2012):
[Equation]
[Equation]
[Equation] where c m = 6 is the number of FLOPs per complex multiply [Table] and c a is the number of FLOPs per complex addition.
For our standard sequence length of n = 1024 , this results in: [Equation] [Equation] [Equation]
[Equation]
A. LMU M ODEL D ETAILS
[Table] 3M     1. 1M 11M     1.