Language Modeling using LMUs: 10x Better Data Efficiency or Improved Scaling Compared to Transformers
Authors: Narsimha Reddy Chilkuri
Eric Hunsberger
Aaron Russell Voelker
Gurshaant Singh Malik
Chris Eliasmith
Venue: ICLR 2022 Submitted (2022)
Keywords: Recurrent Neural Network
Legendre Memory Unit
Natural Language Processing

Abstract: 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 construct a Legendre Memory Unit based model that introduces a general prior for sequence processing and 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 our new architecture attains the same accuracy as transformers with 10x fewer tokens.  We also 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.

Under review as a conference paper at ICLR 2022
LANGUAGE MODELING USING LMUS AND IMPLICIT SELF-ATTENTION
Anonymous authors Paper under double-blind review
ABSTRACT
Recent studies have demonstrated that the performance of transformers on the task of language modeling obeys power-law relationship with model size over six orders of magnitude. While transformers exhibit impressive scaling, their per- formance hinges processing large amounts of data, and their computational and memory requirements grow quadratically with sequence length. Motivated by these considerations, we introduce novel attention module called implicit self-attention and construct a Legendre Memory Unit based model hat exhibits ((n) and O(n n) (or better) dependency for memory and computation re- spectively. 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 trans- formers improve over LSTMs. Additionally, we demonstrate that adding global self-attention complements our architecture and the augmented model improves performance even further.
INTRODUCTION
Self-attention architectures such as the transformer (Vaswani al. 2017) have been extremely suc- cessful in dealing with sequential data and have come to replace LSTMs and other RNN-based methods, especially in the domain Natural Language Processing (Radford al. 2018 2019 De- vlin et al., 2018). Transformers facilitate parallelization within training examples, and this allows them to fully leverage hardware accelerators such GPUs, making training on datasets large as 750GB feasible (Raffel et al.] 2019 Gao et al., 2020). In addition to parallelizing training, self- attention 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-purpest purpose architectures that can be applied to 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-ateention involves the computation of the attention matrix. QKT, which of shape n x n, with 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 approach on the non- parametric Linear Time-Invarian (LTI) component of the Legendre Memory Unit (Voelker et al. 2019). This LTI system, which refer it here the LMU projects a sliding window of the input sequence onto Legendre polynomials to provide temporal representation and compression of the input signal. Although the LTI system an RNN, it has been shown to support both sequential and parallel processing of sequences (Chilkuri & Fliasmith 2021). Another crucial component of our model is modified attention mechanism that operates only on the output of the LMU 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.
PMojectssssss
We refer to the LTI system as the LMU as it is the distinguishing layer of our architectures. The original LMU was defined to include the LTT system well as subsequent nonlinear layer. We have essentially expanded this nonlinear layer include variety of familiar layers.

---

Under review as a conference paper at ICLR 2022
mh
INU
A/Br
A'Br
A/Br
A'Rxx
1771
LIT1 1771 LTI
1
1
Multipyy
Figure 1: (left) Illustration the standard sequential implementation for computing the hidden state m.. The imput is fed into the linear recurrent unit to compute the hidden state m., which, along with 22, then used to compute the next hidden state, and on. (right) Illustration of the time- domain parallel implementation for computing the hidden state m4. The inputs 21 are used to compute the intermediate multiplies, which are then added together to compute the hidden state m4, all without the need for any sequential operations.
Recent work (Kaplan et al. 2020 Gao et al. 2020) has explored the scaling properties of trans- formers the context of autoregressive language modelling. They show that the performance of transformers scales 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 start by reviewing some necessary background details of our model Section [ In Section []] we present our experiments with models containing up to 10 million parameters (or million non non-embedding parameters), which demonstrate smooth power-lav 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.
validate method by studying the scaling properties of our models. To that end,
Section 2 and then present the architectural
2 BACKGROUND: THE LEGENDRE MEMORY UNIT
The non-parametric LTI component of the LMU Voelker & Eliasmith 2018) is the focus of our study. This LTI system is mathematically derived project sliding window of length 0 of the input sequence onto Legendre polynomials. Thus, the two main hyper-parameters to choose when using it are 9 and q. Naturally, if we desire capture the fine-gratd finengrained details of the input sequence, large 0, which sets the length of the sliding window, should be accompanied by large g. the number of Legendre polynomials used approximating the input. We present the state-update equations of the LTT component of the LMU below
m. Am,-1 Br,,
where the Á = A E Rex4 and B A 1 IJB E Rex1 matrices frozen during training, with and B defined as follows:
(2) + 1)
-1 i<j
Arj =
(2)
0
(-1)--541
i i2j
B, 2i 1)(-1)
(3)
Crucially, when needed, the L.TI equation ( above can be evaluated as convolution, in parallel, as shown below (Chilkuri & Eliasmith, 2021):
A Baj.
(4)
j=1
or equivalently, defining
H
BB AB ...] Raxn
(5)
x1)² 1 E n
In
En-1
kn-
(6)
Focusiig one-dimensional inputs for now.
2

---

Under review as a conference paper at ICLR 2022
Input
(n, d)
PPN /loball Attention
LMU
Hidden
(n, dq)
Reshape and Project
M
(n, q% d)
Implicit self-tttention
(n, 3. q' d)
Project FFN
(n, d)
Output
Figure 2: The LMU and implicit self-attention architecture along with output dimensions. In the illustration, n refers to the sequence length, q is the order and g' is the reduced order, and d is the embedding dimension. Normalization layers and skip connections are not shown. One variant uses the FFN component right after the input, and the other variant uses global attention.
the above convolution equation be written element-wise multiplication in the Fourier space as follows:
m.... = (F F(x))
RELATED WORK
Our work falls under the broad category of combining convolution with self-attention. Recent work has demonstrated that combining these two modules can be beneficial for language modeling, speech and other NLP applications (Yang et al. 2019: Wu al. [020] Gulati et al. 2020). For instance, Wu jet 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 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 the models mentioned above, it differs from the previ- ous approaches in three crucial ways: first, we do not learn the convolutional weights, but instead work with the analytically defined 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; finally, we systematically study the scaling properties of our model 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 the several ways of improving efficiency, 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 the k-nearest tokens, technique relies on the LMU to compute optimal compressed representations of sliding windows of input vectors, each of which then used as input to the implicit attention module.

---

Under review as a conference paper at ICLR 2022
4 ARCHITECTURE
In this paper, we modify the architecture presented in Chilkuri & Eliasmith (202T) 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 (æ), 222 ..........))) ) of length n where the individual elements are of dimension T, F.. The most natural way of using an LMU-based model on such sequences is to set 0 = 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 E IR For example, in Section 5 we deal with n = 1024 and 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 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, 0 x 10, which in turn allows us to use a small LMU order, q RJ 5, thus taming the hidden state dimension (see Chilkuri & Eliasmith (202I) for more details). However, enforcing a small sliding window prompts the use of many stacked LMU layers in order to increase the 'receptive field' (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 < q < 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 E Rdxq and not on the vector m RRd9 This is beneficial because while a fully-connected layer needs de 9² 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² q² parameters.
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 cach time- step. which is captured by the M Raxd matrix. Our modified 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 compressed version of the input, which is available at each moment in time and covers a window, determined by 0. Ignoring the bias vectors, normalization layers, and skip-connections, first execute the following sets of operations simultaneously:
Implicit
Q = (L,
K = == ML2
(8)
where L, E R x9, is non-linearity such as gelu, and the matrices Q, K, V are all in R Xx The primary motivation behind this transformation was to reduce the size of the matrix dimension, and we have found the setting q' 9/10 to work well in our experiments, thus the attention matrices contain far fewer elements than n For example, in our largest model we set q = 250, resulting in 25.
Following the computation of Q, K, and V two additional computations result in a d-dimensional vector m:
M' softmax ( QKTVV m = pM',
(9)
(10)
where p E RR1xd'.
In practice, equation X can be made far more efficient computationally, especially during inference, by following the recipe outlined in Section A.I
Feedforward Network We have also found beneficial to include a feedforward network (FFN) Vaswani et al., 2017) before the LMU and after the implicit self-attention block. The FFN compo- nent is defined below:
π(æ) = ax + b,) + bg.
(11)
where W E RRdxd' W2 Rdd xd b Rdd and b2 Rdd
4

---

Under review as a conference paper at ICLR 2022
Table 1: Memory and computation scaling with sequence length during training and inference.
Layer
Memory
Compute
Full Attention
O(n²)
O(n2)
LMU (Parallel)
((n)
(nllan)
LMU (Recurrent)
O(1)
((n)
Table 2: Parameter counts and compute (forward pass) for one layer of the network, per token. The first row indicates the number of FLOPs when following the implementation in Section A Additional background information regarding various implementations of the LMU is provided in Appendix A.2
Operation
Parameters
FLOPs per Token
LMU + Q+ K+
3qq'
3d 3d + (logg + 1) + 6g"] + 6qg'
QKT
2dg²²
M'
2dq 12 + da'
m
q'
2d4'
FFN
2dd'
4dd'
Global Self-Attention We also explore the use of a global self-attention block in place of the FFN component before the LMU. We find 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- ceff-attention is good at prediction with limited context, and the traditional self-attention captures long-range dependencies. We wish to explore the use of efficient self-attention blocks which scale better than (nn²) in the future.
Complexity As shown in Table our architecture employing the parallel LMU along with implicit self-attention has memory requirements that are linear with respect to the sequence length, and has computational requirements that also grow Ilnn. When the LMU, the memory and compute requirements scale as ((1) and O(n) respectively. Recurret implementations, while not efficient GPU architectures for large batch sizes, idcally suited to edge applications, especially with efficient hardware support. Notably, we add global attention to our model, then both compute and memory become quadratic, just like the original transformer.
use the recurrent version of
We also list the number of floating point operations per-token for the (parallel) LMU model in Table 2
EXPERIMENTS
Dataset We train our models on the publicly available internet text dataset called Open WebTexx2 (OWT2) Similar the WebText2 dataset (Radford et al., 2019), OWT2 created using URLs extracted from Reddit submissions with a minimum score of 3 as proxy for quality, and it consists of Reddit submissions from 2005 up until April 2020. After additional filtering, applying the pre- trained GPT2 tokenizer containing 50257 tokens (Radford et al. 2019) results in approximately 8 split 96/3/1%.
billion tokens in total. We use
Training Details We train our models 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 endo f sequence > token. We use learning rate schedule with linear warmup and cosine decay to zero, while also reducing the learning rate on plateau. We chose to train our models to
5

---

Under review as a conference paper at ICLR 2022
Performance vs Parameters
LSTM
Transformer L(N.S)
10
LMU
LMU Attention
LOSS
x x0
10a
10*
Parameters (non embedding)
Figure 3: Cross-entropy scores in nats, averaged across all the tokens in the sequence. Transformers and LSTMs fits are from Kaplan et al. (2020). Our models perform better than Transformers and LSTM models up to million non-embedding parameters.
process a maximum of 13 billion tokens; at batch size of 512, this amounts to training for 25000 steps. Additionally, of the important considerations when doing NLP experiments is the size of the embedding vectors. d. In this work, in order to facilitate fair comparison to the transformer models (Kaplan et al. 2020). we use the following rule determine d:
N
d
24'
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 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 10M, if include all parameters). The cross-entropy results presented in Figure 3 For the transformer models, we list the scores obtained by using the following power-law fit,
N
-0.077
S
0.76
Transformer(M) (o N,
(12)
6.5 1013
Smin (S)
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 obtained from Kaplan al. (2000); the fits were generated by training several transformer models, ranging in size from 768 to 1.5 billion non-embedding parameters, on OpenAl's WebText2 dataset In addition, we compare
against the power-law for LSTM models, also from Kaplan al. 2020), that use 10x more training steps than the transformer and LMU models:
N
LSTM(N)
7.45 1044
Similar to the transformer and LSTM models, we notice that the performance of our models depends strongly scale, with the loss of the LMU model exhibiting the following power-law relationship with respect to N:
N
LMU(N) 1.95 1014
"We note that the LMU models do not make use of positional embedding parameters, and for the trans- former models, the parameter count excludes the positional embedding parameters (given by embed dimension sequence length). For our largest model, the total number of non-parametric weights inside all the LMU layers (given order order order) amounts about 70% of the positional emhedding parameters.

---

Under review as a conference paper at ICLR 2022
Per-token Loss
8.0
Transtomer
7.5
LMU + Attention
Par-to
7.0
6.5
UNU
5.0
6.0
5.5
5.5
55k
5.0
4.5
1M
10-
10a
10a
122
20
107
12
Sequance Indix
Suquemoe trdex
Figure 4: (left) Comparison of per-token loss of LMU model (with global attention) and trans- former model. (right) Per-token loss of an LMU model (without global attention) alongside the transformer's loss.
The LMU model with global self-attention scales as follows:
N
-0.069
LMUg(N) = 3.80 - 1014
remains to be seen whether our models retain this performance advantage when N >> 10%
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 s, it flattens 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.
is also interesting to note that we can compare LMUs and transformers by determining approxi- mately how much training is required for transformer to match LMU loss. Figure demonsttates that our models, trained on 13 billion tokens, have similar scaling to transformers trained on 130 bil- lion tokens. Consequently, the LMU architecture is 10x more data efficient. 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 significant margin.
DISCUSSION
Semi- supervisd learning has proven to be very effective technique in Natural Language Process- ing. General purpose language models pre-trained on a large corpus of text in an unsupervised manner and fine-tuned tasks such as sentiment analysis and question answering often outperform highly task-specific architectures that receive no pre-training. The performance of models on the task of language modelling is thus a crucial metric that indicative of the downstream performance of such models on 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 definitive evidence for the superiority of our LMU architecture. As a result, core objective for future research is to show that the observed trends hold over 6 orders of magnitude, demonstrated by Kaplan et al. 2020) for transformers.
Additionally, we would like to point out three things in relation to the use of Open WebText:::: 1) We made sure to construct or filter the Open WebTett 2 dataset to resemble Open Al's WebText dataset as closely confirmed that the final performance is lower bounded - lower the better - by the power-law and scores reported in Kaplan et al. (020)) 3) Our experiments are reinforced by the results listed in Gao jet al. (2000) and Kim (2021), where the authors use OpenWebText to study the scaling properties of GPP---like transformers.
possible. 2) We trained transformer models on our version of the dataset and
7

---

Under review as a conference paper at ICLR 2022
*offemarce
LNTM JMC Adartion
4x337
Figure 5: Approximately matching the loss between transformers and LMUs requires 10x more training for the transformer. The LMU and Attention model continues to significantly outperform transformers with 10x less training.
Regarding the use the original variant of the transformers as the point of comparison, we note that while efficient transformers (Zaheer et al. 2020 Beltagy et al. 2020 Wang al. 2020), thanks to their lower memory requirements, can be trained on batches of extremely long sequences and this proves advantageous for tasks that benefit from long context, such summarization and question answering. While the efficient transformers tend to outperform the original transformer in such set- tings, 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 compare against models employing full self-attention and not an approximate version. In short, we are demonstrating advantages where efficient transformers do not.
7 CONCLUSION
In this work, we employ the Legendre Memory to construct a model that is well-suited to han- dling long sequences with high-dimensional elements. We apply our architectures model natural language in the infinite data limit, demonstrating that: (I) like the established architectures such as transformers and LSTMS, our models also exhibit power-law relationship between the cross- entropy loss and model size; and (2) at the small-medium scale, models have better scaling properties than other approaches.
8

---

Under review as a conference paper at ICLR 2022
REFERENCES
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint XXxv:2004.05150, 2020.
Narsimha Reddy Chilkuri and Chris Eliasmith. Parallelizing legendre memory unit train- ing. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Confer- ence on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp.
1998-1907. PMLR, 18-24 Jul 2021. URL http: attp:://proceedings wtt.:mmmm..pprrrr htt..igss.......ss111113333 chi wwl.kkkuaa.hhat.
Jacob Devlin, Ming-Wes Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: /80.vv/8/00044055 2018.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling. arXiv preprint 1000:2201.00007, 2020.
Anmol Gulati, James Qin, Chung-Cheng Chiu, Niki Parmar, Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang, Yonghui Wu, et al. Conformer: canvolution-augment transformer for speech recognition. arXiv preprint arXiv:2005.08100, 2020.
Steven G. Johnson and Matteo Frigo. Implementing FFTs in Practice. In C. Sidney Burrus
(ed.), Fast Fourier Transforms. 2012. URL httpp://ts/n/ http:///xxxrrg//////////// wwt.e/rrgg//oooeeee wwwttoootttttttoooooeeeoo
15/
in- Peaan--raciieee
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08366 2020.
Christina
Kim. 2021.
Scaling
laws for language transfer learning. christina.kin URL
htapsss/chrrstiia... wwt.km/mmmm/////. lscalings htw.-g/wwwww-----oo------aaaaaa-
Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improving language under- standing by generative pre-training, 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. arXiv preprint arXiv: 10683, 2019.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey, arXiv preprint ar XXv::2009006732 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Aaron Voelker, Ivana Kajić, and Chris Eliasmith. Legendre memory units: Continuous-time repre- sentation in recurrent neural networks. In Advances in Neural Information Processing Systems, pp. 15544-15553, 2019.
Aaron R Voelker and Chris Eliasmith. Improving spiking dynamical networks: Accurate delays, higher-order synapses, and time cells. Neural computation, 30(3):569-66) 609, 2018.
Sinong Wang, Belinda 7. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint XiX::2006.047666 2020.
Zhanghao Wu, Zhijian Liu, Ji Lin, Yujun Lin, and Song Han. Lite transformer with long long-short range attention. In International Conference on Learning Representations, 2020. ht
//openrevien watt/fooum? d= HKPH

---

Under review as a conference paper at ICLR 2022
Baosong Yang, Longyue Wang, Derek Wong. Lidia S Chao, and Zhaopeng Tu. Convolutional self- attention networks. arXiv preprint 1rXi:11004003177, 2019.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences, In NeurIPS, 2020.
A APPENDIX
A.1 REDUCED-ORDER LMU
When implementing the LMU in these models, we use a parallelizable approach that computes the impulse responses of the LMU (which essentially the Legendre polynomials), and convolve those with the input sequences (either using raw convolution or FFT-based convolution). Specifically, given the g-dimensional impulse response H, we compute the LMU memory state M E 12d8 the current time as
M X H
(13)
where X E 1 is the time-series of previous inputs to the LMU, H E RRxx is the LMU impulse response, and is the convolution operator.
We have found that our models are most expressive when using a value of y that is significantly larger than q', this allows the LMU "remember" the time history with high fidelity, but only 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, transformations as per equation (8. we propose applying the L; transformations directly to the impulse responses
H. L,H
(14)
and then applying these individually to directly compute Q, K, and V:
Q (X Hn)
== H2)
V V=o(X * H3)
(15)
This is mathematically equivalent to the formulation expressed previously, but uses significantly fewer operations per token, particularly when is large or the ratio of q' to q is small:
LLMU-QUK+V 3(dEFFT(g)) 2gq')
(16)
where (g")(g') is given by Equation with q -+y..
A.2 LMU IMPLEMENTATION TRADE-OFFS
The LMU itself is LTI dynamical system, with 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 q2 and q respectively):
Ess 2d + q).
(17)
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 Iiquaations[and]]]] this implementation is able to reduce the complexity from O(g²) ) to Ogg), requiring the following approximate number of operations:
OKK = 6rdg.
(18)
where r is the order the Runge- Kunge-Kuuta method. The disadvantage to this option is that it does not implement the exact same dynamics the ideal system discretized with ZOH, and is less numeri- cally 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
10

---

Under review as a conference paper at ICLR 2022
ZOH), and convolve it with imput in the FFT domain. This implements the exact same dynamics as the ZOH state-space system, but with a complexity that is O(g) rather than O(g²):
d [C(2m) + Cm (2
SFFT
n
[5ogg () + 1) 6q)
(19)
Here, C(n) is the number of FLOPs for radix-2 Cooley- Tukey FFT implementation (Johnson & Frigo 2012):
(nn) 2C
11
+
71 2
Cm + 2ca)
(20)
5n log, n
(21)
where 6n the number FLOPs per complex multiply, and eo is the number of FLOPPs per complex addition. For our standard sequence length of 1024, this results in:
1074 = d(6lq + 55).
(22)
A.3 LMU MODEL DETAILS
Table 3: LMU model details. N refers to the number of non-embedding and trainable parameters; is the embedding dimension; and 0 define LMU's order and the length of the sliding window; is the number of rows in the matrix. We adjust the post-FFN inner ration to obtain the right parameter count for comparison to transformers it's usually set to something in between 1.9 and I.
Total Parameters
Pre-FFN
N
d
q''
Layers
0
Inner Ratio
55k
2.4M
48
3
350
100k
3.3M
65
65
7
3
350
1.5
200k
4.8M
91
90
3
350
1.5
300k
5.9M
112
110
13
3
350
1.5
500k
7.3M
144
150
15
3
350
1.5
IM
11M
204
220
22
3
350
1.5
11


Reviews:
[{"review_id": "wpB5YVFNhK", "reviewer": "Reviewer_9uKE", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "The quadratic memory and time complexity of transformers have resulted in recent efforts to replace the softmax self-attention with efficient variants. This paper presents one such effort, which is very interesting, and could pave the way for a new generation of transformers. The approach is simple and relatively straightforward, and the writing is overall clear.\n\nI have two main concerns about the paper that prevent me from recommending acceptance, both boil down to the same point: I am not convinced that the experiments support the authors interpretation. First, the paper focuses on the main idea of using LMU, but the proposed model contains many different components, that are orthogonal to it. For instance, why is the $L_i$ projection and the following non-linearity needed? Would applying it to the transformer lead to similar gains? The authors say _Our modified 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._ But the same compression can be applied on standard SA as well. Further, the FFN layer after the input was found to be useful by the authors. Did they also add it to their baselines? In general, some ablations would make the paper much stronger and shed more light on the proposed model (on top of what is mentioned above, I would also like to see the effect of the values of q and q').\n\nSecond, the experimental setup is underspecified: the paper's main selling point is the strong performance compared to the transformer, but important details are missing with respect to both the models and the baselines. How many layers does each model have? What is the hidden size? Importantly, the last experiment (Fig. 5), which appears in the paper title, is even less specified: was the same transformer model trained on 130GB of data? Where did this data come from? As a side note, it is also unclear that such small model could leverage 130GB of data, and thus this experiment might not be ideal to test the sample efficiency of the proposed model. I would recommend providing more details about the experimental setup, and justifying the experimental choices.\n", "summary_of_the_paper": "This paper proposes an efficient approximation of the transformer model. The authors replace the matrices in the dot product with a the output of a Legendre Memory Unit (LMU) transformation of the input, which projects a sliding window of the input sequence onto Legendre polynomials. This projection is fixed and not learned, which allows for memory savings. The resulting matrices are also smaller, which leads to further efficiency gains. The idea of using LMUs has been proposed before, and this paper applies them to the transformer. Experiments with language modeling show substantially lower (i.e., better) loss compared to the vanilla transformer model on several model sizes.\n", "summary_of_the_review": "The paper presents a very interesting idea with strong empirical results on language modeling, that could inspire a new generation of transformers. However, there is some uncertainty with respect to the connection between the experiments and the research hypotheses, which makes it challenging to convince oneself of the value of the results. I am looking forward to reading the authors' response.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "byM2oVZvEeg", "reviewer": "Reviewer_ooZ6", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "Strengths\n* Improving the data efficiency in language models is an important problem that so far studies have shown that can be achieved by scaling the size of the model. Having additional ways to improve data efficiency by changing the model design is definitely of interest. \n* Using the Legendre Memory Unit to substitute self-attention in transformers is interesting and has several potential merits: it can reduce the complexity and does not increase the size of the layer.  Although it turns out that additional components need to be introduced for good performance. \n\n\nWeaknesses\n* Something that stands out is the lack of discussion and comparison to related works that employ recurrent formulations of attention. Kernel-based variants of self-attention have a recurrent formulation and lead to linear complexity (see [1,2,3,4]). \n* The rationale behind the architectural choices for the self-attention component is not well explained or empirically verified. I understand that LMU might have limited capacity but this is not specifically discussed and it is unclear how each component contributes to the end performance.  \n* Even though this paper proposes a new efficient transformer, the evaluation does not focus on computational efficiency aspects and comes across as incomplete. Quantifying the latency and scaling with respect to sequence length would be needed for convincing the reader about its usefulness.  \n*  The evaluation focuses on comparing with an empirical law learned on a different experimental configuration and there is a concern about how comparable are the results to the ones obtained in this study and the validity of the conclusions. I don't think it is good practice to use these power laws universally, especially for this purpose. Another limitation is that it is unclear whether the improvement would hold when the size of the model increases; the evaluation is dealing with scaling laws after all. I would suggest the following:\n  * To make the comparison more fair, I would suggest to train transformer models with varying size and derive a power law based on the exact same experimental configuration used for LMU models.\n  *  The claims regarding the 10x better data efficiency are not well supported and I would suggest the authors to compare with transformer models on standard LM benchmarks and potentially to some downstream tasks such as MT to make a stronger case. \n\n\n\nQuestions: \n* How can we be sure that the specific power law derived from a different experimental configuration in Kaplan et al. (2020) corresponds to the empirical law that can be derived for transformers in this particular evaluation setting? It would be great if the authors discuss this or provide some supporting evidence about its correctness. \n* From the evaluation, it seems that comparing to other efficient transformers was not the main focus of the work but people might wonder how well it works compared to alternatives. What recommendation the authors would give for those interested in using it? \n*  What is the contribution of each specific design choice, such as the FFN/global attention and implicit self-attention, to the end performance? An ablation analysis would be most appropriate for quantifying this.  \n* Could the authors provide more details about the evaluation settings and number of parameters of each model for the per-token loss comparison? Is the per-token loss computed on the test set or training set?  Why is the sequence index capped at 10^3? \n\n[1] https://arxiv.org/pdf/2006.16236.pdf\n\n[2] https://arxiv.org/pdf/2009.14794.pdf\n\n[3] https://openreview.net/pdf?id=QtTKTdVrFBB\n\n[4] https://openreview.net/pdf?id=ot2ORiBqTa1", "summary_of_the_paper": "This paper aims to address two issues with standard transformers for language modeling, namely their large data requirements and their computational cost due to self-attention complexity. The main proposal is to replace self-attention in transformers with Legendre Memory Units from Voelker et al (2019) which reduce memory complexity to linear and computation to log linear for their convolutional variant as well as  constant/linear for their recurrent variant respectively. The empirical claim that it makes is that the proposed transformers require 10x fewer tokens to reach the same loss, when comparing to the laws derived from the experimental configuration of Kaplan et al (2020). ", "summary_of_the_review": "Overall, the idea is interesting and the goal of the paper is definitely valuable. The framing and the delivery, however, are greatly lacking. My main concerns are regarding the validity of the evaluation and the lack of experimental details and comparison to alternative efficient variants. For these reasons, I think that the paper is not ready yet because it lacks sufficient evidence for the main claims and requires additional experimental efforts to better demonstrate the effectiveness of the proposed method. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "bXzwmGeH9u2", "reviewer": "Reviewer_hsfs", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "Strengths:\n- The output of the LMU can compress past history at each time step. The use of LMU reduces the self attention complexity from n^2 to q'xq where q is the order of the q Legendre polynomials. This provides advantage to tasks of larger sequence lengths. \n\nWeaknesses:\n- Like discussed in the paper, LMU's implicit self-attention is good at prediction with limited context, while the traditional self-attention can capture long-range dependencies. Also, LMU lacks a mechanism to provide pairwise information, which can be captured by self-attention. This can be illustrated using entailment tasks or translation tasks. \n\n- Empirical results are very limited. Cross-entropy scores are limited because language modeling or a masked language modeling does not reflect how good the model can be at tasks that requires pairwise information. Better to compare with Transformer on down stream tasks with finetuning as well. \n\n- Lacks a comparison with stronger baselines of efficient transformers. Better to compare agains Linformer (low-rank), Synthezer (Random generation), and Primer (learning the primitives).\n\n- Lacks ablation study on the selection of window size and order in the LMU. q is only effective when sequence length is larger than q. For  tasks of smaller sequence length, they won't benefit from the LMU method. \n\n[1] Linformer: Self-Attention with Linear Complexity\n\n[2] Synthesizer: Rethinking Self-Attention in Transformer Models\n\n[3] Primer: Searching for Efficient Transformers for Language Modeling\n", "summary_of_the_paper": "The paper proposed a way to address two limitations of a transformer network. The method is based off a non-parametric Liner Time-Invariant component of the Legendre Memory Unit (LMU), which projects a sliding window of the input sequence onto a Legendre polynomials to provide a temporal representation and compression of the input signal. The newly proposed attention operates only on the output of the LMU at each time step, not across all time steps. The model is validated by studying the scaling properties of the models. The model is demonstrated scaling better than transformers or requires 10x fewer tokens to match the loss of a transformer network. ", "summary_of_the_review": "To sum up, the paper is too immature to publish in ICLR for several reason:\n- First, it lacks a rigorous analysis on complexity compared to vanilla transformer, on various sequence length. Varying the sequence length, the performance of a LMU based method can be changed. \n\n- The paper is weak in baseline. There are many recently work towards building efficient transformers, using low-rank, approximation, kernel based method, and neural architecture search. All of the work provide stronger performance than vanilla transformer. It is better to compare with at least a few of them.\n\n- Perplexity or cross-entropy on a pretraining task such as LM or MLM can be very limited. Plenty of related work (such as gMLP [4]) have been demonstrated useful in the pretraining task, but they fail to outperform transformer in the down stream tasks with finetuning. Many work failed to provide pairwise information as a self-attention can provide. \n\n[4] Pay Attention to MLPs", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "JGSt2_zC2m_", "reviewer": "Reviewer_HHvw", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "Strengths:\n1. Novel and interesting idea, well motivated architecture.\n2. Empirical evaluation over a wide range of model sizes in terms of # parameters.\n\nMajor limitations:\n1. The paper presents an incomplete comparison that plays to the strength of the LMU architecture by comparing only parameter-matched models. Given that a significant fraction of the LMU compute is hidden in non-trainable parameters, a fair comparison should also compare models trained with an equivalent amount of training computation; ideally in terms of total training compute used.\n2. Limited empirical evaluation on 1 language modeling task: also evaluating on tasks other than language modeling (say downstream tasks like the SuperGLUE benchmark, or Machine Translation) would significantly strengthen the claims made in the paper.\n\nQuestions or comments:\n1. The idea of disentangling input dimensions from order to reduce kernel computation costs sounds very similar to depthwise separable convolutions [1].\n2. Is the x-axis in Figure 5 mislabeled? From the caption and the text I would expect it to represent number of training steps or number of tokens trained on.\n\nReferences:\n[1] Xception: Deep Learning With Depthwise Separable Convolutions, Chollet et al.", "summary_of_the_paper": "This work introduces Legendre Memory Units, an approach that utilizes non-parametric linear time-invariant layers to compute representations of a sequence of inputs. The key idea is similar to convolutions, but the weights of the convolution layer were static (non-trainable) and produced a representation of the sequence that can also be represented as a RNN. The LMU is equipped with a implicit self-attention layer; which bypasses the quadratic-time complexity of self-attention by performing attention over a fixed length sequence of hidden states. \n\nThe authors compare the proposed LMU against transformers on a language modeling task with the Webtext-2 dataset. This comparison is conducted over a range of model sizes (ranging from 55K to 1M non-embedding parameters), demonstrating that LMU improves over transformers with a similar number of parameters over the entire range; and extrapolating from the scaling law suggests that the difference would persist at larger model sizes. \n\nBy equipping the LMU with a global self-attention layer (or a transformer self-attention layer), the authors demonstrate additional gains similar to what was observed with a LMU over a vanilla transformer.", "summary_of_the_review": "This work proposes a very interesting application of a novel, parameter-efficient architecture on a large-scale language modeling task. However, the paper presents an incomplete comparison that plays to the strength of the LMU architecture by comparing only parameter-matched models. Given that a significant fraction of the LMU compute is hidden in non-trainable parameters, a fair comparison should also compare models trained with an equivalent amount of training computation; ideally in terms of total training compute used.\n\nAnother weakness pertains to the limited number of tasks used for this comparison; also evaluating on tasks other than language modeling (say downstream tasks like the SuperGLUE benchmark, or Machine Translation) would significantly strengthen the claims made in the paper.\n\nOverall, given the above two limitations I would recommend rejecting the paper. While the proposed architecture is interesting and well motivated, the paper needs stronger empirical validation and a fairer comparison against existing architectures.\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a language modeling architecture based on the RNN cells leveraging Legendre memory units. The proposal is interesting, but as all the reviewers notice, the paper is not ready for the presentation in the top ML conference for several reasons: comparison with weak baselines, shallow or weak analysis of the presented results, insufficient discussion of the related work, etc. Looking forward for all the comments to be addressed by the authors.

In the rebuttal the authors addressed some of the questions but all the reviewers think that the paper is not ready for acceptance and careful rewriting is needed. Recent research on the improved RNN mechanisms suggests that Legendre memory units and related mechanisms might be a gateway to solving several standard issues of training regular RNNs so the topic is definitely of great importance. Thus the authors are highly encouraged to resubmit the paper after making all suggested corrections.
Decision: Reject