# Zoology: Measuring and Improving Recall in Efficient Language Models

###### 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 → \rightarrow ?? . 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. Our code is accessible at: https://github.com/HazyResearch/zoology .

## 1 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 convolutional 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 that these models, which exhibit better asymptotic scaling in input sequence length than attention, can match attention in language modeling quality ( Poli et al., 2023a ; Peng et al., 2023 ; Fu et al., 2023b ) .

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 is still a perplexity gap of up to 2.1 points between state-of-the-art convolution-based architectures and strong Transformer baselines in 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! ⏟ Key-Value ​ It means ​ no worries ⏟ Key-Value ​ for the rest of your days! ​ Hakuna ⏟ Query ​ Matata ⏟ AR Hit ​ means ​ no ⏟ Query → worries ⏟ AR Hit \underbrace{\text{Hakuna Matata!}}_{\mathclap{\textbf{Key-Value}}}\text{It means}\underbrace{\text{no worries}}_{\mathclap{\textbf{Key-Value}}}\text{for the rest of your days!}\underbrace{\text{Hakuna}}_{\mathclap{\textbf{Query}}}\underbrace{\text{Matata}}_{\mathclap{\textbf{AR Hit}}}\text{means}\underbrace{\text{no}}_{\mathclap{\textbf{Query}}}\rightarrow\underbrace{\text{worries}}_{\mathclap{\textbf{AR Hit}}}

We find that errors on “AR Hits” ( e.g. worries above) account for 82% of the perplexity gap to attention on average, despite only representing 6.4 % 6.4\% of all tokens in the Pile dataset. 1 1 1 We measure AR on real data using a simple heuristic: n n -gram tokens that are repeated in context. A 70M parameter Transformer can predict AR Hits better than a 1.4Bn parameter Hyena gated convolution model ( 20 × 20\times larger) (Table 1 , Table 5 ). 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.) ( Section A.1 ). Prior work argues that a 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 as a tool in designing new architectures and prior work has shown that gated convolution architectures match attention on synthetic AR tasks used as proxies for real language modeling ( Fu et al., 2023a ; Poli et al., 2023a ; Lutati et al., 2023 ) . For this reason, the downstream AR perplexity gaps are surprising.

Through measuring recall on real data, we learn the key disparity is that prior synthetic formulations assume there is one query per input, at a fixed position in the sequence, where tokens come from a small vocabulary size ( e.g. | V | < |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-query 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-convolutions process variable sequences using fixed filters defined by the model weights rather than as 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. at a distance of 10 10 tokens for Hakuna Matata and 9 9 for no worries ) in a parameter and FLOPs efficient way compared to attention. Attention achieves input-dependence since it computes all token-to-token interactions when determining how to mix information in the sequence. We formally describe the limitation of gated convolutions via theory and experiments in the rest of the paper.

We first introduce a simple operator, BaseConv , that can provably simulate the class of architectures 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 sequence length (Theorem 4.4 ) while attention can solve Mqar with model dimension independent of sequence length (Proposition 4.3 , Section 4.3 ). 2 2 2 Note that the runtime of Attention is quadratic, but gated convolutions is near linear. In practice, gated-convolutions appear to encode approximate solutions to AR that support only a subset of token-interaction distances (which affects 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 show that input-dependent sequence mixing is important to 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 on our analysis, 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, we show that simply inserting input-dependent operator — e.g., a convolution filter that shifts the sequence based on the bigram positions or sparse attention placed only on repeated bigram positions — to the BaseConv architecture at < 10 % <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 % >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, we analyze the increasingly-popular convolution-based architectures and identify fundamental limitations. We hope the Mqar framework and our analysis of the role of input-dependence for Mqar inform the design of future architectures. We release our code for reproducability: https://github.com/HazyResearch/zoology .

## 2 Background and Preliminaries

In this section, we describe the setting, introduce notation, and discuss important related work. See Appendix A for a 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 of N N tokens 𝒙 = { x 0 , … , x N − 1 } \bm{x}=\{x_{0},...,x_{N-1}\} drawn from a vocabulary C C , the model outputs a probability distribution over C C for each of x i x_{i} given the preceding tokens P ⁡ ( x i | x 0 , … , x i − 1 ) \mathrm{P}(x_{i}|x_{0},...,x_{i-1}) . The language models in this work share the same high-level architecture. First, each token x i x_{i} in the input is embedded in d d -dimensional space yielding a matrix 𝒖 ∈ ℝ N × d \bm{u}\in\mathbb{R}^{N\times d} . Next, 𝒖 \bm{u} is passed through a stack of L L layers, with layer ℓ \ell outputting 𝒖 ℓ ∈ ℝ N × d \bm{u}^{\ell}\in\mathbb{R}^{N\times d} . Finally, the embeddings 𝒖 L \bm{u}^{L} output by the last layer are mapped back to logits over C C with a linear projection. Each layer transforms 𝒖 \bm{u} with a sequence mixer ( e.g. attention) followed by a state mixer ( e.g. MLP). Unless specified, our models adhere to 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 a sequence via a weighted sum. For example, 𝒚 [ i , : ] = ∑ j = 0 N − 1 ω ( i , j ) 𝒖 [ j , : ] ) \bm{y}[i,:]=\sum_{j=0}^{N-1}\omega(i,j)\bm{u}[j,:]) , where ω \omega is a function outputting scalar weights. We study the differences between two classes of sequence mixers, discussed next.

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 𝐐 , 𝐊 , 𝐕 ∈ ℝ N × d \mathbf{Q},\mathbf{K},\mathbf{V}\in\mathbb{R}^{N\times d} . To compute the output 𝒚 \bm{y} given inputs 𝒖 \bm{u} , attention applies the projections to the input: 𝒒 = 𝐐𝐮 \bm{q}=\mathbf{Q}\mathbf{u} , 𝒌 = 𝐊𝐮 \bm{k}=\mathbf{K}\mathbf{u} , 𝒗 = 𝐕𝐮 \bm{v}=\mathbf{V}\mathbf{u} . The projected embeddings are aggregated according to: 𝒚 = softmax ​ ( 1 d ​ 𝒒 ​ 𝒌 ⊤ ) ​ 𝒗 \bm{y}=\text{softmax}(\frac{1}{\sqrt{d}}\bm{q}\bm{k}^{\top})\bm{v} (shown in Fig. 1 ) in 𝒪 ⁡ ( N 2 ​ d ) \mathcal{O}(N^{2}d) time, which is expensive for long sequences (large N N ).

Gated-convolutions. A more efficient alternative to attention is the convolution, which is defined as 𝒚 [ i , : ] = ∑ j = 0 N − 1 𝐤 [ j , : ] ⊙ 𝒖 [ i − j , : ] \bm{y}[i,:]=\sum_{j=0}^{N-1}\mathbf{k}[j,:]\odot\bm{u}[i-j,:] where the kernel 𝐤 ∈ ℝ N × d \mathbf{k}\in\mathbb{R}^{N\times d} is a learnable weight matrix. Convolutions can be computed in time O ⁡ ( N ​ d ​ log ⁡ N ) O(Nd\log N) using the Fast Fourier Transform (FFT) and the convolution theorem: 𝒚 = 𝐮 ∗ 𝐤 = FFT − 1 ​ ( FFT ⁡ ( 𝐮 ) ⊙ FFT ⁡ ( 𝐤 ) ) \bm{y}=\mathbf{u}\ast\mathbf{k}=\mathrm{FFT^{-1}(FFT(\mathbf{u})\odot FFT(\mathbf{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 a 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 input with a transformed version of itself). Gating was first proposed in the context of convolutions by Dauphin et al. (2017) , but more recently it has played a central role in state-of-the-art sub-quadratic language models, some of which claim attention-level 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 Section H.2 for detailed descriptions of these architectures and their similarities. This work analyzes the differences between attention and the broad class of gated convolution mixers.

## 3 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 it (Section 5 ).

### 3.1 Fine-grained analysis of downstream quality

##### Perplexity Gap

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 using the EleutherAI GPT-NeoX training infrastructure ( Gao et al., 2020 ; Andonian et al., 2023 ) . In the main paper, we compare attention to three state-of-the-art gated convolution sequence mixers: H3, Hyena, and RWKV ( Fu et al., 2023a ; Poli et al., 2023a ; Peng et al., 2023 ) . 3 3 3 RWKV is commonly referred to as an RNN. We show it can be viewed as a convolution ( Section H.2.2 ). We further include a pure long-convolution 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 for all architectures see Appendix C .

Across scales, we find that attention outperforms the gated convolutions by at least a third of a perplexity point on average: the minimum gaps are + 2.14 +2.14 , + 0.59 +0.59 , + 0.35 +0.35 PPL 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 ; Hopfield, 1982 ) (see Appendix A for an extended discussion of AR’s history in machine learning). In Section D.1.1 , we provide annotated Pile examples to demonstrate the phenomenon qualitatively.

Quantifying AR performance. It is challenging to derive a quantitative measure of associative recall performance on the Pile because we don’t know which next token predictions in raw text require associative recall. We use a simple heuristic to identify these tokens, which we refer to as AR Hits . An AR Hit is the last token of an n n -gram repeated in context ( e.g. the second occurence of ` ​ ` ​ R ​ i ​ c ​ e ​ " ``Rice" in Fig. 1 ). However, some common n n -grams ( e.g. “of the” ) could have been memorized during training, so we factor in the frequency with which an n n -grams appeared in the training data. This heuristic enables us to scale our analysis to over 10 million tokens of Pile validation 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 a bigram (a pair of consecutive tokens) which previously appeared in context, but ≤ 1250 × \leq 1250\times during training.

2. Other tokens: (93.6% of tokens) Tokens in the final position of a bigram which did not previously appear in context or it appeared > 1,250 >1,250 times during training.

In Figure 1 , we visualize these slices by plotting log-perplexity 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 1 , we also compute the percentage of the difference in perplexity between attention and each model that is due to AR tokens: Δ ​ log ⁡ ( ϕ AR ) ⋅ | T AR | Δ ​ log ⁡ ( ϕ ) ⋅ | T | \frac{\Delta\log(\phi_{\text{AR}})\cdot|T_{\text{AR}}|}{\Delta\log(\phi)\cdot|T|} , where ϕ \phi is the perplexity and T T is the set of tokens in the test set. This quantity can also be interpreted as the fraction of the overall gap that would close if a 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 on the Pile and repeat this analysis (see Table 5 ). 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 20 × \times 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 use a controlled semi-synthetic dataset to measure AR capacity and show that RWKV’s performance degrades sharply as we increase the number of queries in an example while attention performs consistently well Appendix G .

### 3.2 Formalizing the problem: Multi-Query Associative Recall

This gap in associative recall perplexity is very surprising because prior work shows gated-convolutions can perfectly solve a formalized version of the task ( Fu et al., 2023a ; Poli et al., 2023a ; Olsson et al., 2022 ) . In this synthetic task, the input 𝒙 \bm{x} contains a sequence of bigrams representing key-value pairs from a random dictionary followed by a single query token. For example, the correct output for the input below would be 3 3 : A 4 ​ B 3 ⏟ Key-Value ​ C 6 E 2 F 1 C 6 G 8 → B ? ⏟ Query \text{A 4}\underbrace{\text{B 3}}_{\mathclap{\textbf{Key-Value}}}\text{C 6 E 2 F 1 C 6 G 8 }\rightarrow\underbrace{\text{B ?}}_{\mathclap{\textbf{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-Query AR ( Mqar ). We formally define the Mqar problem as follows: 4 4 4 In Section H.7.1 we define a formal, general form of this definition, used in the theoretical analysis.

###### Definition 3.1 (Multi-Query-AR ( Mqar ) \left(\textsc{Mqar}\right) ) .

We are given an input sequence 𝒙 = { x 0 , … , x N − 1 } \bm{x}=\{x_{0},\ldots,x_{N-1}\} where each x i ∈ C x_{i}\in C is a token drawn from a vocabulary of size c = | C | c=|C| . The task is to check, for every query 1 ≤ i < N 1\leq i<N , whether there exists a 0 ≤ j < i 0\leq j<i such that 𝒖 i ≡ 𝒖 j \bm{u}_{i}\equiv\bm{u}_{j} . If so, output 𝒖 j + 1 \bm{u}_{j+1} .

For example, the correct output for input below would be 4, 6, 1, 2, 3 : A 4 B 3 C 6 ​ F 1 ⏟ Key-Value ​ E 2 → A ? C ? ​ F ? ⏟ Query ​ E ? B ? \text{A 4 B 3 C 6}\underbrace{\text{F 1}}_{\mathclap{\textbf{Key-Value}}}\text{E 2}\rightarrow\text{A ? C ?}\underbrace{\text{F ?}}_{\mathclap{\textbf{Query}}}\text{E ? B ?} In Section 4 , we 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 analyzing 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 architectures built from gating and convolutions. This allows us to make general statements that apply to popular gated-convolution 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 to show that solving Mqar with BaseConv (and other gated-convolution architectures) 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 sequence 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, we define our minimal gated-convolution architecture, called BaseConv . Given a function, we would like to know the most efficient model ( e.g. parameters, FLOPs) that can represent 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 operations. 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) it is universal in that it that can simulate any arithmetic circuit 𝒞 \mathcal{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 an input 𝒖 ∈ ℝ N × d \bm{u}\in\mathbbm{R}^{N\times d} , the BaseConv operator for layer ℓ \ell is defined as: 𝒚 \displaystyle\bm{y} : = ( 𝒖 ⋅ 𝑾 ℓ + 𝒃 1 ℓ ) ⏟ Linear Projection ⊙ ( 𝒉 ℓ ∗ 𝒖 + 𝒃 2 ℓ ) ⏟ Convolution \displaystyle:=\underbrace{\left(\bm{u}\cdot\bm{W}^{\ell}+\bm{b}_{1}^{\ell}\right)}_{\mathclap{\textbf{Linear Projection}}}\odot\underbrace{\left(\bm{h}^{\ell}\ast\bm{u}+\bm{b}_{2}^{\ell}\right)}_{\mathclap{\textbf{Convolution}}}\quad (1) where the layer is parameterized by learnable filters 𝒉 ∈ ℝ N × d \bm{h}\in\mathbbm{R}^{N\times d} , a linear projection 𝑾 ℓ ∈ ℝ d × d \bm{W}^{\ell}\in\mathbbm{R}^{d\times d} , and ‘bias’ matrices 𝒃 1 , 𝒃 2 ∈ ℝ N × d \bm{b}_{1},\bm{b}_{2}\in\mathbbm{R}^{N\times d} . The ⊙ \odot is component-wise product and convolution of two matrices is computed as convolution of the corresponding columns.

In our experiments, each BaseConv layer uses 𝒪 ~ ​ ( N ​ d + d 2 ) \tilde{\mathcal{O}}(Nd+d^{2}) parameters 5 5 5 We use 𝒪 ~ ​ ( ⋅ ) \tilde{\mathcal{O}}{(\cdot)} to hide poly-log factors. and can be computed in 𝒪 ~ ​ ( N ​ d 2 ) \tilde{\mathcal{O}}(Nd^{2}) operations. For our theoretical results, we can assume the weight matrix 𝑾 ℓ \bm{W}^{\ell} is restricted to a class of matrices that support near-linear time matrix multiplication ( e.g. Kaleidoscope matrices, see Definition H.3 ). Under this assumption, BaseConv uses 𝒪 ~ ​ ( N ​ d ) \tilde{\mathcal{O}}(Nd) parameters and 𝒪 ~ ​ ( N ​ d ) \tilde{\mathcal{O}}(Nd) FLOPs ( Proposition H.6 ). We now state the equivalency result between arithmetic circuits and BaseConv a “canonical” representation of arithmetic circuits ( Theorem H.21 in Section H.5 ):

###### Theorem 4.2 (Equivalency to Arithmetic Circuits) .

For an arithmetic circuit 𝒞 \mathcal{C} of size s s and depth Δ \Delta that takes 𝐮 ∈ ℝ N × d \bm{u}\in\mathbbm{R}^{N\times d} as input, there exists an equivalent BaseConv operator that uses 𝒪 ~ ​ ( s ​ Δ ) \tilde{\mathcal{O}}(s\Delta) parameters and 𝒪 ~ ​ ( Δ ) \tilde{\mathcal{O}}(\Delta) layers. 6 6 6 The formal statement in the Appendix has a sharper version of this result in terms of the circuit ‘width’.

In other words, any gated convolution model with small number of layers can be simulated by BaseConv with only a (poly)logarithmic 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 (e.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 to 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 BaseConv and Hyena models can simulate each other with only a small constant blowup in parameters ( Proposition H.12 in Section H.5 ).

### 4.2 Theoretical analysis of gated convolution capacity and associative 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 𝐮 ∈ { 0 , 1 } N × 3 ​ c \bm{u}\in\{0,1\}^{N\times 3c} , Attention (even without using soft-max) solves Mqar for 𝐮 \bm{u} using 𝒪 ⁡ ( c 2 ) {\mathcal{O}}(c^{2}) parameters, 𝒪 ⁡ ( N ​ c 2 + N 2 ​ c ) \mathcal{O}(Nc^{2}+N^{2}c) time complexity and 𝒪 ⁡ ( 1 ) {\mathcal{O}}(1) layers.

It is natural to wonder then if all pairwise comparisons among tokens are necessary to solve Mqar . Indeed, in the RAM setting, a sequential algorithm can simply utilize N N logarithmic insertion and membership queries to solve Mqar in subquadratic time. Unfortunately, any model attempting to emulate this would require Ω ⁡ ( N ) \Omega(N) layers. Instead, we observe that we can parallelize this algorithm using dyadic intervals and achieve a depth of 𝒪 ~ ​ ( 1 ) \tilde{\mathcal{O}}(1) ( Proposition H.30 ). We then convert this algorithm into an arithmetic circuit and apply Theorem 4.2 to derive an equivalent BaseConv model. This allows us to prove new upper bounds for BaseConv models applied to Mqar , which improves upon the quadratic time complexity of attention to near-linear runtime at the cost of using poly-log layers ( Theorem H.37 in Section H.7 ).

###### Theorem 4.4 (Data-Independent Filters 7 7 7 We note here that existing architectures also use data-independent convolution filters, meaning the filter is defined as a function of the model parameters, independent of the input. ) .

Given an input 𝐮 ∈ { 0 , 1 } N × 𝒪 ⁡ ( log ⁡ c ) \bm{u}\in\{0,1\}^{N\times\mathcal{O}(\log{c})} to Mqar (where we assume that distinct tokens are embedded into distinct vectors in { 0 , 1 } 𝒪 ⁡ ( log ⁡ c ) \{0,1\}^{\mathcal{O}(\log{c})} ), there exists a BaseConv operator that solves Mqar for 𝐮 \bm{u} using 𝒪 ~ ​ ( N ​ log ⁡ c ) \tilde{\mathcal{O}}(N\log{c}) parameters as well as time complexity and 𝒪 ~ ​ ( 1 ) \tilde{\mathcal{O}}(1) layers.

Nevertheless, the poly-logarithmic number of layers in the above result is undesirable in practice. But, we show that using input-dependent convolution filters, one can get constant many layers (for a sub-class of inputs). Towards that end, we define the interaction distance between a query 𝒒 i \bm{q}_{i} and the matching key 𝒌 j \bm{k}_{j} as i − j i-j . This then allows us to present the corresponding upper bound for data-dependent mixing ( Theorem H.38 in Section H.8 ).

###### Theorem 4.5 (Input-Dependent Filters) .

Given an input 𝐮 ∈ { 0 , 1 } N × c \bm{u}\in\{0,1\}^{N\times c} to Mqar (where we assume that the tokens are embedded as one-hot encoding in { 0 , 1 } c \{0,1\}^{c} and there exists at most t t distinct interaction distances), 8 8 8 Note that the interaction distances can be arbitrary: there is just a bounded number of distinct distances. there exists a BaseConv operator that uses input-dependent kernels to solve the above case of Mqar using 𝒪 ⁡ ( t ⋅ N ​ c ) \mathcal{O}(t\cdot Nc) parameters and 𝒪 ⁡ ( 1 ) {\mathcal{O}}(1) layers.

### 4.3 Empirical analysis of gated convolution capacity and associative recall

In this section, we 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 8,192 , varying model dimension and sequence length from 64 64 to 512 512 . Appendix E provides further details on the formulation and construction of this synthetic task. Following Olsson et al. (2022) , we train two layer models with a Transformer backbone that interleaves sequence mixing and state mixing (MLPs). For each architecture, we sweep four learning rates from log ⁡ ( − 4 ) \log(-4) to log ( − 2 ) } \log(-2)\} for each architecture, and report maximum test accuracy.

Our results, which are summarized in Figure 2 , support two main claims:

Claim 1 (Gated-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 it with near-constant dimensionality. 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 a constant model dimension of 64. In contrast, Mqar does not achieve accuracy > 0.9 >0.9 unless d ≥ N d\geq N .

Claim 2 (Input-dependent filters). Using input-dependent filters in gated-convolution 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 j if matching keys are separated by j j tokens. We evaluate two approaches for constructing this filter: (1) programatically ( i.e. hard-coded comparisons between token ids) or (2) with autocorrelation, which could learn to perform fuzzy-matches (see Section 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 to satisfy causality and using an O ⁡ ( N ​ log ⁡ N ) O(N\log N) filter per gap could be expensive if each gap applies only to a small number of bigrams. A simpler and perhaps more efficient way to solve the problem would be to introduce a small amount of attention 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.

## 5 Closing the Associative Recall Gap

In this section, we evaluate hybrid BaseConv -Attention 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 sub-quadratic scaling. We also show that BaseConv hybrids can outperform attention-only 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 : ℝ N × d → { 0 , 1 } N f:\mathbb{R}^{N\times d}\rightarrow\{0,1\}^{N} . They take as input 𝒖 ∈ ℝ N × d \bm{u}\in\mathbb{R}^{N\times d} and output 𝒚 ∈ ℝ N × d \bm{y}\in\mathbb{R}^{N\times d} : 𝒚 [ i , : ] = softmax ( 1 d 𝒒 [ i , : ] 𝒌 ⊤ ) 𝒗 ⋅ f ( 𝒖 ) [ i ] \bm{y}[i,:]=\mathrm{softmax}(\frac{1}{\sqrt{d}}\bm{q}[i,:]\bm{k}^{\top})\bm{v}\cdot f(\bm{u})[i] (2)

where 𝒒 , 𝒌 , 𝒗 \bm{q},\bm{k},\bm{v} are query, key, and value projections, as in attention. We evaluate four choices for f f :

(1) Full attention. First, we evaluate the performance of full attention hybrids. These hybrids correspond to fixing f ​ ( 𝒖 ) ​ [ i ] = 1 f(\bm{u})[i]=1 for all i i 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 a control, we evaluate sparse attention randomly applied to tokens. This corresponds to a stochastic selection function where f ​ ( 𝒖 ) ​ [ i ] f(\bm{u})[i] is drawn from a Bernoulli. Many prior works have employed sparse attention patterns like this one, which are independent of the input ( Zaheer 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 associative recall hits. Specifically, f ( 𝒙 [ i , : ] ) f(\bm{x}[i,:]) is 1 1 if the token x i x_{i} previously occurred in the sequence. In practice, we compare raw token ids, not token embeddings. f ​ ( 𝒙 ) ​ [ i ] = { 1 if there exists j < i such that x i = x j 0 otherwise f(\bm{x})[i]=\begin{cases}1&\text{if there exists}\quad j<i\quad\text{such that}\quad x_{i}=x_{j}\\ 0&\text{otherwise }\end{cases} (3)

(4) Learned selection. Finally, we prototype a learned selection function f ( 𝒖 ) [ i ] = σ ( 𝒖 [ i , : ] ⋅ 𝑾 ) f(\bm{u})[i]=\sigma(\bm{u}[i,:]\cdot\bm{W}) parameterized as a simple linear layer with sigmoid activation. We fix a hyperparameter k k and select the top- k k tokens with the highest score in each batch. This allows us to compute attention in 𝒪 ⁡ ( n ​ d ​ k ) \mathcal{O}(ndk) time. During training, we add small amount of Gaussian noise to the top- k k calculation to encourage exploration and use an auxiliary loss to encourage sparse selection: ℓ f ​ ( 𝐮 ) = 1 N ​ max ⁡ ( 0 , ∑ i = 1 N f ⁡ ( 𝒖 ) ​ [ i ] − k ) \ell_{f}(\mathbf{u})=\frac{1}{N}\max(0,\sum_{i=1}^{N}f(\bm{u})[i]-k) . This approach is most similar to the recently proposed SeqBoat architecture ( Ren et al., 2023 ) . We discuss the differences in Appendix A .

##### Downstream Evaluations

We evaluate the prototypes on Pile language modeling. We take BaseConv architectures at the 150M and 360M parameter scales and add input-dependent selection to three layers (details in Section C.2 ). Results are in Table 2 . We validate that the prototypes close the overall and AR quality gaps to attention when added to the BaseConv backbone.

At 360M parameters, BaseConv with just 3 3 attention layers can outperform the Transformer, while requiring fewer FLOPs. Our hybrid attention- 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% of the gap between pure BaseConv and attention on the AR slice, in contrast to the random selection control. Learned selection closes 72% a of the gap using just k = 256 k=256 (sub-quadratic) attention positions per example.

## 6 Discussion and Conclusion

We present an extensive analysis of gated convolution architectures in light of their recent popularity. We identify a persisting quality gap between efficient convolution and inefficient attention based architectures, largely due to a 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 on the Pile.

Our results in analyzing gated convolutions go beyond the conventional wisdom that attention is the “ right ” model. A significant amount of work focuses on 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 et 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, as a proxy to study.

## 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 to Together Computer for making this work possible. We gratefully acknowledge the support of DARPA under Nos. FA86501827865 (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, BASF, Accenture, Ericsson, 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) Program, Fannie and John Hertz Foundation, National Science Foundation Graduate Research Fellowship 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.

## Appendix

The appendix includes the following content: 1. Appendix A provides an extended discussion of related work and concepts.

2. Appendix B provides a code implementation of the BaseConv architecture.

3. Appendix C gives details for the experiments, including model architectures and hyperparameters.

4. Appendix D provides additional analysis of how Mqar appears in real data across a variety of language distributions.

5. Appendix E provides a formal definition of the Mqar problem, synthetic construction procedure, and experimental details.

6. Appendix F provides additional synthetic experiments and analysis.

7. Appendix G provides experiments and analysis for how the Mqar gap changes as we scale the gated convolution and attention architectures.s

8. Appendix H gives proofs and additional discussion for the theoretical analysis in our work.

## Appendix A Extended Related Work

We are inspired by and build on prior work in mechanistic interpretability ( Section A.1 ), efficient architecture design ( Section A.2 ), and input-dependent architectures ( Section A.3 ).

### A.1 Mechanistic Interpretability, Synthetic Languages, and Associative Recall

Work in mechanistic interpretability aims to decompose the capabilities of a neural network into human-understandable algorithms that can 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 architectures [ Wang and Eisner, 2016 , White and Cotterell, 2021 , Allen-Zhu and Li, 2023 , Ravfogel et 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 on one particular synthetic task: associative recall . Motivated by psychological models of how humans associate and retrieve information, work in the early days of neural network research focused on developing systems capable of associative recall [ Willshaw et al., 1969 , Feldman et al., 1981 , Hopfield, 1982 ] . For example, Hopfield networks, proposed in 1982, are a recurrent neural network explicitly designed to support associative (“content-addressable”) memory [ Hopfield, 1982 ] . More recently, several notable recurrent neural network mechanisms ( e.g. neural turing machines, LSTMs, and RNNs with fast weights) were evaluated on a synthetic version of associative recall [ Graves et al., 2014 , Ba et al., 2016 , Zhang and Zhou, 2017 , inter alia.] . These works use a formulation of associative recall very similar to 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, at least in 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 convolution architectures. While attention requires compute that scales as 𝒪 ⁡ ( N 2 ) \mathcal{O}(N^{2}) in sequence length N N , convolutions scale as 𝒪 ⁡ ( N ​ log ⁡ N ) \mathcal{O}(N\log N) [ Cooley and Tukey, 1965 ] . Ideal, inference complexity is 𝒪 ⁡ ( 1 ) \mathcal{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-quadratic training and constant inference complexity in sequence length [ Gu et al., 2021 ] . Architectures that use implicit convolutional filters [ Poli et al., 2023a ] , can be converted to an SSM via a simple distillation step [ Massaroli et al., 2023 ] .

### A.3 Input-Dependence in Sequence Modeling Architectures

In an input-dependent sequence model, the way tokens are aggregated across the sequence is controlled 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 et al., 2017 ] achieves input-dependent sequence mixng since the Q Q , K K , and V V terms are achieved via linear combinations of the input x x . For instance, Q = x ​ W Q Q=xW_{Q} for some learned weight matrix W Q W_{Q} .

• 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 input. 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 subsequent work improved upon S4 by changing the parametrization [ Gupta et al., 2022 , Gu et al., 2022 , Mehta 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 is 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, an important task in in-context learning [ Olah, 2022 ] , subsequent work proposed to introduce input-dependence by adding a gating or element-wise multiplication operation to the architecture, where both the inputs to the operator are defined in terms of the input [ Fu et al., 2023a , Poli et al., 2023a ] . The multiplication is generally y = σ ⁡ ( W ​ x ) ⊙ x y=\sigma(Wx)\odot x . The filters remain input-independent. These architectures demonstrate promising results on the associative recall synthetics proposed in prior work Fu et al. [2023a] , Poli et al. [2023a] and provide large downstream improvements over S4. 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 a 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 to 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 architectures 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 n n projections of the input. Liquid S4 Hasani et al. [2022] is also motivated by introducing input-dependence to S4. To accomplish this, the work proposes including correlation terms between 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 3 . 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 3 .

• Recurrent Architectures The pure long-convolution architectures can also be computed as recurrences. Letting s i s_{i} represent the hidden state at time i i , u i u_{i} be the input at time i i , and A , B , C A,B,C be projection matrices, the recurrence computes: s i + 1 = f ⁡ ( A ​ s i + B ​ u i ) s_{i+1}=f(As_{i}+Bu_{i}) y i = C ​ s i y_{i}=Cs_{i} In a linear RNN like S4, the A A , B B , and C C matrices are input-independent and f f is an identify function. However, recent work proposes input-dependent RNNs (e.g. RetNet Sun et al. [2023] ), where a subset of these matrices is input-dependent (e.g. A = x ​ W a A=xW_{a} ). Note that linear attention mechanisms can also be expressed as input-dependent RNNs [ Katharopoulos et al., 2020b ] . While these architectures improve over gated convolutions in Mqar synthetic experiments and downstream quality on 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 ( Section 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 different capabilities. Prior work proposes architectures that hybridize sub-quadratic 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 on language modeling. SeqBoat [ Ren et al., 2023 ] shares closest resemblance to our learned selection module. SeqBoat introduces a different module that learns where to use attention however, their model can end up using full attention at a layer, providing quadratic scaling in the worst case. Different from SeqBoat, we use an auxiliary loss to encourage sparsity in the selected attention positions and select the top- k k positions to remain sub-quadratic.

Overall, our work finds that previously proposed sub-quadratic models do not efficiently solve the Mqar task.

## Appendix B Code Listing

In this section, we include code listings for the BaseConv operator. We begin with the a standard BaseConv implementation that uses an explicitly-parameterized long convolution. Below, we also provide the implementation for BaseConv with an implicitly-parameterized long convolution.

In some of our experiments, we interleave BaseConv layers that use explicit short convolution filters (like those in torch.nn.Conv1d ) and with BaseConv layers that use implicit long convolution filters (like those described in Poli et al. [2023a] ).

Below we include the code for BaseConv with an implicit convolution.

## Appendix C Downstream Experimental Details

We use A100 80GB Nvidia GPUs to run all experiments. We use the reference training infrastructure from https://github.com/EleutherAI/gpt-neox for all pretraining runs. The Pile data is tokenized using the GPT2BPETokenizer and all models see the data in the same order.

Below we provide details on the hyperparameters and settings for training each architecture studied in the paper on the real-world Pile data. In Section C.1 we summarize and justify our method for measuring the quality gap due to associative recall between the convolution architectures and attention. In Section C.2 , we provide details on the pure gated convolution architectures, attention baseline, and hybrid architectures studied in the paper.

### C.1 Measuring the Mqar Gap on Real Language Data

Here we 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 ): 1. Given an input sequence, we identify recurring bigrams (i.e. bigrams that have already appeared in the sequence at a prior position). Since bigrams that appear frequently during training may be memorized by the model, rather than requiring the model to perform recall at inference-time, we only measure AR log-probabilities with respect to bigrams that are seen fewer than a threshold number of times during training. The threshold used in all the experiments in our submission is 1,250 1,250 training occurrences in the 10B tokens of pretraining data.

2. We measure the log-probability assigned to the true bigram completion. This bigram completion is referred to as an AR Hit in our work. This protocol assumes that the model can produce the completion by recalling the prior occurrence of the bigram in the sequence.

3. For the model being evaluated m m , and the attention model M M , we measure the % of the quality gap between m m and M M ascribed to associative recall capability as follows. Let the average log-probability for all AR Hits across validation sequences be l H m l_{H}^{m} and l H M l_{H}^{M} for m m and M M respectively. Let the average log-probabilities of all tokens in the validation sequences be l m l^{m} and l M l^{M} respectively. Let p H p_{H} be the proportion of AR Hit tokens in the validation data. As the final gap ascribed to AR, we report: min ⁡ ( ( l H m − l H M ) ​ p H l m − l M , 1.0 ) \min(\frac{(l_{H}^{m}-l_{H}^{M})p_{H}}{l^{m}-l^{M}},1.0) Shown above, if m m is better than attention ( M M ) overall and M M is better than m m at AR, we ascribe 100% of the gap to AR.

We briefly discuss two important decisions in this protocol. First, we only measure explicit bigrams , i.e. bigrams are identified based on token ids in the sequence. However, intuitively, models may also perform associative recall between related concepts produced by a contextual language model. For instance, language may contain bigrams in which one word is swapped by a synonym. As another example, a model may see a sentence such as “The iPhone is outside my budget so I instead purchased an Android phone. … It was much _” and predict “cheaper” for the blank, recalling that the sequence is discussing cost. Our work does not measure such fuzzy (more abstract) recall instances.

Next, we measure the gap based on log-probabilities rather than perplexity. This is simply because we want to make our metrics independent of the number of tokens in each of the slices of the validation set. Approximately 6.4% of validation tokens are AR Hits with the threshold set to consider bigrams seen less than 1,250 × 1,250\times during training.

### C.2 Gated Convolution Downstream Architectures

We evaluate over 4 4 previously proposed architectures as well as BaseConv , the theoretically “canonical” representation for gated convolutions, introduced in Section 4 , for a total of 14 training runs. Here we provide details on the hyperaparamters and configurations used for training each architecture. We also provide details on the FLOPs computation.

• Attention [ Vaswani et al., 2017 , Touvron et al., 2023 ] We train using the the specifications in Table 6 . The parameters are sourced from the Transformer implementation in https://github.com/EleutherAI/gpt-neox .

• Hyena [ Poli et al., 2023a ] We train using the specifications in Table 8 . The parameters are sourced from the Appendix of Poli et al. [2023a] and the implementation is sourced from the provided reference at https://github.com/HazyResearch/safari .

• H3 [ Fu et al., 2023a ] We train using the specifications in Table 10 . The hyperparameters and implementation use the reference at https://github.com/HazyResearch/H3 .

• RWKV [ Peng et al., 2023 ] We train using the specifications in Table 11 . The parameters are sourced from the Appendix of Peng et al. [2023] and the details provided in the reference implementation at https://github.com/BlinkDL/RWKV-LM . We specifically evaluate RWKV-V4.

• Pure long convolution We train using the specifications in Table 13 . We evaluate a simple long convolution based model with no gating as a reference point. While this is a generic architecture, we use the reference implementation and initialziations/regularizations from recent work Fu et al. [2023c] . The implementation is provided at https://github.com/HazyResearch/safari .

• BaseConv We train using the specifications in Table 15 . The implementation, amounting to 19 lines of PyTorch, is shown in Appendix B .

We provide the equations used to compute the FLOPs for each model, letting D D be the model width, H H the head dimension, L L the depth, N N the sequence length, V V the vocabulary size, and B B the batch size. FLOPs equations for different architecture types are provided in Table 7 (attention), Table 9 (Hyena and adapted for H3), Table 14 (pure long convolution), and Table 16 ( BaseConv ). We compute FLOPs for RWKV as in the Appendix of [ Peng et al., 2023 ] , based on the number of linear layer parameters, plus input and language modeling head FLOPs.

##### Hybrid Architectures

In 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 Section C.2 ). The hybridized (non gated convolution) layers are inserted as 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: • Full attention following the specification of attention in Section C.2 .

• Random selection We evaluate sparse attention randomly applied to tokens as introduced in Section 5 .

• Programmatic selection We apply sparse attention on AR hit tokens as introduced in Section 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 to output high scores if attention should be applied to the token. We can take the top- k k positions to control the computational complexity of the layer. This approach is introduced in Section 5 .

We use these protocols to validate that input-dependence suffices to address the associative recall gap and highlight that there are varied approaches to incorporating input-dependence in an architecture.

## Appendix D Extended Downstream Analysis of Mqar

In this section, we 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 on sources in RedPajama including ArXiv and StackOverflow Together [2023] .

### D.1 Additional Discussion of Mqar in the Pile

The Pile is a widely popular language modeling corpus Gao et al. [2020] . 1. First, in Section D.1.1 , we provide several real examples of how associative recall occurs in the Pile to demonstrate it’s role in language modeling. We color code the tokens to highlight differences in the next token predictions from Attention, RWKV and Hyena models.

2. Next, in Section D.1.2 , we explain where associative recall hits tend to occur in sequences. This is useful to 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 N 2 N^{2} token-to-token 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 b -ist -ro near his apartment : he had been mus -ing about (rwkv= about, attn= on) the … (723 tokens) … the young waitress -’s sigh at the Ma -ison Ber (rwkv= Bl, attn= Ber) -gey (rwkv=-nd, attn=-gey)

The second (rwkv= first, attn= second) section is all about Pixar Fest , and the (rwkv= third, attn= the) final section is all (rwkv= about, attn= all) about Pixar Pier . … (480 tokens) … -If (rwkv=-Disney, attn=-If) there wasn -âĢ -t enough Pixar at Disneyland , Pixar Fest (rwkv= would, attn= Fest) is (rwkv= at, attn= is) coming to the Disneyland Resort on April 13 , 2018 .

David Mus -a P -id -cock is an indigenous English revert to (rwkv=-ant, attn= to) Islam , (rwkv=,, attn= who) who formed the Islamic Party of Britain in 1989 , at London Central Mosque (rwkv= Mosque, attn= University) . (rwkv=., attn= in) …( 711 tokens )… -I have just found out that David Mus (rwkv= Cameron, attn= Mus) -a P (rwkv= (, attn= P) -id (rwkv=-asha, attn=-id) -cock (rwkv=-ham, attn=-cock) is (rwkv= (, attn= is) speaking at a Yorkshire Forum debate today in Bradford with

-The most common map elements (hyena= in, attn= elements) in Sub -Space are prizes , or ” -g (hyena=-points, attn=-g) -reens ” (hyena=-",, attn=") ( -for their green color -). Pri -zes allow players to upgrade their ships and gain (hyena= other, attn= gain) special weapons or abilities (hyena= abilities, attn= upgrades) . While prizes are generally plent -ifully scattered throughout the map (hyena= map, attn= game) , the upgrades or 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 a single unit of measure : energy . Each ship is equipped with a certain amount of energy , from which (hyena= which, attn= the) it must draw its health as 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 .

Thus far , no systematic study has been published on is -olation (hyena=-olation, attn=-ot) via sequential centrif -ug -ation , and no systematic analysis is … (491 tokens) … Fig . 1 (hyena= 1, attn=-Âł) -is -olation via sequential (hyena= centrif, attn= sequential) centrif -ug -ation .

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 ) is named (hyena= the, attn= named) after former (hyena= the, attn= former) J -op -lin Globe publisher and (hyena= and, attn= Fred) Missouri Southern (hyena= State, attn= Southern) board of reg -ents member

#### D.1.2 Distribution of Mqar Hit Positions in Sequences

In Figure 3 , we compute and plot the distances between AR hits and its prior bigram (key-value) occurrence 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 a 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 et 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] .

### D.2 Analyzing Mqar across Varied Language Distributions

We use the HuggingFace sample of the popular RedPajama language modeling corpus Together [2023] to 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 we profile the prevalence of Mqar hits by computing the number of repeated bigram occurrences per sequence. Each sequence is 2,048 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 4 and 5 .

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 domain-specific terminology, groups of related citations, Latex commands, and mathematical equations are frequently repeated 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 are highly variable. The topics and documents in these corpora are relatively diverse and unstructured.

If the placement V V of the hyperplanes is not generic, then the induced subdivision Σ \Sigma is not a triangulation. However, the above bijection is unaffected and, in particular, the fine type of a vertex of C ​ D A CD_{A} can be read off the corresponding facet of Σ \Sigma . We define the crosscut complex c ​ c ​ u ​ t ​ Σ \ ccut{\Sigma} ⊆ 2 [ n ] × [ d ] \subseteq 2^{[n]\times[d]} to be the unique simplicial complex with the same vertices-in-facets incidences as the polyhedral complex Σ \Sigma . The crosscut complex is a standard notion in combinatorial topology and can be defined in more generality (see Björner). The following observation is immediate from the definitions. Proposition D.1 . For V V an ordered sequence of n n points in T d − 1 T^{d-1} let A = A ⁡ ( V ) A=A(V) be the corresponding tropical arrangement and let Σ \Sigma be the regular subdivision of Δ n − 1 × Δ d − 1 \Delta_{n-1}\times\Delta_{d-1} induced by V V . Then the fine cotype ideal f ​ c ​ I \ fcI is Alexander dual to the Stanley-Reisner ideal of c ​ c ​ u ​ t ​ Σ \ ccut{\Sigma} . The crosscut complex encodes the information of which collections of vertices lie in a common face. Hence, the crosscut complex is a purely combinatorial object and does not see the affine structure of the underlying polyhedral complex .

⬇

## Appendix 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 . Synthetic 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 independent of the sequence length N N .

2. Convolutions : The gated convolution models use 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 𝒪 ⁡ ( N ) \mathcal{O}(N) recalls in one forward pass. Our intuition is that it is difficult to efficiently compute all token-interaction distances with such a filter.

3. Recurrences : Recurrent models include a single hidden state as an information bottleneck. We show in Section H.6 that the gated recurrence requires Ω ⁡ ( N ) \Omega(N) bits to solve Mqar for d ≤ N d\leq\sqrt{N} for model dimension d d . Our intuition is that it is difficult to store large numbers of key-value pairs seen in the sequence in low-dimensional hidden states.

Motivated by this analysis, we propose Procedure 1 for generating synthetic Mqar data. This procedure allows toggling two simple properties of the synthetic dataset, specified below, to reveal the differences between the sequence modeling architecture families. We also discuss deviations from the prior work that uses synthetic recall data. 1. Size of the Key-Value Map : The number of unique key-value pairs appearing in each example in the dataset. In Procedure 1 , we use the input parameter D 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 is a single query token that requires recall per example, failing to test whether the architecture can support multiple token-interaction distances. In Procedure 1 , we use the parameters N N (sequence length), D D (size of the key-value map per example), and α \alpha to toggle this property. We enforce that the token-interaction distances appearing in the synthetic examples follow a power-law distribution specified by α \alpha (based on Figure 3 ). Keeping α \alpha and D D constant while varying N 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 \leq 40 tokens.

### E.2 Training Details

We first describe the architectures evaluated on Mqar synthetic data and then the training hyperparameters. We evaluate the following four architecture categories in the experiments: 1. Attention Standard GPT-2 style multi-headed Transformer architecture with learned positional embeddings Brown et al. [2020] . Attention, the core building block of Transformers, is defined in Section 2 . The architecture for synthetics has 1 attention head.

2. Gated convolutions (Hyena Poli et al. [2023a] , RWKV Peng et al. [2023] ). Gating, convolutions, 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 evaluate this architecture (“Chunked RNN”) with chunk sizes of 32 32 and 8 8 in Figure 6 . In RetNet, hidden states are updated as: S n = γ ​ S n − 1 + A n T ​ V n S_{n}=\gamma S_{n-1}+A_{n}^{T}V_{n} Outputs at each timestep are defined as O n = C n ​ S n O_{n}=C_{n}S_{n} where γ ∈ ℝ 1 \gamma\in\mathbb{R}^{1} , A = x ​ W a A=xW_{a} for input x ∈ ℕ × 𝕕 x\in\mathbb{N\times d} and learned weight matrix W a ∈ ℝ d × d W_{a}\in\mathbb{R}^{d\times d} , C = x ​ W c C=xW_{c} for W c ∈ ℝ d × d W_{c}\in\mathbb{R}^{d\times d} , V = x ​ W v V=xW_{v} for W v ∈ ℝ d × d W_{v}\in\mathbb{R}^{d\times d} . RetNet similar to a state-space-model (SSM) Gu et al. [2021] , except for that the matrices A A and C C are input-dependent (i.e. they are functions of the input x x ). We note that RetNet is a special case of the recently proposed Mamba Anonymous [2023a] and GateLoop Anonymous [2023b] architectures, which replace the γ \gamma term in RetNet with yet another input-dependent matrix. The RetNet model is formally defined in Section 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 et al. [2019b] , Qiu et al. [2019] . Consider a window size w w . Sliding window attention permits each token to attend to the prior w w tokens in the sequence. Blocked window attention first splits the sequence into blocks of size w 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 0.1 , warmup duration 10 % 10\% , linear warmup, AdamW optimizer. For each run, we sweep the learning rates in np . logspace ⁡ ( − 4 , − 2 , 4 ) \mathrm{np.logspace(-4,-2,4)} . We train for 64 64 epochs.

• Training duration : The global batch size is 8 8 for input sequence length or model dimension ≥ 512 \geq 512 , 16 16 for input sequence length or model dimension ≥ 256 \geq 256 , and 64 64 otherwise.

• Width and depth : For all synthetic experiments, we use exactly two 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 ).

• 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 100,000 and 3,000 3,000 data points respectively. The data and data order for all runs are constant.

## Appendix F Extended Results on Mqar Across Architectures

In this section, we provide further validation that Mqar is a useful diagnostic tool to explain the behaviors of a wide variety of architecture families. We show: 1. Gated convolutions In the main paper, we 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 N , holding α \alpha and D D constant in Procedure 1 ).

2. Gated recurrences In this section, we show gated recurrrent models use larger dimensionality than attention to solve the task as the number of unique key-value pairs to store ( D D in Procedure 1 ) increases.

3. 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 (intuitively tokens earlier in 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 on Mqar also perform well on downstream real-world AR on the Pile. We use the experimental protocols outlined in Appendix E 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 construct the synthetic dataset following Procedure 1 , we set α = 0.1 \alpha=0.1 and | C | = 8,192 |C|=8,192 for all experiments. We vary N ∈ { 64,128,256,512 } N\in\{64,128,256,512\} to increase the number of token-interaction distances that will appear in the examples and we vary the model dimension ∈ { 64,128,256,512 } \in\{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 α = 0.1 \alpha=0.1 and | C | = 8,192 |C|=8,192 for all experiments. We fix N = 256 N=256 and vary D ∈ { N 32 , N 16 , N 8 , N 4 } D\in\{\frac{N}{32},\frac{N}{16},\frac{N}{8},\frac{N}{4}\} to increase the number of key-value pairs occurring per example. For each number of key-value pairs per example, we additionally vary the model dimension ∈ { 64,128,256 } \in\{64,128,256\} . The results are shown in Fig. 6 . We validate that as the number of associative recall keys and values exceeds the window size, the model again uses increased dimensionality relative to full O ⁡ ( N 2 ) O(N^{2}) attention to solve the task. While Fig. 6 highlights RetNet, we note that it is a special case of the recently released Mamba Anonymous [2023a] and GateLoop Anonymous [2023b] architectures. Corresponding to this experiment, we theoretically show in Section H.6 that the gated recurrence uses Ω ⁡ ( N ) \Omega(N) bits to solve Mqar for d ≤ N d\leq\sqrt{N} . Intuitively, it is difficult to store a 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 1 , we set α = 0.1 \alpha=0.1 and | C | = 8,192 |C|=8,192 for all experiments. We fix N = 256 N=256 and vary D ∈ { N 32 , N 16 , N 8 , N 4 } D\in\{\frac{N}{32},\frac{N}{16},\frac{N}{8},\frac{N}{4}\} to increase the number of key-value pairs occurring per example and the range of token-interaction distances per example. For each number of key-value pairs per example, we additionally vary the model dimension ∈ { 64,128,256 } \in\{64,128,256\} .

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 16 KVs means that there are 16 16 keys and 16 16 values, for 32 32 total tokens. Sliding attention with window size 32 32 performs well up until 16 16 KVs and degrades beyond this point relative to attention models of the same dimensionality. Sliding window attention with window size 8 8 does not perform well on any setting, as expected. Blocked attention also not perform well – intuitively tokens at the beginning of a block lack sufficient prior context to attend to.

### F.2 Demonstrating Mqar as a Tool for Architecture Development

In the main paper, we provide experiments showing that there is gap between the gated convolution and attention models on synthetic Mqar data and correspondingly in the downstream Pile results. In Table 4 , we show that the trends of RetNet and sparse attention on Mqar also carry to the downstream results. We hope Mqar synthetic data may be useful in future architectural development.

## Appendix 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 a 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 is a full perplexity point better on AR than this 1.4Bn Hyena model (Table 1 ). We discuss these results further in ( Section 3 ).

##### Open-source 7B parameter models

We next evaluate the RWKV (gated convolution) Peng et al. [2023] and Llama 2 (attention) pretrained models at the 7B parameter scale Touvron et al. [2023] . 9 9 9 We specifically evaluate the RWKV-Raven model downloaded from https://huggingface.co/docs/transformers/model_doc/rwkv and the Llama 2 model downloaded from https://github.com/facebookresearch/llama . These are both popular models that took a significant amount of effort to train, towards maximizing quality. We find that there is a gap between RWKV and attention at the 7B scale and that it increases as the model needs to conduct more recalls per sequence ( P 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 infrequent bigrams in validation sequences ( Figure 1 , Section 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 a synthetic test to fairly measure the AR capabilities.

2. Evaluation dataset. We construct the synthetic dataset following Algorithm 1 , where each sequence contains P P token pairs, or “bigrams”. Each bigram, containing a “key” token and a “value” token, appears twice in the sequence. On the second occurrence of a “key” token, the model should look back to the prior occurrence to output the corresponding “value” token. We measure the AR perplexity of the model based on its ability to 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 1024 tokens. We measure the AR perplexity when the sequences contain P ∈ { 16 , 32 , 64 , 128 , 256 } P\in\{16,32,64,128,256\} key-value pairs, using 1000 1000 samples per P P value. The tokens that do not contain a key or value are simply filled with a fixed token id (so this token is repeated frequently in the sequence). We plot perplexity for AR and non-AR tokens (fixed token) vs. P P . We find RWKV quality degrades with P P on the AR slice (blue line), while all other lines remain flat. MQAR remains problematic at 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 AR Gap with Scale

In the main paper, we measure the AR perplexity on downstream data based on bigrams that are seen < 1,250 × <1,250\times during pretraining. However, we 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 as attention in-context. In-context learning is defined as learning from examples provided in context [ Xie et al., 2021 ] .

Concretely, the gap between the Hyena / RWKV and attention models at the 350m scale is 1.85 / 1.84 perplexity points when we focus on bigrams seen < 1,250 × <1,250\times during pretraining ( Table 1 ). If we instead focus on bigrams seen 1 × 1\times during pretraining, the gap to attention quality is 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.

## Appendix 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 1 row vector of size k k , given by [ 1 1 … 1 1 ] \begin{bmatrix}1&1&\ldots&1&1\end{bmatrix} , and the all 0 0 row vector of size k k , given by [ 0 0 … 0 0 ] \begin{bmatrix}0&0&\ldots&0&0\end{bmatrix} , as 𝟏 k \bm{1}^{k} and 𝟎 k \bm{0}^{k} , respectively. We also construe the standard basis vector 𝐞 i \mathbf{e}_{i} as a column vector in these notes, and adhere to the following matrix indexing convention: 𝐌 ⁡ [ i , j ] {\bf M}[i,j] is the entry in the i i th row and the j j th column, 𝐌 [ i , : ] ∈ 𝔽 1 × n {\bf M}[i,:]\in\mathbbm{F}^{1\times n} denotes the i i th row, and 𝐌 [ : , j ] ∈ 𝔽 m × 1 {\bf M}[:,j]\in\mathbbm{F}^{m\times 1} denotes the j j th column of 𝐌 ∈ 𝔽 m × n {\bf M}\in\mathbbm{F}^{m\times n} . We then use 𝟏 m × n , 𝟎 m × n ∈ 𝔽 m × 1 \bm{1}^{m\times n},\mathbf{0}^{m\times n}\in\mathbbm{F}^{m\times 1} to denote the matrix of all 1 1 s and 0 0 s, respectively.

Next, we denote the Hadamard product of vectors 𝐮 , 𝐯 ∈ 𝔽 n {\bf u},{\bf v}\in\mathbbm{F}^{n} as 𝐮 ⊙ 𝐯 {\bf u}\odot{\bf v} ; the operation can be extended to matrices by applying the Hadamard product column-wise across the matrices. This is commonly referred to as (element-wise) gating . For vectors 𝐮 , 𝐯 ∈ 𝔽 n {\bf u},{\bf v}\in\mathbbm{F}^{n} , we also denote their linear (or acyclic) convolution as 𝐮 ∗ 𝐯 {\bf u}\ast{\bf v} and cyclic convolution as 𝐮 ⊛ 𝐯 {\bf u}\circledast{\bf v} .

##### Polynomial Notation.

Because convolution is intimately tied to operations on polynomials, it is convenient to use them to discuss the inputs and outputs of gated convolution models. Let us define maps poly , poly ∗ : 𝔽 n → 𝔽 ⁡ [ X ] / ( X n ) \operatorname{poly},\operatorname{poly}^{*}:\mathbbm{F}^{n}\to\mathbbm{F}[X]/(X^{n}) such that poly ⁡ ( 𝒖 ) \displaystyle\operatorname{poly}(\bm{u}) = ∑ i = 0 n − 1 𝒖 [ i ] X i , and poly ∗ ( 𝒖 ) = ∑ i = 0 n − 1 𝒖 [ i ] X n − 1 − i . \displaystyle=\sum_{i=0}^{n-1}\bm{u}[i]X^{i},\text{ and }\operatorname{poly}^{*}(\bm{u})=\sum_{i=0}^{n-1}\bm{u}[i]X^{n-1-i}. This allows us to map between vectors and polynomial. Accordingly, we also define coeff : 𝔽 ⁡ [ X ] / ( X n + 1 ) → 𝔽 n \mathrm{coeff}:\mathbbm{F}[X]/(X^{n+1})\to\mathbbm{F}^{n} as the map converting polynomials back to vectors: coeff ⁡ ( 𝒖 ⁡ ( X ) ) = 𝒖 \mathrm{coeff}(\bm{u}(X))=\bm{u} with 𝒖 ⁡ [ i ] \bm{u}[i] defined as the coefficient in 𝒖 ⁡ ( X ) \bm{u}(X) at degree i i .

These operations allow us to interpret the convolution of vectors in terms of polynomial multiplication [ Heideman and Burrus, 1988 ] . More specifically, we have 𝒖 ∗ 𝒗 \displaystyle\bm{u}\ast\bm{v} = coeff ⁡ ( 𝒖 ⁡ ( X ) ⋅ 𝒗 ⁡ ( X ) mod X n ) , and \displaystyle=\mathrm{coeff}\left(\bm{u}(X)\cdot\bm{v}(X)\mod{X^{n}}\right),\text{ and} 𝒖 ⊛ 𝒗 \displaystyle\bm{u}\circledast\bm{v} = coeff ⁡ ( 𝒖 ⁡ ( X ) ⋅ 𝒗 ⁡ ( X ) mod X n − 1 ) . \displaystyle=\mathrm{coeff}\left(\bm{u}(X)\cdot\bm{v}(X)\mod{X^{n}-1}\right). We can similarly interpret the Hadamard product of vectors 𝒖 ⊙ 𝒗 \bm{u}\odot\bm{v} as the Hadamard product of polynomials 𝒖 ⁡ ( X ) ⊙ 𝒗 ⁡ ( X ) \bm{u}(X)\odot\bm{v}(X) : 𝒖 ⊙ 𝒗 = coeff ⁡ ( 𝒖 ⁡ ( X ) ⊙ 𝒗 ⁡ ( X ) ) = coeff ⁡ ( ∑ i = 0 n − 1 ( 𝒖 ⁡ [ i ] ⋅ 𝒗 ⁡ [ i ] ) ⋅ X i ) . \bm{u}\odot\bm{v}=\mathrm{coeff}\left(\bm{u}(X)\odot\bm{v}(X)\right)=\mathrm{coeff}\left(\sum_{i=0}^{n-1}(\bm{u}[i]\cdot\bm{v}[i])\cdot X^{i}\right).

##### Arithmetic Circuit Notation.

We briefly introduce the notation of arithmetic circuits [ Volkovich, 2016 ] , the focus of Section H.5 . An arithmetic circuit 𝒞 \mathcal{C} with variables X ≜ { x 1 , x 2 , … , x n } X\triangleq\{x_{1},x_{2},\ldots,x_{n}\} over a field 𝔽 \mathbbm{F} is interpreted as a directed acyclic graph, where the input nodes are labelled by either the variables from X X or constants from 𝔽 \mathbbm{F} and the internal nodes are labelled by + + or × \times 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, or ‘wires’ which will be intersected by a horizontal ‘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:

###### Definition H.1 .

An arithmetic circuit 𝒞 \mathcal{C} is an ( n , s , Δ , w ) (n,s,\Delta,w) -circuit if 𝒞 \mathcal{C} is an n n -variate arithmetic circuit of size s s and of depth at most Δ \Delta , and width w w .

##### Model Notation.

Now we introduce the notation we will be using for defining layers. In what follows, we denote 𝒖 ∈ ℝ N × d \bm{u}\in\mathbbm{R}^{N\times d} as the model input; N N as the sequence length; L L as the number of stacked layers, indexed by ℓ \ell ; and d d as the input (embedding) dimension.

#### H.1.2 Summary of the Results

The outline of our results are as follows: In Section H.2 we introduce gated convolution models and define H3, Hyena, RWKV, RetNet, and BaseConv . In Section H.3 we introduce a set of primitive operations that BaseConv can implement. We use these as tools in the subsequent proofs. Then, in Section H.5 , we show that a general arithmetic circuit of size s s and degree at most Δ \Delta 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 a specific proposal (e.g. Hyena).

Next, we study the representational power that gated convolutions and attention use to solve Mqar . In Section H.7 , we derive a BaseConv model inspired from dyadic intervals to show the dimensionality for gated convolutions (with input-independent filters) solve Mqar . Next, we analyze the dimensionality for a BaseConv model with data-dependent kernels to solve Mqar in Section H.8 .

### H.2 Gated Attention-Free 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 ′ ) \left(N,L,d,N^{\prime},d^{\prime}\right) .

###### Definition H.2 .

An ( N , L , d , N ′ , d ′ ) − Gated Convolution Model \left(N,L,d,N^{\prime},d^{\prime}\right)-\text{\text{Gated Convolution Model}} is a stacked sequence to sequence model with L L layers such that:

1. Input and output are N × d N\times d 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 ′ × d ′ N^{\prime}\times d^{\prime} matrices and output N ′ × d ′ N^{\prime}\times d^{\prime} matrices. We refer to the tuple ( N ′ , d ′ ) (N^{\prime},d^{\prime}) as the inner dimension of the model.

We define the Hyena, RWKV, RetNet, and BaseConv layers to make step 2 more concrete. We also assume that the input 𝒖 ∈ ℝ N × d \bm{u}\in\mathbbm{R}^{N\times d} is embedded into 𝒖 ′ ∈ ℝ N ′ × d ′ \bm{u}^{\prime}\in\mathbbm{R}^{N^{\prime}\times d^{\prime}} such that

𝒖 ′ ​ [ n , t ] = { 𝒖 ⁡ [ n , t ] if ​ n < N , t < d 0 otherwise. \bm{u}^{\prime}[n,t]=\begin{cases}\bm{u}[n,t]\ \ \text{ if }n<N,\ t<d\ \\ 0\ \ \text{ otherwise. }\end{cases}

The output from the last layer 𝒛 ∈ ℝ N ′ × d ′ \bm{z}\in\mathbbm{R}^{N^{\prime}\times d^{\prime}} is transformed into output 𝒚 ∈ R N × d \bm{y}\in R^{N\times d} by extracting the top left N × d N\times d entries in 𝒛 \bm{z} .

Next, we define the class of weight matrices that we will use in the linear projections in the models:

###### Definition H.3 .

The linear projection 𝙻𝚒𝚗𝚎𝚊𝚛 m , m {\tt Linear}_{m,m} has its matrix representation as the weight matrix 10 10 10 I.e. 𝙻𝚒𝚗𝚎𝚊𝚛 m , m ​ ( 𝒛 ) = 𝑾 ⋅ 𝒛 {\tt Linear}_{m,m}(\bm{z})=\bm{W}\cdot\bm{z} . 𝐖 ∈ ℝ m × m \mathbf{W}\in\mathbb{R}^{m\times m} taken to be a K-matrix for 𝐖 ∈ ( ℬ ​ ℬ ∗ ) poly- ​ log ⁡ m poly- ​ log ⁡ m \mathbf{W}\in(\mathcal{B}\mathcal{B}^{*})^{\text{poly-}\log{m}}_{\text{poly-}\log{m}} ( Definition H.17 ). Consequently, each matrix 𝐖 \mathbf{W} has 𝒪 ~ ​ ( m ) \tilde{\mathcal{O}}(m) parameters and runtime for matrix vector multiplication, and allows us to represent general linear transformations with low-depth linear arithmetic circuits [ Dao et al., 2020 ] . We will also need linear maps 𝙻𝚒𝚗𝚎𝚊𝚛 m , n {\tt Linear}_{m,n} for m < n m<n , where each take the corresponding square matrices from 𝙻𝚒𝚗𝚎𝚊𝚛 m , m {\tt Linear}_{m,m} and note that such matrices has 𝒪 ~ ​ ( n ) \tilde{\mathcal{O}}(n) parameters and runtime for matrix vector multiplication.

###### Remark H.4 .

We note that the weight matrix 𝐖 ∈ ℝ d × d \mathbf{W}\in\mathbb{R}^{d\times d} above is taken to be a dense matrix in the experiments. However, for our theoretical results, we 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 et al. [2023] .

#### H.2.1 The Hyena Layer

We will now outline the Hyena layer [ Poli et al., 2023b ] . Hyena takes a sequence 𝒖 ∈ ℝ N × d \bm{u}\in\mathbbm{R}^{N\times d} as input and produces L + 1 L+1 projections p 1 , … , p L , v p^{1},\ldots,p^{L},v by passing 𝒚 \bm{y} though a linear layer and applying a short convolution afterwards. The algorithm then recursively performs a point-wise multiplication of the projection with the convolution of the filter h l h^{l} with the previous output. We summarize this process in Algorithm 3 .

###### Remark H.5 .

In Algorithm 3 , L 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 recursive step does not make a difference.

Henceforth, we will refer to a model consisting of L L Hyena layers is a gated convolution model with associated tuple ( N , L , d , N , ( L + 1 ) ​ d ) \left(N,L,d,N,(L+1)d\right) as ( N , L , d , N , ( L + 1 ) ​ d ) − Hyena \left(N,L,d,N,(L+1)d\right)-\text{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 et al., 2021 ] ), it can also be viewed as a convolution. Here, we present the convolutional view of RWKV. We will see that it is closely related to the Hyena layer.

RWKV takes a sequence 𝒖 ∈ ℝ N × d \bm{u}\in\mathbbm{R}^{N\times d} as input and applies a short convolution. Next, it produces three projections q , k , v q,k,v by passing 𝒖 \bm{u} through a linear layer. Then, it performs a convolution sandwiched by element-wise multiplication. We summarize this process in Algorithm 5 .

In practice, the short convolution filter 𝒉 s \bm{h}_{s} is restricted to a length two filter with 𝒉 s ​ [ 0 ] = μ \bm{h}_{s}[0]=\mu and 𝒉 s ​ [ 1 ] = 1 − μ \bm{h}_{s}[1]=1-\mu , where μ ∈ [ 0 , 1 ] \mu\in[0,1] is learned parameter. In the RWKV paper, the short convolution is referred to as a “time shift”.

In practice, the long convolution filter 𝒉 \bm{h} is also restricted to 𝒉 [ : , t ] = e w ⁡ ( t − 1 ) \bm{h}[:,t]=e^{w(t-1)} , where w ∈ ℝ w\in\mathbb{R} is a learnable parameter that controls how quickly the magnitudes of the filter decreases as t t grows.

To summarize, the differences between Hyena and RWKV are: (1) Hyena applies the short convolution after the linear projection whereas RWKV applies it before , (2) RWKV includes non-linear activations (sigmoid and softmax) while Hyena does not, (3) RWKV and Hyena use different parameterizations for the convolutional filters, and (4) Hyena recursively performs a point-wise multiplication of the projections with the convolution filter whereas RWKV performs this operation only once (though, in practice, Hyena uses a recurrence depth of 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 𝒖 ∈ ℝ N × d \bm{u}\in\mathbbm{R}^{N\times d} and project it using learnable weight matrices. We then compute the states 𝒛 n \bm{z}^{n} recurrently as in line 6 and the output 𝙾𝚞𝚝 [ n , : ] {\tt Out}[n,:] as in line 7 for each n ∈ [ N ] n\in[N] (see Algorithm 6 below).

#### H.2.4 BaseConv

Finally, we introduce the BaseConv here as follows: 𝐘 = ( 𝒖 ​ 𝐖 + 𝒃 1 ) ⊙ ( 𝒖 ∗ 𝒉 + 𝒃 2 ) , \mathbf{Y}=(\bm{u}\mathbf{W}+\bm{b}_{1})\odot(\bm{u}\ast\bm{h}+\bm{b}_{2}), (4) with input 𝒖 ∈ ℝ N ′ × d ′ \bm{u}\in\mathbbm{R}^{N^{\prime}\times d^{\prime}} , weight matrix 𝐖 ∈ ℝ d ′ × d ′ \mathbf{W}\in\mathbb{R}^{d^{\prime}\times d^{\prime}} and bias matrices 𝒃 i ∈ ℝ N ′ × d ′ \bm{b}_{i}\in\mathbb{R}^{N^{\prime}\times d^{\prime}} defining linear projections of the input sequence, and 𝒉 ∈ ℝ N ′ × d ′ \bm{h}\in\mathbbm{R}^{N^{\prime}\times d^{\prime}} is the a set of the d ′ d^{\prime} mixed length filters.

The corresponding pseudocode for BaseConv is as follows:

Note that the convolution in Algorithm 7 is not limited to causal convolution and allows circular convolution as well. For simplicity, we use ∗ \ast here and will disambiguate when required. We will start by specifying the parameter count and complexity of a single layer of BaseConv below.

###### Proposition H.6 .

A single layer of BaseConv requires O ~ ​ ( N ′ ​ d ′ ) \tilde{O}(N^{\prime}d^{\prime}) parameters and runtime.

###### Proof.

We know the parameters and runtime for the linear part of BaseConv via Definition H.3 is 𝒪 ~ ​ ( N ′ ​ d ′ ) \tilde{\mathcal{O}}(N^{\prime}d^{\prime}) . Further, each convolution operation requires 𝒪 ⁡ ( N ′ ) \mathcal{O}(N^{\prime}) parameters and 𝒪 ⁡ ( N ′ ​ log ⁡ N ′ ) \mathcal{O}(N^{\prime}\log{N^{\prime}}) runtime, and we employ d ′ d^{\prime} such operations. Finally the Hadamard product step takes O ⁡ ( n ​ d ) O(nd) time. ∎

Similar to Hyena, we will refer to a model consisting of L L BaseConv layers as ( N , L , d , N ′ , d ′ ) − BaseConv \left(N,L,d,N^{\prime},d^{\prime}\right)-\text{{BaseConv}} . In our experiments, we extend BaseConv by adding an MLP after Algorithm 7 . For simplicity we will denote BaseConv ​ ( 𝒖 , 𝒉 , 𝒃 1 , 𝒃 2 ) \textsc{BaseConv}\left({\bm{u},\bm{h},\bm{b}_{1},\bm{b}_{2}}\right) as BaseConv ​ ( 𝒖 , 𝒉 ) \textsc{BaseConv}\left({\bm{u},\bm{h}}\right) when 𝒃 1 = 𝒃 2 = 𝟎 \bm{b}_{1}=\bm{b}_{2}={\bf 0} .

We will now show that there exists a BaseConv model that can emulate each of the basic operations in Algorithm 7 .

###### Lemma H.7 .

The functions Linear d , d ​ ( 𝐮 ) \texttt{Linear}_{d,d}\left(\bm{u}\right) ( Definition H.3 ), with d , d ′ d,d^{\prime} defined as in Algorithm 7 , convolution with filter 𝐡 ∈ ℝ N × d \bm{h}\in\mathbbm{R}^{N\times d} , and element-wise gating can be computed with Algorithm 7 via a ( N , 1 , d , N , d ′ ) − BaseConv \left(N,1,d,N,d^{\prime}\right)-\text{{BaseConv}} .

###### Proof.

For each operation from Definition H.2 and Algorithm 7 : 1. For any input 𝒖 ∈ ℝ N ′ × d ′ \bm{u}\in\mathbbm{R}^{N^{\prime}\times d^{\prime}} , Linear d , d ′ ​ ( 𝒖 ) \texttt{Linear}_{d,d^{\prime}}\left(\bm{u}\right) with matrix representation 𝑾 ∈ ℝ N ′ × d ′ \bm{W}\in\mathbbm{R}^{N^{\prime}\times d^{\prime}} can be performed by a single BaseConv layer computing BaseConv ​ ( 𝒚 , 𝑾 , 𝒉 , 𝒃 1 , 𝒃 2 ) \textsc{BaseConv}\left({\bm{y},\bm{W},\bm{h},\bm{b}_{1},\bm{b}_{2}}\right) with 𝒃 1 \bm{b}_{1} and 𝒃 2 \bm{b}_{2} being the matrix of all 0 0 s and all 1 1 s, respectively while and the convolution with the zero filter. That is, we have 𝐘 = ( 𝒖 ​ 𝐖 + 𝟎 N ′ × d ′ ) ⊙ ( 𝒖 ∗ 𝟎 N ′ × d ′ + 𝟏 N ′ × d ′ ) = ( 𝒖 ​ 𝐖 ) ⊙ 𝟏 N ′ × d ′ = 𝒖 ​ 𝐖 = Linear d , d ′ ​ ( 𝒖 ) . \mathbf{Y}=(\bm{u}\mathbf{W}+\mathbf{0}^{N^{\prime}\times d^{\prime}})\odot(\bm{u}\ast\bm{0}^{N^{\prime}\times d^{\prime}}+\mathbf{1}^{N^{\prime}\times d^{\prime}})=(\bm{u}\mathbf{W})\odot\mathbf{1}^{N^{\prime}\times d^{\prime}}=\bm{u}\mathbf{W}=\texttt{Linear}_{d,d^{\prime}}\left(\bm{u}\right).

2. For any input 𝒖 ∈ ℝ N × d \bm{u}\in\mathbbm{R}^{N\times d} , convolution with filter 𝒉 ∈ ℝ N × d \bm{h}\in\mathbbm{R}^{N\times d} can be performed by a single BaseConv layer computing BaseConv ​ ( 𝒚 , 𝑾 , 𝒉 , 𝒃 1 , 𝒃 2 ) \textsc{BaseConv}\left({\bm{y},\bm{W},\bm{h},\bm{b}_{1},\bm{b}_{2}}\right) where 𝐖 , 𝒃 2 {\bf W},\bm{b}_{2} are all zeroes, and 𝒃 1 \bm{b}_{1} is the matrix of all 1s so that we get 𝐘 = ( 𝒖 ​ 𝟎 N ′ × d ′ + 𝟏 N ′ × d ′ ) ⊙ ( 𝒖 ∗ 𝒉 + 𝟎 N ′ × d ′ ) = 𝟏 N ′ × d ′ ⊙ ( 𝒖 ∗ 𝒉 ) = 𝒖 ∗ 𝒉 . \mathbf{Y}=(\bm{u}\mathbf{0}^{N^{\prime}\times d^{\prime}}+\mathbf{1}^{N^{\prime}\times d^{\prime}})\odot(\bm{u}\ast\bm{h}+\mathbf{0}^{N^{\prime}\times d^{\prime}})=\mathbf{1}^{N^{\prime}\times d^{\prime}}\odot(\bm{u}\ast\bm{h})=\bm{u}\ast\bm{h}.

3. We may compute element-wise gating between matrices 𝒖 , 𝒗 ∈ ℝ N × d \bm{u},\bm{v}\in\mathbbm{R}^{N\times d} , where 𝒗 \bm{v} is some fixed factor, with a single layer computing BaseConv ( 𝒚 , , ) 𝟎 N ′ × d ′ , 𝒆 0 , 𝒗 , 𝟎 N ′ × d ′ \textsc{BaseConv}\left({\bm{y},,}\right)\mathbf{0}^{N^{\prime}\times d^{\prime}},{\bm{e}_{0},\bm{v},\mathbf{0}^{N^{\prime}\times d^{\prime}}} where 𝒆 1 \bm{e}_{1} is the identity filter, respectively, by Definition H.2 . 𝐘 = ( 𝒖 ​ 𝟎 N ′ × d ′ + 𝒗 ) ⊙ ( 𝒖 ∗ 𝒆 0 + 𝟎 N ′ × d ′ ) = 𝒗 ⊙ 𝒖 . \mathbf{Y}=(\bm{u}\mathbf{0}^{N^{\prime}\times d^{\prime}}+\bm{v})\odot(\bm{u}\ast\bm{e}_{0}+\mathbf{0}^{N^{\prime}\times d^{\prime}})=\bm{v}\odot\bm{u}.

∎

### H.3 Primitives

In this section, we will establish some additional basic primitives that we expect a gated convolution model to emulate: shift , remember \texttt{shift},\texttt{remember} and add . We specify them below: 1. Shift an sequential input of length N N up or down by s s entries:

shift_up( 𝒚 , s \bm{y},s ), shift_down( 𝒚 , s \bm{y},s ) • Input: 𝒚 ∈ ℝ N × d \bm{y}\in\mathbbm{R}^{N\times d} , s ≥ 0 s\geq 0 • Output: 𝒛 ∈ ℝ N × d \bm{z}\in\mathbbm{R}^{N\times d} where 𝒛 + = shift_down ​ ( 𝒚 , s ) \bm{z}^{+}=\texttt{shift\_down}(\bm{y},s) and 𝒛 − = shift_up ​ ( 𝒚 , s ) \bm{z}^{-}=\texttt{shift\_up}(\bm{y},s) 𝒚 ≡ ( ← 𝒚 0 → ← 𝒚 i − 1 → ← 𝒚 i → ← 𝒚 N − 1 → ) 𝒛 + ≡ ( ← 𝟎 → ← 𝟎 → ← 𝒚 0 → ← 𝒚 N − 1 − s → ) 𝒛 − ≡ ( ← 𝒚 s → ← 𝒚 N − 1 → ← 𝟎 → ← 𝟎 → ) \bm{y}\equiv\begin{pmatrix}\leftarrow\bm{y}_{0}\rightarrow\\ \hline\cr\\ \vdots\\ \hline\cr\\ \leftarrow\bm{y}_{i-1}\rightarrow\\ \hline\cr\\ \leftarrow\bm{y}_{i}\rightarrow\\ \hline\cr\\ \vdots\\ \hline\cr\\ \leftarrow\bm{y}_{N-1}\rightarrow\end{pmatrix}\qquad\bm{z}^{+}\equiv\begin{pmatrix}\leftarrow\bm{0}\rightarrow\\ \hline\cr\\ \vdots\\ \hline\cr\\ \leftarrow\bm{0}\rightarrow\\ \hline\cr\\ \leftarrow\bm{y}_{0}\rightarrow\\ \hline\cr\\ \vdots\\ \hline\cr\\ \leftarrow\bm{y}_{N-1-s}\rightarrow\end{pmatrix}\qquad\bm{z}^{-}\equiv\begin{pmatrix}\leftarrow\bm{y}_{s}\rightarrow\\ \hline\cr\\ \vdots\\ \hline\cr\\ \leftarrow\bm{y}_{N-1}\rightarrow\\ \hline\cr\\ \leftarrow\bm{0}\rightarrow\\ \hline\cr\\ \vdots\\ \hline\cr\\ \leftarrow\bm{0}\rightarrow\end{pmatrix}

2. Add a sequence 𝒙 ∈ ℝ n × d \bm{x}\in\mathbbm{R}^{n\times d} to a running sum 𝑺 ∈ ℝ n × d \bm{S}\in\mathbbm{R}^{n\times d} for some 2 ​ n ≤ N ′ 2n\leq N^{\prime} with both 𝒙 \bm{x} and 𝑺 \bm{S} contained as subvectors in 𝒚 ∈ ℝ N × d \bm{y}\in\mathbbm{R}^{N\times d} . add n ( 𝒚 : 𝒙 , 𝑺 ) \texttt{add}_{n}(\bm{y}:\bm{x},\bm{S}) : • Input: sequence 𝒚 \bm{y} containing 𝒙 , 𝑺 ∈ ℝ n × d \bm{x},\bm{S}\in\mathbbm{R}^{n\times d} for 2 ​ n ≤ N 2n\leq N such that 𝒚 [ 0 : n − 1 ] ≡ 𝒙 , 𝒚 [ n : 2 n − 1 ] ≡ 𝑺 \bm{y}[0:n-1]\equiv\bm{x},\bm{y}[n:2n-1]\equiv\bm{S} and 𝒚 [ 2 n : N − 1 ] ≡ 𝟎 N − 2 ​ n \bm{y}[2n:N-1]\equiv\bm{0}^{N-2n} . • Output: 𝒛 ∈ ℝ N × d \bm{z}\in\mathbbm{R}^{N\times d} containing the sum 𝒚 + 𝑺 \bm{y}+\bm{S} such that 𝒚 [ 0 : n − 1 ] ≡ 𝟏 n , 𝒚 [ n : 2 n − 1 ] ≡ 𝑺 + 𝒙 \bm{y}[0:n-1]\equiv\bm{1}^{n},\bm{y}[n:2n-1]\equiv\bm{S}+\bm{x} and 𝒛 [ 2 n : N − 1 ] ≡ 𝟎 N − 2 ​ n \bm{z}[2n:N-1]\equiv\bm{0}^{N-2n} . 𝒚 ≡ ( ← 𝒙 → ← 𝑺 → ⟵ 𝟎 ⟶ ← 𝟎 → ) 𝒛 ≡ ( ← 𝟏 n → ← 𝑺 + 𝒙 → ⟵ 𝟎 ⟶ ← 𝟎 → ) \bm{y}\equiv\begin{pmatrix}\leftarrow\bm{x}\rightarrow\\ \hline\cr\\ \leftarrow\bm{S}\rightarrow\\ \hline\cr\\ \longleftarrow\bm{0}\longrightarrow\\ \hline\cr\\ \vdots\\ \hline\cr\\ \leftarrow\bm{0}\rightarrow\\ \end{pmatrix}\qquad\bm{z}\equiv\begin{pmatrix}\leftarrow\bm{1}^{n}\rightarrow\\ \hline\cr\\ \leftarrow\bm{S}+\bm{x}\rightarrow\\ \hline\cr\\ \longleftarrow\bm{0}\longrightarrow\\ \hline\cr\\ \vdots\\ \hline\cr\\ \leftarrow\bm{0}\rightarrow\\ \end{pmatrix}

3. Remember 𝒗 ∈ ℝ m × d \bm{v}\in\mathbbm{R}^{m\times d} as part of a sequence of input 𝒚 ∈ ℝ N × d \bm{y}\in\mathbbm{R}^{N\times d} while performing gated convolution only on 𝒙 ∈ ℝ n × d \bm{x}\in\mathbbm{R}^{n\times d} for some m , n ≤ N × d m,n\leq N\times d . remember n , m , s , t ( 𝒚 : 𝒙 , 𝒗 , 𝒉 , 𝒑 ) \texttt{remember}_{n,m,s,t}(\bm{y:x,v,h,p}) : • Input: sequence 𝒚 ∈ ℝ N × d \bm{y}\in\mathbbm{R}^{N\times d} containing 𝒙 ∈ ℝ n × d , 𝒗 ∈ ℝ m × d \bm{x}\in\mathbbm{R}^{n\times d},\bm{v}\in\mathbbm{R}^{m\times d} , and modifiers 𝒑 , 𝒉 ∈ ℝ n × d \bm{p,h}\in\mathbbm{R}^{n\times d} such that 𝒚 [ 0 : n − 1 ] ≡ 𝒙 , 𝒚 [ n + s : n + s + m − 1 ] ≡ 𝒗 \bm{y}[0:n-1]\equiv\bm{x},\bm{y}[n+s:n+s+m-1]\equiv\bm{v} and 𝒚 ⁡ [ i ] = 0 \bm{y}[i]=0 otherwise with 𝒙 ∗ 𝒉 ∈ ℝ ( n + s ) × d \bm{x}\ast\bm{h}\in\mathbbm{R}^{(n+s)\times d} and 𝒗 ∗ 𝒉 ∈ ℝ ( m + t ) × d \bm{v}\ast\bm{h}\in\mathbbm{R}^{(m+t)\times d} . • Output: 𝒛 ∈ ℝ N × d \bm{z}\in\mathbbm{R}^{N\times d} containing ( 𝒑 ⊙ ( 𝒙 ∗ 𝒉 ) ) ∈ ℝ ( n + s ) × d \left(\bm{p}\odot(\bm{x}\ast\bm{h})\right)\in\mathbbm{R}^{(n+s)\times d} , such that: 𝒚 ≡ ( ← 𝒙 → 𝟎 s ← 𝒗 → 𝟎 ) 𝒛 ≡ ( ← 𝒑 ⊙ ( 𝒙 ∗ 𝒉 ) → ← 𝒗 → 𝟎 ) \bm{y}\equiv\begin{pmatrix}\leftarrow\bm{x}\rightarrow\\ \hline\cr\\ \bm{0}^{s}\\ \hline\cr\\ \leftarrow\bm{v}\rightarrow\\ \hline\cr\\ \vdots\\ \hline\cr\\ \bm{0}\end{pmatrix}\qquad\qquad\bm{z}\equiv\begin{pmatrix}\leftarrow\bm{p}\odot(\bm{x}\ast\bm{h})\rightarrow\\ \hline\cr\\ \leftarrow\bm{v}\rightarrow\\ \hline\cr\\ \vdots\\ \hline\cr\\ \bm{0}\end{pmatrix}

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 \left(N,L,d,N^{\prime},d^{\prime}\right)-\text{{BaseConv}} model with a small constant L L .

###### Proposition H.8 (The Shift Primitive) .

For any 𝐲 ∈ ℝ N × d \bm{y}\in\mathbbm{R}^{N\times d} , there exist ( N , 1 , d , N , d ) − BaseConv \left(N,1,d,N,d\right)-\text{{BaseConv}} and ( N , 3 , d , N , d ) − BaseConv \left(N,3,d,N,d\right)-\text{{BaseConv}} that computes shift_down ​ ( 𝐲 , s ) \texttt{shift\_down}({\bm{y}},{s}) and shift_up ​ ( 𝐲 , s ) \texttt{shift\_up}({\bm{y}},{s}) for any s ≤ N s\leq N .

###### Proof.

Define the following kernel dependent on s ≤ N s\leq N 𝒉 s [ n , : ] ≡ { 𝟏 d if ​ n = s + 1 𝟎 d otherwise . \bm{h}_{s}[n,:]\equiv\begin{cases}\bm{1}^{d}&\text{if }n=s+1\\ \bm{0}^{d}&\text{otherwise}.\end{cases} We now deal with the down and up shifts separately: 1. We define 𝐖 := 𝟎 N × d , 𝒃 1 := 𝟏 N × d , 𝒃 2 := 𝟎 N × d {\bf W}:=\mathbf{0}^{N\times d},\bm{b}_{1}:=\mathbf{1}^{N\times d},\bm{b}_{2}:=\mathbf{0}^{N\times d} . Then, for input 𝒚 ∈ ℝ N × d \bm{y}\in\mathbbm{R}^{N\times d} , BaseConv ​ ( 𝒚 , 𝟎 N × d , 𝒉 s , 𝟏 N × d , 𝟎 N × d ) \textsc{BaseConv}\left({\bm{y},\mathbf{0}^{N\times d},\bm{h}_{s},\mathbf{1}^{N\times d},\mathbf{0}^{N\times d}}\right) for BaseConv in Algorithm 7 is given by equation 4 as 𝐘 ≡ 𝒚 ∗ 𝒉 s . \displaystyle{\bf Y}\equiv\bm{y}\ast\bm{h}_{s}. Now, to perform shift_down ( 𝒚 , s ) (\bm{y},s) , we note that 𝐘 [ : , j ] = 𝒚 [ : , j ] ∗ 𝒉 s [ : , j ] \displaystyle{\bf Y}[:,j]=\bm{y}[:,j]\ast\bm{h}_{s}[:,j] = coeff ( 𝒚 [ : , j ] ( X ) ⋅ 𝒉 s [ : , j ] ( X ) ) \displaystyle=\mathrm{coeff}(\bm{y}[:,j](X)\cdot\bm{h}_{s}[:,j](X)) = coeff ⁡ ( ( ∑ i = 0 N − 1 𝒚 ⁡ [ i , j ] ⋅ X i ) ⋅ X s mod X N ) \displaystyle=\mathrm{coeff}\left(\left(\sum_{i=0}^{N-1}\bm{y}[i,j]\cdot X^{i}\right)\cdot X^{s}\mod{X^{N}}\right) = coeff ⁡ ( ∑ i = 0 N − 1 𝒚 ⁡ [ i , j ] ⋅ X i + s mod X N ) \displaystyle=\mathrm{coeff}\left(\sum_{i=0}^{N-1}\bm{y}[i,j]\cdot X^{i+s}\mod{X^{N}}\right) = coeff ⁡ ( ∑ i = s N − 1 + s 𝒚 ⁡ [ i − s , j ] ⋅ X i mod X N ) \displaystyle=\mathrm{coeff}\left(\sum_{i=s}^{N-1+s}\bm{y}[i-s,j]\cdot X^{i}\mod{X^{N}}\right) = coeff ⁡ ( ∑ i = s N − 1 𝒚 ⁡ [ i − s , j ] ⋅ X i ) , \displaystyle=\mathrm{coeff}\left(\sum_{i=s}^{N-1}\bm{y}[i-s,j]\cdot X^{i}\right), which implies that we exactly get what is specified in the output.

2. We again define 𝐖 := 𝟎 N × d , 𝒃 1 := 𝟏 N × d , 𝒃 2 := 𝟎 N × d {\bf W}:=\mathbf{0}^{N\times d},\bm{b}_{1}:=\mathbf{1}^{N\times d},\bm{b}_{2}:=\mathbf{0}^{N\times d} . Then, for input 𝒚 ∈ ℝ N × d \bm{y}\in\mathbbm{R}^{N\times d} , BaseConv ​ ( 𝒚 , 𝟎 N × d , 𝒆 0 , 𝟏 N × d , 𝟎 N × d ) \textsc{BaseConv}\left({\bm{y},\mathbf{0}^{N\times d},\bm{e}_{0},\mathbf{1}^{N\times d},\mathbf{0}^{N\times d}}\right) for BaseConv in Algorithm 7 is given in equation 4 as 𝐘 0 ≡ 𝒚 ⊛ 𝒆 0 . \displaystyle{\bf Y}_{0}\equiv\bm{y}\circledast\bm{e}_{0}. Now, to perform shift_up ( 𝒚 , s ) (\bm{y},s) , as before, we first apply the circular convolution to reverse the input 𝐘 0 [ : , j ] = 𝒚 [ : , j ] ⊛ 𝒆 0 \displaystyle{\bf Y}_{0}[:,j]=\bm{y}[:,j]\circledast\bm{e}_{0} = coeff ⁡ ( ∑ i = 0 N − 1 𝒚 ⁡ [ N − 1 − i , j ] ⋅ X i ) , \displaystyle=\mathrm{coeff}\left(\sum_{i=0}^{N-1}\bm{y}[N-1-i,j]\cdot X^{i}\right), We then apply 𝒀 1 ≡ shift_down ​ ( 𝒀 0 , s ) \bm{Y}_{1}\equiv\texttt{shift\_down}(\bm{Y}_{0},s) to get 𝒀 1 [ : , j ] \displaystyle\bm{Y}_{1}[:,j] ≡ coeff ⁡ ( ∑ i = s N − 1 𝒀 0 ​ [ N − 1 − ( i − s ) , j ] ⋅ X i ) , \displaystyle\equiv\mathrm{coeff}\left(\sum_{i=s}^{N-1}\bm{Y}_{0}[N-1-(i-s),j]\cdot X^{i}\right), ≡ coeff ⁡ ( ∑ i = s N − 1 𝒀 0 ​ [ N − 1 − i + s , j ] ⋅ X i ) . \displaystyle\equiv\mathrm{coeff}\left(\sum_{i=s}^{N-1}\bm{Y}_{0}[N-1-i+s,j]\cdot X^{i}\right). Finally, we apply another circular convolution with the identity filter to replace N − 1 − i N-1-i with i i to get 𝐘 2 [ : , j ] = 𝒀 1 [ : , j ] ⊛ 𝒆 0 \displaystyle{\bf Y}_{2}[:,j]=\bm{Y}_{1}[:,j]\circledast\bm{e}_{0} = coeff ⁡ ( ∑ i = 0 N − 1 𝒚 ⁡ [ i + s , j ] ⋅ X i ) , \displaystyle=\mathrm{coeff}\left(\sum_{i=0}^{N-1}\bm{y}[i+s,j]\cdot X^{i}\right), Here, we note that we can compute both of these primitives in one and three layers, respectively (see Lemma H.11 ).

∎

Now, we present a BaseConv model with two layers that implements the add n ( 𝒚 : 𝒙 , 𝑺 ) \texttt{add}_{n}({\bm{y}:\bm{x}},{\bm{S}}) , the purpose of which is to add some window of computation 𝒙 \bm{x} to a running sum 𝑺 \bm{S} .

###### Proposition H.9 (The Running Sum Primitive) .

For any 𝐱 , 𝐒 ∈ ℝ n × d \bm{x,S}\in\mathbbm{R}^{n\times d} contained in some 𝐲 ∈ ℝ N × d \bm{y}\in\mathbbm{R}^{N\times d} , there exists a ( N , 2 , d , N , d ) − BaseConv \left(N,2,d,N,d\right)-\text{{BaseConv}} that computes add n ( 𝐲 : 𝐱 , 𝐒 ) \texttt{add}_{n}({\bm{y}:\bm{x}},{\bm{S}}) for BaseConv as in Algorithm 7 .

###### Proof.

We will show this for d ′ = 1 d^{\prime}=1 and the general case follows as we will explain at the end. We now specify the two layers that we use 𝒛 1 \displaystyle\bm{z}^{1} ≡ BaseConv ​ ( 𝒚 , 𝟎 N × 1 , 𝒉 1 , 𝒃 1 1 , 𝟎 N × 1 ) ≡ 𝒃 1 1 ⊙ ( 𝒉 1 ∗ 𝒚 ) \displaystyle\equiv\textsc{BaseConv}\left({\bm{y},\mathbf{0}^{N\times 1},\bm{h}^{1},\bm{b}^{1}_{1},\mathbf{0}^{N\times 1}}\right)\equiv\bm{b}^{1}_{1}\odot\left(\bm{h}^{1}\ast\bm{y}\right) 𝒛 \displaystyle\bm{z} ≡ BaseConv ​ ( 𝒛 1 , 𝟎 N × 1 , 𝒉 2 , 𝒃 1 2 , 𝒃 1 2 ) ≡ 𝒃 1 2 ⊙ ( 𝒉 2 ∗ 𝒚 + 𝒃 2 2 ) , \displaystyle\equiv\textsc{BaseConv}\left({\bm{z}^{1},\mathbf{0}^{N\times 1},\bm{h}^{2},\bm{b}^{2}_{1},\bm{b}^{2}_{1}}\right)\equiv\bm{b}^{2}_{1}\odot\left(\bm{h}^{2}\ast\bm{y}+\bm{b}^{2}_{2}\right), where we will specify the kernels as we go along. Let us start by defining the kernel and the bias for the first layer as 𝒉 1 ≡ ( 𝒆 0 𝒆 0 𝟎 n ⋯ 𝟎 n ) , 𝒃 1 ≡ ( 𝟎 n 𝟏 n 𝟎 n ⋯ 𝟎 n ) . \bm{h}^{1}\equiv\begin{pmatrix}\bm{e}_{0}\\ \hline\cr\\ \bm{e}_{0}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}^{n}\end{pmatrix},\quad\quad\bm{b}_{1}\equiv\begin{pmatrix}\bm{0}^{n}\\ \hline\cr\\ \bm{1}^{n}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}^{n}\end{pmatrix}. Let us first compute 𝒉 1 ∗ 𝒚 \bm{h}^{1}\ast\bm{y} as follows: 𝒉 1 ​ ( X ) ⋅ 𝒚 ​ ( X ) \displaystyle\bm{h}^{1}(X)\cdot\bm{y}(X) = ( X n + 1 ) ⋅ ( 𝑺 ⁡ ( X ) ⋅ X n + 𝒙 ⁡ ( X ) ) \displaystyle=(X^{n}+1)\cdot(\bm{S}(X)\cdot X^{n}+\bm{x}(X)) = 𝑺 ⁡ ( X ) ⋅ X 2 ​ n + ( 𝑺 + 𝒙 ) ​ ( X ) ⋅ X n + 𝒙 ⁡ ( X ) . \displaystyle=\bm{S}(X)\cdot X^{2n}+(\bm{S}+\bm{x})(X)\cdot X^{n}+\bm{x}(X). We then have 𝒛 1 ≡ 𝒃 1 1 ⊙ ( 𝒉 1 ∗ 𝒚 ) ≡ ( 𝟎 n 𝟏 n 𝟎 n ⋯ 𝟎 n ) ⊙ ( 𝒙 𝑺 + 𝒙 𝑺 ⋯ 𝟎 n ) ≡ ( 𝟎 n 𝑺 + 𝒙 𝟎 n ⋯ 𝟎 n ) \bm{z}_{1}\equiv\bm{b}_{1}^{1}\odot\left(\bm{h}^{1}\ast\bm{y}\right)\equiv\begin{pmatrix}\bm{0}^{n}\\ \hline\cr\\ \bm{1}^{n}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}^{n}\end{pmatrix}\odot\begin{pmatrix}\bm{x}\\ \hline\cr\\ \bm{S}+\bm{x}\\ \hline\cr\\ \bm{S}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}^{n}\end{pmatrix}\equiv\begin{pmatrix}\bm{0}^{n}\\ \hline\cr\\ \bm{S}+\bm{x}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}^{n}\end{pmatrix}

##### Resetting for Next Phase.

We now use the next layer to reset for the next phase. Here, we need the first vector to be 𝟏 n \bm{1}^{n} in order to start adding the next vector. We thus use the kernel and the biases 𝒉 2 , 𝒃 1 2 , 𝒃 2 2 \bm{h}^{2},\bm{b}^{2}_{1},\bm{b}^{2}_{2} defined as 𝒉 2 ≡ ( 𝒆 0 𝟎 n 𝟎 n ⋯ 𝟎 n ) , 𝒃 1 2 ≡ ( 𝟏 n 𝟏 n 𝟎 n ⋯ 𝟎 n ) , 𝒃 2 2 ≡ ( 𝟏 n 𝟎 n 𝟎 n ⋯ 𝟎 n ) . \bm{h}^{2}\equiv\begin{pmatrix}\bm{e}_{0}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}^{n}\end{pmatrix},\quad\quad\bm{b}^{2}_{1}\equiv\begin{pmatrix}\bm{1}^{n}\\ \hline\cr\\ \bm{1}^{n}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}^{n}\end{pmatrix},\quad\quad\bm{b}^{2}_{2}\equiv\begin{pmatrix}\bm{1}^{n}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}^{n}\end{pmatrix}. Explicitly, for the second layer, we compute the result of the convolution in terms of polynomials as follows: 𝒉 2 ​ ( X ) ⋅ 𝒛 1 ​ ( X ) = 1 ⋅ ( 𝑺 + 𝒙 ) ​ ( X ) ⋅ X n = ( 𝑺 + 𝒙 ) ​ ( X ) ⋅ X n . \displaystyle\bm{h}^{2}(X)\cdot\bm{z}^{1}(X)=1\cdot(\bm{S}+\bm{x})(X)\cdot X^{n}=(\bm{S}+\bm{x})(X)\cdot X^{n}. Thus, the output for the second layer is given by 𝒛 ≡ 𝒃 1 2 ⊙ ( 𝒉 2 ∗ 𝒛 1 + 𝒃 2 2 ) ≡ ( 𝟏 n 𝟏 n 𝟎 n ⋯ 𝟎 n ) ⊙ ( ( 𝟎 n 𝑺 + 𝒙 𝟎 n ⋯ 𝟎 n ) + ( 𝟏 n 𝟎 n 𝟎 n ⋯ 𝟎 n ) ) ≡ ( 𝟏 n 𝟏 n 𝟎 n ⋯ 𝟎 n ) ⊙ ( 𝟏 n 𝑺 + 𝒙 𝟎 n ⋯ 𝟎 n ) ≡ ( 𝟏 n 𝑺 + 𝒙 𝟎 n ⋯ 𝟎 n ) . \bm{z}\equiv\bm{b}_{1}^{2}\odot\left(\bm{h}^{2}\ast\bm{z}^{1}+\bm{b}^{2}_{2}\right)\equiv\begin{pmatrix}\bm{1}^{n}\\ \hline\cr\\ \bm{1}^{n}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}^{n}\end{pmatrix}\odot\left(\begin{pmatrix}\bm{0}^{n}\\ \hline\cr\\ \bm{S}+\bm{x}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}^{n}\end{pmatrix}+\begin{pmatrix}\bm{1}^{n}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}^{n}\end{pmatrix}\right)\equiv\begin{pmatrix}\bm{1}^{n}\\ \hline\cr\\ \bm{1}^{n}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}^{n}\end{pmatrix}\odot\begin{pmatrix}\bm{1}^{n}\\ \hline\cr\\ \bm{S}+\bm{x}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}^{n}\end{pmatrix}\equiv\begin{pmatrix}\bm{1}^{n}\\ \hline\cr\\ \bm{S}+\bm{x}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}^{n}\end{pmatrix}. Therefore, we have used two BaseConv layers to add 𝒙 \bm{x} to the running sum 𝑺 \bm{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 to d > 1 d>1 . ∎

Next, we show that a 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 𝐱 ∈ ℝ n × d , 𝐯 ∈ ℝ m × d \bm{x}\in\mathbbm{R}^{n\times d},\bm{v}\in\mathbbm{R}^{m\times d} contained in some 𝐲 ∈ ℝ N × d \bm{y}\in\mathbbm{R}^{N\times d} for some n + m + s + t ≤ N n+m+s+t\leq N so that for 𝐡 ∈ ℝ n × d \bm{h}\in\mathbbm{R}^{n\times d} and 𝐩 ∈ ℝ ( n + s ) × d \bm{p}\in\mathbbm{R}^{(n+s)\times d} with 𝐱 ∗ 𝐡 ∈ ℝ ( n + s ) × d \bm{x}\ast\bm{h}\in\mathbbm{R}^{(n+s)\times d} and 𝐯 ∗ 𝐡 ∈ ℝ ( m + t ) × d \bm{v}\ast\bm{h}\in\mathbbm{R}^{(m+t)\times d} , there exists a ( N , 5 , d , N , d ) − BaseConv \left(N,5,d,N,d\right)-\text{{BaseConv}} that computes 𝚛𝚎𝚖𝚎𝚖𝚋𝚎𝚛 ( 𝐲 : 𝐱 , 𝐯 , 𝐡 , 𝐩 ) {\tt remember}(\bm{y}:\bm{x},\bm{v},\bm{h},\bm{p}) for BaseConv as in Algorithm 7 .

###### Proof.

We will again show this for d ′ = 1 d^{\prime}=1 and the general case should follow. We now specify the first two layers that we use 𝒛 1 \displaystyle\bm{z}^{1} ≡ BaseConv ​ ( 𝒚 , 𝟎 N × 1 , 𝒉 1 , 𝒃 1 1 , 𝟎 N × d ) ≡ 𝒃 1 1 ⊙ ( 𝒉 1 ∗ 𝒚 ) \displaystyle\equiv\textsc{BaseConv}\left({\bm{y},\mathbf{0}^{N\times 1},\bm{h}^{1},\bm{b}^{1}_{1},\mathbf{0}^{N\times d}}\right)\equiv\bm{b}^{1}_{1}\odot\left(\bm{h}^{1}\ast\bm{y}\right) 𝒛 2 \displaystyle\bm{z}^{2} ≡ BaseConv ​ ( 𝒛 1 , 𝟎 N × 1 , 𝒉 2 , 𝒃 1 2 , 𝟎 N × d ) ≡ 𝒃 1 2 ⊙ ( 𝒉 2 ∗ 𝒚 ) , \displaystyle\equiv\textsc{BaseConv}\left({\bm{z}^{1},\mathbf{0}^{N\times 1},\bm{h}^{2},\bm{b}^{2}_{1},\mathbf{0}^{N\times d}}\right)\equiv\bm{b}^{2}_{1}\odot\left(\bm{h}^{2}\ast\bm{y}\right), The kernel 𝒉 1 \bm{h}^{1} and the bias 𝒃 1 1 \bm{b}^{1}_{1} for the first layer are then given by 𝒉 1 ≡ ( 𝒉 𝟎 m 𝒆 s + t 𝟎 n 𝟎 n ⋯ 𝟎 n ) , 𝒃 1 1 ≡ ( 𝒑 𝟎 m + t 𝟎 n 𝟎 s 𝟏 m ⋯ 𝟎 n ) . \bm{h}^{1}\equiv\begin{pmatrix}\bm{h}\\ \hline\cr\\ \bm{0}^{m}\\ \hline\cr\\ \bm{e}_{s+t}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}^{n}\\ \end{pmatrix},\quad\quad\bm{b}_{1}^{1}\equiv\begin{pmatrix}\bm{p}\\ \hline\cr\\ \bm{0}^{m+t}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \bm{0}^{s}\\ \hline\cr\\ \bm{1}^{m}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}^{n}\\ \end{pmatrix}. where recall that 𝒙 ∗ 𝒉 ∈ ℝ ( n + s ) × d \bm{x}\ast\bm{h}\in\mathbbm{R}^{(n+s)\times d} and 𝒗 ∗ 𝒉 ∈ ℝ ( m + t ) × d \bm{v}\ast\bm{h}\in\mathbbm{R}^{(m+t)\times d} .

We now want to first specify the result of applying the first kernel: ( 𝒉 1 ∗ 𝒚 ) \displaystyle\left(\bm{h}^{1}\ast\bm{y}\right) = coeff ⁡ ( ( 𝒉 ⁡ ( X ) + X n + m + s + t ) ⋅ ( 𝒗 ⁡ ( X ) ⋅ X n + s + 𝒙 ⁡ ( X ) ) ) \displaystyle=\mathrm{coeff}\left((\bm{h}(X)+X^{n+m+s+t})\cdot\left(\bm{v}(X)\cdot X^{n+s}+\bm{x}(X)\right)\right) = coeff ⁡ ( 𝒉 ∗ 𝒗 ⁡ ( X ) ⋅ X n + s + 𝒉 ∗ 𝒙 ⁡ ( X ) + 𝒗 ⁡ ( X ) ⋅ X 2 ​ n + 2 ​ s + m + t + 𝒙 ⁡ ( X ) ⋅ X n + m + s + t ) \displaystyle=\mathrm{coeff}\left(\bm{h}\ast\bm{v}(X)\cdot X^{n+s}+\bm{h}\ast\bm{x}(X)+\bm{v}(X)\cdot X^{2n+2s+m+t}+\bm{x}(X)\cdot X^{n+m+s+t}\right) We then have 𝒛 1 ≡ 𝒃 1 1 ⊙ ( 𝒉 1 ∗ 𝒚 ) ≡ ( 𝒑 𝟎 m + t 𝟎 n 𝟎 s 𝟏 m ⋯ 𝟎 n ) ⊙ ( 𝒉 ∗ 𝒙 𝒉 ∗ 𝒗 𝒙 𝟎 s 𝒗 ⋯ 𝟎 n ) ≡ ( 𝒑 ⊙ ( 𝒉 ∗ 𝒙 ) 𝟎 m + t 𝟎 n 𝟎 s 𝒗 ⋯ 𝟎 n ) . \bm{z}_{1}\equiv\bm{b}_{1}^{1}\odot\left(\bm{h}^{1}\ast\bm{y}\right)\equiv\begin{pmatrix}\bm{p}\\ \hline\cr\\ \bm{0}^{m+t}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \bm{0}^{s}\\ \hline\cr\\ \bm{1}^{m}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}^{n}\\ \end{pmatrix}\odot\begin{pmatrix}\bm{h}\ast\bm{x}\\ \hline\cr\\ \bm{h}\ast\bm{v}\\ \hline\cr\\ \bm{x}\\ \hline\cr\\ \bm{0}^{s}\\ \hline\cr\\ \bm{v}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}^{n}\\ \end{pmatrix}\equiv\begin{pmatrix}\bm{p}\odot(\bm{h}\ast\bm{x})\\ \hline\cr\\ \bm{0}^{m+t}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \bm{0}^{s}\\ \hline\cr\\ \bm{v}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}^{n}\\ \end{pmatrix}. We now describe the second kernel 𝒉 2 \bm{h}^{2} and the bias matrix 𝒃 1 2 \bm{b}^{2}_{1} as follows: 𝒉 2 ≡ ( 𝒆 0 𝟎 m + t 𝒆 0 𝟎 n + s 𝟎 m ⋯ 𝟎 ) , 𝒃 1 2 ≡ ( 𝟎 n + s 𝟎 m + t 𝟎 n 𝟏 n + s 𝟏 m ⋯ 𝟎 ) \bm{h}^{2}\equiv\begin{pmatrix}\bm{e}_{0}\\ \hline\cr\\ \bm{0}^{m+t}\\ \hline\cr\\ \bm{e}^{0}\\ \hline\cr\\ \bm{0}^{n+s}\\ \hline\cr\\ \bm{0}^{m}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}\end{pmatrix},\quad\quad\bm{b}_{1}^{2}\equiv\begin{pmatrix}\bm{0}^{n+s}\\ \hline\cr\\ \bm{0}^{m+t}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \bm{1}^{n+s}\\ \hline\cr\\ \bm{1}^{m}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}\end{pmatrix} This yields the following convolution computation: 𝒉 2 ⊙ 𝒛 1 \displaystyle\bm{h}^{2}\odot\bm{z}^{1} ≡ coeff ⁡ ( ( X m + n + s + t + 1 ) ⋅ ( 𝒗 ⁡ ( X ) ⋅ X 2 ​ n + 2 ​ s + m + t + ( 𝒑 ⊙ ( 𝒉 ∗ 𝒙 ) ) ​ ( X ) ) ) \displaystyle\equiv\mathrm{coeff}\left(\left(X^{m+n+s+t}+1\right)\cdot\left(\bm{v}(X)\cdot X^{2n+2s+m+t}+\left(\bm{p}\odot(\bm{h}\ast\bm{x})\right)(X)\right)\right) ≡ coeff ⁡ ( 𝒗 ⁡ ( X ) ⋅ X 3 ​ n + 3 ​ s + 2 ​ m + 2 ​ t + 𝒗 ⁡ ( X ) ⋅ X 2 ​ n + 2 ​ s + m + t CLOSE \displaystyle\equiv\mathrm{coeff}(\bm{v}(X)\cdot X^{3n+3s+2m+2t}+\bm{v}(X)\cdot X^{2n+2s+m+t} + ( 𝒑 ⊙ ( 𝒉 ∗ 𝒙 ) ) ( X ) ⋅ X m + n + s + t + ( 𝒑 ⊙ ( 𝒉 ∗ 𝒙 ) ) ( X ) ) \displaystyle\quad+\left(\bm{p}\odot(\bm{h}\ast\bm{x})\right)(X)\cdot X^{m+n+s+t}+\left(\bm{p}\odot(\bm{h}\ast\bm{x})\right)(X)) Thus we have 𝒛 2 ≡ 𝒃 2 1 ⊙ ( 𝒉 2 ∗ 𝒛 1 ) ≡ ( 𝟎 n + s 𝟎 m + t 𝟎 n 𝟏 n + s 𝟏 m ⋯ 𝟎 ) ⊙ ( 𝒑 ⊙ ( 𝒉 ∗ 𝒙 ) 𝟎 m + t 𝟎 n 𝒑 ⊙ ( 𝒉 ∗ 𝒙 ) 𝒗 ⋯ 𝟎 ) ≡ ( 𝟎 n + s 𝟎 m + t 𝟎 n 𝒑 ⊙ ( 𝒉 ∗ 𝒙 ) 𝒗 ⋯ 𝟎 ) \bm{z}^{2}\equiv\bm{b}_{2}^{1}\odot\left(\bm{h}^{2}\ast\bm{z}^{1}\right)\equiv\begin{pmatrix}\bm{0}^{n+s}\\ \hline\cr\\ \bm{0}^{m+t}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \bm{1}^{n+s}\\ \hline\cr\\ \bm{1}^{m}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}\end{pmatrix}\odot\begin{pmatrix}\bm{p}\odot(\bm{h}\ast\bm{x})\\ \hline\cr\\ \bm{0}^{m+t}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \bm{p}\odot(\bm{h}\ast\bm{x})\\ \hline\cr\\ \bm{v}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}\end{pmatrix}\equiv\begin{pmatrix}\bm{0}^{n+s}\\ \hline\cr\\ \bm{0}^{m+t}\\ \hline\cr\\ \bm{0}^{n}\\ \hline\cr\\ \bm{p}\odot(\bm{h}\ast\bm{x})\\ \hline\cr\\ \bm{v}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}\end{pmatrix} We now shift this up by 2 ​ n + s + m + t 2n+s+m+t entries using the primitive operation defined in Proposition H.8 that costs three additional layers so that we end up with 𝒛 ≡ ( 𝒑 ⊙ ( 𝒉 ∗ 𝒖 ) 𝒗 ⋯ 𝟎 ) \bm{z}\equiv\begin{pmatrix}\bm{p}\odot(\bm{h}\ast\bm{u})\\ \hline\cr\\ \bm{v}\\ \hline\cr\\ \cdots\\ \hline\cr\\ \bm{0}\end{pmatrix} Again, we note that the only operations we perform and are convolutions and Hadamard product and they generalize in the obvious way to d > 1 d>1 . ∎

Finally, we show that these primitives may be composed by ’stacking’ models with matching inner dimension ( N ′ , d ′ ) (N^{\prime},d^{\prime}) .

###### Lemma H.11 .

For f , g : ℝ N × d → ℝ N × d f,g:\mathbbm{R}^{N\times d}\to\mathbbm{R}^{N\times d} that have ( N , L 1 , d , N ′ , d ′ ) (N,L_{1},d,N^{\prime},d^{\prime}) and ( N , L 2 , d , N ′ , d ′ ) (N,L_{2},d,N^{\prime},d^{\prime}) BaseConv models then their composition f ∘ g f\circ g has an ( N , L 1 + L 2 , d , N ′ , d ′ ) (N,L_{1}+L_{2},d,N^{\prime},d^{\prime}) BaseConv model which can be computed by performing their models in succession, or ’stacking’.

###### Proof.

This result follows from noting that for any f ⁡ ( 𝒖 ) f(\bm{u}) which requires L 1 L_{1} layers to compute and that we can compute f ∘ g ⁡ ( 𝒖 ) = g ⁡ ( f ⁡ ( 𝒖 ) ) f\circ g(\bm{u})=g(f(\bm{u})) using the BaseConv model with L 2 L_{2} layers, yielding L 1 + L 2 L_{1}+L_{2} layers in total. ∎

#### H.3.1 BaseConv -Hyena Equivalence

We show that the equivalence between BaseConv 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 𝐮 ∈ ℝ N × d \bm{u}\in\mathbbm{R}^{N\times d} and ( N , L , d , N ′ , d ) − (N,L,d,N^{\prime},d)- Hyena such that 𝐳 Hyena ≡ Hyena ⁡ ( 𝐮 ) \bm{z}_{\mathrm{Hyena}}\equiv\mathrm{Hyena}(\bm{u}) with a set of filters 𝐡 ℓ \bm{h}^{\ell} and linear projections 𝐩 ℓ \bm{p}^{\ell} as per Definition H.3 for ℓ ∈ [ L ] \ell\in[L] , there exists a ( N , 5 ​ L , d , N ′ + N , d ) (N,5L,d,N^{\prime}+N,d) - BaseConv model such that 𝐳 Hyena ≡ BaseConv ​ ( 𝐮 ) \bm{z}_{\mathrm{Hyena}}\equiv\textsc{BaseConv}(\bm{u}) .

Similarly, for any input 𝐮 BaseConv ∈ ℝ N × d \bm{u}_{\textsc{BaseConv}}\in\mathbbm{R}^{N\times d} and ( N , L , d , N ′ , d ) − (N,L,d,N^{\prime},d)- Coyote such that 𝐳 BaseConv ≡ BaseConv ​ ( 𝐮 ) \bm{z}_{\textsc{BaseConv}}\equiv\textsc{BaseConv}(\bm{u}) with a set of filters 𝐡 ℓ \bm{h}^{\ell} for ℓ ∈ [ L ] \ell\in[L] , there exists a series of Hyena layers such that we have Hyena ⁡ ( Hyena ⁡ ( … ​ Hyena ​ ( 𝒖 BaseConv , 𝒉 ) ) ) ⏟ L ​ layers ≡ 𝒛 BaseConv . \underbrace{\mathrm{Hyena}\left(\mathrm{Hyena}\left(\ldots\mathrm{Hyena}(\bm{u}_{\textsc{BaseConv},\bm{h}})\right)\right)}_{L\text{ layers}}\equiv\bm{z}_{\textsc{BaseConv}}.

###### Proof.

For the input 𝒖 Hyena ∈ ℝ N × d \bm{u}_{\mathrm{Hyena}}\in\mathbbm{R}^{N\times d} , the output of the ℓ \ell th layer 𝒛 Hyena ℓ ∈ ℝ N ′ × d ′ \bm{z}^{\ell}_{\mathrm{Hyena}}\in\mathbbm{R}^{N^{\prime}\times d^{\prime}} for Hyena is given by (see Algorithm 3 ) 𝒛 Hyena ℓ ≡ 𝒑 Hyena ℓ ⊙ ( 𝒉 l ∗ 𝒛 ℓ − 1 ) , \displaystyle\bm{z}^{\ell}_{\mathrm{Hyena}}\equiv\bm{p}^{\ell}_{\mathrm{Hyena}}\odot(\bm{h}^{l}\ast\bm{z}^{\ell-1}), where 𝒑 Hyena ℓ ≡ Linear ​ ( 𝒖 Hyena ) ∈ ℝ N ′ × d \bm{p}^{\ell}_{\mathrm{Hyena}}\equiv\texttt{Linear}(\bm{u}_{\mathrm{Hyena}})\in\mathbbm{R}^{N^{\prime}\times d} . Now, using the original input 𝒖 Hyena ∈ ℝ N × d \bm{u}_{\mathrm{Hyena}}\in\mathbbm{R}^{N\times d} to Hyena, we define the following input for BaseConv using one layer: 𝒖 BaseConv ≡ ( 𝒖 Hyena 𝟎 ( N ′ − N ) × d 𝒖 Hyena ) \bm{u}_{\textsc{BaseConv}}\equiv\begin{pmatrix}\bm{u}_{\mathrm{Hyena}}\\ \hline\cr\\ \bm{0}^{(N^{\prime}-N)\times d}\\ \hline\cr\\ \bm{u}_{\mathrm{Hyena}}\end{pmatrix} Then, we simply use the 𝚛𝚎𝚖𝚎𝚖𝚋𝚎𝚛 N , N , N ′ − N , N ′ − N ( 𝒖 BaseConv : 𝒖 Hyena , 𝒖 Hyena , 𝒉 Hyena ℓ , 𝒑 Hyena ℓ ) {\tt remember}_{N,N,N^{\prime}-N,N^{\prime}-N}(\bm{u}_{\textsc{BaseConv}}:\bm{u}_{\mathrm{Hyena}},\bm{u}_{\mathrm{Hyena}},\bm{h}^{\ell}_{\mathrm{Hyena}},\bm{p}^{\ell}_{\mathrm{Hyena}}) primitive for BaseConv . Consequently, this allows us to “remember” the input 𝒖 Hyena \bm{u}_{\mathrm{Hyena}} in the output of the previous BaseConv layer 𝒛 BaseConv ℓ − 1 \bm{z}^{\ell-1}_{\textsc{BaseConv}} . We then use this to retrieve 𝒑 Hyena ℓ ≡ linear ⁡ ( 𝒖 Hyena ) \bm{p}^{\ell}_{\mathrm{Hyena}}\equiv\mathrm{linear}(\bm{u}_{\mathrm{Hyena}}) with the projection used for BaseConv given by p BaseConv ℓ ≡ Linear ​ ( 𝒛 BaseConv ℓ − 1 ) ≡ ( 𝟏 N × d 𝒑 Hyena ℓ . ) p^{\ell}_{\textsc{BaseConv}}\equiv\texttt{Linear}(\bm{z}^{\ell-1}_{\textsc{BaseConv}})\equiv\begin{pmatrix}\bm{1}^{N\times d}\\ \hline\cr\\ \bm{p}^{\ell}_{\mathrm{Hyena}}.\end{pmatrix} Overall, the output of the ℓ \ell th layer for BaseConv is given by 𝒛 BaseConv ℓ ≡ ( 𝒑 Hyena ℓ ⊙ ( 𝒉 Hyena ℓ ∗ 𝒖 Hyena ) 𝟎 ( M − N ) × d 𝒖 Hyena ) ≡ ( 𝒛 Hyena ℓ 𝟎 ( M − N ) × d 𝒖 Hyena ) \bm{z}^{\ell}_{\textsc{BaseConv}}\equiv\begin{pmatrix}\bm{p}^{\ell}_{\mathrm{Hyena}}\odot\left(\bm{h}^{\ell}_{\mathrm{Hyena}}\ast\bm{u}_{\mathrm{Hyena}}\right)\\ \hline\cr\\ \bm{0}^{(M-N)\times d}\\ \hline\cr\\ \bm{u}_{\mathrm{Hyena}}\end{pmatrix}\equiv\begin{pmatrix}\bm{z}^{\ell}_{\mathrm{Hyena}}\\ \hline\cr\\ \bm{0}^{(M-N)\times d}\\ \hline\cr\\ \bm{u}_{\mathrm{Hyena}}\end{pmatrix} Hence, we can reproduce the output of the ℓ \ell th layer of Hyena using five layers of BaseConv after augmenting the input and using the remembering primitive ( Proposition H.10 ) with internal dimension N ′ + N N^{\prime}+N .

Now, for the input 𝒖 BaseConv ∈ ℝ N × d \bm{u}_{\textsc{BaseConv}}\in\mathbbm{R}^{N\times d} , the output of the ℓ \ell th layer for BaseConv is given by z BaseConv ℓ ≡ Linear ​ ( z BaseConv ℓ − 1 ) ⊙ conv ⁡ ( h l , z BaseConv ℓ − 1 ) . \displaystyle z^{\ell}_{\textsc{BaseConv}}\equiv\texttt{Linear}(z^{\ell-1}_{\textsc{BaseConv}})\odot\mathrm{conv}(h^{l},z^{\ell-1}_{\textsc{BaseConv}}). Here, we show inductively that simply using ℓ \ell -many Hyena models recursively simulates z BaseConv ℓ z^{\ell}_{\textsc{BaseConv}} . For ℓ = 1 \ell=1 , we have Hyena ⁡ ( 𝒖 BaseConv , 𝒉 ) ≡ Linear ​ ( 𝒖 BaseConv ) ⊙ ( 𝒉 1 ∗ 𝒖 BaseConv ) ≡ 𝒛 BaseConv 1 . \mathrm{Hyena}(\bm{u}_{\textsc{BaseConv},\bm{h}})\equiv\texttt{Linear}(\bm{u}_{\textsc{BaseConv}})\odot(\bm{h}^{1}\ast\bm{u}_{\textsc{BaseConv}})\equiv\bm{z}^{1}_{\textsc{BaseConv}}. We now assume that ( ℓ − 1 ) (\ell-1) -many recursive Hyena models produce z BaseConv ( ℓ − 1 ) z^{(\ell-1)}_{\textsc{BaseConv}} . For the ℓ \ell th layer, we then have Hyena ⁡ ( Hyena ⁡ ( … ​ Hyena ​ ( 𝒖 BaseConv , 𝒉 ) ) ) \displaystyle\mathrm{Hyena}\left(\mathrm{Hyena}\left(\ldots\mathrm{Hyena}(\bm{u}_{\textsc{BaseConv},\bm{h}})\right)\right) ≡ \displaystyle\equiv\quad Hyena ⁡ ( z BaseConv ( l − 1 ) ) \displaystyle\mathrm{Hyena}\left(z^{(l-1)}_{\textsc{BaseConv}}\right) ≡ \displaystyle\equiv\quad linear ⁡ ( z BaseConv ( l − 1 ) ) ⊙ conv ⁡ ( h l , z BaseConv ( l − 1 ) ) \displaystyle\mathrm{linear}\left(z^{(l-1)}_{\textsc{BaseConv}}\right)\odot\mathrm{conv}\left(h^{l},z^{(l-1)}_{\textsc{BaseConv}}\right) ≡ \displaystyle\equiv\quad z BaseConv ℓ . \displaystyle z^{\ell}_{\textsc{BaseConv}}. ∎

### H.4 Linear Arithmetic Circuits

In this section we show the relation between linear arithmetic circuits and BaseConv . We recall a few definitions from [ Dao et al., 2020 ] .

###### Definition H.13 (Linear Arithmetic Circuit [ Bürgisser et al., 1996 ] ) .

An arithmetic circuit is called a linear arithmetic circuit if it only uses addition, subtraction and scalar multiplication. Further, every multiplication has a fixed constant from 𝔽 \mathbbm{F} as at least one of its two inputs. In other words, all gates in the circuit are linear functions of their inputs (i.e. of the form a ​ x + b ​ y ax+by for fixed constants a , b ∈ 𝔽 a,b\in\mathbbm{F} ).

###### Definition H.14 (Butterfly Matrices [ Dao et al., 2020 ] ) .

A butterfly factor of size k ≥ 2 k\geq 2 (denoted as 𝐁 k \mathbf{B}_{k} ) is a matrix of the form 𝐁 k = \mathbf{B}_{k}= [ 𝐃 1 𝐃 2 𝐃 3 𝐃 4 ] \left[\begin{array}[]{ll}\mathbf{D}_{1}&\mathbf{D}_{2}\\ \mathbf{D}_{3}&\mathbf{D}_{4}\end{array}\right] where each 𝐃 i \mathbf{D}_{i} is a k 2 × k 2 \frac{k}{2}\times\frac{k}{2} diagonal matrix. We restrict k k to be a power of 2 .

A butterfly factor matrix of size n n with block size k k (denoted as 𝐁 k ( n ) \mathbf{B}_{k}^{(n)} ) is a block diagonal matrix of n k \frac{n}{k} (possibly different) butterfly factors of size k k : 𝐁 k ( n ) = diag ⁡ ( [ 𝐁 k ] 1 , [ 𝐁 k ] 2 , … , [ 𝐁 k ] n k ) \mathbf{B}_{k}^{(n)}=\operatorname{diag}\left(\left[\mathbf{B}_{k}\right]_{1},\left[\mathbf{B}_{k}\right]_{2},\ldots,\left[\mathbf{B}_{k}\right]_{\frac{n}{k}}\right) Finally, a butterfly matrix of size n n (denoted as 𝐁 ( n ) \mathbf{B}^{(n)} ) is a matrix that can be expressed as a product of butterfly factor matrices: 𝐁 ( n ) = 𝐁 n ( n ) ​ 𝐁 n 2 ( n ) ​ … ​ 𝐁 2 ( n ) \mathbf{B}^{(n)}=\mathbf{B}_{n}^{(n)}\mathbf{B}_{\frac{n}{2}}^{(n)}\ldots\mathbf{B}_{2}^{(n)} . Equivalently, we may define 𝐁 ( n ) \mathbf{B}^{(n)} recursively as a matrix that can be expressed in the following form: 𝐁 ( n ) = 𝐁 n ( n ) ​ [ [ 𝐁 ( n 2 ) ] 1 0 0 [ 𝐁 ( n 2 ) ] 2 ] \mathbf{B}^{(n)}=\mathbf{B}_{n}^{(n)}\left[\begin{array}[]{cc}{\left[\mathbf{B}^{\left(\frac{n}{2}\right)}\right]_{1}}&0\\ 0&{\left[\mathbf{B}^{\left(\frac{n}{2}\right)}\right]_{2}}\end{array}\right] (Note that [ 𝐁 ( n 2 ) ] 1 \left[\mathbf{B}^{\left(\frac{n}{2}\right)}\right]_{1} and [ 𝐁 ( n 2 ) ] 2 \left[\mathbf{B}^{\left(\frac{n}{2}\right)}\right]_{2} may be different.)

From Definition H.14 , we observe that size n n butterfly factor is comprised of three vectors 𝒅 , 𝒅 + , 𝒅 − ∈ ℝ n \bm{d,d^{+},d^{-}}\in\mathbbm{R}^{n} such that 𝒅 \displaystyle\bm{d} = ( diag − 1 ⁡ ( 𝐃 1 ) , diag − 1 ⁡ ( 𝐃 4 ) ) , \displaystyle=\left(\operatorname{diag}^{-1}\left({\bf D}_{1}\right),\operatorname{diag}^{-1}\left({\bf D}_{4}\right)\right), 𝒅 + \displaystyle\bm{d}^{+} = ( 𝟎 n 2 , diag − 1 ⁡ ( 𝐃 2 ) ) , and \displaystyle=\left(\bm{0}^{\tfrac{n}{2}},\operatorname{diag}^{-1}\left({\bf D}_{2}\right)\right),\text{ and } 𝒅 − \displaystyle\bm{d}^{-} = ( diag − 1 ⁡ ( 𝐃 3 ) , 𝟎 n 2 ) , \displaystyle=\left(\operatorname{diag}^{-1}\left({\bf D}_{3}\right),\bm{0}^{\tfrac{n}{2}}\right),

where diag − 1 ⁡ ( 𝐃 ) : ℝ n × n ↦ ℝ n \operatorname{diag}^{-1}({\bf D}):\mathbbm{R}^{n\times n}\mapsto\mathbbm{R}^{n} is the mapping from diagonal matrices to the vector of its diagonal entries. Let us define 𝐃 𝟏 , 𝐃 𝟐 , 𝐃 𝟑 ∈ ℝ n × n {\bf D_{1},D_{2},D_{3}}\in\mathbbm{R}^{n\times n} as diag ⁡ ( 𝒅 ) , diag ⁡ ( 𝒅 + ) \operatorname{diag}\left(\bm{d}\right),\operatorname{diag}\left(\bm{d^{+}}\right) , and diag ⁡ ( 𝒅 − ) \operatorname{diag}\left(\bm{d^{-}}\right) respectively. Then we note that

𝐃 1 ≡ [ 𝐃 1 𝟎 𝟎 𝐃 4 ] 𝐃 2 ​ 𝑺 n 2 ≡ [ 𝟎 𝐃 2 𝟎 𝟎 ] 𝑺 n 2 ​ 𝐃 3 ≡ [ 𝟎 𝟎 𝐃 3 𝟎 ] {\bf D}_{1}\equiv\left[\begin{array}[]{ll}\mathbf{D}_{1}&\mathbf{0}\\ \mathbf{0}&\mathbf{D}_{4}\end{array}\right]\qquad{\bf D}_{2}\bm{S}^{\tfrac{n}{2}}\equiv\left[\begin{array}[]{ll}\mathbf{0}&\mathbf{D}_{2}\\ \mathbf{0}&\mathbf{0}\end{array}\right]\qquad\bm{S}^{\tfrac{n}{2}}{\bf D}_{3}\equiv\left[\begin{array}[]{ll}\mathbf{0}&\mathbf{0}\\ \mathbf{D}_{3}&\mathbf{0}\end{array}\right] (5)

where 𝐒 k ∈ 𝔽 n × n {\bf S}^{k}\in\mathbbm{F}^{n\times n} is a shift matrix for i ∈ [ n / 2 ] i\in[n/2] . This gives us the following proposition:

###### Proposition H.15 .

For any powers of 2, n = k ≥ 2 n=k\geq 2 , any butterfly factor matrix 𝐁 k ( n ) {\bf B}^{(n)}_{k} is equivalent to 𝐁 k ( n ) = 𝐒 k 2 ​ 𝐃 3 + 𝐃 2 ​ 𝐒 n 2 + 𝐃 1 {\bf B}^{(n)}_{k}={\bf S}^{\tfrac{k}{2}}{\bf D}_{3}+{\bf D}_{2}{\bf S}^{\tfrac{n}{2}}+{\bf D}_{1} where 𝐃 𝟑 , 𝐃 𝟐 , 𝐃 𝟏 , 𝐒 𝐧 𝟐 {\bf D_{3},D_{2},D_{1},S^{\frac{n}{2}}} are defined as in equation 5 .

We use Proposition H.15 to show that butterfly matrices can easily be computed by BaseConv .

###### Lemma H.16 .

For any n , d ≥ 2 n,d\geq 2 , k ≥ 1 k\geq 1 , and arbitrary vector 𝐱 ∈ ℝ n ​ d \bm{x}\in\mathbbm{R}^{nd} : (1) there exists a ( N , L , d , N ′ , d ′ ) − BaseConv \left(N,L,d,N^{\prime},d^{\prime}\right)-\text{{BaseConv}} that can represent 𝐁 k ( n ​ d ) ⋅ 𝒙 {\bf B}^{(nd)}_{k}\cdot\bm{x} with N = n , N ′ = 𝒪 ⁡ ( N ) N=n,N^{\prime}=\mathcal{O}(N) , L = 𝒪 ⁡ ( 1 ) L=\mathcal{O}(1) , and d ′ = 𝒪 ⁡ ( d ) d^{\prime}=\mathcal{O}(d) , and

(2) there exists a ( N , L , d , N ′ , d ′ ) − BaseConv \left(N,L,d,N^{\prime},d^{\prime}\right)-\text{{BaseConv}} that can represent 𝐁 ( n ​ d ) ⋅ 𝒙 {\bf B}^{(nd)}\cdot\bm{x} with N = n , N ′ = 𝒪 ⁡ ( N ) N=n,N^{\prime}=\mathcal{O}(N) , L = 𝒪 ⁡ ( log ⁡ n ​ d ) L=\mathcal{O}(\log nd) , and d ′ = 𝒪 ⁡ ( d ) d^{\prime}=\mathcal{O}(d) .

###### Proof.

(1) Given 𝒙 ∈ ℝ n ​ d \bm{x}\in\mathbbm{R}^{nd} , construct 𝒖 ∈ ℝ n × d \bm{u}\in\mathbbm{R}^{n\times d} where 𝒙 \bm{x} is the row-major form of 𝒖 \bm{u} . We show that BaseConv can compute 𝐁 n ​ d ⋅ 𝒙 {\bf B}_{nd}\cdot\bm{x} column by column.

Let 𝐀 = 𝐒 k 2 ​ 𝐃 3 ′ , 𝐂 = 𝐃 2 ′ ​ 𝐒 n 2 {\bf A}={\bf S}^{\tfrac{k}{2}}{\bf D}^{\prime}_{3},{\bf C}={\bf D}^{\prime}_{2}{\bf S}^{\tfrac{n}{2}} , and 𝐃 = 𝐃 1 {\bf D}={\bf D}_{1} for 𝐃 i , 𝐒 k 2 ∈ ℝ n ​ d × n ​ d {\bf D}_{i},{\bf S}^{\tfrac{k}{2}}\in\mathbbm{R}^{nd\times nd} for 1 ≤ i ≤ 3 1\leq i\leq 3 as defined in Proposition H.15 . We take 𝒅 1 = 𝟏 n ​ d ​ 𝐃 , 𝒅 2 = 𝟏 n ​ d ​ 𝐂 2 , 𝒅 3 = 𝟏 n ​ d ​ 𝐀 \bm{d}_{1}=\bm{1}^{nd}{\bf D},\ \bm{d}_{2}=\bm{1}^{nd}{\bf C}_{2},\bm{d}_{3}=\bm{1}^{nd}{\bf A} , which extracts the diagonal entries of 𝐃 i {\bf D}_{i} . With this we construct 𝐃 i ′ ∈ ℝ n × d {\bf D}^{\prime}_{i}\in\mathbbm{R}^{n\times d} where 𝒅 i \bm{d}_{i} is the row major form of 𝑫 i ′ \bm{D}^{\prime}_{i} . This implies that 𝐃 i ​ 𝒙 ≡ 𝐃 i ′ ⊙ 𝒖 . {\bf D}_{i}\bm{x}\equiv{\bf D}^{\prime}_{i}\odot\bm{u}.

Then we can decompose 𝐁 n ​ d ⋅ 𝒙 {\bf B}_{nd}\cdot\bm{x} into 𝐁 n ​ d ​ 𝒙 ≡ 𝐃 1 ⊙ 𝒖 + 𝐃 2 ⊙ 𝒖 + 𝐃 3 ⊙ 𝒖 . {\bf B}_{nd}\bm{x}\equiv{\bf D}_{1}\odot\bm{u}+{\bf D}_{2}\odot\bm{u}+{\bf D}_{3}\odot\bm{u}.

By Lemma H.7 , each Hadamard product 𝐀 ⊙ 𝒖 , 𝐁 ⊙ 𝒖 , 𝐂 ⊙ 𝒖 {\bf A}\odot\bm{u},{\bf B}\odot\bm{u},{\bf C}\odot\bm{u} can be trivially be performed with a single layer BaseConv model. Let each of these model outputs be denoted 𝒚 𝟏 , 𝒚 𝟐 , 𝒚 𝟑 \bm{y_{1},y_{2},y_{3}} , respectively. Finally all that remains is to compute the 𝒚 𝟏 + 𝒚 𝟐 + 𝒚 𝟑 \bm{y_{1}+y_{2}+y_{3}} . We achieve this using layers of add primitives 11 11 11 Recall that add n ( 𝒚 : 𝒙 , 𝑺 ) \texttt{add}_{n}(\bm{y}:\bm{x},\bm{S}) adds the subvector 𝒙 \bm{x} to 𝑺 \bm{S} for the input 𝒚 \bm{y} . : add n ( 𝒚 1 : 𝒚 1 , 𝟎 ) \displaystyle\texttt{add}_{n}(\bm{y}^{1}:\bm{y}_{1},\bm{0}) add n ( 𝒚 2 : 𝒚 2 , 𝒚 1 ) \displaystyle\texttt{add}_{n}(\bm{y}^{2}:\bm{y}_{2},\bm{y}_{1}) add n ( 𝒚 3 : 𝒚 3 , 𝒚 1 + 𝒚 2 ) , \displaystyle\texttt{add}_{n}(\bm{y}^{3}:\bm{y}_{3},\bm{y}_{1}+\bm{y}_{2}), where using by Proposition H.9 and Lemma H.11 , this requires six more layers, and we get 𝒚 3 ≡ 𝒚 𝟏 + 𝒚 𝟐 + 𝒚 𝟑 ≡ 𝐁 n ​ d ​ 𝒙 . \bm{y}^{3}\equiv\bm{y_{1}+y_{2}+y_{3}}\equiv{\bf B}_{nd}\bm{x}. Then we can construct the ( N , L , d , N ′ , d ′ ) − BaseConv \left(N,L,d,N^{\prime},d^{\prime}\right)-\text{{BaseConv}} as desired with L = O ⁡ ( 1 ) L=O(1) layers.

(2) From Definition H.14 , 𝐁 ( n ​ d ) = 𝐁 n ​ d ( n ​ d ) ​ 𝐁 n ​ d 2 ( n ​ d ) ​ … ​ 𝐁 2 ( n ​ d ) \mathbf{B}^{(nd)}=\mathbf{B}_{nd}^{(nd)}\mathbf{B}_{\frac{nd}{2}}^{(nd)}\ldots\mathbf{B}_{2}^{(nd)} . From (1), BaseConv can compute any butterfly matrix by simulating the log ⁡ ( n ​ d ) \log(nd) butterfly factor matrices which comprise 𝐁 ( n ​ d ) \mathbf{B}^{(nd)} . With Lemma H.11 , this creates a BaseConv with 5 ⋅ log ⁡ ( n ​ d ) = 𝒪 ⁡ ( log ⁡ ( n ​ d ) ) 5\cdot\log(nd)=\mathcal{O}(\log(nd)) layers. Lemma H.11

∎

Butterfly matrices comprise the kaleidoscope hierarchy, which we define below:

###### Definition H.17 (The Kaleidoscope Hierarchy [ Dao et al., 2020 ] ) .

• Define ℬ \mathcal{B} as the set of all matrices that can be expressed in the form 𝐁 ( n ) {\bf B}^{(n)} (for some n n ).

• Define ( ℬℬ ∗ ) \left(\cal{BB}^{*}\right) as the set of matrices 𝐌 {\bf M} of the form 𝐌 = 𝐌 1 ​ 𝐌 2 ∗ {\bf M}={\bf M}_{1}{\bf M}_{2}^{*} for some 𝐌 𝟏 , 𝐌 2 ∈ ℬ {\bf M_{1}},{\bf M}_{2}\in\mathcal{B} .

• Define ( ℬℬ ∗ ) w \left(\cal{BB}^{*}\right)^{w} as the set of matrices 𝐌 {\bf M} that can be expressed as 𝐌 = 𝐌 w ​ … ​ 𝐌 2 ​ 𝐌 1 {\bf M}={\bf M}_{w}\ldots{\bf M}_{2}{\bf M}_{1} , with each 𝐌 i ∈ ( ℬℬ ∗ ) ​ ( 1 ≤ i ≤ w ) . {\bf M}_{i}\in\left(\cal{BB}^{*}\right)(1\leq i\leq w). (The notation w w represents width.)

• Define ( ℬℬ ∗ ) e w \left(\cal{BB}^{*}\right)^{w}_{e} as the set of n × n n\times n matrices 𝐌 {\bf M} that can be expressed as 𝐌 = 𝐒𝐄𝐒 ⊤ {\bf M}={\bf SES}^{\top} for some e ​ n × e ​ n en\times en matrix 𝐄 ∈ ( ℬℬ ∗ ) w {\bf E}\in\left(\cal{BB}^{*}\right)^{w} , where 𝐒 ∈ 𝔽 n × e ​ n = [ 𝐈 n 0 … 0 ] ] {\bf S}\in\mathbbm{F}^{n\times en}=\begin{bmatrix}{\bf I}_{n}&0&\ldots&0\end{bmatrix}] (i.e. 𝐌 {\bf M} is the upper-left corner of 𝐄 {\bf E} ). (The notation e e represents expansion relative to n n .)

We similarly show how BaseConv can simulate any kaleidoscope matrix.

###### Lemma H.18 .

Given n , d ≥ 2 n,d\geq 2 , e > 0 e>0 for any n ​ d × n ​ d nd\times nd matrix 𝐌 ∈ ( ℬℬ ∗ ) e w {\bf M}\in\left(\cal{BB}^{*}\right)^{w}_{e} , and 𝐱 ∈ ℝ n ​ d \bm{x}\in\mathbbm{R}^{nd} there exists a ( N , L , d , N ′ , d ′ ) − BaseConv \left(N,L,d,N^{\prime},d^{\prime}\right)-\text{{BaseConv}} that can represent 𝐌 ⋅ 𝐱 {\bf M}\cdot\bm{x} with N = n , L = 𝒪 ⁡ ( w ​ log ⁡ ( e ​ n ​ d ) ) N=n,L=\mathcal{O}(w\log(end)) , N ′ = e ​ n , N^{\prime}=en, and d ′ = d \ d^{\prime}=d .

###### Proof.

By Definition H.17 , 𝐌 {\bf M} can be decomposed with respect to size e ​ n ​ d × e ​ n ​ d end\times end matrix 𝐄 = 𝐄 1 ⋅ 𝐄 2 ⋯ 𝐄 w . {\bf E}={\bf E}_{1}\cdot{\bf E}_{2}\cdots{\bf E}_{w}.

Further, any 𝐄 i ∈ ( ℬℬ ∗ ) {\bf E}_{i}\in\left(\cal{BB}^{*}\right) can be expressed as a product of 2 ​ log ⁡ e ​ n ​ d 2\log{end} butterfly factor matrices. Then by Lemma H.16 and Lemma H.11 we can compute 𝐄 i ​ 𝒙 ′ {\bf E}_{i}\bm{x^{\prime}} in by stacking 2 ​ log ⁡ e ​ n ​ d 2\log{end} ( n , d , L , e ​ n , d ) − BaseConv \left(n,d,L,en,d\right)-\text{{BaseConv}} models each with L = 𝒪 ⁡ ( 1 ) L=\mathcal{O}(1) . Because 𝐄 {\bf E} has width w w , Lemma H.11 implies that composing with each 𝐄 i {\bf E}_{i} for 1 ≤ i ≤ w 1\leq i\leq w constructs a final model with 𝒪 ⁡ ( w ​ log ⁡ ( e ​ n ​ d ) ) \mathcal{O}(w\log(end)) layers. ∎

Finally, the kaleidoscope hierarchy is related to linear arithmetic circuits via the following result. We note that in [ Dao et al., 2020 ] it is assumed that w = s w=s , yet inspection of the proof yields the following stronger result:

###### Theorem H.19 ( [ Dao et al., 2020 ] ) .

Let 𝐌 {\bf M} be an n × n n\times n matrix such that multiplication of 𝐌 {\bf M} times an arbitrary vector 𝐮 \bm{u} can be be represented as ( n , s , Δ , w ) (n,s,\Delta,w) -linear arithmetic circuit 𝒞 \mathcal{C} . Then, 𝐌 ∈ ( ℬℬ ∗ ) 𝒪 ⁡ ( w / n ) 𝒪 ⁡ ( Δ ) {\bf M}\in\left(\cal{BB}^{*}\right)^{\mathcal{O}(\Delta)}_{\mathcal{O}(w/n)} .

We combine Theorem H.19 and Lemma H.18 to show that BaseConv can compute any linear arithmetic circuit with polylogarithmic factors in Δ \Delta .

###### Corollary H.20 .

For any ( n ​ d , s , Δ , w ) (nd,s,\Delta,w) -linear arithmetic circuit 𝒞 \mathcal{C} that can be represented by a matrix 𝐌 ∈ ℝ n ​ d × n ​ d {\bf M}\in\mathbbm{R}^{nd\times nd} multiplied by a vector 𝐱 ∈ ℝ n ​ d \bm{x}\in\mathbbm{R}^{nd} , there exists an equivalent ( n , Δ ′ , d , w , d ) − BaseConv \left(n,\Delta^{\prime},d,w,d\right)-\text{{BaseConv}} with Δ ′ = 𝒪 ⁡ ( Δ ​ log ⁡ ( w ) ) \Delta^{\prime}=\mathcal{O}(\Delta\log(w)) such that 𝐌 ​ 𝐱 = BaseConv ​ ( 𝐮 , 𝐡 ) {\bf M}\bm{x}=\textsc{BaseConv}\left({\bm{u},\bm{h}}\right) where 𝐱 \bm{x} is the row major form of 𝐮 ∈ ℝ n × d \bm{u}\in\mathbbm{R}^{n\times d} .

### 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 ( n ​ d , s , Δ , w ) (nd,s,\Delta,w) -arithmetic circuit 𝒞 \mathcal{C} , there exists an equivalent ( N , Δ ′ , d , N ′ , d ′ ) − BaseConv \left(N,\Delta^{\prime},d,N^{\prime},d^{\prime}\right)-\text{{BaseConv}} with N = n , Δ ′ = 𝒪 ⁡ ( Δ ​ log ⁡ w ) N=n,\Delta^{\prime}=\mathcal{O}(\Delta\log{w}) , N ′ = 𝒪 ⁡ ( w ) , d ′ = d N^{\prime}=\mathcal{O}(w),d^{\prime}=d that simulates 𝒞 \mathcal{C} .

For the reader’s convenience, we begin with a proof sketch and then provide the details afterwards.

###### Proof Sketch of Theorem H.21 .

Let us layer 𝒞 \mathcal{C} so that each layer 𝒞 ℓ \mathcal{C}_{\ell} for ℓ ∈ [ L 𝒞 ] \ell\in[L_{\mathcal{C}}] either only has linear gates or multiplication gates, where the number of such layers L 𝒞 = 𝒪 ⁡ ( Δ ) L_{\mathcal{C}}=\mathcal{O}(\Delta) . The composition of all 𝒞 ℓ \mathcal{C}_{\ell} layers results in 𝒞 \mathcal{C} . We use 𝒛 ℓ + 1 ∈ ℝ w \bm{z}^{\ell+1}\in\mathbbm{R}^{w} to denote the output of the ℓ \ell -th layer 𝒞 ℓ \mathcal{C}_{\ell} which feeds as the input to the ( ℓ + 1 ) (\ell+1) -th layer 𝒞 ℓ + 1 \mathcal{C}_{\ell+1} . Here, we note that if we can simulate each 𝒞 ℓ \mathcal{C}_{\ell} with BaseConv , then we can simulate the entire layered circuit 𝒞 \mathcal{C} due to Lemma H.11 .

Now, if the layer 𝒞 ℓ lin \mathcal{C}_{\ell}^{\mathrm{lin}} is a linear layer (with only addition gates), then it can be represented by a matrix 𝑴 ∈ ℝ w × w \bm{M}\in\mathbbm{R}^{w\times w} multiplied by 𝒛 ℓ ∈ ℝ w \bm{z}^{\ell}\in\mathbbm{R}^{w} (We can append with 0 0 s if necessary so that the input from the previous gates can be written as w w -length vector). Thus, we can apply Corollary H.20 to simulate 𝒞 ℓ lin \mathcal{C}_{\ell}^{\mathrm{lin}} with an equivalent ( n , log ⁡ w , d , 𝒪 ⁡ ( log ⁡ w ) , d ) − BaseConv \left(n,\log{w},d,\mathcal{O}(\log{w}),d\right)-\text{{BaseConv}} model.

Next, if 𝒞 ℓ mult \mathcal{C}_{\ell}^{\mathrm{mult}} instead consists of only the multiplication gates. Then, we note here that the output 𝒛 ℓ \bm{z}^{\ell} may not exactly equal the input to 𝒞 ℓ mult \mathcal{C}_{\ell}^{\mathrm{mult}} . Nevertheless, we can apply a 𝒪 ⁡ ( w ) \mathcal{O}(w) sparse linear map 𝐑 ∈ ℝ w × w {\bf R}\in\mathbbm{R}^{w\times w} so that 𝐑 ​ 𝒛 ℓ {\bf R}\bm{z}^{\ell} yields vectors 𝒗 1 , 𝒗 2 , \bm{v}^{1},\bm{v}^{2}, and 𝒗 3 \bm{v}^{3} , where 𝒗 1 \bm{v}_{1} constitutes the “first” input to all the multiplication gates and 𝒗 2 \bm{v}_{2} constitutes all the “second” inputs while 𝒗 3 \bm{v}^{3} consists of all entries needed as inputs in the subsequent layers. That is, for the i i th gate in 𝒞 ℓ mult \mathcal{C}_{\ell}^{\mathrm{mult}} , we compute 𝒗 i 1 ⋅ 𝒗 i 2 \bm{v}^{1}_{i}\cdot\bm{v}^{2}_{i} . This implies that for all the gates in 𝒞 ℓ mult \mathcal{C}_{\ell}^{\mathrm{mult}} , we can simply compute 𝒗 1 ⊙ 𝒗 2 \bm{v}_{1}\odot\bm{v}_{2} . To this end, we use the 𝚛𝚎𝚖𝚎𝚖𝚋𝚎𝚛 {\tt remember} primitive with constant number of layers from Proposition H.10 to define a ( n , 𝒪 ⁡ ( log ⁡ w ) , d , w , d ) − BaseConv \left(n,\mathcal{O}(\log{w}),d,w,d\right)-\text{{BaseConv}} model that remembers 𝒗 3 \bm{v}^{3} while performing the Hadamard product of 𝒗 1 \bm{v}^{1} with 𝒗 2 \bm{v}^{2} .

Overall, we can then collect all the resulting BaseConv layers and compose them as in Lemma H.11 to simulate 𝒞 \mathcal{C} . Overall, the number of layers used is given by 𝒪 ⁡ ( Δ ​ log ⁡ w ) \mathcal{O}(\Delta\log{w}) while the internal dimension remains fixed at w w . ∎

Using the outline in the proof sketch, we now delve into a detailed analysis of the arithmetic circuit 𝒞 \mathcal{C} to an equivalent BaseConv model.

###### Proof of Theorem H.21 .

Let 𝒖 ∈ ℝ n ​ d \bm{u}\in\mathbbm{R}^{nd} be the input to the arithmetic circuit 𝒞 \mathcal{C} with depth Δ \Delta and width w w . We begin by rearranging the circuit into layers of addition and multiplication gates. That is, each layer 𝒞 ℓ \mathcal{C}^{\ell} has either all addition gates or all multiplication gates. This allows us to readily apply the results from Section H.4 . Note here that we can assert that the number of such layers L 𝒞 = 𝒪 ⁡ ( Δ ) L_{\mathcal{C}}=\mathcal{O}(\Delta) . Moreover, we construe the input to each such circuit as a vector of length w {w} by appending with extra 0s if necessary so that the composition of the layers results in a circuit equivalent to 𝒞 \mathcal{C} . See Fig. 8 for an example of such decomposition

Let 𝒛 ℓ ∈ ℝ w \bm{z}^{\ell}\in\mathbbm{R}^{w} denote the input to each layer 𝒞 ℓ \mathcal{C}_{\ell} for ℓ ∈ [ L 𝒞 ] \ell\in[L_{\mathcal{C}}] with 𝒛 1 ≡ ( 𝒖 , 𝟎 w − n ​ d ) \bm{z}^{1}\equiv(\bm{u},\bm{0}^{w-nd}) . It is important to note here that we may not have 𝒛 ℓ + 1 ≡ 𝒞 ℓ ​ ( 𝒛 ℓ ) \bm{z}^{\ell+1}\equiv\mathcal{C}_{\ell}(\bm{z}^{\ell}) since the inputs for gates at the ( ℓ + 1 ) (\ell+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 𝒞 ℓ lin \mathcal{C}_{\ell}^{\mathrm{lin}} denote an arbitrary linear layer which only contains the addition gates that takes in 𝒛 ℓ \bm{z}^{\ell} as input. We know that there exists a matrix 𝐌 ∈ ℝ w × w ′ {\bf M}\in\mathbbm{R}^{w\times w^{\prime}} such that we have 𝒞 ℓ lin ​ ( 𝒛 ℓ ) ≡ 𝐌 ​ 𝒛 ℓ . \mathcal{C}_{\ell}^{\mathrm{lin}}(\bm{z}^{\ell})\equiv{\bf M}\bm{z}^{\ell}. Here, note here that we may need entries from the vector 𝒛 ℓ \bm{z}^{\ell} in subsequent layers. Let s s be the total number of such entries, then for at each index for such entries i ∈ [ w ] i\in[w] , we have the corresponding index i s ∈ [ s ] i_{s}\in[s] . We then append the i i th standard row vector into the matrix 𝐌 {\bf M} to get the matrix 𝐌 ′ ∈ ℝ w × ( w ′ + s ) {\bf M}^{\prime}\in\mathbbm{R}^{w\times(w^{\prime}+s)} so that we have ( 𝐌 ′ ​ 𝒛 ℓ ) ​ [ j ] = { ( 𝐌 ​ 𝒛 ℓ ) ​ [ j ] if ​ j < w ′ 𝒛 ℓ ​ [ i ] if ​ j = w ′ + i s \left({\bf M}^{\prime}\bm{z}^{\ell}\right)[j]=\begin{cases}\left({\bf M}\bm{z}^{\ell}\right)[j]&\text{if }j<w^{\prime}\\ \bm{z}^{\ell}[i]&\text{if }j=w^{\prime}+i_{s}\end{cases} Here, note that we must have w ′ + s ≤ w w^{\prime}+s\leq w as w w is the width of the circuit. If needed, we then append the matrix 𝐌 ′ {\bf M}^{\prime} with all-zero rows 𝟎 w \bm{0}^{w} so that we get the matrix 𝐌 ′′ ∈ ℝ w × w {\bf M}^{\prime\prime}\in\mathbbm{R}^{w\times w} . Note here that we have thus preserved all entries needed in subsequent layers and incorporated the output of the ℓ \ell th layer with the output 𝐌 ′′ ​ 𝒛 ℓ ≡ 𝒛 ℓ + 1 {\bf M}^{\prime\prime}\bm{z}^{\ell}\equiv\bm{z}^{\ell+1} serving as the input of the ( ℓ + 1 ) (\ell+1) -th layer. Applying Corollary H.20 then yields an equivalent ( n , 𝒪 ⁡ ( log ⁡ w ) , d , w , d ) − BaseConv \left(n,\mathcal{O}(\log{w}),d,w,d\right)-\text{{BaseConv}} model.

##### Multiplication Gates.

Next, we deal with the layer 𝒞 mult ℓ \mathcal{C}^{\ell}_{\mathrm{mult}} of multiplication gates by collecting the two inputs to each multiplication gates as 𝒗 1 ℓ , 𝒗 2 ℓ \bm{v}_{1}^{\ell},\bm{v}_{2}^{\ell} . Note that for the input 𝒛 ℓ ∈ ℝ w \bm{z}^{\ell}\in\mathbbm{R}^{w} from the previous layer, we will again have entries that need to be used in the subsequent layers that we will denote as 𝒗 3 ℓ \bm{v}_{3}^{\ell} . We thus need to permute the entries of 𝒛 ℓ \bm{z}^{\ell} to get the vector [ 𝒗 1 ℓ : : 𝒗 2 ℓ : : 𝒗 3 ℓ ] [\bm{v}_{1}^{\ell}::\bm{v}_{2}^{\ell}::\bm{v}_{3}^{\ell}] so that we can remember 𝒗 3 ℓ \bm{v}_{3}^{\ell} while taking the Hadamard product 𝒗 1 ℓ ⊙ 𝒗 2 ℓ \bm{v}_{1}^{\ell}\odot\bm{v}_{2}^{\ell} . To this end, we can achieve this permutation of entries using a 𝒪 ⁡ ( w ) \mathcal{O}(w) -sparse linear matrix 12 12 12 Each row of 𝐑 ℓ {\bf R}^{\ell} has exactly one non-zero entry. 𝐑 ℓ ∈ ℝ w × w {\bf R}^{\ell}\in\mathbbm{R}^{w\times w} with which is equivalently represented by an ( w , w , 1 , w ) (w,w,1,w) -linear arithmetic circuit that simply moves the appropriate wires from the previous layer. This can again be achieved by a equivalent ( n , 𝒪 ⁡ ( log ⁡ w ) , d , w , d ) − BaseConv \left(n,\mathcal{O}(\log{w}),d,w,d\right)-\text{{BaseConv}} model. That is, 𝐑 ℓ ​ 𝒛 ℓ {\bf R}^{\ell}\bm{z}^{\ell} has the following form: 𝐑 ℓ ​ 𝒛 ℓ ≡ ( 𝒗 1 ℓ 𝒗 2 ℓ 𝒗 3 ℓ ) . {\bf R}^{\ell}\bm{z}^{\ell}\equiv\begin{pmatrix}\bm{v}_{1}^{\ell}\\ \hline\cr\\ \bm{v}_{2}^{\ell}\\ \hline\cr\\ \bm{v}_{3}^{\ell}\end{pmatrix}. Next, we can now define a ( n , 1 , d , n , d ) − BaseConv \left(n,1,d,n,d\right)-\text{{BaseConv}} which extracts 𝒗 1 ℓ \bm{v}_{1}^{\ell} as the projection with the input to the remember primitive given by 𝒚 ℓ ≡ ( 𝒗 2 ℓ , 𝒗 3 ℓ , 𝟎 ) \bm{y}^{\ell}\equiv(\bm{v}_{2}^{\ell},\bm{v}_{3}^{\ell},\bm{0}) . We then specify the remember primitive as 𝚛𝚎𝚖𝚎𝚖𝚋𝚎𝚛 ( 𝒚 ℓ : 𝒗 2 ℓ , 𝒗 3 ℓ , 𝒆 0 , 𝒗 1 ℓ ) {\tt remember}\left(\bm{y}^{\ell}:\bm{v}_{2}^{\ell},\bm{v}_{3}^{\ell},\bm{e}_{0},\bm{v}_{1}^{\ell}\right) which computes the following 𝒛 ℓ + 1 \displaystyle\bm{z}^{\ell+1} ≡ ( 𝒗 1 ℓ ⊙ ( 𝒆 0 ∗ 𝒗 2 ℓ ) 𝒗 3 ℓ 𝟎 ) \displaystyle\equiv\begin{pmatrix}\bm{v}_{1}^{\ell}\odot(\bm{e}_{0}\ast\bm{v}_{2}^{\ell})\\ \hline\cr\\ \bm{v}_{3}^{\ell}\\ \hline\cr\\ \bm{0}\end{pmatrix} ≡ ( 𝒗 1 ℓ ⊙ 𝒗 2 ℓ 𝒗 3 ℓ 𝟎 ) \displaystyle\equiv\begin{pmatrix}\bm{v}_{1}^{\ell}\odot\bm{v}_{2}^{\ell}\\ \hline\cr\\ \bm{v}_{3}^{\ell}\\ \hline\cr\\ \bm{0}\end{pmatrix} ≡ 𝒞 mult ℓ ​ ( 𝒛 ℓ ) ∈ ℝ w . \displaystyle\equiv\mathcal{C}^{\ell}_{\mathrm{mult}}(\bm{z}^{\ell})\in\mathbbm{R}^{w}. Using Proposition H.10 , we know that this requires a ( n , 𝒪 ⁡ ( log ⁡ w ) , d , w , d ) − BaseConv \left(n,\mathcal{O}(\log{w}),d,w,d\right)-\text{{BaseConv}} which remembers the entries that will serve as inputs in subsequent layers while performing the Hadamard product 𝒗 1 ⊙ 𝒗 2 \bm{v}^{1}\odot\bm{v}^{2} .

Overall, we can now stack all the resulting BaseConv layers and compose them as in Lemma H.11 to simulate 𝒞 \mathcal{C} . Overall, the number of layers blows up by 𝒪 ⁡ ( Δ ​ log ⁡ w ) \mathcal{O}(\Delta\log{w}) while the internal dimension remains fixed at w w . ∎

### H.6 The Retnet Reduction and Space Complexity for AR

We are now in a position to show that there exists an equivalent BaseConv model that can simulate a Retnet layer from Algorithm 6 . Recall that for an input sequence 𝒖 ∈ ℝ N × d \bm{u}\in\mathbbm{R}^{N\times d} , we project the input using weight matrices 𝐖 A , 𝐖 C , 𝐖 V ∈ ℝ d × d {\bf W}_{A},{\bf W}_{C},{\bf W}_{V}\in\mathbbm{R}^{d\times d} to 𝐀 , 𝐂 , 𝐕 ∈ ℝ N × d {\bf A},{\bf C},{\bf V}\in\mathbbm{R}^{N\times d} and compute the following recurrence (see line 6 in Algorithm 6 ): 𝒛 n ≡ γ 𝒛 n − 1 + ( 𝐀 [ n , : ] ) ⊤ 𝐕 [ n , : ] , \bm{z}^{n}\equiv\gamma\bm{z}^{n-1}+\left({\bf A}[n,:]\right)^{\top}{\bf V}[n,:], (6) where 𝒛 n ∈ ℝ d × d \bm{z}^{n}\in\mathbbm{R}^{d\times d} . We unroll this recurrence to express each state 𝒛 n \bm{z}^{n} in terms of 𝒖 \bm{u} with coefficients given by γ \gamma and the weight matrices 𝐖 A , 𝐖 C , 𝐖 V {\bf W}_{A},{\bf W}_{C},{\bf W}_{V} as follows: 𝒛 0 \displaystyle\bm{z}^{0} ≡ ( ( 𝒖 𝐖 A ) [ 0 , : ] ) ⊤ ( 𝒖 𝐖 V ) [ 0 , : ] \displaystyle\equiv\left(\left(\bm{u}{\bf W}_{A}\right)[0,:]\right)^{\top}\left(\bm{u}{\bf W}_{V}\right)[0,:] ≡ ( 𝒖 [ 0 , : ] 𝐖 A ) ⊤ ( 𝒖 [ 0 , : ] ) 𝐖 V \displaystyle\equiv\left(\bm{u}[0,:]{\bf W}_{A}\right)^{\top}\left(\bm{u}[0,:]\right){\bf W}_{V} ≡ 𝐖 A ⊤ ( 𝒖 [ 0 , : ] ) ⊤ ( 𝒖 [ 0 , : ] ) 𝐖 V \displaystyle\equiv{\bf W}_{A}^{\top}\left(\bm{u}[0,:]\right)^{\top}\left(\bm{u}[0,:]\right){\bf W}_{V} ≡ 𝐖 A ⊤ ( 𝒖 ⊤ [ : , 0 ] 𝒖 [ 0 , : ] ) 𝐖 V . \displaystyle\equiv{\bf W}_{A}^{\top}\left(\bm{u}^{\top}[:,0]\bm{u}[0,:]\right){\bf W}_{V}. Similarly, we have 𝒛 1 \displaystyle\bm{z}^{1} ≡ γ 𝒛 0 + 𝐖 A ⊤ ( 𝒖 ⊤ [ : , 1 ] 𝒖 [ 1 , : ] ) 𝐖 V \displaystyle\equiv\gamma\bm{z}^{0}+{\bf W}_{A}^{\top}\left(\bm{u}^{\top}[:,1]\bm{u}[1,:]\right){\bf W}_{V} ≡ γ ( 𝐖 A ⊤ ( 𝒖 ⊤ [ : , 1 ] 𝒖 [ 0 , : ] ) 𝐖 V ) + 𝐖 A ⊤ ( 𝒖 ⊤ [ : , 1 ] 𝒖 [ 1 , : ] ) 𝐖 V \displaystyle\equiv\gamma\left({\bf W}_{A}^{\top}\left(\bm{u}^{\top}[:,1]\bm{u}[0,:]\right){\bf W}_{V}\right)+{\bf W}_{A}^{\top}\left(\bm{u}^{\top}[:,1]\bm{u}[1,:]\right){\bf W}_{V} ≡ 𝐖 A ⊤ ( γ ( 𝒖 ⊤ [ : , 0 ] 𝒖 [ 0 , : ] ) + 𝒖 ⊤ [ : , 1 ] 𝒖 [ 1 , : ] ) 𝐖 V \displaystyle\equiv{\bf W}_{A}^{\top}\left(\gamma\left(\bm{u}^{\top}[:,0]\bm{u}[0,:]\right)+\bm{u}^{\top}[:,1]\bm{u}[1,:]\right){\bf W}_{V} We can then generalize this to n ∈ [ N ] n\in[N] to get 𝒛 n ≡ 𝐖 A ⊤ ( ∑ i = 0 n γ n − i 𝒖 ⊤ [ : , i ] 𝒖 [ i , : ] ) 𝐖 V . \bm{z}^{n}\equiv{\bf W}_{A}^{\top}\left(\sum_{i=0}^{n}\gamma^{n-i}\bm{u}^{\top}[:,i]\bm{u}[i,:]\right){\bf W}_{V}. (7) The above helps us infer that 𝒛 n \bm{z}^{n} can be expressed as a polynomial in { 𝒖 ⁡ [ 0 , 0 ] , 𝒖 ⁡ [ 0 , 1 ] , … , 𝒖 ⁡ [ N − 1 , d − 1 ] } \{\bm{u}[0,0],\bm{u}[0,1],\ldots,\bm{u}[N-1,d-1]\} . For each polynomial, there exists an arithmetic circuit that computes the polynomial in a natural way, whence we can apply Theorem H.21 to get an equivalent BaseConv model.

###### Corollary H.22 .

For a RetNet model with 𝒪 ⁡ ( d 2 ) \mathcal{O}(d^{2}) parameters and N N layers, there exists an equivalent BaseConv model that uses 𝒪 ⁡ ( N ​ d ) \mathcal{O}(Nd) parameters and 𝒪 ⁡ ( N ​ log ⁡ d ) \mathcal{O}(N\log{d}) layers.

###### Proof.

We will start by describing the arithmetic circuit that computes the states 𝒛 n ∈ ℝ d × d \bm{z}^{n}\in\mathbbm{R}^{d\times d} here. Indeed, each state requires exactly two alternative layers with only multiplication gates and only addition gates, respectively. First, we note that for each n ∈ N n\in N , we have 𝒛 n ≡ γ 𝒛 n − 1 + ( 𝐀 [ n , : ] ) ⊤ 𝐕 [ n , : ] . \bm{z}^{n}\equiv\gamma\bm{z}^{n-1}+\left({\bf A}[n,:]\right)^{\top}{\bf V}[n,:]. Since we can compute 𝒛 0 \bm{z}^{0} in 𝒪 ⁡ ( 1 ) \mathcal{O}(1) layers and computing γ ⋅ 𝒛 n − 1 \gamma\cdot\bm{z}^{n-1} needs exactly one layer with each entry from 𝒛 n − 1 \bm{z}^{n-1} serving as an input to a multiplication gate along with γ \gamma , the depth of the circuit is 𝒪 ⁡ ( N ) \mathcal{O}(N) with width d 2 d^{2} . We can thus apply Theorem H.21 to get an ( 𝒪 ⁡ ( N ) , 𝒪 ⁡ ( N ​ log ⁡ d ) , d 2 , 𝒪 ⁡ ( d 2 ) , 𝒪 ⁡ ( d 2 ) ) − (\mathcal{O}(N),\mathcal{O}(N\log{d}),d^{2},\mathcal{O}(d^{2}),\mathcal{O}(d^{2}))- BaseConv model that can compute 𝒛 n \bm{z}^{n} for each n ∈ N n\in N . ∎

##### 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 { 𝒌 i , 𝒗 i } i = 0 n − 1 \{\bm{k}_{i},\bm{v}_{i}\}_{i=0}^{n-1} along with a query 𝒒 \bm{q} appended at the end as input and the goal is to output 𝒗 i \bm{v}_{i} if 𝒒 = 𝒌 i \bm{q}=\bm{k}_{i} for some i ∈ [ 0 , N − 1 ] i\in[0,N-1] . Indeed, we will first provide a lower-bound for any model that purports to 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 𝒙 ∈ { 0 , 1 } n \bm{x}\in\{0,1\}^{n} and Bob has an index i ∈ [ n ] i\in[n] , and the goal for the players is to output the i i -th entry 𝒙 i \bm{x}_{i} . Moreover, we also require the communication to be one-way : only Alice is allowed to send a single message to Bob and Bob needs to output the answer. We will make use of the following lower-bound result.

###### Theorem H.23 ( [ Jayram et al., 2008 ] ) .

The one-way randomized communication complexity 13 13 13 The randomized communication complexity of function f f is defined as min π ⁡ ‖ π ‖ \min_{\pi}\left\lVert\pi\right\rVert , where π \pi ranges over all randomized protocols that can solve f f with probability of success at least 2 / 3 2/3 . of the index problem for sending an n n -length bit string is Ω ⁡ ( n ) \Omega(n) .

We now use Theorem H.23 to provide a lower bound on the number of bits required by the Retnet model to solve AR.

###### Corollary H.24 .

The RetNet model requires Ω ⁡ ( N ) \Omega(N) -bits to solve AR for d ≤ N d\leq\sqrt{N} .

###### Proof.

Consider an instance ( 𝒙 , i ) (\bm{x},i) of the index problem with 𝒙 ∈ { 0 , 1 } N \bm{x}\in\{0,1\}^{N} . We now describe the corresponding instance of the AR problem: { i , 𝒙 i } i = 0 N − 1 , i . \{i,\bm{x}_{i}\}_{i=0}^{N-1},i. (8) Next, consider the following one-way protocol for solving the index problem using the RetNet model. Alice with their access of 𝒙 ∈ { 0 , 1 } N \bm{x}\in\{0,1\}^{N} generate an input for AR (without the query) as in equation 8 . Alice then runs the RetNet model on { i , 𝒙 i } i = 0 N − 1 \{i,\bm{x}_{i}\}_{i=0}^{N-1} and sends the memory content of running the RetNet model to Bob. This should include the output 𝒛 N − 1 \bm{z}^{N-1} as we can reasonably assume that both have access weight matrices 𝐖 A , 𝐖 C , 𝐖 V {\bf W}_{A},{\bf W}_{C},{\bf W}_{V} and the scalar γ \gamma . Since we assume that this model solves AR, the output 𝙾𝚞𝚝 [ N , : ] = 𝒙 i {\tt Out}[N,:]=\bm{x}_{i} should contain the associated value of i i . Here, Bob can compute 𝙾𝚞𝚝 [ N , : ] {\tt Out}[N,:] by using the memory content sent by Alice along with the term γ ​ 𝒛 N − 1 \gamma\bm{z}^{N-1} : 𝒙 i = 𝙾𝚞𝚝 [ N , : ] = 𝐂 [ N , : ] 𝒛 N = 𝐂 [ N , : ] ( γ 𝒛 N − 1 + ( 𝐀 [ N , : ] ) ⊤ 𝐕 [ N , : ] ) . \bm{x}_{i}={\tt Out}[N,:]={\bf C}[N,:]\bm{z}^{N}={\bf C}[N,:]\left(\gamma\bm{z}^{N-1}+({\bf A}[N,:])^{\top}{\bf V}[N,:]\right). That is, the total number of bits that are communicated in this protocol is O ⁡ ( d 2 ) O(d^{2}) . For d ≤ N d\leq\sqrt{N} , have shown that a one-way communication protocol exists for solving the index problem exists that uses o ⁡ ( N ) o(N) communication complexity. This contradicts Theorem H.23 and hence, we conclude that the ReNet model solving AR also needs Ω ⁡ ( N ) \Omega(N) bits. ∎

### H.7 The Multiple-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 Definition 3.1 to a slightly more general problem: Suppose we are given an input sequence 𝒖 [ 0 ⋯ N − 1 ] ≜ { ( 𝒌 0 , 𝒗 0 , 𝒒 0 ) , … , ( 𝒌 N 3 − 1 , 𝒗 N 3 − 1 , 𝒒 N 3 − 1 ) } \bm{u}[0\cdots N-1]\triangleq\left\{\left(\bm{k}_{0},\bm{v}_{0},\bm{q}_{0}\right),\ldots,\left(\bm{k}_{\frac{N}{3}-1},\bm{v}_{\frac{N}{3}-1},\bm{q}_{\frac{N}{3}-1}\right)\right\} with each 𝒌 i , 𝒗 i , 𝒒 i ∈ C \bm{k}_{i},\bm{v}_{i},\bm{q}_{i}\in C is a token drawn from a vocabulary of size c = | C | c=|C| . Our goal is then to check, for each 1 ≤ i ≤ N 3 − 1 1\leq i\leq\frac{N}{3}-1 , whether there exists 0 ≤ j < i 0\leq j<i such that 𝒒 i ≡ 𝒌 j \bm{q}_{i}\equiv\bm{k}_{j} , and if so, output 𝒗 j \bm{v}_{j} . Here, we note that it suffices to have d ≈ log ⁡ ( c ) d\approx\log(c) so that 𝒌 i , 𝒗 i , 𝒒 i \bm{k}_{i},\bm{v}_{i},\bm{q}_{i} is embedded in { 0 , 1 } d \{0,1\}^{d} . However, we will specify the specific embedding being used for the results below. Here, we construe the tokens 𝒌 i , 𝒒 i \bm{k}_{i},\bm{q}_{i} and 𝒗 i \bm{v}_{i} to be the keys , the queries , and the associated values . Indeed, it might be helpful to think of the input 𝒖 \bm{u} as a 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 a sequence of length N 3 \frac{N}{3} : 𝒖 [ 0 ⋯ N − 1 ] := { 𝒙 0 , … , 𝒙 N − 1 } \bm{u}[0\cdots N-1]:=\{\bm{x}_{0},\ldots,\bm{x}_{N-1}\} , where each 𝒙 i ∈ C \bm{x}_{i}\in C . The goal of Definition 3.1 is then to check, for each 1 ≤ i < N − 1 1\leq i<N-1 , whether there exists 0 ≤ j < i 0\leq j<i such that 𝒙 i ≡ 𝒙 j \bm{x}_{i}\equiv\bm{x}_{j} , and if so, output 𝒙 j + 1 \bm{x}_{j+1} , and continue otherwise. We can reduce this problem to the above general formulation by taking the following sequence of tuples as the input { ( 𝒙 i , 𝒙 i + 1 , 𝒙 i ) } \{(\bm{x}_{i},\bm{x}_{i+1},\bm{x}_{i})\} .

###### 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 it 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 𝒖 ∈ { 0 , 1 } N × d \bm{u}\in\{0,1\}^{N\times d} has designated indices for the keys, queries, and values in the sequence. We gather these indices below: 𝒦 \displaystyle\mathcal{K} = { i ∈ { 0 , … , N − 1 } | i ≡ 0 mod 3 } , \displaystyle=\{i\in\{0,\ldots,N-1\}\lvert\ i\equiv 0\mod{3}\}, (9) 𝒱 \displaystyle\mathcal{V} = { i ∈ { 0 , … , N − 1 } | i ≡ 1 mod 3 } , \displaystyle=\{i\in\{0,\ldots,N-1\}\lvert\ i\equiv 1\mod{3}\}, 𝒬 \displaystyle\mathcal{Q} = { i ∈ { 0 , … , N − 1 } | i ≡ 2 mod 3 } , . \displaystyle=\{i\in\{0,\ldots,N-1\}\lvert\ i\equiv 2\mod{3}\},. The input 𝒖 ∈ ℝ N × d \bm{u}\in\mathbbm{R}^{N\times d} to Attention for d = 3 ​ c d=3c is then given by 𝒖 [ i , : ] ≡ { [ 𝒌 i : 𝟎 c : 𝟎 c ] if ​ i ∈ 𝒦 [ 𝟎 c : 𝒗 i : 𝟎 c ] if ​ i ∈ 𝒱 [ 𝟎 c : 𝟎 c : 𝒒 i ] if ​ i ∈ 𝒬 \bm{u}[i,:]\equiv\begin{cases}[\bm{k}_{i}:\bm{0}^{c}:\bm{0}^{c}]&\text{ if }i\in\mathcal{K}\\ [\bm{0}^{c}:\bm{v}_{i}:\bm{0}^{c}]&\text{ if }i\in\mathcal{V}\\ [\bm{0}^{c}:\bm{0}^{c}:\bm{q}_{i}]&\text{ if }i\in\mathcal{Q}\end{cases} Here, each 𝒌 i , 𝒗 i , 𝒒 i \bm{k}_{i},\bm{v}_{i},\bm{q}_{i} is embedded as a one-hot encoding in { 0 , 1 } c \{0,1\}^{c} .

Without softmax, the output for an attention layer 𝐎 ∈ ℝ N × d {\bf O}\in\mathbbm{R}^{N\times d} is given by 𝐎 ≡ ( 𝐐𝐊 ⊤ ) ​ 𝐕 , {\bf O}\equiv\left({\bf Q}{\bf K}^{\top}\right){\bf V}, (10) where 𝐐 , 𝐊 , 𝐕 ∈ ℝ N × d {\bf Q},{\bf K},{\bf V}\in\mathbbm{R}^{N\times d} are defined as 𝒖 ​ 𝐖 Q , 𝒖 ​ 𝐖 K , 𝒖 ​ 𝐖 V \bm{u}{\bf W}_{Q},\bm{u}{\bf W}_{K},\bm{u}{\bf W}_{V} for 𝒖 ∈ ℝ N × d \bm{u}\in\mathbbm{R}^{N\times d} . Instead of position embeddings, we use ALiBi, a popular technique that biases the attention scores 𝐐𝐊 ⊤ {\bf Q}{\bf K}^{\top} with a lower-triangular Toeplitz matrix 𝐁 ∈ ℝ N × N {\bf B}\in\mathbb{R}^{N\times N} [ Press et al., 2021 ] . The values in this matrix are controlled by a fixed hyperparameter so they do not count towards the number of parameters in the model.

###### Proposition H.27 .

Given an input 𝐮 ∈ { 0 , 1 } N × d \bm{u}\in\{0,1\}^{N\times d} (encoded as in Remark H.26 ) where d = 3 ​ c d=3c , Attention with linear biases (even without using soft-max) solves Mqar for 𝐮 \bm{u} using 𝒪 ⁡ ( c 2 ) {\mathcal{O}}(c^{2}) parameters, 𝒪 ⁡ ( N ​ c 2 + N 2 ​ c ) {\mathcal{O}}(Nc^{2}+N^{2}c) time complexity and 𝒪 ⁡ ( 1 ) {\mathcal{O}}(1) layers.

###### Proof.

We use two layers of attention. We will start by specifying the projection matrices for the first layer 𝐖 Q 1 , 𝐖 K 1 , 𝐖 V 1 ∈ ℝ d × d {\bf W}_{Q}^{1},{\bf W}^{1}_{K},{\bf W}^{1}_{V}\in\mathbbm{R}^{d\times d} as: 𝐖 K 1 ≡ 𝐖 Q 1 ≡ 𝟎 , 𝐖 V 1 ≡ ( 𝟎 𝟎 𝟎 𝟎 𝐈 𝐜 × 𝐜 𝟎 𝟎 𝟎 𝟎 ) {\bf W}^{1}_{K}\equiv{\bf W}^{1}_{Q}\equiv{\bf 0},\quad{\bf W}_{V}^{1}\equiv\begin{pmatrix}{\bf 0}&{\bf 0}&{\bf 0}\\ {\bf 0}&{\bf I_{c\times c}}&{\bf 0}\\ {\bf 0}&{\bf 0}&{\bf 0}\end{pmatrix}

Above, 𝐖 𝐕 𝟏 {\bf W_{V}^{1}} is meant to isolate the 𝒗 i \bm{v}_{i} embeddings. For the first layer, we then have 𝐐 1 ≡ 𝐊 1 ≡ 𝟎 {\bf Q}^{1}\equiv{\bf K}^{1}\equiv{\bf 0} , and 𝐕 1 [ i , : ] := 𝒖 1 [ i , : ] 𝐖 V 1 ≡ { [ 𝟎 c : 𝒗 i : 𝟎 c ] if ​ i ∈ 𝒱 , 𝟎 d otherwise , {\bf V}^{1}[i,:]:=\bm{u}^{1}[i,:]{\bf W}_{V}^{1}\equiv\begin{cases}[\bm{0}^{c}:\bm{v}_{i}:\bm{0}^{c}]&\text{if }i\in\mathcal{V},\\ {\bf 0}^{d}&\text{otherwise}\end{cases}, where 𝒦 , 𝒬 , 𝒱 \mathcal{K},\mathcal{Q},\mathcal{V} are defined as in equation 9 . The output for the first layer is given by the following: 𝐎 1 [ i , : ] \displaystyle{\bf O}^{1}[i,:] = ( ( 𝐐𝐊 ⊤ + 𝐁 1 ) 𝐕 ) [ i , : ] \displaystyle=\left(\left({\bf Q}{\bf K}^{\top}+{\bf B}^{1}\right){\bf V}\right)[i,:] = 𝐁 1 𝐕 [ i , : ] \displaystyle={\bf B}^{1}{\bf V}[i,:] = { [ 𝟎 c : 𝒗 i : 𝟎 c ] if ​ i ∈ 𝒦 𝟎 d otherwise , \displaystyle=\begin{cases}[\bm{0}^{c}:\bm{v}_{i}:\bm{0}^{c}]&\text{if }i\in\mathcal{K}\\ \bm{0}^{d}&\text{otherwise}\end{cases}, The 𝐐𝐊 𝐓 {\bf QK^{T}} is ignored ( ≡ 𝟎 \equiv{\bf 0} ) and we isolate the shifted sequence by setting the bias matrix 𝐁 {\bf B} appropriately. In particular, the last equality follows from the fact that 𝐁 1 {\bf B}^{1} is an up-shift matrix that shifts each row of 𝐕 {\bf V} by 1. We apply the residual connection at the end of the standard Transformer block to insert the 𝒌 i \bm{k}_{i} adjacent to the 𝒗 i \bm{v}_{i} . For the second layer, the input 𝒖 2 ∈ ℝ N × d \bm{u}^{2}\in\mathbbm{R}^{N\times d} is given by 𝒖 2 [ i , : ] ≡ { [ 𝒌 i : 𝒗 i : 𝟎 c ] if ​ i ∈ 𝒦 [ 𝟎 c : 𝒗 i : 𝟎 c ] if ​ i ∈ 𝒱 [ 𝟎 c : 𝟎 c : 𝒒 i ] if ​ i ∈ 𝒬 , \bm{u}^{2}[i,:]\equiv\begin{cases}[\bm{k}_{i}:\bm{v}_{i}:\bm{0}^{c}]&\text{if }i\in\mathcal{K}\\ [\bm{0}^{c}:\bm{v}_{i}:\bm{0}^{c}]&\text{if }i\in\mathcal{V}\\ [\bm{0}^{c}:\bm{0}^{c}:\bm{q}_{i}]&\text{if }i\in\mathcal{Q}\end{cases}, Further, we take the following projection matrices: 𝐖 K 2 ≡ ( 𝐈 c × c 𝟎 𝟎 𝟎 𝟎 𝟎 𝟎 𝟎 𝟎 ) , 𝐖 Q 2 ≡ ( 𝟎 𝟎 𝟎 𝟎 𝟎 𝟎 𝐈 c × c 𝟎 𝟎 ) , 𝐖 V 2 ≡ ( 𝟎 𝟎 𝟎 𝐈 c × c 𝟎 𝟎 𝟎 𝟎 𝟎 ) {\bf W}^{2}_{K}\equiv\begin{pmatrix}{\bf I}_{c\times c}&{\bf 0}&{\bf 0}\\ {\bf 0}&{\bf 0}&{\bf 0}\\ {\bf 0}&{\bf 0}&{\bf 0}\end{pmatrix},{\bf W}^{2}_{Q}\equiv\begin{pmatrix}{\bf 0}&{\bf 0}&{\bf 0}\\ {\bf 0}&{\bf 0}&{\bf 0}\\ {\bf I}_{c\times c}&{\bf 0}&{\bf 0}\end{pmatrix},{\bf W}^{2}_{V}\equiv\begin{pmatrix}{\bf 0}&{\bf 0}&{\bf 0}\\ {\bf I}_{c\times c}&{\bf 0}&{\bf 0}\\ {\bf 0}&{\bf 0}&{\bf 0}\end{pmatrix} We then have the following matrices as input to attention after applying projection: 𝐐 2 [ i , : ] \displaystyle{\bf Q}^{2}[i,:] : = 𝒖 2 [ i , : ] 𝐖 2 K ≡ { [ 𝒒 i : 𝟎 2 ​ c ] if ​ i ∈ 𝒬 , 𝟎 d otherwise , \displaystyle:=\bm{u}^{2}[i,:]{\bf W}^{2}_{K}\equiv\begin{cases}[\bm{q}_{i}:\bm{0}^{2c}]&\text{if }i\in\mathcal{Q},\\ {\bf 0}^{d}&\text{otherwise}\end{cases}, (11) 𝐊 2 [ i , : ] \displaystyle{\bf K}^{2}[i,:] : = 𝒖 2 [ i , : ] 𝐖 2 Q ≡ { [ 𝒌 i : 𝟎 2 ​ c ] if ​ i ∈ 𝒦 , 𝟎 d otherwise , \displaystyle:=\bm{u}^{2}[i,:]{\bf W}^{2}_{Q}\equiv\begin{cases}[\bm{k}_{i}:\bm{0}^{2c}]&\text{if }i\in\mathcal{K},\\ {\bf 0}^{d}&\text{otherwise}\end{cases}, 𝐕 2 [ i , : ] \displaystyle{\bf V}^{2}[i,:] : = 𝒖 2 [ i , : ] 𝐖 2 V ≡ { [ 𝒗 i : 𝟎 2 ​ c ] if ​ i ∈ 𝒦 ∪ 𝒱 , 𝟎 d otherwise . \displaystyle:=\bm{u}^{2}[i,:]{\bf W}^{2}_{V}\equiv\begin{cases}[\bm{v}_{i}:\bm{0}^{2c}]&\text{if }i\in\mathcal{K}\cup\mathcal{V},\\ {\bf 0}^{d}&\text{otherwise}\end{cases}. Here, note that the values have been shifted to the corresponding key position. Next, we compute the term in the parenthesis for the second layer as ( 𝐐 2 ​ 𝐊 2 ⊤ ) ​ [ i , j ] = ( 𝐐 2 ​ 𝐊 2 ⊤ ) ​ [ i , j ] \displaystyle\left({\bf Q}^{2}{\bf K}^{2\top}\right)[i,j]=\left({\bf Q}^{2}{\bf K}^{2\top}\right)[i,j] = ⟨ 𝐐 2 [ i , : ] , 𝐊 2 ⊤ [ : , j ] ⟩ \displaystyle=\langle{\bf Q}^{2}[i,:],{\bf K}^{2\top}[:,j]\rangle = ⟨ 𝐐 2 [ i , : ] , 𝐊 2 [ j , : ] ⟩ \displaystyle=\langle{\bf Q}^{2}[i,:],{\bf K}^{2}[j,:]\rangle = { ⟨ 𝒒 i , 𝒌 j ⟩ if ​ i ∈ 𝒬 , j ∈ 𝒦 0 otherwise \displaystyle=\begin{cases}\langle\bm{q}_{i},\bm{k}_{j}\rangle&\text{if }i\in\mathcal{Q},j\in\mathcal{K}\\ 0&\text{otherwise}\end{cases} = { 1 if ​ i ∈ 𝒬 , j ∈ 𝒦 , 𝒒 i ≡ 𝒌 j ≡ 𝒆 k ​ for some k 0 otherwise . \displaystyle=\begin{cases}1&\text{if }i\in\mathcal{Q},j\in\mathcal{K},\bm{q}_{i}\equiv\bm{k}_{j}\equiv\bm{e}_{k}\text{ for some $k$}\\ 0&\text{otherwise}.\end{cases} Finally, we compute the output as follows: 𝐎 2 [ i , : ] \displaystyle{\bf O}^{2}[i,:] = ( ( 𝐐 2 𝐊 2 ⊤ + 𝐁 2 ) 𝐕 2 ) [ i , : ] \displaystyle=\left(\left({\bf Q}^{2}{\bf K}^{2\top}+{{\bf B}^{2}}\right){\bf V}^{2}\right)[i,:] = ( ( 𝐐 2 𝐊 2 ⊤ + 𝟎 ) 𝐕 2 ) [ i , : ] \displaystyle=\left(\left({\bf Q}^{2}{\bf K}^{2\top}+{{\bf 0}}\right){\bf V}^{2}\right)[i,:] = ( 𝐐 2 𝐊 2 ⊤ ) [ i , : ] ⋅ 𝐕 2 \displaystyle=\left({\bf Q}^{2}{\bf K}^{2\top}\right)[i,:]\cdot{\bf V}^{2} = ∑ j = 0 N − 1 ( 𝐐 2 𝐊 2 ⊤ ) [ i , j ] ⋅ 𝐕 2 [ j , : ] \displaystyle=\sum_{j=0}^{N-1}\left({\bf Q}^{2}{\bf K}^{2\top}\right)[i,j]\cdot{\bf V}^{2}[j,:] = ∑ j ∈ 𝒦 ( 𝐐 2 𝐊 2 ⊤ ) [ i , j ] ⋅ [ 𝒗 j : 𝟎 2 ​ c ] \displaystyle=\sum_{j\in\mathcal{K}}\left({\bf Q}^{2}{\bf K}^{2\top}\right)[i,j]\cdot[\bm{v}_{j}:\bm{0}^{2c}] = { [ 𝒗 j : 𝟎 2 ​ c ] if ​ j ∈ 𝒦 , i ∈ 𝒬 , 𝒒 i ≡ 𝒌 j 𝟎 d otherwise , \displaystyle=\begin{cases}[\bm{v}_{j}:\bm{0}^{2c}]&\text{if }j\in\mathcal{K},i\in\mathcal{Q},\bm{q}_{i}\equiv\bm{k}_{j}\\ \bm{0}^{d}&\text{otherwise}\end{cases}, where we use the fact that for each index j ∈ 𝒦 j\in\mathcal{K} , the matrix 𝐕 2 {\bf V}^{2} contains the associated value from equation 11 . Thus, for each query 𝒒 i \bm{q}_{i} , we solve the associated value problem yielding a match for the j j th key.

In total, we only need 𝒪 ⁡ ( c 2 ) \mathcal{O}(c^{2}) -many parameters to perform these multiplications and the linear bias in the first layer is a hyperparameter that is static and unlearned; the time complexity comes from the multiplication of 𝐐𝐊 ⊤ {\bf QK^{\top}} in 𝒪 ⁡ ( N 2 ​ c ) {\mathcal{O}}(N^{2}c) , and projections in 𝒪 ⁡ ( N ​ c 2 ) {\mathcal{O}}(Nc^{2}) . Finally, we only need 𝒪 ⁡ ( 1 ) \mathcal{O}(1) layers for this solution. ∎

In the sequel, we develop a parallel algorithm to solve the multiple-query associative recall problem with 𝒪 ⁡ ( N ​ d ⋅ log 2 ⁡ N ) \mathcal{O}(Nd\cdot\log^{2}{N}) work complexity and 𝒪 ⁡ ( d ⋅ log 2 ⁡ N ) \mathcal{O}(d\cdot\log^{2}{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 ) \tilde{\mathcal{O}}(1) layers and 𝒪 ~ ​ ( N ​ d ) \tilde{\mathcal{O}}(Nd) parameters.

#### H.7.3 Initial Attempt: A Sequential Algorithm

We will first discuss the algorithm that simply uses an associative array to solve the multiple-query associative recall problem. Specifically, we want to use a data structure that allows for logarithmic insertion and membership query. Here, we do not specify a choice but data structures including self-balancing binary search trees which allow for 𝒪 ⁡ ( log ⁡ N ⋅ d ) \mathcal{O}(\log{N}\cdot d) 𝚒𝚗𝚜𝚎𝚛𝚝 {\tt insert} and 𝚏𝚒𝚗𝚍 {\tt find} operations for d d -bit entries should be sufficient.

###### Proposition H.28 .

Algorithm 9 solves the multiple-query associative recall problem ( Mqar ) \left(\textsc{Mqar}\right) in 𝒪 ⁡ ( d ​ N ​ log ⁡ N ) \mathcal{O}(dN\log{N}) time for an input sequence 𝐮 ∈ { 0 , 1 } N × d \bm{u}\in\{0,1\}^{N\times d} .

###### Proof.

For any i ∈ { 0 , … , N 3 − 1 } i\in\{0,\ldots,\frac{N}{3}-1\} , we know that both insertion and lookup operations take 𝒪 ⁡ ( log ⁡ ( i ) ⋅ d ) \mathcal{O}(\log(i)\cdot d) time. Overall, the runtime of the algorithm is ∑ i = 0 N 3 − 1 𝒪 ⁡ ( log ⁡ ( i ) ⋅ d ) = 𝒪 ⁡ ( log ⁡ ( N ! ) ⋅ d ) = 𝒪 ⁡ ( N ​ log ​ N ⋅ d ) . \sum_{i=0}^{\frac{N}{3}-1}\mathcal{O}(\log(i)\cdot d)=\mathcal{O}(\log(N!)\cdot d)=\mathcal{O}(N\log{N}\cdot d). ∎

#### 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 to Algorithm 9 , it may be the case that the arithmetic circuit has a large number of layers Ω ⁡ ( N ) \Omega(N) . Unfortunately, this would imply that the resulting BaseConv model may have near quadratic complexity in N N . Instead, we now initiate our effort into designing a BaseConv model with both small enough number of parameters and number of layers. Here, we will first subdivide the problem using dyadic intervals into 𝒪 ⁡ ( N ) \mathcal{O}(N) subproblems and reduce each such subproblem into a multiple search problem [ Akl and Meijer, 1990 ] . To this end, we briefly introduce the multiple search problem below. Given two array of numbers A ≜ a 0 ≤ … ≤ a n − 1 A\triangleq a_{0}\leq\ldots\leq a_{n-1} and B ≜ ( b 0 ≤ … ≤ b m − 1 ) B\triangleq(b_{0}\leq\ldots\leq b_{m-1}) with n ≤ m n\leq m , for each a j ∈ A a_{j}\in A , the goal is to find the smallest element in B B that is larger than or equal to a j a_{j} . The multiple search problem is solved by a parallel binary search ( pbs ) algorithm in [ Akl and Meijer, 1990 ] with work complexity 𝒪 ⁡ ( n ⋅ log ⁡ m ) \mathcal{O}(n\cdot\log{m}) and time 𝒪 ⁡ ( log ⁡ n ​ log ⁡ m ) \mathcal{O}(\log{n}\log{m}) . Specifically, for sorted arrays A [ 0 ⋯ n − 1 ] A[0\cdots n-1] and B [ 0 ⋯ m − 1 ] B[0\cdots m-1] , 𝚙𝚋𝚜 {\tt pbs} constructs the array C [ 0 ⋯ n − 1 ] C[0\cdots n-1] defined as C ⁡ [ i ] ≜ { min 0 ≤ j < m { j | A [ i ] ≤ B [ j ] } if ​ A ​ [ i ] ≤ B ⁡ [ m − 1 ] m otherwise. C[i]\triangleq\begin{cases}\min_{0\leq j<m}\{j\lvert\ A[i]\leq B[j]\}&\text{if }A[i]\leq B[m-1]\\ m&\text{otherwise.}\end{cases} (12) The algorithm itself runs in exclusive-read exclusive-write (EREW) PRAM model—no two processors 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.

Let Σ \Sigma be the set { 0 , 1 } \{0,1\} and denote the set of binary strings of size n n as Σ n \Sigma^{n} . We define prefix ⁡ ( 𝒙 ) \mathrm{prefix}(\bm{x}) for n n -bit strings as the set of all initial substrings of 𝒙 ∈ Σ n \bm{x}\in\Sigma^{n} which includes the empty string and 𝒙 \bm{x} itself. Next, let dec : { 0 , 1 } n → ℕ \mathrm{dec}:\{0,1\}^{n}\to\mathbbm{N} be the decimal representation of an n n -bit string 𝒙 \bm{x} with 𝒙 ⁡ [ 0 ] \bm{x}[0] denoting the least significant bit. We also use 𝚜𝚘𝚛𝚝 ⁡ ( A ) {\tt sort}(A) as a procedure that sorts an array A A . Finally, wlog, we assume that N N is a power of 2 2 . We are now ready to present a parallel algorithm that solves the multiple-query associative recall problem below.

###### Remark H.29 .

In lines 8 and 11 , 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 17 .

#### H.7.5 Correctness and Complexity

###### Proposition H.30 .

Algorithm 11 solves the multiple-query associative recall problem with work complexity 𝒪 ⁡ ( N ​ d ⋅ log 2 ⁡ N ) \mathcal{O}(Nd\cdot\log^{2}{N}) and time 𝒪 ⁡ ( d ⋅ log 2 ⁡ N ) \mathcal{O}(d\cdot\log^{2}{N}) .

###### Proof.

The correctness of pbs implies the correctness of Algorithm 11 if we can show that, for each 1 ≤ i < N 3 1\leq i<\frac{N}{3} , we check for a match among the keys { 𝒌 j } j ∈ [ i − 1 ] \{\bm{k}_{j}\}_{j\in[i-1]} . To this end, for each 1 ≤ i < N 3 1\leq i<\frac{N}{3} , let the set of all iterator indices associated with an index i i be defined as K i ≜ { ( k , x ) | i ∈ J k 𝒙 } K_{i}\triangleq\{(k,x)\lvert i\in J_{k}^{\bm{x}}\} with J k 𝒙 J_{k}^{\bm{x}} as noted in line 10 . Then, we define the corresponding set for keys as ℐ i ≜ ⋃ ( k , 𝒙 ) ∈ K i I k 𝒙 ​ with I k 𝒙 s defined as in line 7 . \mathcal{I}_{i}\triangleq\bigcup_{(k,\bm{x})\in K_{i}}I_{k}^{\bm{x}}\text{ with $I_{k}^{\bm{x}}$s defined as in line \ref{step: i-def}}. That is, for all the calls to pbs-key-values that i i is part of (given by K i K_{i} ) where the algorithm checks for a match among the keys in ℐ i \mathcal{I}_{i} , it then suffices to show that ℐ i = [ i − 1 ] \mathcal{I}_{i}=[i-1] .

Here, first note that if some index j ∈ I k 𝒙 ⊆ ℐ i j\in I_{k}^{\bm{x}}\subseteq\mathcal{I}_{i} for some 𝒙 ∈ Σ log ⁡ ( N 3 ) − k \bm{x}\in\Sigma^{\log\left(\frac{N}{3}\right)-k} , then, by definition, 𝒙 ∈ prefix ⁡ ( bin ⁡ ( j ) ) \bm{x}\in\mathrm{prefix}(\mathrm{bin}(j)) . Here, let 𝒙 1 := 𝒙 | 𝒙 ⁡ [ log ⁡ ( N 3 ) − k ] = 1 \bm{x}^{1}:=\bm{x}\lvert_{\bm{x}[\log\left(\frac{N}{3}\right)-k]=1} where we set the ( log ⁡ ( N 3 ) − k ) (\log\left(\frac{N}{3}\right)-k) -th index of 𝒙 \bm{x} to be 1. Consequently, as we have i ∈ J k 𝒙 i\in J^{\bm{x}}_{k} for the same k k and 𝒙 \bm{x} as in I k 𝒙 I_{k}^{\bm{x}} ( cf. line 7 ), we must have 𝒙 1 ∈ prefix ⁡ ( bin ⁡ ( i ) ) \bm{x}^{1}\in\mathrm{prefix}(\mathrm{bin(i)}) . Thus we get j < i j<i , whence we can claim that ℐ i ⊆ [ i − 1 ] \mathcal{I}_{i}\subseteq[i-1] .

For the other direction, for any i i , let b b denote the position of the most significant bit in bin ⁡ ( i ) \mathrm{bin}(i) which differs from bin ⁡ ( j ) \mathrm{bin}(j) for any j ∈ [ i − 1 ] j\in[i-1] . Then, there must exist a binary string that is in the prefix set of both bin ⁡ ( i ) \mathrm{bin}(i) and bin ⁡ ( i ) \mathrm{bin}(i) . That is, there exists 𝒙 ∈ prefix ⁡ ( bin ⁡ ( i ) ) ∩ prefix ⁡ ( bin ⁡ ( j ) ) \bm{x}\in\mathrm{prefix}(\mathrm{bin}(i))\cap\mathrm{prefix}(\mathrm{bin}(j)) with 𝒙 ∈ Σ b \bm{x}\in\Sigma^{b} . Thus, we then must have bin ⁡ ( j ) ∈ I log ⁡ ( N 3 ) − b 𝒙 \mathrm{bin}(j)\in I^{\bm{x}}_{\log\left(\frac{N}{3}\right)-b} and bin ⁡ ( i ) ∈ J log ⁡ ( N 3 ) − b 𝒙 \mathrm{bin}(i)\in J^{\bm{x}}_{\log\left(\frac{N}{3}\right)-b} with 𝒙 \bm{x} as the corresponding witness. Hence, we have [ i − 1 ] ⊆ ℐ i [i-1]\subseteq\mathcal{I}_{i} .

Overall, we have shown that ℐ i = [ i − 1 ] \mathcal{I}_{i}=[i-1] . Since this holds for all 1 ≤ i < N 3 1\leq i<\frac{N}{3} , we can conclude that Algorithm 11 solves the multiple-query associative recall problem.

Next, it is easy to see that we execute lines 7 to 17 ∑ k = 0 log ⁡ ( N 3 ) − 1 N 3 2 k + 1 \sum_{k=0}^{\log\left(\frac{N}{3}\right)-1}\frac{\frac{N}{3}}{2^{k+1}} -many times. We note that sorting n n values each of size d d can be done with work complexity n ​ log ⁡ n ⋅ d n\log{n}\cdot d . We note that, at each instance, we are sorting sorting 2 k 2^{k} values. Meanwhile, remembering 2 k 2^{k} sorted permutation of indices can be done in linear time using arrays. Moreover, each call to pbs-key-values has n = m = 2 k n=m=2^{k} which has work complexity n ​ log ⁡ m n\log{m} . Finally, we know that the work complexity of lines 19 to 21 is 𝒪 ⁡ ( N ) \mathcal{O}(N) . Thus, the overall work complexity of Algorithm 11 is d ⋅ ∑ k = 0 log ⁡ ( N 3 ) − 1 N 3 2 k + 1 ​ 𝒪 ​ ( 2 k ⋅ log ⁡ 2 k ) = d ⋅ 𝒪 ⁡ ( N ) ⋅ ∑ k = 0 log ⁡ N / 3 − 1 𝒪 ⁡ ( k ) = 𝒪 ⁡ ( N ​ d ⋅ log 2 ⁡ N ) . d\cdot\sum_{k=0}^{\log\left(\frac{N}{3}\right)-1}\frac{\frac{N}{3}}{2^{k+1}}\mathcal{O}(2^{k}\cdot\log{2^{k}})=d\cdot\mathcal{O}(N)\cdot\sum_{k=0}^{\log{N/3}-1}\mathcal{O}(k)=\mathcal{O}(Nd\cdot\log^{2}{N}). (13)

We will now analyze the depth of Algorithm 11 . We know that the depth of computation for Algorithm 10 is 𝒪 ⁡ ( log ⁡ n ​ log ⁡ m ) \mathcal{O}(\log{n}\log{m}) for input sizes. Moreover, we have 𝒪 ⁡ ( 1 ) \mathcal{O}(1) depth for the computation in 19 to 21 as each entry in 𝚒𝚍𝚡 {\tt idx} can have at most one entry. Since the nested for loops iterating over k k s and the associated 𝒙 \bm{x} s runs in parallel, the depth of Algorithm 11 is dominated by the largest depth among all calls to pbs-key-values and to sort . The largest such call to pbs-key-values is of size n = m = 2 log ⁡ ( N 3 ) − 1 = N / 6 n=m=2^{\log\left(\frac{N}{3}\right)-1}=N/6 which yields a depth of d ⋅ log 2 ⁡ N 3 d\cdot\log^{2}{\frac{N}{3}} . Moreover, using sorting networks ( Definition H.34 ), we know that the largest depth is for sorting N 6 \frac{N}{6} values of size d d given by d ⋅ Θ ⁡ ( log ⁡ N ) d\cdot\Theta(\log{N}) ( Lemma H.35 ). Thus, we can conclude that Algorithm 11 takes 𝒪 ⁡ ( d ⋅ log 2 ⁡ N ) \mathcal{O}(d\cdot\log^{2}{N}) , time where N N is the length of the input. ∎

#### H.7.6 Conversion to Arithmetic Circuit

We will convert Algorithm 11 to an arithmetic circuits modularly. In particular, after writing out 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 Algorithm 10 as pbs-key-values as well with the input gates comprising of each entry from 𝒒 [ s ⋯ t ] \bm{q}[s\cdots t] and 𝒌 [ x ⋯ y ] \bm{k}[x\cdots y] and the i i -th output gate yielding the value C ⁡ [ i ] C[i] 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 a device with inputs 𝒙 \bm{x} and 𝒚 \bm{y} and outputs 𝒙 ′ \bm{x}^{\prime} and 𝒚 ′ \bm{y}^{\prime} , that performs the following function: 𝒙 ′ \displaystyle\bm{x}^{\prime} = min ⁡ ( 𝒙 , 𝒚 ) , \displaystyle=\min(\bm{x},\bm{y}), 𝒚 ′ \displaystyle\bm{y}^{\prime} = max ⁡ ( 𝒙 , 𝒚 ) . \displaystyle=\max(\bm{x},\bm{y}).

Using comparators, we can use bit-wise XOR and AND to define the result of the comparisons in lines 6 and 11 as the following fixed variables: ℓ x \displaystyle\ell_{x} : = 𝟙 { 𝒒 [ mid ] ≤ 𝒌 x } , \displaystyle:=\mathbbm{1}\{\bm{q}[\mathrm{mid}]\leq\bm{k}_{x}\}, (14) g y \displaystyle g_{y} : = 𝟙 { 𝒒 [ mid ] > 𝒌 y } , \displaystyle:=\mathbbm{1}\{\bm{q}[\mathrm{mid}]>\bm{k}_{y}\}, z \displaystyle z : = bin-search ​ ( { 𝒌 i } i = s t , 𝒒 ⁡ [ mid ] ) . \displaystyle:=\texttt{bin-search}({\{\bm{k}_{i}\}_{i=s}^{t},\ \bm{q}[\mathrm{mid}]}). This then allows us to infer the index of the key array for each recursive call to pbs-key-values in lines 9 , 14 , 19 , and 20 from Algorithm 10 . Specifically, let z s z_{s} and z t − 1 z_{t}-1 denote the starting and ending indices for the keys as inputs to the recursive calls in Algorithm 10 ) below: pbs-key-values ( 𝒒 [ mid + 1 ⋯ t ] , 𝒌 [ z t ⋯ y ] ) ; (lines 9 and 20 ) \displaystyle\texttt{pbs-key-values}(\displaystyle\bm{q}[\mathrm{mid}+1\cdots t],\ \bm{k}[z_{t}\cdots y]);\text{ (lines \ref{step: rec-1} and \ref{step: right})} (15) pbs-key-values ( 𝒒 [ s ⋯ mid − 1 ] , 𝒌 [ x ⋯ z s − 1 ] ) ; (lines 14 and 19 ) \displaystyle\texttt{pbs-key-values}(\bm{q}[s\cdots\mathrm{mid}-1],\ \bm{k}[x\ \cdots z_{s}-1]);\text{ (lines \ref{step: rec-2} and \ref{step: left})} (16) Here, z t z_{t} and z s z_{s} can assume values dependent on the results of the comparisons in lines 6 and 11 . Specifically, we have z t \displaystyle z_{t} = ℓ x ⋅ x + ( 1 − ℓ x ) ​ ( 1 − g y ) ⋅ z = { x if ​ 𝒒 ​ [ mid ] ≤ 𝒌 x ​ (line 9 ) z if ​ 𝒒 ​ [ mid ] ∈ ( 𝒌 x , 𝒌 y ] ​ (line 20 ) 0 otherwise , \displaystyle=\ell_{x}\cdot x+(1-\ell_{x})(1-g_{y})\cdot z=\begin{cases}x&\text{if }\bm{q}[\mathrm{mid}]\leq\bm{k}_{x}\text{ (line \ref{step: rec-1})}\\ z&\text{if }\bm{q}[\mathrm{mid}]\in(\bm{k}_{x},\bm{k}_{y}]\text{ (line \ref{step: right})}\\ 0&\text{otherwise}\end{cases}, z s \displaystyle z_{s} = g y ⋅ ( y + 1 ) + ( 1 − ℓ x ) ​ ( 1 − g y ) ⋅ z = { y + 1 if ​ 𝒒 ​ [ mid ] > 𝒌 y ​ (line 14 ) z if ​ 𝒒 ​ [ mid ] ∈ ( 𝒌 x , 𝒌 y ] ​ (line 19 ) 0 otherwise . \displaystyle=g_{y}\cdot(y+1)+(1-\ell_{x})(1-g_{y})\cdot z=\begin{cases}y+1&\text{if }\bm{q}[\mathrm{mid}]>\bm{k}_{y}\text{ (line \ref{step: rec-2})}\\ z&\text{if }\bm{q}[\mathrm{mid}]\in(\bm{k}_{x},\bm{k}_{y}]\text{ (line \ref{step: left})}\\ 0&\text{otherwise}\end{cases}. Here, z s z_{s} or z t z_{t} getting a value 0 0 signifies that the branch is dead, and we do not execute the recursive call.

Finally, let the arrays C t [ mid + 1 ⋯ t ] C_{t}[\mathrm{mid}+1\cdots t] and C s [ s ⋯ mid − 1 ] C_{s}[s\cdots\mathrm{mid}-1] denote the outputs to the recursive calls in equation 15 and equation 16 , respectively. We can then succinctly express the outputs for each index of the output array C C as C ⁡ [ i ] = { ℓ x ⋅ x + z s ⋅ ( C 1 ​ [ i ] ) i ∈ [ s ⋯ mid − 1 ] g y ⋅ ( y + 1 ) + z t ⋅ ( C 2 ​ [ i ] ) i ∈ [ mid + 1 ⋯ t ] ℓ x ⋅ x + ( 1 − ℓ x ) ​ ( 1 − g y ) ⋅ z + g y ⋅ ( y + 1 ) i = mid C[i]=\begin{cases}\ell_{x}\cdot x+z_{s}\cdot(C_{1}[i])&i\in[s\cdots\mathrm{mid}-1]\\ g_{y}\cdot(y+1)+z_{t}\cdot(C_{2}[i])&i\in[\mathrm{mid}+1\cdots t]\\ \ell_{x}\cdot x+(1-\ell_{x})(1-g_{y})\cdot z+g_{y}\cdot(y+1)&i=\mathrm{mid}\\ \end{cases} (17)

We can thus state the circuit schematically in Fig. 9 .

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. 9 .

###### Lemma H.32 ( Cormen et al. [2022] ) .

For binary strings 𝐱 , 𝐲 ∈ Σ d \bm{x},\bm{y}\in\Sigma^{d} of length d d , there exists a comparison network of size 𝒪 ⁡ ( d ) \mathcal{O}(d) , width 𝒪 ⁡ ( d ) \mathcal{O}(d) , and depth 𝒪 ⁡ ( log ⁡ d ) \mathcal{O}(\log{d}) .

We use this lemma to compute the complexity of the arhmetic circuit from Fig. 9 .

###### Proposition H.33 .

There exists an ( 𝒪 ⁡ ( ( n + m ) ​ d ) , 𝒪 ⁡ ( n ​ d ⋅ ( log ⁡ m + log ⁡ n ) ) , 𝒪 ⁡ ( log ⁡ n ⁡ ( log ⁡ m + log ⁡ n ​ log ​ d ) ) , 𝒪 ⁡ ( n ​ d ) ) (\mathcal{O}((n+m)d),\mathcal{O}(nd\cdot(\log{m}+\log{n})),\mathcal{O}(\log{n}(\log{m}+\log{n}\log{d})),\mathcal{O}(nd)) -arithmetic circuit 14 14 14 Recall that a ( n , s , Δ , w ) (n,s,\Delta,w) -arithmetic circuit is an n n -variate circuit with size s s , depth at most Δ \Delta , and width w w . equivalent to pbs-key-values ( Algorithm 10 ) with inputs 𝐪 \bm{q} and 𝐤 \bm{k} of lengths n n and m m with d d -bit entries.

###### Proof.

The size of the circuit for pbs-key-values should equal the work-complexity of Algorithm 10 but we also need to account for the comparison gates in equation 14 . Further, the circuit for binary search also has a size of 𝒪 ⁡ ( d ⋅ n ) \mathcal{O}(d\cdot n) instead of 𝒪 ⁡ ( d ⋅ log ⁡ n ) \mathcal{O}(d\cdot\log{n}) work as in Algorithm 10 . Using Lemma H.32 , along with the fact that the comparison gates and the binary search are used at most 𝒪 ⁡ ( log ⁡ n ) \mathcal{O}(\log{n}) times, we deduce that we are adding 𝒪 ⁡ ( n ​ d ​ log ⁡ n ) \mathcal{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 pbs-key-values is 𝒪 ⁡ ( d ​ n ​ log ⁡ m + n ​ d ​ log ⁡ n ) \mathcal{O}(dn\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 𝒪 ⁡ ( log ⁡ n ​ log ⁡ d ) \mathcal{O}(\log{n}\log{d}) , and finally, at most n n processors are used for the parallel algorithm in Algorithm 10 , which yields the width of the circuit. ∎

##### Circuit for Parallel-MQAR .

Now, we will call the circuit for Parallel-MQAR with the same name while the input gates contain the inputs of Algorithm 11 . Indeed, we can directly “translate” Algorithm 11 to an arithmetic circuit as the values for I k 𝒙 I_{k}^{\bm{x}} and I k 𝒙 I_{k}^{\bm{x}} for each 𝒙 \bm{x} and k k are predetermined from N N . Thus, we start by placing the corresponding sorting networks which feeds into the pbs-key-values ( 𝒒 [ s ⋯ t ] , 𝒌 [ x ⋯ y ] , n , m ) \texttt{pbs-key-values}\left(\bm{q}[s\cdots t],\ \bm{k}[x\cdots y],n,m\right) circuit for Algorithm 10 in Fig. 9 so that the output values from the calls to pbs-key-values result in the checks as in line 17 of Algorithm 11 . That is, we get outputs C k ​ [ J permuted k ​ 𝒙 ​ ( dec ⁡ ( i ) ) ] C_{k}[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))] from each call to pbs-key-values . We can then use a comparison gate to check if this value equals 2 k 2^{k} , and if not, we have found a match C k ​ [ J permuted k ​ 𝒙 ​ ( dec ⁡ ( i ) ) ] C_{k}[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))] for the query 𝒒 i \bm{q}_{i} which results in the output of the associated value 𝒗 C k ​ [ J permuted k ​ 𝒙 ​ ( dec ⁡ ( i ) ) ] + 1 \bm{v}_{C_{k}[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))]+1} , exactly as in Algorithm 11 . That is, we first define the following variable as the output of the comparison gate: 𝒄 dec ⁡ ( i ) k := 𝟙 { C k [ J permuted k ​ 𝒙 ( dec ( i ) ) ] ≠ 2 k } . \bm{c}_{\mathrm{dec}(i)}^{k}:=\mathbbm{1}\{C_{k}[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))]\neq 2^{k}\}. (18) Here, as C k ​ [ J permuted k ​ 𝒙 ​ ( dec ⁡ ( i ) ) ] ≠ 2 k C_{k}[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))]\neq 2^{k} implies that I permuted k ​ 𝒙 ​ ( C k ​ [ J permuted k ​ 𝒙 ​ ( dec ⁡ ( i ) ) ] ) I_{\mathrm{permuted}}^{k\bm{x}}\left(C_{k}[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))]\right) equals the index of the matching key 𝒌 j \bm{k}_{j} corresponding to the query 𝒒 i \bm{q}_{i} , the i i th output is then simply given by 𝒄 dec ⁡ ( i ) k ⋅ I permuted k ​ 𝒙 ​ ( C k ​ [ J permuted k ​ 𝒙 ​ ( dec ⁡ ( i ) ) ] ) \bm{c}_{\mathrm{dec}(i)}^{k}\cdot I_{\mathrm{permuted}}^{k\bm{x}}\left(C_{k}[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))]\right) , where the 0 0 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. 10 .

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 A be an array with d d -bit entries of size n n . Then, one can implement a sorting network to sort the array A A with size 𝒪 ⁡ ( d ⋅ n ​ log ⁡ n ) \mathcal{O}(d\cdot n\log{n}) and depth 𝒪 ⁡ ( log ⁡ d ​ log ⁡ n ) \mathcal{O}(\log{d}\log{n}) .

###### Proposition H.36 .

There exists an ( N ⋅ d , 𝒪 ⁡ ( N ​ d ⋅ log 2 ⁡ N ) , 𝒪 ⁡ ( log ⁡ d ​ log 2 ​ N ) , 𝒪 ⁡ ( N ​ d ​ log ​ N ) ) (N\cdot d,\mathcal{O}(Nd\cdot\log^{2}{N}),\mathcal{O}(\log{d}\log^{2}{N}),\mathcal{O}(Nd\log{N})) -arithmetic circuit that solves the multiple-query associative recall problem.

###### Proof.

We note here that for each k k , there are N 3 2 k + 1 \frac{\frac{N}{3}}{2^{k+1}} parallel problems of size 2 k 2^{k} for both the sorting networks and the 𝚙𝚋𝚜 − 𝚔𝚎𝚢 − 𝚟𝚊𝚕𝚞𝚎𝚜 {\tt pbs-key-values} circuit. Using Lemma H.35 , the cumulative size of these sorting networks is 𝒪 ⁡ ( d ⋅ N ​ log 2 ​ N ) \mathcal{O}(d\cdot N\log^{2}{N}) (see equation 13 ) with overall depth 𝒪 ⁡ ( log ⁡ d ​ log ⁡ N ) \mathcal{O}(\log{d}\log{N}) .

Similarly, the next layer again runs ∑ k = 0 log ⁡ N 3 − 1 N 3 2 k + 1 \sum_{k=0}^{\log{\frac{N}{3}}-1}\frac{\frac{N}{3}}{2^{k+1}} -many circuits for 𝚙𝚋𝚜 − 𝚔𝚎𝚢 − 𝚟𝚊𝚕𝚞𝚎𝚜 {\tt pbs-key-values} each of which has size 𝒪 ⁡ ( 2 k ​ d ​ ( log ⁡ 2 k + log ⁡ 2 k ) ) = O ⁡ ( d ⋅ 2 k ​ log ⁡ 2 k ) \mathcal{O}(2^{k}d(\log{2^{k}}+\log{2^{k}}))=O(d\cdot 2^{k}\log{2^{k}}) , depth 𝒪 ⁡ ( log 2 ⁡ 2 k ​ log ​ d ) \mathcal{O}(\log^{2}{2^{k}}\log{d}) and width 𝒪 ⁡ ( 2 k ) \mathcal{O}(2^{k}) ( Proposition H.33 ). Again, the cumulative size of this layer is given by 𝒪 ⁡ ( N ​ d ⋅ log 2 ⁡ N ) \mathcal{O}(Nd\cdot\log^{2}{N}) (see equation 13 ). Since we run each of these circuits in parallel, the depth of this layer is again 𝒪 ⁡ ( log ⁡ d ​ log 2 ⁡ ( N ) ) \mathcal{O}(\log{d}\log^{2}(N)) while the width is 𝒪 ⁡ ( N ⋅ log ⁡ N ) \mathcal{O}(N\cdot\log{N}) .

Finally, we perform N 3 ​ log ⁡ N 3 \frac{N}{3}\log{\frac{N}{3}} comparisons at the end of d d -bit strings in parallel which results in size 𝒪 ⁡ ( N ​ log ⁡ N ⋅ d ) \mathcal{O}(N\log{N}\cdot d) , depth 𝒪 ⁡ ( log ⁡ d ) \mathcal{O}(\log{d}) and width 𝒪 ⁡ ( N ​ log ⁡ N ⋅ d ) \mathcal{O}(N\log{N}\cdot d) ( Lemma H.32 ). Therefore, the resulting arithmetic circuit has size 𝒪 ⁡ ( d ⋅ N ​ log 2 ​ N + N ​ d ⋅ log 2 ⁡ N + N ​ log ⁡ N ⋅ d ) = O ⁡ ( N ​ d ​ log 2 ​ N ) \mathcal{O}(d\cdot N\log^{2}{N}+Nd\cdot\log^{2}{N}+N\log{N}\cdot d)=O(Nd\log^{2}{N}) , depth 𝒪 ⁡ ( log ⁡ d ​ log 2 ​ N ) \mathcal{O}(\log{d}\log^{2}{N}) and width 𝒪 ⁡ ( N ​ d ​ log ⁡ N ) \mathcal{O}(Nd\log{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 a corresponding BaseConv model that solves the multiple-query associative recall problem with 𝒪 ~ ​ ( N ​ log ⁡ c ) \tilde{\mathcal{O}}(N\log{c}) parameters and 𝒪 ⁡ ( 1 ) ~ \tilde{\mathcal{O}(1)} layers.

###### Theorem H.37 .

There exists a ( N ​ d , 𝒪 ~ ​ ( 1 ) , 𝒪 ~ ​ ( 1 ) , 𝒪 ~ ​ ( N ​ d ) , 𝒪 ~ ​ ( 1 ) ) − BaseConv \left(Nd,\tilde{\mathcal{O}}(1),\tilde{\mathcal{O}}(1),\tilde{\mathcal{O}}(Nd),\tilde{\mathcal{O}}(1)\right)-\text{{BaseConv}} solves the multiple-query associative recall problem.

###### Proof.

Directly applying Theorem H.21 yields a BaseConv model with the number of layers 𝒪 ⁡ ( log ⁡ d ⋅ log 2 ⁡ N ⋅ log ⁡ N ​ d ​ log ​ N ) = 𝒪 ⁡ ( 1 ) \mathcal{O}(\log{d}\cdot\log^{2}{N}\cdot\log{Nd\log{N}})=\mathcal{O}(1) layers while the claim on the input and inner dimensions follow trivially. ∎

### H.8 Data-Dependent Convolutions

#### H.8.1 Introduction

In this section, we are again concerned with solving the multiple-query associative recall problem ( Section H.7.1 ). However, in contrast to Section H.7.4 , which yields a circuit that is unchanged and works for all inputs, we instead take the viewpoint of adapting the sequence mixing weights ( Section 2 ) with respect to the particular sequence that the model receives as input. More specifically, we take the distance between the tokens in the sequence as a measure for designing data-dependent convolutions.

##### Setup.

To formally setup the problem, as in our discussion of designing a parallel algorithm, we consider the following problem description of the multiple-query associative recall problem. Suppose we are given an input 𝒖 [ 0 ⋯ N − 1 ] ≜ { ( 𝒌 0 , 𝒗 0 , 𝒒 0 ) , … , ( 𝒌 N 3 − 1 , 𝒗 N 3 − 1 , 𝒒 N 3 − 1 ) } \bm{u}[0\cdots N-1]\triangleq\left\{\left(\bm{k}_{0},\bm{v}_{0},\bm{q}_{0}\right),\ldots,\left(\bm{k}_{\frac{N}{3}-1},\bm{v}_{\frac{N}{3}-1},\bm{q}_{\frac{N}{3}-1}\right)\right\} with each 𝒌 i , 𝒗 i , 𝒒 i ∈ C \bm{k}_{i},\bm{v}_{i},\bm{q}_{i}\in C . Here, each token is embedded using the standard one-hot encoding in { 0 , 1 } c \{0,1\}^{c} (i.e. we assume d = c d=c ). 15 15 15 Our arguments do need c = d c=d . However we do not need d = N d=N but we made this simplification for ease of presentation. Our goal is again to check, for each 1 ≤ i ≤ N 3 − 1 1\leq i\leq\frac{N}{3}-1 , whether there exists 0 ≤ j < i 0\leq j<i such that 𝒒 i ≡ 𝒌 j \bm{q}_{i}\equiv\bm{k}_{j} , and if so, output 𝒗 j \bm{v}_{j} . Here, we define the interaction distance between the i i th query 𝒒 i \bm{q}_{i} and the matching key 𝒌 j \bm{k}_{j} as i − j i-j . We then also assume that number of distinct interaction distances is bounded by t t . We can identify these distances using an autocorrelation [ Chatfield, 1995 ] , which has an elegant underlying 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. Let 𝒖 ~ ​ [ t ] := 𝒖 ​ [ − t ] \tilde{\bm{u}}[t]:=\bm{u}[-t] , then the cross correlation of two vectors 𝒖 \bm{u} and 𝒗 \bm{v} is given by 𝒖 ⋆ 𝒗 ≜ 𝒖 ~ ∗ 𝒗 . \bm{u}\star\bm{v}\triangleq\tilde{\bm{u}}\ast\bm{v}. The autocorrelation of a vector 𝒖 ∈ ℝ n \bm{u}\in\mathbbm{R}^{n} is the cross correlation of 𝒖 \bm{u} with itself. Moreover, in terms of polynomials, we have 𝒖 ~ ​ ( X ) = X n − 1 ⋅ 𝒖 ⁡ ( 1 / X ) \tilde{\bm{u}}(X)=X^{n-1}\cdot\bm{u}(1/X) . Thus, in analogy with our interpretation of convolution in terms of polynomial multiplication, we characterize the autocorrelation of a vector 𝒖 ∈ ℝ n \bm{u}\in\mathbbm{R}^{n} , given by 𝐰 ∈ ℝ n {\bf w}\in\mathbbm{R}^{n} as follows: 𝐰 = coeff ⁡ ( 𝒖 ⁡ ( X ) ⋅ 𝒖 ~ ​ ( X ) mod X n − 1 ) . {\bf w}=\mathrm{coeff}\left(\bm{u}(X)\cdot\tilde{\bm{u}}(X)\mod{X^{n}-1}\right).

#### 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 Input-Dependent 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 t t distinct shifts of the sequence that result in highly overlapping values (e.g. matches between query and keys in our associative recall setting). We can then construct convolution filters that perform each of these t t shifts. That is, we define a function 𝚃𝚘𝚙 {\tt Top} such that 𝚃𝚘𝚙 ⁡ ( u ⋆ u , t ) {\tt Top}\left(u\star u,t\right) returns a list of the top t t shifts { s ℓ } ℓ ∈ [ t ] \{s_{\ell}\}_{\ell\in[t]} . We then use these top t t distances { s ℓ } ℓ ∈ [ t ] \{s_{\ell}\}_{\ell\in[t]} to define the following two kernels: 𝐡 k ​ ( X ) \displaystyle{\bf h}^{k}(X) ≡ ∑ ℓ ∈ [ t ] X s ℓ + ( ℓ − 1 ) ⋅ N 3 , \displaystyle\equiv\sum_{\ell\in[t]}X^{s_{\ell}+(\ell-1)\cdot\frac{N}{3}}, (19) 𝐡 v ​ ( X ) \displaystyle{\bf h}^{v}(X) ≡ ∑ ℓ ∈ [ t ] X s ℓ − 1 + ( ℓ − 1 ) ⋅ N 3 . \displaystyle\equiv\sum_{\ell\in[t]}X^{s_{\ell}-1+(\ell-1)\cdot\frac{N}{3}}. Here, we note that we only only have two kernels as we will assume N ′ = t ​ N 3 N^{\prime}=t\frac{N}{3} and the shift will be done in ”parallel.” Obviously, one can instead define t t distinct shift kernels but then there is a cost of 𝒪 ⁡ ( t ) \mathcal{O}(t) in the number of layers.

##### Projections.

We define the following projections 𝐊 , 𝐐 , 𝐕 ∈ { 0 , 1 } N × d {\bf K},{\bf Q},{\bf V}\in\{0,1\}^{N\times d} of the input that we shall use below using equation 9 . 𝐊 [ i , : ] \displaystyle{\bf K}[i,:] : = { 𝒖 [ i , : ] if ​ i ∈ 𝒦 , 𝟎 d otherwise , \displaystyle:=\begin{cases}\bm{u}[i,:]&\text{if }i\in\mathcal{K},\\ {\bf 0}^{d}&\text{otherwise}\end{cases}, (20) 𝐐 [ i , : ] \displaystyle{\bf Q}[i,:] : = { 𝒖 [ i , : ] if ​ i ∈ 𝒬 , 𝟎 d otherwise , \displaystyle:=\begin{cases}\bm{u}[i,:]&\text{if }i\in\mathcal{Q},\\ {\bf 0}^{d}&\text{otherwise}\end{cases}, 𝐕 [ i , : ] \displaystyle{\bf V}[i,:] : = { 𝒖 [ i , : ] if ​ i ∈ 𝒱 , 𝟎 d otherwise , \displaystyle:=\begin{cases}\bm{u}[i,:]&\text{if }i\in\mathcal{V},\\ {\bf 0}^{d}&\text{otherwise}\end{cases},

Finally, we present the BaseConv model that solves the multiple-query associative recall problem with input-dependent kernels using 𝒪 ⁡ ( 1 ) \mathcal{O}(1) layer and 𝒪 ⁡ ( t ⋅ N ​ d ) \mathcal{O}(t\cdot Nd) -many parameters.

###### Theorem H.38 .

There exists a BaseConv model with gated data-dependent convolutions that solves the multiple-query associative recall problem on inputs from { 0 , 1 } 3 ​ N × c \{0,1\}^{3N\times c} with the total number of distinct interaction distances bounded by t t in 𝒪 ⁡ ( 1 ) \mathcal{O}(1) layers and 𝒪 ⁡ ( t ⋅ N ​ c ) \mathcal{O}(t\cdot Nc) total parameters.

###### Proof.

We note that we have the input dimension d = c d=c . Now, for any input sequence 𝒖 ∈ { 0 , 1 } N × d \bm{u}\in\{0,1\}^{N\times d} , we get the input-dependent kernels as in equation 19 using autocorrelation of the input. We will now outline the following computations for the BaseConv layers: 𝒚 \displaystyle\bm{y} = 𝙻𝚒𝚗𝚎𝚊𝚛 𝐐 ​ ( 𝒖 ) ⊙ ( 𝐡 K ∗ 𝙻𝚒𝚗𝚎𝚊𝚛 𝐊 ​ ( 𝒖 ) ) \displaystyle={\tt Linear}_{{\bf Q}}(\bm{u})\odot\left({\bf h}^{K}*{\tt Linear}_{{\bf K}}(\bm{u})\right) = 𝐐 ⊙ ( 𝐡 K ∗ 𝐊 ) \displaystyle={\bf Q}\odot\left({\bf h}^{K}*{\bf K}\right) (21) 𝒛 \displaystyle\bm{z} = 𝙻𝚒𝚗𝚎𝚊𝚛 𝐄 ​ ( 𝒚 ) ⊙ ( 𝐡 V ∗ 𝙻𝚒𝚗𝚎𝚊𝚛 𝐕 ​ ( 𝒖 ) ) \displaystyle={\tt Linear}_{{\bf E}}(\bm{y})\odot\left({\bf h}^{V}*{\tt Linear}_{{\bf V}}(\bm{u})\right) = 𝐄 ⊙ ( 𝐡 V ∗ 𝐕 ) , \displaystyle={\bf E}\odot\left({\bf h}^{V}*{\bf V}\right), (22) where we have the linear projections 𝙻𝚒𝚗𝚎𝚊𝚛 𝐐 ​ ( 𝒖 ) = 𝐐 {\tt Linear}_{{\bf Q}}(\bm{u})={\bf Q} , 𝙻𝚒𝚗𝚎𝚊𝚛 𝐊 ​ ( 𝒖 ) = 𝐊 {\tt Linear}_{{\bf K}}(\bm{u})={\bf K} , 𝙻𝚒𝚗𝚎𝚊𝚛 𝐕 ​ ( 𝒖 ) = 𝐕 {\tt Linear}_{{\bf V}}(\bm{u})={\bf V} and 𝙻𝚒𝚗𝚎𝚊𝚛 𝐄 ​ ( 𝒚 ) = 𝐄 {\tt Linear}_{{\bf E}}(\bm{y})={\bf E} defined as 𝐄 [ i , : ] := 𝙻𝚒𝚗𝚎𝚊𝚛 𝐄 ( 𝒚 ) [ i , : ] = { 𝟏 d if ​ ∃ j ∈ [ d ] ​ such that ​ 𝒚 ​ [ i , j ] = 1 𝟎 d otherwise . {\bf E}[i,:]:={\tt Linear}_{{\bf E}}(\bm{y})[i,:]=\begin{cases}{\bf 1}^{d}&\text{if }\exists\ j\in[d]\text{ such that }\bm{y}[i,j]=1\\ {\bf 0}^{d}&\text{otherwise}\end{cases}. (23) Here, we will first present the argument for the special case when we have t = 1 t=1 as that will help us elucidate the general case. To this end, as the kernels from equation 19 for t = 1 t=1 are given by 𝐡 k ​ ( X ) \displaystyle{\bf h}^{k}(X) ≡ X s 1 ; \displaystyle\equiv X^{s_{1}}; (24) 𝐡 v ​ ( X ) \displaystyle{\bf h}^{v}(X) ≡ X s 1 − 1 , \displaystyle\equiv X^{s_{1}-1}, we observe that convolving with these kernels 𝐡 ∗ 𝒚 {\bf h}\ast\bm{y} is equivalent to operating with the following primitives ( Section H.3 ): 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝒚 , s 1 ) ; \displaystyle\tt{\displaystyle shift\_down}(\bm{y},s_{1}); (25) 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝒚 , s 1 − 1 ) . \displaystyle\tt{\displaystyle shift\_down}(\bm{y},s_{1}-1). We note that we shift down instead of shifting up as the index of the top-left entry is ( 0 , 0 ) (0,0) . We can then write down the computations performed in equation 21 and equation 22 as follows: 𝒚 \displaystyle\bm{y} = 𝐐 ⊙ 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐊 , s 1 ) \displaystyle={\bf Q}\odot{\tt shift\_down}({\bf K},s_{1}) (26) 𝒛 \displaystyle\bm{z} = 𝐄 ⊙ 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐕 , s 1 − 1 ) , \displaystyle={\bf E}\odot{\tt shift\_down}({\bf V},s_{1}-1), (27) We begin by examining 𝒚 \bm{y} below: 𝒚 [ i , : ] \displaystyle\bm{y}[i,:] = ( 𝐐 ⊙ 𝚜𝚑𝚒𝚏𝚝 _ 𝚍𝚘𝚠𝚗 ( 𝐊 , s 1 ) ) [ i , : ] \displaystyle=\left({\bf Q}\odot{\tt shift\_down}({\bf K},s_{1})\right)[i,:] = 𝐐 [ i , : ] ⊙ 𝚜𝚑𝚒𝚏𝚝 _ 𝚍𝚘𝚠𝚗 ( 𝐊 , s 1 ) [ i , : ] \displaystyle={\bf Q}[i,:]\odot{\tt shift\_down}({\bf K},s_{1})[i,:] (28) = ( { 𝒖 [ i , : ] if ​ i ∈ 𝒬 𝟎 d otherwise ) ⊙ ( { 𝒖 [ i − s 1 , : ] if ​ i − s 1 ∈ 𝒦 , 𝟎 d otherwise ) \displaystyle=\left(\begin{cases}\bm{u}[i,:]&\text{if }i\in\mathcal{Q}\\ {\bf 0}^{d}&\text{otherwise}\end{cases}\right)\odot\left(\begin{cases}\bm{u}[i-s_{1},:]&\text{if }i-s_{1}\in\mathcal{K},\\ {\bf 0}^{d}&\text{otherwise}\end{cases}\right) (29) = { 𝒖 [ i , : ] ⊙ 𝒖 [ i − s 1 , : ] if ​ i ∈ 𝒬 ​ and ​ i − s 1 ∈ 𝒦 , 𝟎 d otherwise \displaystyle=\begin{cases}\bm{u}[i,:]\odot\bm{u}[i-s_{1},:]&\text{if }i\in\mathcal{Q}\text{ and }i-s_{1}\in\mathcal{K},\\ {\bf 0}^{d}&\text{otherwise}\end{cases} Here, we use the fact that the Hadamard product is row-independent in equation 28 , and the definitions of the projections from equation 20 in equation 29 . Examining the j j th entry, we get 𝒖 ⁡ [ i , j ] ⊙ 𝒖 ⁡ [ i − s 1 , j ] \displaystyle\bm{u}[i,j]\odot\bm{u}[i-s_{1},j] = { 1 if ​ i ∈ 𝒬 , i − s 1 ∈ 𝒦 ​ and ​ 𝒒 i ≡ 𝒌 i − s 1 ≡ 𝐞 j 0 otherwise. \displaystyle=\begin{cases}1&\text{if }i\in\mathcal{Q},i-s_{1}\in\mathcal{K}\text{ and }\bm{q}_{i}\equiv\bm{k}_{i-s_{1}}\equiv{\bf e}_{j}\\ 0&\text{otherwise.}\end{cases} That is, we can express 𝒚 [ i , : ] \displaystyle\bm{y}[i,:] = { 𝐞 j if ​ i ∈ 𝒬 , i − s 1 ∈ 𝒦 ​ and ​ 𝒒 i ≡ 𝒌 i − s 1 ≡ 𝐞 j 𝟎 d otherwise . \displaystyle=\begin{cases}{\bf e}_{j}&\text{if }i\in\mathcal{Q},i-s_{1}\in\mathcal{K}\text{ and }\bm{q}_{i}\equiv\bm{k}_{i-s_{1}}\equiv{\bf e}_{j}\\ {\bf 0}^{d}&\text{otherwise}\end{cases}. (30) Consequently, as per the definition in equation 23 , we get 𝐄 [ i , : ] = { 𝟏 d if ​ i ∈ 𝒬 , i − s 1 ∈ 𝒦 ​ and ​ 𝒒 i ≡ 𝒌 i − s 1 𝟎 d otherwise \displaystyle{\bf E}[i,:]=\begin{cases}{\bf 1}^{d}&\text{if }i\in\mathcal{Q},i-s_{1}\in\mathcal{K}\text{ and }\bm{q}_{i}\equiv\bm{k}_{i-s_{1}}\\ {\bf 0}^{d}&\text{otherwise}\end{cases} (31) We can now finally specify the output 𝒛 \bm{z} from equation 27 as follows: 𝒛 [ i , : ] \displaystyle\bm{z}[i,:] = ( 𝐄 ⊙ 𝚜𝚑𝚒𝚏𝚝 _ 𝚍𝚘𝚠𝚗 ( 𝐕 , s 1 − 1 ) ) [ i , : ] \displaystyle=\left({\bf E}\odot{\tt shift\_down}({\bf V},s_{1}-1)\right)[i,:] = 𝐄 [ i , : ] ⊙ 𝚜𝚑𝚒𝚏𝚝 _ 𝚍𝚘𝚠𝚗 ( 𝐕 , s 1 − 1 ) [ i , : ] \displaystyle={\bf E}[i,:]\odot{\tt shift\_down}({\bf V},s_{1}-1)[i,:] (32) = ( { 𝟏 d if ​ i ∈ 𝒬 , i − s 1 ∈ 𝒦 ​ and ​ 𝒒 i ≡ 𝒌 i − s 1 𝟎 d otherwise ) ⊙ ( { 𝒖 [ i − s 1 + 1 , : ] if ​ i − s 1 + 1 ∈ 𝒱 , 𝟎 d otherwise ) \displaystyle=\left(\begin{cases}{\bf 1}^{d}&\text{if }i\in\mathcal{Q},i-s_{1}\in\mathcal{K}\text{ and }\bm{q}_{i}\equiv\bm{k}_{i-s_{1}}\\ {\bf 0}^{d}&\text{otherwise}\end{cases}\right)\odot\left(\begin{cases}\bm{u}[i-s_{1}+1,:]&\text{if }i-s_{1}+1\in\mathcal{V},\\ {\bf 0}^{d}&\text{otherwise}\end{cases}\right) (33) = { 𝒖 [ i − s 1 + 1 , : ] if ​ i ∈ 𝒬 , i − s 1 ∈ 𝒦 , i − s 1 + 1 ∈ 𝒱 ​ and ​ 𝒒 i ≡ 𝒌 i − s 1 𝟎 d otherwise \displaystyle=\begin{cases}\bm{u}[i-s_{1}+1,:]&\text{if }i\in\mathcal{Q},i-s_{1}\in\mathcal{K},i-s_{1}+1\in\mathcal{V}\text{ and }\bm{q}_{i}\equiv\bm{k}_{i-s_{1}}\\ {\bf 0}^{d}&\text{otherwise}\end{cases} = { 𝒗 i − s 1 if ​ 𝒒 i ≡ 𝒌 i − s 1 𝟎 d otherwise \displaystyle=\begin{cases}\bm{v}_{i-s_{1}}&\text{if }\bm{q}_{i}\equiv\bm{k}_{i-s_{1}}\\ {\bf 0}^{d}&\text{otherwise}\end{cases} (34) Again, we use the fact that the Hadamard product is row-independent in equation 32 , and the definitions of the projections from equation 20 in equation 33 . Overall, we have solved associative recall for all queries that have interaction distance exactly equal to s 1 s_{1} .

In order to generalize this to arbitrary t ≤ N 3 t\leq\frac{N}{3} , we first increase the internal dimension so that the input to the kernels 𝒖 ′ ∈ ℝ ( N ⋅ t ) × d \bm{u}^{\prime}\in\mathbbm{R}^{(N\cdot t)\times d} in equation 19 and the projections 𝐊 ′ , 𝐐 ′ , 𝐕 ′ ∈ ℝ ( N ⋅ t ) × d {\bf K^{\prime},Q^{\prime},V^{\prime}}\in\mathbbm{R}^{(N\cdot t)\times d} are given by 𝒖 ′ ≡ ( 𝟎 d 𝟎 d 𝒖 ) , 𝑲 ′ ≡ ( 𝟎 d 𝟎 d 𝐊 ) , 𝐐 ′ ≡ ( 𝐐 𝐐 𝐐 ) , 𝐕 ′ ≡ ( 𝟎 d 𝟎 d 𝐕 ) , \bm{u}^{\prime}\equiv\begin{pmatrix}{\bf 0}^{d}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf 0}^{d}\\ \hline\cr\\ \bm{u}\end{pmatrix},\bm{K}^{\prime}\equiv\begin{pmatrix}{\bf 0}^{d}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf 0}^{d}\\ \hline\cr\\ {\bf K}\end{pmatrix},{\bf Q}^{\prime}\equiv\begin{pmatrix}{\bf Q}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf Q}\\ \hline\cr\\ {\bf Q}\end{pmatrix},{\bf V}^{\prime}\equiv\begin{pmatrix}{\bf 0}^{d}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf 0}^{d}\\ \hline\cr\\ {\bf V}\end{pmatrix}, We then observe that for 𝐡 ℓ k ​ ( X ) := X s ℓ {\bf h}^{k}_{\ell}(X):=X^{s_{\ell}} and 𝐡 ℓ v ​ ( X ) := X s ℓ − 1 {\bf h}^{v}_{\ell}(X):=X^{s_{\ell}-1} , we have 𝐡 k ​ ( X ) \displaystyle{\bf h}^{k}(X) ≡ ∑ ℓ ∈ [ t ] 𝐡 ℓ k ​ ( X ) ⋅ X ( ℓ − 1 ) ⋅ N 3 , \displaystyle\equiv\sum_{\ell\in[t]}{\bf h}^{k}_{\ell}(X)\cdot X^{(\ell-1)\cdot\frac{N}{3}}, 𝐡 v ​ ( X ) \displaystyle{\bf h}^{v}(X) ≡ ∑ ℓ ∈ [ t ] 𝐡 ℓ k ​ ( X ) ⋅ X ( ℓ − 1 ) ⋅ N 3 . \displaystyle\equiv\sum_{\ell\in[t]}{\bf h}^{k}_{\ell}(X)\cdot X^{(\ell-1)\cdot\frac{N}{3}}. In analogy with equation 25 , we can then equivalently write ( 𝐡 K ∗ 𝐊 ) \displaystyle\left({\bf h}^{K}\ast{\bf K}\right) ≡ ( 𝐡 t K 𝐡 2 K 𝐡 1 K ) ∗ ( 𝟎 d 𝟎 d 𝐊 ) \displaystyle\equiv\begin{pmatrix}{\bf h}^{K}_{t}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf h}^{K}_{2}\\ \hline\cr\\ {\bf h}^{K}_{1}\end{pmatrix}\ast\begin{pmatrix}{\bf 0}^{d}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf 0}^{d}\\ \hline\cr\\ {\bf K}\end{pmatrix} ≡ ( 𝐡 t K ∗ 𝐊 𝐡 2 K ∗ 𝐊 𝐡 1 K ∗ 𝐊 ) \displaystyle\equiv\begin{pmatrix}{\bf h}^{K}_{t}\ast{\bf K}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf h}^{K}_{2}\ast{\bf K}\\ \hline\cr\\ {\bf h}^{K}_{1}\ast{\bf K}\end{pmatrix} ≡ ( 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐊 , s t ) 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐊 , s 2 ) 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐊 , s 1 ) ) . \displaystyle\equiv\begin{pmatrix}{\tt shift\_down}({\bf K},s_{t})\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\tt shift\_down}({\bf K},s_{2})\\ \hline\cr\\ {\tt shift\_down}({\bf K},s_{1})\end{pmatrix}. Similarly, we also have ( 𝐡 V ∗ 𝐕 ′ ) \displaystyle\left({\bf h}^{V}\ast{\bf V}^{\prime}\right) ≡ ( 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐕 , s t − 1 ) 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐕 , s 2 − 1 ) 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐕 , s 1 − 1 ) ) . \displaystyle\equiv\begin{pmatrix}{\tt shift\_down}({\bf V},s_{t}-1)\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\tt shift\_down}({\bf V},s_{2}-1)\\ \hline\cr\\ {\tt shift\_down}({\bf V},s_{1}-1)\end{pmatrix}. That is, the argument for t = 1 t=1 now applies to each of the t t shifts as we now have ( cf. equation 21 ) 𝒚 ′ \displaystyle\bm{y}^{\prime} ≡ 𝐐 ′ ⊙ ( 𝐡 V ∗ 𝐕 ) \displaystyle\equiv{\bf Q}^{\prime}\odot\left({\bf h}^{V}\ast{\bf V}\right) ≡ ( 𝐐 𝐐 𝐐 ) ⊙ ( 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐕 , s t − 1 ) 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐕 , s 2 − 1 ) 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐕 , s 1 − 1 ) ) \displaystyle\equiv\begin{pmatrix}{\bf Q}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf Q}\\ \hline\cr\\ {\bf Q}\end{pmatrix}\odot\begin{pmatrix}{\tt shift\_down}({\bf V},s_{t}-1)\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\tt shift\_down}({\bf V},s_{2}-1)\\ \hline\cr\\ {\tt shift\_down}({\bf V},s_{1}-1)\end{pmatrix} ≡ ( 𝐐 ⊙ 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐕 , s t − 1 ) 𝐐 ⊙ 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐕 , s 2 − 1 ) 𝐐 ⊙ 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐕 , s 1 − 1 ) ) \displaystyle\equiv\begin{pmatrix}{\bf Q}\odot{\tt shift\_down}({\bf V},s_{t}-1)\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf Q}\odot{\tt shift\_down}({\bf V},s_{2}-1)\\ \hline\cr\\ {\bf Q}\odot{\tt shift\_down}({\bf V},s_{1}-1)\end{pmatrix} ≡ ( 𝒚 t 𝒚 2 𝒚 1 ) , \displaystyle\equiv\begin{pmatrix}\bm{y}_{t}\\ \hline\cr\\ \vdots\\ \hline\cr\\ \bm{y}_{2}\\ \hline\cr\\ \bm{y}_{1}\\ \end{pmatrix}, where, for each ℓ ∈ [ t ] \ell\in[t] , we have ( cf. equation 30 ) 𝒚 ℓ [ i , : ] ≡ { 𝐞 j if i ∈ 𝒬 , i − s ℓ ∈ 𝒦 and 𝒒 i ≡ 𝒌 i − s ℓ ≡ 𝐞 j , 𝟎 d otherwise . \bm{y}_{\ell}[i,:]\equiv\begin{cases}{\bf e}_{j}&\text{if }i\in\mathcal{Q},i-s_{\ell}\in\mathcal{K}\text{ and }\bm{q}_{i}\equiv\bm{k}_{i-s_{\ell}}\equiv{\bf e}_{j},\\ {\bf 0}^{d}&\text{otherwise}\end{cases}. We then analogously get 𝐄 ′ {\bf E}^{\prime} as follows: 𝐄 ′ ≡ 𝙻𝚒𝚗𝚎𝚊𝚛 𝐄 ​ ( 𝒚 ′ ) ≡ ( 𝙻𝚒𝚗𝚎𝚊𝚛 𝐄 ​ ( 𝒚 t ) 𝙻𝚒𝚗𝚎𝚊𝚛 𝐄 ​ ( 𝒚 2 ) 𝙻𝚒𝚗𝚎𝚊𝚛 𝐄 ​ ( 𝒚 1 ) ) ≡ ( 𝐄 t 𝐄 2 𝐄 1 ) , {\bf E}^{\prime}\equiv{\tt Linear}_{{\bf E}}(\bm{y}^{\prime})\equiv\begin{pmatrix}{\tt Linear}_{{\bf E}}(\bm{y}_{t})\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\tt Linear}_{{\bf E}}(\bm{y}_{2})\\ \hline\cr\\ {\tt Linear}_{{\bf E}}(\bm{y}_{1})\end{pmatrix}\equiv\begin{pmatrix}{\bf E}_{t}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf E}_{2}\\ \hline\cr\\ {\bf E}_{1}\end{pmatrix}, where, for each ℓ ∈ [ t ] \ell\in[t] , we have ( cf. equation 31 ) 𝐄 ℓ [ i , : ] = { 𝟏 d if ​ i ∈ 𝒬 , i − s ℓ ∈ 𝒦 ​ and ​ 𝒒 i ≡ 𝒌 i − s ℓ 𝟎 d otherwise {\bf E}_{\ell}[i,:]=\begin{cases}{\bf 1}^{d}&\text{if }i\in\mathcal{Q},i-s_{\ell}\in\mathcal{K}\text{ and }\bm{q}_{i}\equiv\bm{k}_{i-s_{\ell}}\\ {\bf 0}^{d}&\text{otherwise}\end{cases} The output in the general case is then given by 𝒛 ′ \displaystyle\bm{z}^{\prime} ≡ 𝐄 ′ ⊙ ( 𝐡 V ∗ 𝐕 ′ ) \displaystyle\equiv{\bf E}^{\prime}\odot\left({\bf h}^{V}*{\bf V}^{\prime}\right) ≡ ( 𝐄 t 𝐄 2 𝐄 1 ) ⊙ ( 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐕 , s t − 1 ) 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐕 , s 2 − 1 ) 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐕 , s 1 − 1 ) ) \displaystyle\equiv\begin{pmatrix}{\bf E}_{t}\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf E}_{2}\\ \hline\cr\\ {\bf E}_{1}\end{pmatrix}\odot\begin{pmatrix}{\tt shift\_down}({\bf V},s_{t}-1)\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\tt shift\_down}({\bf V},s_{2}-1)\\ \hline\cr\\ {\tt shift\_down}({\bf V},s_{1}-1)\end{pmatrix} ≡ ( 𝐄 t ⊙ 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐕 , s t − 1 ) 𝐄 2 ⊙ 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐕 , s 2 − 1 ) 𝐄 1 ⊙ 𝚜𝚑𝚒𝚏𝚝 ​ _ ​ 𝚍𝚘𝚠𝚗 ​ ( 𝐕 , s 1 − 1 ) ) \displaystyle\equiv\begin{pmatrix}{\bf E}_{t}\odot{\tt shift\_down}({\bf V},s_{t}-1)\\ \hline\cr\\ \vdots\\ \hline\cr\\ {\bf E}_{2}\odot{\tt shift\_down}({\bf V},s_{2}-1)\\ \hline\cr\\ {\bf E}_{1}\odot{\tt shift\_down}({\bf V},s_{1}-1)\end{pmatrix} ≡ ( 𝒛 t 𝒛 2 𝒛 1 ) , \displaystyle\equiv\begin{pmatrix}\bm{z}_{t}\\ \hline\cr\\ \vdots\\ \hline\cr\\ \bm{z}_{2}\\ \hline\cr\\ \bm{z}_{1}\\ \end{pmatrix}, where, for each ℓ ∈ [ t ] \ell\in[t] , we have ( cf. equation 34 ) 𝒛 ℓ [ i , : ] ≡ { 𝒗 i − s ℓ if ​ 𝒒 i ≡ 𝒌 i − s ℓ 𝟎 d otherwise \bm{z}_{\ell}[i,:]\equiv\begin{cases}\bm{v}_{i-s_{\ell}}&\text{if }\bm{q}_{i}\equiv\bm{k}_{i-s_{\ell}}\\ {\bf 0}^{d}&\text{otherwise}\end{cases} Finally, we define the last output layer to compute 𝒛 o ​ u ​ t ≡ 𝙻𝚒𝚗𝚎𝚊𝚛 sum ​ ( 𝒛 ′ ) ≡ ∑ ℓ ∈ [ t ] 𝒛 ℓ \bm{z}_{out}\equiv{\tt Linear}_{\mathrm{sum}}(\bm{z}^{\prime})\equiv\sum_{\ell\in[t]}\bm{z}_{\ell} so that we have 𝒛 o ​ u ​ t [ i , : ] ≡ { 𝒗 i − s ℓ if ​ 𝒒 i ≡ 𝒌 i − s ℓ ​ for some ​ ℓ ∈ [ t ] 𝟎 d otherwise \bm{z}_{out}[i,:]\equiv\begin{cases}\bm{v}_{i-s_{\ell}}&\text{if }\bm{q}_{i}\equiv\bm{k}_{i-s_{\ell}}\text{ for some }\ell\in[t]\\ {\bf 0}^{d}&\text{otherwise}\end{cases} To recall, we retrieved the top t t interaction distances of the input 𝒖 \bm{u} using auto-correlation and defined the corresponding convolution kernels (equation 19 ). We then shifted djown the keys 𝐊 {\bf K} using the first kernel and gated with the corresponding queries 𝐐 {\bf Q} so that we got a match exactly when there exists a key that is at s ℓ s_{\ell} interaction distance from the corresponding query. After “smearing” this match to get 𝐄 {\bf E} , we used it as a mask to retrieve the value in the next layer. Overall, since we have t t atomic kernels that perform t t shifts with each of these kernels using 𝒪 ⁡ ( N ​ d ) \mathcal{O}(Nd) parameters, we can conclude that the output solves the associative recall problem for all queries with exactly ℓ \ell interaction distance from the corresponding keys for all ℓ ∈ [ t ] \ell\in[t] using 𝒪 ⁡ ( 1 ) \mathcal{O}(1) layers and 𝒪 ⁡ ( t ⋅ N ​ c ) \mathcal{O}(t\cdot Nc) parameters as we have d = c d=c . ∎

## Instructions for reporting errors

We are continuing to improve HTML versions of papers, and your feedback helps enhance accessibility and mobile support. To report errors in the HTML that will help us improve conversion and rendering, choose any of the methods listed below:

Click the "Report Issue" ( ) button, located in the page header.

Tip: You can select the relevant text first, to include it in your report.

Our team has already identified the following issues . We appreciate your time reviewing and reporting rendering errors we may not have found yet. Your efforts will help us improve the HTML versions for all readers, because disability should not be a barrier to accessing research. Thank you for your continued support in championing open access for all.

Have a free development cycle? Help support accessibility at arXiv! Our collaborators at LaTeXML maintain a list of packages that need conversion , and welcome developer contributions .