Long Range Language Modeling via Gated State Spaces
Authors: Harsh Mehta
Ankit Gupta
Ashok Cutkosky
Behnam Neyshabur
Venue: ICLR 2023 poster (2023)
Keywords: Long range language modeling
language modeling
state space models

Abstract: State space models have shown to be effective at modeling long range dependencies, specially on sequence classification tasks. In this work we focus on autoregressive sequence modeling over English books, Github source code and ArXiv mathematics articles. Based on recent developments around the effectiveness of gated activation functions, we propose a new layer named \textit{Gated State Space} (GSS) and show that it trains significantly faster than the diagonal version of S4 (i.e. DSS) on TPUs, is fairly competitive with several well-tuned Transformer-based baselines and exhibits zero-shot generalization to longer inputs while being straightforward to implement. Finally, we show that leveraging self-attention to model local dependencies improves the performance of GSS even further.

Published as a conference paper at ICLR 2023
LONG RANGE LANGUAGE MODELING VIA GATED STATE SPACES
Harsh Mehta
Ankit Gupta
Google Research harshm@google.com
IBM Research" htt.:pppaiiitaaauuu@gmiilllllo
Ashok Cutkosky
Behnam Neyshabur
Boston University
Deepmind
ashok@cutkosky.com
neyshaburQgoogle.coom
ABSTRACT
State space models have shown to be effective at modeling long range dependencies, especially on sequence classification tasks. In this work we focus on autoregressive sequence modeling over English books, Github source code and ArXiv mathematics articles. Based on recent developments around the effectiveness of gated activation functions, we propose new layer named Gated State Space (GSS) and show that it trains significantly faster than the diagonal version of S4 (i.e. DSS) on TPUs, competitive with several well-tuned Transformer-basee baselines and exhibits zero-shot generalization to longer inputs while being straightforward implement. Finally, we show that leveraging self-attention model local dependencies improves the performance of GSS even further.
INTRODUCTION
Modeling long range dependencies on sequential data is crucial step towards closing the gap with human-level performance on many tasks. Attention based models like Transformer (Vaswani et al., 2017) have proven to be strong choice of backbone architecture for considerable number of tasks across modalities and scale (Devlin et al., 2019; Brown et al., 2020; Dosovitskiy et al., 2021). Vanilla Sett-Heead-Attemiie famously incurs 52 ?2(L²) penalty in modeling a sequence of length L. This is prohibitive at best for tasks where the model is required to capture long range dependencies from various parts of the input. Over the years, a variety of improvements have been proposed to alleviate this quadratic complexity (Tay et al., 2020; Choromanski et al., 2021; Ramsauer et al., 2021; Wang et al., 2020; Katharopoulos et al., 2020; Peng et al., 2021; Ainslie et al., 2020; Zaheer et al., 2020; Beltagy et al., 2020; Dai et al., 2020; Kitaev et al., 2020; Vyas et al., 2020).
On somewhat orthogonal direction, aitention-free models based on state spaces, such as S4 (Gu et al., 2022a) and DSS (Gupta et al., 2022), have shown remarkable improvements on Long Range Arena (LRA) (Tay et al., 2021), benchmark designed with long range modeling its focus and consists of diverse tasks with 1k- 1k-16k sequence length across modalities. These models require careful initialization, originally borrowing ideas from the theory of HiPPO matrices (Voelker et al., 2019; Gu et al., 2020), to achieve good results on LRA.
In this work, explore and extend the of state space models by focusing solely the task of autoregressive sequence modeling (Brown et al., 2020; Rae et al., 2021; Chowdhery et al., 2022; Zhang al., 2022; Hoffmann al., 2022: Srivastava et al., 2022). Several key properties endowed by the state space model family makes it particularly attractive in the context of language modeling. First, it reduces the ()L2) 2 complexity
input sequence length to O(L log L). This complexity results from the use of Fast Fourier Transform (FFT) (Cooley & Tukey, 1965) for performing convolutions, which will be described in detail in later sections. Second, the state space model is fully parallelizable in the length dimension. This is an arguably subtle but an important property at training time. Note
Work done Tei-Aviv University

---

Published as a conference paper at ICLR 2023
Input
Dense
Dense
def ges (x, F=4096, 1-4096, E-1024, B-256):
shortcut, x x, norm (x)
Gated State Space
dense (x, F, activation='gelu" = dense (x, H, activation='gelu y dss (u, H, I)
Space
# yhl, ..., yhh are linear in whl, uhb uc dense (y,
Dense
= dense (uc V, E) return shorteut
Dense
Figure I: (a) Our proposed Gated State Space (GSS) layer, (b) Pseudocode for GSS (full implementa- tion in $A.2).
that transformers are also fully parallelizable, a worthy advantage over traditional RNNs for modeling sequences, which otherwise incurs only an O(L) penalty. While this parallelism is useful at training time, it may also curse at inference time where decoding every token requires attending to the whole past. The ideal model is parallelizable at training time but incurs a small constant cost (per decoded token) at inference time. This brings us to the final point. Due to the inherent convolution- recurrence equivalence of the state space model, it can be made accumulate state and unroll like an RNN at inference time without any approximations.
Despite these attractive properties, found that current state space models (c.g. S4. DSS) run slower than we expected on TPUs, our accelerator of choice. We take this opportunity to modify the architecture to reduce dimensionality of specific bottleneck operations. Our proposed changes borrow from a well-supported empirical success of gating units (Shazeer, 2020). Specifically, Hua et al. (2022) observed that replacing the Ferward layer in the Transformer with gating units allows for a reduced dimensionality when mixing tokens along the length dimension using self-attention. We extend the use of gating units state space model family and observe that, even in our context, the use of gating units allows for reduction in dimensionality when performing FFT operations, which we observed to be the main bottleneck behind slow training. Furthermore, somewhat contrary to observations made by S4 and DSS authors, we found the performance on language modeling tasks to much less sensitive to initialization: only the scale and structural aspects of initialization of state space variables were important and not the exact values. We were able to successfully train the model while initializing the state space variables randomly. This departs from the reliance of the design on the theory of HiPPO matrices, which led the S4 model to employ several numerical linear algebra tricks to able to make it work. Combining both of these contributions, we propose a layer named Gated State Space (GSS) (Figure 1), which we empirically verified to be 2-3x faster than DSS while keeping the perplexity on several language modeling benchmarks (Table 1).
Going one step further, we also perform comparison with well-tuned and performant baselines reported in Block Recurrent Transformers (Hutchins et al., 2022), on several long range language modeling benchmarks over modalities such English books, raw source code from Github and LaTeX source of ArXiv mathematics articles. As detailed in Table 2, while our GSS model currently lags behind on some tasks when compared in the fixed-parameter setting, it competitive in the fixed-compute setting where we measure compute as the amount of TPUv4 hours spent training, which is good proxy for the cost of training that model. Furthermore, we also experimented with hybrid model in which we sparingly interleave Transformer layers (having local attention) in a GSS stack to allow for a richer modeling of short range interactions. To our delight, this further improves performance at (roughly) no extra training cost, both in terms of parameters and compute.
In our experiments we train on sequences of length at most 4k, but evaluate on a wide range of sequence lengths up to 65k. The performance actually improves the sequence length is increased, suggesting that GSS utilizes extra context despite not being trained with this context. Further, at inference time, state space models including GSS are quite efficient since decoding happen in recurrent mode (as much as 60x better in the case S4 (Gu et al., 2022a)). Though, the hybrid model which also uses local attention complicates this advantage bit.
2

---

Published as a conference paper at ICLR 2023
In summary, we propose GSS, simple-to topple--t-implement alternative to S4 and DSS which trains 2- 3x faster, and is competitive with Transformer-bas based baselines on long-range language modeling benchmarks.
RELATED WORK
In recent years, attention- based models have emerged dominant technique for sequence modeling, achieving remarkable improvements in a wide range of tasks, starting in NLP (Vaswani et al., 2017; Devlin al., 2019; Radford et al., 2019; Liu et al., 2019), then moving other classical machine learning areas such as computer vision (Dosovitskiy et al., 2021) and now to the physical sciences (Avsec al., 2021; Jumper et al., 2021). In brief, an attention layer takes as input three matrices K, Q, V in IR Lxd which should usually be thought of as length L lists of d-dimensional vectors. The attention layer then outputs Y = =(QK' V V E
I fxd where a indicates a row-wise softmax operation. In the popular self-attention variant, K, Q and are themselves learned functions of a single input sequence X = (s,,..., ..) E R Lxd Unfortunately, this requires O time and space due to the need to construct the matrix QK T E RR
The surge in popularity of attention has engendered corresponding surge in interest on methods for increasing the context length while controlling the O(L²) computational cost. Broadly speaking, most approaches fall into two camps: those that attempt to "linearize" attention, and those that sparsify the attention matrix QK The first camp exploits the fact that, when softmax is removed, we have (QK")r = Q K V), where now KTV V E dRd and so the whole operation is only linear in L rather than quadratic. This idea is the underlying principle behind methods such as the Performer (Choromanski et al., 2021), Linear Attention (Katharopoulos et al., 2020), Random Feature Attention (Peng et al., 2021) cosFormer (Qin al., 2022). In the other camp, one simply sparsifies the matrix QK limit computation and storage requirements. This camp includes BigBird (Zaheer et al., 2020), GMAT (Gupta & Berant, 2020), Longformer (Beltagy al., 2020), and Blockwise self-attention (Qiu et al., 2020), Reformer (Kitaev et al., 2020) and Perceiver AR (Hawthorne et al., 2022). These approaches can be viewed as a trade-off between performance and computation time (or memory).
Despite its current empirical success, the attention mechanism is not the only approach for modeling sequence data, nor even necessarily the most natural. For example, classical recursive or state space layer operates on an input (Zi...., II) R 4x by defining a sequence of "states" st+1 = T(s,, and returning the output sequence y/ E(s,) where T and E are learned "transition" and "emission" functions. Assuming T and E can be computed efficiently, this requires only O(L) time in theory. Through the years, many different possibilities for T and E have appeared, such as the simple RNN, which sets 7 and E to both be MIPs, or more complicated LSTM layer (Hochreiter & Schmidhuber, 1997). Nevertheless, the performance of state space models has in many cases been quickly surpassed by attention.
There have been several recent efforts recapture the favorable properties of state space models while maintaining or improving upon the performance of attention models, such as the Transformer-XL. (Dai et al., 2019), SRU++ (Lei, 2021), 1), or Block-RRecurren Transformers (Hutchins et al., 2022). Of particular relevance to this paper, alternative models based on linear dynamical systems have shown great promise (Gu al., 2022a; Gupta al., 2022; Gu et al., 2022b). Rather than attempting to "fix" attention, these are classical state space models where both T and F. are linear. An important insight in these works is to recast the model as a convolution with very large kernel, and then leverage the convolution theorem to compute the sequence ye in (o)))) time using Fast Fourier Transform (Cooley & Tukey, 1965). While seemingly worse than OL), this operation is easily parallclizable and so in practice is significantly faster. Moreover, in situation in which the log( L)) factor is onerous, one may still fall back to the serial OLL) algorithm. Our approach will build on these works.
3 METHOD
We start by reviewing the necessary background on state spaces required to fully describe our model 1). We then formally define our GSS model in $3.2 and the GSS--Trnsformer-HHyrr in $3.3.
3

---

Published as a conference paper at ICLR 2023
3.1 STATE SPACE PRELIMINARIES
While in this work we are interested in modeling sequences of vectors, let us first review how state space models define sequence-to-sequeece map for 1-D sequences.
State Spaces A discretized state space, parameterized by a state matrix RN*N vectors B E IR Nx1 ,C' I2" IxN and a sample time A E Rso defines sequence-to-sequence map from input e R to output (yo....
Bon....,,--111 =
..0............................. ,.L-11 E Rf via the recurrence',
Tk Ara-1 Buk
yk Crr
(1)
A B = (A- I)A C=C. Assuming x-1 0 for simplicity, the above recurrence can be explicitly unrolled
A CA uk-j
Vk
Σ
(2)
3=0
For convenience, the scalars C.A' k B are gathered to define the SSM kernel K E R as
K (CB,CAB,..., =(
(Ce ( (e -A T)A=BB Dosk<L>
where the last equality follows by substituting the values of A, B,C from Equation 1. Hence,
Yk Σ K Uk-s
1=
where K3 denotes the value of the kernel at position Given input sequence E TR. it is possible to compute the output E TR. sequentially via the recurrence in Equation 1. While this property highly desirable for autoregressive decoding, sequential computation is prohibitively slow to train with long inputs and, instead, Equation can be used to compute all elements of in parallel, provided we have already computed K.
Computing y from u and K is easy. Given an input sequence RR_ and the SSM kernel K TR naively using Equation 4 for computing y would require O(I²²) multiplications. Fortunately, this can be done much more efficiently by observing that for the univariate polynomials
K(=)
Σ K K and u(z) 1=0
Σ 1-0
yk is the coefficient of ² in the polynomial K(2) u(()) i.e. all yk 's can be computed simultaneously by multiplying two degree L polynomials. It is well-known that this can be done in O log( L)) time via Fast Fourier Transform (FFT) (Cormen et al., 2009). We denote this fast computation of Equation 4 via the discrete convolution
K * u.
(5)
Diagonal State Spaces The challenging part is computing itself as involves computing L distinct matrix powers (Equation 3). Gupta et al. (2022) observed that the state matrix A can be assumed to be diagonal without loss in performance, thereby allowing straighforward computation of K. Their DSS model (DSS from hereon) assumes A to be a diagonal matrix and assumes B (1) DSS has parameters Are, Aim E R.N. C. C and A R. The diagonal of (i.e. (AI,...,AN)) is computed as -clementwise-xpp(A.. + Ajm where V-I
Discrretizatonn used in Equation I relies on zero order hold (ZOH) which assumes constant value between sampled times. We also experimented with other discretization methods such as bilinear transform but did not much difference in performance.
4

---

Published as a conference paper at ICLR 2023
and A is computed exp( oip(Aing) Ryo. For this parameterization, the kernel (Equation 3) can be computed as a matrix-vector product
K (C ((e - isiss 1IN elementwwise-xxp(Psal
(6)
where Pi,k A and denotes elementwise multiplication.
This formally defines linear sequence-tee peouuenne-tt-eequeee to-sequence map for 1-D sequences. In the case of sequences of H-ddimensional vectors, state space models are applied individually on the H features as follows.
Each DSS layer receives a length-L sequence u of H H-dimensional vectors as input, i.e., u E IR HxL and produces output y E IR #xx. The parameters of the layer are Ares Aim IR Alog R and E " For each coordinate
1,..., H, a state space kernel K₁ E R computed for Are: Aim, (Atog) C₁, via Equation 6. The output Uh E R for coordinate is computed from UA. E R and K2, using Equation 5.
For batch size B, sequence length L and hidden size H, the DSS layer requires O(N time to compute the kernels, and O O(BHL Log(()) time for the discrete convolution. In S4 and DSS, the authors suggest to use N 64 as default and for this choice of N, the time taken to compute the kernels becomes an important factor, specially for small batches.
3.2 GSS LAYER
We build the idea of the recently introduced Gated Attention Unit (GAU) (Hua al., 2022) and replace the Q(T,) T,²) attention used in GAU by a further simplified DSS layer ($3.1). Gating allows our model to be contextualized over reduced dimensionality and the use of state spaces provides it with superior contextualizing abilities, while enjoying O(L log L) complexity out the box.
Concretely, our Gated State Space ("GSS") layer maps a sequence X R LxE
E R LxE
to an output O E
U =(WW) E IXH
== VXX(), IXF
(7)
Y DSS(U) R LXH
U Ucottxx
YW, e R Rxx
(8)
O (Ucontes V)W, E R LXE
where * denotes elementwise multiplication and á denotes GELU (Hendrycks & Gimpel, 2016). Note that, contrary to Hua al. (2022), in our experiments, we did not observe much benefit from using RELU or Swish activations instead GELU.
In Equation 8, contextualized representation of U is formed via DSS in time. In most of our experiments we used = 5/4 = 256, = 512 and F' = 4096. Similar to GAU, gating allows us to perform a weaker contextualization using the state space layer by reducing the dimension of the input to DSS, by 4x in our case. This offers a much needed increase throughput since observed that FFT ops were the main bottleneck (measured TPUs). As shown in Table 1, this provides more than 3x speedups over the vanilla DSS block described in $3.1.
log L)
Simpified DSS used in GSS layer In Equation instead of directly using DSS as described in $3.1, we made several key simplifications based on our exploratory results.
As described in $3.1, S4 and DSS train a separate for each the H coordinates resulting in the creation of H separate P matrices in Equation 6. In the DSS layer used in GSS, we decided to eliminate the use of A by fixing it as This reduces the compute required for creating the kernels and makes the kernel computation simpler.
Secondly, we found that randomly initializing the parameters corresponding to works just as well as the Skew-Hippo initialization suggested in (Gupta et al., 2022). Importantly, we parametrize both real and imaginary parts of A in log space, so that the random initializations span many orders of magnitude ($A.2). The effectiveness of random initialization is in contrast to the findings of Gu et al. (2022a) and Gupta et al. (2022) who reported their respective initializations the state space parameters to be critical to the model performance. We do however note that the experiments in
5

---

Published as a conference paper at ICLR 2023
setting of large-scale language modeling conducted orders of magnitude larger scale of compute than what is used in the tasks considered these works. Moreover, the data modalities we consider in this work are different from the tasks considered in these works (e.g. Long Range Arena) which are specifically designed to stress test long-range reasoning.
3.3 GSS- DASTTRANSFORMER HYBRRD
Conceptually, GSS looks fairly different from the current workhorse of machine learning; the Transformer architecture. Given this, it not immediately clear if one is decidedly better than the other or if they each provide some orthogonal benefits. If its the latter, one might wonder if there are synergies between these architectures which can be exploited create a hybrid model which is stronger than either one of them individually. To that end, we also consider conceptually simple hybrid between GSS and Transformer where we sparingly interleave traditional Transformer blocks with GSS layers. Despite its glaring simplicity, shown in Table 2, observed that it shows consistent and significant improvements on all tasks.
used sequence lengths large enough to be prohibitive for traditional Transformer layers. To get around this restriction at the Transformer layers used in hybrid model, chunk their inputs into non-overlapping chunks of length 512 and run the Transformer layer on each of them independently.
Chunking long inputs In all our experiments
4 RESULTS
We conduct experiments with GSS 4 different datasets, LMIB, PG19, ArXiv and Github, each of them varying qualitatively with another in terms of modality and average document length.
LMIB is standard and reasonably big dataset (1B tokens) where each training example consists of a short sentence (Chelba et al., 2014). This is different from rest of the datasets which consists of a much larger sequence of tokens per example. Although, our primary goal is measure GSS's ability to capture long range dependencies, we include results on LMIB benchmark and compare with vanilla Transformer baseline, which is hard do for larger sequence lengths.
PG19 dataset is constructed from extracting a large collection of full-length books from Project Gutenberg (Rae et al., 2020). All extracted books were written before 1919 and only contain tokens from the English language. Over the years PG-19 has become a standard benchmark for measuring progress on long range dependency modeling over text.
ArXív Math dataset was recently collected by Wu et al. (2022) and contains LaTeX source for articles focusing on Mathematics. Even though articles are typically shorter than full-length books, as in PG19, since LaTeX source can have many special characters, typical sub-piece vocabularies are not very effective and tends to produce examples of similar size, as measured by number of tokens. is possible to train custom sub-piece vocabulary for this dataset, but we stick with the vocabulary used by both Wu et al. (2022) and Hutchins et al. (2022) for fair comparison.
Github was also first collected and used by Wu et al. (2022). It is corpus of raw source code collected from several Github repositories with open-source licences. The dataset contains code from several programming languages, including C, C++, Java, Python, Go and Typescript. In this case individual files can be small but code repositories are typically is organized so that code can he reused across file boundaries. Thus, for every repository, all the files were concatenated produce single document.
4.1 COMPARING DSS AND GSS
As shown in Table 1, in our first set of results, we compare DSS and GSS models both in terms of perplexity and throughput, as measure by steps per second on all 4 datasets. For the PG-19, Arxiv and Github datasets, all models were trained with 4k sequence length at training time. However, at evaluation time, we present results for large range of sequence lengths. Despite being a (slightly) smaller model, GSS improves perplexity while being 2-3x faster depending on these datasets. In contrast, for LMIB, we train and evaluate with sequence length of 512, since the dataset only contains short sentences, most examples comfortably fully fitting in the context window.

---

Published as a conference paper at ICLR 2023
Eval Sequence Length
Dataset
Model
Params
Throughput (steps/sec)
Perplexity
512
4k
16k
65k
Transformer
182M
6.6
13.51
LM1B
DSS
188M
1.8
13.59
GSS
190M
5.6
13.26
Transformer
182M
6.6
13.00
PG-19
DSS
209M
1.8
14.51
13.53
692.1
OOM
GSS
192M
5.3
14.01
12.84
12.94
12.47
Transformer
182M
6.6
3,27
Arxiv
DSS
209M
3.65
3.13
284.6
OOM 2.75
GSS
192M
5.3
3.57
3.08
3.08
Transformer
182M
6.6
2.88
Github
DSS
209M
1.8
3.65
3.21
242.7
OOM
GSS
192M
5.3
2.68
2.35
2.31
2.12
Table Comparison of DSS and GSS models fixed-param setting. We consistently find that GSS outperforms DSS (with hyperparameters taken from Gupta et al. (2022)) while being 2-3x faster on all tasks. We train both GSS and DSS with sequence length 4k except for LMIB. Training sequence length is 512 for the Transformer baseline for all datasets.
Furthermore, we observe significant generalization over changes in sequence length on all the datasets. Not only does the performance not degrade when increasing sequence length, actually improves quite bit! This suggests that the model is effective at utilizing the extra context even though the model was not trained with that amount of context. Note that used default hyperparameters suggested by Gupta et al. (2022) for initializing state space variables for DSS. But, for GSS, since we made several structural changes, we retuned the hyperparameters related initialization on PG19 alone and used them for all the datasets. Since we had access to evaluation metrics for all sequence lengths, length generalization factored in for our choice of hyperparameters, which not true for DSS. It may be possible to similar length generalization even with vanilla DSS if initialization hyperparameters were retuned.
Similar to language modeling experiments in (Gu et al., 2022a), every block of DSS baseline consists of DSS layer followed by GLU (Dauphin et al., 2017) and a Feedforward layer similar to the one used in Transformer block with GELU activation (Hendrycks & Gimpel, 2016). DSS baseline consists 12 layers and an embedding dimension 1024. For GSS, we increased the number of layers to to match the parameter count.
4.2 COMPARISON WITH OTHER BASELINES
In this section, we turn our attention towards apples-to-apples comparison of GSS versus well-tuned baselines on these datasets. For complete picture of the cost of training these models, we report both the number of parameters and time spent training as measured by TPUv4 hours. For baselines, we selected the best performing models reported (Hutchins al., 2022) for every dataset and compare with GSS model both in fixed-paran and fixed-compute settings.
GSS models GSS consists of 16 layers and an embedding dimension of 1024. We also consider larger variant with 32 layers as denoted by GSS-L. For GSS-Hybrid model, we used vanilla Transformer blocks at every 4th layer starting with the 2nd layer. Since GSS layers are inherently position aware, using them for the 1st layer eschews any need of explicit position embeddings typically used with otherwise position invariant Transformer blocks. Thus, barring position aware nature of GSS layers, don't 't use any kind of explicit position embedding bias in models. For the Transformer blocks used in hybrid models, we use multi-head self attenttioo with 8 heads, each with size 128.
7

---

Published as a conference paper at ICLR 2023
Eval Sequence Length
Dataset
Model
Params TPUv4 hours
Perplexity
512
16k
65k
Rec:fixed
196M
0.8k
11.55
PG-19
Feedback:lstm:singge
196M+
0.8k+
11.31
Feedback: fieedsckkfixed.skipp
196M+
0.8k+
11.24
GSS
(this work)
192M
0.5k
14.01
12.84
12.94
12.47
GSS-L
352M
0.8k
12.48
11.33
11.16
11.12
GSS-Hybrid-L
373M
0.8k
11.45
10.52
10.44
10.1
Rec:fixed:skip
196M
0.8k
2.36
Arxiv
Feedback:
196M+
0.8k+
2.33
Feedback.fixed:skjj
196M+
0.8k+
2.36
GSS
192M
0.5k
3.57
3.08
3.08
2.75
GSS-L
352M
0.8k
3.29
2.71
2.72
2.51
GSS-Hybrid-L.
373M
0.8k
2.94
2.51
18.27
125.2
sec:fixed:skpp
196M
3k
2.04
Github
Feedback:lstm:singg
196M+
3k+
2.07
fix
196M+
3k+
2.16
GSS
192M
0.5k
2.68
2.35
2.31
2.12
GSS-L
352M
1.8k
2.31
1.99
2.20
2.28
GSS-Hybrid-L.
373M
1.8k
2.34
1.88
1.74
2.09
Table 2: Comparing GSS variants with best-performing models from Hutchins et al. (2022) in both fixed-param and fixed-compute settings. While, in the fixed-param setting, GSS currently lags behind performance of best Block Recurrent Transformer baseline, it is fairly competitive in the fixed-compute setting (as measured by total TPUv4 hours taken to complete training). For the larger model GSS-L used for fixed compute comparison, we simply double the layers from 16 to 32 keeping everything else fixed. In addition, due the inherent recurrent view of state space model family, decoding at inference time is much faster than Transformer based baselines. For block-recurrent baselines, adding feedback increases both parameter count and training time, stick with conservative estimates derived from Hutchins et al. (2022) for both, which we denote by '+'. Hutchins al. (2022) report log2 of perplexity which we convert raw perplexity in this table. Param count for all the models include embedding layers as well. A comparison with FLASH (Hua al., 2022) is included Appendix (Table 4).
Baselines We considered 3 eigh-performing baselines and numbers reported in (Hutchins et al., 2022). Block Recurrent Transformer leverages recurrence over blocks Transformer layer to model long range dependencies. Hutchins et al. (2022) performed comprehensive exploration of open design space of incorporating recurrence across blocks. Somewhat surprisingly, Rec:fixe:skip. which accumulates recurrent state vector better than more complicated gating designs. Another variation which performed well with Block Recurrence the idea of adding feedback mechanism blocks similar to Feedback Transformer (Fan et al., 2020). Note that feedback mechanism makes training more expensive due to additional cross-attention modules and corresponding paramaters.
exponential moving average time performs
Fixed- Fixed-param param comparison As shown in Table 2, we see that GSS variants come very close but quite beat the strongest block recurrent baseline in the fxee-param setting. In this case, we are comparing GSS model which has roughly 192M parameters (including embeddings) with the baselines all which have around 196M parameters. Even though the parameter count fixed, we that GSS runs faster than block recurrent baselines, likely due the fact that all the layers be completely parallelized unlike the the recurrence in the baselines which run sequential fashion.
fixed-compute comparison Since the GSS model runs faster than the baselines, also train with versions larger than GSS such that the training time (as measured by total TPUv4 hours) matches the time taken by the baselines. We simply double the number of layers from 16 to 32 to construct GSS-L. As expected, adding more parameters improves perplexity numbers on the eval set all the datasets. Moreover, find that the GSS-Hybrid versions of the model outperform the best baseline
8

---

Published as a conference paper at ICLR 2023
Model
Params
Layers
WLP
Vocabulary
Transtormer-Xi XL (Rae et al., 2020)
36
36.3
32k
Compressive Transformer (Rae et al., 2020)
36
33.6
32k
Routing Transformer (Roy al., 2021)
490M
22
33.2
98k
Perceiver AR (Hawthorne et al., 2022)
974.6M
60
28.9
32k
Block Recurrent Trans (Hutchins et al., 2022)
1.3B
24
26.5
32k
341M
32
32,9
32k
68 M
64
30.2
32k
GSS-Hybrid (ours)
IB
96
28.0
32k
1.4B
128
27.3
32k
Table 3: Comparison of GSS with other competitive models on PG-19 across multiple scales as measured by word-level perplexity (WLP) on test set of PG-19. We evaluated WLP on 16k sequence length for GSS-Hyybrid set of models. Model size numbers for Perceiver AR (Hawthorne et al., 2022) and Routing Transformer (Roy et al., 2021) are taken from Hutchins al. (2022).
model on PG-19 and Github datasets. We do see significant improvements for Arxiv dataset as well but unfortunately not enough to be stronger than the baseline. We think this may be resolved by the use of a vocabulary more suited to Arxiv symbols. On the flip side, we can no longer do a fair stick with the vocabularies
comparison with token level perplexity if change the vocabulary, used by the baselines for this study.
Length generalization We train all variants of GSS with sequence length of 4k but evaluate different lengths I E [512, 4k, 16k, 65k.]. On PG-19, we see significant length generalization across the board. Not only the performance doesn't 't degrade as the sequence length is increased but gets significantly better for all model variants. On Arxiv and Github, the situation is a little more complicated. For smaller models, we still see length generalization but it tends to degrade when either the model is made bigger or the dataset has lot of noise (as indicated by variation in perplexity metric over time). How to robustly achieve length generalization an interesting research question on it is own and we believe one can design interventions which can lead to further improvements, which we don't explore here.
Note that block recurrent baselines, with or without feedback mechanism, process documents in sequential fashion such that recurrent state from previous segment from the document is passed the next segment, with backward pass truncated to the current segment. This means that, even though the segment sequence length is set to 4k, block recurrent models have (arguably weak) access to almost the entire past. Thus, perplexity comparison at sequence length 4k is slightly unfair towards GSS models since they do not employ such state caching mechanisms.
4.3 SCALING UP
In this section, we further scale up the GSS model focusing only on PG19 dataset. As shown in Table 3, even though GSS does not beat current state of the art on fixed-param basis, we consistently observe significant gains from simply scaling up, similar to popular Transformer based language models like GPT-3 (Brown et al., 2020) and PaLM (Chowdhery al., 2022). Furthermore, GSS also outperforms Routing Transformer (Roy et al., 2021) and Perceiver AR (Hawthome et al., 2022) when comparing similar sized models.
CONCLUSION We introduce GSS, a general purpose sequence model which leverages gated units and trains significantly faster as shown several language modeling benchmarks. Further comparison with well tuned Transformer baselines suggests that GSS is fairly competitive in fixed-compute setting. We further show that hybrid model constructed by interleaving GSS and Transformer improves performance even further. We however recognize that perplexity numbers alone are not satisfying for gauging overall model performance. As part of future work, will explore scaling GSS even further and conducting few-shot evaluation on standardized benchmarks, comparing it against other popular large language models.

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGMENTS
We are grateful to the developers of Jax and Flax libraries. In addition, we would like to thank DeLesley Hutchins and Imanol Schlag for answering our questions. sharing the datasets and helping us in general with details of Block Recurrent Transformer paper. Finally, we are grateful to Ethan Dyer for discussions on long context modeling, Walid Krichene, Albert Gu, John Anderson and anonymous reviewers for providing feedback on an earlier draft of this paper. Ashok Cutkosky acknowledges support from NSF grant CCF-2211718.
BIBLIOGRAPHY
Daniel Adiwardana, Minh-Thang Luong, David R. So, Jamie Hall, Noah Fiedel, Romal Thoppilan, Zi Yang, Apoorv Kulshreshtha, Gaurav Nemade, Yifeng Lu, and Quoe V. Le. Towards human-like open-domain chatbot. ArXiv, abs/2001.09977 2020.
Joshua Ainslie, Santiago Ontanon, Chris Alberti, Philip Pham, Anirudh Ravula, and Sumit Sanghai. Etc: Encoding long and structured data in transformers. ArXiv preprint, abs/2004.08483, 2020.
URL /ttps::/arxiv...r org/ abs/20000/0// 2004. 08483.
Ziga Avsec, Vikram Agarwal, Daniel Visentin, Joseph R Ledsam, Agnieszka Grabska-Barwinskk. Kyle R Taylor, Yannis Assael, John Jumper, Pushmeet Kohli, and David R Kelley. Effective gene expression prediction from sequence by integrating long-range interactions. Nature methods, 18 109:1139-220,, 2021.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. ArXiv preprint, abs/2004. 005550, 2020. URL https: http//aaxvvv....../..e.. abs 0200//00000050
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agar- wal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCan- dlish, Alec Radford, Iyy Sutskever, and Dario Amodei. Language models are few-shot learners. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Ties Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL https: http:///ooooeeeeneeeeoeooo neurips wttpppper/002////////eocp /aahh/ 96741677777777 Bbwbaaccllccccccceecccaacbaaabla Abstttt html.
Ciprian Chelba, Tomas Mikolov, Mike Schuster, Qi Ge, T. Brants, Phillip Todd Koehn, and Tony Robinson. One billion word benchmark for measuring progress in statistical language modeling. In INTERSPEECH, 2014.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamás Sarlós, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, David Ben- jamin Belanger, Lucy J. Colwell, and Adrian Weller. Rethinking attention with performers. In 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL StRLtttts////peeeeeeeee http:////eeeeeieeeeee net/ forumorrummommm ?id= Ua 62ukkwwR.
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek B Rao, Parker Barnes, Yi Tay, Noam M. Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Benton C. Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier García, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Oliveira Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi
10

---

Published as a conference paper at ICLR 2023
Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathleen S. Meier- Hellstern. Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways. ArXiv preprint, abx/2204.02311, 2022. URL https: httppp///axxxv..ro/////o org/ /bs/ 2204. 02311.
James W. Cooley and John W. Tukey. An algorithm for the machine calculation of complex fourier series. Mathematics of Computation, 19:297-301, 1965.
Thomas H. Cormen, Charles E. Leiserson, Ronald L.. Rivest, and Clifford Stein. Introduction to Algorithms. The MIT Press, 3rd edition, 2009.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdinov. T: Attentive language models beyond fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 2978-2988, Florence, Italy, 2019. Association for Computational Linguistics. doi: 10.1865//1//119-122 URL httph:ttp///mmm htt/a/aaaaannnoacaaaaaaaaaaaleee org/ P12 1285.
Zihang Dai, Guokun Lai, Yiming Yang, and Quoc Le. Funnel-transformer: Filtering out sequential redundancy for efficient language processing. Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Pro- cessing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6- 12, 2020, virtual, 2020. URL https:::ppocceeenggggg..///e neur neurips cc/ paper/2020/han hww.///aahhaaaaaaa cd2 6954 690 644eee444222222e11555555
Abstract .tmll.
Yann Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In Proceedings of the 34th International Conference Machine Learning Volume 70, ICML'CI 17, pp. 933-941. JMLR.org, 2017.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pp. 4171-4186 Minneapolis, Minnesota, 2019. Association for Computational Linguistics. doi: 10. 10.18653/v1/N1191422. URL https:// aclanthology org/N19g 9
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition scale. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https: https:/ /tpp///eeeeeeeeeeeeeeeeee/ee net forum?ids VicbFaNTTy:
Angela Fan, Thibaut Lavril, Edouard Grave, Armand Joulin, and Sainbayar Sukhbaatar. Addressing some limitations of transformers with feedback memory. ArXiv preprint, abs/2002.09402, 2020. URL https:://////..... //pp:://arxas http/axxvvv.////// org/ abs/ abs/2202. 09402.
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan- -ienn Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL https: http://proeeeeng...o.o/ nwurips. cct.ccc/ppcae cc/paper/////e/eecce 2020/ / /wwhh/ahhhhatt
6 wtap://oddddieeeeabaaaaattteeeee
html.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In International Conference Learning Representations, 2022a. URL https: //openreview net/ forum?
Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models. In Annual Conference Neural Information Processing Systems (NeurIPS), 2022b.
Ankit Gupta and Jonathan Berant. Gmat: Global memory augmentation for transformers. ArXiv preprint, hhs/2000.032244 2020. URL https: //axiiv http://wwwwww..g/////////// /006 03274.
11

---

Published as a conference paper at ICLR 2023
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces effective as structured state spaces. In Annual Conference on Neural Information Processing Systems (NeurIPS), 2022.
Curtis Hawthorne, Andrew Jaegle, Cätälina Cangea, Sehastian Borgeaud, Charlie Nash, Ma- teusz. Malinowski, Sander Dieleman, Oriol Vinyals, Matthew Botvinick, lan Simon, Hannah Sheahan, Neil Zeghidour, hean-Baptiste Alayrac, Joao Carreira, and Jesse Engel. General- purpose, long-context autoregressive modeling with perceiver AR. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceed- ings of the 39th International Conference ings of Machine Learning Research, pp. 8535-8558. PMLR, 17-23 Jul 2022. URL https: //proceedings.mlle press / vv 62/ http://ll/tttttttteeeeeeeeeeeeee hwwthorne22 html.
Machine Learning, volume 162 of Proceed-
Dan Hendrycks and Kevin Gimpel. Bridging nonlinearities and stochastic regularizers with gaussian error linear units. ArXiv preprint, abs/1606.08415,, 2016. URL httpa://axxiv. org/abs/////ooaaoooo/..ooo..osas 1606. 015.5.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):
1735-1780, 1997.
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hennigan, Eric Noland, Katie Millican, George van den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Elsen, Jack W. Rae, Oriol Vinyals, and L. Sifre. Training compute-optimal large language models. ArXiv preprint, 155522233115556 2022. URL / wiw...vvv.oo. org/ alss 2203 15556.
https:
Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc V. Le. Transformer quality in linear time. In Proceedings of the 39th International Conference on Machine Learning (ICML), 2022. URL https::/arxxt org/ wws.aas/////ooooo.o/ooooo.oooa.r 2 10447.
DeLesley Hutchins, Imanol Schlag. Yuhuai Wu, Ethan Dyer, and Behnam Neyshabur. Block-recurrent transformers. ArXiv preprint, abs/2203.07852. 2022. URL https: https::/arssv.og///eee 07852.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Židck, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature, 596(7873):583 2021.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 5156-5165. PMLR, 2020. URL http: //ppoceedings... mlr press vt1p/vl11/////// http://aaharpoolsssaees html..
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http: // ///axxiv org/ abs 412...22 6980.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL https: /ttp://///eeeeeeeeeeeeeee net forum?id= vKg
Tao Lei. When attention meets fast recurrence: Training language models with reduced compute. In EMNLP, 2021.
Yinhan Liu, Myle Ott. Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. ArXiv preprint, abs/1907.11699 2019. URL https: /wtp//aaivvv/v..///.a//e.aaaaaaa www.ii..rrrrbbb....e.. 907. 692.
12

---

Published as a conference paper at ICLR 2023
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A. Smith, and Lingpeng Kong. Random feature attention. In 9th International Conference on Learning Represen- tations, ICLR 2021, Virtull Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https: htttss://peeneeee ew. net/ forum? F Te
Zhen Qin, Weixuan Sun, Hui Deng, Dongxu Li, Yunshen Wei, Baohong Lv, Junjie Yan, Lingpeng Kong, and Yiran Zhong. cosformer: Rethinking softmax in attention. In Interaational Confer- ence on Learning Representations, 2022. URL https: htttps://operreview www.:/eee/foumm/////e forum?id- B1 B18C@xx2Up4.
Jiezhong Qiu, Hao Ma, Omer Levy, Wen-tau Yih, Sinong Wang, and Jie Tang. Blockwise self- attention for long document understanding. In Findings the Association for Computational Linguistics: EMNLP 2020, pp. 3555-2565, Online, 2020. Association for Computational Linguis- tics. URL https: /aclantooannotoyoooooooooao htwpoooooyyogg//22/2 org/2020 fiindings-emmn 232.
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.
Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Chloe Hillier, and Timothy P. Lillicrap. Compressive transformers for long-range sequence modelling. In 8th International Conference on Learning Representations, ICLR 2020, Addis Abuba, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL htttps:://ppereeeeeee
net/ forum?idee Sarrummii=Sy KikSDDH.
Jack W. Rae, Sebastian Borgeaud, Trevor Cai, Katie Millican, Jordan Hoffmann, Francis Song, John Aslanides, Sarah Henderson, Roman Ring, Susannah Young, Eliza Rutherford, Tom Hennigan, Jacob Menick, Albin Cassirer, Richard Powell, George van den Driessche, Lisa Anne Hendricks, Maribeth Rauh, Po-Sen Huang, Amelia Glaese, Johannes Welbl, Sumanth Dathathri, Saffron Huang, Jonathan Uesato, John F. J. Mellor, Irina Higgins, Antonia Creswell, Nathan McAleese, Amy Wu, Erich Elsen, Siddhant M. Jay Jayakumar, Elena Buchatskaya, David Budden, Esme Suther- land, Karen Simonyan, Michela Paganini, L. Sifre, Lena Martens, Xiang Lorraine Li, Adhiguna Kuncoro, Aida Nematzadeh, Elena Gribovskaya, Domenic Donato, Angeliki Lazaridou, Arthur Mensch, Jean- Baptiste Lespiau, Maria Tsimpoukelli, N. K. Grigorev, Doug Fritz, Thibault Sottiaux, Mantas Pajarskas, Tobias Pohlen, Zhitao Gong. Daniel Toyama, Cyprien de Masson d' Autume, Yu- jia Li, Tayfun Terzi, Vladimir Mikulik, Igor Babuschkin, Aidan Clark, Diego de Las Casas, Aurelia Guy, Chris Jones, James Bradbury, Matthew G. Johnson, Blake A. Hechtman, Laura Weidinger, Ia- son Gabriel, William S. Isaac, Edward Lockhart, Simon Osindero, Laura Rimell, Chris Dyer, Oriol Vinyals, Kareem W. Ayoub, Jeff Stanway, L. L. Bennett, Demis Hassabis, Koray Kavukcuoglu, and Geoffrey Irving. Scaling language models: Methods, analysis & insights from training gopher. ArXiv preprint, abs/21 12. 446, 2021. URL https:ttt/// arxiv. org/ abs / 2112, 11 1446 6.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter Liu. Exploring the limits of transfer learning with unified text-to- text transformer. J. Mach. Learn. Res., 11:140:1-140:67, 2020. URL http: /ttp///mmm. httpm/mrroooooo.oo /rg/ papers/v21/20-00 html.
Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Lukas Gruber, Markus Holzleitner, Thomas Adler, David P. Kreil, Michael K. Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. Hopfield networks is all you need. In 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL
net. forum? d=
Aurko Roy, Mohammad Saffar, Ashish Vaswani, and David Grangier. Efficient content-hased sparse attention with routing transformers. Transactions of the Association for Computational Linguistics, 9:53-68, Feb 2021. ISSN 2307-387X. doi: a_00033, URL http:
//dx doi. wrg. 11 62 62 /tcc 000333.
Noam M. Shazeer. Glu variants improve transformer. ArXiv preprint, https: / /tt::///rxxii aaariiiiiiiv///ii org.obgg//..... a /200..0022222 05202.
2020. URI
Aarohi Srivastava, Abhinav Rastogi, Abhishek B Rao, Abu Awal Md Shoeb, Abubakar Abid, Adam Fisch, Adam R. Brown, Adam Santoro, Aditya Gupta, Adrià Garriga-Alonse Agnieszka Kluska, Aitor Lewkowycz, Akshat Agarwal, Alethea Power, Alex Ray, Alex Warstadt, Alexander W.
13

---

Published as a conference paper at ICLR 2023
Kocurek, Ali Safaya, Ali Tazarv, Alice Xiang, Alicia Parrish, Allen Nie, Aman Hussain, Amanda Askell, Amanda Dsouza, Ameet Annasaheb Rahane, Anantharaman S. Iyer, Anders Johan An- dreassen, Andrea Santilli, Andreas Stuhlmuller, Andrew M. Dai, Andrew D. La, Andrew Kyle Lampinen, Andy Zou, Angela Jiang. Angelica Chen, Anh Vuong, Animesh Gupta, Anna Gottardi, Antonio Norelli, Anu Venkatesh, Arash Cholamidavoodi, Arfa Tabassum, Arul Menezes, Arun Kirubarajan, Asher Mullokandov, Ashish Sabharwal, Austin Herrick, Avia Efrat, Aykut Erdem, Ayla Karakacs, Bridget R. Roberts, Bao Sheng Loe, Barret Zoph, Bartlomiej Bojanowski, Batuhan Ozyurt, Behnam Hedayatmia, Behnam Neyshabur, Benjamin Inden, Benno Stein, Berk Ekmekci, Bill Yuchen Lin, Blake Stephen Howald, Cameron Diao, Cameron Dour, Catherine Stinson, Cedrick Argueta, C'esar Ferri Ram irez, Chandan Singh, Charles Rathkopf, Chenlin Meng, Chitta Baral, Chiyu Wu, Chris Callison-Buree
Chris Waites, Christian Voigt. Christopher D. Manning, Christopher Potts, Cindy Tatiana Ramirez, Clara Rivera, Clemencia Siro, Colin Raffel, Courtney Ashcraft, Cristina Garbacea, Damien Sileo, Daniel H Garrette, Dan Hendrycks, Dan Kilman, Dan Roth, Daniel Freeman, Daniel Khashabi, Daniel Levy, Daniel Gonz' alez, Danny Hernandez, Danqi Chen, Daphne Ippolito, Dar Gilboa, David Dohan, D. Drakard, David Jurgens, Debajyoti Datta, Deep Ganguli, Denis Emelin, Denis Kleyko, Deniz Yuret, Derek Chen, Derek Tam, Dieuwke Hupkes, Diganta Misra, Dilyar Buzan, Dimitri Coelho Mollo, Diyi Yang, Dong-Ho Lee, Ekate- rina Shutova, Ekin Dogus Cubuk, Elad Segal, Eleanor Hagerman, Elizabeth Barnes, Blizabeth P. Donoway, Ellie Pavlick, Emanuele Rodolà, Emma FC Lam, Eric Chu, Eric Tang. Erkut Erdem, Ernie Chang. Ethan A. Chi. Ethan Dyer, Ethan Jerzak, Ethan Kim, Eunice Engefu Manyasi, Evgenii Zheltonozhskii, Fan Xia, Fatemeh Siar, Fernando Mart inez-Plumed, Francesca Happ's e. François Chollet, Frieda Rong, Gaurav Mishra, Genta Indra Winata, Gerard de Melo, Germán Kruszewski, Giambattista Parascandolo, Giorgio Mariani, Gloria Wang. Gonzalo Jaimovicch-L'oper Gregor Betz, Guy Gur-Ari, Hana Galijasevic, Han Sol Kim, Hannah Rashkin, Hanna Hajishirzi, Harsh Mehta, Hayden Bogar, Henry Shevlin, Hinrich Schütze, Hiromu Yakura, Hongming Zhang. Hubert Wong, Ian Aik-Soon Ng. Isaac Noble, Jaap Jumelet, Jack Geissinger, John Kernion, Jacob Hilton, Jaehoon Lee, Jaime Fernnndez Fisac, J. Brooker Simon, James Koppel, James Zheng, James Zou, Jan Koco' n, Jana Thompson, Jared Kaplan, Jarema Radom, Jascha Dickstiic,,,,,, Jason Phang, Jason Wei, Jason Yosinski, Jekaterina Novikova, Jelle Bosscher, Jenni Marsh, Jeremy Kim, Jeroen Taal, Jesse Engel, Jesujoha Oluwadara Alabi, Jiacheng Xu, Jiaming Song, Jillian Tang. Jane W Waweru, John Burden, John Miller, John U. Balis, Jonathan Berant, Jorg Frohberg, Jos Rozen, José Hermández-Orallo Joseph Boudeman, Joseph Jones, Joshua B. Tenenbaum, Joshua S. Rule, Joyce Chua, Kamil Kanclerz, Karen Liveseu, Karl Krauth, Karthik Gopalakrishnan, Ka- terina Ignatyeva, Katja Markert, Kaustubh D. Dhole, Kevin Gimpel, Kevin Ochieng' Omondi, McDonell, Kory Wallace Mathewson, Kristen Chiafullo, Ksenia Shkaruta, Kumar Shridhar, Kyle Kyle Richardson, Laria Reynolds, Leo Gao. Li Zhang, Liam Dugan, Lianhui Qin, Lidia Contreras- Ochando, Louis-Philippe Morency, Luca Moschella, Luca Lam, Lucy Noble, Ludwig Schmidt, Luheng He, Luis Oliveros Col' on, Luke Metz, Lutfi Kerem eSenel, Maarten Bosma, Maarten Sap, Maartje ter Hoeve. Madotto Andrea, Maheen Saleem Farooqi. Manaal Faruqui. Mantas Mazeika, Marco Baturan, Marco Marelli, Marco Maru, M Quintana, Marie Tolkiehn, Mario Giulianelli, Martha Lewis, Martin Potthast, Matthew Leavitt, Matthias Hagen, M'aty' as Schubert, Medina Baitemirova, Melissa Arnaud, Melvin Andrew McElrath, Michael A. Yee, Michael Cohen, Mi Gu, Michael I. Ivanitskiy, Michael Starritt, Michael Strube, Michal Swkedrowski, Michele Bevilacqua, Michihiro Yasunaga Mihir Kale, Mike Cain, Mimee Xu, Mirac Suzgun, Monica Tiwari. Mohit Bansal, Moin Aminnaseri, Mor Geva, Mozhdeh Gheini, T Mukund Nanyun Peng. Nathan Chi, Nayeon Lee, Neta Gur-Ari Krakover, Nicholas Cameron, Nicholas S. Roberts, Nicholas Doiron, Nikita Nangia, Niklas Deckers, Niklas Muennighoff, Nitish Shirish Keskar, Niveditha Iyer, Noah Constant, Noah Fiedel, Nuan Wen, Oliver Zhang, Omar Agha, Omar Elbaghdadi, Omer Levy, Owain Evans, Pablo Casares, Parth Doshi, Pascale Fung, Paul Pu Liang, Paul Vicol, Pegah Alipoormolabashi, Peiyuan Liao, Percy Liang, Peter W. Chang, Peter Eckersley, Phu Mon Htut, Pi-Bei Hwang, P. Milkowski, Piyush S. Patil, Pouya Pezeshkpour, Priti Oli, Qiaozhu Mei, QING LYU, Qinlang Chen, Rabin Banjade, Rachel Etta Rudolph, Raefer Gabriel, Rahel Habacker, Ram' on Risco Delgado, Raphaël Millière, Rhythm Garg, Richard Barnes, Rif A. Saurous, Riku Arakawa, Robbe Raymaekers, Robert Frank, Rohan Sikand, Roman Novak, Roman Sitelew, Ronan Lebras, Rosanne Liu, Rowan Jacobs, Rui Zhang, Ruslan Salakhutdinov, Ryan Chi, Ryan Lee, Ryan Stovall, Ryan Techan, Rylan Yang, Sahib J. Singh, Saif M. Mohammad. Sajant Anand, Sam Dillavou, Sam Shleifer. Sam Wiseman, Samuel Gruetter. Sam Bowman, Samuel S. Schoenholz, Sanghyun Han, Sanjeev K Kwatra, Sarah A. Rous, Sarik Ghazarian, Sayan Ghosh, Sean Casey, Sebastian Bischoff, Sebastian Gehrmann, Sebastian Schuster, Sepideh Sadeghi, Shadi Sameh
14

---

Published as a conference paper at ICLR 2023
Hamdan, Sharon Zhou, Shashank Srivastava, Sherry Shi, Shikhar Singh, Shima Asaadi, Shixi- ang Shane Gu, Shubh Pachchigar, Shubham Toshniwal. Shyam Upadhyay, Shyamolima Debnath, Siamak Shakeri, Simon Thormeyer, Simone Melzi, Siva Reddy, Sneha Priscilla Makini, Soo hwan Lee, Spencer Bradley Torene, Sriharsha Hatwar, Stanislas Dehaene, Stefan Divic, Stefano Ermon, Stella Rose Biderman, Stephanie C. Lin, Stephen Prasad, Steven T. Piantadosi, Stuart M. Shieber, Summer Misherghi, Svetlana Kiritchenko, Swaroop Mishra, Tal Linzen, Tal Schuster, Tao Li, Tao Yu, Tariq A. Ali, Tatsuo Hashimoto, Te-Lin Wu, Theo Desbordes, Theodore Rothschild, Thomas Phan, Tianle Wang, Tiberius Nkinyili,, Timo Schick, T. N. Kornev, Timothy Dolleen-Lawton Tunduny, Tobias Gerstenberg. Trenton Chang, Trishala Neeraj. Tushar Khot, Tyler O. Shultz, Uri Shaham, Vedant Misra, Vera Demberg. Victoria Nyamai, Vikas Raunak, Vinay V. Ramasesh, Vinay Uday Prabhu, Vishakh Padmakumar, Vivek Srikumar, William Fedus, William Saunders, William Zhang, W Vossen, Xiang Ren, Xiaoyu FF Tong, Xinyi Wu, Xudong Shen, Yadollah Yaghoobzadeh, Yair Lakretz, Yang Song, Yasaman Bahri, Ye Ji Choi, Yichi Yang. Yiding Hao, Yifu Chen, Yonatan Belinkov, Yu Hou, Yu Hou, Yushi Bai, Zachary Seid, Zhao Xinran, Zhuoye Zhao, Zi Fu Wang, Zijie J. Wang, Zirui Wang, Ziyi Wu, Sahib Singh, and Uri Shaham. Beyond the imitation game: Quantifying and extrapolating the capabilities of language models. ArXiv, 0bs/2206.04615, 2022.
Titus
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey, ArXiv preprint, abs/2009 06732, 2020. URL https: //ap://aaivvvv org/ abs/////////// 22..2000..00000 06732.
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham. Jinfeng Rao, Liu Yang. Sebastian Ruder, and Donald Metzler. Long range arena : benchmark for efficient transformers. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. /
5998-6008, 2017. URL. https: http://pooocedeeeeeeee neurnnursnnnnnnnne Abstract -aaaasaat....csaaaaaaaaattt..tt. html.
neuript ps
Aaron Voelker, Ivana Kajic, and Chris Eliasmith. Legendre memory units: Continuous-time time representation in recurrent neural networks. In Hanna M. Wallach, Hugo Larochelle. Alina Beygelzimer, Florence d' Buce Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Pro- cessing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 15544-15553, 2019. URL httpt: http:///pooccece...eeecooco htwpeeeeddinnn...uipieeeeee nww.ips...p...e wwterripscccppeeee httpp/ape/2/9//haaaeeee 95285338nee711ee5111883333333132
html.
Apoorv Vyas, Angelos Katharopoulos, and François Fleuret. Fast transformers with clustered attention. In Hugo Larochelle, Marc Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference Neural Information Processing Systems 2020, NeurIPS 2020, December 6- 12, 2020, virtual, 2020. URL https: http:////oocceeeeeeeeeee/ / Abstttectbbbttttteeatt html.
neurips wwtpppppp//2220000 /
Sinong Wang, Belinda Li, Madian Khabsa, Han Fang. and Hao Ma. Linformer: Self-attention with linear complexity. ArXiv preprint, abs/2006.04768, 2020. URL httpt:////a/ //ar:i/a//iivv///a www.wi.wwwwaaa....aa.aoo 2006 04768.
Yuhuai Wu, Markus Norman Rabe, DeLesley Hutchins, and Christian Szegedy. Memorizing transformers. In International Conference Learning Representations, 2022. URL https:
/ //openreview net wtwp:/vwwwww.eeetotmnneeeeee.e.e
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang. Li Yang. and Amr Ahmed. Big bird: Transformers for longer sequences. In Hugo Larochelle, Mare Aurelio Ranzato, Raia Hadsell,
15

---

Published as a conference paper at ICLR 2023
Maria-Florina Balcan, and Hsuan-Tien Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL 2020 ///00//aa////aa c851 2d1 22a22249222255552
.ew.eeiiddinnnnnn........... www.ipee.ic..... wwwpe/ppp///////e////eeeeeeeeeee ract.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li. Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Pumit Singh Koura, Anjali Sridhar, Tianlu Wang. and Luke Zettlemoyer. Opt: Open pre-trained transformer language models. ArXiv preprint, abs/2205 01068, 2022. URL https: //wvivv../vve//////////eaeovvav org/abs 2205. 01 068.
16

---

Published as a conference paper at ICLR 2023
A
SUPPLEMENTAL MATERIAL
A.1 FAST CONVOLUTION VIA FFT
For u, C 1xl. the Circular Convolution Theorem states that,
11-1
(FFFF FFT,()))
1
Ug
circulant
WL-1
ug
where * denotes elementwise multiplication. As FFT, invFFT can be done O(L log L) time this provides fast algorithm for circulant matrix-vector product. In practice, linear systems can often be expressed as a circulant matrix-vector product and is also true in the case of Equation 4 which can be equivalently expressed
[y0 9z-1 J1221. [K] |0 ..... O]ixrr circulant( uz-1 10 ... 2])21x21.
A.2 IMPLEMENTATION OF GSS
def log_ iniiiiiiiiizerr (min, max)
def init (shape)
return random. uni form (shape) (log (max) (min) + log (min)
return init
def oin (H, L, N=512);
# Lambda_re, Lambda im: (N)
# C_ro, C_im: [H N]
Lambda LLmbda re.exp() + 1i*Iambdda im. exp
[N]
с rere 1i im
[8 NJ
s (Lambda * arange (L) view (1, (1) .exp
[N 6
с C (lambda exp() 1) / lambda
[H N]
return einsum hn, nl->hl', C, S) real
[B L
def das (u, H, D):
norm (u)
# compute H state space kernels
simplifiedd_ss__ermee (8, 1}
K_f rfft (K, pad_te-2+L)
rfft (u, pad to=2+t
y irfft (K_ u_t :i]
# param D: []
return y
def gss (8, F-4096, L-4096, E-1024, H=256); shortcut, x x, norm (x)
dense (x, F, activation='gelu))
u dense (x, H, "ativattio=/geell)) y (u, H, 1.) uc dense (Y, F)
о dense (ue V, E) return o shortcut
Figure 2: Pseudocode of GSS ($3.2).
Lambda re and Lambda im are initialized with log_initing je__iitaaalzer set to (min=le-3, max=le0) and (min=le-5, max=le2) respectively. We tuned these parameters on PG-19 on 4k sequence length and used them for all our experiments. We recommend re-tuning these the training sequence length is changed.
17

---

Published as a conference paper at ICLR 2023
Eval Sequence Length
Dataset
Model
Params
TPUv4 hours
Perplexity
512
4k
16k
65k
FLASH (ours)
340M
12.88
11.42
PG-19
GSS-L
352M
0.8k
12.48
11.33
11.16
11.12
GSS-Hybrid-L.
373M
0.8k
11.45
10.52
10.44
10.1
FLASH (ours)
340M
1k
3.22
2.48
Arxiv
GSS-L
352M
0.8k
3.29
2.71
2.72
2.51
GSS- DSS-Hybrid-L
373M
0.8k
2.94
2.51
18.27
125.2
FLASH (ours)
340M
2.5k
2.65
1.90
Github
GSS-L
352M
1.8k
2.31
1.99
2.20
2.28
GSS-Hybrid-L
373M
1.8k
2.34
1.88
1.74
2.09
Table 4: Comparing GSS variants with our re-implementation of the best-performing models from Hua et al. (2022) in both fixed-paran and fixed-compute settings. Param count for all the models include embedding layers as well.
B DATA DETAILS
We do token level modeling on all the datasets and report resulting perplexity numbers on a heldout of examples. Perplexity numbers are obtained using teacher forcing (or parallel mode) where the correct output from the heldout set is used for decoding the next token at each position. For a fair comparison with the baselines we considered, we keep the vocabularies consistent as used by the baselines models. Specifically, we used custom trained 30k sized sentence centencc--piceeeeee vocab for LM1B, T5 vocab with 32k tokens PG19 (Raffel et al., 2020) and Meena vocab with 32k tokens (Adiwardana et al., 2020) for both ArXiv and Github datasets.
C COMPARISON WITH FLASH
In this section we compare GSS with similar sized (both param and compute) FLASH baseline (Hua al., 2022) which also incorporates gated activation. As shown in Table 4, we see trend similar to the comparison with Block Recurrent Transformer where GSS-Hybrid outperforms FLASH baseline on PG-19 and Github but lags behind Arxiv dataset. In addition, we see that GSS-L slighly more compute efficient than FLASH since it achieves better performance of PG-19 and Github while requiring less compute.
FLASH hy typeeaarass. We re-implemented FLASH baseline in our setup while being faithful to the architecture hyperparameters reported in Hua et (2022). In order to fair comparison we made small number of changes, 1) we set the number of layers in FLASH to 32, which matches the GSS baseline and 2) we increase "expansion rate" from to 3 to match the parameter count of GSS. Other than that, repurposed all well-tuned hyperparameters as is from Hua et al. (2022), except for the learning rate which we increased linearly, proportional to the number of tokens per batch. More precisely, Hua et al. (2022) set learning rate to 0.0007 with 28 tokens per batch. Since GSS was trained with 2 tokens per batch, in order to be fair we doubled the learning rate to 0.0014 for the FLASH baseline for PG-19 and Arxiv datasets. Although, due lot of noise the Github dataset, we stick with the lower learning rate of 0.0007 since a high learning rate of 0.0014 led to the model diverging half way through the training run.
D ADDITIONAL ANALYSIS
In this section, we perform two ablations GSS and GSS-Hybrid.
Varying kernel length. In Table 5, we report how the performance of GSS and GSS-Hybrid changes the length of kernel used in simplified DSS is varied. In all of our experiments far,
18

---

Published as a conference paper at ICLR 2023
Eval Sequence Length
Model
Kernel Length
Params
TPUv4 hours
Perplexity
512
16k
65k
64
192M
0.52k
14.81
14.00
14.21
13.56
256
192M
0.49k
14.29
13.36
13.36
12.95
GSS
1024
192M
0,48k
14.22
13.10
12.80
13.45
4096
192M
0.55k
14.01
12.84
12.94
12.47
64
203M
0.55k
12.87
12.90
12.01
12.09
256
203M
0.52k
12.90
12.70
11.88
11.96
GSS-Hybrid
1024
203M
0.55k
12.90
11.86
11.73
11.41
4096
203M
0.56k
12.82
11.78
11.66
11.70
Table 5: Comparing GSS when kernel length is varied on PG-19.
Eval Sequence Length
Model
Attention Length
Params
TPUv4 hours
Perplexity
512
4k
16k
65k
GSS
N/A
192M
0.55k
14.01
12.84
12.94
12.47
64
203M
0.54k
13.66
12.67
12.22
12.19
128
203M
0.51k
13.44
12.35
12.00
12.02
GSS-Hybrid
256
203M
0.54k
13.15
12.09
12.11
11.99
512
203M
0.56k
12.82
11.78
11.66
11.70
Table 6: Comparing GSS when attention chunk length is varied on PG-19.
set the kernel length to the same as input sequence length. This means that the state space layer accumulates information from all of the past tokens. Intuitively, decreasing the kernel length shortens the context window to past tokens but should train faster. As shown in Table we see that both GSS and GSS-Hybrid trains faster we believe suffers from padding issues on TPUs since it may be smaller than what can be used in single instruction). Another thing to note is that, GSS GSS-Hybrid seems degrade much less than vanilla GSS as kernel length is decreased. For instance, as kernel length is decreased from 4k to 1k, GSS-Hybrid degrades by 0.08 points but GSS degrades by 0.26 points. This suggests that, intermediate Transformer layers may be helping with modeling context lengths beyond the chunk length of 512 that limited to. We leave more thorough exploration of this interesting direction future work.
decrease context length (barring kernel length of 64, which
Varying attention chunk length. fully understand the properties of intermediate attention layers in GSS- oSS-Hybrid, we additionally ablate the attention chunk length. As shown in Table 6, as expected, performance significantly degrades
the chunk length is decreased and the values smoothly interpolates between chunk size of 512 and vanilla GSS model. It is interesting see that even a modest chunk size of 64 leads to better performance than vanilla GSS model.
E ADDITIONAL COMPARISON OF DSS vs GSS
In order to do a fairer comparison between DSS and GSS, ran additional study shown in Table 7 where we report numbers for an additional GSS model for which we reduced the number of layers from 16 to 12 and increased the MLP size from 4096 to 6144 in order roughly match the parameter count. We find that GSS-12L is faster yet almost as competitive as GSS (16 layers) on all 3 datasets.
19

---

Published as a conference paper at ICLR 2023
Eval Sequence Length
Dataset
Model
Params
Throughput (steps/sec)
Perplexity
512
4k
16k
65k
DSS
209M
1.8
14.51
13.53
692.1
OOM
PG-19
GSS
192M
5.3
14.01
12.84
12.94
12.47
GSS-12L
209M
5.7
14.07
12.96
12.71
12.36
DSS
209M
1.8
3.65
3.13
284.6
OOM
Arxiv
GSS
192M
5.3
3.57
3.08
3.08
2.75
GSS-12L
209M
5.7
3.57
3.06
3.06
2.8
DSS
209M
1.8
3.65
3.21
242.7
0OM
Github
GSS
192M
5.3
2.68
2.35
2.31
2.12
GSS-12L
209M
5.7
2.67
2.33
2.23
2.33
Table Comparison of DSS, GSS and SSS-12L.
F TRAINING DETAILS
F.1 TRAINING DETAILS FOR RESULTS IN TABLE
All the models in Table were trained with 2 tokens per batch and 125k total training steps. We make sure change the batch size function of the sequence length that number of tokens in the batch remains the same. For example, for LMIB we set batch size to 1024 and sequence length 512 but for of the datasets we use batch size of 128 and sequence length of 4k. For datasets with longer documents, we also considered increasing the sequence length even further. Intuitively, training on longer sequences would help the model learn longer range dependencies better. On the flip side, it makes the optimization a bit more challenging due to large number of correlated tokens per batch and would even likely result in overfitting. Since GSS able to generalize beyond the length was trained on in most cases, we found sequence length of 4k to be reasonable middle ground for this set of experiments.
F.2 TRAINING DETAILS FOR RESULTS IN TABLE
Similar to Section 4.3, unless otherwise mentioned, all the son-baseline models trained using 64 TPUv4 cores for 125k steps. We use batch size of 128 and 4k sequence length at training time, with total of 2 19 tokens per batch. We increase the batch size to 256 for the Github dataset (with token count of 20 per batch) since observed lot of noise in metrics. This is consistent with observations made by Hutchins et al. (2022).
We used Adam optimizer (Kingma & Ba, 2015) and tuned the base learning rate over grid of e [0.0064, 0.0032, 0.0016. 0.0008]. We also employ linear warmup for 1k steps and cosine decay until le-6. We also observed better performance by using higher than typical weight decay rate of 0.1, other than that we did not employ any additional regularization techniques, including dropout. Note that similar (Gu et al., 2022a: Gupta 2022), used constant learning rate of 1e-3 and set weight decay rate to 0.0 for state space parameters part of GSS Layer In addition, we clip the gradient to norm 1.0 before passing the optimizer. Both of these helped with certain instabilities we observed in our preliminary experiments.
F.3 TRAINING DETAILS FOR RESULTS IN TABLE 3
For these experiments, we build on our best performing model from Section 4.2 on PG-19 i.c. GSS- Hybrid-L. In order to scale up GSS, we simply increase the number of layers in the model. We also increased the batch size to 256 (from 128) while using 4k sequence length at training time, with total of 220 tokens per batch, keeping everything else exactly the same.
20


Reviews:
[{"review_id": "q-VmVJKJh5", "reviewer": "Reviewer_LYuN", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "", "summary_of_the_paper": "This paper introduces GSS, a general purpose sequence model which leverages gated units and trains significantly faster as shown on several language modeling benchmarks. This simple-to-implement alternative to S4 and DSS which trains 2-3 times faster, and is competitive with Transformer-based baselines on long-range language modeling benchmarks.", "summary_of_the_review": "The paper is well written and easy following but the experimental part lacks analysis. I think a minor revise is needed.", "strength_and_weaknesses": "Strength:\n1. The motivation of modeling long range dependencies is clear.\n2. The paper is in a good shape and easy following.\n\nWeaknesses:\n1. The experimental part lacks of analysis on the proposed Gated State Spaces to show the insight. This part is not convincing.\n2. Some minor issue in paper. For instance, I found it is hard to find the alignment on figure and code in Figure 1.\n", "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": "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": "WvDjzrs2TV", "reviewer": "Reviewer_Rso2", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "", "summary_of_the_paper": "This paper proposed to leverage the gating structure in gated attention unit (GAU) in state space models, named Gated State Space (GSS). By reducing the dimension of the state space module in GSS which requires FFT to compute the output, GSS is faster than the diagonal state space (DSS) model.\n\nExperiments were conducted on four language modeling benchmarks, and the authors demonstrated the effectiveness and efficiency of GSS comparing with DSS and other Transformer-based baselines.", "summary_of_the_review": "To sum up, I am worried that the paper needs more experiments and analysis before publication.", "strength_and_weaknesses": "Strengths:\nThe design of GSS is well-motivated and the paper writing is clear and easy to follow. The authors also clearly claimed their contributions.\n\nWeaknesses:\nThe main concerns are from experiments. I found that the experimental setup makes the comparisons unfair, which makes it difficult to position GSS among other neural models in language modeling. I elaborated my concerns in the following questions.\n", "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": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "TZpMYXBESc", "reviewer": "Reviewer_Z4KM", "summary": "", "questions": "", "limitations": "", "rating": 6, "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": "", "summary_of_the_paper": "This paper proposes a new layer named Gated State Space (GSS) that combines the gate design in Gated Attention Unit (GAU) with a simplified version of Diagonal State Space (DSS) layer. Authors further propose a hybrid model that combines GSS with sparingly interleaved Transformer blocks. The resulted model trains significantly faster than the original DSS and can extrapolate to longer sequences at test time. The proposed model is evaluated on 4 long range language modeling tasks, comparing with a series of strong baseline models.", "summary_of_the_review": "This paper is an established work that is based on some previous works, such as GAU and DSS, speeding up training compared to DSS. However, the empirical improvements seem not be very strong compared to the baseline methods. But it's a good paper that contributes to the state space models and connects with Transformer layers.", "strength_and_weaknesses": "**Strength:**\n1. The paper is very clearly written and easy to understand.\n2. Authors find that the main bottleneck of DSS comes from the high contextualized input dimension, as the FFT is performed independently on each dimensions of the sequence. With the GAU gating design, GSS enables a 4x reduction in this dimension and increases the throughput when trained on TPUs.\n3. GSS can extrapolate naturally to unseen sequence length at test time due to the recurrent nature of DSS layer, which also allows faster inference speed.\n\n\n**Weaknesses**:\n1. GSS is largely motivated by GAU, but GAU is not included as a baseline method in this paper. The hybrid model with interleaved Transformer layer is closely related to the SRU++ (When Attention Meets Fast Recurrence: Training Language Models with Reduced Compute) paper, but it is not considered as a baseline method either.\n\n**Questions**:\n1. GSS is only evaluated on the long range language modeling tasks, which is fair enough. But have you evaluated GSS on any other tasks or commonly used LM datasets, e.g. Wikitext-103 and LRA?\n2. Have you measured the training speed on GPUs?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "Kn_FIHF7scq", "reviewer": "Reviewer_GxKh", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "", "summary_of_the_paper": "The attention mechanism in transformers scales quadratically with the input size that prohibits its application to long sequences. This work proposes a variant of S4 that can scale to long sequences. The key change is to apply S4 to a space of reduced dimensionality before projecting it back to the original space inspired by a recent work Gated Attention Units. This change reduces the time of fast Fourier transform operations which is a bottleneck on TPUs, resulting in a 3X speedup. In addition, this work also found that random initialization works as well as Hippo initialization.  Experiments on long-form text generation datasets found that while the proposed model slightly underperforms baselines such as block-recurrent transformer at the same parameter count, the proposed model is better at a fixed training cost, especially when interleaved with transformer layers. Besides, the model trained on shorter context window sizes is able to generalize to utilize more context at test time.", "summary_of_the_review": "I think this paper is marginally below the acceptance threshold, due to a few concerns: first, at the same parameter count this work seems to underperform baselines; second, I have some concerns about the argument based on limiting the training cost as mentioned above; third, the final model is a hybrid between S4 and transformer layers, and I'm not sure if it will be actually used due to 1) the hybrid architecture is not as easy to implement and tune as pure S4 layers or pure transformer layers; 2) the speedups are only measured on TPUs which is less widely used than GPUs; 3) the recent lesson we learned in training large LMs that attention cost is negligible as we scale the model size. That being said, this paper has its merits in that it fixes a compute bottleneck and observes good empirical speedups. Therefore, I think this paper is marginally below the acceptance threshold.", "strength_and_weaknesses": "Strengths:\n1. Shows that a simple change that applies S4 to a space of reduced dimensionality improves throughput by a factor of two to three.\n2. It's interesting that a model trained on shorter context sizes can learn to utilize longer contexts at inference time.\n\nWeaknesses:\n1. The main argument is that while the proposed model underperforms the baseline block-recurrent transformer at similar number of parameters, it's better at the same training cost. However, it is not clear to me that the baseline is optimized for limited training cost scenarios: for example, does training the baseline for a shorter period of time significantly degrades its performance in terms of PPL? In this regard, I think showing a curve that plots validation PPL against training time might be more convincing. But even so, the comparison is still not fair if the baseline is not tuned to arrive at the Pareto frontier for the given compute (such as by changing model size, learning rate, and other hyperparameters)\n2. The model architecture seems to be optimized for a particular hardware architecture TPUs (which is not that widely accessible compared to GPUs). Does the increased throughput generalize to different hardware architectures such as GPUs?\n3. Recent works on training large LMs (such as OPT) show that the attention operation only takes negligible time for large models, and the dominating term is the feedforward operations. Is it still meaningful to address the quadratic attention complexity issue as hardware gets more and more powerful and we can run larger and larger models?\n4. Since Gated Attention Units (Hua et al 2022) motivated this work, why don't you compare to it as a baseline? Besides, in Table 1, can you compare to Transformers on other datasets as well (for the 512 setting) to get a sense of how much gain we can get by leveraging more context?\n\nTypos:\n1. Page 6 last paragraph: we present results for on a large range of sequence lengths", "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": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster