Hence, we can reproduce the output of the $\ell$th layer of Hyena using five layers of  after augmenting the input and using the remembering primitive (<ref>) with internal dimension $N'+N$.

Now, for the input $\bm{u}_{\Coyote} \in \R^{N \times d}$, the output of the $\ell$th layer for  is given by

\begin{align*}
&z^{\ell}_{\Coyote} \equiv \texttt{Linear}(z^{\ell - 1}_{\Coyote}) \odot \mathrm{conv}(h^l, z^{\ell - 1}_{\Coyote}).
\end{align*}

Here, we show inductively that simply using $\ell$-many Hyena models recursively simulates $z^{\ell}_{\Coyote}$. For $\ell = 1$, we have

\[
\mathrm{Hyena}(\bm{u}_{\Coyote, \bm{h}}) \equiv \texttt{Linear}(\bm{u}_{\Coyote}) \odot (\bm{h}^{1}\ast \bm{u}_{\Coyote}) \equiv \bm{z}^{1}_{\Coyote}.
\]

We now assume that $(\ell-1)$-many recursive Hyena models produce $z^{(\ell-1)}_{\Coyote}$. For the $\ell$th layer, we then have

\begin{align*}
&\mathrm{Hyena}\paren{\mathrm{Hyena}\paren{\ldots \mathrm{Hyena}(\bm{u}_{\Coyote, \bm{h}})}} \\
\equiv\quad  &\mathrm{Hyena}\paren{z^{(l-1)}_{\Coyote}}\\
\equiv\quad  &\mathrm{linear}\paren{z^{(l-1)}_{\Coyote}} \odot \mathrm{conv}\paren{h^l, z^{(l-1)}_{\Coyote}} \\
\equiv\quad  &z^{\ell}_{\Coyote}.
\end{align*}

The gated convolution step of <ref> gives the following result with respect to input $\hyenaInput \in \R^{N \times d}$, and convolution filters $\bm{h} \in  \R^{N \times d}$

\begin{align*}
\bm{y}_H = \bm{p}^{1} \odot \paren{\bm{h} \ast \bm{z}^{0}}
\end{align*}

where  $\bm{p}^{1} = \Projection{\hyenaInput,\convFilter_s}[:,:,0]$ and $\bm{z}^{0} = \Projection{\hyenaInput,\convFilter_s}[:,:,1]$ as in <ref>. Specifically we have

\[\bm{p}^{1}= \paren{\convFilter_s \ast \LinearD{\headDim}{2\headDim}{\bm{u}}}[:,0:\headDim-1] \qquad \qquad \bm{z}^{0}= \paren{\convFilter_s \ast \LinearD{\headDim}{2\headDim}{\bm{u}}}[:,
\headDim:2\headDim-1]\]

We note that $\bm{y}_H$ can be computed from $\hyenaInput$ via the following sequence of operations:

* Compute $\bm{z} =\convFilter_s \ast \LinearD{\headDim}{2\headDim}{\bm{u}}$ where $\bm{z} \in \R^{\inputLength \times 2\headDim}$, giving us

\[ \bm{z} =  \begin{pmatrix}
\bm{p}^1\\
\hline \\
\bm{z}^0
\end{pmatrix} \]

* Compute $\shiftN\texttt{\_down}({\bm{z}},{\inputLength})$ to get

\[ \bm{z}_{\shiftN} =  \begin{pmatrix}
\bm{0}\\
\hline \\
\bm{p}^1
\end{pmatrix} \]

Where is $\bm{z}_{\shiftN}$ used later on?

* Compute $\copyPrimitive{\bm{z}}{\bm{p^1,z^0,h,p^1}}$ There seem to be two issues here. (i) How do you get $\bm{z}^0$ as input from $\bm{z}_{\shiftN}$ (which presumably the input to this step?). (ii) The way I interpreted the remember operator the last two inputs are independent of $\bm{y}$: so we cannot have $\bm{p}^1$ as the last input since it depends on $\bm{y}$to get

\[ \bm{y} =  \begin{pmatrix}
\bm{p}^1 \\
\hline \\
\bm{p}^{1} \odot \paren{\bm{h} \ast \bm{z}^{0}}
\end{pmatrix}, \]

as desired.

By <ref>, completing each step requires $\coyoteTuple{\inputLength}{1}{\headDim}{2\inputLength}{\innerD}$.
With <ref>, we can stack Since $\bm{z}_{\shiftN}$ does not feed into the remember step, I do not see how you can stack things. these models to define $\coyoteTuple{\inputLength}{3}{\headDim}{2\inputLength}{\innerD}$ such that its output $\bm{y}_C = \bm{y}_H$ as desired.

Now we show that for any $\bm{y}_C = \SHyena{\hyenaInput}{\bm{h}}$ defined on $\coyoteTuple{\inputLength}{1}{\headDim}{\innerN}{\innerD}$, there exists a $\hyenaTuple{\inputLength}{1}{\headDim}{N}{2\headDim}$ such that
$\bm{y}_C = \bm{y}_H$. By <ref>, for any $\bm{u} \in \R^{N \times d}$ the output for $\SHyenaName$ with no bias matrices is given as

\begin{align*}
&\bm{y}_C \equiv \tbf{W}\hyenaInput \odot \paren{\convFilter \ast \bm{u}}.
\end{align*}

Because the final step of <ref> is a gated convolution, we may get $\bm{y}_C = \bm{y}_H$ by setting $\Projection{\hyenaInput, \convFilter_s} = \tbf{W}\hyenaInput$, thereby ignoring the convolution step in $\ProjectionN$. Then the statement follows.

§.§ Linear Arithmetic Circuits

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

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 $\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 $ax + by$ for fixed constants $a, b \in \F$).

A butterfly factor of size $k \geq 2$ (denoted as $\mathbf{B}_k$ ) is a matrix of the form $\mathbf{B}_k=$ $\left[\begin{array}{ll}\mathbf{D}_1 & \mathbf{D}_2 \\ \mathbf{D}_3 & \mathbf{D}_4\end{array}\right]$ where each $\mathbf{D}_i$ is a $\frac{k}{2} \times \frac{k}{2}$ diagonal matrix. We restrict $k$ to be a power of 2 .

A butterfly factor matrix of size $n$ with block size $k$ (denoted as $\mathbf{B}_k^{(n)}$ ) is a block diagonal matrix of $\frac{n}{k}$ (possibly different) butterfly factors of size $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$ (denoted as $\mathbf{B}^{(n)}$ ) is a matrix that can be expressed as a product of butterfly factor matrices: $\mathbf{B}^{(n)}=\mathbf{B}_n^{(n)} \mathbf{B}_{\frac{n}{2}}^{(n)} \ldots \mathbf{B}_2^{(n)}$. Equivalently, we may define $\mathbf{B}^{(n)}$ recursively as a matrix that can be expressed in the following form:

\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 $\left[\mathbf{B}^{\left(\frac{n}{2}\right)}\right]_1$ and $\left[\mathbf{B}^{\left(\frac{n}{2}\right)}\right]_2$ may be different.)

From <ref>, we observe that size $n$ butterfly factor is comprised of three vectors $\bm{d, d^+, d^-} \in \R^n$ such that

\begin{align*}
\bm{d} &= \paren{\diag^{-1}\paren{\tbf{D}_1} , \diag^{-1}\paren{\tbf{D}_4}}, \\
\bm{d}^{+} &= \paren{\bm{0}^{\tfrac n2}, \diag^{-1}\paren{\tbf{D}_2}}, \text{ and } \\
\bm{d}^{-} &= \paren{\diag^{-1}\paren{\tbf{D}_3} , \bm{0}^{\tfrac n2}},
\end{align*}

where $\diag^{-1}(\tbf{D}): \R^{n \times n} \mapsto \R^n$ is the mapping from diagonal matrices to the vector of its diagonal entries. Let us define $\tbf{D_1, D_2, D_3} \in \R^{n \times n}$ as $\diag\paren{\bm{d}}, \diag\paren{\bm{d^+}}$, and $\diag\paren{\bm{d^-}}$ respectively. Then we note that

\begin{equation}
\label{eq: butterfly-split}
\tbf{D}_1 \equiv \left[\begin{array}{ll}\mathbf{D}_1 & \mathbf{0} \\ \mathbf{0} & \mathbf{D}_4\end{array}\right] \qquad \tbf{D}_2\bm{S}^{\tfrac n2} \equiv \left[\begin{array}{ll}\mathbf{0} & \mathbf{D}_2 \\ \mathbf{0} & \mathbf{0}\end{array}\right] \qquad \bm{S}^{\tfrac n2}\tbf{D}_3 \equiv \left[\begin{array}{ll}\mathbf{0} & \mathbf{0} \\ \mathbf{D}_3 & \mathbf{0}\end{array}\right]
\end{equation}

where $\tbf{S}^k \in \F^{n \times n}$ is a shift matrix for $i \in [n/2]$. This gives us the following proposition:

For any powers of 2, $n = k \geq 2$,
any butterfly factor matrix $\tbf{B}^{(n)}_k$ is equivalent to

\[ \tbf{B}^{(n)}_k= \tbf{S}^{\tfrac k2}\tbf{D}_3 +\tbf{D}_2\tbf{S}^{\tfrac n2} + \tbf{D}_1\]

where $\tbf{D_3, D_2, D_1, S^{\frac{n}{2}}}$ are defined as in (<ref>).

We use <ref> to show that butterfly matrices can easily be computed by $\SHyenaName$ .

For any $n,\headDim \geq 2$, $k \geq 1$, and arbitrary vector $\bm{x} \in \R^{n\headDim}$:

* there exists a $\coyoteTupleN$ that  can represent $\tbf{B}^{(nd)}_{k}\cdot \bm{x}$ with $N = n, \innerN = \calO(\inputLength)$, $\hyenaDepth = \calO(1)$, and $\innerD = \calO(\headDim)$, and

* there exists a $\coyoteTupleN$ that  can represent $\tbf{B}^{(nd)}\cdot \bm{x}$ with $N = n, \innerN = \calO(\inputLength)$, $\hyenaDepth = \calO(\log nd)$, and $\innerD = \calO(\headDim)$.

* Given $\bm{x} \in \R^{n\headDim}$, construct $\hyenaInput \in \R^{n\times d}$
where $\bm{x}$ is the row-major form of $\hyenaInput$. We show that $\Coyote$ can compute $\tbf{B}_{nd}\cdot \bm{x}$ column by column.

Let $\tbf{A}=\tbf{S}^{\tfrac k2}\tbf{D}'_3, \tbf{C} = \tbf{D}'_2\tbf{S}^{\tfrac n2}$, and $\tbf{D} = \tbf{D}_1$ for  $\tbf{D}_i,\tbf{S}^{\tfrac{k}{2}} \in \R^{nd \times nd}$ for $1\le i\le 3$
as defined in <ref>. We take $\bm{d}_1 = \bm{1}^{nd}\tbf{D}, \ \bm{d}_2 = \bm{1}^{nd}\tbf{C}_2, \bm{d}_3 = \bm{1}^{nd}\tbf{A}$, which extracts the diagonal entries of $\tbf{D}_i$. With this we construct $\tbf{D}'_i \in \R^{n\times d}$ where $\bm{d}_i$ is the row major form of $\bm{D}'_i$. This implies that

\[
\tbf{D}_i\bm{x} \equiv \tbf{D}'_i \odot \bm{u}.
\]

Then we can decompose   $\tbf{B}_{nd}\cdot \bm{x}$ into

\[\tbf{B}_{nd}\bm{x} \equiv \tbf{D}_1 \odot \bm{u} +  \tbf{D}_2 \odot \bm{u} + \tbf{D}_3 \odot \bm{u}.\]

By <ref>, each Hadamard product $\tbf{A} \odot \bm{u}, \tbf{B} \odot \bm{u}, \tbf{C} \odot \bm{u}$ can be trivially be performed with a single layer $\Coyote$ 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[Recall that $\texttt{add}_{n}(\bm{y}: \bm{x}, \bm{S})$ adds the subvector $\bm{x}$ to $\bm{S}$ for the input $\bm{y}$.]:

\[
\begin{aligned}
&\texttt{add}_{n}(\bm{y}^1: \bm{y}_1, \bm{0})\\
&\texttt{add}_{n}(\bm{y}^2: \bm{y}_2, \bm{y}_1)\\
&\texttt{add}_{n}(\bm{y}^3: \bm{y}_3, \bm{y}_1+\bm{y}_2),
\end{aligned}
\]

where using by <ref> and <ref>, this requires six more layers, and we get

\[
\bm{y}^3 \equiv \bm{y_1 + y_2 + y_3} \equiv \tbf{B}_{nd}\bm{x}.
\]

Then we can construct the $\coyoteTupleN$ as desired with $L = O(1)$ layers.

* From <ref>, $\mathbf{B}^{(nd)}=\mathbf{B}_{nd}^{(nd)} \mathbf{B}_{\frac{nd}{2}}^{(nd)} \ldots \mathbf{B}_2^{(nd)}$. From (1), $\Coyote$ can compute any butterfly matrix by simulating the $\log(nd)$ butterfly factor matrices which comprise $\mathbf{B}^{(nd)}$. With <ref>, this creates a $\Coyote$ with $5 \cdot \log(nd) = \calO( \log(nd))$ layers. <ref>

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

* Define $\calB$ as the set of all matrices that can be expressed in the form $\tbf{B}^{(n)}$ (for some $n$).

* Define $\paren{\cal{BB}^*}$ as the set of matrices $\tbf{M}$ of the form $\tbf{M} = \tbf{M}_1\tbf{M}_2^*$ for some $\tbf{M_1}, \tbf{M}_2 \in \calB$.

* Define $\paren{\cal{BB}^*}^w$ as the set of matrices $\tbf{M}$ that can be expressed as $\tbf{M} = \tbf{M}_w\ldots \tbf{M}_2\tbf{M}_1$, with each $\tbf{M}_i \in \paren{\cal{BB}^*} (1 \le i \le w).$ (The notation $w$ represents width.)

* Define $\paren{\cal{BB}^*}^w_e$ as the set of $n \times n$ matrices $\tbf{M}$ that can be expressed as $\tbf{M} = \tbf{SES}^{\top}$ for some $en \times en$ matrix $\tbf{E} \in \paren{\cal{BB}^*}^w$, where $\tbf{S} \in \F^{n \times en} = \begin{bmatrix}
\tbf{I}_n & 0 &\ldots & 0
\end{bmatrix}]$ (i.e. $\tbf{M}$ is the upper-left corner of $\tbf{E}$). (The notation $e$ represents expansion relative to $n$.)

We similarly show how $\Coyote$ can simulate any kaleidoscope matrix.

Given $n,\headDim \geq 2$,  $e > 0$ for any $nd \times nd$ matrix $ \tbf{M} \in \paren{\cal{BB}^*}^{w}_{e}$, and $\bm{x} \in \R^{n\headDim}$  there exists a $\coyoteTupleN$ that  can represent $\tbf{M}\cdot \bm{x}$
with $N = n, \hyenaDepth = \calO(w\log(end))$,
$\innerN= en,$
and $\  \innerD=d$.

By <ref>, $\tbf{M}$ can be decomposed with respect to size $end \times end$ matrix

\[\tbf{E} = \tbf{E}_1 \cdot \tbf{E}_2 \cdots \tbf{E}_{w}.\]

Further, any $\tbf{E}_i \in \paren{\cal{BB}^*}$ can be expressed as a product of $2 \log {end}$
butterfly factor matrices. Then by <ref> and <ref> we can compute  $\tbf{E}_i\bm{x'}$ in by stacking $2 \log {end}$ $\coyoteTuple{n}{d}{L}{en}{d}$ models each with $L =\calO(1)$. Because $\tbf{E}$ has width $w$, <ref> implies that composing with each $\tbf{E}_i$ for $1 \leq i \leq w$ constructs a final model with $\calO(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$, yet inspection of the proof yields the following stronger result:

Let $\tbf{M}$ be an $n \times n$ matrix such that multiplication of $\tbf{M}$ times an arbitrary vector $\bm{u}$ can be
be represented as $(n,s,\circuitDegree,w)$-linear arithmetic circuit $\calC$. Then, $\tbf{M} \in \paren{\cal{BB}^*}^{\calO(\circuitDegree)}_{\calO(w/n)}$.

We combine <ref> and <ref> to show that $\Coyote$ can compute any linear arithmetic circuit with polylogarithmic factors in $\circuitDegree$.

For any $(nd,s,\circuitDegree,w)$-linear arithmetic circuit $\calC$ that can be represented by a matrix $\tbf{M} \in \R^{n\headDim \times n\headDim}$ multiplied by a vector $\bm{x} \in \R^{n\headDim}$, there exists an equivalent $\coyoteTuple{n}{\circuitDegree'}{\headDim}{w}{d}$ with
$\circuitDegree'=\calO(\Delta\log(w))$ such that  $\tbf{M}\bm{x} = \SHyena{\hyenaInput}{\convFilter}$ where $\bm{x}$ is the row major form of $\hyenaInput \in \R^{n \times d}$.

§.§ General Arithmetic Circuits

We are now ready to prove the result that yields the equivalency between arithmetic circuits and .

For any $(n\headDim,s,\circuitDegree,w)$-arithmetic circuit $\calC$, there exists an equivalent $\coyoteTuple{\inputLength}{\circuitDegree'}{\headDim}{\innerN}{\innerD}$ with $N = n, \circuitDegree'=\calO(\Delta\log{w})$, $N'=\calO(w), d'= d$ that simulates $\calC$.

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

Let us layer $\calC$ so that each layer $\calC_{\ell}$ for $\ell \in [L_{\calC}]$ either only has linear gates or multiplication gates, where the number of such layers $L_{\calC} = \calO(\Delta)$. The composition of all $\calC_{\ell}$ layers results in $\calC$. We use $\bm{z}^{\ell+1} \in \R^{w}$ to denote the output of the $\ell$-th layer $\calC_{\ell}$ which feeds as the input to the $(\ell+1)$-th layer $\calC_{\ell+1}$. Here, we note that if we can simulate each $\calC_{\ell}$ with , then we can simulate the entire layered circuit $\calC$ due to <ref>.

Now, if the layer $\calC_{\ell}^{\mathrm{lin}}$ is a linear layer (with only addition gates), then it can be represented by a matrix $\bm{M} \in \R^{w \times w}$ multiplied by $\bm{z}^{\ell} \in \R^{w}$ (We can append with $0$s if necessary so that the input from the previous gates can be written as $w$-length vector).
Thus, we can apply <ref> to simulate $\calC_{\ell}^{\mathrm{lin}}$ with an equivalent $\coyoteTuple{n}{\log{w}}{\headDim}{\calO(\log{w})}{d}$ model.

Next, if $\calC_{\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 $\calC_{\ell}^{\mathrm{mult}}$. Nevertheless, we can apply a $\calO(w)$ sparse
linear map $\tbf{R} \in \R^{w\times w}$ so that $\tbf{R}\bm{z}^{\ell}$ yields vectors $\bm{v}^1, \bm{v}^2,$ and $\bm{v}^3$, where $\bm{v}_1$ constitutes the “first" input to all the multiplication gates and $\bm{v}_2$ constitutes all the “second" inputs while $\bm{v}^3$ consists of all entries needed as inputs in the subsequent layers. That is, for the $i$th gate in $\calC_{\ell}^{\mathrm{mult}}$, we compute $\bm{v}^1_i \cdot \bm{v}^2_i$. This implies that for all the gates in $\calC_{\ell}^{\mathrm{mult}}$, we can simply compute $\bm{v}_1\odot \bm{v}_2$. To this end, we use the ${\tt remember}$ primitive with constant number of layers from <ref> to define a $\coyoteTuple{n}{\calO(\log{w})}{\headDim}{w}{d}$ model that remembers $\bm{v}^3$ while performing the Hadamard product of $\bm{v}^1$ with $\bm{v}^2$.

Overall, we can then collect all the resulting  layers and compose them as in <ref> to simulate $\calC$. Overall, the number of layers used is given by $\calO(\Delta\log{w})$ while the internal dimension remains fixed at $w$.

Using the outline in the proof sketch, we now delve into a detailed analysis of the arithmetic circuit $\calC$ to an equivalent model.

Let $\bm{u} \in \R^{nd}$ be the input to the arithmetic circuit $\calC$ with depth $\Delta$ and width $w$. We begin by rearranging the circuit into layers of addition and multiplication gates. That is, each layer $\calC^{\ell}$ has either all addition gates or all multiplication gates. This allows us to readily apply the results from <ref>. Note here that we can assert that the number of such layers $L_{\calC} = \calO(\Delta)$. Moreover, we construe the input to each such circuit as a vector of length ${w}$ by appending with extra 0s if necessary so that the composition of the layers results in a circuit equivalent to $\calC$. See <ref> for an example of such decomposition

every picture/.style=line width=0.75pt

(229.67,105.17) .. controls (229.67,97.71) and (235.71,91.67) .. (243.17,91.67) .. controls (250.62,91.67) and (256.67,97.71) .. (256.67,105.17) .. controls (256.67,112.62) and (250.62,118.67) .. (243.17,118.67) .. controls (235.71,118.67) and (229.67,112.62) .. (229.67,105.17) – cycle ;

(285,104.33) .. controls (285,96.97) and (290.97,91) .. (298.33,91) .. controls (305.7,91) and (311.67,96.97) .. (311.67,104.33) .. controls (311.67,111.7) and (305.7,117.67) .. (298.33,117.67) .. controls (290.97,117.67) and (285,111.7) .. (285,104.33) – cycle ;

(357.67,103.17) .. controls (357.67,95.71) and (363.71,89.67) .. (371.17,89.67) .. controls (378.62,89.67) and (384.67,95.71) .. (384.67,103.17) .. controls (384.67,110.62) and (378.62,116.67) .. (371.17,116.67) .. controls (363.71,116.67) and (357.67,110.62) .. (357.67,103.17) – cycle ;

(413,103.33) .. controls (413,95.97) and (418.97,90) .. (426.33,90) .. controls (433.7,90) and (439.67,95.97) .. (439.67,103.33) .. controls (439.67,110.7) and (433.7,116.67) .. (426.33,116.67) .. controls (418.97,116.67) and (413,110.7) .. (413,103.33) – cycle ;

(264.67,183.17) .. controls (264.67,175.71) and (270.71,169.67) .. (278.17,169.67) .. controls (285.62,169.67) and (291.67,175.71) .. (291.67,183.17) .. controls (291.67,190.62) and (285.62,196.67) .. (278.17,196.67) .. controls (270.71,196.67) and (264.67,190.62) .. (264.67,183.17) – cycle ;

(320,182.33) .. controls (320,174.97) and (325.97,169) .. (333.33,169) .. controls (340.7,169) and (346.67,174.97) .. (346.67,182.33) .. controls (346.67,189.7) and (340.7,195.67) .. (333.33,195.67) .. controls (325.97,195.67) and (320,189.7) .. (320,182.33) – cycle ;

(383.67,183.17) .. controls (383.67,175.71) and (389.71,169.67) .. (397.17,169.67) .. controls (404.62,169.67) and (410.67,175.71) .. (410.67,183.17) .. controls (410.67,190.62) and (404.62,196.67) .. (397.17,196.67) .. controls (389.71,196.67) and (383.67,190.62) .. (383.67,183.17) – cycle ;

(247.67,261.17) .. controls (247.67,253.71) and (253.71,247.67) .. (261.17,247.67) .. controls (268.62,247.67) and (274.67,253.71) .. (274.67,261.17) .. controls (274.67,268.62) and (268.62,274.67) .. (261.17,274.67) .. controls (253.71,274.67) and (247.67,268.62) .. (247.67,261.17) – cycle ;

(303,261.33) .. controls (303,253.97) and (308.97,248) .. (316.33,248) .. controls (323.7,248) and (329.67,253.97) .. (329.67,261.33) .. controls (329.67,268.7) and (323.7,274.67) .. (316.33,274.67) .. controls (308.97,274.67) and (303,268.7) .. (303,261.33) – cycle ;

(365.67,261.17) .. controls (365.67,253.71) and (371.71,247.67) .. (379.17,247.67) .. controls (386.62,247.67) and (392.67,253.71) .. (392.67,261.17) .. controls (392.67,268.62) and (386.62,274.67) .. (379.17,274.67) .. controls (371.71,274.67) and (365.67,268.62) .. (365.67,261.17) – cycle ;

(243.17,118.67) – (278.17,169.67) ;

(298.33,118) – (278.17,169.67) ;

(298.33,118) – (397.17,169.67) ;

(243.17,118.67) – (333.33,169) ;

(426.33,116.67) – (333.33,169) ;

(371.17,117.33) – (397.17,169.67) ;

(243.17,118.67) .. controls (229.33,186.83) and (229.33,189.83) .. (261.17,247.67) ;

(333.33,195.67) – (316.33,248) ;

(397.17,196.67) – (379.17,247.67) ;

(278.17,196.67) – (263.5,246.75) ;

(333.33,195.67) – (379.17,247.67) ;

(426.33,116.67) .. controls (425.33,152.83) and (427,252.25) .. (425,305.25) ;

(233.5,66.92) – (244.17,92.67) ;

(256.5,67.92) – (244.17,92.67) ;

(287.5,64.92) – (298.17,90.67) ;

(310.5,65.92) – (298.17,90.67) ;

(360,63.42) – (370.67,89.17) ;

(383,64.42) – (370.67,89.17) ;

(416,63.92) – (426.67,89.67) ;

(439,64.92) – (426.67,89.67) ;

(261.17,274.67) – (261.5,307.08) ;

(316.33,274.67) – (316.67,307.08) ;

(379.17,274.67) – (379.5,307.08) ;

(278.17,196.67) – (316.33,248) ;

(323,41.4) node [anchor=north west][inner sep=0.75pt]    $\dotsc $;

(232,96.73) node [anchor=north west][inner sep=0.75pt]    $\times $;

(289,95.73) node [anchor=north west][inner sep=0.75pt]    $\times $;

(361,93.73) node [anchor=north west][inner sep=0.75pt]    $\times $;

(416,93.73) node [anchor=north west][inner sep=0.75pt]    $\times $;

(270,173.73) node [anchor=north west][inner sep=0.75pt]    $+$;

(326,172.73) node [anchor=north west][inner sep=0.75pt]    $+$;

(391,175.73) node [anchor=north west][inner sep=0.75pt]    $+$;

(251,251.73) node [anchor=north west][inner sep=0.75pt]    $\times $;

(306,251.73) node [anchor=north west][inner sep=0.75pt]    $\times $;

(369,251.73) node [anchor=north west][inner sep=0.75pt]    $\times $;

(331,309.29) node [anchor=north west][inner sep=0.75pt]    $\dotsc $;

(454.67,94.96) node [anchor=north west][inner sep=0.75pt]    $\mathcal{C}^{\ell}_{\mathrm{mult}}$;

(453.33,174.96) node [anchor=north west][inner sep=0.75pt]    $\mathcal{C}^{\ell+1}_{\mathrm{lin}}$;

(452.67,252.29) node [anchor=north west][inner sep=0.75pt]    $\mathcal{C}^{\ell+2}_{\mathrm{mult}}$;

An example decomposition of an arithmetic circuit as layers of only addition or multiplication gates.

Let $\bm{z}^{\ell} \in \R^{w}$ denote the input to each layer $\calC_{\ell}$ for $\ell \in [L_{\calC}]$ with $\bm{z}^1 \equiv (\bm{u}, \bm{0}^{w-nd})$. It is important to note here that we may not have $\bm{z}^{\ell+1} \equiv \calC_{\ell}(\bm{z}^{\ell})$ since the inputs for gates at the $(\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 $\calC_{\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 $\tbf{M} \in \R^{w \times w'}$ such that we have

\[
\calC_{\ell}^{\mathrm{lin}}(\bm{z}^{\ell}) \equiv \tbf{M}\bm{z}^{\ell}.
\]

Here, note here that we may need entries from the vector $\bm{z}^{\ell}$ in subsequent layers. Let $s$ be the total number of such entries, then for at each index for such entries $i \in [w]$, we have the corresponding index $i_s \in [s]$. We then append the $i$th standard row vector into the matrix $\tbf{M}$ to get the matrix $\tbf{M}' \in \R^{w \times (w'+s)}$ so that we have

\[
\paren{\tbf{M}'\bm{z}^{\ell}}[j] = \begin{cases}
\paren{\tbf{M}\bm{z}^{\ell}}[j] &\text{if } j < w'\\
\bm{z}^{\ell}[i] &\text{if } j = w'+i_s
\end{cases}
\]

Here, note that we must have $w'+s \le w$ as $w$ is the width of the circuit. If needed, we then append the matrix $\tbf{M}'$ with all-zero rows $\bm{0}^w$ so that we get the matrix $\tbf{M}'' \in \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 $\tbf{M}''\bm{z}^{\ell} \equiv \bm{z}^{\ell+1}$ serving as the input of the $(\ell+1)$-th layer. Applying <ref> then yields an equivalent $\coyoteTuple{n}{\calO(\log{w})}{\headDim}{w}{d}$ model.

Multiplication Gates.
Next, we deal with the layer $\calC^{\ell}_{\mathrm{mult}}$ of multiplication gates by collecting the two inputs to each multiplication gates as $\bm{v}_1^{\ell}, \bm{v}_2^{\ell}$. Note that for the input $\bm{z}^{\ell} \in \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 $\bm{v}_3^{\ell}$. We thus need to permute the entries of $\bm{z}^{\ell}$ to get the vector $[\bm{v}_1^{\ell} :: \bm{v}_2^{\ell} :: \bm{v}_3^{\ell}]$ so that we can remember $\bm{v}_3^{\ell}$ while taking the Hadamard product $\bm{v}_1^{\ell} \odot \bm{v}_2^{\ell}$. To this end, we can achieve this permutation of entries using a $\calO(w)$-sparse linear matrix[Each row of $\tbf{R}^{\ell}$ has exactly one non-zero entry.] $\tbf{R}^{\ell} \in \R^{w \times w}$ with which is equivalently represented by an  $(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 $\coyoteTuple{n}{\calO(\log{w})}{\headDim}{w}{d}$ model. That is, $\tbf{R}^{\ell}\bm{z}^{\ell}$ has the following form:

\[
\tbf{R}^{\ell}\bm{z}^{\ell} \equiv
\begin{pmatrix}
\bm{v}_1^{\ell}\\
\hline\\
\bm{v}_2^{\ell}\\
\hline\\
\bm{v}_3^{\ell}
\end{pmatrix}.
\]

Next, we can now define a $\coyoteTuple{n}{1}{\headDim}{n}{d}$ which extracts $\bm{v}_1^{\ell}$ as the projection with the input to the remember primitive given by $\bm{y}^{\ell} \equiv (\bm{v}_2^{\ell}, \bm{v}_3^{\ell}, \bm{0})$. We then specify the remember primitive as ${\tt remember}\paren{\bm{y}^{\ell}: \bm{v}_2^{\ell}, \bm{v}_3^{\ell}, \bm{e}_0, \bm{v}_1^{\ell}}$ which computes the following

\begin{align*}
\bm{z}^{\ell+1}
&\equiv \begin{pmatrix}
\bm{v}_1^{\ell} \odot (\bm{e}_0 \ast \bm{v}_2^{\ell})\\
\hline\\
\bm{v}_3^{\ell}\\
\hline\\
\bm{0}
\end{pmatrix}    \\
&\equiv \begin{pmatrix}
\bm{v}_1^{\ell} \odot \bm{v}_2^{\ell} \\
\hline\\
\bm{v}_3^{\ell} \\
\hline\\
\bm{0}
\end{pmatrix}    \\
&\equiv \calC^{\ell}_{\mathrm{mult}}(\bm{z}^{\ell}) \in \R^{w}.
\end{align*}

Using <ref>, we know that this requires a $\coyoteTuple{n}{\calO(\log{w})}{\headDim}{w}{d}$ which remembers the entries that will serve as inputs in subsequent layers while performing the Hadamard product $\bm{v}^1 \odot \bm{v}^2$.

Overall, we can now stack all the resulting   layers and compose them as in <ref> to simulate $\calC$. Overall, the number of layers blows up by $\calO(\Delta\log{w})$ while the internal dimension remains fixed at $w$.

§.§ The Retnet Reduction and Space Complexity for AR

We are now in a position to show that there exists an equivalent  model that can simulate a Retnet layer from <ref>. Recall that for an input sequence $\bm{u} \in \R^{N \times d}$, we project the input using weight matrices $\tbf{W}_A, \tbf{W}_C, \tbf{W}_V\in \R^{d \times d}$ to $\tbf{A}, \tbf{C}, \tbf{V} \in  \R^{N \times d}$ and compute the following recurrence (see line <ref> in <ref>):

\begin{equation}
\label{eq: retnet-recurrence}
\bm{z}^{n} \equiv \gamma \bm{z}^{n-1} + \paren{\tbf{A}[n,:]}^{\top}\tbf{V}[n,:],
\end{equation}

where $\bm{z}^n \in \R^{d \times d}$. We unroll this recurrence to express each state $\bm{z}^n$ in terms of $\bm{u}$ with coefficients given by $\gamma$ and the weight matrices $\tbf{W}_A, \tbf{W}_C, \tbf{W}_V$ as follows:

\begin{align*}
\bm{z}^{0}
&\equiv \paren{\paren{\bm{u}\tbf{W}_A}[0,:]}^{\top}\paren{\bm{u}\tbf{W}_V}[0,:]\\
&\equiv \paren{\bm{u}[0,:]\tbf{W}_A}^{\top}\paren{\bm{u}[0,:]}\tbf{W}_V\\
&\equiv \tbf{W}_A^{\top}\paren{\bm{u}[0,:]}^{\top}\paren{\bm{u}[0,:]}\tbf{W}_V\\
&\equiv \tbf{W}_A^{\top} \paren{\bm{u}^{\top}[:,0] \bm{u}[0,:]}\tbf{W}_V.
\end{align*}

Similarly, we have

\begin{align*}
\bm{z}^{1}
&\equiv \gamma \bm{z}^{0} + \tbf{W}_A^{\top} \paren{\bm{u}^{\top}[:,1] \bm{u}[1,:]}\tbf{W}_V\\
&\equiv \gamma \paren{\tbf{W}_A^{\top} \paren{\bm{u}^{\top}[:,1] \bm{u}[0,:]}\tbf{W}_V} + \tbf{W}_A^{\top} \paren{\bm{u}^{\top}[:,1] \bm{u}[1,:]}\tbf{W}_V \\
&\equiv \tbf{W}_A^{\top} \paren{\gamma\paren{\bm{u}^{\top}[:,0] \bm{u}[0,:]} + \bm{u}^{\top}[:,1] \bm{u}[1,:]}\tbf{W}_V \\
\end{align*}

We can then generalize this to $n \in [N]$ to get

\begin{equation}
\label{eq: retnet-poly}
\bm{z}^{n} \equiv \tbf{W}_A^{\top} \paren{\sum_{i = 0}^{n} \gamma^{n-i}\bm{u}^{\top}[:,i] \bm{u}[i,:]}\tbf{W}_V.
\end{equation}

The above helps us infer that $\bm{z}^n$ can be expressed as a polynomial in $\{\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 <ref> to get an equivalent  model.

For a RetNet model with $\calO(d^2)$ parameters and $N$ layers, there exists an equivalent model that uses $\calO(Nd)$ parameters and $\calO(N\log{d})$ layers.

We will start by describing the arithmetic circuit that computes the states $\bm{z}^{n} \in \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 \in N$, we have

\[
\bm{z}^n \equiv \gamma\bm{z}^{n-1} + \paren{\tbf{A}[n,:]}^{\top}\tbf{V}[n,:].
\]

Since we can compute $\bm{z}^0$ in $\calO(1)$ layers and computing $\gamma \cdot \bm{z}^{n-1}$ needs exactly one layer with each entry from $\bm{z}^{n-1}$ serving as an input to a multiplication gate along with $\gamma$, the depth of the circuit is $\calO(N)$ with width $d^2$. We can thus apply <ref> to get an $(\calO(N), \calO(N\log{d}), d^2, \calO(d^2), \calO(d^2))-$ model that can compute $\bm{z}^n$ for each $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 $\{\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 $\bm{v}_i$ if $\bm{q} = \bm{k}_i$ for some $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 $\bm{x} \in \{0,1\}^n$ and Bob has an index $i \in [n]$, and the goal for the players is to output the $i$-th entry $\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.

The one-way randomized communication complexity[The randomized communication complexity of function $f$ is defined as $\min_{\pi} \norm{\pi}$, where $\pi$ ranges over all randomized protocols that can solve $f$ with probability of success at least $2/3$.] of the index problem for sending an $n$-length bit string is $\Omega(n)$.

We now use <ref> to provide a lower bound on the number of bits required by the Retnet model to solve AR.

The RetNet model requires $\Omega(N)$-bits to solve AR for $d \le \sqrt{N}$.

Consider an instance $(\bm{x}, i)$ of the index problem with $\bm{x} \in \{0,1\}^N$. We now describe the corresponding instance of the AR problem:

\begin{equation}
\label{eq: ar-index}
\{i, \bm{x}_i\}_{i = 0}^{N-1}, i.
\end{equation}

Next, consider the following one-way protocol for solving the index problem using the RetNet model. Alice with their access of $\bm{x} \in \{0,1\}^N$ generate an input for AR (without the query) as in (<ref>). Alice then runs the RetNet model on $\{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 $\bm{z}^{N-1}$ as we can reasonably assume that both have access weight matrices $\tbf{W}_A, \tbf{W}_C, \tbf{W}_V$ and the scalar $\gamma$. Since we assume that this model solves AR, the output ${\tt Out}[N, :] = \bm{x}_i$ should contain the associated value of $i$. Here, Bob can compute ${\tt Out}[N, :]$ by using the memory content sent by Alice along with the term $\gamma \bm{z}^{N-1}$:

\[
\bm{x}_i = {\tt Out}[N, :] = \tbf{C}[N, :]\bm{z}^N = \tbf{C}[N, :]\paren{\gamma \bm{z}^{N-1} +  (\tbf{A}[N,:])^{\top}\tbf{V}[N, :]}.
\]

That is, the total number of bits that are communicated in this protocol is $O(d^2)$. For $d \le \sqrt{N}$, have shown that a one-way communication protocol exists for solving the index problem exists that uses $o(N)$ communication complexity. This contradicts <ref> and hence, we conclude that the ReNet model solving AR also needs $\Omega(N)$ bits.

§.§ The Multiple-Query Associative Recall Problem

§.§.§ Introduction

In this section, we consider a general version of the associative recall problem [Ba et al., 2016].

Next, we will redefine the multiple-query associative recall problem ($\Task$) from <ref> to a slightly more general problem:

Suppose we are given an input sequence
\bm{u}[0 \cdots N-1] \triangleq \left\{\paren{\bm{k}_0, \bm{v}_0, \bm{q}_0}, \ldots, \paren{\bm{k}_{\frac{N}{3}-1}, \bm{v}_{\frac{N}{3}-1}, \bm{q}_{\frac{N}{3}-1}}\right\}
with each $\bm{k}_i, \bm{v}_i, \bm{q}_i \in C$ is a token drawn from a vocabulary of size $c = |C|$.
Our goal is then to check, for each $1 \le i \le \frac{N}{3}-1$, whether there exists $0 \le j < i$ such that $\bm{q}_i \equiv \bm{k}_j$, and if so, output $\bm{v}_{j}$.

Here, we note that it suffices to have $d \approx \log(c)$ so that $\bm{k}_i, \bm{v}_i, \bm{q}_i$ is embedded in $\{0,1\}^{d}$. However, we will specify the specific embedding being used for the results below.
Here, we construe the tokens $\bm{k}_i, \bm{q}_i$ and $\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 <ref>, considers a sequence of length $\frac N3$: $\bm{u}[0 \cdots N-1]:= \{\bm{x}_0, \ldots, \bm{x}_{N-1}\}$, where each $\bm{x}_i\in C$. The goal of <ref> is then to check, for each $1 \le i < N-1$, whether there exists $0 \le j < i$ such that $\bm{x}_i \equiv \bm{x}_j$, and if so, output $\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 $\{(\bm{x}_i, \bm{x}_{i+1}, \bm{x}_i)\}$.

As noted above, this version is more general than <ref>. Thus, the results proven in the sequel, which are proven for the above general   can be ported (with constant blowup in parameters) to get the results corresponding results for <ref>.

§.§.§ ${\Task}$ Solution via Attention

Before describing how  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.

We note that the input for the multiple-query associative recall problem $\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:

\begin{equation}
\label{eq: kqv-indices}
\begin{aligned}
\calK &= \{i \in \{0, \ldots, N-1\} \lvert\ i \equiv 0 \mod{3}\},\\
\calV &= \{i \in \{0, \ldots, N-1\} \lvert\ i \equiv 1 \mod{3}\},\\
\calQ &= \{i \in \{0, \ldots, N-1\} \lvert\ i \equiv 2 \mod{3}\},.
\end{aligned}
\end{equation}

The input $\bm{u} \in \R^{N \times d}$ to Attention for $d =3c$ is then given by

\[
\bm{u}[i,:] \equiv \begin{cases}
[\bm{k}_i:\bm{0}^c:\bm{0}^c] &\text{ if } i \in \calK \\
[\bm{0}^c:\bm{v}_i:\bm{0}^c] &\text{ if }  i \in \calV\\
[\bm{0}^c:\bm{0}^c:\bm{q}_i] &\text{ if }  i \in \calQ
\end{cases}
\]

Here, each $\bm{k}_i, \bm{v}_i, \bm{q}_i$ is embedded as a one-hot encoding in $\{0,1\}^{c}$.

Without softmax, the output for an attention layer $\tbf{O} \in \R^{N \times d}$ is given by

\begin{equation}
\label{eq: attention}
\tbf{O} \equiv \paren{\tbf{Q}\tbf{K}^{\top}}\tbf{V},
\end{equation}

where $\tbf{Q},\tbf{K},\tbf{V} \in \R^{N \times d}$
are defined as $\bm{u}\tbf{W}_Q, \bm{u}\tbf{W}_K, \bm{u}\tbf{W}_V$ for $ \bm{u} \in \R^{N \times d}$.
Instead of position embeddings, we use ALiBi, a popular technique that biases the attention scores $\tbf{Q}\tbf{K}^{\top}$ with a lower-triangular Toeplitz matrix $\tbf{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.
$\texttt{ALiBi-without-softmax}\paren{\hyenaInput[0\cdots N-1], \tbf{O}_{\mathrm{prev}}[0\cdots N-1], \tbf{B})}$
[1]
Input sequence $\hyenaInput[0 \cdots N-1] \triangleq \{\paren{\bm{k}_i, \bm{v}_i, \bm{q}_i}\}_{i=0}^{\frac{N}{3}-1}
$ with each $\bm{k}_i, \bm{v}_i, \bm{q}_i \in \{0,1\}^{3c}$, previous layer's output $\tbf{O}_{\mathrm{prev}}[0\cdots N-1] \in \R^{N \times 3c}$, and linear bias $\tbf{B} \in \R^{N \times N}$.
Add $\bm{u}_{\mathrm{curr}} \gets \hyenaInput + \tbf{O}_{\mathrm{prev}}$ as an input to this layer.
$\tbf{K}, \tbf{Q}, \tbf{V}\gets \bm{u}_{\mathrm{curr}}\tbf{W}_Q, \bm{u}_{\mathrm{curr}}\tbf{W}_K, \bm{u}_{\mathrm{curr}}\tbf{W}_V$.
$\tbf{O} \gets \paren{\tbf{Q}\tbf{K}^{\top} + \tbf{B}}\tbf{V}$
$\tbf{O}$ as the output of this layer.

Given an input $\hyenaInput\in \{0,1\}^{N \times d}$ (encoded as in <ref>) where $d = 3c$,
Attention with linear biases (even without using soft-max) solves $\Task$ for $\hyenaInput$ using ${\calO}(c^2)$ parameters, ${\calO}(Nc^2 + N^2c)$ time complexity and ${\calO}(1)$ layers.

We use two layers of attention.
We will start by specifying the projection matrices for the first layer $\tbf{W}_Q^1,\tbf{W}^1_K,\tbf{W}^1_V \in \R^{d \times d}$ as:

\[
\tbf{W}^1_K \equiv \tbf{W}^1_Q \equiv \tbf{0},\quad
\tbf{W}_V^1 \equiv \begin{pmatrix}
\tbf{0} & \tbf{0} & \tbf{0} \\
\tbf{0} & \tbf{I_{c \times c}} & \tbf{0} \\
\tbf{0} & \tbf{0} & \tbf{0}
\end{pmatrix}
\]

Above, $\tbf{W_V^1}$ is meant to isolate the $\bm{v}_i$ embeddings. For the first layer, we then have $\tbf{Q}^1 \equiv
\tbf{K}^1 \equiv \tbf{0}$, and

\begin{equation*}
\label{eq: kqv-projections}
\tbf{V}^1[i,:] := \bm{u}^1[i,:] \tbf{W}_V^1 \equiv
\begin{cases}
[\bm{0}^{c}:\bm{v}_i:\bm{0}^c] &\text{if }i \in \calV,\\
\tbf{0}^d &\text{otherwise}
\end{cases},
\end{equation*}

where $\calK, \calQ, \calV$ are defined as in (<ref>). The output for the first layer is given by the following:

\begin{align*}
\tbf{O}^1[i,:]
&= \paren{\paren{\tbf{Q}\tbf{K}^{\top}+\tbf{B}^1}\tbf{V}}[i,:] \\
&= \tbf{B}^1\tbf{V}[i,:] \\
\begin{cases}
[\bm{0}^{c}:\bm{v}_i:\bm{0}^c] &\text{if }i \in \calK\\
\bm{0}^d &\text{otherwise}
\end{cases},
\end{align*}

The $\tbf{QK^T}$ is ignored ($ \equiv \tbf{0}$) and we isolate the shifted sequence by setting the bias matrix $\tbf{B}$ appropriately. In particular, the last equality follows from the fact that $\tbf{B}^1$ is an up-shift matrix that shifts each row of $\tbf{V}$ by 1.
We apply the residual connection at the end of the standard Transformer block to insert the $\bm{k}_i$ adjacent to the $\bm{v}_i$. For the second layer, the input $\bm{u}^2 \in \R^{N \times d}$ is given by

\[
\bm{u}^2[i,:] \equiv \begin{cases}
[\bm{k}_i: \bm{v}_i:\bm{0}^{c}] &\text{if }i \in \calK\\
[\bm{0}^{c}: \bm{v}_i:\bm{0}^{c}] &\text{if }i \in \calV\\
[\bm{0}^{c}: \bm{0}^{c}:\bm{q}_i] &\text{if }i \in \calQ
\end{cases},
\]

Further, we take the following projection matrices:

\[
\tbf{W}^2_K \equiv\begin{pmatrix}
\tbf{I}_{c \times c} & \tbf{0} & \tbf{0} \\
\tbf{0} & \tbf{0} & \tbf{0} \\
\tbf{0} & \tbf{0} & \tbf{0}
\end{pmatrix},
\tbf{W}^2_Q \equiv \begin{pmatrix}
\tbf{0} & \tbf{0} & \tbf{0} \\
\tbf{0} & \tbf{0} & \tbf{0} \\
\tbf{I}_{c \times c} & \tbf{0} & \tbf{0}
\end{pmatrix},
\tbf{W}^2_V \equiv \begin{pmatrix}
\tbf{0} & \tbf{0} & \tbf{0} \\
\tbf{I}_{c \times c} & \tbf{0} & \tbf{0} \\
\tbf{0} & \tbf{0} & \tbf{0}
\end{pmatrix}
\]

We then have the following matrices as input to attention after applying projection:

\begin{equation}
\begin{aligned}
\tbf{Q}^2[i,:] &:= \bm{u}^2[i,:]\tbf{W}^2_K \equiv \begin{cases}
[\bm{q}_i:\bm{0}^{2c}] &\text{if }i \in \calQ,\\
\tbf{0}^d &\text{otherwise}
\end{cases},\\
\tbf{K}^2[i,:] &:=\bm{u}^2[i,:]\tbf{W}^2_Q \equiv  \begin{cases}
[\bm{k}_i:\bm{0}^{2c}] &\text{if }i \in \calK,\\
\tbf{0}^d &\text{otherwise}
\end{cases},\\
\tbf{V}^2[i,:] &:= \bm{u}^2[i,:]\tbf{W}^2_V \equiv
\begin{cases}
[\bm{v}_i:\bm{0}^{2c}] &\text{if }i \in \calK \cup \calV,\\
\tbf{0}^d &\text{otherwise}
\end{cases}.
\end{aligned}
\label{eq: shifted-value}
\end{equation}

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

\begin{align*}
\paren{\tbf{Q}^2\tbf{K}^{2\top}}[i,j] = \paren{\tbf{Q}^2\tbf{K}^{2\top}}[i,j]
&= \angles{\tbf{Q}^2[i,:], \tbf{K}^{2\top}[:,j]} \\
&= \angles{\tbf{Q}^2[i,:], \tbf{K}^2[j,:]} \\
&= \begin{cases}
\angles{\bm{q}_i, \bm{k}_j} &\text{if }i \in \calQ, j \in \calK \\
0 &\text{otherwise}
\end{cases}\\
&= \begin{cases}
1 &\text{if }i \in \calQ, j \in \calK, \bm{q}_i \equiv \bm{k}_j \equiv \bm{e}_k\text{ for some $k$}\\
0 &\text{otherwise}.
\end{cases}
\end{align*}

Finally, we compute the output as follows:

\begin{align*}
\tbf{O}^2[i,:]
&= \paren{\paren{\tbf{Q}^2\tbf{K}^{2\top}+{\tbf{B}^2}}\tbf{V}^2}[i,:] \\
&= \paren{\paren{\tbf{Q}^2\tbf{K}^{2\top}+{\tbf{0}}}\tbf{V}^2}[i,:] \\
&= \paren{\tbf{Q}^2\tbf{K}^{2\top}}[i,:]\cdot\tbf{V}^2  \\
&= \sum_{j = 0}^{N-1} \paren{\tbf{Q}^2\tbf{K}^{2\top}}[i,j]\cdot \tbf{V}^2[j,:]  \\
&= \sum_{j \in \calK} \paren{\tbf{Q}^2\tbf{K}^{2\top}}[i,j]\cdot [\bm{v}_j:\bm{0}^{2c}] \\
\begin{cases}
[\bm{v}_j:\bm{0}^{2c}] &\text{if }j \in \calK, i \in \calQ,  \bm{q}_i \equiv \bm{k}_j \\
\bm{0}^d &\text{otherwise}
\end{cases},
\end{align*}

where we use the fact that for each index $j \in \calK$, the matrix $\tbf{V}^2$ contains the associated value from (<ref>).    Thus, for each query $\bm{q}_i$, we solve the associated value problem yielding a match for the $j$th key.

In total, we only need $\calO(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 $\tbf{QK^{\top}}$ in ${\calO}(N^2c)$, and projections in ${\calO}(Nc^2)$. Finally, we only need $\calO(1)$ layers for this solution.

In the sequel, we develop a parallel algorithm to solve the multiple-query associative recall problem with  $\calO(Nd\cdot\log^2{N})$ work complexity and $\calO(d \cdot \log^2{N})$ time. We then convert the algorithm into a  model via the route of arithmetic circuits, which then solves the multiple-query associative recall problem with $\tilde{\calO}(1)$ layers and $\tilde{\calO}(Nd)$ parameters.

§.§.§ 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 $\calO(\log{N}\cdot d)$ ${\tt insert}$ and ${\tt find}$ operations for $d$-bit entries should be sufficient.

$\texttt{Sequential-MQ-AR}\paren{\hyenaInput}[0\cdots N-1]$
[1]
Input sequence $\hyenaInput[0 \cdots N-1] \triangleq \{\paren{\bm{k}_i, \bm{v}_i, \bm{q}_i}\}_{i=0}^{\frac{N}{3}-1}
$ with each $\bm{k}_i, \bm{v}_i, \bm{q}_i \in \{0,1\}^{\headDim}$.
Initialize an associative array with ${\tt insert}$ and ${\tt find}$ and an output array $\texttt{out} \gets [].$
$i \in \{0, \ldots, \frac{N}{3}-1$}
$(\bm{k}_j, \bm{v}_j) \gets {\tt find}(\bm{q}_i)$ Query for $\bm{q}_i$ in the data structure.
$\bm{k}_j$ is not null
Add $\bm{v}_{j}$ to $\texttt{out}$.
${\tt insert}(\bm{k}_i, \bm{v}_{i})$ Add the key-value pair to the data structure.

<ref> solves the multiple-query associative recall problem $\paren{\Task}$ in $\calO(dN\log{N})$ time for an input sequence $\bm{u} \in \{0,1\}^{\inputLength \times \headDim}$.

For any $i \in \{0,\ldots, \frac{N}{3}-1\}$, we know that both insertion and lookup operations take $\calO(\log(i)\cdot d)$ time. Overall, the runtime of the algorithm is

\[\sum_{i=0}^{\frac{N}{3}-1} \calO(\log(i)\cdot d) = \calO(\log(N!)\cdot d) = \calO(N \log{N}\cdot d).\]

§.§.§ 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 <ref> will lead to a  model that solves the multiple-query associative recall problem. With respect to <ref>, it may be the case that the arithmetic circuit has a large number of layers $\Omega(N)$. Unfortunately, this would imply that the resulting  model may have near quadratic complexity in $N$. Instead, we now initiate our effort into designing a  model  with both small enough number of parameters and number of layers. Here, we will first subdivide the problem using dyadic intervals into $\calO(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 \triangleq a_0\le \ldots \le a_{n-1}$ and $B \triangleq (b_0\le \ldots \le b_{m-1})$ with $n \le m$, for each $a_j \in A$, the goal is to find the smallest element in $B$ that is larger than or equal to $a_j$.

The multiple search problem is solved by a parallel binary search (pbs) algorithm in [Akl and Meijer, 1990] with work complexity $\calO(n \cdot \log{m})$ and time $\calO(\log{n}\log{m})$. Specifically, for sorted arrays $A[0 \cdots n-1]$ and $B[0 \cdots m-1]$, ${\tt pbs}$ constructs the array $C[0\cdots n-1]$ defined as

\begin{equation}
C[i] \triangleq
\begin{cases}
\min_{0 \le j < m} \{j\lvert\ A[i] \le B[j]\} &\text{if } A[i] \le B[m-1]\\
m &\text{otherwise.}
\end{cases}
\label{eq: pbs-output}
\end{equation}

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.
$\texttt{pbs-key-values}\paren{\bm{q}[s \cdots t],\ \bm{k}[x \cdots y], n, m}$
[1]
sorted arrays $\bm{q}[s \cdots t]:= \{\bm{q}_i\}_{i = s}^t,\ \bm{k}[x\cdots y]:= \{(j, \bm{k}_j)\}_{j = x}^y$.
Initialize $n$ processors denoted $P_0, P_1, \ldots, P_{n-1}$
$\{$Sequential steps are assumed to be executed by $P_s$.$\}$
Initialize the output array $C := [m]_{i=s}^t$.
$s \le t$
$\mathrm{mid} \gets \floor{(s+t)/2}$
$\bm{q}[\mathrm{mid}] \le \bm{k}[x][1]$
$i := s$ to $\mathrm{mid}$ in parallel
$C[i] \gets j$ Step executed in parallel by $P_i$
$\texttt{pbs-key-values}\paren{\bm{q}[\mathrm{mid}+1 \cdots t],\ \bm{k}[x \cdots y]}$
$\bm{q}[\mathrm{mid}] > \bm{k}[y][1]$
$i := \mathrm{mid}$ to $t$ in parallel
$C[i] \gets {y+1}$ Step executed in parallel by $P_i$
$\texttt{pbs-key-values}\paren{\bm{q}[s \cdots \mathrm{mid}-1],\ \bm{k}[x \cdots y]}$
C[mid] is determined using sequential binary search
$z \gets \min_{x \le j \le y} \{j\lvert\ \bm{q}[\mathrm{mid}] \le \bm{k}[j][1]\}$
$C[\mathrm{mid}] \gets z$
do steps <ref> and <ref> in parallel
$\texttt{pbs-key-values}\paren{\bm{q}[s \cdots \mathrm{mid}-1],\ \bm{k}[x \cdots z-1]}$]
$\texttt{pbs-key-values}\paren{\bm{q}[\mathrm{mid}+1 \cdots t],\ \bm{k}[z \cdots y]}$

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

$\texttt{Parallel-MQAR}\paren{\hyenaInput[0 \cdots N-1]}$
[1]
Input sequence $\hyenaInput[0 \cdots N-1] \triangleq \{\paren{\bm{k}_i, \bm{v}_i, \bm{q}_i}\}_{i=0}^{\frac{N}{3}-1}
$ with each $\bm{k}_i, \bm{v}_i, \bm{q}_i \in \{0,1\}^{\headDim}$.
Initialize $\frac{N}{3}\log\paren{\frac{N}{3}}$ processors denoted $P_0, \ldots, P_{\frac{N}{3}\log\paren{\frac{N}{3}}-1}$.
Initialize the index and output array $\texttt{idx, val} \gets [].$
$k := \{0, \ldots, \log\paren{\frac{N}{3}}-1\}$
$\bm{x} := \{\bm{x} \in \Sigma^{\log\paren{\frac{N}{3}}-k}\lvert\ \bm{x}[\log\paren{\frac{N}{3}}-k-1] = 0\}$
{All the steps below are executed in parallel by $\{\{\{P_{i}^{\bm{x}, k}\}_{\bm{x}}\}_{i \in [0, 2k-1]}\}_{k}$}
$I_k^{\bm{x}}\gets \{\bm{y} \in \Sigma^{\log\paren{\frac{N}{3}}} \lvert\ \bm{x} \in \mathrm{prefix}(\bm{y})\}.$
$\bm{k}_{\mathrm{sorted}}^{k\bm{x}}, I_{\mathrm{permuted}}^{k\bm{x}} \gets {\tt sort}\paren{\{\bm{k}_{\mathrm{dec}(i)}\}_{i \in I^{\bm{x}}_k}}$
$I_{\mathrm{permuted}}^{k\bm{x}} := \{(j, \mathrm{dec}(i))\lvert\  \bm{k}_{\mathrm{sorted}}^{k\bm{x}}[j] \equiv \bm{k}_{\mathrm{dec}(i)}\}$
$\bm{x}[\log\paren{\frac{N}{3}}-k-1] \gets 1$
$J_k^{\bm{x}} \gets \{\bm{y} \in \Sigma^{\log\paren{\frac{N}{3}}} \lvert\ \bm{x} \in \mathrm{prefix}(\bm{y})\}.$
$\bm{q}_{\mathrm{sorted}}^{k\bm{x}}, J_{\mathrm{permuted}}^{k\bm{x}} \gets {\tt sort}\paren{\{\bm{q}_{\mathrm{dec}(j)}\}_{j \in J^{\bm{x}}_k}}$  $J_{\mathrm{permuted}}^{k\bm{x}} := \{(\mathrm{dec}(j), k)\lvert\  \bm{q}_{\mathrm{sorted}}^{k\bm{x}}[k] \equiv \bm{q}_{\mathrm{dec}(j)}\}$

$C_k \gets \texttt{pbs-key-values}\paren{\bm{q}_{\mathrm{sorted}}^{k\bm{x}},\ \bm{k}_{\mathrm{sorted}}^{k\bm{x}}, 2^{k}, 2^{k} } $
$j \in  J^{\bm{x}}_k$
$C_k[\mathrm{dec}(j)] \neq 2^k$
$\bm{c}^{k\bm{x}}_j \gets C_k[J_{\mathrm{permuted}}^{k\bm{x}} (\mathrm{dec}(j))]$
$\bm{c}^{k\bm{x}}_j \neq 2^{k}$ cf. (<ref>)
Add $I_{\mathrm{permuted}}^{k\bm{x}} (\bm{c}^{k\bm{x}}_j)$ to $\texttt{idx}[\mathrm{dec}(j)]$.

$i \in \{1, \ldots, \frac{N}{3}-1\}$
{Executed in parallel by $P_i$.}
$\exists\ j \in {\tt idx}[i]$
Add $\bm{v}_{j+1}$ to $\texttt{val}$

In lines <ref> and <ref>, 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 <ref>.

§.§.§ Correctness and Complexity

<ref> solves the multiple-query associative recall problem with work complexity $\calO(Nd \cdot \log^2{N})$ and time $\calO(d \cdot \log^2{N})$.

The correctness of pbs implies the correctness of <ref> if we can show that, for each $1 \le i < \frac{N}{3}$, we check for a match among the keys $\{\bm{k}_j\}_{j \in [i-1]}$. To this end, for each $1 \le i < \frac{N}{3}$, let the set of all iterator indices associated with an index $i$ be defined as  $K_i \triangleq \{(k,x)\lvert i \in J_k^{\bm{x}}\}$ with $J_k^{\bm{x}}$ as noted in line <ref>. Then, we define the corresponding set for keys as
\(
\calI_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$ is part of (given by $K_i$) where the algorithm checks for a match among the keys in $\calI_i$, it then suffices to show that $\calI_i = [i-1]$.

Here, first note that if some index $j \in I_k^{\bm{x}}\subseteq \calI_i$ for some $\bm{x} \in \Sigma^{\log\paren{\frac{N}{3}}-k}$, then, by definition, $\bm{x} \in \mathrm{prefix}(\mathrm{bin}(j))$. Here, let $\bm{x}^1 := \bm{x}\lvert_{\bm{x}[\log\paren{\frac{N}{3}}-k] = 1}$ where we set the $(\log\paren{\frac{N}{3}}-k)$-th index of $\bm{x}$ to be 1. Consequently, as we have $i \in J^{\bm{x}}_k$ for the same $k$ and $\bm{x}$ as in $I_k^{\bm{x}}$ (cf. line <ref>), we must have $\bm{x}^1 \in \mathrm{prefix}(\mathrm{bin(i)})$. Thus we get $j < i$, whence we can claim that $\calI_i \subseteq [i-1]$.

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

Overall, we have shown that $\calI_i = [i-1]$. Since this holds for all $1 \le i < \frac{N}{3}$, we can conclude that <ref> solves the multiple-query associative recall problem.

Next, it is easy to see that we execute lines <ref> to <ref> $\sum_{k=0}^{\log\paren{\frac{N}{3}}-1} \frac{\frac{N}{3}}{2^{k+1}}$-many times.
We note that sorting $n$ values each of size $d$ can be done with work complexity $n\log{n}\cdot d$. We note that, at each instance, we are sorting sorting $2^k$ values. Meanwhile, remembering $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$ which has work complexity $n \log{m}$. Finally, we know that the work complexity of lines <ref> to <ref> is $\calO(N)$. Thus, the overall work complexity of <ref> is

\begin{equation}
\label{eq: work-complexity}
d \cdot \sum_{k=0}^{\log\paren{\frac{N}{3}}-1} \frac{\frac{N}{3}}{2^{k+1}} \calO(2^k\cdot \log{2^k}) = d \cdot \calO(N) \cdot \sum_{k=0}^{\log{N/3}-1} \calO(k) = \calO(Nd \cdot \log^2{N}).
\end{equation}

We will now analyze the depth of <ref>. We know that the depth of computation for <ref> is $\calO(\log{n}\log{m})$ for input sizes. Moreover, we have $\calO(1)$ depth for the computation in <ref> to <ref> as each entry in ${\tt idx}$ can have at most one entry. Since the nested for loops iterating over $k$s and the associated $\bm{x}$s runs in parallel, the depth of <ref> is dominated by the largest depth among all calls to
and to sort. The largest such call to  is of size $n = m = 2^{\log\paren{\frac{N}{3}}-1} = N/6$ which yields a depth of $d \cdot \log^2{\frac{N}{3}}$. Moreover, using sorting networks (<ref>), we know that the largest depth is for sorting $\frac{N}{6}$ values of size $d$ given by $d \cdot \Theta(\log{N})$ (<ref>). Thus, we can conclude that <ref> takes $\calO(d \cdot \log^2{N})$,
time where $N$ is the length of the input.

§.§.§ Conversion to Arithmetic Circuit

We will convert <ref> to an arithmetic circuits modularly. In particular, after writing out an explicit circuit for <ref>, we will uses this circuit as a black-box along with circuits for sorting networks.

Circuit for .
We will denote the corresponding arithmetic circuit for <ref>  as  as well with the input gates comprising of each entry from $\bm{q}[s \cdots t]$ and $\bm{k}[x \cdots y]$ and the $i$-th output gate yielding the value $C[i]$ as in <ref>.

Here, we first convert the comparisons for the if statements in <ref>. To this end, we briefly introduce comparators.

A comparator is a device with inputs $\bm{x}$ and $\bm{y}$ and outputs $\bm{x}'$ and $\bm{y}'$, that performs the following function:

\[
\begin{aligned}
\bm{x}'  &=  \min(\bm{x}, \bm{y}),\\
\bm{y}'  &=  \max(\bm{x}, \bm{y}).
\end{aligned}
\]

Using comparators, we can use bit-wise XOR and AND to  define the result of the comparisons in lines <ref> and <ref> as the following fixed variables:

\begin{equation}
\label{eq: comparison-gates}
\begin{aligned}
\ell_x &:= \mathbbm{1}\{\bm{q}[\mathrm{mid}] \leq \bm{k}_{x}\},\\
g_{y} &:= \mathbbm{1}\{\bm{q}[\mathrm{mid}] > \bm{k}_{y}\},\\
z &:= \texttt{bin-search}({\{\bm{k}_{i}\}_{i=s}^{t} ,\ \bm{q}[\mathrm{mid}]}).
\end{aligned}
\end{equation}

This then allows us to infer the index of the key array for each recursive call to $\texttt{pbs-key-values}$ in lines <ref>, <ref>, <ref>, and <ref> from <ref>. Specifically, let $z_s$ and $z_t - 1$ denote the starting and ending indices for the keys as inputs to the recursive calls in <ref>) below:

\begin{align}
&\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})}
\label{eq: rec-1}\\
&\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})}
\label{eq: rec-2}
\end{align}

Here, $z_t$ and $z_s$ can assume values dependent on the results of the comparisons in lines <ref> and <ref>. Specifically, we have

\begin{align*}
z_t &= \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 &= 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}.
\end{align*}

Here, $z_s$ or $z_t$ getting a value $0$ signifies that the branch is dead, and we do not execute the recursive call.

Finally, let the arrays $C_t[\mathrm{mid} +1\cdots t]$ and $C_s[ s\cdots \mathrm{mid} -1]$ denote the outputs to the recursive calls in (<ref>) and (<ref>), respectively. We can then succinctly express the outputs for each index of the output array $C$ as

\begin{equation}
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}
\end{equation}

We can thus state the circuit schematically in <ref>.

Now, before accounting for the complexity of the circuit for <ref>, we must first assert the complexity of the comparators that we use in <ref>.

every picture/.style=line width=0.75pt

(184,117.01) – (255.33,117.01) – (255.33,134.51) – (184,134.51) – cycle ;

(395,118.4) – (470,118.4) – (470,134.84) – (395,134.84) – cycle ;

(270,119.09) – (380,119.09) – (380,147.01) – (270,147.01) – cycle ;

(206,90) – (206,117.01) ;
[shift=(206,117.01), rotate = 270] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(411,86) – (206,117.01) ;
[shift=(206,117.01), rotate = 341] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(411,86) – (314,119.01) ;
[shift=(314,119.01), rotate = 325] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(411,86) – (425,118.01) ;
[shift=(425,118.01), rotate = 195] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(347,88.01) – (425,118.01) ;
[shift=(425,118.01), rotate = 245] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(206,134) – (206,158.51) ;
[shift=(206,158.51), rotate = 270] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(425,136) – (425,158.51) ;
[shift=(425,158.51), rotate = 270] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(196.75,167.51) .. controls (196.75,162.54) and (200.89,158.51) .. (206,158.51) .. controls (211.11,158.51) and (215.25,162.54) .. (215.25,167.51) .. controls (215.25,172.48) and (211.11,176.51) .. (206,176.51) .. controls (200.89,176.51) and (196.75,172.48) .. (196.75,167.51) – cycle ;

(415.75,167.51) .. controls (415.75,162.54) and (419.89,158.51) .. (425,158.51) .. controls (430.11,158.51) and (434.25,162.54) .. (434.25,167.51) .. controls (434.25,172.48) and (430.11,176.51) .. (425,176.51) .. controls (419.89,176.51) and (415.75,172.48) .. (415.75,167.51) – cycle ;

(211.75,301.51) .. controls (211.75,296.54) and (215.89,292.51) .. (221,292.51) .. controls (226.11,292.51) and (230.25,296.54) .. (230.25,301.51) .. controls (230.25,306.48) and (226.11,310.51) .. (221,310.51) .. controls (215.89,310.51) and (211.75,306.48) .. (211.75,301.51) – cycle ;

(496.75,258.51) .. controls (496.75,253.54) and (500.89,249.51) .. (506,249.51) .. controls (511.11,249.51) and (515.25,253.54) .. (515.25,258.51) .. controls (515.25,263.48) and (511.11,267.51) .. (506,267.51) .. controls (500.89,267.51) and (496.75,263.48) .. (496.75,258.51) – cycle ;

(326,258.51) .. controls (326,253.54) and (330.14,249.51) .. (335.25,249.51) .. controls (340.36,249.51) and (344.5,253.54) .. (344.5,258.51) .. controls (344.5,263.48) and (340.36,267.51) .. (335.25,267.51) .. controls (330.14,267.51) and (326,263.48) .. (326,258.51) – cycle ;

(285.75,219.51) .. controls (285.75,214.54) and (289.89,210.51) .. (295,210.51) .. controls (300.11,210.51) and (304.25,214.54) .. (304.25,219.51) .. controls (304.25,224.48) and (300.11,228.51) .. (295,228.51) .. controls (289.89,228.51) and (285.75,224.48) .. (285.75,219.51) – cycle ;

(133.75,252.51) .. controls (133.75,247.54) and (137.89,243.51) .. (143,243.51) .. controls (148.11,243.51) and (152.25,247.54) .. (152.25,252.51) .. controls (152.25,257.48) and (148.11,261.51) .. (143,261.51) .. controls (137.89,261.51) and (133.75,257.48) .. (133.75,252.51) – cycle ;

(285.75,219.51) – (216,205.51) ;
[shift=(285.75,219.51), rotate = 190] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(418,206.51) – (304.25,219.51) ;
[shift=(304.25,219.51), rotate = 350] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(315,147) – (331,203.51) ;

(339,225.51) – (335.25,249.51) ;
[shift=(335.25,249.51), rotate = 280] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(331,203.51) .. controls (338,193.51) and (346,225.51) .. (339,225.51) ;

(295,228.51) – (326,258.51) ;
[shift=(326,258.51), rotate = 220] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(206,134) – (143,243.51) ;
[shift=(143,243.51), rotate = 230] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(425,136) – (506,249.51) ;
[shift=(506,249.51), rotate = 235] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(133,231.51) – (143,243.51) ;
[shift=(143,243.51), rotate = 300] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(525,231.51) – (506,249.51) ;
[shift=(506,249.51), rotate = 317] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(215.25,167.51) – (235,149.51) ;
[shift=(215.25,167.51), rotate = 317] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(400,148.51) – (415.75,167.51) ;
[shift=(415.75,167.51), rotate = 230] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(430.75,300.51) .. controls (430.75,295.54) and (434.89,291.51) .. (440,291.51) .. controls (445.11,291.51) and (449.25,295.54) .. (449.25,300.51) .. controls (449.25,305.48) and (445.11,309.51) .. (440,309.51) .. controls (434.89,309.51) and (430.75,305.48) .. (430.75,300.51) – cycle ;

(506,267.51) – (449.25,300.51) ;
[shift=(449.25,300.51), rotate = 325] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(335.25,267.51) – (430.75,300.51) ;
[shift=(430.75,300.51), rotate = 200] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(230.25,301.51) – (335.25,267.51) ;
[shift=(230.25,301.51), rotate = 340] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(211.75,301.51) – (143,261.51) ;
[shift=(211.75,301.51), rotate = 210] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(221,310.51) – (221,352.51) ;
[shift=(221,352.51), rotate = 270] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(440,309.51) – (440,353.51) ;
[shift=(440,353.51), rotate = 270] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(162.42,469.51) .. controls (162.42,464.54) and (166.56,460.51) .. (171.67,460.51) .. controls (176.78,460.51) and (180.92,464.54) .. (180.92,469.51) .. controls (180.92,474.48) and (176.78,478.51) .. (171.67,478.51) .. controls (166.56,478.51) and (162.42,474.48) .. (162.42,469.51) – cycle ;

(162.75,420.51) .. controls (162.75,415.54) and (166.89,411.51) .. (172,411.51) .. controls (177.11,411.51) and (181.25,415.54) .. (181.25,420.51) .. controls (181.25,425.48) and (177.11,429.51) .. (172,429.51) .. controls (166.89,429.51) and (162.75,425.48) .. (162.75,420.51) – cycle ;

(337.17,326.51) .. controls (337.17,321.54) and (341.31,317.51) .. (346.42,317.51) .. controls (351.53,317.51) and (355.67,321.54) .. (355.67,326.51) .. controls (355.67,331.48) and (351.53,335.51) .. (346.42,335.51) .. controls (341.31,335.51) and (337.17,331.48) .. (337.17,326.51) – cycle ;

(90.94,352.51) – (340,352.51) – (340,370.51) – (90.94,370.51) – cycle ;

(350,353.51) – (625.67,353.51) – (625.67,368.51) – (350,368.51) – cycle ;

(143,261.51) – (337.17,326.51) ;
[shift=(337.17,326.51), rotate = 200] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(355.67,326.51) – (434,307.51) ;
[shift=(355.67,326.51), rotate = 345] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(172,429.51) – (171.67,460.51) ;
[shift=(171.67,460.51), rotate = 270] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(171.67,371.51) – (172,411.51) ;
[shift=(172,411.51), rotate = 270] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(528.75,420.51) .. controls (528.75,415.54) and (532.89,411.51) .. (538,411.51) .. controls (543.11,411.51) and (547.25,415.54) .. (547.25,420.51) .. controls (547.25,425.48) and (543.11,429.51) .. (538,429.51) .. controls (532.89,429.51) and (528.75,425.48) .. (528.75,420.51) – cycle ;

(537.67,371.51) – (538,411.51) ;
[shift=(538,411.51), rotate = 270] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(435,205.51) .. controls (748,229.51) and (675,404.51) .. (547.25,420.51) ;
[shift=(547.25,420.51), rotate = 345] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(206,175.51) – (206,198.51) ;
[shift=(206,198.51), rotate = 270] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(196.75,207.51) .. controls (196.75,202.54) and (200.89,198.51) .. (206,198.51) .. controls (211.11,198.51) and (215.25,202.54) .. (215.25,207.51) .. controls (215.25,212.48) and (211.11,216.51) .. (206,216.51) .. controls (200.89,216.51) and (196.75,212.48) .. (196.75,207.51) – cycle ;

(198,205.51) .. controls (-104,246.51) and (78.33,400.51) .. (162.75,420.51) ;
[shift=(162.75,420.51), rotate = 200] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(426,175.51) – (426,198.51) ;
[shift=(426,198.51), rotate = 270] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(416.75,207.51) .. controls (416.75,202.54) and (420.89,198.51) .. (426,198.51) .. controls (431.11,198.51) and (435.25,202.54) .. (435.25,207.51) .. controls (435.25,212.48) and (431.11,216.51) .. (426,216.51) .. controls (420.89,216.51) and (416.75,212.48) .. (416.75,207.51) – cycle ;

(235.33,177.91) node [anchor=north west][inner sep=0.75pt]  [font=]  $1$;

(215,203) – (235.33,187.91) ;
[shift=(215,203), rotate = 325] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(383.33,176.91) node [anchor=north west][inner sep=0.75pt]  [font=]  $1$;

(417,203) – (390.33,186.91) ;
[shift=(417,203), rotate = 210] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(538,429.51) – (538.67,459.51) ;
[shift=(538.67,459.51), rotate = 270] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(346.42,335.51) – (344.67,510.51) ;
[shift=(344.67,510.51), rotate = 270] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(143,261.51) .. controls (60.33,351.51) and (61.33,412.51) .. (162.42,469.51) ;
[shift=(162.42,469.51), rotate = 215] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(171.67,478.51) – (171.33,509.51) ;
[shift=(171.33,509.51), rotate = 270] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(529.42,468.51) .. controls (529.42,463.54) and (533.56,459.51) .. (538.67,459.51) .. controls (543.78,459.51) and (547.92,463.54) .. (547.92,468.51) .. controls (547.92,473.48) and (543.78,477.51) .. (538.67,477.51) .. controls (533.56,477.51) and (529.42,473.48) .. (529.42,468.51) – cycle ;

(506,267.51) .. controls (687,371.51) and (660,384.51) .. (547.92,468.51) ;
[shift=(547.92,468.51), rotate = 320] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(538.67,477.51) – (538.33,508.51) ;
[shift=(538.33,508.51), rotate = 270] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(184,120) node [anchor=north west][inner sep=0.75pt]  [font=]  $\bm{q}[\mathrm{mid}] \leq \bm{k}_{x}$;

(396,120) node [anchor=north west][inner sep=0.75pt]  [font=]  $\bm{q}[\mathrm{mid}] > \bm{k}_{y}$;

(275,118.09) node [anchor=north west][inner sep=0.75pt]  [font=] [align=left] ;

(270,128) node [anchor=north west][inner sep=0.75pt]  [font=]  $({\{\bm{k}_{i}\}_{i=s}^{t} ,\ \bm{q}[\mathrm{mid}]})$;

(-70, 70.0) node [anchor=north west][inner sep=0.75pt]  [font=]   $\texttt{pbs-key-values}\paren{\bm{q}[s \cdots t],\ \bm{k}[x \cdots y], n, m}$;

(-75,65.5) – (175,65.5) – (175,85.5) – (-75,85.5) – cycle ;

(-100,60.0) – (680,60.0) – (680,550) – (-100,550) – cycle ;

(357,70.0) node [anchor=north west][inner sep=0.75pt] [font=]   $\bm{q}_{s}$;

(397,67.0) node [anchor=north west][inner sep=0.75pt]  [font=]  $\bm{q}[\mathrm{mid}]$;

(458,70.0) node [anchor=north west][inner sep=0.75pt]  [font=]  $\bm{q}_{t}$;

(375,78.0) node [anchor=north west][inner sep=0.75pt]  [font=]  $.....$;

(202,70.0) node [anchor=north west][inner sep=0.75pt]   [font=] $\bm{k}_{x}$;

(331,70.0) node [anchor=north west][inner sep=0.75pt] [font=]   $\bm{k}_{y}$;

(225,78.0) node [anchor=north west][inner sep=0.75pt] [font=]   $.......................$;

(433,78.0) node [anchor=north west][inner sep=0.75pt]  [font=]  $.....$;

(121,222.4) node [anchor=north west][inner sep=0.75pt]  [font=]  $x$;

(526,228.4) node [anchor=north west][inner sep=0.75pt]  [font=]  $y+1$;

(286.75,210.91) node [anchor=north west][inner sep=0.75pt]    $\times $;

(326.75,249.91) node [anchor=north west][inner sep=0.75pt]    $\times $;

(134.75,243.91) node [anchor=north west][inner sep=0.75pt]    $\times $;

(497.75,250.91) node [anchor=north west][inner sep=0.75pt]    $\times $;

(235.33,137.91) node [anchor=north west][inner sep=0.75pt]  [font=]  $-1$;

(383.33,136.91) node [anchor=north west][inner sep=0.75pt]  [font=]  $-1$;

(213,293.4) node [anchor=north west][inner sep=0.75pt]    $+$;

(432,292.4) node [anchor=north west][inner sep=0.75pt]    $+$;

(205,340.91) node [anchor=north west][inner sep=0.75pt] [font=] $z_{s}$;

(90.94,356.07) node [anchor=north west][inner sep=0.75pt]  [font=] [align=left] ($\displaystyle \bm{q}[ s\cdots \mathrm{mid} -1,\ \bm{k}[ z_{s} \cdots y]$);

(350.94,355.07) node [anchor=north west][inner sep=0.75pt]  [font=] [align=left] ($\displaystyle \bm{q}[\mathrm{mid} +1\cdots t] ,\ \bm{k}[ x\ \cdots z_{t} -1]$);

(208,137.4) node [anchor=north west][inner sep=0.75pt]  [font=]  $\ell_{x}$;

(410,137.4) node [anchor=north west][inner sep=0.75pt]  [font=]  $g_{y}$;

(306,148.4) node [anchor=north west][inner sep=0.75pt]  [font=]  $z$;

(446,340.91) node [anchor=north west][inner sep=0.75pt]  [font=]  $z_{t} $;

(339,318.4) node [anchor=north west][inner sep=0.75pt]    $+$;

(163.75,412.91) node [anchor=north west][inner sep=0.75pt]    $\times $;

(529.75,412.91) node [anchor=north west][inner sep=0.75pt]    $\times $;

(172.25,489.91) node [anchor=north west][inner sep=0.75pt]    $\dotsc $;

(521,522.4) node [anchor=north west][inner sep=0.75pt]    $...........$;

(198,159.4) node [anchor=north west][inner sep=0.75pt]    $\times$;

(198,199.4) node [anchor=north west][inner sep=0.75pt]    $+$;

(417,159.4) node [anchor=north west][inner sep=0.75pt]    $\times$;

(417,199.4) node [anchor=north west][inner sep=0.75pt]    $+$;

(164,462.4) node [anchor=north west][inner sep=0.75pt]    $+$;

(530,460.4) node [anchor=north west][inner sep=0.75pt]    $+$;

(509.25,488.91) node [anchor=north west][inner sep=0.75pt]    $\dotsc $;

(323,515.4) node [anchor=north west][inner sep=0.75pt]  [font=]  $C[\mathrm{mid}]$;

(455,515.4) node [anchor=north west][inner sep=0.75pt]  [font=]  $C[\mathrm{mid} -1]$;

(120,515.4) node [anchor=north west][inner sep=0.75pt]  [font=]  $C[ s]$;

(205,515.4) node [anchor=north west][inner sep=0.75pt]  [font=]  $C[\mathrm{mid} -1]$;

(148,522.4) node [anchor=north west][inner sep=0.75pt]    $...........$;

(575,515.06) node [anchor=north west][inner sep=0.75pt]  [font=]  $C[ t]$;

$\texttt{pbs-key-values}\paren{\bm{q}[s \cdots t],\ \bm{k}[x \cdots y], n, m}$ as a circuit with recursive calls and subprocedures as “black boxes.”

For binary strings $\bm{x}, \bm{y} \in \Sigma^{d}$ of length $d$, there exists a comparison network of size $\calO(d)$, width $\calO(d)$, and depth $\calO(\log{d})$.

We use this lemma to compute the complexity of the arhmetic circuit from <ref>.

There exists an $(\calO((n+m)d), \calO(nd\cdot(\log{m}+\log{n})), \calO(\log{n}(\log{m}+\log{n}\log{d})), \calO(nd))$-arithmetic circuit[Recall that a $(n,s,\Delta, w)$-arithmetic circuit is an $n$-variate circuit with size $s$, depth at most $\Delta$, and width $w$.] equivalent to pbs-key-values (<ref>) with inputs $\bm{q}$ and $\bm{k}$ of lengths $n$ and $m$ with $d$-bit entries.

The size of the circuit for  should equal the work-complexity of <ref> but we also need to account for the comparison gates in (<ref>). Further, the circuit for binary search also has a size of $\calO(d \cdot n)$ instead of $\calO(d \cdot \log{n})$ work as in <ref>. Using <ref>, along with the fact that the comparison gates and the binary search are used at most $\calO(\log{n})$ times, we deduce that we are adding $\calO(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  is $\calO(dn\log{m}+nd\log{n})$. Further, the depth of the circuit here is determined by the runtime of <ref> along with the depth of the comparison gates and binary search $\calO(\log{n}\log{d})$, and finally, at most $n$ processors are used for the parallel algorithm in <ref>, which yields the width of the circuit.

Circuit for .
Now, we will call the circuit for  with the same name while the input gates contain the inputs of <ref>. Indeed, we can directly “translate" <ref> to an arithmetic circuit as the values for $I_{k}^{\bm{x}}$ and $I_{k}^{\bm{x}}$ for each $\bm{x}$ and $k$ are predetermined from $N$. Thus, we start by placing the corresponding sorting networks which feeds into the $\texttt{pbs-key-values}\paren{\bm{q}[s \cdots t],\ \bm{k}[x \cdots y], n, m}$ circuit for <ref> in <ref> so that the output values from the calls to  result in the checks as in line <ref> of <ref>. That is, we get outputs $C_k[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))]$ from each call to . We can then use a comparison gate to check if this value equals $2^k$, and if not, we have found a match $C_k[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))]$ for the query $\bm{q}_i$ which results in the output of the associated value $\bm{v}_{C_k[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))]+1}$, exactly as in <ref>. That is, we first define the following variable as the output of the comparison gate:

\begin{equation}
\label{eq: comparison-gate-neq}
\bm{c}_{\mathrm{dec}(i)}^k := \mathbbm{1}\{C_k[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))] \neq 2^k\}.
\end{equation}

Here, as $C_k[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))] \neq 2^k$ implies that $I_{\mathrm{permuted}}^{k\bm{x}}\paren{C_k[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))]}$ equals the index of the matching key $\bm{k}_j$ corresponding to the query $\bm{q}_i$, the $i$th output is then simply given by $\bm{c}_{\mathrm{dec}(i)}^k \cdot I_{\mathrm{permuted}}^{k\bm{x}}\paren{C_k[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}(i))]}$, where the $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:

Sorting networks are circuits with gates and wires where the gates of the circuit are comparators (<ref>) connecting two wires. Each such circuit can perform sorting on a fixed number of values.

We can then show the circuit schematically as in <ref>.

(22.33,134.58) .. controls (22.33,130.67) and (25.5,127.51) .. (29.4,127.51) – (129.93,127.51) .. controls (133.84,127.51) and (137,130.67) .. (137,134.58) – (137,155.78) .. controls (137,159.68) and (133.84,162.84) .. (129.93,162.84) – (29.4,162.84) .. controls (25.5,162.84) and (22.33,159.68) .. (22.33,155.78) – cycle ;

(143.96,136.04) .. controls (143.96,132.62) and (146.74,129.84) .. (150.16,129.84) – (240.8,129.84) .. controls (244.22,129.84) and (247,132.62) .. (247,136.04) – (247,154.64) .. controls (247,158.07) and (244.22,160.84) .. (240.8,160.84) – (150.16,160.84) .. controls (146.74,160.84) and (143.96,158.07) .. (143.96,154.64) – cycle ;

(20,131.18) .. controls (20,126.57) and (23.73,122.84) .. (28.33,122.84) – (276.67,122.84) .. controls (281.27,122.84) and (285,126.57) .. (285,131.18) – (285,156.18) .. controls (285,160.78) and (281.27,164.51) .. (276.67,164.51) – (28.33,164.51) .. controls (23.73,164.51) and (20,160.78) .. (20,156.18) – cycle ;

(410.33,137.98) .. controls (410.33,134.04) and (413.53,130.84) .. (417.47,130.84) – (505.87,130.84) .. controls (509.81,130.84) and (513,134.04) .. (513,137.98) – (513,159.38) .. controls (513,163.32) and (509.81,166.51) .. (505.87,166.51) – (417.47,166.51) .. controls (413.53,166.51) and (410.33,163.32) .. (410.33,159.38) – cycle ;

(520.33,137.78) .. controls (520.33,133.95) and (523.44,130.84) .. (527.27,130.84) – (613.07,130.84) .. controls (616.9,130.84) and (620,133.95) .. (620,137.78) – (620,158.58) .. controls (620,162.41) and (616.9,165.51) .. (613.07,165.51) – (527.27,165.51) .. controls (523.44,165.51) and (520.33,162.41) .. (520.33,158.58) – cycle ;

(401,132.84) .. controls (401,127.87) and (405.03,123.84) .. (410,123.84) – (669.33,123.84) .. controls (674.3,123.84) and (678.33,127.87) .. (678.33,132.84) – (678.33,159.84) .. controls (678.33,164.81) and (674.3,168.84) .. (669.33,168.84) – (410,168.84) .. controls (405.03,168.84) and (401,164.81) .. (401,159.84) – cycle ;

(313.33,54.84) – (111.23,122.21) ;
[shift=(109.33,122.84), rotate = 341.57] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(371.33,55.84) – (577.42,119.26) ;
[shift=(579.33,119.84), rotate = 197.1] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(50,163) – (62.62,227.55) ;
[shift=(63,229.51), rotate = 258.94] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(111,164) – (64.18,227.9) ;
[shift=(63,229.51), rotate = 306.23] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(12.7,235.27) .. controls (12.7,231.46) and (15.79,228.37) .. (19.6,228.37) – (234.14,228.37) .. controls (237.95,228.37) and (241.04,231.46) .. (241.04,235.27) – (241.04,255.95) .. controls (241.04,259.76) and (237.95,262.84) .. (234.14,262.84) – (19.6,262.84) .. controls (15.79,262.84) and (12.7,259.76) .. (12.7,255.95) – cycle ;

(5,231.04) .. controls (5,225.96) and (9.12,221.84) .. (14.2,221.84) – (276.8,221.84) .. controls (281.88,221.84) and (286,225.96) .. (286,231.04) – (286,258.64) .. controls (286,263.72) and (281.88,267.84) .. (276.8,267.84) – (14.2,267.84) .. controls (9.12,267.84) and (5,263.72) .. (5,258.64) – cycle ;

(405,236.16) .. controls (405,231.94) and (408.42,228.51) .. (412.65,228.51) – (623.35,228.51) .. controls (627.57,228.51) and (631,231.94) .. (631,236.16) – (631,259.12) .. controls (631,263.35) and (627.57,266.77) .. (623.35,266.77) – (412.65,266.77) .. controls (408.42,266.77) and (405,263.35) .. (405,259.12) – cycle ;

(402,232.96) .. controls (402,228.07) and (405.97,224.11) .. (410.86,224.11) – (681.14,224.11) .. controls (686.03,224.11) and (690,228.07) .. (690,232.96) – (690,259.53) .. controls (690,264.42) and (686.03,268.39) .. (681.14,268.39) – (410.86,268.39) .. controls (405.97,268.39) and (402,264.42) .. (402,259.53) – cycle ;

(480.27,168.99) – (496.46,226.59) ;
[shift=(497,228.51), rotate = 254.3] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(557.67,168.51) – (498.42,227.1) ;
[shift=(497,228.51), rotate = 315.32] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(521.6,267.74) – (522.64,294.51) ;
[shift=(522.72,296.51), rotate = 267.77] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(143.74,329.38) .. controls (143.74,326.15) and (146.36,323.53) .. (149.59,323.53) – (294.87,323.53) .. controls (298.1,323.53) and (300.71,326.15) .. (300.71,329.38) – (300.71,346.92) .. controls (300.71,350.15) and (298.1,352.76) .. (294.87,352.76) – (149.59,352.76) .. controls (146.36,352.76) and (143.74,350.15) .. (143.74,346.92) – cycle ;

(135,325.5) .. controls (135,320.81) and (138.81,317) .. (143.5,317) – (330.5,317) .. controls (335.19,317) and (339,320.81) .. (339,325.5) – (339,351.01) .. controls (339,355.7) and (335.19,359.51) .. (330.5,359.51) – (143.5,359.51) .. controls (138.81,359.51) and (135,355.7) .. (135,351.01) – cycle ;

(199,265) – (199.94,325.51) ;
[shift=(200,327.51), rotate = 268.24] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(166,353) – (166,419.51) ;
[shift=(166,421.51), rotate = 270] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(199,265) .. controls (13,300.51) and (-28.17,396.68) .. (155.83,431.68) ;
[shift=(155.83,431.68), rotate = 190.77] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(155.83,431.68) .. controls (155.83,426.06) and (160.39,421.51) .. (166,421.51) .. controls (171.61,421.51) and (176.17,426.06) .. (176.17,431.68) .. controls (176.17,437.29) and (171.61,441.84) .. (166,441.84) .. controls (160.39,441.84) and (155.83,437.29) .. (155.83,431.68) – cycle ;

(166,441.84) – (165.37,478.84) ;
[shift=(165.33,480.84), rotate = 270.98] [color=rgb, 255:red, 0; green, 0; blue, 0   ][line width=0.75]    (10.93,-3.29) .. controls (6.95,-1.4) and (3.31,-0.3) .. (0,0) .. controls (3.31,0.3) and (6.95,1.4) .. (10.93,3.29)   ;

(1.33,7) – (699.33,7) – (699.33,517.51) – (1.33,517.51) – cycle ;

(9,14) – (215.33,14) – (215.33,30.51) – (9,30.51) – cycle ;

(218,40.4) node [anchor=north west][inner sep=0.75pt]  [font = ]  $( \bm{k}_{1} ,\ \bm{v}_{1} ,\ \bm{q}_{1}) ,\ \dotsc \dotsc \dotsc \dotsc ,\ ( \bm{k}_{\frac{N}{3}-1} ,\ \bm{v}_{\frac{N}{3}-1} ,\ \bm{q}_{\frac{N}{3}-1})$;

(260,158.4) node [anchor=north west][inner sep=0.75pt]    $\dotsc $;

(250,126.4) node [anchor=north west][inner sep=0.75pt]  [font=]  $k=0$;

(642.63,163.4) node [anchor=north west][inner sep=0.75pt]    $\dotsc $;

(623.63,128.4) node [anchor=north west][inner sep=0.75pt]  [font=]  $k=\log \frac{N}{3}$;
(643.63,138.4) node [anchor=north west][inner sep=0.75pt]  [font=]  $-1$;

(300,146.4) node [anchor=north west][inner sep=0.75pt]    $\dotsc \dotsc \dotsc \dotsc $;

(453.72,134.88) node [anchor=north west][inner sep=0.75pt]  [font=] [align=left] sort;

(419,145.09) node [anchor=north west][inner sep=0.75pt]  [font=]  $(\{\bm{q}_{\mathrm{dec}( j)}\}_{j\in J_{k}^{x}})$;

(556.35,132.88) node [anchor=north west][inner sep=0.75pt]  [font=] [align=left] sort;

(526.33,148.64) node [anchor=north west][inner sep=0.75pt]  [font=]  $(\{\bm{k}_{\mathrm{dec}( i)}\}_{i\in I_{k}^{x}})$;

(250.37,226.4) node [anchor=north west][inner sep=0.75pt]  [font=]  $k=0$;

(185.72,135.88) node [anchor=north west][inner sep=0.75pt]  [font=] [align=left] sort;

(150,145.09) node [anchor=north west][inner sep=0.75pt]  [font=]  $(\{\bm{k}_{\mathrm{dec}( i)}\}_{i\in I_{0}^{x}})$;

(210,288.4) node [anchor=north west][inner sep=0.75pt]    $\dotsc $;
(178,273.4) node [anchor=north west][inner sep=0.75pt]  [font=]  $C_{0}$;

(143.46,334.75) node [anchor=north west][inner sep=0.75pt]  [font=]  $C_{0}[J_{\mathrm{permuted}}^{k\bm{x}}(\mathrm{dec}( i))] \neq 2^{0}$;

(170,363.4) node [anchor=north west][inner sep=0.75pt]  [font=]  $\bm{c}^{0}_{\mathrm{dec}( i)}$;

(303,325.51) node [anchor=north west][inner sep=0.75pt]  [font=]  $i \in I^{\bm{x}}_0$;

(350,335.51) node [anchor=north west][inner sep=0.75pt]  [font=]  $k \in \{0, \ldots, \log{\frac{N}{3}}-1\}$;

(350,350.51) node [anchor=north west][inner sep=0.75pt]    $\dotsc \dotsc \dotsc \dotsc $;

(157,425.74) node [anchor=north west][inner sep=0.75pt]    $\times $;

(181,479.4) node [anchor=north west][inner sep=0.75pt]    $\dotsc \dotsc \dotsc \dotsc $;

(60.79,389.64) node [anchor=north west][inner sep=0.75pt]    $\dotsc $;

(37,360.4) node [anchor=north west][inner sep=0.75pt]  [font=]  $I_{\mathrm{permuted}}^{k\bm{x}}\paren{\bm{c}^{0}_{\mathrm{dec}( i)}}$;

(161,486.74) node [anchor=north west][inner sep=0.75pt]  [font=]  $\bm{v}_{j}$;

(185,419.4) node [anchor=north west][inner sep=0.75pt]    $\dotsc \dotsc \dotsc \dotsc $;

(125,16.4) node [anchor=north west][inner sep=0.75pt]  [font=]  $(\hyenaInput[ 0\cdots N-1])$;

(11,17) node [anchor=north west][inner sep=0.75pt]  [font=] [align=left] Parallel-General-AR;

(67.82,135.53) node [anchor=north west][inner sep=0.75pt]  [font=] [align=left] sort;

(33.79,144.6) node [anchor=north west][inner sep=0.75pt]  [font=]  $(\{\bm{q}_{\mathrm{dec}( j)}\}_{j\in J_{0}^{x}})$;

(304,262.4) node [anchor=north west][inner sep=0.75pt]    $\dotsc \dotsc \dotsc \dotsc $;

(659.86,260.66) node [anchor=north west][inner sep=0.75pt]    $\dotsc $;

(633.51,227.46) node [anchor=north west][inner sep=0.75pt]  [font=]  $k=\log{\frac{N}{3}}$;
(655.51,237.46) node [anchor=north west][inner sep=0.75pt]  [font=]  $-1$;

(481.85,232.75) node [anchor=north west][inner sep=0.75pt]  [font=] [align=left] pbs-key-values;

(405.97,247.65) node [anchor=north west][inner sep=0.75pt]  [font=]  $(\mathrm{\tt sort}\{\bm{q}_{\mathrm{dec}( j)}\}_{j\in J_{k}^{x}} ,\mathrm{\tt sort}\{\bm{k}_{\mathrm{dec}( i)}\}_{i\in I_{k}^{x}} \ )$;

(82.61,233.05) node [anchor=north west][inner sep=0.75pt]  [font=] [align=left] pbs-key-values;

(20.16,246.73) node [anchor=north west][inner sep=0.75pt]  [font=]  $(\mathrm{\tt sort}\{\bm{q}_{\mathrm{dec}( j)}\}_{j\in J_{k}^{x}} ,\mathrm{\tt sort}\{\bm{k}_{\mathrm{dec}( i)}\}_{i\in I_{k}^{x}} \ )$;

$(\hyenaInput[0 \cdots N-1])$ as a circuit that includes sorting networks and the circuit for  as subroutines.

We now dilineate the complexity of the circuit, starting with the complexity of the sorting networks.

Let $A$ be an array with $d$-bit entries of size $n$. Then, one can implement a sorting network to sort the array $A$ with size $\calO(d \cdot n\log{n})$ and depth $\calO(\log{d}\log{n})$.

There exists an $(\inputLength\cdot d, \calO(Nd\cdot \log^2{N}), \calO(\log{d}\log^2{N}), \calO(Nd\log{N}))$-arithmetic circuit that solves the multiple-query associative recall problem.

We note here that for each $k$, there are $\frac{\frac{N}{3}}{2^{k+1}}$ parallel problems of size $2^k$ for both the sorting networks and the ${\tt pbs-key-values}$ circuit. Using <ref>, the cumulative size of these sorting networks is $\calO(d \cdot N\log^2{N})$ (see (<ref>)) with overall depth $\calO(\log{d}\log{N})$.

Similarly, the next layer again runs $\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 $\calO(2^kd(\log{2^k} + \log{2^k})) = O(d\cdot 2^k\log{2^k})$, depth $\calO(\log^2{2^k}\log{d})$ and width $\calO(2^k)$ (<ref>). Again, the cumulative size of this layer is given by $\calO(Nd\cdot\log^2{N})$ (see (<ref>)).
Since we run each of these circuits in parallel, the depth of this layer is again $\calO(\log{d}\log^2(N))$ while the width is $\calO(N \cdot \log{N})$.

Finally, we perform $\frac{N}{3}\log{\frac{N}{3}}$ comparisons at the end of $d$-bit strings in parallel which results in size $\calO(N\log{N}\cdot d)$, depth $\calO(\log{d})$ and width $\calO(N\log{N}\cdot d)$ (<ref>). Therefore, the resulting arithmetic circuit has size $\calO(d \cdot N\log^2{N} + Nd \cdot \log^2{N} + N\log{N}\cdot d) = O(Nd\log^2{N})$, depth $\calO(\log{d}\log^2{N})$ and width $\calO(Nd\log{N})$.

§.§.§ The Resulting  Model

As we have an arithemtic circuit for solving the multiple-query associative recall problem, we can now invoke <ref> to claim that there is a corresponding  model that solves the multiple-query associative recall problem with $\tilde{\calO}(N\log{c})$ parameters and $\tilde{\calO(1)}$ layers.

There exists a $\coyoteTuple{\inputLength d}{\tilde{\calO}(1)}{\tilde{\calO}(1)}{\tilde{\calO}(N d)}{\tilde{\calO}(1)}$ solves the multiple-query associative recall problem.

Directly applying <ref> yields a  model with the number of layers $\calO(\log{d} \cdot \log^2{N} \cdot \log{Nd\log{N}}) = \calO(1)$ layers while the claim on the input and inner dimensions follow trivially.

§.§ Data-Dependent Convolutions

§.§.§ Introduction

In this section, we are again concerned with solving the multiple-query associative recall problem (<ref>). However, in contrast to <ref>, which yields a circuit that is unchanged and works for all inputs, we instead take the viewpoint of adapting the sequence mixing weights (<ref>) 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.

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
\bm{u}[0 \cdots N-1] \triangleq \left\{\paren{\bm{k}_0, \bm{v}_0, \bm{q}_0}, \ldots, \paren{\bm{k}_{\frac{N}{3}-1}, \bm{v}_{\frac{N}{3}-1}, \bm{q}_{\frac{N}{3}-1}}\right\}
with each $\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$ (i.e. we assume $d=c$).[Our arguments do need $c=d$. However we do not need $d=N$ but we made this simplification for ease of presentation.] Our goal is again to check, for each $1 \le i \le \frac{N}{3}-1$, whether there exists $0 \le j < i$ such that $\bm{q}_i \equiv \bm{k}_j$, and if so, output $\bm{v}_{j}$.

Here, we define the interaction distance between the $i$th query $\bm{q}_i$ and the matching key $\bm{k}_j$ as $i-j$. We then also assume that number of distinct interaction distances is bounded by $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.

We introduce autocorrelation convolutions. Let $\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 $\bm{u} \in \R^{n}$ is the cross correlation of $\bm{u}$ with itself. Moreover, in terms of polynomials, we have
$\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 $\bm{u} \in \R^n$, given by $\tbf{w} \in \R^n$ as follows:

\begin{equation*}
\label{eq: auto-cor}
\tbf{w} = \coeff\paren{\bm{u}(X) \cdot \tilde{\bm{u}}(X) \mod{X^{n}-1}}.
\end{equation*}

§.§.§  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$ 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$ shifts.
That is, we define a function ${\tt Top}$ such that ${\tt Top}\paren{u \star u, t}$ returns a list of the top $t$ shifts $\{s_{\ell}\}_{\ell \in [t]}$. We then use these top $t$ distances $\{s_{\ell}\}_{\ell \in [t]}$ to define the following two kernels:

\begin{equation}
\label{eq: data-dependent-kernels}
\begin{aligned}
\tbf{h}^{k}(X) &\equiv \sum_{\ell \in [t]} X^{s_{\ell} + (\ell -1) \cdot \frac{N}{3}},\\
\tbf{h}^{v}(X) &\equiv \sum_{\ell \in [t]} X^{s_{\ell} - 1 + (\ell -1)\cdot \frac{N}{3}}.\\
\end{aligned}
\end{equation}

Here, we note that we only only have two kernels as we will assume $N'=t\frac{N}{3}$ and the shift will be done in "parallel." Obviously, one can instead define $t$ distinct shift kernels but then there is a cost of $\calO(t)$ in the number of layers.

We define the following projections $\tbf{K}, \tbf{Q}, \tbf{V} \in \{0,1\}^{N\times d}$ of the input that we shall use below using (<ref>).

\begin{equation}
\label{eq: kqv-projections}
\begin{aligned}
\tbf{K}[i,:] &:=
\begin{cases}
\bm{u}[i,:] &\text{if }i \in \calK,\\
\tbf{0}^d &\text{otherwise}
\end{cases},\\
\tbf{Q}[i,:] &:= \begin{cases}
\bm{u}[i,:] &\text{if }i \in \calQ,\\
\tbf{0}^d &\text{otherwise}
\end{cases},\\
\tbf{V}[i,:] &:= \begin{cases}
\bm{u}[i,:] &\text{if }i \in \calV,\\
\tbf{0}^d &\text{otherwise}
\end{cases},\\
\end{aligned}
\end{equation}

Finally, we present the  model that solves the multiple-query associative recall problem with input-dependent kernels using $\calO(1)$ layer and $\calO(t \cdot Nd)$-many parameters.

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

We note that we have the input dimension $d = c$. Now, for any input sequence $\bm{u} \in \{0,1\}^{N\times d}$, we get the input-dependent kernels as in (<ref>) using autocorrelation of the input. We will now outline the following computations for the  layers:

\begin{align}
\bm{y}  &= {\tt Linear}_{\tbf{Q}}(\bm{u}) \odot \paren{\tbf{h}^{K} * {\tt Linear}_{\tbf{K}}(\bm{u})}\nonumber \\
&= \tbf{Q} \odot \paren{\tbf{h}^{K} * \tbf{K}} \label{eq: y-data-ind}\\
\bm{z}  &= {\tt Linear}_{\tbf{E}}(\bm{y}) \odot \paren{\tbf{h}^{V} * {\tt Linear}_{\tbf{V}}(\bm{u})}\nonumber \\
&= \tbf{E} \odot \paren{\tbf{h}^{V} * \tbf{V}},  \label{eq: z-data-ind}
\end{align}

where we have the linear projections ${\tt Linear}_{\tbf{Q}}(\bm{u}) = \tbf{Q}$, ${\tt Linear}_{\tbf{K}}(\bm{u}) = \tbf{K}$, ${\tt Linear}_{\tbf{V}}(\bm{u}) = \tbf{V}$ and ${\tt Linear}_{\tbf{E}}(\bm{y}) = \tbf{E}$ defined as

\begin{equation}
\label{eq: E-def}
\tbf{E}[i,:] := {\tt Linear}_{\tbf{E}}(\bm{y})[i,:] =
\begin{cases}
\tbf{1}^d &\text{if }\exists\ j \in [d]\text{ such that }\bm{y}[i,j] = 1\\
\tbf{0}^{d} &\text{otherwise}
\end{cases}.
\end{equation}

Here, we will first present the argument for the special case when we have $t = 1$ as that will help us elucidate the general case. To this end, as the kernels from (<ref>) for $t=1$ are given by

\begin{equation}
\label{eq: data-dependent-kernels-t1}
\begin{aligned}
\tbf{h}^{k}(X) &\equiv X^{s_{1}};\\
\tbf{h}^{v}(X) &\equiv X^{s_{1} - 1},
\end{aligned}
\end{equation}

we observe that convolving with these kernels $\tbf{h} \ast \bm{y}$ is equivalent to operating with the following primitives (<ref>):

\begin{equation}
\label{eq: data-dependent-kernels-prim1}
\begin{aligned}
&{\tt shift\_down}(\bm{y}, s_{1});\\
&{\tt shift\_down}(\bm{y}, s_{1}-1).
\end{aligned}
\end{equation}

We note that we shift down instead of shifting up as the index of the top-left entry is $(0,0)$. We can then write down the computations performed in (<ref>) and (<ref>) as follows:

\begin{align}
\bm{y}  &= \tbf{Q} \odot {\tt shift\_down}(\tbf{K}, s_{1}) \label{eq: y-data-ind-t1}\\
\bm{z}  &= \tbf{E} \odot {\tt shift\_down}(\tbf{V}, s_{1}-1),  \label{eq: z-data-ind-t1}
\end{align}

We begin by examining $\bm{y}$ below:

\begin{align}
\bm{y}[i,:]
&= \paren{\tbf{Q} \odot {\tt shift\_down}(\tbf{K}, s_1)}[i,:]\nonumber \\
&= \tbf{Q}[i,:] \odot {\tt shift\_down}(\tbf{K}, s_1)[i,:] \label{eq: gated-row-independency}  \\
&= \paren{
\begin{cases}
\bm{u}[i,:] &\text{if }i \in \calQ\\
\tbf{0}^d &\text{otherwise}
\end{cases}
} \odot \paren{
\begin{cases}
\bm{u}[i-s_1,:] &\text{if }i - s_1 \in \calK,\\
\tbf{0}^d &\text{otherwise}
\end{cases}
} \label{eq: key-indices-shift}\\
&= \begin{cases}
\bm{u}[i,:] \odot \bm{u}[i-s_1,:] &\text{if }i \in \calQ \text{ and }i-s_1 \in \calK,\\
\tbf{0}^d &\text{otherwise}
\end{cases} \nonumber
\end{align}

Here, we use the fact that the Hadamard product is row-independent in (<ref>), and the definitions of the projections from (<ref>) in (<ref>). Examining the $j$th entry, we get

\begin{align*}
\bm{u}[i,j] \odot \bm{u}[i-s_1,j] &=
\begin{cases}
1 & \text{if }i \in \calQ, i-s_1 \in \calK\text{ and }\bm{q}_{i} \equiv \bm{k}_{i-s_1} \equiv \tbf{e}_j\\
0 & \text{otherwise.}
\end{cases}
\end{align*}

That is, we can express

\begin{align}
\bm{y}[i,:]
&= \begin{cases}
\tbf{e}_j &\text{if }i \in \calQ, i-s_1 \in \calK\text{ and }\bm{q}_{i} \equiv \bm{k}_{i-s_1} \equiv \tbf{e}_j\\
\tbf{0}^d &\text{otherwise}
\end{cases} \label{eq: y-t1}.
\end{align}

Consequently, as per the definition in (<ref>), we get

\begin{align}
\tbf{E}[i,:] =
\begin{cases}
\tbf{1}^d &\text{if } i \in \calQ, i-s_1 \in \calK\text{ and }\bm{q}_{i} \equiv \bm{k}_{i-s_1} \\
\tbf{0}^{d} &\text{otherwise}
\end{cases} \label{eq: E-t1}
\end{align}

We can now finally specify the output $\bm{z}$ from (<ref>) as follows:

\begin{align}
\bm{z}[i,:]
&= \paren{\tbf{E} \odot {\tt shift\_down}(\tbf{V}, s_1-1)}[i,:]\nonumber \\
&= \tbf{E}[i,:] \odot {\tt shift\_down}(\tbf{V}, s_1-1)[i,:] \label{eq: gated-row-independency-2}  \\
&= \paren{
\begin{cases}
\tbf{1}^d &\text{if } i \in \calQ, i-s_1 \in \calK\text{ and }\bm{q}_{i} \equiv \bm{k}_{i-s_1} \\
\tbf{0}^{d} &\text{otherwise}
\end{cases}
} \odot \paren{
\begin{cases}
\bm{u}[i-s_1+1,:] &\text{if }i - s_1 + 1 \in \calV,\\
\tbf{0}^d &\text{otherwise}
\end{cases}
} \label{eq: value-indices-shift}\\
&= \begin{cases}
\bm{u}[i-s_1+1,:] &\text{if }i \in \calQ, i-s_1 \in \calK, i - s_1 + 1 \in \calV \text{ and }\bm{q}_{i} \equiv \bm{k}_{i-s_1}\\
\tbf{0}^d &\text{otherwise}
\end{cases} \nonumber\\
&= \begin{cases}
\bm{v}_{i-s_1} &\text{if }\bm{q}_{i} \equiv \bm{k}_{i-s_1}\\
\tbf{0}^d &\text{otherwise}
\end{cases} \label{eq: output-t1}
\end{align}

Again, we use the fact that the Hadamard product is row-independent in (<ref>), and the definitions of the projections from (<ref>) in (<ref>). Overall, we have solved associative recall for all queries that have interaction distance exactly equal to $s_1$.

In order to generalize this to arbitrary $t \le \frac{N}{3}$, we first increase the internal dimension so that the input to the kernels $\bm{u}' \in \R^{(N\cdot t)\times d}$ in (<ref>) and the projections $\tbf{K',Q',V'} \in \R^{(N\cdot t)\times d}$ are given by

\[
\bm{u}' \equiv
\begin{pmatrix}
\tbf{0}^d\\
\hline \\
\vdots\\
\hline\\
\tbf{0}^d\\
\hline\\
\bm{u}
\end{pmatrix},
\bm{K}' \equiv
\begin{pmatrix}
\tbf{0}^d\\
\hline \\
\vdots\\
\hline\\
\tbf{0}^d\\
\hline\\
\tbf{K}
\end{pmatrix},
\tbf{Q}' \equiv
\begin{pmatrix}
\tbf{Q}\\
\hline \\
\vdots\\
\hline\\
\tbf{Q}\\
\hline\\
\tbf{Q}
\end{pmatrix}, \tbf{V}' \equiv
\begin{pmatrix}
\tbf{0}^d\\
\hline \\
\vdots\\
\hline\\
\tbf{0}^d\\
\hline\\
\tbf{V}
\end{pmatrix},
\]

We then observe that for $\tbf{h}^k_{\ell}(X) := X^{s_{\ell}}$ and $\tbf{h}^v_{\ell}(X) := X^{s_{\ell} - 1}$, we have

\[
\begin{aligned}
\tbf{h}^{k}(X) &\equiv \sum_{\ell \in [t]} \tbf{h}^k_{\ell}(X)\cdot X^{(\ell -1) \cdot \frac{N}{3}},\\
\tbf{h}^{v}(X) &\equiv \sum_{\ell \in [t]} \tbf{h}^k_{\ell}(X)\cdot X^{(\ell -1)\cdot \frac{N}{3}}.
\end{aligned}
\]

In analogy with (<ref>), we can then equivalently write

\begin{align}
\paren{\tbf{h}^{K} \ast \tbf{K}}
\begin{pmatrix}
\tbf{h}^{K}_{t}\\
\hline \\
\vdots\\
\hline\\
\tbf{h}^{K}_{2}\\
\hline\\
\tbf{h}^{K}_{1}
\end{pmatrix}
\ast
\begin{pmatrix}
\tbf{0}^d\\
\hline \\
\vdots\\
\hline\\
\tbf{0}^d\\
\hline\\
\tbf{K}
\end{pmatrix} \nonumber \\
&\equiv \begin{pmatrix}
\tbf{h}^{K}_{t} \ast \tbf{K}\\
\hline \\
\vdots\\
\hline\\
\tbf{h}^{K}_{2} \ast \tbf{K}\\
\hline\\
\tbf{h}^{K}_{1} \ast \tbf{K}
\end{pmatrix}
\nonumber \\
&\equiv \begin{pmatrix}
{\tt shift\_down}(\tbf{K}, s_{t}) \\
\hline \\
\vdots\\
\hline\\
{\tt shift\_down}(\tbf{K}, s_{2})\\
\hline\\
{\tt shift\_down}(\tbf{K}, s_{1})
\end{pmatrix}
\nonumber.
\end{align}

Similarly, we also have

\begin{align}
\paren{\tbf{h}^{V} \ast \tbf{V}'}
&\equiv \begin{pmatrix}
{\tt shift\_down}(\tbf{V}, s_{t}-1) \\
\hline \\
\vdots\\
\hline\\
{\tt shift\_down}(\tbf{V}, s_{2}-1)\\
\hline\\
{\tt shift\_down}(\tbf{V}, s_{1}-1)
\end{pmatrix}
\nonumber.
\end{align}

That is, the argument for $t=1$ now applies to each of the $t$ shifts as we now have (cf. (<ref>))

\begin{align*}
\bm{y}'
&\equiv \tbf{Q}' \odot \paren{\tbf{h}^{V} \ast \tbf{V}}  \\
&\equiv \begin{pmatrix}
\tbf{Q}\\
\hline \\
\vdots\\
\hline\\
\tbf{Q}\\
\hline\\
\tbf{Q}
\end{pmatrix} \odot \begin{pmatrix}
{\tt shift\_down}(\tbf{V}, s_{t}-1) \\
\hline \\
\vdots\\
\hline\\
{\tt shift\_down}(\tbf{V}, s_{2}-1)\\
\hline\\
{\tt shift\_down}(\tbf{V}, s_{1}-1)
\end{pmatrix}  \\
&\equiv \begin{pmatrix}
\tbf{Q} \odot {\tt shift\_down}(\tbf{V}, s_{t}-1) \\
\hline \\
\vdots\\
\hline\\
\tbf{Q} \odot {\tt shift\_down}(\tbf{V}, s_{2}-1)\\
\hline\\
\tbf{Q} \odot {\tt shift\_down}(\tbf{V}, s_{1}-1)
\end{pmatrix} \\
&\equiv \begin{pmatrix}
\bm{y}_t \\
\hline \\
\vdots\\
\hline\\
\bm{y}_2 \\
\hline\\
\bm{y}_1 \\
\end{pmatrix},
\end{align*}

where, for each $\ell \in [t]$, we have (cf. (<ref>))

\[
\bm{y}_{\ell}[i,:] \equiv \begin{cases}
\tbf{e}_j &\text{if }i \in \calQ, i-s_{\ell} \in \calK\text{ and }\bm{q}_{i} \equiv \bm{k}_{i-s_{\ell}} \equiv \tbf{e}_j,\\
\tbf{0}^d &\text{otherwise}
\end{cases}.
\]

We then analogously get $\tbf{E}'$ as follows:

\[
\tbf{E}' \equiv {\tt Linear}_{\tbf{E}}(\bm{y}') \equiv
\begin{pmatrix}
{\tt Linear}_{\tbf{E}}(\bm{y}_t)\\
\hline \\
\vdots\\
\hline\\
{\tt Linear}_{\tbf{E}}(\bm{y}_2)\\
\hline\\
{\tt Linear}_{\tbf{E}}(\bm{y}_1)
\end{pmatrix}  \equiv  \begin{pmatrix}
\tbf{E}_t\\
\hline \\
\vdots\\
\hline\\
\tbf{E}_2\\
\hline\\
\tbf{E}_1
\end{pmatrix},
\]

where, for each $\ell \in [t]$, we have (cf. (<ref>))

\[
\tbf{E}_{\ell}[i,:] =
\begin{cases}
\tbf{1}^d &\text{if } i \in \calQ, i-s_{\ell} \in \calK\text{ and }\bm{q}_{i} \equiv \bm{k}_{i-s_{\ell}} \\
\tbf{0}^{d} &\text{otherwise}
\end{cases}
\]

The output in the general case is then given by

\begin{align}
\bm{z}'
&\equiv \tbf{E}' \odot \paren{\tbf{h}^{V} * \tbf{V}'} \nonumber\\
&\equiv \begin{pmatrix}
\tbf{E}_t\\
\hline \\
\vdots\\
\hline\\
\tbf{E}_2\\
\hline\\
\tbf{E}_1
\end{pmatrix} \odot \begin{pmatrix}
{\tt shift\_down}(\tbf{V}, s_{t}-1) \\
\hline \\
\vdots\\
\hline\\
{\tt shift\_down}(\tbf{V}, s_{2}-1)\\
\hline\\
{\tt shift\_down}(\tbf{V}, s_{1}-1)
\end{pmatrix}  \nonumber \\
&\equiv \begin{pmatrix}
\tbf{E}_t \odot {\tt shift\_down}(\tbf{V}, s_{t}-1) \\
\hline \\
\vdots\\
\hline\\
\tbf{E}_2 \odot {\tt shift\_down}(\tbf{V}, s_{2}-1)\\
\hline\\
\tbf{E}_1 \odot {\tt shift\_down}(\tbf{V}, s_{1}-1)
\end{pmatrix}  \nonumber \\
&\equiv \begin{pmatrix}
\bm{z}_t \\
\hline \\
\vdots\\
\hline\\
\bm{z}_2 \\
\hline\\
\bm{z}_1 \\
\end{pmatrix}  \nonumber,
\end{align}

where, for each $\ell \in [t]$, we have (cf. (<ref>))

\[
\bm{z}_{\ell}[i,:] \equiv
\begin{cases}
\bm{v}_{i-s_{\ell}} &\text{if }\bm{q}_{i} \equiv \bm{k}_{i-s_{\ell}}\\
\tbf{0}^d &\text{otherwise}
\end{cases}
\]

Finally, we define the last output layer to compute
\(
\bm{z}_{out}\equiv {\tt Linear}_{\mathrm{sum}}(\bm{z}') \equiv \sum_{\ell \in [t]} \bm{z}_{\ell}
\)
so that we have

\[
\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]\\
\tbf{0}^d &\text{otherwise}
\end{cases}
\]

To recall, we retrieved the top $t$ interaction distances of the input $\bm{u}$ using auto-correlation and defined the corresponding convolution kernels ((<ref>)). We then shifted djown the keys $\tbf{K}$ using the first kernel and gated with the corresponding queries $\tbf{Q}$ so that we got a match exactly when there exists a key that is at $s_{\ell}$ interaction distance from the corresponding query. After “smearing” this match to get $\tbf{E}$, we used it as a mask to retrieve the value in the next layer. Overall, since we have $t$ atomic kernels that perform $t$ shifts with each of these kernels using $\calO(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 $\ell \in [t]$ using $\calO(1)$ layers and $\calO(t \cdot Nc)$ parameters as we have $d = c$.

Attention Training Settings

73M     125M     360M

Optimizer     3cAdam

Optimizer momentum     3c$\beta_1, \beta_2=0.9, 0.95$

Optimizer eps     3c$1e-8$

Precision      3cBFloat16

Learning rate decay     3cCosine

Learning rate (min, base)     3c8e-5, 8e-4

Global batch size     3c256

Training iterations     3c20000

Warmup Duration (Linear)     3c0.01

Weight decay     3c0.1

Num Layers     6     12     24

Hidden Size     704      768     1024

FFN Width     3c4

Position Embeddings     3cRotary

Weight Tying      3cTrue

Number of Heads ($H$)     8     12     16

Attention FLOPs Computation

Input Layer     $B \times V \times N \times D$

Sequence Mixer $\mathbf{Q}, \mathbf{K}, \mathbf{V}$ Projections     $B \times N \times D \times D \times 3$

Sequence Mixer Attention     $B \times H \times H \times D + H \times N \times N + B \times N \times N \times D$

Sequence Mixer Output Projection     $B \times N \times D \times D$

Channel Mixer (FFN Width $4$)      $B \times D \times D \times 8 \times \frac{2}{3} \times N$

Language Modeling Head     $B \times V \times N \times D$

Hyena [Poli et al., 2023] Training Settings

72M     158M     358M

Optimizer     3cAdam

Optimizer momentum     3c$\beta_1, \beta_2=0.9, 0.98$

Precision      3cBFloat16

Learning rate decay     3cCosine

Learning rate (min, base)     3c(1e-5, 8e-4)

Global batch size     3c256

Training iterations     3c20000

Warmup Duration (Linear)     3c0.01

Weight decay     3c0.1

Num Layers     8     18     24

Hidden Size     768      864     1024

FFN Width     3c2

Position Embeddings     3cNone

Weight Tying      3cTrue

Short Conv. Filter Size     3c3

Exp. Mod. Decay (Fast, Slow)     3c0.3, 1.2

Filter Sine Freq. (w)     3c14

Filter Order     3c64

Filter Inner MLP     3c2

Filter Weight Decay     3c0

Hyena FLOPs Computation

Input Layer     $B \times V \times N \times D$

Sequence Mixer Input Projection     $B \times N \times D \times D \times 3 + B \times N \times 9 \times D$

Sequence Mixer Long Convolution     $10 \times N \times \log(N) \times D \times B$

Sequence Mixer Short Convolution     $3 \times B \times N \times D$

Sequence Mixer Implicit MLP (Order $64$)     $D \times 64$

Sequence Mixer Output Projection     $B \times N \times D \times D$

Channel Mixer (FFN Width $2$)      $B \times D \times D \times 2 \times 2 \times N$

Language Modeling Head     $B \times V \times N \times D$

H3 [Fu et al., 2023] Training Settings

72M     168M     357M

Optimizer     3cAdam

Optimizer momentum     3c$\beta_1, \beta_2=0.9, 0.95$

Precision      3cBFloat16

Learning rate decay     3cCosine

Learning rate (min, base)     3c(8e-5, 8e-4)

Global batch size     3c256

Training iterations     3c20000

Warmup Duration (Linear)     3c0.01

Weight decay     3c0.1

Weight Tying      3cTrue

Num Layers     8     19     33

Hidden Size     720      864     1024

FFN Width     3c2

Position Embeddings     3cNone

State Space Model State Size     3c64

RWKV [Peng et al., 2023] Training Settings

72M     169M     351M

Optimizer     3cAdam

Optimizer momentum     3c$\beta_1, \beta_2=0.9, 0.999$

Optimizer eps     3c$1e-8$

Precision      3cBFloat16

Learning rate decay     3cCosine

Learning rate (min, base)     3c1e-5, 8e-4

Global batch size     3c256

Training iterations     3c20000

Warmup Duration (Linear)     3c0.01

Weight decay     3c0.1

Weight Tying      3cFalse

Num Layers     6     12     20

Hidden Size     624     768     984

Position Embeddings     3cNone

Initialization     3cFrom Reference Impl.

RetNet [Sun et al., 2023] Training Settings

Optimizer     1cAdam

Optimizer momentum     1c$\beta_1, \beta_2=0.9, 0.98$

Optimizer eps     1c1.0e-8

Precision      1cBFloat16

Learning rate decay     1cLinear

Learning rate (min, base)     1c(1e-5, 8e-4)

Global batch size     1c256

Training iterations     1c20000

Warmup Duration (Linear)     1c0.01

Weight decay     1c0.05

Normalization     1cLayernorm

Weight Tying      1cTrue

Num Layers     12

Hidden Size     768

Value Hidden Size     1280

Window Size     128

Position Embeddings     1cxPos (relative position embeddings)

Initialization     DeepNet

Simple Long Convolution [Fu et al., 2023] Training Settings

76M     128M     360M

Optimizer     3cAdam

Optimizer momentum     3c$\beta_1, \beta_2=0.9, 0.95$

Precision      3cBFloat16

Learning rate decay     3cLinear

Learning rate (min, base)     3c8e-5, 8e-4

Global batch size     3c256

Training iterations     3c20000

Warmup Duration (Linear)     3c0.01

Weight decay     3c0.1

Num Layers     6     12     24

Hidden Size     704      864     1024

FFN Width     3c4

Position Embeddings     3c-

Weight Tying      3cTrue

Channels     3c1

Lam     3c0.001

Kernel Dropout     3c0.1

Kernel LR     3c$5e-5$

Activation     3cGeLU

Exponential Modulation     3cTrue

Simple Long Convolution FLOPs Computation

Input Layer     $B \times V \times N \times D$

Sequence Mixer Long Convolution     $10 \times N \times \log(N) \times D \times B$

Sequence Mixer Output Projection     $B \times N \times D \times D$

Channel Mixer (FFN Width $4$)      $B \times D \times D \times 8 \times \frac{2}{3} \times N$

Language Modeling Head     $B \times V \times N \times D$

Training Settings

168M     354M

Optimizer     2cAdam

Optimizer momentum     2c$\beta_1, \beta_2=0.9, 0.95$

Precision      2cBFloat16

Learning rate decay     2cCosine

Learning rate (min, base)     2c8e-5, 8e-4

Global batch size     2c256

Training iterations     2c20000

Warmup Duration (Linear)     2c0.01

Weight decay     2c0.1

Num Layers     30     48

Hidden Size       852     1080

FFN Width     2c2

Position Embeddings     2c-

Weight Tying      2cTrue

Short Conv. Filter Size     2c3

Exp. Mod. Decay (Fast, Slow)     2c0.3, 1.2

Filter Sine Freq. (w)     2c14

Filter Order     2c64

Filter Inner MLP     2c2

Filter Weight Decay     2c0

FLOPs Computation

Input Layer     $B \times V \times N \times D$

Sequence Mixer Long Convolution     $10 \times N \times \log(N) \times 0.5(D) \times B$

Sequence Mixer Short Convolution     $B \times N \times 0.5(D)$

Sequence Mixer Implicit MLP (Order $64$)     $0.5(D) \times 64$

Sequence Mixer Linear Projection     $B \times N \times D \times D$

Channel Mixer (FFN Width $2$)      $B \times D \times D \times 2 \times 2 \times N$

Language Modeling Head     $B \times V \times N \times D$