Zoology: Measuring and Improving  Recall in Efficient Language Models
Authors: Simran Arora
Sabri Eyuboglu
Aman Timalsina
Isys Johnson
Michael Poli
James Zou
Atri Rudra
Christopher Re
Venue: ICLR 2024 poster (2024)
Keywords: nlp
language models
representation learning
in-context learning
Primary area: representation learning for computer vision, audio, language, and other modalities
TL;DR: We show fundamental differences between attention-based language models and increasingly-popular gated-convolution based ones.

Abstract: Attention-free language models that combine gating and convolutions are growing in popularity due to their efficiency and increasingly competitive performance. To better understand these architectures, we pretrain a suite of 17 attention and gated-convolution language models, finding that SoTA gated-convolution architectures still underperform attention by up to 2.1 perplexity points on the Pile. In fine-grained analysis, we find 82% of the gap is explained by each model's ability to recall information that is previously mentioned in-context, e.g. "Hakuna Matata means no worries Hakuna Matata it means no" -> ??. On this task, termed "associative recall", we find that attention outperforms gated-convolutions by a large margin: a 70M parameter attention model outperforms a 1.4 billion parameter gated-convolution model on associative recall. This is surprising because prior work shows gated convolutions can perfectly solve synthetic tests for AR capability.  To close the gap between synthetics and real language, we develop a new formalization of the task called multi-query associative recall (MQAR) that better reflects actual language. We perform an empirical and theoretical study of MQAR that elucidates differences in the parameter-efficiency of attention and gated-convolution recall. Informed by our analysis, we evaluate simple convolution-attention hybrids and show that hybrids with input-dependent sparse attention patterns can close 97.4% of the gap to attention, while maintaining sub-quadratic scaling. Code is at: https://github.com/HazyResearch/zoology.

Published as a conference paper at ICLR 2024
ZOOLOGY: MEASURING AND IMPROVING RECALL IN EFFICIENT LANGUAGE MODELS
Simran Arora', Sabri Eyuboglu; Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, Christopher Ré
Department of Computer Science, Stanford University Stanford University
(simran, eyuboglu, poli, jamesz, chrismre) oerrianrreccsssstfoorr... fisysjohn, atri]ebuffalo. edu (atimall siimmmlss@purduue edu
ABSTRACT
Attention-free language models that combine gating and convolutions growing in popularity due to their efficiency and increasingly competitive performance. To better understand these architectures, we pretrain a suite of 17 attention and gated- convolution language models. finding that SoTA convolconnolution architectures still underperform attention by up to 2.1 perplexity points on the Pile. In fine- grained analysis, we find 82% of the gap is explained by each model's ability to recall information that previously mentioned in-context, e.g. Hakuna Matata means no worries Hakuna Matata it means - ??. On this task, termed as- sociative recall, we find that attention outperforms gated-convolutions by a large margin: 70M parameter attention model outperforms 1.4 billion parameter cateed-oovvolttion model on associative recall. This is surprising because prior work shows gated convolutions bility. To close the gap between synthetics and real language, we develop a new formalization of the task called multi-query associative recall (MQAR) that better reflects actual language. We perform an empirical and theoretical study of MQAR that elucidates differences in the parameter-efficiincyy of attention and gated- convolution recall. Informed by our analysis, we evaluate simple convolution- attention hybrids and show that hybrids with input-dependent sparse attention patterns can close 97.4% of the gap to attention, while maintaining sub-quadratic scaling. Code at: https: / / gi thub. com/ Ha Hazzzzzy Mazzyesearchh zool ogy.
perfectly solve synthetic tests for AR capa-
INTRODUCTION
Two advances gating and long convolutions have catalyzed a wave of excitement around gated- convolution language models (Fu et al., 2023a; Ma et al., 2022; Wang et al., 2022; Poli et al., 2023a, inter alia.). These architectures combine gating (i.e. element-wise multiplication) with long convolu- tional filters (i.e. the length of the sequence) to enable interactions between distant tokens (Dauphin et al., 2017; Gu et al., 2021). Recent work suggests these models, which exhibit better asymptotic scaling in input sequence length than attention, can match attention in language modeling quality.
We pretrain and evaluate 17 language models across 4 scales (70M 1.4Bn) and 5 architectures on the same data and infrastructure setup. Surprisingly, we find that there still a perplexity gap of up to 2.1 points between state-of-the--at convolution-based architectures and strong Transformer baselines language modeling on the Pile (Table 1). Through fine-grained analysis, we find a single, simple capability is responsible for much of the gap: recalling information seen in-context. Consider the example below, where some tokens can be predicted by recalling an earlier association:
Hakuna Matata! It means no worries for the rest of your days! Hakuna Matata means no
worries
Value
Key-Value
Query
Hit
Query
We find that errors "AR Hits" (e.g. worries above) account for 82% of the perplexity gap to attention on average, despite only representing 6.4% of all tokens in the Pile dataset. 70M param- eter Transformer can predict AR Hits better than 1.4Bn parameter Hyena gated convolution model
Equul contribution, Random ordering by coin

---

Published as a conference paper at ICLR 2024
Assoc. recall hit
Other tokens
Attention
Aggregation is dependent on input
Alsension
Bigram occurs in-context
First occurrence of bigram The contest starts with the qualifictioo phose, which takes oloce oneerdeding (Phee)
Thegoommon (uzzordlys a bird
Projection uw+
Attention sef (QVV
-
qk
with lorge range Though
compoct, common buzzord
Kin Tim Rice in March 2018.. inter March 201 Tem Rice haid..
logiperplexity)
Gated-Conv
Agg Aggregation controlled by
Convolution Karnel
Convolution k U
yy
Projection bW+b
kernel weights
10'
10"
10"
10'
10"
10"
met 19m 2018
Iaterin March Kim ofter Nim had.
Occurences in training data
Occurences in training data RWKV
Attention
Hyena
The associative recall gap. We stratify Pile validation data for models from each architec- ture class by whether or not the predicted token a previously seen bigram in the example context. We plot validation perplexity versus the bigram's 's frequency in the training data. We can clearly see that the gap is localized to examples where bigrams occur, and are seen rarely during in training.
Figure 1:
(20x larger) (Table 1, Table 6). The AR gap persists at the 7Bn parameter scale when comparing RWKV and Llama-2 (Appendix G).'
This task, associative recall (AR), has a long history in machine learning (Graves et al., 2014; Ba et al., 2016, inter alia.) (Appendix 1). Prior work argues that model's ability to perform AR is predictive of in-context learning quality (Elhage et al., 2021; Olsson et al., 2022). As a result, AR has been adopted convolution architectures match attention on synthetic AR tasks used as proxies for real language modeling (Fu et al., 2023a; Poli et al., 2023a; Lutati al., 2023). For this reason, the downstream AR perplexity gaps are surprising.
tool in designing new architectures and prior work has shown that gated
Through measuring recall on real data, we learn the key disparity is that prior synthetic formulations assume there is one query per input, a fixed position in the sequence, where tokens come from a small vocabulary size (e.g. [V| < 50, less than model dimension). Yet, language modeling often requires performing multiple recalls (e.g. for both "Hakuna Matata" and "no worries" above, in a single forward pass), at varying positions, with tokens from a large vocabulary (larger than model dimension). We thus propose the study of multi-querry AR (MQAR). Compared to the prior AR formulations, MQAR better captures the persisting quality gaps on synthetic and real world data (Section 4). However, it is not clear why MQAR elucidates the gap.
We formalize the the MQAR gap across the architecture classes. Gated-connolutions process variable sequences using fixed filters defined by the model weights rather than functions of the input data (See Figure 1). We find it is inefficient for gated-convolutions to perform the variable distance token-to-token interactions (e.g.
distance of 10 tokens for Hakuna Matata and 9 for no worries) in parameter and FLOPs efficient way compared to attention. Attention achieves input-dependence since it computes all loken-to-token interactions when determining how to mix information in the sequence. We formally describe the limitation of gated convolutions via theory and experiments the rest of the paper.
We first introduce a simple operator, BASECONV, that can provably simulate the class of archi- tectures built from gating and convolution primitives. This includes architectures such as H3, Hyena, RWKV, and RetNet. We show with theory and experiments that the model dimension for BASECONv (and thus the aforementioned architectures) to solve MQAR grows with the input quence length (Theorem 4.4) while attention can solve MQAR with model dimension independent of sequence length (Proposition 4.3, Section 2. In practice, gated-convolutions appear to encode approximate solutions to AR that support only a subset of token-interaction distances (which af- fects the ability to identify matching keys given MQAR queries). While theoretically analyzing deep learning architectures is challenging (Hahn, 2020; Merrill et al., 2022; Keles et al., 2023), the fact that gating and convolutions are polynomial operations facilitates our precise analysis.
We measure AR on real data using a simple hcuristic: n-gram tokens that are repeated in context.
Note that the runtime of Attention is quadratic, hut gated convolutions near linear.
2

---

Published as a conference paper at ICLR 2024
Overall
Slices
of gap due
Model
Param (M)
TFLOPs
AR Hits
Other Tokens
AR Hits
Attention
125
2.46
11.01 (2.40)
2.16 (0.77)
12.45 (2.52)
Long Conv
128
1.74
16.98 (2.83)
25.62 (3.24)
16.46 (2.80)
40.1%
H3
168
2.55
12,06 (2,49)
6.75 (1.91)
12.60 (2.53)
88.4%
Hyena
158
2.41
11,60 (2.45)
5.00 (1.61)
12.28 (2.51)
100.0%
RWKV
169
2.08
11.64 (2.45)
5.70 (1.74)
12.29 (2.51)
100.0%
Attention
360
6.23
9,44
1.98 (0.69)
10.62 (2.36)
Long Conv
360
4.08
13.13 (2.57)
13.27 (2.59)
13.12 (2.57)
40.5%
H3
357
4.85
10.38 (2.34)
4.81 (1.57)
11.00 (2.40)
65.8%
Hyena
358
5.03
10.07 (2.31)
3.83 (1.34)
10.75 (2.38)
98.2%
RWKV
351
4.31
9.79 (2.28)
3.82 (1.34)
10.51 (2.35)
100.0%
Table 1: Language modeling validation perplexity on the Pile. After pretraining on 10B tokens of Pile data, we report log perplexity with negative log-likelihood in parentheses. We report overall scores, and for the AR vs. non-AR token slices defined Section 3. FLOPs are computed for inputs of 2048 tokens based on the equations in Appendix C. A table with additional results is in Table 4.
We show that input-dependent sequence mixing is important solve MQAR efficiently. The scaling for gated convolutions with input-independent convolutions is undesirable so we next ask which architectural choices close the gap. We show that data-dependent sequence mixing helps an architecture solve MQAR efficiently (Theorem 4.5). The model needs to adapt the sequence mixing weights based on the token-interaction distances required for each new example.
Several architectural modifications could satisfy the input-dependence property. Based anal- ysis, we evaluate minimal modifications that only add input-dependent operations on exact-match repeated bigram tokens (our heuristic for measuring tokens that require recall). Note that language models often need to perform recall between fuzzier substrings (e.g. synonymous bigrams) or higher-dimensional concepts. However, c.g., convolution filter that shifts the sequence based on the bigram positions or sparse attention placed only on repeated bigram positions
show that simply inserting input-dependent operator
the BASECONV architecture < 10% of layers suffices to outperform the Transformer baseline on Pile language modeling (Section 5) and succeed on MQAR synthetic tasks (Section 4.3). Moreover, this closes 80% of the MQAR perplexity gap on the Pile validation data. Finally, we prototype solutions that learn the positions at which to use input-dependent operations, validating that they also close the gap to attention on the Pile.
In this work, analyze the increasingly- popular convolution-bas based architectures and identify fun- damental limitations. We hope our analysis informs the design of future efficient architectures.
2 BACKGROUND AND PRELIMINARIES
In this section, we describe the setting, introduce notation, and discuss important related work. See Appendix A for broader discussion of related work.
Language modeling. We study auto-regressive language models trained on the task of next token prediction (Gao et al., 2020). Given a sequence N tokens æ = {To (E....... IN-13 drawn from a vocabulary C, the model outputs probability distribution over C for each of s; given the preceding tokens P (o.. 2;-1). The language models in this work share the same high-level architecture. First, each token v; in the input is embedded in d-dimensional space yielding matrix u N Next, u is passed through a stack of L layers, with layer & outputting 4 E N Finally, the embeddings - output by the last layer are mapped back to logits over C with linear projection. Each layer transforms u with a sequence mixer (e.g. attention) followed by a state mixer (e.g. MLP). Unless specified, our models adhere the implementation details of the LLaMA architecture (except the sequence mixer, which we vary throughout) (Touvron et al., 2023).
Sequence mixers. Our work evaluates how the choice of sequence mixer affects the quality and behavior of language models. Most sequence mixers aggregate the token embeddings in sequence via a weighted sum. For example, :] Σj=o w ()(j, where w is function outputting scalar weights. We study the differences between two classes of sequence mixers, discussed next.
3

---

Published as a conference paper at ICLR 2024
Attention. (Vaswani et al., 2017) We review attention, the de facto language model sequence mixer. An attention layer is parameterized by three learnable projection matrices Q Q, V E R xxx xd To compute the output y given imputs u. attention applies the projections to the input: q Qu,
= Ku, v- = Vu. The projected embeddings are aggregated according to: y = softmax gk ) (shown in Fig. 1) in (N)2d) time, which is expensive for long sequences (large N).
Gated-comvolttion. A more efficient alternative to attention is the convolution, which is defindd as y|i Eドز = O uli j.:] where the kernel k E R NNxd is a learnable weight matrix. Convolutions can be computed in time O(Nd log N) using the Fast Fourier Transform (FFT) and the convolution theorem: y = u * k- = FFT (FFT(u) @ FFT(k)) (Cooley and Tukey, 1965). While purely convolutional architectures match or outperform attention in certain domains (e.g. vision (Tay et al., 2022; Gu et al., 2021), audio (Goel et al., 2022), and time-series (Zhang et al,, 2023)), they trail by large margin on language (Fu et al., 2023a).
Recent work has closed much of this gap by combining convolutions with gating (i.e. elementwise multiplication of the imput with a transformed version of itself Dauphin et al. (2017)). Some state- of-the-art art sub-quadratic language models claim attention- quality (Fu et al., 2023a; Wang et al., 2022; Poli et al., 2023a; Peng et al., 2023; Zhai et al., 2021; Fu et al., 2023b, inter alia). Though they may appear different on the surface, these sub-quadratic architectures can all be expressed in terms of convolutions and gating (see Appendix H.2 for detailed descriptions). This work analyzes the differences between attention and the broad class of gated convolution mixers.
IDENTIFYING THE ASSOCIATIVE RECALL PROBLEM
In this section, we measure the perplexity gap between gated convolutions and attention and show that single skill termed associative recall accounts for 82% of the gap on average. This is surprising because prior work shows gated convolutions solve a synthetic version of associative recall perfectly. Informed by our analysis, we define a new synthetic formulation that better reflects real data. This task facilitates our analysis of why the gap occurs (Section 4) and how to fix (Section 5).
3.1 FINE-GRAINED ANALYSIS OF DOWNSTREAM QUALITY
We pretrain a suite of large language models with different sequence mixers across 3 scales (70M-360M) for 10B tokens on the standard Pile language modeling setting us- ing the EleutherAI GPT-NeoX training infrastructure (Gao et al., 2020; Andonian et al., 2023). In the main paper, we compare attention to three sfate-of-the H3, Hyena, and RWKV (Fu et al., 2023a; Poli et al., 2023a; Peng et al., 2023).² We further in- clude pure oongcconvolution model to underscore the importance of gating. In Appendix F, we include results on additional sequence mixers (Hasani et al., 2022; Sun et al., 2023). We use a strong Transformer attention baseline with rotary embeddings and SwiGLU MLPs following the Llama architecture (Touvron et al., 2023), We also use this strong training recipe when training the attention-free sequence mixers. For experimental details and hyperparameters, see Appendix C.
Perplexity Gap
gated convolution sequence mixers:
Across scales, we find that attention outperforms the gated convolutions. The minimum quality gaps are +2.14, +0.59, +0.35 PPI. at 70M, 160M, and 360M parameter scales, respectively. We report overall test perplexity in Table 1. Though these gaps are relatively small on average, models may still perform very differently on different subsets of data (Eyuboglu et al., 2022).
Associative Recall Perplexity To better understand the differences between attention and gated convolutions, we perform a fine-grained analysis of next token predictions and observe that convolution-based models struggle to recall associations previously seen in context. For example, in Fig. 1, the model must recall the association between "Tim" and the last name "Rice". Following a long line of prior work, we call this skill associative recall (AR) (Willshaw et al., 1969; Hop- field, 1982) (see Appendix A for extended discussion of AR's history in machine learning). In Appendix D.1.1, we provide annotated Pile examples to demonstrate the phenomenon qualitatively.
Quantifying AR performance. It is challenging derive a quantitative measure of associative recall performance on the Pile because we don't know which next token predictions in raw text require recall. We propose simple heuristic to identify these tokens, which we refer to as AR Hits, that enables us to scale our analysis to over 10 million tokens of Pile validation data. An AR Hit is
RWKV is commonly referred an RNN. We show can he viewed as convolution (Appendix H.2.2).
4

---

Published as a conference paper at ICLR 2024
the last token of n-gram repeated in context (e.g. the second occurence of "Rice" in Fig. I). However, some common n a-grams (e.g. "of the") could have been memorized during training, so we factor in the frequency with which an n-grams appeared in the training data. We stratify Pile tokens into two slices based on this heuristic and report perplexity on each in Table 1:
1. AR Hits:
(6.4% of tokens) Tokens in the final position of bigram (a pair of consecutive tokens) which previously appeared in context, but - x during training.
2. Other tokens: (93.6% of tokens) Tokens in the final position of bigram which did not previously appear in context or it appeared >> 1, 250 times during training.
In Figure I, we visualize these slices by plotting log- peg-peelyx against the frequency with which bigrams appear during training. Strikingly, the gap between attention and gated convolutions is the largest on AR hits with the fewest occurrences. On the other tokens, there is no gap.
In Table we also compute the percentage of the difference in perplexity between attention and Llog(dxn)-IZA where ø is the perplexity and T is the set of each model that is due to AR tokens: A (7T)
tokens in the test set. This quantity can also be interpreted as the fraction of the overall gap that would close if model matched attention on the AR slice. We find that the AR slice accounts for 82% of the average quality gap between the gated convolutions and attention.
To evaluate the AR capacity of larger models, we train two 1.4 billion parameter attention and Hyena models for 50 billion tokens the Pile and repeat this analysis (see Table 6). Strikingly, a 70 million parameter attention model is a full perplexity point better in the AR slice than this 1.4B Hyena model that is 20x its size (2.41 vs. 3.43 ppl.). In Appendix G, we also evaluate open-source RWKV and attention models trained up to 7 billion parameters. We show that RWKV's performance degrades sharply relative attention as we increase the number of recall-queries in example Appendix G.
3.2 FORMALIZING THE PROBLEM: MULTI- MULTI-QUERY ASSOCIATIVE RECALL
This gap in associative recall perplexity is very surprising because prior work shows gated- convolutions can perfectly solve formalized version of the task (Fu et al., 2023a; Poli et al., 2023a; Olsson al., 2022). In this synthetic task, the input æ contains sequence of bigrams representing key-value pairs from a random dictionary followed by single query token. For example, the correct output for the input below would be 3:
A 4 B C6E2FIC6GS
B
Key-Value
Query
Gated convolutions (e.g. H3, Hyena, RWKV) can solve this task perfectly for most sequence lengths.
These conclusions are inconsistent with our findings on the Pile, as described above, so we ask how this formulation of AR differs from the way AR manifests in real language. We identify a major difference. In real world inputs, the language model often needs to perform multiple associative recalls in a single forward pass, at varying positions in the sequence (e.g. "Tim Rice" " and "March 2018" in Fig. 1. We refer to this as Multi Multi-Query AR (MQAR), and the formal definition follows:
Definition 3.1 (Multi-Query-AR (MQAR)). We are given an input sequence x = (ro...., IN- where each 34 C' is a token drawn from vocabulary of size
|C]. The task is to check, for every query 1 < i < N, whether there exists a 0 < j < i such that u, Uj. If so, output uz+1-
For example, the correct output for input below would be 4, 6, 1, 2, 3:
B 3 F1E2 E2 + ? F? E E
Key-Value
Query
In Section 4, use MQAR to explain the quality gap between gated convolutions and attention.
4 EXPLAINING THE ASSOCIATIVE RECALL PROBLEM
In this section, we provide an explanation for the gap in associative recall performance by ana- lyzing the formal MQAR task theoretically and empirically. In Section 4.1, we define a simple gated-convolution architecture, called BASECONV, which we show can simulate a broad class of
"In Appendix H.7.1 we define formal, general form of this definition, used in the theoretical analysis.
5

---

Published as a conference paper at ICLR 2024
architectures built from gating and convolutions. This allows us to make general statements that ap- ply to popular gated- coneolconvoottion architectures like Hyena, RWKV, or H3. In Section 4.2, we show that there exist theoretical solutions to MQAR that could in principle be learned by BASECONV, and we analyze their complexity in terms of model width and depth. In Section 4.3, we use experiments on synthetic data show that solving MQAR with BASECONV (and other gated-convolution archi- tectures) requires model dimension to scale linearly with the sequence length. In contrast, attention solves MQAR consistently in our experiments with model dimension scaling independently of quence length. These empirical scaling laws provide a potential explanation for the AR gap and, alongside our theoretical analysis, point to the potential solutions discussed in Section 5.
4.1 BASECONV: A MINIMAL GATED CONVOLUTION OPERATOR
In this section, define our minimal gated-convolution architecture, called BASECONV. Given function, we would like to know the most efficient model (e.g. parameters, FLOPs) that can repre- sent the solution. In this work, we show we can precisely reason about this question for representing polynomial functions with gated convolutions as gating and convolutions are both polynomial op- erations. The standard model defining computational complexity for polynomials is by the size of the smallest arithmetic circuit that can compute the polynomial. We define the BASECONV gated convolution operator which is exciting because (1) is universal in that that can simulate any arith- metic circuit C (with only a poly-log blowup in the corresponding parameters) and (2) it is simple to implement efficiently (19 lines of pure PyTorch including imports, see Appendix B).
Definition 4.1 (BASECONV Operator). Given input Rxd, the BASECONV operator for layer is defined as:
y == uu W + b() (h' u bb)
(I)
Linear Projection
Convolution
where the layer is parameterized by learnable filters h E R Nxd, linear projection W E R
and 'bias' matrices b1, b2 E R The is component-win product and convolution of two matrices is computed as convolution of the corresponding columns.
In our experiments, each BASECONV layer uses ( + d²) parameters and can be computed in (NNe) operations. For our theoretical results, we can assume the weight matrix W' is restricted to class of matrices that support near-linear time matrix multiplication (e.g. Kaleidoscope matri- ces, see Definition H.3). Under this assumption, BASECONv uses Ö(Nd) parameters and ((Nd) FLOPs (Proposition H.6). We now state the equivalency result between arithmetic circuits and BASECONV "canonical" representation of arithmetic circuits (Theorem H.21 Appendix H.5):
Theorem 4.2 (Equivalency to Arithmetic Circuits). For arithmetic circuit C of size and depth A that takes IR. Nxd as input, there exists equivalent BASECONv operator that uses (ss)) parameters and Ó(A) layers. 6
In other words, any gated convolution model with small number of layers can be simulated by BASECONV with only a logarttogaii blowup in parameters and layers. We note that arithmetic circuits are a very well studied computation model in computational complexity Bürgisser et al. (1996). Many well-known efficient algorithms on matrices (c.g. the FFT or the current best known matrix-matrix multiplication algorithm) in fact give small arithmetic circuits. However, arithmetic circuits are inherently discrete objects - we cannot learn them via gradient descent. Theorem 4.2 shows that (up to poly-log loss in parameters), we can instead learn over BASECONV models. This result generalizes a similar result from Dao et al. (2020) for the special class of linear functions: we generalize the earlier result the class of all polynomials.
For specific gated convolution layers, we can get rid of the poly-logarithmic factor blowup-we observe in the appendix that BASECONY and Hyena models can simulate each other with only a small constant blowup in parameters (Proposition H.12 in Appendix H.5).
'We use Õ-) to hide poly-log factors.
'The formal statement in the Appendix has sharper version of this result in terms of the circuit 'width".

---

Published as a conference paper at ICLR 2024
Sequence Length:
Sequence Length:
Sequence Length: 258
Sequence Length:
1.00
Attention
0.75
Hyena
0.50
RWKV
0.25
H3
BaseConv
0.00
Atlention hybrid
Accuracy
0.75
Autocorrelation
Fillers
025
Programmatic
Filters
64
258 Model dimension
128
512
258 Model dimension
512
258 Model dimension
128
512
64
128
512
Model dimension
Figure 2: The x-axis is the model dimension and the y-axis is accuracy on MQAR. (Top, Claim 1) Gated convolutions and attention. We observe the gated convolutions require larger dimension- ality than attention solve the task. (Bottom, Claim 2) Input-dependent filters. Models with input-dependent aggregation achieve improved scaling over the galed-convolutions.
4.2 THEORETICAL ANALYSIS OF GATED CONVOLUTION CAPACITY AND RECALL
In this section, we provide theoretical MQAR solutions that could in principle be learned by each architecture and analyze their complexity in terms of model width and depth. First, we note that attention solves MQAR with parameters independent of sequence length (Proposition H.27).
Proposition 4.3 (Attention). Given an input u E {0, 1] x3x Attention (even without using soft- max) solves MQAR for using (ce²) parameters, O(Ne² ( +N²c) time complexity and ((1) layers.
It is natural to wonder then if all pairwise comparisons among tokens are necessary to solve MQAR. Indeed, in the RAM setting, sequential algorithm can simply utilize logarithmic insertion and membership queries to solve MQAR in subquadratic time. Unfortunately, any model attempting to emulate this would require (2N layers. Instead, we observe that we can parallelize this algorithm using dyadic intervals and achieve a depth of Õ(1) (Proposition H.30). We then convert this algo- rithm into arithmetic circuit and apply Theorem 4.2 derive an equivalent BASECONY model. This allows us to prove new upper bounds for BASECONY models applied to MQAR, which im- proves upon the quadratic time complexity of attention to near-linear runtime the cost of using poly-log layers (Theorem H.37 in Appendix H.7).
Theorem 4.4 (Input- Independdd Filters?), Given an input E {0, 1 x to MQAR (where we assume that distinct tokens are embedded into distinct vectors in {0, (1logg))) there exists BASECONV operator that solves MQAR for u using Ô log e) parameters as well as time com- plexity and Õ(1) layers.
Nevertheless, the holy-logarithmic number of layers in the above result is undesirable in practice. But, show that using input-depender convolution filters, one get constant many layers (for a sub-class of inputs). Towards that end, we define the interaction distance between a query and the matching key k, as i j. This then allows us to present the corresponding upper bound for detandependdnt mixing (Theorem H.38 in Appendix H.8).
Theorem 4.5 (Input- Filters). Given an input E {0, 1] to MQAR (where we assume that the tokens embedded one-hot encoding in 13 and there exists most t distinct interaction distances), 8 there exists a BASECONV operator that uses input-dependent kernels to solve the above case of MQAR using Ott Ne) parameters and O(1) layers.
4.3 EMPIRICAL ANALYSIS OF GATED CONVOLUTION CAPACITY AND ASSOCIATIVE RECALL
In this section, measure empirically how model dimension must scale in order for different sequence mixers to solve MQAR.
Setup We train and evaluate models on a synthetic MQAR with vocabulary size 8, 192, across model dimension and sequence length (64 512). Appendix E provides further details the
Mcaning the filter is defined as a function of the model parameters, independent of the input.
Note that the interaction distances he arbitrary: there is just bounded number of distinct distances.

---

Published as a conference paper at ICLR 2024
struction of this synthetic task. Following Olsson et al. (2022), train two layer models with Transformer backbone that interleaves sequence mixing and state mixing (MLPs). For each archi- tecture, we sweep four learning rates from log( (4) to log(-2) and report maximum test accuracy.
Our results, which are summarized in Figure 2, support two main claims:
Claim (Gnted-convolutions and attention). Gated-convolution models with two layers require model dimension to scale at least linearly in sequence length in order to solve associative recall, while attention models can solve with near-constant dimensionalityy We compare attention and BASECONV as well as three popular instantiations of gated-convolution architectures: RWKV, H3, and Hyena (Peng et al., 2023; Fu et al., 2023c; Poli et al., 2023a). In the top row of Fig. 2, attention solves MQAR perfectly at all sequence lengths using constant model dimension of 64. In contrast, MQAR does not achieve accuracy > 0.9 unless 2 N.
Claim (Input-dependent filters). Using input-dependent filters in galed-comvolation models can close some of the gap to attention. In Theorem 4.5, we show that BASECONV with input-dependent filters could solve MQAR with improved scaling. In this solution, we construct a filter that spikes at position j if matching keys are separated by j tokens. We evaluate two approaches for con- structing this filter: (1) programatically (i.e. hard-coded comparisons between token ids) or (2) with autocorrelation, which could learn to perform fuzzy-matches (see Appendix H.8). In the bottom row of Fig. 2, we see that BASECONV with programmatic input-dependent filters achieves near- constant scaling in model dimension and that BASECONV with autocorrelation input-dependent filters achieves improved scaling over BASECONV with input-independent filters.
These input-dependent filters cannot easily be made satisfy causality and using O(N log N) filter per gap could be expensive each gap applies only to a small number of bigrams. A simpler and perhaps more efficient way solve the problem would be to introduce a small amount of atten- tion to an otherwise BASECONV model (Fu et al., 2023a). As the first natural baseline, we evaluate an attention hybrid on synthetic MQAR and show that it achieves improved scaling in Figure 2. Next, we put these empirical and theoretical insights into practice on Pile language modeling.
CLOSING THE ASSOCIATIVE RECALL GAP
In this section, we evaluate hybrid BASECONV-Attentee models that leverage different sparsity patterns. We show that hybrids with input-dependent sparsity patterns can close most of the gap to attention, while maintaining -ubdquadratic scaling. We also show that BASECONY hybrids can out- perform iitentiion-nll models by up to a full perplexity point, all while being dramatically simpler to implement and analyze than prior hybrids (Fu et al., 2023a) (see implementation in Appendix B. We describe the architectures in Section 5 and results on Pile language modeling in Section 5.
Sparse BASECoNV-Attention Hybrids We evaluate hybrids composed primarily of BASECONV layers and three attention layers (6.3% of layers at 354M parameters and 10% at 168M parameters). We augment the attention layers with operators that selectively apply attention to some tokens based on a selection function f IR Xxd + 10,1) They take as input u E Nxd and output y E R.Vxd
y i softmax V gli all. (kTT) (o))))()(ii
where q, k, v query, key, and value projections, in attention. We evaluate four choices for f:
(1) Full attention. First, we evaluate the performance of full attention hybrids. These hybrids correspond to fixing f fu)][i] [i] 1 for all in Eq. (2). Prior work has shown full attention hybrids to be effective, but they do not explain why supplementing gated convolutions with attention is necessary (Fu et al., 2023a). Based on our findings in Section 3, we hypothesize that sparse attention applied only to associative recall hits may suffice.
(2) Random selection. As control, we evaluate sparse attention randomly applied to tokens. This corresponds to a stochastic selection function where [u)[[i] is drawn from Bernoulli. Many prior works have employed sparse attention patterns like this one, which are independent of the input (Za- heer et al., 2020; Child et al., 2019a; Beltagy et al., 2020, inter alia.).
(3) Programmatic selection. Next, to evaluate our hypothesis that attention is needed for associative recall, we prototype a programmatic selection function that selects only those tokens that might be
8

---

Published as a conference paper at ICLR 2024
Overall
Slices
Model
Param (M)
TFLOPs
AR Hits
Other Tokens
Attention
125
2.46
11.01 (2.40)
2.16 (0.77)
12.45 (2.52)
BASECONV
168
2.46
12.90 (2.56)
8.68 (2.16)
13.29 (2.59)
+ Random selection
162
2.44
12.13 (2.50)
5.33 (1.67)
12.83 (2.55)
+ Programmatic selection
2.44
11.06 (2.40)
2.70 (0.99)
12.19 (2.50)
+ Learndd selection
166
2.44
11.06 (2.40)
3.20 (1.16)
12.14 (2.50)
+ Full attention
166
2.58
9.57 (2.26)
2.01 (0.70)
10.76 (2.38)
Attention
360
6.23
9.44 (2.25)
1.98 (0.69)
10.62 (2.36)
BASECONV
354
4.81
11.01 (2.40)
5.98 (1.79)
11.52 (2.44)
+ Random selection
365
5.06
12.94 (2.56)
6.17 (1.82)
13.62 (2.61)
+ Programmatic selection
365
5,06
9.54 (2.26)
2.35 (0.86)
10,50 (2.35)
+ Learned selection
351
5,06
9.59 (2.26)
2.61 (0.96)
10.58 (2.36)
+ Full attention
351
5.10
8.59 (2.15)
1.95 (0.67)
8,91 (2.19)
Table 2: Language model perplexity on slices of the PILE. We evaluate Hyena and BASECONY with Hybridization and Selective look-up 160 and 355M parameters.
associative recall hits. Specifically, f( :I) is 1 the token z, previously occurred in the sequence. In practice, we compare raw token ids, not token embeddings.
[x) {1
if there exists j << such that
jjj
(3)
otherwise
(4) Learned selection. Finally, we prototype learned selection function f(u)]i]] :]] W) parameterized as simple linear layer with sigmoid activation. We fix hyperparameter & and select the top-k tokens with the highest score in each batch. This allows us to compute at- tention in (nddk) time. During training, we add small amount of Gaussian noise to the top- calculation to encourage exploration and use an auxiliary loss to encourage sparse selection: (y( = N max(0, Σ [(u)](]] k). This approach most similar to the recently proposed SeqBoat architecture (Ren al., 2023). We discuss the differences Appendix A.
1
Downstream Evaluations We evaluate the prototypes on Pile language modeling, We take BASECONY architectures at the 150M and 360M parameter scales and add input-dependent selec- tion to three layers (details in Appendix C.2). Results are in Table 2. We validate that the prototypes close the overall and AR quality gaps attention when added the BASECONV backbone.
At 360M parameters, BASECONY with just 3 attention layers can outperform the Transformer, while requiring fewer FLOPs. Our hybrid attention-Basse BASECONV models outperform attention only models by 0.85 perplexity points while enabling an 18% reduction in total FLOPs vs attention. However, this uses full quadratic attention. We next show that sparse attention localized to potential AR tokens, is also sufficient to close the gap, validating our insights on the role of input-dependence for MQAR. At 360M parameters, programmatic selection closes 85% the gap between pure BASECONY and attention the AR slice. in contrast the random selection control. Learned selection closes 72% a of the gap using just 256 (sub-quadratic) attention positions per example.
DISCUSSION AND CONCLUSION
We present an extensive analysis of gated convolution architectures in light of their recent popu- larity. We identify persisting quality gap between efficient convolution and inefficient attention based architectures, largely due single failure mode associative recall. We design a new multi- query associative recall (MQAR) analysis tool, which correlates with downstream AR quality. We theoretically and empirically explain the gap is due to insufficient data-dependent mixing in gated convolutions and we show minimal architectures that close the gap the Pile.
Our results in analyzing gated convolutions go beyond the conventional wisdom that attention the "right" model. A significant amount of work focuses improving the efficiency of attention (Dao et al., 2022; Dao, 2023; Katharopoulos et al., 2020a) and theoretically studying the exact power of attention (Hahn, 2020; Merrill et al., 2022; Keles al., 2023). Attention is often used as the goalpost for what is needed downstream. We hope our contributions highlight the value of MQAR, and more broadly tasks tied to real language modeling, proxy study.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
We thank Tri Dao, Daniel Fu, Neel Guha, Stefano Massaroli, Eric Nguyen, and Michael Zhang for helpful feedback and discussion during this work. We are grateful Together Computer for making this work possible. We gratefully acknowledge the support of DARPA under Nos. FA86501 827865 (SDH) and FA86501827882 (ASED); NIH under No. U54EB020405 (Mobilize). NSF under Nos. CCF1763315 (Beyond Sparsity), CCF1563078 (Volume to Velocity), and 1937301 (RTML); ONR under No. N000141712266 (Unifying Weak Supervision); the Moore Foundation, NXP, Xilinx, LETI-CEA, Intel, IBM, Microsoft, NEC, Toshiba, TSMC, ARM, Hitachi, BASE, Accenture, Er- iesson, Qualcomm, Analog Devices, the Okawa Foundation, American Family Insurance, Google Cloud, Microsoft Azure, Swiss Re, Brown Institute for Media Innovation, Department of Defense (DoD) through the National Defense Science and Engineering Graduate Fellowship (NDSEG) Pro- gram, Fannie and John Hertz Foundation, National Science Foundation Graduate Research Fel- lowship Program, Texas Instruments Stanford Graduate Fellowship in Science and Engineering, and members of the Stanford DAWN project: Teradata, Facebook, Google, Ant Financial, NEC, VMWare, and Infosys. The U.S. Government is authorized to reproduce and distribute reprints for Governmental purposes notwithstanding any copyright notation thereon. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views, policies, or endorsements, either expressed or implied, of DARPA, NIH, ONR, or the U.S. Government. AR's work is supported by NSF grant# CCF-2247014. IJ's work is supported by an NSF Graduate Fellowship.
REFERENCES
Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. Hungry Hungry Hippos: Towards language modeling with state space models. In International Confer- ence on Learning Representations, 2023a.
Xuezhe Ma. Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Zettlemoyer Luke. Mega: Moving average equipped gated attention. arXiv preprint arXiv:2209. 10655, 2022.
Junxiong Wang, Jing Nathan Yan, Albert Gu, and Alexander M Rush. Pretraining without attention. arXiv preprint 12 2022.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. arXiv preprint 2rXiv:2302.100S66 2023a.
Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In International conference on machine learning, pages 933-941. PMLR, 2017.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.0033 2021.
Alex Graves, Greg Wayne, and Ivo Danihelka. 17X.v:1410.5401, 2014.
Neural turing machines. arXiv preprint
Jimmy Ba, Geoffrey Hinton, Volodymyr Mnih, Joel Z. Leibo, and Catalin Ioneseu. Using fast weights to attend to the recent past. Advances in neural information processing systems, 29, 2016.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 1, 2021.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209.1 1895, 2022.
Shahar Lutati, Itamar Zimerman, and Lior Wolf. Focus your attention (with adaptive iir filters), 2023.
10

---

Published as a conference paper at ICLR 2024
Michael Hahn. Theoretical limitations of self-attention in neural sequence models. In Transactions of the Association for Computational Linguistics, volume 8, 2020.
William Merrill, Ashish Sabharwal, and Noah A. Smith. Saturated transformers are constant-depth threshold circuits. Transactions of the Association for Computational Linguistics, volume 10, 2022.
Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, and Chinmay Hegde. On the computational complexity of self-attention. In 34th International Conference on Algorithmic Learning Theory, volume 201, page 1-23, 2023.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy, The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00077, 2020.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, and Shruti Bhosale. Llama Open foundation and fine-tuned chat models. arXiv:2307.092888 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. volume 30, 2017.
James W Cooley and John W Tukey. An algorithm for the machine calculation of complex fourier series. Mathematcc of computation, 19(90))297-301 1965.
Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey. ACM Computing Surveys, 55(6):1-28, 2022.
Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré. It's raw! audio generation with state- space models. Proceedings of the 39 th International Conference on Machine Learning, 2022.
Michael Zhang, Khaled Saab, Michael Poli, Tri Dao, Karan Goel, and Christopher Ré. Effectively modeling time series with simple discrete state spaces. International Conference on Learning Representations, 2023.
Bo Peng. Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Przemys- law Kazienko, Jan Kocon, and Jiaming et al. Kong. Rwkv: Reinventing mns for the transformer era. arXiv:2305. /3048, 2023.
Shuangfei Zhai, Walter Talbott, Nitish Srivastava, Chen Huang, Hanlin Goh, Ruixiang Zhang, and Josh Susskind. An attention free transformer. arXiv preprint arXiv:2105.14103, 2021.
Daniel Y. Fu, Simran Arora, Jessica Grogan, Isys Johnson, Sabri Eyuboglu, Armin W. Thomas, Benjamin Spector, Michael Poli, Atri Rudra, and Christopher Ré. Monarch mixer: A simple sub-quadratic gemm-based architecture, 2023b.
Alex Andonian, Quentin Anthony, Stella Biderman, Sid Black, Preetham Gali, Leo Gao, Eric Hallahan, Josh Levy-Kramer, Connor Leahy, Lucas Nestler, Kip Parker, Michael Pieler, Jason Phang, Shivanshu Purohit, Hailey Schoelkopf, Dashiell Stander, Tri Songz, Curt Tigges, Ben- jamin Thérien, Phil Wang, and Samuel Weinbach. GPT-NeoX: Large Scale Autoregressive Lan- guage Modeling in PyTorch, 9 2023. URL https: /www github. com/eleutherail gpt-neox.
Ramin Hasani, Mathias Lechner, Tsun-Huang Wang, Makram Chahine, Alexander Amini, and Daniela Rus. Liquid structural state-space models. arXiv preprint arX1v:2209.1295/ 2022.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue. Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models, 2023.
Sabri Eyuboglu. Maya Varma, Khaled Saab, Jean- Ben-t Delbrouck, Christopher Lee-Messer, Jared Dunnmon, James Zou, and Christopher Ré. Domino: Discovering systematic errors with cross- modal embeddings. In International Conference Learning Representations, March 2022.

---

Published as a conference paper at ICLR 2024
David J Willssaw, O Peter Buneman, and Hugh Christopher Longuet-Higgins. Non-holographic associative memory. Nature, 222(5197):960 962, 1969.
John J Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 79(8):255-2258, 1982.
Bürgisser, T. Lickteig, M. Clausen, and A. Shokrollahi. Algebraic Complexity Theory. Grundlehren der mathematischen Wissenschaften. Springer Berlin Heidelberg, 1996. ISBN 9783540605829. URL https: /books....//// google www/books//...ccccoooooccco...oo foop
Tri Dao, Nimit S Sohoni, Albert Gu, Matthew Eichhorn, Amit Blonder, Megan Leszezynski, Atri Rudra, and Christopher Ré. Kaleidoscope: An efficient, learnable representation for all structured linear maps. arXiv preprint (rXiv:20122119660 2020.
Daniel Y. Fu, Elliot L. Epstein, Eric Nguyen, Armin W. Thomas, Michael Zhang, Tri Dao, Atri Rudra, and Christopher Ré. Simple harcware-efficiet long convolutions for sequence modeling. arXiv preprint arXiv:2302.06644 2023c.
Manzil Zaheer, Guru Guruganesh, Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and et al. Big bird: Transformers for longer sequences. Proceedings of NeurlPS, 2020.
Rewon Child, Scott Gray, Alee Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint UrXiv:194.110509 2019a.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint urXiv:2004.0511 (550000.5555,, 2020.
Liliang Ren, Yang Liu, Shuohang Wang, Yichong Xu, Chenguang Zhu, and ChengXiang Zhai. Sparse modular activation for efficient sequence modeling. Sparse Modular Activation for Dffi- cient Sequence Modeling, 2023.
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. FlashAttention: Fast and memory-effiiient exact attention with IO-awareness. In Advances Neural Information Process- ing Systems, 2022.
Tri Dao. FlashAttention: Faster attention with better parallelism and work partitioning. 2023.
A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers rans: Fast autoregressive transformers with linear attention. In Proceedings of the International Conference on Machine Learning (ICML), 2020a. URL
org/ abss2006 16236.
Chris Olah. Mechanistic interpretability, variables, the importance of interpretable bases: informal note some intuitions related to mechanistic interpretability, 2022. URL https: / //transformer- www.ciistt....oe .pp.s.ppp........... pubs mubb2002 ht////echhhneeeeeee intth-in-------- essay index.niix. html.
Alethea Power, Yuri Burda, Harri Edwards, Igor Babuschkin, and Vedant Misra. Grokking: Gener- alization beyond overfitting small algorithmic datasets. arXiv:2201.02177 2022.
Nick Cammarata, Shan Carter, Gabriel Goh, Chris Olah, Michael Petrov, Ludwig Schubert, Chelsea Voss, Ben Egan, and Swee Kiat Lim. Thread: circuits. Distill, 5(3):e24, 2020.
Dingquan Wang and Jason Eisner. The galactic dependencies treebanks: Getting more data by synthesizing new languages. Transactions of the Association for Computational Linguistics, 4: 491-505, 2016.
Jennifer C White and Ryan Cotterell. Examining the inductive bias of neural language models with artificial languages. arXiv preprint arXiv:2106.0100 2021.
Zeyuan Allen-Zhu and Yuanzhi Li. Physics of language models: Part 1, context-free grammar. arXiv preprint arXiv:2305.13673 2023.
Shauli Ruvfogel, Yoav Goldberg, and Tal Linzen. Studying the inductive biases of mns with syn- thetic variations of natural languages. arXiv preprint 1003v:9993.06444 2019.
12

---

Published as a conference paper at ICLR 2024
Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma. An explanation of in-context learning as implicit bayesian inference. arXiv preprint arXiv:2111.02080, 2021.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint (.Xi::2001.0445//1 2020.
JA Feldman, GE Hinton, and JA Anderson. Parallel models of associative memory, 1981.
Wei Zhang and Bowen Zhou. Learning to update auto-associative memory in recurrent neural net- works for improving sequence memorization. arXiv preprint arXiv: 179.v6770, 2017.
Stefano Massaroli, Michael Poli, Daniel Y Fu, Hermann Kumbong, David Romero, Rom Par- nichukun, Aman Timalsina, Quinn McIntyre, Beidi Chen, Atri Rudra, Ce Zhang, Christopher Ré, Stefano Ermon, and Yoshua Bengio. Laughing hyena distillery: Extracting compact recurrences from convolutions. 2023.
David W. Romero, Anna Kuzina, Erik J. Bekkers, Jakub M. Tomezak, and Mark Hoogendoorn. Ckconv: Continuous kernel convolution for sequential data. 2022.
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are effective as structured state spaces, 2022.
Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models, 2022.
Harsh Mchta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language model- ing via gated state spaces, 2022.
Jimmy T. H. Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling, 2023.
Eric Nguyen, Michael Poli. Marjan Faizi, Armin Thomas, Callum Birch-Sykes Michael Wornow, Aman Patel, Clayton Rabideau, Stefano Massaroli, Yoshua Bengio, Stefano Ermon, Stephen A. Baccus, and Chris Ré. Hyenadna: Long-range genomic sequence modeling at single nucleotide resolution, 2023.
Transformers revolutionized ai. what will replace them?, https: /wwww//ww/////////////// forbes cww/sites/ wwwm:/ttee//rrbbee//////////////
2023.
URL
6ed698269c1
Jue Wang, Wentao Zhu, Pichao Wang, Xiang Yu, Linda Liu, Mohamed Omar, and Raffay Hamid. Selective structured state-spaces for long-form video understanding. In CVPR, 2023.
Brandon Yang, Gabriel Bender, Quoe V. Le, and Jiquan Ngiam. Condconv: Conditionally
parametrized convolutions for efficient inference. In 33rd Conference on Neural Information Processing Systems (NeurIPS 2019), 2020.
Chrysoula Kosma, Giannis Nikolentzos, and Michalis Vazirgiannis. Time lime-parameterized convolu- tional neural networks for irregularly sampled time series. 2023.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mnss: Fast autoregressive transformers with linear attention, 2020b.
Together. Redpajama: An open source recipe to reproduce Ilama training dataset, 2023. URL https:/ / thu.. com/ www.ttthecoooteeeeecoot RReeerr/eePPPaae jama-Data
Sainbayar Sukhbaatar, Edouard Grave, Piotr Bojanowski, and Armand Joulin. Adaptive attention span in transformers, Association of Computational Linguistics, 2019.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
13

---

Published as a conference paper at ICLR 2024
Anonymous, Mamba: Linear-time sequence modeling with selective state spaces. 2023a. URL
httpp::/coenneevieww .nt.
Anonymous, Gateloop: Fully data-controlled linear recurrence for sequence modeling. 2023b. URL
httts:://gpereeviee wwt. http://eenneeeeeeeeeeeeedddeeeee http://pdffd==000gggee
Rewon Child, Scott Gray, Alec Radford, and llya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv: (0te::994..00550000 10509, 2019b.
Jiezhong Qiu, Hao Ma, Omer Levy, Scott Wen tau Yih, Sinong Wang, and Jie Tang. Blockwise self-attention for long document understanding, arXiv preprint arXiv:1911.02972. 2019.
Michael T Heideman and C Sidney Burrus. Multiplicatiive complexity, convolution, and the DFT. Springer, 1988.
Ilya Volkovich. A guide learning arithmetic circuits. In Conference on Learning Theory, pages 1540-1561. PMLR, 2016.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. Proceedings of the 40th International Conference on Machine Learning, 2023b.
Thathachar S Jayram, Ravi Kumar, and Dandapani Sivakumar. The one-way communication com-
plexity of hamming distance. Theory of Computing, 4(1):129-135, 2008.
Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108. 12409, 2021.
Selim G Akl and Henk Meijer. Parallel binary search. IEEE Transactions on Parallel & Distributed Systems, 1(02):247-250) 1990.
Thomas H Cormen, Charles E Leiserson, Ronald L Rivest, and Clifford Stein. Introduction to algorithms. MIT press, 2022.
Miklós Ajtai, János Komlós, and Endre Szemerédi. An 0 (n log n) sorting network. In Proceedings of the fifteenth annual ACM symposium on Theory of computing, pages 1-9, 1983.
Chris Chatfield The analysis of time series: An introduction, fifth edition. 1995.
14

---

Published as a conference paper at ICLR 2024
APPENDIX
The appendix includes the following content:
Appendix provides an extended discussion of related work and concepts.
2. Appendix provides a code implementation of the BASECONV architecture.
3. Appendix C gives details for the experiments, including model architectures and hyperpa- rameters.
4. Appendix provides additional analysis of how MQAR appears in real data across a variety of language distributions.
5. Appendix E provides formal definition of the MQAR problem, synthetic construction procedure, and experimenta details.
6. Appendix F provides additional synthetic experiments and analysis.
7. Appendix G provides experiments and analysis for how the MQAR gap changes scale the gated convolution and attention architectures.s
8. Appendix H gives proof's and additional discussion for the theoretical analysis in our work.
A EXTENDED RELATED WORK
We are inspired by and build prior work in mechanistic interpretability (Appendix A.1), efficient architecture design (Appendix A.2), and input-dependent architectures (Appendix A.3).
A.1 MECHANISTIC INTERPRETABILITY. SYNTHETIC LANGUAGES, AND RECALL
Work in mechanistic interpretability aims to decompose the capabilities of a neural network into human-understanddalle algorithms that be attributed to specific parameters in the model (Olah, 2022; Power et al., 2022; Elhage et al., 2021; Cammarata et al., 2020). Some of these works, use synthetic data to validate mechanistic interpretations of neural networks (Olsson et al., 2022). This relates to a broader line of work using synthetic languages to study language model architec- tures (Wang and Eisner, 2016; White and Cotterell, 2021; Allen-Zhu and Li, 2023; Ravfogel al., 2019; Xie et al., 2021). Mechanistic design puts mechanistic interpretations to use in designing new architectures and learning algorithms. Several works in architecture research have used synthetic tasks to validate architecture designs (Kitaev et al., 2020).
Our work is focused one particular synthetic task: associative recall. Motivated by psychological models of how humans associate and retrieve information, work in the carly days of neural network research focused developing systems capable of associative recall (Willshaw al., 1969; Feld- man et al., 1981; Hopfield, 1982). For example, Hopfield networks, proposed in 1982, are a recurrent neural network explicitly designed support associative "'conent--addressable"" memory (Hop- field, 1982). More recently, several notable recurrent neural network mechanisms (e.g. neural turing machines, LSTMs, and RNNs with fast weights) were evaluated a synthetic version of associa- tive recall (Graves et al., 2014; Ba et al., 2016; Zhang and Zhou, 2017, inter alia.). These works use formulation of associative recall very similar the single-query associative recall in our work. Since the rise of large language models, several works have argued that LLMs ability to perform in-context learning is due, least part, to the associative recall capabilities of attention (Elhage et al., 2021; Olsson et al., 2022).
A.2 EFFICIENT LANGUAGE MODELING ARCHITECTURES
We first briefly review the efficiency motivations for the recent excitement around gated convolu- tion architectures. While attention requires compute that scales as (N2) in sequence length N, convolutions scale O(N log N) (Cooley and Tukey, 1965). Ideal, inference complexity is O(1) in sequence length, as provided by recurrent neural networks. State-space models can be computed either as a convolution or recurrence, to achieve both sub- lub-quadratic training and constant inference complexity in sequence length (Gu et al., 2021). Architectures that use implicit convolutional filters (Poli al., 2023a), can be converted to SSM via a simple distillation step (Massaroli et al., 2023).
15

---

Published as a conference paper at ICLR 2024
A.3 INPUT-DEPENDENCE IN SEQUENCE MODELING ARCHITECTURES
In an input-dependent sequence model, the way tokens are aggregated across the sequence is con- trolled by the data, not just the model parameters. We highlight several prior works related to our study that explore input-dependent sequence models.
Attention (Vaswani al., 2017) achieves suput-dependent sequence mixng since the Q, K, and V terms are achieved via linear combinations of the input æ. For instance, Q = zWo for some learned weight matrix WQ-
Input-Independent Convolution Architectures. Given the quadratic scaling of attention, Gu et al. (2021) propose S4, an architecture that use long convolutions to process the in- put. CKConv Romero et al. (2022) uses a convolution filter that is implicitly parametrized by an MLP of the relative positions of the observations in a time-series. A line of sub- sequent work improved upon S4 by changing the parametrization (Gupta et al., 2022: Gu et al., 2022: Mchta et al., 2022: Ma et al., 2022; Fu et al., 2023c) or changing the SISO convolutions to MIMO convolutions (Smith et al., 2023). The convolution filter in each of these architectures input-independent The linear convolution layer alone cannot perform MQAR (Fu et al., 2023a).
Input-Dependence via Gating. Since pure convolution architectures with input- independent filters struggle to perform associative recall, important task in in-context learning (Olah, 2022), subsequent work proposed to introduce input-dependence by adding gating or element-wise multiplication operation puts to the operator are defined in terms of the input (Fu et al., 2023a; Poli et al., 2023a). The multiplication generally y = ((x) O x. The filters remain input-independent. These architectures demonstrate promising results on the associative recall synthetics pro- posed in prior work Fu et al. (2023a); Poli et al. (2023a) and provide large downstream improvements over S4.
the architecture, where both the in-
Prior work suggests the models match attention in quality (Poli et al., 2023a; Fu et al., 2023b; Peng et al., 2023) and this has led to increasing use of these architectures (Nguyen et al., 2023; toe, 2023), however our results suggest there is still a sizable gap to attention. We show the gap is largely due to the models' recall abilities and capture this behavior in novel synthetic task. We theoretically show that dimension scales in sequence length when solving MQAR with gated convolutions.
Selective S4 Wang et al. (2023) learns which positions to mask prior to passing the input the next long convolution (S4) layer. This resembles gating (which also controls which information flows forwards to the next layer).
Input-Dependent Convolution Architectures Prior work introduces convolution archi- tectures with input-dependent filters (Yang et al., 2020; Kosma et al., 2023). For instance, Yang et al. (2020) parametrizes the convolution filter as a linear combination of projec- tions of the input. Liquid S4 Hasani et al. (2022) is also motivated by introducing input- dependence S4. To accomplish this, the work proposes including correlation terms be- tween the input tokens during state mixing. We evaluate Liquid S4 and find the architecture lags attention on the Pile by 4.4 perplexity points in Table
Overall
Slices
% of gap due to
Model
Param (M)
AR Hits
Other Tokens
Attention
125
12.37 (2.52)
2.32 (0.84)
14.04 (2.64)
Liquid S4
145
16.80 (2.82)
22.75 (3.12)
16.42 (2.80)
52.6%
Table 3: Validation perplexity on the Pile after pretraining on 5B tokens. We report log perplexity with negative log-likelihood in parentheses. AR vs. non-AR token slices are defined in Section
Recurrent Architectures The pure long- long-convolution architectures can also be computed as recurrences. Letting 8; represent the hidden state time i, u; be the input at time i, and A, B, C be projection matrices, the recurrence computes:
8;+1 s(As + Bua)
16

---

Published as a conference paper at ICLR 2024
y1 Cs.
a linear RNN like S4, the A, B, and matrices are input-independent and f is an identify function. However, recent work proposes input-dependent RNNs (e.g. RetNet Sun et al. (2023)), where subset of these matrices input-dependent (e.g. A = xWa) Note that linear attention mechanisms can also be expressed input-dependent RNNs (Katharopoulos et al., 2020b). While these architectures improve over gated convolutions in MQAR synthetic experiments and downstream quality the Pile (Appendix F), we observe and prove that the required RNN hidden state dimensionality grows with the number of key-value pairs that the model needs to recall in the sequence (Appendix H.6). In contrast, attention complexity does not scale with the number of key-value pairs to recall.
Hybrid Architectures Finally, we can combine architectural components that provide dif- ferent capabilities. Prior work proposes architectures that hybridize sub- rua--uuaratic layers and attention (Fu et al., 2023a; Ma et al., 2022; Ren et al., 2023), but does motivate this choice from a mechanistic design perspective. Further, H3 without attention Fu et al. (2023a) and MEGA (Ma et al., 2022), which uses blocked attention, underperform attention language modeling.
SeqBoat (Ren et al., 2023) shares closest resemblance to our learned selection module. SeqBoat introduces different module that learns where model can end up using full attention at a layer, providing quadratic scaling in the worst case. Different from SeqBoat, attention positions and select the top-k positions to remain sub-quadratic.
attention however, their
auxiliary loss encourage sparsity in the selected
Overall, our work finds that previously proposed sub-quadratic models do not efficiently solve the MQAR task.
B CODE LISTING
In this section, we include code listings for the BASECONY operator. We begin with the a standard BASECONY implementation that uses selliiityyparrmmettriinn long convolution. Below, also provide the implementation for BASECONVWith an
long convolution.
import torch
def fit_ccnv (u: forch.Tensor k: =**
Args:
u (teror] Opatch_size, d_model, seq_len)
k (torch.Tensorl) (d_model, 1_max)
Return:
y (tnsor) (batch size, d modee, seq_
10
11
seglen u. shape [-1]
fft_size = seqlen
13
k_ torch. fft. rfft (k, n=fft_size) / fft_size
14
uf torch. fft. rfft tu cdtype=k,dtype n-fft_size)
15
terch. irfft (u_f * k_f, n=fft_size, norm="forward" [..., seqien]
return
17
18
class BaseConv /to........ouuuu.e.....
19
def iinit_(seff d_model: int, 1_max: int, **kwargs)
super () ()_er((n_i___)))
22
self.d_model, 1max = doddel, 1_ax,
23
self.projection = torch. nn. Linear (self.d m silf.d_mode!).
self. = torch nn. nn. (torch. randn (self. d model, 1 i_max) requires_grad-True
25
26
def forward(self,, u:
27
Args:
Tensor) (batch_size, d_model, seg_len)
17

---

Published as a conference paper at ICLR 2024
Return:
y (torch, Tensor) (batch_size, dnodel, seq_len)
32
.*"
33
Lccon fft_conv (u. transpese (1, 2), self.filter) www.ffilleett.raspsssssste... transposo (2, 2) u_proj self.projectee (u)
y u_conv u_proj return y
36
Listing 1: Explicit BASECONV implementation. Implementation of the BaseConv layer with long convolutions. The implemtnation is 19 lines excluding comments and
whitespace.
In some of experiments, interleave BASECONv layers that explicit short convolution filters (like those torch. mn Convld) and with BAseCoNVlayers that use implicit long con- volution filters (like those described in Poli et al. (2023a)).
Below we include the code for BASECONY with an implicit convolution.
class PositicnalEnteedding (na. Module):
def _iniiinin_t____
emb__di dim: int, seq_len: int, **kwargs)
"""Complex exponential positional embeddings for implicit long convolution filters.
super ().ertt..nntt_... init_ torch, (0, 1, seg [None, :, None] 1,
bands (emb_dim - 1)
t_rescaled torch. www.pace (0, seq_len - 1, len) [None,
None]
math.pi rescaled seq len 1, L, torch. linspace (1e-4, bands 1, bands) [None, None]
10
2 forch.exp(-13 w)
11
torch.cat ( [t, r.rll,, 2. imamg),, dim=-1) self.2 nn. Parammett (2, roquires_srad-Falee
14
def (self,, 1):
15
return self.2[;, :b]
16
class BaseImplicitconv (nn. Modull))
18
19
BaseConv with implicit filter parameterized by an MLP.
20
Args:
model (int): Number of expected features in input and output. max (int): The maximum sequence length.
25
24
é emb (int, optional): Dimension of the positional embeddings. Must be odd and $Ageos to 3 (time, sine and cosine). Defaults to 3. h_iideen (int, optional): The number of features in the hidden layer of the MIP, Defaults =**
26
27
def init_ (self, :_model: int, m int, d_mmb: int-3, d hidden: int 16,): **.
29
30
Long convolution with implicit filter parameterized by an MLP.
31
super () Qi_init_
self.pos_emb PositionalEnbeddigg (d_emb, 1_max)
34
self.filter_miy = nn. nn Sww.eeuuntiall (nn. Linear (d_emb, d_hidden), torch .n. .en. nn. Lineir (d hidden, d_model))
self.projection = torch nn. Linear (d_model, d model)
37
38
def forward u: forch.Tensoa targs, **wargg)
Args: (torch, Tensor) : fbatch_size, seg len, d_madel)
18

---

Published as a conference paper at ICLR 2024
Return:
43
y Tensor) (batch_size, seg_len, d_nodel)
44
45
filter self. self.filter.nipseelf
(u. shape (1])) transposo (1,
2)
c fft_conv f (u.transostet 2), filter) transpose to
stype=u.dttype)
47
_proj self. self.projectiontuu
48
cenov u_proj return
Listing 2: Implicit BASECONV implementation. Implementation of the BASECONv layer with parmmtttiyypr long convolutions. (The implementation is 34 lines excluding comments and whitespace).
C DOWNSTREAM EXPERIMENTAL DETAILS
We use A100 80GB Nvidia GPUs run all experiments. We the reference training infrastructure from https: : //gitttus com/ EleutherAlle hwwpthherrrggttttnneereeeee tokenized using the GPT2BPETokenizer and all models see the data in the same order.
for all pretraining runs. The Pile data
Below we provide details on the hyperparameters and settings for training each architecture studied in the paper on the real-world Pile data. Appendix C.I we summarize and justify our method for measuring the quality gap due to associative recall between the convolution architectures and attention. In Appendix C.2, we provide details on the pure gated convolution architectures, attention baseline, and hybrid architectures studied in the paper.
C.I MEASURING THE MQAR GAP ON REAL LANGUAGE DATA
Here summarize our method for computing the amount of quality gap between the convolution and attention models that is ascribed to associative recall capability (e.g.. in Table 1):
Given an input sequence, we identify recurring bigrams (i.e. bigrams that have already appeared in the sequence prior position). Since bigrams that appear frequently during training may be memorized by the model, rather than requiring the model to perform re- call at inference-time, are seen fewer than threshold number of times during training. The threshold used in all the experiments in our submission is 250 training occurrences in the 10B tokens of pretraining data.
only measure AR log-rrobabilities with respect to bigrams that
2. We measure the log- probability assigned to the true bigram completion. This bigram com- pletion is referred produce the completion by recalling the prior occurrence of the bigram the sequence.
AR Hit in our work. This protocol assumes that the model can
3. For the model being evaluated m, and attention model M, we measure the of the quality gap between and M ascribed to associative recall capability follows. Let the average log-pno probaabibt for all AR Hits across validation sequences be 17 and 1 for m and M respectively. Let the average log-pue probaabbiliiss of all tokens in the validation sequenees be 1a and MM respectively. Let ри be the proportion of AR Hit tokens in the validation data. As the final gap ascribed to AR, report:
(mT 1M PH 1.0)
min
Im
IM
Shown above, if is better than attention (M) overall and M is better than AR, ascribe 100% of the gap to AR.
We briefly discuss two important decisions this protocol. First, we only measure explicit bigrams, iec bigrams are identified based token ids in the sequence. However, intuitively, models may also perform associative recall between related concepts produced by a contextual language model. word is swapped by synonym. As
For instance, language may contain bigrams in which another example, a model may see a sentence such as "The iPhone outside my budget so instead purchased an Android phone. ... It was much and predict "cheaper" for the blank, recalling
19

---

Published as a conference paper at ICLR 2024
that the sequence is discussing cost. Our work does not measure such fuzzy (more abstract) recall instances.
Next, we measure the gap based on pogpbabbabllities rather than perplexity. This is simply because we want to make our metrics independent of the number of tokens in each of the slices the vali- dation set. Approximately 6.4% of validation tokens are AR Hits with the threshold set consider bigrams seen less than 1, x during training.
C.2 GATED CONVOLUTION DOWNSTREAM ARCHITECTURES
We evaluate over 4 previously proposed architectures as well as BASECONV, the theoretically "canonical" representation for gated convolutions, introduced in Section 4, for a total of 14 train- ing runs. Here we provide details on the hyperaparamters and configurations used for training cach architecture. We also provide details the FLOPs computation.
Attention (Vaswani et al., 2017; Touvron et al., 2023) We train using the the specifi- cations in Table 7. The parameters sourced from the Transformer implementation in htte/Elluttttrree wew..x..pppto.ooooooooomm neox.
hittps://githie
Hyena (Poli al., 2023a) We train using the specifications in Table 11. The parameters are sourced from the Appendix of Poli et al. (2023a) and the implementation is sourced from the provided reference at https: //////ggg github..hh....g com/ http///yyeeeeeerc/ssseeeeeee saf.eaaraaasaaai.....a..
H3 (Fu et al., 2023a) We train using the specifications in Table 13. The hyperparameters and implementation use the reference https: H3.
com/ Ha syaecchh
RWKV (Peng et al., 2023) We train using the specifications in Table 14. The parameters are sourced from the Appendix of Peng et al. (2023) and the details provided in the ref- erence implementation at https: / thub. com/ B1 winknL/ RWKV-LM. We specifi- cally evaluate RWKV-V4.
Pure long convolution We train using the specifications in Table 16. We evaluate simple long convolution based model with no gating as a reference point. While this is generic architecture, we use the reference implementation and intiallzationsrrgalaaizaiioossss from recent work Fu et al. (2023c). The implementation is provided at https: //gt thub. com/Hazyyeee conn/Haayyeeearrchheeee safari.
BASECONY We train using the specifications in Table 9. The implementation, amounting to 19 lines of PyTorch, is shown in Appendix B.
We provide the equations used to compute the FLOPSs for each model, letting D be the model width, H the head dimension, L the depth, N the sequence length, V the vocabulary size, and B the batch size. FLOPs equations for different architecture types are provided in Table 8 (attention). Table 12 (Hyena and adapted for H3), ?? (pure long convolution). and Table 10 (BASECONV). We compute FLOPs for RWKV in the Appendix of (Peng et al., 2023), based the number of linear layer parameters, plus input and language modeling head FLOPs.
the main paper and appendix, we evaluate a series of architectures with a hybrid of gated convolution and non gated convolution layers. For these architectures, we use the same number of layers as the pure gated convolution architecture (as reported in Appendix C.2). The hybridized (non gated convolution) layers inserted replacements of the gated convolution layer. For each architecture, we simply evenly intersperse the two types of layers. We evaluate with the following replacement layers in this work:
Hybrid Architectures
Full attention following the specification of attention in Appendix C.2.
Random selection We evaluate sparse attention randomly applied to tokens introduced in Section 5.
Programmatic selection We apply sparse attention AR hit tokens as introduced in Sec- tion 5. When processing the input sequence, we can causally determine if a bigram of raw token ids is repeated in the sequence to construct the attention pattern.
Learned selection We learn the positions on which to use attention by introducing a simple linear layer with sigmoid activation to the architecture that is trained output high scores
20

---

Published as a conference paper at ICLR 2024
if attention should be applied to the token. We take the top- positions to control the computational complexity of the layer. This approach is introduced in Section
We use these protocols to validate that input-dependence suffices address the associative recall gap and highlight that there are varied approaches to incorporating input-dependence in an architecture.
D EXTENDED DOWNSTREAM ANALYSIS OF MQAR
In this section, provide additional analysis of how MQAR manifests in real language data. We extend our discussion of associative recall on the Pile Gao et al. (2020). We also perform analysis sources in RedPajama including ArXiv and StackOverflow Together (2023).
D.1 ADDITIONAL. DISCUSSION MQAR THE PILE
The Pile is widely popular language modeling corpus Gao al. (2020).
1. First, in Appendix D.1.I, we provide several real examples of how associative recall curs in the Pile to demonstrate it's role in language modeling. We color code the tokens highlight differences in the next token predictions from Attention, RWKV and Hyena models.
2. Next, in Appendix D.1.2, we explain where associative recall hits tend to occur in se- quences. This is useful guide the design of sequence mixers - if associative recall tends to occur in specific places, the input-dependent sequence mixer may not need to compute all 2 token-to -oknetto-ooken interactions for every sequence.
D.1.1 PILE EXAMPLES
For the examples below, the legend is: tokens are colored as both correct, both incorrect, Attention correct and Hyena incorrect, Attention incorrect and Hyena correct. For tokens where the models disagree, the predictions of each model are provided in parentheses.
while lunch -ing at the Ma -ison Ber -gey -ist -ro near his apartment : he
had been mus -ing about (rwkv= about, attn= on) th. ... (723 tckens) the young waitress -'s sigh at the Ma iison Ber (rwkv= Bl, attn= Ber) -gey (rwkv=-nd, attn=-gey)
Example D.I: Comparing next token predictions of 350M parameter Attention model and 350M parameter RWKY model. In this example, the models need to perform associative recall to correctly predict "Bergey" in the 4-gram Ma-ison Ber-gey. The previous mention of the 4-gram was more than 700 tokens earlier in the passage.
The second (rwkv= first, attn= second) section is all about Pixar Fest and the (rwkv= thi rd, attn= the) final section all (rwkv= about, attn= all) about Pixar Pier ... (480 tokens) If (rwkv=-Disney, attn=-If) there wasn -áG -t enough Pixar at Disney- land Pixar Fest (rwkv- would, attn- Fest) is (rwky- at, attn- is) coming to the Disneyland Resort on April 13, 2018
Example D.2: Comparing next token predictions of a 350M parameter Attention model and 350M parameter RWKV model. In this example, the models need to perform associative recall to correctly predict "Fest" in the bigram Pixar Fest.
21

---

Published as a conference paper at ICLR 2024
David Mus -a P -id -cock is an indigenous English re- vert to (Iwkv=-ant, attn= to) Islam (rwkv=,, attn= who) who formed the Islamic Party of Britain in 1989 London Central (rwkv=., attn= in) (,11 tokens)... -I have just found out that David Mus (zwkv= Cameron, attn= Mus) -a (zrwkv= (, attn= P) -id (rwky=-asha, attn=-id)
Mosque (rwkv= Mosque, attn= University)
-cock (rwky=-ham, attn=-cock) is (rwkv= ( attn= is) speaking at a Yorkshire Forum debate today in Bradford with
Example D.3: Comparing next token predictions of a 350M parameter Attention model and 350M parameter RWKV model. In this example, the models need to perform associative recall to correctly predict a middle and last name. Note that the middle and last names span several tokens.
-The most common map elements (hyena= in, attn= e elements) in Sub -Space are prizes or (hhyena=-pointt. attn=-g) -reens (hyena=-", .. attn-") (-for their green color Pri Pri-zes allow players to up- grade their ships and gain (hyena- other, attn= gain) special weapons or abilities (hyena= abilities, attn= upgrades) ally plent -ifully scattered throughout the map (ayena= map, attn= game), the upgrades abilities (hyena= bonuses, attn= abilities) they award are randomly selected by the zone -Energy -Rather than dealing with ammunition counts and hit points separately, Sub -Space combines both of these elements into single unit of measure: energy Each ship is equipped with a certain amount of energy from which (hyena- which attn- the) must draw its health well as its weapons power... ... (1, 526 tokens)... Speed Zone proved to be less popular than the Jack -pot -/ Running. Chaos, or "-flag" zone games and support was discontinued shortly after Sub -Space (hyena=-space, attn=-Space) went to retail
While prizes are gener-
Example D.4: Comparing next token predictions of a 350M parameter Attention model and 355M parameter Hyena model. In this example, the models need perform associative recall to correctly predict the name of the game "SubSpace". Note that both models correctly perform the recall when there is short gap between tokens, but only Attention correctly performs recall when the gap is greater than 1,000 tokens.
Thus far olation (hyena=-olation, attn=-ot) via sequential centrif -ug ation and no systematic analysis is ... (491 tokens)... ... Fig. (hyena- 1, attn=-A_11 -is -olation via sequential (hyena- centrif, attn- sequential) centrif -gg -ation.
systematic study has been published is
Example D.5: Comparing next token predictions of a 350M parameter Attention model and 355M parameter Hyena model. In this example, the models need perform associative recall to correctly predict the technique "sequential centrifugation".
Miss -ouri (hyena=-ouri, attn=-iss) Southern has had 14 Major
League Baseball draft selections since the draft began in 1965 ... (149 tokens)... Fred G Hughes Stadium (hyena= Stadium, attn= Field)
-opened in 1975 after former (hyena= the, attn= former) publisher and (hyena= and, attn= Fred) Missouri em (hyena= State, attn= Southern) board of reg ents member
is named (hyena the, attn= named)
-op -lin Globe
South-
22

---

Published as a conference paper at ICLR 2024
Example D.6: Comparing next token predictions of 350M parameter Attention model and 355M parameter Hyena model. In this example, the models need to perform associative recall to correctly predict the name of the university "Missouri Southern".
D.1.2 DISTRIBUTION MQAR HIT POSITIONS SEQUENCES
In Figure 3, we compute and plot the distances between AR hits and its prior bigram (key-value) oc- currence in the sequence across the Pile training data. The distances follow a power law distribution where most AR hits are within 100 token positions from the prior bigram occurrence, and long tail of AR hits requires long-range interactions.
Implications This suggests that architectures which compute token-to-token interactions within windows, where the window size is less than the full sequence length, may suffice to handle MQAR in real data. For instance, sliding window attention may help with AR (Beltagy et al., 2020; Sun al., 2023, inter alia.). Prior work leverages the observation that local attention captures most of the token dependencies to improve efficiency Sukhbaatar et al. (2019).
Bigram Distance Frequency in Training Data
400
Frequency (10e6)
300
200
100
50
0
0
100 200
300
400
500
Gap between bigram occurrences
Figure 3: Across the Pile training data, we measure the distances between n-grams and their prior occurrences in the provided context. We plot the frequency across distances, finding follows power law distribution.
D.2 ANALYZING MQAR ACROSS VARIED LANGUAGE DISTRIBUTIONS
We use the HuggingFace sample of the popular RedPajama language modeling corpus Together (2023) understand how the prevalence of MQAR hits and and ways in which the hits appear vary across language distributions. We analyze text sequences from the following sources: ArXiv papers, Books, C4, Common Crawl, GitHub, Stack Exchange, and Wikipedia.
First profile the prevalence of MQAR hits by computing the number of repeated bigram oc- currences per sequence. Each sequence is 2, 048 tokens in length and exclude bigrams containing NLTK stop words or punctuation from the computation. The prevalence of hits and the distances between hit tokens and the prior bigram occurrence in context vary across distributions as shown in Figure and
Arxiv, Github, and Stack Exchange contain relatively structured or richly formatted langauge data. We observe that the prevalence of MQAR hits is relatively high in these three sources in Figure 4. We inspect why these documents contain many bigrams. In Arxiv, we find dooaiin-speciin terminology, groups of related citations, Latex commands, and mathematical equations within the same document. In GitHub and Stack Exchange, we find function and variable names are frequently reused. Meanwhile, the C4, CC, Books, and Wikipedia bigrams highly variable. The topics and documents in these corpora are relatively diverse and unstructured.
frequently repeated
23

---

Published as a conference paper at ICLR 2024
Proportion of Bigram Tokens by Source
Proportion of Bigram Tokens
0.25
0.20
0.15
0.10
0.05
0.00
arxiv
books
github
stack exchange
wiki
Figure 4: MQAR hits by sub-source of the RedPajama language corpus Together (2023). Hits are measured using the repeated bigrams in each sequence of length.
books
250000
40000
43200
30000
350000
20003
50000
3.0
200
330
400
500
10c
200
300
305
100
200
300
400
500
stack eachange
50000
1.75
250000
12000
30003
1.25
1.50
200000
20000
0.75
30000
0.25
30000
0.50
18.....
Distance
Distance
Figure 5: We measure the distance between n-grams and their prior occurrences in the provided context across sequences from each source. We plot the frequency across distances, finding follows power law distribution.
24

---

Published as a conference paper at ICLR 2024
If the placement V of the hyperplanes is not generic, then the induced subdivision E is not a triangulation. However, the above bijection is unaffected and, in particular, the fine type of vertex of CDA can be read off the corresponding facet of Σ. We define the crosscut complex ccutE C 2 x [4] to be the unique simplicial complex with the same vertices-in-facets incidences the polyhedral complex Σ. The crosscut complex is a standard notion in combinatorial topology and can be defined in more generality (see Björner). The following observation immediate from the definitions.
Proposition D.I. For V an ordered sequence of a points in T les = be the corresponding tropical arrangement and let E be the regular subdivision of An x Ad-1 induced by V. Then the fine cotype ideal fel is Alexander dual to the Stanley-Reisner ideal of ccutE.
The crosscut complex encodes the information of which collections of vertices lie in common face. Hence, the crosscut complex is purely combinatorial object and does not see the affine structure of the underlying polyhedral complex.
Example D.7: ArXiv sequence where the initial bigram occurrence is highlighted in red and the repeated bigram with the MQAR hit is highlighted in blue.
/+*
Aads a directed edge to the graph from ptl to pt2. Precondition:
Both
GeographicPoine cPoints have al already been added to the graph
@param from The starting point of the edge
@param to The ending point of the edge
@param roadName The name of the road @param roadType The type of the road @param length The length of the road, in km the points have not already
10
@throws IllegalärgunnntEcett been
11
* added as nodes to the graph, if any of the arguments is null, or if the
length is less than 0.
13
*/
14
public void addEdge (Goographccccii from, GeographicPoint to, String roadName, String roadType, double length) throws 11llgga
15
{
17
18
if (from mull ll to null II roadName null
19
11 roadType == null II length 0 II !mapNodes. containsKey (from)
I1
21
throw new Illegallagumentteooppoooool
22
23
if (ImapNodes. get (from) (ta))
HapNodes.gen (from) adhNeighhor (to, roadName, roadType, length
};
26
numEdges++;
Example D.8: GitHub sequence where substrings such as IIlleaaArgummenEExxptinnn"" "roadType"" (and other variable names), "GeographicPoinn" (and other data types), "containsKey" (and other function calls) are repeated throughout.
25

---

Published as a conference paper at ICLR 2024
E SYNTHETIC MQAR EXPERIMENTAL DETAILS
In this paper, we propose MQAR as a useful tool to help explain gaps between three popular sequence modeling layers Transformers or Sparse Transformers, Convolutions, and Recurrences. In this section, we detail the motivation behind the design of MQAR and include extended experiments on additional architectures. We provide a procedure for generating MQAR synthetic data, which can help in the development of new architectures.
E.1 MQAR GENERATION PROCEDURE
Here we provide additional discussion on the properties of our MQAR synthetic data. The objective of the synthetic analysis is to help explain the differences in language modeling behavior between different classes of language modeling architectures, as observed on real-world data Section 3. Syn- thetic recall tasks were used in the development of Hyena Poli et al. (2023a) and H3 Fu et al. (2023a), and Olsson et al. (2022) to study Transformer in-context learning behavior.
1. Attention: Attention can perform recall trivially (Proposition 4.3). The bound is indepen- dent of the sequence length N.
2. Convolutions: The gated convolution models input-independent filters. We show in Section 4 that the dimensionality for solving recall grows with the input sequence length. Real language modeling can require performing O(N) recalls in one forward pass. Our intuition is that it is difficult to efficiently compute all inten-interaction distances with such filter.
3. Recurrences: Recurrent models include single hidden state an information bottleneck. We show in Appendix H.6 that the gated recurrence requires (2(NV) bits to solve MQAR for d < VN for model dimension d. Our intuition is that is difficult to store large numbers of key-v pairs seen in the sequence in low-dimensional hidden states.
Motivated by this analysis, we propose Procedure for generating synthetic MQAR data. This procedure allows toggling the differences between the sequence modeling architecture families. We also discuss deviations from the prior work that uses synthetic recall data.
simple properties of the synthetic dataset, specified below, reveal
1. Size of the Key-Value Map: The number of unique key-value pairs appearing in each example in the dataset. In Procedure I, we use the input parameter D to toggle this value. 2. Number of Gaps in the Data: The number of unique token-interaction distances required to perform the recall task. Prior work assumes there single query token that requires recall per example, failing to test whether the architecture can support multiple token- interaction distances. In Procedure I, we use the parameters N (sequence length), D (size of the key-value map per example), and to toggle this property. We enforce that the token- interaction distances appearing in the synthetic examples follow a power-law distribution specified by (based on Figure 3). Keeping and D constant while varying N changes the number of unique token-interaction distances appearing in the example. Finally, noting that the vocabulary size appears in the theoretical bounds for all three architecture families (Section 4), we increase the vocabulary size to be much larger than the model dimension as typically seen in language modeling (30k- 50k tokens). Prior work uses vocab sizes < 40 tokens.
Algorithm 1 MQAR Synthetic Data Generation Procedure
Input:
Vocabulary C, Sequence length N, Power-law parameter Number of Key-Value Pairs D Output: Synthetic sequenee
1: Let the first half of C' be keys K and the second half be values V.
2: Pair each key token & e K with a random value token v E V.
3: Sub-select D random key-value pairs to include in the sequence.
4: Place the D key-value pairs at the start of the sequence (i.e. consuming the first 2D positions).
5: Place second occurrence of each d E D a distance from the first occurrence in the sequence. The distance for each d E D is selected at random from the positions [2D..N) where the probability of choosing each position follows the power law distribution specified by a.
6: Output the synthetic sequence.
E.2 TRAINING DETAILS
We first describe the architectures evaluated on MQAR synthetic data and then the training hyperpa- rameters. We evaluate the following four architecture categories in the experiments:
26

---

Published as a conference paper at ICLR 2024
1. Attention Standard GPT-2 style multi-headed Transformer architecture with learned posi- tional embeddings Brown et al. (2020). Attention, the core building block of Transformers, is defined in Section 2. The architecture for synthetics has attention head.
2. Gated convolutions (Hyena Poli et al. (2023a), RWKV Peng et al. (2023)). Gating, con- volutions, and the class of gated convolutions are defined in Section 2.
3. Gated recurrences (RetNet Sun et al. (2023)). RetNet proposes computing attention over chunks of the sequence and combining this with a recurrence over the sequence. We eval- uate this architecture ("Chunked RNN") with chunk sizes of 32 and 8 in Figure 6. In RetNet, hidden states are updated as:
Sn = S + A
Outputs at each timestep are defined as
On Sn
where y TR.. A = zWa for input x E N x and learned weight matrix Wa E Rdxd, C= xWW for We E Rdxd zW, for W, E Rdxdd
RetNet similar to a state-space-model (SSM) Gu et (2021), except for that the matrices A and are input-dependent (i.e. they are functions of the input æ). We note that Ret- Net is special case of the recently proposed Mamba Anonymous (2023a) and GateLoop Anonymous (2023b) architectures, which replace the y term in RetNet with yet another input-dependent matrix.
The RetNet model is formally defined in Appendix H.2.3.
4. Sparse attention Several works have proposed methods for efficiently computing attention Tay et al. (2022). We evaluate two classical methods: sliding window attention Beltagy et al. (2020); Zaheer et al. (2020) and blocked window attention Child al. (2019b); Qiu et al. (2019).
Consider window size w. Sliding window attention permits each token to attend to the prior tokens in the sequence. Blocked window attention first splits the sequence into blocks of size w tokens, and then computes standard causal attention within the block.
For all synthetic runs in the main paper and the appendix, we use the following training protocol: Optimizer and schedule: Weight decay 0.1, warmup duration 10%, linear warmup, AdamW optimizer. For each run, we sweep the learning rates in np.logspace(- (44, -2, 4). We train for 64 epochs.
The global batch size is 8 for input sequence length or model dimen- sion 512, 16 for input sequence length or model dimension 2 256, and 64 otherwise. Width and depth: For all synthetic experiments, we use exactly two-four layers (each with one sequence mixer and one MLP, interleaved with layer normalization). The model dimension, sequence length. and number of KV pairs are varied based on the relevant experiment (see Appendix F).
Training duration:
Position information: No position embeddings are used for the pure convolution runs (input-dependent nor input-independent filters). Position embeddings are used for the runs with any attention variant.
Data: We train and evaluate each model on 100, 000 and 3, 000 data points respectively. The data and data order for all runs are constant.
F EXTENDED RESULTS ON MQAR ACROSS ARCHITECTURES
In this section, we provide further validation that MQAR is useful diagnostic tool to explain the behaviors of a wide variety of architecture families. We show:
1. Gated convolutions In the main paper, claim the gated convolution models with input-independent filters use larger dimensionality than attention as the number of token- interaction distances required for the task increases (N, holding a and D constant in Pro- cedure l).
2. Multi-head gated convolutions Recent work Fu et al. (2023a); Massaroli et al. (2023); Fu et al. (2023b) uses multi-head convolutions. This extension groups the model dimension into heads (of dimension d, = & for d model dimension and / heads), computes the outer product between different input projections in each head (resulting in R N xd, xd,, tensors), passes the results through the long convolution, and multiplies with another projection of
27

---

Published as a conference paper at ICLR 2024
Example:
Example
KVs Per Example:
KVs Example: 64
1.00
0.75
Architecture
0.50
Attention
0.25
Rether Ehunked IMN (52)
Mtyon
0.00
Rethet - Chunked RAN (8)
236
128 Model Dimension
Model Dimension
Model Dimension
Model Dimension
KVs Per Example:
KVss Per Exemple:
KVs Per Example:
KVs Per Exempler 64
1.00
Accuracy
0.75
Architecture
0.50
Block Attn. (32)
Black Alto. (8)
0.25
0.00
Sliding Attr (81
256 64
64
128
128 Model Dimensian
Modal Dimension
Model Dimension
Model Dimension
Figure 6: MQAR quality numbers of key-value pairs per example. (Top) We focus on the evaluation the input-dependent gated recurrent RetNet (Sun al., 2023) architecture. The architecture combines chunked attention with recurrence and we evaluate at two different chunk sizes (8 and 32). For reference, we include gated convolution (Hyena) and attention. (Bottom) We evaluate efficient attention variants sliding window and blocked window attention at two window sizes (8 and 32).
vary the model dimension for synthetic datasets that contain different
the inputs return RNxxs then concatenates. The full algorithm is provided in Fu al. (2023b) (Appendix Algorithm 1).
3. Gated recurrences In this section, show gated recurrrent models use larger dimension- ality than attention to solve the task as the number of unique key-value pairs to store (D in Procedure 1) increases.
4. Sparse attention In this section, we validate that sliding window attention helps close the MQAR gap when the token-interaction distances are relatively short and degrades on longer distances. Meanwhile, blocked window attention struggles to close the gap (intu- itively tokens earlier a block are able to interact with few other tokens in the sequence). Concretely, in Appendix D, we observe that repeated bigrams in the Pile and RedPajama corpora tend to occur within neighboring tokens in the context.
We finally show the architectures that perform well MQAR also perform well on downstream real- world AR on the Pile. We use the experimental protocols outlined in Appendix for the synthetic experiments and Appendix C for the downstream experiments in this section.
F.1 SYNTHETIC EXPERIMENTS
Experiment 1: Increasing the number of token-interaction distances per example. To con- struet the synthetic dataset following Procedure 1, we set = 0.1 and |C| 8, 192 for all experi- ments. We vary N {64, 128, 256, 512) to increase the number of token-interaction distances that will appear in the examples and we vary the model dimension E {64, 128, 256, 512]. The results of this experiment are shown in Fig. 2, validating that the gated convolutiosn use larger dimensionality than attention as the number of token-interaction distances required for the task increases.
Experiment 2: Increasing the number of key-value pairs occurring per example. To construct the synthetic dataset following Procedure 1, we set a = 0.1 and |C| = 8, 192 for all experiments. N { 32 T6 8 ] to increase the number key-value pairs oc- additionally vary the
We fix N 256 and vary curring per example. For each number of key-value pairs per example, model dimension E (64, 128, 256). The results are shown in Fig. 6. We validate that as the num- ber of associative recall keys and values exceeds the window size, the model again uses increased dimensionality relative to full O(N2) attention to solve the task. While Fig. 6 highlights RetNet, note that it is special case of the recently released Mamba Anonymous (2023a) and GateLoop Anonymous (2023b) architectures. Corresponding to this experiment, we theoretically show in Ap- pendix H.6 that the gated recurrence uses Q(N) bits to solve MQAR for d < / VN. Intuitively, it is difficult store large number of key-value pairs in low dimensional hidden states.
Experiment 3: Increasing the range of token-interaction distancs per example. To construct the synthetic dataset following Procedure
set 0.1 and |C| = 192 for all experiments. ] increase the number of key-value pairs occurring
We fix N 256 and vary D E { 32
N
N
28

---

Published as a conference paper at ICLR 2024
Sequence Length; 64
Sequence Length: 128 Sequence Length: 256 Sequence Length: 512
1.00
Accuracy
0.50
Hyena w/ 1-Head
Hyena w/ 2-Heads
0.25
Hyena
0.00
Hyena 8 8-Heads
64
Model dimension
Model dimension
Model dimension
Model dimension
Figure MQAR quality as we vary the model dimension. We consider number of convolution heads
{1,2,4, 8) and observe more heads improves scaling (at the cost of additional FLOPs).
Overall
Slices
gap due AR Hits
Model
Param (M)
TFLOPs
AR Hits
Other Tokens
Attention
1.52
12.99 (2.56)
2.41 (0.88)
14.76 (2.69)
Long Conv
76
1.20
20.28 (3.01)
40.25 (3.70)
19.25 (2.96)
44.4%
H3
72
1.33
15.78 (2.76)
13.86 (2.63)
15.94 (2.77)
63.2%
Hyena
72
1.34
15.13 (2.72)
9.00 (2.20)
15.74 (2.76)
60.8%
RWKV
72
1.89
16.10
14.11 (2.65)
16.26 (2.79)
57.9%
Attention
125
2.46
11.01 (2.40)
2.16 (0.77)
12.45 (2.52)
Long Conv
128
1.74
16.98 (2.83)
25.62 (3.24)
16.46 (2.80)
H3
168
2.55
12.06 (2.49)
6.75 (1.91)
12.60 (2.53)
88.4%
RWKV
169
2.08
11.64 (2.45)
5.70 (1.74)
12.29 (2.51)
100.0%
RetNet (128)
152
2.60
(2.41)
3.01 (1.10)
12.45 (2.55)
100%
Hyena
160
2.41
11.60 (2.45)
5.00 (1.61)
12.28 (2.51)
100%
Slide Attn. (64)
159
2.28
11.57 (2.45)
(1.65)
(2.51)
100%
+ Slide Attn. (256)
159
2.29
10.65 (2.37)
11.52 (2.45)
100%
Attention
360
6,23
9.44 (2.25)
1.98 (0.69)
10.62 (2.36)
Long Conv
360
4.08
13.13 (2.57)
13.27 (2.59)
13.12 (2.57)
40.5%
H3
357
4.85
10.38 (2.34)
4.81 (1.57)
11.00 (2.40)
65.8%
Hyena
358
5,03
10.07 (2.31)
3,83 (1.34)
10.75 (2.38)
98.2%
RWKV
4.31
9.79 (2.28)
3.82 (1.34)
10.51 (2.35)
100.0%
Table 4: Language modeling validation perplexity on the Pile. After pretraining on 10B tokens of Pile data, we report log perplexity with negative log-likelihood in parentheses. We report overall scores, and for the AR vs. non-AR token slices defined in Section 3. FLOPs are computed for inputs of 2048 tokens based on equations Appendix
per example and the range of pairs per example, we additionally vary the model dimension E {64, 128, 256].
distances per example. For each number of key-value
The results are shown in Fig. 6, validating that sliding window attention closes the gap to attention when the number of KV pairs is within the window size. I.e., 16 KVs means that there are 16 keys and 16 values, for 32 total tokens. Sliding attention with window size 32 performs well up until 16 KVs and degrades beyond this point relative to attention models of the same dimensionality, Sliding window attention with window size does not perform well on any setting, as expected. Blocked attention also not perform well intuitively tokens the beginning of block lack sufficient prior context to attend to.
Experiment 4: Increasing the number of convolution heads. In Figure 7, vary the number of heads in the gated convolution layers of Hyena following Algorithm 1 in Fu et al. (2023b). We fixed model dimension.
observe that the models with more heads achieve higher recall quality This provides empirical evidence for theoretical arguments in Massaroli et al. (2023) that suggest multi-head convolutions increase the effective dimension and enable solving associative recall with reduced model dimension vs. single-head convolutions. However, additional heads increase the FLOP count and we continue to observe large gaps relative to attention-level recall quality.
F.2 DEMONSTRATING MQAR AS TOOL FOR ARCUITECTRRE DEVELOPMENT
In Table 4, Table we show that the trends of RetNet, sparse attention, and multi-headed gated long convolutions MQAR also carry to the downstream Pile results. We believe the MQAR synthetic may be useful future architectural development.
29

---

Published as a conference paper at ICLR 2024
Overall
Slices
% of gap due to
Model
Param (M)
TFLOPs
AR Hits
Other Tokens
AR Hits
Attention
125
2.46
11.01 (2.40)
2.16 (0.77)
12.45 (2.52)
Attention
360
6.23
9.44 (2.25)
1.98 (0.69)
10.62 (2.36)
Hyena 14)
2.41
11.60 (2.45)
5,00 (1.61)
12.28 (2.51)
100%
Hyena (w = 2)
158
2,41
11.47 (2.44)
4.63 (1.53)
12.30 (2.51)
100%
4 Heads
150
2,47
11.05 (2.40)
3,47 (1.24)
12.07 (2.49)
100%
Hyena (10 = 14)
358
5.03
10.07 (2.31)
3.83 (1.34)
10.75 (2.38)
98.2%
Hyena (w = 2)
358
5,03
9.86 (2.29)
3,46 (1.24)
10.69 (2.37)
100%
+4 Heads
358
5.18
9.48 (2.25)
2.78 (1.02)
10.42 (2.34)
100%
Table 5: Language modeling validation perplexity on the Pile. After pretraining on 10B tokens of Pile data, report log perplexity with negative log-likelihood in parentheses. We report overall scores, and for the AR vs. non-AR token slices defined Section 3. FLOPs are computed for inputs of 2048 tokens based the equations in Appendix C. Here evaluate Hyena at two different hyperparameter settings, sign frequencies w (w = 14 the default in Poli et al. (2023a)), and the multi-head setting (described in Algorithm 1 (Fu et al., 2023b)).
G MQAR PERPLEXITY GAP AND MODEL SIZE
In this section, we investigate how the associative recall gap changes as we increase the model size to billion parameters and beyond. Overall, the results suggest that the larger models improve in memorizing bigrams from the training data, but do not rapidly improve in in-context learning (perplexity on rare bigrams).
Below, we pretrain at the 1.4Bn parameter scale and observe that the 70M parameter attention model full perplexity point better AR than this 1.4Bn Hyena model (Table 1). We discuss these results further in (Section 3).
Overall
Slices
% of gap due to AR Hits
Model
Param (B)
TFLOPs
AR Hits
Other Tokens
Attention
1.4
1.52
8.19 (2.10)
1.91 (0.65)
9.86 (2.29)
Hyena
1.4
1.20
9.65 (2.27)
3.43 (1.23)
11.01 (2.40)
40.3%
Table Large-scale language model validation perplexity on the Pile. After pretraining on 50B tokens of Pile data, we report log perplexity with negative log- ligeliill in parentheses. We report overall scores, and for the AR vs. non-AR token slices defined Section FLOPs are computed for inputs of 2048 tokens based on the equations in Appendix C.
Open-source 7B parameter models We next evaluate the RWKV (gated convolution) Peng al. (2023) and Llama 2 (attention) pretrained models at the 7B parameter scale Touvron et al. (2023).' These are both popular models that took significant amount of effort train, towards maximizing quality. We find that there is a gap between RWKV and attention at the 7B scale and that it increases the model needs to conduct more recalls per sequence (P below). We summarize the experimental protocol below.
1. Justifying the experimental protocol. Since frequent bigrams may just be memorized by the model and not require in-context recall, our work measures AR quality on infre- quent bigrams in validation sequences (Figure Appendix C.1). We do not have access to the custom training data mixtures used in training RWKV or Llama 2 to measure bigram frequencies, so we use synthetic test to fairly measure the AR capabilities,
2. Evaluation dataset. We construct the synthetic dataset following Algorithm 1, where each sequence contains P token pairs, or "bigrams". Each bigram, containing a "key" token and "value" token, appears twice in the sequence. On the second occurrence of "key" token, the model should look back to the prior occurrence to output the corresponding
We specifically evaluate the RWKY-Raven model downloaded from https: https://buuggee htt.b/gggnggac..... github. wttthbb.com/faccbookkeeeecc 1acc// rwkv and the I.lmma 2 model downloaded from https:
30

---

Published as a conference paper at ICLR 2024
Perplexity of RWKV and LLAMA models at 7B Parameters
RWKV AR
175
LLAMA AR
150
LLAMA Non-AR
125
Perplexity
100
75
50
25
0-
10
20
30
50
Number of Key-Value Pairs per Sequence (P)
Figure 8: Perplexity of RWKV 7B and Llama2 7B parameter models on synthetic MQAR data as function of the number of recalls required per input / number of key-value pairs per input (P). Inputs are constructed using each model's full vocabulary.
"value" token. We measure the AR perplexity of the model based its ability predict the correct "value" token as the next word for these repeated keys. The sequences are constructed using the models' vocabulary tokens.
3. Evaluation details. We evaluate (inference only, no training) on sequence lengths of 1024 tokens. We measure the AR perplexity when the sequences contain {16, 32, 64, 128, 256} key-value pairs, using 1000 samples per P value. The tokens that do not contain a key or value are simply filled with a fixed token id (so this token repeated frequently in the sequence). We plot perplexity for AR and non-AR tokens (fixed token) vs. P. We find RWKV quality degrades with P on the AR slice (blue line), while all other lines remain flat. MQAR remains problematic scale.
Overall our results suggest that simply scaling the model size may not close the MQAR quality gap between the gated convolutions and attention.
Measuring the AR Gap at Scale In the main paper, we measure the AR perplexity on downstream 1, 250 during pretraining. However, hypothesize that larger models may memorize a larger number of bigram pairs seen in the training dataset, but do not rapidly gain the capability to perform associative recall as well attention in-context. In-context learning is defined as learning from examples provided in context (Xie et al., 2021).
data based bigrams that
Concretely, the gap between the Hyena / RWKV and attention models at the 350m scale 1.85 / 1.84 perplexity points when we focus on bigrams seen < 1, 250x during pretraining (Table 1). If instead focus on bigrams seen during pretraining, the gap to attention quality 12.0 / 13.2 perplexity points respectively. The gated convolutions appear to struggle in the regime of rare bigrams that require the model to use the context.
H DETAILS ON THEORETICAL ANALYSIS
This section provides proofs and extensions to the theoretical results in the main paper.
H.1 PRELIMINARIES AND NOTATION
H.1.1 NOTATION
We denote the all 1 row vector of size k given by [1 1 1 1]. and the all 0 row vector of O 0]. as 1 and 0*, respectively. We also construe the standard basis vector e; as a column vector in these notes, and adhere to the following matrix indexing convention: and the jth column, M[i, E FF xn denotes the ith row, and M E F= denotes the jth column of M E ÆmxX, We then use 1 xx 0mxn e F X to denote the matrix of all 1s and Os, respectively.
size k. given by [0 0
M[ijj] is the entry in the ith
31

---

Published as a conference paper at ICLR 2024
Next, we denote the Hadamard product of vectors u, E F" uov; the operation can be extended to matrices by applying the Hadamard product column- across the matrices. This is commonly referred to as (element-wise) gating. For vectors u. FF,, also denote their linear (or acyclic) convolution as u v and eyclic convolution as u @ V.
Polynomial Notation. Because convolution is intimately tied to operations on polynomials, is convenient to use them to discuss the inputs and outputs of gated convolution models. Let us define maps poly, poly' F + (Xx) such that
n-1 (oly(u)= Σ uliix', and poly"
Σ u1iix-11-
This allows us to map between vectors and polynomial. Accordingly, we also define coeff : F F|X]/(X*++) IF the map converting polynomials back vectors: coeff coef((u(X)) with uli defined as the coefficient in (XX) at degree i.
These operations allow to interpret the convolution of vectors in terms of polynomial multiplica- tion (Heideman and Burrus, 1988). More specifically, we have
coell (u(X) +XX) mod X"), and t & t = coeff (u(X) v(X) mod - 1)-
the Hadamard product of We can similarly interpret the Hadamard product of vectors polynomials (XX)= 3 v(X):
u ④ e coeff (u(X) o x(X) = coeff Σ (ul) ml)) X
Arithmetic Circuit Notation. We briefly introduce the notation of arithmetic circuits (Volkovich, 2016), the focus of Appendix H.5. An arithmetic circuit C with variables X A T2,....... n over field IF is interpreted as a directed acyclic graph, where the input nodes are labelled by either the variables from X or constants from IF and the internal nodes are labelled by + or with the output being the polynomial computed at the output node.
We shall also refer to the size of the circuit as the number of nodes, the depth of the circuit as the length of the longest path between an input node and the output node, and the width of the circuit as the number of parallel operations in the circuit, "cut' through the circuit. Moreover, the degree of a circuit is defined as the degree of the polynomial computed by the circuit. We summarize this with the following definition:
'wires' which will be intersected by horizontal
Definition H.1. An arithmetic circuit C is an (n, s, A, wiicircuit if C is an n-variate arithmetic circuit of size and of depth at most A, and width w.
Model Notation. Now we introduce the notation we will be using for defining layers. In what follows, we denote u N xd as the model input; N as the sequence length; as the number of stacked layers, indexed by &; and as the input (embedding) dimension.
H.1.2 SUMMARY OF THE RESULTS
The outline of our results are as follows: In Appendix H.2 we introduce gated convolution models and define H3, Hyena, RWKV, RetNet, and BASECONV. In Appendix H.3 we introduce a set of primitive operations that BASECONV can implement. We use these tools in the subsequent proofs. Then, in Appendix H.5, we show that a general arithmetic circuit of size a and degree at most A can be simulated by BASECONV. We use the above results to show all models built from gating and convolutions can be simulated by BASECONV. This helps us analyze the class of gated convolutions, beyond specific proposal (e.g. Hyena).
Next, we study the representational power that gated convolutions and attention use to solve MQAR In Appendix H.7, we derive BASECONY model inspired from dyadic intervals show the di- mensionality for gated convolutions (with input-independent filters) solve MQAR. Next, we analyze the dimensionality for BASECONV model with data-dependent kernels solve MQAR in Ap- pendix H.8.
GATED ATTENTION-FRES MODELS We now present formal definitions of gated convolution and recurrence models with respect to the 5-tuple of parameters (N, L, d, N d''.')
H.2
32

---

Published as a conference paper at ICLR 2024
Definition H.2. An (N, L, N' NN,, d') -attdd Convolution Model is a stacked sequence sequence
model with L layers such that:
1. Input and output are N x matrices,
2. Each layer's operations consist of element-wise gating, convolution, linear projection, and
3. All the individual gated convolution layers take in N' x d' matrices and output N' d' matrices. We refer to the tuple (N', d") as the inner dimension of the model.
We define the Hyena, RWKV, RetNet, and BASECONV layers make step more concrete. We N' such that also assume that the input IN Nxd is embedded into u' R
u' [n, [u(n, i t<d otherwise.
The output from the last layer = E R N'Nx xd' is transformed into output y R by extracting the top left N x d entries in z.
Next, we define the class weight matrices that we will use in the linear projections in the models: Definition H.3. The linear projection Linearm,m has its matrix representation the weight ma- W E mR taken to be a K-matrix for W (BB*) poly- log
poly poly- log m (Definition H.17). Conse-
quently, each matrix W has ((m) parameters and runtime for matrix vector multiplication, and allows us to represent general linear transformations with low-depth linear arithmetic circuits (Dao al., 2020). We will also need linear maps Linearm,n for m n, where each take the correspond- ing square matrices from Linearm,n and note that such matrices has Ó(n) parameters and runtime for matrix vector multiplication.
Remark H.4. We note that the weight matrix W E Rdxd above is taken to be dense matrix in the experiments. However, for theoretical results, restrict the linear projection in our models as per Definition H.3. In the rest of the paper, unless mentioned otherwise all linear projections will follow Definition H.3.
Next, we define three popular gated convolution models that we study in our work: Hyena Poli et al. (2023a), RWKV Peng et al. (2023), and RetNet Sun al. (2023).
H.2.1 THE HYENA LAYER
We will now outline the Hyena layer (Poli al., 2023b). Hyena takes sequence E RNx as input and produces L ++1 projections .'....,''' , v by passing though a linear layer and applying short convolution afterwards. The algorithm then recursively performs point-wise multiplication of the projection with the convolution of the filter h! with the previous output. We summarize this process in Algorithm
Algorithm Projection (u, h)
Input: Input sequence u E IR Nxd, a short convolution filter h E RRN.
1: In parallel for 0 < n < N =|n + Linear (w(n, :)) so that é E IR NxL++1)
2: In parallel for 0 < / < (L+ + 1)d: 2|: t] <-
3: Reshape and split 2 E IR Nx(++1) into p' ......, pd v. where p¹, p u Nxd for & [Z].
4: return pournpp.... p!.... p'....,''
Algorithm 3 Hyena (u, h. h2)
Input: Input sequence u E R Nxd set of convolution filters h² h²...., ... h² Nxd short convolu- tion filter h E N.
1: p ,'. p" v + Projection(u, h.).
2: 00 2-
3: for forl=11..
do
4:
In parallel for 0 < d: 1. p' [t [k] x (
5: return
LL.e. Linearm
(=) W.
33

---

Published as a conference paper at ICLR 2024
Remark H.5. In Algorithm 3, L is used to denute the nuber of recursive applications of the Hadamard product and convolutions, not the number of layers Note that asymptotically, the re- cursive step does not make difference.
Henceforth, we will refer to model consisting of L. Hyena layers is a gated convolution model with associated tuple (N, L, d, N, (I+ 1)d) as (N, I,,d, N, (T, 1)d) Hyena.
H.2.2 THE RWKV LAYER
We will now describe the RWKV layer (Peng et al., 2023). RWKV is typically referred to as an RNN, but, like some other recurrent models (e.g. S4 (Gu al., 2021)), can also be viewed convolution. Here, we present the convolutional view of RWKV. We will see that it is closely related the Hyena layer.
RWKV takes sequence E IR Nxd as input and applies a short convolution. Next, it produces three projections q, k, v by passing u through linear layer. Then, performs convolution sandwiched by element-wise multiplication. We summarize this process in Algorithm 5.
Algorithm RwKVProjection
Input: Input sequence E IR Rxd a short convolution filter h IR
1: In parallel for S f < d: _[; I + u[;
2: In parallel for 0 S N Im, Linearasd ($/n,:) so that I Nx3d
3: Reshape and split R. Nx3d into q, e. Rxd.
4: return q, k, v.
In practice, the short convolution filter h, is restricted length two filter with h.[0] μ and h.[1] = where E [0, 1] is learned parameter. In the RWKV paper, the short convolution is referred as "time shift".
Algorithm 5 RWKV (u, h ht)
Input: Input sequence u E R N xd set of convolution filters h' h'...., ... h" IR Nxd short convolu- tion filter h, E R.
1: q, k, U + Projection(( hg).
2: In parallel for S f d: [. (g)::)) o (h'1; []+ (softmax (A) o v)(;t(])
3: return
In practice, the long convolution filter h is also restricted to h[::0] = c=((-1), where w E R is learnable parameter that controls how quickly the magnitudes of the filter decreases grows.
To summarize, the differences between Hyena and RWKV are: (1) Hyena applies the short convo- lution after the linear projection whereas RWKV applies before, (2) RWKV includes non-linear activations (sigmoid and softmax) while Hyena does not, (3) RWKV and Hyena use different pa- rameterizations for the convolutional filters, and (4) Hyena recursively performs point-wise multi- plication of the projections with the convolution filter whereas RWKV performs this operation only once (though, in practice, Hyena uses a recurrence depth one, making them equivalent in this regard).
H.2.3 THE RETNET LAYER
In this section, we introduce the RetNet layer Sun et al. (2023). To this end, we take an input sequence R Nxd and project it using learnable weight matrices. We then compute the states x" recurrently as in line and the output Out[n, :] as in line 6 for each n E [N] (see Algorithm below).
34

---

Published as a conference paper at ICLR 2024
Algorithm 6 RetNet (u, h h.)
Input: Input sequence E R RNx scalar y E R, learnable weight matrices WA. We, Wv e Rdxd
1: A, C.V+ + wW wWa c, u W so that A, C, C
2: Initialize the output Out E R Nxd
3: Initialize the state x + (A] Vl
4: for S N do
5: 2" f- 72"-1==-- (A[n,:)) V/n, for z" E R
6: Outin, | f :|="
7: refurn Out
H.2.4 BASECONV
Finally, we introduce the BASECONv here as follows:
Y = (uW + b1) @ (u h b2),
(4)
with input u E VR xd weight matrix W E TR'''dd xd' and bias matrices b, E NR xd' defining linear projections of the input sequence, and h E IR N''x' is the set of the d' mixed length filters.
The corresponding pseudocode for BASECONY is as follows:
Algorithm/ 7 BASECONV (u, W, b1, h b2)
Input: Input sequence
R N'x' linear mapping W Rd''d' that W (BB*) (BB*) log d' log
convolution filter h E I ( 'd'' bias matrices b1, b2 E IN N'd'"
1: In parallel for S n < N' æ[n, : Lineara.d (u(n,:])
2: In parallel for 0 5 f < d' 2[:,t] h[;,t] f) * u(::t)
3: In parallel for 0 < \ d' y y < (ml + by (:(:d)) (2|,,4 b2(3s4)- See equation
4: return
Note that the convolution in Algorithm 7 is not limited to causal convolution and allows circular convolution as well. For simplicity, we use here and will disambiguate when required. We will start by specifying the parameter count and complexity a single layer of BASECONv below. Proposition H.6. A single layer of BASECONV requires O(N'd') parameters and runtime.
Proof. We know the parameters and runtime for the linear part of BASECONY via Definition H.3 d(N'd'. Further, each convolution operation requires O(N') parameters and O(N' log N') runtime, and we employ d' such operations. Finally the Hadamard product step takes ((nd) time.
Similar to Hyena, will refer to model consisting of L BASECONY layers as (N, L, d, N', d') BASECONV. our experiments, we extend BASECONY by adding an MLP after Algorithm For simplicity we will denote BASECONV (u, h, by, b2) as BASECONV (u, h) when b b2 = 0. We will now show that there exists BASECONV model that can emulate each of the basic operations in Algorithm 7.
Lemma H.7. The functions Lineard,a (u) (Definition H.3), with d, d' defined as in Algorithm 7, convolution with filter h E TR Nxd and element-wise gating can be computed with Algorithm 7 via (N, 1, N,d') BASECONV.
Proof. For each operation from Definition H.2 and Algorithm 7:
1. For any input E R. x Linearaa (u) with matrix representation W E N'x' can performed by single BASECONV layer computing BASECONY (y, W, h, b,, b2) with bi and by being the matrix of all Os and all 1s. respectively while and the convolution with the zero filter. That is, we have
Y (uW +0 xd (uu0 xd' +1 N'd (uW)@1 N'xd uW Liinarra (u).
35

---

Published as a conference paper at ICLR 2024
2. For any imput RNx, convolution with filter h R. Nxd can be performed by single BASECONY layer computing BASECONV (y, W, h, b1, by) where W, by are all zeroes, and b is the matrix of all Is so that get
(n0 N' xd' + N' xd'
@ (u N' xd = N' xd @ (u h) = u h.
3. We may compute element-wist wise gating between matrices u, v E R
Nxd where e is some fixed factor, with single layer computing BASECONv (v.,) ON''x' xd' is the identity filter, respectively, by Definition H.2.
Co, v, O xdd where
Y = (uo xd' + @ (u * eo + 0 xd = ) u.
H.3 PRIMITIVES
In this section, we will establish some additional basic primitives that we expect a gated convolution model emulate: shift, remember and add. We specify them below:
1. Shift an sequential input of length N up or down by entries:
shift (y,8),
(y, 8)
Input: y
R Rxd, s 2
Output: 2 E Nxd where z+ shi shift.down(,,) and z
shift.up(y,s)
ys
< Vi-1
+- 0 +
yN
Yi
yo
0
UN
--
+- 0
2. Add a sequence x E R to a running sum S E R for some 2n < N' with both æ and S contained as subvectors in y R Nxd.
add,, (y x S):
Input: sequence y containing æ; S E Rxd for 2n & N such that y|0 n n- x. y/n 2n
S and y N 1] 0 2-2 2n
Output: E IR xd containing the sum y + S such that y(0 1] = 1", y/n
2n = and 2 222 N- 1] = ~2n.
1
+ S
S +
T
0
0
36

---

Published as a conference paper at ICLR 2024
3. Remember e E Rmxd as part of sequence of input y IR Nxd while performing gated convolution only on x E R Rnx for some m, n & N x d.
remember m.m.s, (y: : 2, U. h, p):
Input: sequence y E RNxd containing a E Rexd v E Rm and modifiers
p, h Roxd such that y(0 - 1] x. yln +
+
1] = and
yli] = 0 otherwise with * h E R(m+s)xd and v * h E (m+ttxdd
Output: x RNx containing (p (x h)) R ((*+s)xd, such that:
h2)
0
f- 1
v
y =
x
0
These primitives are building blocks of our proofs in the sequel. We will show that each of these primitives can be solved by some (N, L, d, N', d') BASECONv model with a small constant Proposition H.8 (The Shift Primitive). For any E RNxx there exist (N, 1, d, N, d) BASECONV and (N, 3, d, N.d) d) BASECONV that computes sh ft_down(() 8) and sh ft.up(y, 8) for any S N.
Proof. Define the following kernel dependent on 8 S
1 od otherwise.
if 8+1
[, :|
We now deal with the down and up shifts separately:
1. We define W == ONxd by == 1 N by := ONxd Then, for input y E Nxd, BASECONv (y, 0xx hg, Nxxd (xd) for BASECONV in Algorithm 7 is given by equa- tion 4 as
y h.
Now, to perform shi shifft.downg..... s), we note that
Y = yl helts = cocff (ooffyyl (x)) hx(:))((x))
N-1
cocff
Σ
vit, X
X² mod x-
1=0
coeff
Σ
y(i x mod X-
N-1+8
coeff
Σ
y/i -8,1] x mod x-v
N Σ
coeff
yli
8. x
which implies that exactly get what is specified in the output.
2. We again define W = xd xd by == 1Nd b2 = xxd Then, for input y E Nxd, BASECONV (y. Nxd eo, 1 xxd 0 xxx for BASECONY Algorithm 7 is given in equa-
tion 4 as
Yo y Co-
37

---

Published as a conference paper at ICLR 2024
Now, to perform shift.up(y, s), as before, first apply the circular convolution to reverse the input
N
Yo[..j] y[: @ coefl
1- 1 x²
We then apply Yi shi ft down
Y0,8) 8) to get
N-1
Y1 [:]: = coeff
Yg/N (i- x
= coeff
- - x
Finally, apply another circular convolution with the identity filter replace N with i to get
N
Y2l. = ki @ co cocff
s, - x
vii+
==0
Here, we note that we can compute both of these primitives in one and three layers, respec- tively (see Lemma H.ll).
Now, we present BASECONV model with two layers that implements the add. x. S), the purpose of which is to add some window of computation Proposition H.9 (The Running Sum Primitive). For any x, S E R Rxd contained some e RNxd there exists (N,2, d, BASECONv that computes add,(y x. S) for BASECONY as in Algorithm 7.
rumning sum S.
Proof. We will show this for d' and the general case follows will explain at the end. We now specify the two layers that we use
z' = BASECONV (v, Nx1 h b), N = b. о (h' y) BASECONV (z Nx1 h b2, b2) bj ④ (h y b9)
where we will specify the kernels as we go along. Let start by defining the kernel and the bias for the first layer as
E0
0
ep
1
h
0
b
0t
0
0
Let us first compute ht y follows:
h 1²(X) v(X) = +1) (S(X) X + x(X))
=S(X) X (S++)) x x(X).
We then have
0"
0
1
S +
S
2, b @ (h+ 4)
0
S
0
0t
0
0
38

---

Published as a conference paper at ICLR 2024
Resetting for Next Phase. We the first vector to be 1" in order to start adding the next vector. We thus use the kernel and the biases h? bj. bi defined
the next layer to reset for the next phase. Here, need
Ep
0
1
0
h
0
bj=
0"
0"
Explicitly, for the second layer, we compute the result of the convolution in terms of polynomials as follows:
h (x) 2 (X) = 1. (S + xx) X = (S + (X) x".
Thus, the output for the second layer given by
0"
1
1
1
1"
S
0
1
S +
S +
blo +3)
0
0
0t
0
0
0"
0%
0
0
0%
0
0
Therefore, have used two BASECONV layers add the running S and reset for the next phase. Here, we note that the only operations we perform and are convolutions and Hadamard product and they generalize in the obvious way d >
Next, we show that five layer BASECONv model can perform gated convolution on windows of the input (without changing the rest of the input).
Proposition H.10 (The Remembering Primitive). For any I E Ruxd E I Rax contained in some y E IR Nxd for some n + m + 8 + t S so that for h E Ruxd and p E R ((++a)xd with wth R (n+ and u* h IR. (m(+t))dd there exists (N, d, d) -BASECONY that computes remember(() (y: x, v. h p) for BASECONV as in Algorithm 7.
Proof. We will again show this for d' and the general should follow. We now specify the first two layers that we use
= BASECONV (y. Nx1 h bj. 0 Nxd) = bj o (h * +gy
z² = BASECONV (z 0 x1 h²2 bZ 0 xxd
D o (h3 +
The kernel h and the bias bj for the first layer are then given by
h
p
0m
m
e
0
h
0
b
0
0
1 m
0
0
39

---

Published as a conference paper at ICLR 2024
where recall that h E R(++s)xd and v h R (m+
We now want to first specify the result of applying the first kernel:
(h' y) coeff ((h(X) + x (v(X) x + xx()))
coefl (h (XX) X + h x(X) o(X) 2 2
+ xx) x
We then have
P
h x
p 3 (h * x:)
0 g
h v
0
0
0"
21 = bj (h' v)
0*
0*
0
T22
0
0
0
We now descrihe the second kernel h and the bias matrix bf as follows:
eg
0 +
0
0 --
e
0
h
0
6
1 n2+
0
122
0
0
This yields the following convolution computation:
h ② z = coell ((xm + +) (v(X). 2 + (pp 2 (h * x))) (X)))
(X)) x3n++6+2m+22 v(X). x2n+2o-m+t
+ (p (h * 2:)) (X) xmtntote + (p @ (h * x:) (X))
Thus we have
0"
p 3
(h
0
m=+
0 m+
0th
0
0"
z² = by. o (h2 x!)
nー
p +
(h *
(h x)
m
v
0
0
40

---

Published as a conference paper at ICLR 2024
We now shift this up by 2n+ m-+t entries using the primitive operation defined in Proposition H.8 that costs three additional layers so that we end up with
о h u
Again, we note that the only operations we perform and are convolutions and Hadamard product and they generalize the obvious way to 1.
Finally, we show that these primitives may be composed by 'stacking' models with matching inner dimension (N', dd).
Lemma H.11. For f.g: Nxd + Rxd that have (N, LN,L1,d, d,,, N',',d)) d') and (N, L2, d, N', d') BASECONY models then their composition f o g has an (N, L2, N' d') BASECONY model which can be computed by performing their models in succession, or 'stacking'.
Proof. This result follows from noting that for any f(u) which requires L1 layers compute and that we can compute f g(u) = ((f()) using the BASECONY model with L2 layers, yielding L1 + L2 layers in total.
H.3.1 BASECONV-HYENA EQUIVALENCE
We show that the equivalence between BASECONY and Hyena by showing that each layer can simulate the other's computation using a constant number of layers. Proposition H.12. For any input E RNxd and (N, L, d, N' dy-Hyena such that ZHyema Hyena(w) u) with set of filters he and linear projections p' as per Definition H.3 for E [L], there exists a (N, 5L, d, N' + N, BRAseCoNvmoder such that Hyenna = BASECONV(u)..
Similarly, for any input UBASECONY E R. Nxx and (N, L, d, N', Cly-Coyote such that #BAsECONY = BASECONV(LL) with set of filters h' for & [L], there exists series of Hyena layers such that have
Hyena (Hyena (... Hyena(
ZBASECONV.
layers
Proof. For the input UHyena R Nxd the output of the éth layer zyynn E IR N' xd' for Hyena is given by (see Algorithm 3)
(h' 2
Hyena PHyena
where Pilyena = Linear ( Hynnn)) E IR N'd. Now, using the original input WHenna E I Nxd to Hyena, we define the following input for BASECoNvusing one layer:
Uh Hyena
-Nxd
UBASECONV
UHyena
Then, we simply use the primitive for BASECONV. Consequently, this allows us to "remember" the input U in
WWW (BBASCCONV -N
UHyena Uyenen hftyenas Pityena)
the output of the previous BASECoNvlayer BASECoNN We then use this retrieve Pityna lincar( (incar(uHyeaaa) with the projection used for BASECONV given by
2-1
1 Nxd
2 PBAsECoNy = Linear(z f-1 (x
=
Pflyena
41

---

Published as a conference paper at ICLR 2024
Overall, the output of the éth layer for BASECONV is given by
PHyena ① Hyens
WHyenn)
Hyena
O
(M)
BASECONY
Hyenn
Uyeena
Hence, we can reproduce the output of the 4th layer of Hyena using five layers of BASECONY after augmenting the input and using the remembering primitive (Proposition H.10) with internal dimension N' + N. Now, for the imput UBASECONY Nxd the output the 9th layer for BASECONV given by
BASECONY = Linear BASECONv) O conv (h', 2-1 =BBSECONN)
Here, we show inductively that simply using E-many Hyena models recursively simulates of BBASECONV- For = 1, have
Hyyna((BBASCCON) (BASECONV, h.) Linear (UBASECONV) LinearruBBsECom/
(h' UBBSECONN =BASECONY
now assume that (E - m)-y recursive Hyena models produce "BASEConv- For ?th layer, we then have
(f-1)
Hyena (Hyena (... ((rena(upAssccow..)))
Hyena
(1-1) "BASECONV)
linear (BASECONV) L "BASECONV'
(t-1) conv (h', BBstcow))
H.4 LINEAR ARITIMETIC CIRCUITS
In this section show the relation between linear arithmetic circuits and BASECONV. We recall few definitions from (Dao et al., 2020).
Definition H.13 (Linear Arithmetic Circuit (Bürgisser al., 1996)). An arithmetic circuit is called a linear arithmetic circuit every multiplication has a fixed constant from IF as at least one of its two inputs. In other words, all gates in the circuit are lincar functions of their inputs (i.e. of the form az + by for fixed constants a, E IF).
only uses addition, subtraction and scalar multiplication. Further,
Definition H.14 (Butterfly Matrices (Dao et al., 2020)). A butterfly factor of size k: 2 2 (denoted
DI
D2 D,
as B is matrix the form BN =
where each D, is x diagonal matrix. We
D3
restrict i to be power of 2.
butterfly factor matrix of size n with block size (denoted as B(a)) ) is a block diagonal matrix of (possibly different) butterfly factors of size k:
B (m) = diag BBah .Bll.....,
[Bblg)
Finally, butterfly matrix of size n (denoted as B(n) is matrix that can be expressed as a product of
butterfly factor matrices: B (m) B B ((m) ... ", Equivalently, may define B (")) recursively
as matrix that can be expressed in the following form:
[
(n) = B(n)
(Note that [B()] and [ [B()] may be different.)
42

---

Published as a conference paper at ICLR 2024
From Definition H.14, d, d , E R." such that
observe that size butterfly factor is comprised of three vectors
d = (diag (D1), diag (D:)).
d 0 2 diag -1 (Dz) and
d
diag -1 (D3) 0
where diag (D): R + R" is the mapping from diagonal matrices to the vector of its diagonal entries. Let us define D2. D3 E R=*0 as diag (d) diag (d+), and diag (d respectively. Then we note that
D1 0
a
0 D2
0
D
D2S²
S D3
(S)
D4
000
0
D3
where S E EF is shift matrix for [n/2]. This gives us the following proposition:
Proposition H.15. For any powers 2f n = k: 2 2, any butterfly factor matrix ((()) is equivalent
B = S D3 D DzS +D1
where D3. D2, D are defined as in equation 5.
We Proposition H.15 show that butterfly matrices can easily be computed by BASECONV. Lemma H.16. For any n. d 2 2, k 2 1, and arbitrary vector a E Rdd
there exists (N, L, d, N' d') -BASECONV that represent B (nd) O(N), L (L), and d' = O(d), and
k
with N n, N'
(2) there exists a , L, d, N', -oddd'--oaal [00]BBBSCCOV that can represent ((nd), x with N = n1, N' O(N), L = =Olog nd), and d' O(d).
Proof.
(1) Given æ E Rad construct u E Rxd where æ is the row-major form of u. We show that BASECONV can compute Bnd & column by column.
& Let A = Dz C D'SS2, and = D2 for D,, S E Rddxnd for s as defined in Proposition H.15. We take dj 1 DD, d2 1 nd C2, d3 = A, which extracts the diagonal entries of D,. With this we construct D E R nxd where d, is the row major form of D. This implies that
D D o u.
Then we can decompose Bnd a into
BRa = O
By Lemma H.7, each Hadamard product A B. u, Co can be trivially be per- formed with a single layer BASECONV model. Let each of these model outputs be denoted
y/, 4/, Us. respectively. Finally all that remains to compute the f + V22 + ys. We achieve this using layers of add primitives!!
add, 91, 0)
add, (y² 3/2 31)
add, (y² ys y1 + y2),
where using by Proposition H.9 and Lemma H. this requires six more layers, and we get
y 31 3/2 y3 Brdw.
Then we can construct the (N, L, d, N'. d') BASECONV as desired with O(1) layers.
(2) From Definition H. (nd) B nd
(nd)
B (nd) B, (nd) From (1), BASECONV compute
any butterfly matrix by simulating the log(nd) butterfly factor matrices which comprise B (nd), With Lemma H.11, this creates BASECONV with 5 log(nd) = O(logndd)
layers. Lemma H.11
Recall that add, (y S) adds the suhvector for the input y.
43

---

Published as a conference paper at ICLR 2024
Butterfly matrices comprise the kaleidoscope hierarchy, which define below:
Definition H.17 (The Kaleidoscope Hierarchy (Dao al., 2020)).
Define as the set of all matrices that can be expressed in the form B(n) (for some n).
Define (BB*) as the set of matrices M of the form M M,M, for some M1. M2 E B.
Define (BB*) as the set of matrices M that can be expressed as M = M M,MM,
with each M; E (BB") (1 Σ i S w). (The notation represents width.)
Define (BB")) )e the of
matrices M that be expressed M SES for 0]] (i.e. M is the
some X matrix E E (BBB) where S E FRXen [I, 0
upper-left corner of E). (The notation e represents expansion relative to n.) simulate any kaleidoscope matrix.
We similarly show how BASECONV
Lemma H.18. Given n, d 2 2, > foror any nd x nd matrix M E (BB" )a and x E Rnd there exists a (N, L, d, N', dd' BASECONY that can represent M with N n, O(w log(end)),
N' = en, and d' = d.
Proof. By Definition H.17, M can be decomposed with respect size end end matrix
E
E E2 ......... Ew-
Further, any E, (BB") can be expressed product of 2 log end butterfly factor matrices. Then by Lemma H.16 and Lemma H. we can compute E,x' in by stacking 2 log end (n,, d, L, en, d) BASECONV models each with
(1). Because E has width w. Lemma H.11 implies that composing with each E; for 1 S i S w constructs final model with O(w log(end)) layers.
L
Finally, the kaleidoscope hierarchy is related linear arithmetic circuits via the following result. We note that in (Dao et al., 2020) it is assumed that 8, yet inspection of the proof yields the following stronger result: Theorem H.19 (Dao et al., 2020)). Let M be an x matrix such that multiplication of M times an arbitrary vector u can be be represented as (n, s. A, w--linear arithmetic circuit C. Then,
M E (BB") (BB")) ((A)
We combine Theorem 1919 and Lemma H.18 to show that BASECONV can compute any linear arithmetic circuit with polylogarithmic factors in A.
Corollary H.20. For any (nd, s, A, w)-linear arithmetic circuit C that can be represented by matrix M E Rddxn multiplied by vector æ E Rnd, there exists an equivalent (n, A', d, w, d) BASECONV with A' = O(A (o)()) such that Mæ BASECONY (u, h) where x is the row major form of u E Raxd
H.5 GENERAL ARITHMETIC CIRCUITS
We are now ready to prove the result that yields the equivalency between arithmetic circuits and BASECONV.
Theorem H.21. For any (nd, A, -ritthmmet circuit C, there exists an equivalent (N, d, N', d') ( BASECONN with N n, A' O(A log w). N' O(w), d' that simulates C.
For the reader's convenience, we begin with proof sketch and then provide the details afterwards.
Proof Sketch of Theorem H.21. Let us layer C so that each layer C2 for & E [Tc]] either only has linear gates or multiplication gates, where the number of such layers I.c = ( The composition of all C2 layers results in C. We use 2+1 IR." to denote the output of the é-th layer C₁ which feeds as the input to the (E + 1)-th layer C++1- Here, we note that if we can simulate each C2 with BASECONV, then can simulate the entire layered circuit C due to Lemma H.11.
Now, the layer Clin is a linear layer (with only addition gates), then it can be represented by a matrix M E IR" " multiplied by x E R" (We can append with Os if necessary so that the input from the previous gates can be written as wollnggth vector). Thus, we can apply Corollary H.20 to simulate Clin with an equivalent (n, log w., d. Olog w), d) - BASECONV model.
Next, if Chmuls instead consists of only the multiplication gates. Then, we note here that the output z' may not exactly equal the input to Crmult. Nevertheless, we can apply a O(w) sparse linear map R E Re so that Rz yields vectors v², v², and v², where V1 constitutes the "first" input to all the multiplication gates and U2 constitutes all the "second" inputs while &a consists of all entries needed as inputs in the subsequent layers. That is, for the ith gate in Cr we compute o! u,2 This implies that for all the gates in Crmilt, we can simply compute i @ U2. To this end, the remember primitive with constant number of layers from Proposition H.10 to define
44

---

Published as a conference paper at ICLR 2024
(n, (log w), d, w.d) BASECONV model that remembers va while performing the Hadamard product u with v². Overall,
then collect all the resulting BASECONV layers and compose them in Lemma H.11 to simulate C. Overall, the number of layers used given by O log w) while the internal dimension remains fixed w.
L
Using the outline in the proof sketch, we now delve into a detailed analysis of the arithmetic circuit C an equivalent BAsECoNvmodel.
Proof of Theorem H.21. Let u & R be the input to the arithmetic circuit C with depth A and width w. We begin by rearranging the circuit into layers of addition and multiplication gates. That is, each layer C has either all addition gates or all multiplication gates. This allows us to readily apply the results from Appendix H.4. Note here that we can assert that the number of such layers Le (()). Moreover, we construe the input to each such circuit as a vector of length by appending with extra Os if necessary so that the composition of the layers results in circuit equivalent C. See Fig. for example of such decomposition
X
CE
mult
+
ct+1
cl+2
mult
Figure 9: An example decomposition of an arithmetic circuit as layers only addition or multipli- cation gates.
Let 2 E R denote the input each layer C. for E E [Lc] with x1 (u, 0w-2d). It important to note here that we may not have f+1 Cz(z²) since the inputs for gates at the (e ( 1)-th layer may come from any of the previous layers. We now handle the case of layers with addition gates, multiplication gates, and the intermediate stage separately.
Addition Gates. Let Clin denote an arbitrary linear layer which only contains the addition gates that takes in z as input. We know that there exists a matrix M E IR "x such that we have
cliu M
Here, note here that we may need entries from the vector z in subsequent layers. Let be the total number of such entries, then for each index for such entries [w]. have the corresponding index i, E [s]. We then append the ith standard row vector into the matrix M to get the matrix M' E R Rx (wx so that we have
(M's9 (M²) Di [1]]]]
if Hj<u w/
wi=
Here, note that we must have w' + < w is the width of the circuit. If needed, we then append the matrix M' with zero rows 0 so that we get the matrix M E "xx,, Note here that we have thus preserved all entries needed in subsequent layers and incorporated the output of the 4th layer with the output M z' = e+1 serving as the input of the (E + 1) layer. Applying Corollary H.20 then yields equivalent (n, Oloo w), w. d) BASECONV model.
45

---

Published as a conference paper at ICLR 2024
Multiplication Gates. Next, deal with the layer Cunut of multiplication gates by collecting the two imputs each multiplication gates as vf, vz. Note that for the imput x R from the previous layer, we will again have entries that need to be used the subsequent layers that we will denote vg. We thus need permute the entries of 2f to get the vector [ef 2 eg] so that we can remember os while taking the Hadamard product of @ V To this end, we can achieve this permutation of entries using (ssa-se linear matrix R E R=xx with which is equivalently represented by an (w, w, 1, w) -lineer arithmetic circuit that simply moves the appropriate wires from the previous layer. This can again be achieved by a equivalent (n, Olog w), d, w, d) BASECONY model. That is, R xf has the following form:
R z'2²= =
Next, we can now define (n, L, BASECONY which extracts of as the projection with the input to the remember primitive given by y = (vg. ij 0). We then specify the remember primitive as remember (y vn, uyg eo, vf) which computes the following
vuf O (en * U.!
241
v O w
rg
Culte (2²) E R
Using Proposition H010, we know that this requires (n, Ollog w), d, w, d) BASECONY which remembers the entries that will serve as inputs subsequent layers while performing the Hadamard product v @ 2
Overall, we can now stack all the resulting BASECONV layers and compose them as Lemma H. 11 to simulate C. Overall, the number of layers blows up by O(A log w) while the internal dimension remains fixed at w.
H.6 THE RETNET REDUCTION AND SPACE COMPLEXITY FOR AR
We are now in a position to show that there exists an equivalent BASECONY model that can simulate Retnet layer from Algorithm 6. Recall that for an input sequence E R xd we project the input using weight matrices W A, Wc, Wy E Raxd A, C, V NR and compute the following recurrence (see line 5 in Algorithm 6):
--2-1 + (A] i[n,:],
where nd E Rddd We unroll this recurrence to express each state z" in terms of with coefficients given by y and the weight matrices W., Wc, Wy follows:
=0 ((uW a) [0 (uWv) [0,:]
= (u[0,:]W A) T (u[0.:) :) Wv
WT [)) (u Wy
wT (u [. (0.) Wv.
Fach of R has exactly one non-zero entry.
46

---

Published as a conference paper at ICLR 2024
Similarly, we have
z' 72 + W (u [1 []](1,)) Wv
=T (u []](0,)) Wv) + w (u T []][;;)))) Wv
= =WI (7 (u ()) + u I. []],,)) Wv
We can then generalize this to n E [N] to get
n
z" W Σ In ".. Wv.
The above helps us infer that z" can expressed as a polynomial [u[0, 0], 200, 1] ............ [N
1, d - 1]}. For each polynomial, there exists arithmetic circuit that computes the polynomial in natural way, whence we can apply Theorem H.21 to get an equivalent BASECONY model.
Corollary H.22. For RetNet model with (dd²) parameters and N layers, there exists an equiv- alent BASECONYmode! that uses (NN)) parameters and O(N log d) layers.
Proof. We will start by describing the arithmetic circuit that computes the states 2" Radd here. Indeed, each state requires exactly two alternative layers with only multiplication gates and only addition gates, respectively. First, note that for each N, we have
==>="1+- + (A (n, :)' V[n,:].
Since we can compute = in (1) layers and computing 7 -==1 needs exactly one layer with each entry from zn serving as an input to a multiplication gate along with % the depth of the circuit is ((N) with width d² We can thus apply Theorem H.21 to get an (O(N), O(N log d), d2 (dd), ((d))))BBBBECCOO N.
model that can compute z" for each E
O
The Space Complexity of AR In the context of the Retnet Model and associative recall, it is worth exploring the space complexity of the associative recall (AR) problem:
The AR problem takes key-value pairs {kg, v, v along with query q appended
at the end as input and the goal is to output if k; for some E [0. N 1V-17 -1].
Indeed, we will first provide a power-bound for any model that purports solve AR. To this end, we will require a randomized communication complexity lower bound result for the index problem:
The index problem has two agents, Alice and Bob, where Alice has a string x E {0, 1)"h and Bob has an index i E [v). and the goal for the players is to output the i-th entry CG, Moreover, we also require the communication to one-way: only Alice is allowed to send single message to Bob and Bob needs to output the answer.
We will make use of the following lower-bound result. Theorem H.23 (JJayram et al., 2008)). The one-way randomized communication complexity of the index problem for sending an n-length bit string is M(n)..
We now use Theorem H.23 to provide a lower bound on the number of bits required by the Retnet model solve AR.
Corollary H.24. The RetNet model requires Q(N) to solve AR d S VN.
Proof. Consider an instance (x, i) of the index problem with æ E {0, 1J. We now describe the corresponding instance of the AR problem:
{ x, 1 i.
(8)
Next, consider the following one-way protocol for solving the index problem using the RetNet model. Alice with their access of x E (0, 1 generate an input for AR (without the query) as in equation 8. Alice then runs the RetNet model on {i, J. and sends the memory content of running the RetNet model Bob. This should include the output 2 as we can reasonably assume that both have access weight matrices WA, Wc, Wy and the scalar 7y Since we assume
The randomized communication complexity of function f is defined as min, |T]] where v ranges over all randomized protocols that solve f with probability of success least 2/3.
47

---

Published as a conference paper at ICLR 2024
that this model solves AR, the output [ut[N, z, should contain the associated value of i. Here, Bob can compute Out N, by using the memory content sent by Alice along with the term y N=1
2; Out N. = =] C|N.:] (yz + (A[N, V(NN
That is, the total number of bits that are communicated in this protocol is (cd)). ). For d S VN, have shown that a one-way communication protocol exists for solving the index problem exists that uses o(N) communication complexity. This contradicts Theorem H.23 and hence, we conclude that the ReNet model solving AR also needs (2(N) bits.
H.7 THE SILTIPLE-QUERY ASSOCIATIVE RECALL PROBLEM
H.7.1 INTRODUCTION
In this section, we consider a general version of the associative recall problem (Ba et al., 2016).
Setup. Next, we will redefine the multiple-query associative recall problem (MQAR) from Defini- tion 3.1 to a slightly more general problem:
Suppose we are given an input sequence u(o...N N - 1] { (Ku, uu. qu). ..... ( } with each k, 21 qi E C is a token drawn from vocabulary of size |C|. Our goal is then to check, for each < i < 1, whether there exists 0 output v) Here, we note that it suffices to have d z log(c) so that k; v; qi is embedded in (0, 1}d. However, will specify the specific embedding being used for the results below. Here, construe the tokens k; qi and v; to be the keys, the queries, and the associated values. Indeed, might be helpful think of the input u as streaming sequence of key-value pairs for which we sequentially employ standard associative recall for every key that shows up in the sequence so far. To see that the above generalizes Definition 3.1, considers sequence of length N ul0- 1 f#o...... ................................ xN where each E C. The goal of Definition 3.1 is then to check, for each isic N 1, whether there exists 0 S j < such that æ; = 3; and if so, output + and continue otherwise. We can reduce this problem to the above general formulation by taking the following sequence of tuples as the input (<x,, Ii+1, x.)).
i such that q: = ky, and if so,
Remark H.25. As noted above, this version is more general than Definition 3.1. Thus, the results proven in the sequel, which are proven for the above general MQAR can be ported (with constant blowup in parameters) to get the results corresponding results for Definition 3.1.
H.7.2 MQAR SOLUTION VIA ATTENTION
Before describing how BASECONV solves the multiple query associative recall problem, we discuss how Attention solves trivially using pairwise inner-products. To this end, we will specify how the input is presented to attention.
Remark H.26. We note that the input for the multiple-query associative recall problem u e {0, Nd has designated indices for the keys, queries, and values in the sequence. We gather these indices below:
E (0. ..... N- --] mod 3},
mod 3}, @ = E E (0, NO 1 i= = mod 3}..
V = E (0. ................................ - i
(9)
The input u E Nxd Attention for d is then given by
[k. : 0 09]
if E KC if V
u : =
00 0. 0
[ 0 q.]
ifie
Here, each ki, Vi q: is embedded as a one-hot encoding in (0, 135. Without softmax, the output for attention layer O E IR is given by
O (QKT) V,
(10)
where Q K, V E Nxd defined as W Q. uW K, uWy for u Nxd Instead of position embeddings, ALiBi, popular technique that biases the attention scores QK with lower- triangular Toeplitz matrix B E N (Press et al., 2021). The values in this matrix are controlled by fixed hyperparameter they do not count towards the number of parameters in the model.
48

---

Published as a conference paper at ICLR 2024
Algorithm WLLBB-witthutt
(u(0.. N 1).0, prey 00... ... N 1].
Input: Input sequence u(0 u N {Kk vi, qi) q:) J== with each k; Vi. {0, 1 133e previous layer's s output O prev (.... N 1] E IR" N x3c and linear bias B E IR N
Add Heurr <- + prev as an input to this layer.
2: K, Q, V &- Ueurr Wo: Q, Ueur W K, Ucure Wy.
3: + (QKT B)
4: return O the output of this layer.
Proposition H.27 Given an imput u E {0, 1] (encoded as in Remark H.26) where d = 3c, Attention with linear biases (even without using soft-max) solves MQAR for using O(e²) param- eters, ONN + 22c time complexity and (1) layers.
Proof. We use two layers attention. We will start by specifying the projection matrices for the first layer W Wk. W Rdxd as:
0
0
0 )
W = W W
0
Lexe
0
0
Above, W1 meant isolate the vs embeddings. For the first layer, we then have Q K 0. and
110- u.:0 : if od
v v' u li, W
otherwise
where IC, 2, V are defined as in equation 9. The output for the first layer is given by the following:
0 = ((QKT + +B') v) [t,
= BB Vii.:
[0 09 if otherwise
0d
The QKT is ignored (= 0) and we isolate the shifted sequence by setting the bias matrix B appro- priately. In particular, the last equality follows from the fact that B is up-shift matrix that shifts each row of V by 1. We apply the residual connection at the end of the standard Transformer block insert the k; adjacent to the vs. For the second layer, the input u² e Nxd is given by
[k; U. 0
if K
[0 : 0. 0
if
0 0
q.
if
Further, we take the following projection matrices:
Lexc
0
0
0
0
0
0
W
K
=
0
W =
0
w
=
L
0
0 0/
0
Ixxe 0
0
0
0
We then have the following matrices as input to attention after applying projection:
[
02c
if E Q,
Q --
0d
otherwise"
[As 02e] d
i
:= |
(1I)
otherwise
[0. dd
02e
ifie
=
otherwise
49

---

Published as a conference paper at ICLR 2024
Here, note that the values have been shifted to the corresponding key position. Next, compute the term in the parenthesis for the second layer as
(Q² [i, (Q²K²T) [.,]] = (Q) K 2T [:,]])
(i, K2 [i:1))
(qi, kj) if E Q.j E
otherwise
1 if Q.j K, k, ex for some k
0 otherwise.
Finally, we compute the output as follows:
0 [ (QQ²K?T 2T + Bd) V2) [..] ((QRK2 + 0)) v²) [
(Q'K²T) [i, V2
(Q* [i,5] v.
(Q²K2T)
[i.5] [0, : 021
JEK
[0 020 0d
j kk
otherwise
where we use the fact that for each index j K. the matrix 2 contains the associated value from equation 11. Thus, for each query q., we solve the associated value problem yielding a match for the jth key.
In total, we only need m-many parameters to perform these multiplications and the linear bias in the first layer hyperparameter that is static and unlearned; the time complexity comes from the multiplication of QK in O(N2c), and projections in O(Ne).. Finally, we only need ((1) layers for this solution.
In the sequel, we develop parallel algorithm to solve the multiple-query associative recall problem with O(Nd log" N) work complexity and O(d log" N) time. We then convert the algorithm into a BASECONV model via the route of arithmetic circuits, which then solves the multiple-query associative recall problem with (1) layers and ((Nd) parameters.
H.7.3 INITIAL ATTEMPT: A SEQUENTIAL ALGORITHM
We will first discuss the algorithm that simply uses associative arra to solve the multiple-quer associative recall problem. Specifically, we want to use a data structure that allows for logarithmic insertion and membership query. Here, do not specify choice but data structures including self-balancing binary search trees which allow for Olog d) insert and find operations for d-bit entries should be sufficient.
Algorithm Sequential. MM--MRQQ-------- AR (w) 10... 1]
Input: Input sequence u(0 N - A {(k,, vi, q:) 1=0 with each k,, Ve, qi (0. 1]d. 1: Initialize an associative array with insert and f find and an output array out < [.
2: for E {O...., (k,, v,) + find qi)
do
3:
Query for qi in the data structure.
4: if' k2 is not null then
5: Add vy to out.
6: insert(k,, v.)
> Add the key-value pair to the data structure.
7: return out.
Proposition H.28. Algorithm solves the multiple-ques associative recall problem (MQAR) in OdN log N) time for input sequence {0, 1) Nxd"
50

---

Published as a conference paper at ICLR 2024
Proof. For any E {O..... - 1/, know that both insertion and lookup operations take Ollog(i) d) time. Overall, the runtime of the algorithm is
Σ Ollog(i) d) Ollog(N)) d) O(N log N d).
OoV
H.7.4 ALGORITHM VIA PARALLEL BINARY SEARCH
Our plan is to convert the algorithm for solving the multiple-query associative recall problem in the RAM model into an arithmetic circuit, which by Theorem H.21 will lead to a BASECONV model that solves the multiple-query associative recall problem. With respect Algorithm 9, it may be the case that the arithmetic circuit has a large number of layers (NN). Unfortunately, this would imply that the resulting BASECONY model may have near quadratic complexity in N. Instead, we now initiate our effort into designing BASECONV model with both small enough number of parameters and number of layers. Here, we will first subdivide the problem using dyadic intervals into ( subproblems and reduce each such subproblem into a multiple search problem (Akl and Meijer, 1990). To this end, briefly introduce the multiple search problem below.
Given two array of numbers - ag S ... az-1 and B - (by ... S bm-1) with S m, for each aj E A, the goal is to find the smallest element in B that is larger than or equal to dj. The multiple search problem is solved by parallel binary search (pbs) algorithm in (Akl and Meijer, 1990) with work complexity On log m) and time Ollog n log m). Specifically, for sorted arrays A(0... 1] and B(0
1]. pbs constructs the array C'0 ..... 1] defined as
Cli] [minosi<mn Lil A[] < BDD if A S B(m 1]
(12)
otherwise.
The algorithm itself runs in exclusive-read exclusive-write (EREW) PRAM model no two proces- sors are allowed to read from or write into the same memory location at the same time. We now augment the algorithm copied from (Akl and Meijer, 1990) for our purposes below.
Algorithm 10 pbs- veluevllevessee (g(s sils...t]. t), A(x 1/. n,m) m)
Input: sorted arrays q(s-..-)) = (q:)/ k(x-. == {G, k 1: Initialize n processors denoted Po, PI....,PP Pa-1 executed by P..}
(Sequential steps are assumed to be
2: Initialize the output array C = [m]]..
3: if S then
4: mid [ss t)/2]
5: if q[mid]] < [1] then
6: for i : to mid in parallel do
7:
Cii] +
b Step executed in parallel by P.
8:
pbs key- (gimic 11...) f), klz.
9:
else
10:
if glmid] k[y][]] then
11:
for 1 == mid to in parallel do Cii +
12:
Step executed in parallel by P.
13:
pbs -ey--alues values (q(s -......mm mid - 1], k|r (f(x...y)
14:
else
C[mid] is determined using sequential binary search
15:
= min qimid] s {[]][1]]
16:
C[mid] +
17:
do steps 18 and 19 in parallel
18:
pba-key-val ues "q(s...mis 1}, k..... 1)) pbs-key-ue laae--valles lue (qimid 1...f]. a(2...y))
19:
20:
return C.
Let Σ be the set {0. 1} and denote the set of binary strings of size n as Σn, We define prefix( for n-bit strings the of all initial substrings of as Σn which includes the empty string and
51

---

Published as a conference paper at ICLR 2024
itself. Next, let dec 13 N be the decimal representation of an n-bit string with x x00] denoting the least significant bit. We also use sort(A) as a procedure that sorts an array A. Finally, wlog, assume that N is a power 2. We now ready to present parallel algorithm that solves the multiple-query associative recall problem below.
Algorithm 11 Paral Parallel-Moas (u
11)
N Input: Input sequence u/0--. N 1] - {kk, v, 9:) 1=0 with each k, vs qi E {0, 13
1=0
1: Initialize
3 log ( processors denoted Po....: : P. lox -)11-
2: Initialize the index and output array idx, val
3: for k : {0....,. log 2 1} do
4: for x {xc Sl 3) $) x log ) k 1J = 0} do
5: {All the steps below are executed in parallel by iite. " 1abe E prefix(y)}.
1hth + t )log( 5
6:
7:
kkk I
sorted ppormtted f sort ((Adeo()) Herre)
Tke permuted ((j,dc(()) ko sorteadli [o] (
8:
x0 æo1og (N) k
61 +) J <- fy Slog( x E pr))
9:
10:
Isorted: Jher
ce
pprruted f- sort
(ort(((gg (qdoo()) (Ee
J Jkn permuted = {(dec(j), k km
Adec(j)
11:
Ck < pbs -ey Lkey-val va values (goorted, kkk
socted 2k 2*)
12:
for j E J do
13:
if Ck dec[jj] + 2 then
14:
g + Cs (peeprratted ( (dec(j)))
15:
if ko + 2 then
D cf. equation 12
Add hk permuted (c)²) to idx[dec(j)].
17: for i E (1,..., 1) do
18: (Executed in parallel by p..}
if idx then
20: Add Uz+1 to val
21:
return val.
Remark H.29. In lines 7 and 10, we keep track of the sorted permutation of the indices of keys and queries, respectively. This helps us in the retrieval of the index of the matching key as in line 16.
H.7.5 CORRECTNESS AND COMPLEXITY
Proposition H.30. Algorithm 11 solves the mudtiple-query associative recall problem with work complexity O(Nd log" N) and time O(d Pldd N).
Proof. The correctness of pbs implies the correctness of Algorithm 11 if we can show that, for each <
N we check for a match among the keys (kz))æ{{{))- To this end, for N. let the set of all iterator indices associated with an index be defined as each S <
3
K; 2)) } with J as noted in line 9. Then, we define the corresponding set for keys as I, & UIkaeeK, IE with Ie' defined as in line 6. That is, for all the calls to pbs -key-aalees that is part of (given by K2) where the algorithm checks for match among the keys in I,, it then suffices to show that Iz [ 1].
Here, first note that if some index E 17th C I4 for some 2o ) then, by definition, E prefix(bin(j)). Here, let x == =l } - where we set the (log ( 1--k k))th index of
to be 1. Consequently, as we have E J for the same & and æ as in 1 (cf. line 6). we must have x' E prrefixbbi(ii) Thus we get j < i, whence we can claim that I₁ S 1 -
For the other direction, for any let b denote the position of the most significant bit in bin(i) which differs from bin(j) for any E [i - 1]. Then, there must exist a binary string that is in the prefix set of both bin(i) and bi)( That is, there exists x prefix(bin(i)) O prefix(bin(j)) with E..
52

---

Published as a conference paper at ICLR 2024
Thus, then must have bin(j) E _ lug( 3 and bin(t)
J log( x
with as the corresponding
witness. Hence, we have l 1] - Iz. Overall, we have shown that I, = [i I]. Since this holds for all 1 S
N
we can conclude that
Algorithm 11 solves the multiple-query associative recall problem.
2k=0 10 2 many times. We note that sorting values each of size d can be done with work complexity n log n-gndd dd We note that, at each instance, we are sorting sorting 2 values. Meanwhile, remembering 2 sorted permutation of indices can be done in linear time using arrays. Moreover, each call to pbs-key-value, has
Next, it is easy to that execute lines 6 16
n m 2 which has work complexity n log m. Finally, we know that the work complexity of lines 18 to 20 is (NN). Thus, the overall work complexity of Algorithm 11 is
log( + --1
N 3
log N/3-1
d
O2 2
log 2*h d. O(N).
O(k) O(Nd log N). (13)
We will now analyze the depth of Algorithm 11. We know that the depth of computation for Algo- rithm 10 Ollog n log m) for input sizes. Moreover, we have O(1) depth for the computation in 18 to 20 as each entry in idx can have at most one entry. Since the nested for loops iterating over ks and the associated æs runs in parallel, the depth of Algorithm 11 dominated by the largest depth among all calls to pbs -----ey-aalllleeee values and to sort. The largest such call to pbs pbs-key- -ey--eyyllllleeee-ee values of size n = m = 2 ( 5 networks (Definition H.34), we know that the largest depth is for sorting 6
N/6 which yields a depth of loog N Moreover, using sorting 3 N values of size d given
by d Olog N) (Lemma H.35). Thus, can conclude that Algorithm 11 takes Old llog N), time where is the length of the input.
H.7.6 CONVERSION TO ARITHMETIC CIRCUIT
We will convert Algorithm 11 arithmetic circuits modularly. In particular, after writing an explicit circuit for Algorithm 10, we will uses this circuit as a black-box along with circuits for sorting networks.
Circuit for pbs-key-values. We will denote the corresponding arithmetic circuit for Algo- rithm 10 as pbs key-ey--- Keey-values as well with the input gates comprising each entry from g and k<(x...) and the i-th output gate yielding the value C'il as in Algorithm 10. Here, we first convert the comparisons for the if statements in Algorithm 10. To this end, we briefly introduce comparators.
Definition H.31 (Comparators). A comparator is device with inputs æ and y and outputs w' and y', that performs the following function:
min(x, y),
y' max(x; y).
Using comparators, use bit-wise XOR and AND to define the result of the comparisons in lines 5 and 10 as the following fixed variables:
& [g[mid] < kg).
9, = [q[mid] > ky),
(14)
2 = b sinnsearcc((kk)))
q(mid]).
This then allows to infer the index of the key arry for each recursive call to abs-key-values in lines 8, 13, 18, and 19 from Algorithm 10. Specifically, let 2, and 2 - denote the starting and ending indices for the keys inputs to the recursive calls in Algorithm 10) below:
pbs-key-values (qiid [mid +1- .....f) t), k(2 ... y)); -e):; (lines and 19)
(15)
pbs-key-values(g) mid - 1], k(x ...2, 2. 1]): (lines 13 and 18)
(16)
53

---

Published as a conference paper at ICLR 2024
Here, 24 and & can assume values dependent the results of the comparisons in lines 5 and 10. Specifically, we have
if q [midd] I k2 (line 8)
24 = x (1-)))(11 -2
if q [midd] E (Nez kyl (line 19).
otherwise
y +
if q (q(mid]] ky (line 13) q [mid] E (kzz kyy (line 18). otherwise
2, = (y (x) -2
2
Here, getting value 0 signifies that the branch is dead, and we do not execute the recursive call.
Finally, let the arryys 11 E] and [s mid 1] denote the outputs the recursive calls equation 15 and equation 16, respectively, We can then succinctly express the outputs for each index of the output array C' as
le x + 2. (Ci(t))
E [s ... mid - 1]
cli =
Oy (y + + (C2lil)
i E mid + -.. 1
(17)
lo x (1 (1- 9v) + 9y (2 1
i mid
We can thus state the circuit schematically in Fig. 10. Now, before accounting for the complexity of the circuit for Algorithm 10, we must first assert the complexity of the comparators that we use in Fig. 10.
plase key valuew (gle- w).
ky ... fiid
glmid]
bin sea Fon ((k.) Vime g/mid])
gimd
1+1
bbs-key-velae900 mid 1, A/s,
pac-kay-volucs-o mid 1 6 8 1
21
Cls) Ciuid 1]
C|mid]]
C' mid 1I...........C. C(t)
Figure 10: pbs-key-values (q]s...t], t), kjr... vl,n, m) as circuit with recursive calls and subprocedures as "black boxes."
Lemma H.32 (Cormen et (2022)). For binary strings x. y Ed of length d, there exists comparison network of size (d), width O(d), and depth (log d). We use this lemma to compute the complexity of the arhmetic circuit from Fig. 10.
54

---

Published as a conference paper at ICLR 2024
Proposition H.33. There exists an (O((n + m)d), O(nd (logm + log (log ((lg + log n log d)), (nd)) ( arithmetic circuit" equivalent to pbs-key-values (Algorithm 10) with inputs q and is of lengths and with d-bit entries.
Proof. The size of the circuit for pbs-key-values should equal the cork-complexity of Algo- rithm 10 but we also need to account for the comparison gates in equation 14. Further, the circuit for binary search also has size of (d n) instead of Old log work as in Algorithm 10. Us- ing Lemma H.32, along with the fact that the comparison gates and the binary search are used at most (log n) times, we deduce that we are adding O(nd log n)) size to the circuit in addition to the work complexity of the parallel algorithm. Thus, the overall size of the arithmetic circuit for -ees--eevvaaae
is Odn log m + nd log n.). Further, the depth of the circuit here is determined by the runtime of Algorithm 10 along with the depth of the comparison gates and binary search Olog n log d). and finally, at most n processors are used for the parallel algorithm in Algorithm 10, which yields the width of the circuit.
Now, we will call the circuit for name while the input gates contain the inputs of Algorithm 11. Indeed, we can directly "translate" Algorithm 11 an arithmetic circuit as the values for _ and 1 for each æ and k are predeter- mined from N. Thus, we start by placing the corresponding sorting networks which feeds into the pbs-key-values (q(s (qls...t), t]. k(zz.. y), n, m) circuit for Algorithm 10 in Fig. 10 so that the output values from the calls to pbs-key-values result in the checks as in line 16 of Algo- rithm 11. That is, we get outputs CR [ permmteed (dec(i)) from each call to pbs-key- lues.
Circuit for Parallel-MOAR,
with the same
We can then use a comparison gate to check if this value equals 2*, and if not, we have found a match C2 permmpeed (d)c for the query qi which results the output of the associated value "C permuted (dec exactly in Algorithm 11. That is, we first define the following variable as the output of the comparison gate:
Ar caee(s) == 11 11C [ (dec(i)) 4 2*
(18)
permuted (C& (her pa (dec()))) equals the index of the matching key k; corresponding to the query q., the ith output is then simply given by
Here, as SerSpermttad (dec(i)) * 2* implies that pke
dee(s)) perrmtedd [ (decm where the output implies that there does not exist a matching key.
Here, we also briefly introduce the the sorting networks that we use to sort the keys and queries: Definition H.34 (Informal). Sorting networks are circuits with gates and wires where the gates of the circuit are comparators (Definition H.31) connecting two wires. Each such circuit can perform sorting on a fixed number of values.
We can then show the circuit schematically as in Fig. 11.
We now dilineate the complexity of the circuit, starting with the complexity of the sorting networks. Lemma H.35 (Ajtai et al. (1983)). Let A be an array with d-bit entries of size Then, one can implement a sorting network to sort the array A with size Old n log n) and depth O(lo d log n). Proposition H.36. There exists an (N d, O(N log N), Olog d logg doN arithmetic circuit that solves the multiple-query associative recall problem.
Proof. We note here that for each kt, there are उ parallel problems of size 2$ for both the sorting networks and the pbs- key- values circuit. Using Lemma H.35, the cumulative size of these sorting networks is Odl- N llo2 N) (see equation 13) with overall depth O d log N). Similarly, the next layer again runs Eleg $ of which has size O(2*d(logg 2* + log 2*) (22 (Proposition H.33), Again, the cumulative size of this layer is given by O log .N) (see equation 13). Since we run each of these circuits in parallel, the depth of this layer is again Olog dl log"((N) while the width is O(N log N).
N 2 -any circuits for pbs key values each O(d 2 1og2*), depth Ollog" 2 log d) and width
N Finally, we perform I log 3 comparisons at the end of d-bit strings in parallel which results in size O(N log N d), depth Olog d) and width O(N log N d) (Lemma H.32). Therefore, the resulting
N
MRecall that (n. A, 20 arrithmeeii circuit is an n-variate circuit with size depth at most A, and width
w.
55

---

Published as a conference paper at ICLR 2024
Parallae liee G-neenaal Bateera-AAR(M AR( N 1))
(K1 oi q1).
(Ae 1 11 4
1'
sost (196e()))segg)
sest ((hddets)) heerz)
==
seet ((Qdee(s) beep))
2025 ((Erp)
log
pbs-key-values (sart( be sp. sort (kot(
pbo pec-xcy-values (aort (Adee(s) ye sort (kdec(t())))
k= lop
Co
i
pormuntd (d))
log N
hhe permated (ceet)))
dee(s)
Figure 11:
Parallel Paraalle-MMAR(((O-... 1)) as circuit that includes sorting networks and the circuit for pba-key-values subroutines.
arithmetic circuit has size O(d N log" + Nd log² N + log d) = O(Ndlog" N), depth (log d log N) and width ONN N).
H.7.7 THE RESULTING BASECONV MODEL.
As we have an arithemtic circuit for solving the multiple-query associative recall problem, we can now invoke Theorem H.21 to claim that there is corresponding BASECONV model that solves the multiple-quest query associative recall problem with Ö(N log parameters and O(1) layers.
Theorem H.37. There exists a Nd,, Õ(1), Õ(1), Ö(Ne), 0(1)) BASECONV solves the multiple- query associative recall problem.
Proof. Directly applying Theorem H.21 yields a BASECONV model with the number of layers Olog d log² N log log N) O(1) layers while the claim on the input and inner dimensions follow trivially.
H.8 ANTA-DEPENDENT CONVOLUTIONS
H.8.1 INTRODUCTION
In this section, we are again concerned with solving the multiple-query associative recall problem (Appendix H.7.1). However, in contrast Appendix H.7.4, which yields circuit that unchanged and works for all inputs, we instead take the viewpoint of adapting the sequence mixing weights (Section 2) with respect the particular sequence that the model receives as input. More specifically,
56

---

Published as a conference paper at ICLR 2024
take the distance between the tokens in the sequence measure for designing data-dependent convolutions.
Setup.
To formally setup the problem, as in our discussion of designing a parallel algorithm,
consider the following problem description of the multiple-query associative recall problem. Suppose
we
are given an input w(O...N N
1]
{ (koo, vo. qo).
-g-1)) with each k; vi qi C. Here,
each token is embedded using the standard one-hot encoding in {0, 1] (i.e. we assume c). 15 Our goal again check, for each <
1, whether
there exists 0 S j < such that qi = kj. and if so, output f/
Here, we define the interaction distance between the ith query qi and the matching key k2 as i j. We then also assume that number of distinct interaction distances is bounded by f. We can identify these distances using autocorrelation (Chatfield, 1995), which has an elegant un- derlying formulation. We will briefly introduce the relevant mathematical machinery in the context of elucidating the data-dependent model that we seek to develop in the sequel.
Autocorrelations We introduce autocorrelation convolutions. correlation of two vectors and is given by
Let ūi{ := ul-[], then the cross
The aulocorrelation of a vector u E R is the cross correlation of u with itself. Moreover, in terms of polynomials, have ū(X)= X u11/)). Thus, in analogy with our interpretation of convolution in terms of polynomial multiplication, we characterize the autocorrelation of vector R given by w E R as follows:
coeff (u(X) ū(X) mod X" 11
H.8.2 BASECONV WITH KERNELS GENERATED USING AUTO-CORRELATION
We are now ready to describe the model that solves the multiple-query associative recall problem using data-dependent kernels derived using auto-correlations.
The Imput-Dependee Kernels. There are several potential strategies to identify the best token- interaction distances for an input using a convolution. Here we will focus on an autocorrelation convolution for exposition. Autocorrelation allows us to identify the top / distinct shifts of the sequence that result in highly overlapping values (c.g. matches between query and keys our associative recall setting). We then construct convolution filters that perform each of these t shifts. That is, we define a function Top such that Top (u * u, f) returns a list of the top t shifts {se hejs)- We then use these top f distances {se hee(s) define the following two kernels:
hx(x)=
Σ (2
SE({]
(19)
hx(X)= = Σ x
EE({]
Here, we note that we only only have two kernels we will assume N' = N and the shift will be done in "parallel." Obviously, one can instead define distinct shift kernels but then there is cost of (t) in the number of layers.
Projections. We define the following projections K,Q, Ve E {0, 1 Xxd of the input that we shall use below using equation 9.
[[ii:]]
if e K, otherwise
K
0d
u](;;:]
if Ω otherwise
Qi,: ==
(20)
0d
[i];:]
if E V, otherwise"
Vii. = 0d
Our arguments do need d. However we do not need d = N but we made this simplification for case of presentation.
57

---

Published as a conference paper at ICLR 2024
Finally, we present the BASECONV model that solves the multiple-query associative recall problem with depputddeeedentt kernels using (11) layer and O Nd)-man -any parameters.
Theorem H.38. There exists BASECONV model with gated data-dependent convolutions that solves the multiple-query associative recall problem on inputs from {0. 1 1 xe with the total num- ber of distinct interaction distances bounded by t in (1)) layers and O Ne) total parameters.
Proof. We note that we have the input dimension d = c. Now, for any input sequence u E {0, 13 Nxdxd we get the input-dependent kernels in equation 19 using autocorrelation of the in- put. We will now outline the following computations for the BASECONV layers:
y Linearq (u) (h* Lineark(u))
Q @ K K)
(21)
2 = Linearg (y) @ (h" * *Linearv(u))
Eo
(22)
where we have the linear projections Linearq(u) = Q, Lineark = K, Linear /(u) = V and Linear E defined as
1d
j E [d] such that y(i.
E i | == Linear
(23)
0
otherwise
Here, we will first present the argument for the special case when we have t = as that will help us elucidate the general case. To this end, as the kemnels from equation 19 for f 1 are given by
h (x) = X h (X) = X
(24)
we observe that convolving with these kernels h y is equivalent to operating with the following primitives (Appendix H.3):
shift_down (1,81);
(25)
shift_down (y, 1).
We note that shift down instead of shifting up as the index of the top-left entry is (0,0). We can then write down the computations performed in equation 21 and equation 22 as follows:
Q @ shift Seett_domm(K,s,
(26)
z shift (V,s1 1),
(27)
We begin by examining y below:
ylū,: :] = (Q o shift.down(K, s1)) |i
Q @ shift_down(K K, 81 [i
(28)
[ if Q
uli s1 if - if
(29)
otherwise
otherwise
Juli, o uli s1,:] if Q and i $1 K, 00d
otherwise
Here, we use the fact that the Hadamard product is row-independent in equation 28, and the defini- tions of the projections from equation 20 in equation 29. Examining the jth entry, we get
1 if E Q, 1ii-s E K. and qi ki-s; Cj {
u [ uli s1.1l
otherwise.
That is, we can express
c, y
if Q s1
and q: = ki-s = c
(30)
0d
otherwise
Consequently, as per the definition in equation 23, get
[1d
if Q,i si $1 E K and q, = k;-si
Bli, || =
(31)
dd otherwise
58

---

Published as a conference paper at ICLR 2024
We can now finally specify the output from equation 27 as follows:
=|t,:| (E shift.down 1)) [i
Eii, shift down - [)
(32)
= (for
d if Q. - E K and q: = ki-
2 ( - +1.
if $1 + 1 E V.
otherwise
otherwise
(33)
ulli 0dd
81 + 1,d
if if 81 E - 81 + 1 and q: = k
otherwise
Vi-81 Odd
if qf ki-s otherwise
(34)
Again, we use the fact that the Hadamard product is row-independent in equation 32, and the defini- tions of the projections from equation 20 equation 33. Overall, we have solved associative recall for all queries that have interaction distance exactly equal to
In order to generalize this arbitrary S N. 3 we first increase the internal dimension so that the input to the kernels u E xd equation 19 and the projections K' , V E (N-t)xd are given by
0d
0
Q
004
u'
K'
Q'
V
0d
0t
Q
0
u
We then observe that for hx):: * and h²(X)= X we have
hx(X)= Σ bx() (1e-1)
8E/A
h²(X)= Σ hy(()) (((-1)-
SE
59

---

Published as a conference paper at ICLR 2024
In analogy with equation 25, we can then equivalently write
h K
04
(h*
K)
h
0d
h K
K
h
h K
h K shiff_down. (ooown(KK00 8t)
K, 82)
shift.dome (iitt.donnKK.. 81))
Similarly, we also have
shift down Somn(V,,, V, s; 1)
hh V')
shift down (V 82 - 1)
shift down (V, 81 1)
60

---

Published as a conference paper at ICLR 2024
That is, the argument for now applies to each of the t shifts now have (cf. equation 21)
* V) y' Q' ② (h
Q
shift.down V
Q
shift V, 88
shift_dowm. (. .11 1)
shift_down V 1)
Q o shift_dowe down (V 1)
shift_down (V,8) (iniadownnW...S.S 1)
yt
3/1
where, for each & [f], we have (ef. equation 30)
ej if - KC and q: = kk-ss = Zej
ye
0d
otherwise
We then analogously get E' follows:
Linearg(y:)
E
E' = Linearg(y'))
Linearg(y2)
Linear
E
where, for each E E [t], we have (ef. equation 31)
if Q, of K. and q: kk-at
Eg
od
otherwise
61

---

Published as a conference paper at ICLR 2024
The output in the general case is then given by
z'
E
(h' V")
E
shift down (V, 1)
E2
Chift_downr V 1)
E
shift_down V $1 1)
E, O saift_down V - 1)
Eg o shift V, - 1)
Ej @ shift down V88 11-11
2x
22
21
where, for each l [4], we have (ef. equation 34)
zr(i,
Ui-o;
0a
otherwise
Finally, we define the last output layer to compute Znut = Linear >{e(() Zz so that we have
Vi-ez if q. = ki-s for some l [t]
Zout [i
O otherwise
To recall, we retrieved the top interaction distances of the input u using auto-correlation and defined the corresponding convolution kernels (equation 19). We then shifted djown the keys K using the first kernel and gated with the corresponding queries Q so that we got match exactly when there exists key that is at sę interaction distance from the corresponding query. After "smearing" this match to get E, we used as a mask to retrieve the value in the next layer. Overall, since we have atomic kernels that perform í shifts with each of these kernels using O(Nd) parameters, we can conclude that the output solves the associative recall problem for all queries with exactly interaction distance from the corresponding keys for all i [t] using ((1) layers and O(t Ne)
parameters as we have d = c.
62

---

Published as a conference paper at ICLR 2024
Table 7: Attention Training Settings
73M
125M
360M
Optimizer
Adam
Optimizer momentum Optimizer eps
81, 32 0.9, 0.95 le - 8
Precision
BFloat16
Learning rate decay
Cosine
Learning rate (min, base)
8e-5, 8e-4
Global batch size
256
Training iterations
20000
Warmup Duration (Linear) Weight decay
0.01
0.1
Num Layers
12
24
Hidden Size
704
768
1024
FFN Width
4 Rotary True
Position Embeddings Weight Tying
Number of Heads (H)
12
16
Table 8: Attention FLOPs Computation
Equation
Input Layer
Sequence Mixer Q, K. V Projections
B N D D x 3
Sequence Mixer Attention
BXHxH D H x N N X
Sequence Mixer Output Projection
Bx N x D D
Channel Mixer (FFN Width
B D D X 8 X
Language Modeling Head
63

---

Published as a conference paper at ICLR 2024
Table 9: BASECONY Training Settings
168M
354M
Optimizer
Adam
Optimizer momentum
B1 82 0.9, 0.95
Precision
BFloat16
Learning rate decay Learning rate (min, base)
Cosine
8e-5, 8e-4
Global batch size
256
Training iterations
20000
Warmup Duration (Linear)
0.01
Weight decay
0.1
Num Layers
30
48
Hidden Size
852
1080
FFN Width
2
Position Embeddings
Weight Tying
True
Short Conv. Filter Size
3
Exp. Mod. Decay (Fast, Slow)
0.3, 1.2
Filter Sine Freq. (w)
14
Filter Order
64
Filter Inner MLP
2
Filter Weight Decay
0
Table 10: BASECONV FLOPs Computation
Equation
Input Layer
B V N X D
Sequence Mixer Long Convolution
10 x N log(N) X 0.5(D) B
Sequence Mixer Short Convolution
x N X 0.5(D)
Sequence Mixer Implicit MLP (Order 64) Sequence Mixer Linear Projection
0.5(D) 64
Channel Mixer (FFN Width 2)
B x D 2
Language Modeling Head
B BxVxNxD
64

---

Published as a conference paper at ICLR 2024
Table 11: Hyena (Poli et al., 2023a) Training Settings
72M
158M
358M
Optimizer
Adam
Optimizer momentum
B1, 82 0.9, 0.98
Precision
BFloat16
Learning rate decay
Cosine
Learning rate (min, base)
(le-5, 8e-4)
Global batch size
256
Training iterations
20000
Warmup Duration (Linear)
0.01
Weight decay
0.1
Num Layers
18
24
Hidden Size
768
864
1024
FFN Width
2 None True
Position Embeddings
Weight Tying
Short Conv. Filter Size
3
Exp. Mod. Decay (Fast, Slow)
0.3, 1.2
Filter Sine Freq. (w)
14
Filter Order
64
Filter Inner MLP
2
Filter Weight Decay
0
Table 12: Hyena FLOPs Computation
Equation
Input Layer
B X D
Sequence Mixer Input Projection
B N D D 3 B D
Sequence Mixer Long Convolution Sequence Mixer Short Convolution Sequence Mixer Implicit MIP (Order 644)
10 N X log N D
3 x B N D
D 64
Sequence Mixer Output Projection
N N
Channel Mixer (FFN Width 2)
D
Language Modeling Head
BxVxNxD
65

---

Published as a conference paper at ICLR 2024
Table 13: H3 (Fu et al., 2023c) Training Settings
72M
168M
357M
Optimizer
Adam
Optimizer momentum
31. 82 0.9, 0.95
Precision
BFloat16
Learning rate decay Learning rate (min, base)
Cosine
(8c-5, 8c-4)
Global batch size
256
Training iterations
20000
Warmup Duration (Linear)
0.01
Weight decay
0.1
Weight Tying
True
Num Layers
19
33
Hidden Size
720
864
1024
FFN Width
2 None
Position Embeddings
State Space Model State Size
64
Table 14: RWKV Peng al. (2023) Training Settings
72M
169M
351M
Optimizer
Adam
Optimizer momentum Optimizer eps
81. 82
0.9, 0.999
Precision
BFloat16
Learning rate decay Learning rate (min, base)
Cosine
le-5, 8e-4
Global batch size
256
Training iterations
20000
Warmup Duration (Linear)
0.01
Weight decay
0.1
Weight Tying
False
Num Layers
20
Hidden Size
624
768
984
Position Embeddings
None
Initialization
From Reference Impl.
66

---

Published as a conference paper at ICLR 2024
Table 15: RetNet (Sun al., 2023) Training Settings
152M
Optimizer Optimizer momentum Optimizer eps Precision
Adam
B1, .2 0.9, 0.98
1.0e-8 BFloat16
Learning decay Learning rate (min, base)
Linear
(le-5, 8c-4)
Global batch size
256
Training iterations
20000
Warmup Duration (Linear)
0.01
Weight decay
0.05
Normalization
Layernorm
Weight Tying
True
Num Layers
12
Hidden Size
768
Value Hidden Size
1280
Window Size
128
Position Embeddings Initialization
xPos (relative position embeddings) DeepNet
Table 16: Simple Long Convolution Fu et al. (2023c) Training Settings
76M
128M
360M
Optimizer
Adam
Optimizer momentum Precision
$1. B2 = 0.9, 0.95
BFloat16
Learning rate decay Learning rate (min, base)
Linear
8e-5, 8e-4
Global batch size
256
Training iterations
20000
Warmup Duration (Linear)
0.01
Weight decay
0.1
Num Layers
6
12
24
Hidden Size
704
864
1024
FFN Width
Position Embeddings
Weight Tying
True
Channels
Lam
0.001
Kernel Dropout
0.1
Kernel LR
5e 5
Activation
GeLU
Exponential Modulation
True
67


Reviews:
[{"review_id": "LL4kNOsUMV", "reviewer": "Reviewer_9Su6", "summary": "The authors study the performance gap between Gated Convolution Models (GCM) and Transformers on language modeling. They show that the lacking capability of GCM on the Associative Recall (AR) task can explain most of the perplexity gap in the real world scenario. They further demonstrate two architecture modifications to GCM, i.e., adding extra attention layers or selectively looking-up for associative tokens, can bridge the performance gap. A theoretical bound of the number of parameters required for GCM to solve the Multiple Query AR task is also derived.", "questions": "According to Table 1, it seems that as the number parameters increase the performance gap between GCMs and Transformers is shrinking. Can you explain this phenomenon? Is it possible that GCMs may outperform Transformers on a larger scale such as with 700M or 1B parameter counts?\n\nCan you provide mathematical formulas for calculating the perplexity scores of AR hit tokens mentioned in Table 1 and Table2?\n\nHow is the selectively look-up mechanism exactly implemented to produce the numbers in Table 2? Is the attention based look-up trainable or not trainable? How is it added to the GCMs layer? Can you provide a series of formulas to describe a GCM layer that is equipped with the proposed selectively look-up mechanism?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper provides a comprehensive study of the associative recall capability of neural language models with different neural architectures, and examines its impact on the next token prediction performance of the models on real world data.\n\nThe authors empirically demonstrate that boosting the associative recall capability of GCMs can mostly bridge its performance gap with the attention-based model under the scale of 360M parameters.\n\nThe authors derive a theoretical scaling bound for data-independent GCMs to solve AR, and validate it with synthetic data.", "weaknesses": "The novelty of the paper is limited. The lack of AR ability of State Space Models (which is a special kind of long convolution model with embedded recurrency) has been analyzed in the H3 paper [1] through synthetic data. The proposed architectural modification of hybridization is a simple replication of the Hybrid-H3.\n\nThe scale of the experiments is limited. The authors only empirically examine their hypothesis for models with the size up to 360M number of parameters. It is not clear whether their claims still holds empirically given the shrinking trend of the performance gap that can already be observed under the current setting. \n\nThe paper does not provide important technical details for reproducibility. The implementation details of the proposed modification of selectively look-up is missing.\n\nThe research problem that the authors are trying to solve has been alleviated with existing [1,2] or emerging solutions [3,4]. The authors do not examine the empirical AR ability of data-dependant convolutions by claiming technical difficulties, but there does exist data-dependant SSMs, such as Liquid-S4 [2], that support causal language modeling. Not to mention the latest GCM, Monarch Mixer [3], that also supports causality. On hybridization, a previous work [4] on dynamic input subset selection for attention modules has also been proposed for efficiently combining SSMs with attention. The authors should consider comparing the proposed architectural modifications with these works to avoid being outdated upon publication.\n\n---\n[1] Hungry Hungry Hippos: Towards Language Modeling with State Space Models (ICLR 2023)\n\n[2] Liquid Structural State-Space Models (ICLR 2023)\n\n[3] Monarch Mixer: A Simple Sub-Quadratic GEMM-Based Architecture (NeurIPS 2023)\n\n[4] Sparse Modular Activation for Efficient Sequence Modeling (NeurIPS 2023)", "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": "", "extra_scores": {}}, {"review_id": "luL4VyJwP5", "reviewer": "Reviewer_xdZW", "summary": "This paper demonstrates that convolutional LMs are worse at bigram associative recall (AR) than transformers on real data by demonstrating that the test AR perplexity is sensitive to the number of occurences of the given bigrams in the training set. This possibly suggests that convolutional LMs are worse at in-context learning than previously thought in other works that test on synthetic tasks. This work also proposes a new synthetic task called a multiquery associative recall (MQAR), and uses new theoretical insights to devise a simple convolution-only alternative competitive to attention on AR.", "questions": "(wrote under weaknesses)", "limitations": "", "rating": 3, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "Using the bigram frequency and the test perplexity of real data was insightful.", "weaknesses": "There's a lot of typos and confusing writing. It's hard to properly understand all the theoretical claims of the work. For starters, Proposition 4.3 should say `u \\in \\{ 0, 1 \\}^{N \\times \\log c}`. The description of definition 3.1 could be improved; a simple example would be quite helpful. In appendix, `N` sometimes means the entire sequence length or the number of triplets in MQAR; this confusion is exacerbated by the fact that the meaning changes in the same theorem/proof.\n\nPage 6 states \"[e]ach BaseConv layer uses O(Nd) parameters and ... O(Nd) operations\". I believe it uses `O(Nd +d^2)` parameters and `O(d N log N + Nd^2)` FLOPs? This makes me believe that the rest of the theoretical results may have to be carefully revisited by the authors.\n\nPage 26. Proof of C.19. It's unclear how `Q[i, :]` can be set to a zero vector when `i \\notin Q` (also boldface 0 suffices to express a vector; no need for a superscript d; also, d is defined as `log(C)` in the same page but `C` is a set. `d = log(c)` small c since `c := |C|`.) and similarly for `K` and `V`, because QKV is actually a linear projection of the input. Linear projection can not implement this non-linear operation of masking some of the activations out. There's a sentence that reads `... where Q, K, V ... are positional embeddings ...`. They are different projections of u.\n\nThe calculation of the percentage of gap due to AR hits could be better motivated and justified in appendix.\n\nOther minor typos:\n - `u * k = FFT^{-1} (FFT(u) FFT(k))` (should drop the convolution operation in frequency domain).\n - Page 5: Other tokens: ... `1,000` => `1250`.", "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": "", "extra_scores": {}}, {"review_id": "tffThhYAox", "reviewer": "Reviewer_5EW4", "summary": "The authors systematically study the impact of neural architecture on language modeling performance. The authors identify a persisting quality gap between convolution and attention networks. Specifically, they identify a single failure mode, i.e., multi-query associative recall (MQAR), and demonstrate that convolution networks fall short in this. To verify the impact of this gap, the authors conduct systematic studies and provide both empirical and theoretical evidence. Moreover, the authors present strategies for migrating this gap.", "questions": "As to the training stability and the sensitivity to the training hyper-parameters, I'm wondering how the proposed look-up method and the hybrid method compare with the attention network.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The studied problem is important and may have a big impact. The pinpointed failure model (i.e., associative recall) is novel and reasonable. Both empirical and theoretical studies are conducted to support the argument. To further demonstrate the impact of the analyses, the authors examine two alternative strategies, which support the intuition of the author.", "weaknesses": "The proposed attention hybrid method seems to perform well in the experiment. However, it is not clear how it would perform on a larger scale.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)