Zoology: Measuring and Improving Recall in Efficient Language Models — part 7/14

Cleaned-text code points [93491, 108041). Bibliography omitted here; source.md includes it.

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.