Input  ⬇ Who scored the first touchdown of the game? Hoping to rebound from
their loss to the Patriots, the Raiders stayed at home for a Week 16 duel with
the Houston Texans. Oakland would get the early lead in the first quarter as
quarterback JaMarcus Russell completed a 20-yard touchdown pass to rookie wide
receiver Chaz Schilens. The Texans would respond with fullback Vonta Leach
getting a 1-yard touchdown run, yet the Raiders would answer with kicker
Sebastian Janikowski getting a 33-yard and a 30-yard field goal. Houston would
tie the game in the second quarter with kicker Kris Brown getting a 53-yard
and a 24-yard field goal. Oakland would take the lead in the third quarter
with wide receiver Johnnie Lee Higgins catching a 29-yard touchdown pass from
Russell, followed up by an 80-yard punt return for a touchdown. The Texans
tried to rally in the fourth quarter as Brown nailed a 40-yard field goal, yet
the Raiders’ defense would shut down any possible attempt. Hoping to rebound
from their loss to the Patriots, the Raiders stayed at home for a Week 16 duel
with the Houston Texans. Oakland would get the early lead in the first quarter
as quarterback JaMarcus Russell completed a 20-yard touchdown pass to rookie
wide receiver Chaz Schilens. The Texans would respond with fullback Vonta
Leach getting a 1-yard touchdown run, yet the Raiders would answer with kicker
Sebastian Janikowski getting a 33-yard and a 30-yard field goal. Houston would
tie the game in the second quarter with kicker Kris Brown getting a 53-yard
and a 24-yard field goal. Oakland would take the lead in the third quarter
with wide receiver Johnnie Lee Higgins catching a 29-yard touchdown pass from
Russell, followed up by an 80-yard punt return for a touchdown. The Texans
tried to rally in the fourth quarter as Brown nailed a 40-yard field goal, yet
the Raiders’ defense would shut down any possible attempt. The first touchdown
of the game was scored by ⬇ Chaz Schilens

##  Appendix G Theoretical results

We begin by setting notation.

##### Notation.

We will be denoting the all $1$ row vector of size $k$, given by
$\begin{bmatrix}1&1&\ldots&1&1\end{bmatrix}$, and the all $0$ row vector of
size $k$, given by $\begin{bmatrix}0&0&\ldots&0&0\end{bmatrix}$, as
$\bm{1}^{k}$ and $\bm{0}^{k}$, respectively. We will also construe the
standard basis vector $\mathbf{e}_{i}$ as a column vector in these notes, and
adhere to the following matrix indexing convention: ${\bf M}[i,j]$ is the
entry in the $i$th row and the $j$th column, ${\bf
M}[i,:]\in\mathbb{F}^{1\times n}$ denotes the $i$th row, and ${\bf
M}[:,j]\in\mathbb{F}^{m\times 1}$ denotes the $j$th column of ${\bf
M}\in\mathbb{F}^{m\times n},$ where $\mathbb{F}$ is a field and the reader can
substitute $\mathbb{F}$ for $\mathbb{R}$ for convenience. We then use
$\bm{1}^{m\times n},\mathbf{0}^{m\times n}\in\mathbb{F}^{m\times 1}$ to denote
the matrix of all $1$s and $0$s, respectively.

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

We also recall the definition of BaseConv for the reader’s convenience:

######  Definition G.1 (BaseConv [23]).

Given an input sequence ${\bm{u}}\in\mathbb{R}^{N\times d},$ where $N$ is the
sequence length and $d$ is the model dimension, a learned weight matrix
${\bm{W}}^{B}\in\mathbb{R}^{d\times d}$ and biases
${\bm{B}}^{B},{\bm{B}}^{K}\in\mathbb{R}^{N\times d}$ and a matrix of
convolution filters ${\bm{K}}\in\mathbb{R}^{N\times d}$, a BaseConv layer
computes the following:

$\bm{z}^{\texttt{BaseConv}}:=({\bm{u}}{{\bm{W}}}^{B}+{\bm{B}}^{B})\odot\left({{\bm{K}}}\ast{\bm{u}}+{\bm{B}}^{K}\right)\in\mathbb{R}^{N\times
d},$ (8)

where the convolutions are applied across the input length $N$.

We will need the following “$5$-tuple" notation for BaseConv model:

######  Definition G.2.

An $\left(N,0pt,d,\tilde{N},\tilde{d}\right)-\text{-}\texttt{BaseConv}$ is a
stacked sequence to sequence model with $L$ layers such that:

  1. 1.

input and output are $N\times d$ matrices,

  2. 2.

each layer corresponds to the a BaseConv layer as defined in Definition G.1,
and

  3. 3.

all the individual gated convolution layers take in $\tilde{N}\times\tilde{d}$
matrices and output $\tilde{N}\times\tilde{d}$ matrices. We refer to the tuple
$(\tilde{N},\tilde{d})$ as the _inner dimension_ of the model.

We also assume that the input $\bm{u}\in\mathbb{R}^{N\times d}$ is embedded
into $\bm{u}^{\prime}\in\mathbb{R}^{\tilde{N}\times\tilde{d}}$ such that

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

The output from the last layer
${\bm{z}}\in\mathbb{R}^{\tilde{N}\times\tilde{d}}$ is transformed into output
$\bm{y}\in R^{N\times d}$ by extracting the top left $N\times d$ entries in
${\bm{z}}$.

######  Definition G.3.

An MLP layer is map $\mathbb{R}^{N\times d}\to\mathbb{R}^{N\times d}$ defined
via matrices ${\bm{W}}^{1},{\bm{W}}^{2}\in\mathbb{R}^{d\times d}$ and “bias"
matrices ${\bm{B}}^{1},{\bm{B}}^{2}\in\mathbb{R}^{N\times d}$ as follows:

$\texttt{MLP}(\bm{u})=\mathrm{ReLU}(\bm{u}{\bm{W}}^{1}+{\bm{B}}^{1}){\bm{W}}^{2}+{\bm{B}}^{2}.$

###  G.1 JRT Lower Bounds for BaseConv

First, we formally define JRT prompts below.

######  Definition G.4 (JRT Prompts).

For any model $\mathcal{M}$ with input $\bm{u}\in\mathbb{R}^{N\times d}$, a
JRT prompt for input $\bm{u}$ is the repeated input
$\bm{u}^{\mathrm{JRT}}\in{\mathbb{R}^{2N\times d}}$ given by

$\bm{u}^{\mathrm{JRT}}[i,:]:=\begin{cases}\bm{u}[i,:]&\text{if }i<N\\\
\bm{u}[i-N,:]&\text{otherwise.}\\\ \end{cases}$

####  G.1.1 Lower Bound on the Number of Layers for AR

In this section, we will provide a lower bound on the number of layers needed
to solve the standard associative recall problem with JRT prompts. We formally
recall the associative recall 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]$.

We also 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 use the following well-known lower bound for the index problem.

######  Theorem G.5 ([80]).

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

We will now mirror the argument from [7, Theorem F.4] to show that the lower
bound on the number of layers for a BaseConv model solving AR still holds for
JRT prompts.

######  Theorem G.6.

Given a JRT prompt ${\bm{u}}^{\mathrm{JRT}}\in\\{0,1\\}^{2N\times d}$ for
input ${\bm{u}}\in\\{0,1\\}^{N\times d}$ to the AR problem with any encoding
such that $\log{c}\leq d\leq 2^{(\log{N})^{1-\epsilon}}$ for $\epsilon>0$, and
$c$ possible tokens from the vocabulary with $c\leq N$, a data-independent
BaseConv model with model parameters taking $O(\log{N})$ bits needs
$\Omega(\epsilon\log\log{N})$ layers to solve AR.

###### Proof.

Given a BaseConv model $\mathcal{M}$ solving AR, regardless of the input
length $N$, we know that there exists an equivalent polynomial
$P({\bm{u}}^{\mathrm{JRT}})$ of degree at most $2^{L}$ that solves AR for any
${\bm{u}}^{\mathrm{JRT}}\in\\{0,1\\}^{2N\times d}$, where $L$ denotes the
number of layers.999See the proof of [7, Theorem F.4] for justification. Now,
take the instance $(\bm{x},i)$ of the index problem with
$\bm{x}\in\\{0,1\\}^{N}$ and the corresponding JRT prompt of the AR problem as
before

${\bm{u}}^{\mathrm{JRT}}:=\\{j,\bm{x}_{j}\\}_{j=0}^{N-1},i,\\{j,\bm{x}_{j}\\}_{j=0}^{N-1},i$
(9)

Next, we build the following one-way protocol for solving the index problem
using the BaseConv model from the hypothesis that it solves AR. Alice with
their access of $\bm{x}\in\\{0,1\\}^{N}$ will again generate a JRT input
${\bm{u}}^{\mathrm{JRT}}$ for AR (without the query) as in equation 9. More
specifically, Alice takes the values
${\bm{a}}:={\bm{u}}^{\mathrm{JRT}}[0:N-2,:]\equiv{\bm{u}}^{\mathrm{JRT}}[N:2N-2,:]\in\\{0,1\\}^{2(N-1)\times
d}$ while leaving out the query
${\bm{q}}:={\bm{u}}^{\mathrm{JRT}}[N-1,:]={\bm{u}}^{\mathrm{JRT}}[2N-1,:]$,
and substitutes these known $2(N-1)d$ values to define the following
polynomial:

$\ Q^{\mathrm{JRT}}({\bm{q}})=P({\bm{a}},{\bm{q}},{\bm{a}},{\bm{q}}).$ (10)

Crucially, $Q^{\mathrm{JRT}}$ is still a polynomial in $d$ variables,
corresponding to the values
${\bm{u}}^{\mathrm{JRT}}[N-1,:]={\bm{u}}^{\mathrm{JRT}}[2N-1,:]$ that Bob has
and trivially has degree $D\leq 2^{L}$. As in the proof of [7, Theorem F.4],
Alice can run the model $\mathcal{M}$, retrieve the coefficients of
$Q^{\mathrm{JRT}}$, and send it to Bob. Since we assume that $P$ solves AR,
Bob can take the coefficients of $Q^{\mathrm{JRT}}$ and substitute
${\bm{u}}^{\mathrm{JRT}}[N-1,:]={\bm{u}}^{\mathrm{JRT}}[2N-1,:]$ to
$Q^{\mathrm{JRT}}$ to compute $P({\bm{u}}^{\mathrm{JRT}})$ which is the value
${\bm{x}}_{i}$.

Moreover, the polynomial $Q^{\mathrm{JRT}}$ that Alice sends still has at most
$d^{2^{L}}$ coefficients as each term in $Q^{\mathrm{JRT}}$ can have degree at
most $2^{L}$. If each such coefficient has $B$ bits, then using theorem G.5,
the total number of bits being communicated must satisfy $B\cdot
d^{2^{L}}\geq\Omega(N)$. This follows from the fact that if $B\cdot
d^{2^{L}}\leq o(N)$, then since the associated value of $i$ in equation 9 is
the answer to the indexing problem, we have shown that a one-way communication
protocol for solving the index problem uses $o(N)$ communication complexity,
which then contradicts theorem G.5. This is the same equation we get in the
proof of [7, Theorem F.4], which yields the following lower bound on the
number of layers:

$\displaystyle L$
$\displaystyle\geq{\log\left(\frac{\log{N}-\log{B}}{(\log{N})^{1-\epsilon}}\right)}.$
(11)

Recall here that the model parameters are assumed to be $O(\log{N})$ bits, so
any coefficient in $Q^{\mathrm{JRT}}$ should have absolute value at most
$\left(2^{O(\log{N})}\cdot 2Nd\right)^{2^{L}}$ as each coefficient can be a
product of at most $2Nd$ variables. That is, for some $\alpha>0$, we have the
following bound on each coefficient:

$2^{B}\leq(2\cdot N^{\alpha+1}d)^{2^{L}}\leq(2N^{(\alpha+2)})^{2^{L}}$

where the last equality uses the fact that $d\leq
2^{\log{N}^{(1-\epsilon)}}\leq N$. We thus have

$\log(B)\leq\log(\alpha+2)+L+\log\log\left(2N\right).$ (12)

Substituting equation 12 to equation 11, we get

$\displaystyle
L\geq\log\left(\frac{\log{N}-\log(\alpha+2)-L-\log\log\left(2N\right)}{(\log{N})^{1-\epsilon}}\right)$
(13)

Now, if $L>\log\log{2N}$, we are done. Otherwise, if
$L\leq\log\log\left(2N\right)$, then we can substitute this to equation 13 to
get

$\displaystyle L$
$\displaystyle\geq\log\left(\frac{\log{N}-\log(\alpha+2)-2\log\log\left(2N\right)}{(\log{N})^{1-\epsilon}}\right)$
$\displaystyle=\log\left(\log{N}-\log(\alpha+2)-2\log\log{2N}\right)-(1-\epsilon)\log\log{N}$
(14)

We now claim that first term in equation 14 satisfies the following:

$\log\left({\log{N}-\log(\alpha+2)-2\log\log\left(2N\right)}\right)\geq(1-\frac{\epsilon}{2})\log\log{N}.$
(15)

To see this, note that, for sufficiently large enough $N$, the following
holds:

$\displaystyle\frac{\log{N}}{2}\geq\log(\alpha+2)+2\log\log\left(2N\right),$

hence, we get

$\log\left({\log{N}-\log(\alpha+2)-2\log\log\left(2N\right)}\right)\geq\log\left(\frac{\log{N}}{2}\right)\geq\log\log{N}-1\geq(1-\frac{\epsilon}{2})\log\log{N}.$

This proves the claim in equation 15. Finally, using equation 15, equation 14
leads to the following:

$L\geq(1-\frac{\epsilon}{2})\log\log{N}-(1-\epsilon)\log\log{N}=\frac{\epsilon}{2}\log\log{N},$

which still provides the lower bound $L=\Omega(\epsilon\log\log{N})$, as
desired. ∎

####  G.1.2 Lower Bounds for MQAR with $d=\log_{2}{c}$

Next, we present lower bounds for the mulitple-query associative recall (MQAR)
problem which generalizes the AR problem [23]. To this end, we recall the
definition of MQAR below.

> Suppose we are given an input sequence $\bm{u}[0\cdots
> 3N-1]\triangleq\left\\{\left(\bm{k}_{0},\bm{v}_{0},\bm{q}_{0}\right),\ldots,\left(\bm{k}_{{N}-1},\bm{v}_{{N}-1},\bm{q}_{{N}-1}\right)\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\leq
> i\leq{N}-1$, whether there exists $0\leq j<i$ such that
> $\bm{q}_{i}\equiv\bm{k}_{j}$, and if so, output $\bm{v}_{j}$.

We now present the following lower bound from [7] for the MQAR problem
$d=\log_{2}{c}$ to encode all $c$ possible tokens from $C$ using the natural
binary encoding, which also holds for JRT input. This is because the result
(Theorem F.5) in [7] is derived using Lemma 5.1 in [7] (degree of multilinear
polynomial computed by BaseConv in terms of its number of layers) and Lemma
5.2 in [7] (degree of multilinear polynomial for the MQAR problem), both of
which are independent of the input length $N$.

######  Theorem G.7.

A data-independent BaseConv model needs $\log(2d)$-layers to solve
$\mathrm{MQAR}$ with a JRT prompt $\bm{u}\in\\{0,1\\}^{2\cdot 3N\times d}$ for
the original input $\bm{u}\in\\{0,1\\}^{3N\times d}$ with $d=\log_{2}(c)$.

####  G.1.3 Lower Bounds for MQAR via the Equality (EQ) Problem

[7] also contains lower bounds on the number of layers solving MQAR due to the
lower bounds on the equality problem (EQ), where we define the equality
problem (EQ) as checking whether the two encodings are equal:
$\bm{u}_{1}\equiv\bm{u}_{2}$ for an input pair $\bm{u}_{1},\bm{u}_{2}$ where
each $\bm{u}_{i}$ is a token drawn from a vocabulary of size $c=|C|$ and
embedded in $\\{0,1\\}^{d}$.

We next show that any model with JRT prompts solving MQAR also solves EQ.

######  Proposition G.8.

Any model $M_{\mathrm{MQAR}}$ that solves MQAR with JRT prompt also solves EQ
using the same number of layers.

###### Proof.

If there exists a model $\textsc{M}_{\mathrm{MQAR}}$ that solves MQAR using
$L$ layers with JRT prompt, then for an arbitrary input instance for EQ given
by $\bm{u}_{1},\bm{u}_{2}\in\mathbb{R}^{2\times d}$, we can produce the
following input instance for MQAR:
$\bm{u}:=\\{(\bm{u}_{1},\mathbbm{1},\bm{u}_{1}),(\bm{u}_{2},\mathbbm{1},\bm{u}_{2}),(\bm{u}_{1},\mathbbm{1},\bm{u}_{1}),(\bm{u}_{2},\mathbbm{1},\bm{u}_{2})\\}$
and solve EQ using $L$ layers with $\textsc{M}_{\mathrm{MQAR}}$ returning
$\mathbbm{1}$ iff there is a match. ∎

Due to proposition G.8, we obtain the following corollary.

######  Corollary G.9.

Any lower bound $\overline{L}$ on the number of layers $L$ of BaseConv to
solving EQ is also a lower bound on the number of layers required for solving
$\mathrm{MQAR}$ with JRT prompts.

The lower bounds for the EQ problem in [7] depends on showing that the
polynomial $P$ representing EQ in $p$-hot encoding has $\deg(P)\geq 2p$, which
does not depend on the sequence length (Proposition F.5). Since corollary G.9
also holds in the JRT setting, we inherit the lower following lower bound for
BaseConv solving MQAR in the $p$-hot encoding setting, which we recall here
for the reader’s convenience.

######  Definition G.10 ($p$-Hot Encoding).

We define the $p$-hot encoding to be the collection of embeddings for a token
$\bm{x}_{t}$ with $0\leq t<c$ such that we express $t$ in base
$\sqrt[p]{c}:(t_{0},..,t_{p-1})\in[0,\sqrt[p]{c})^{p}$ and represent each
$t_{i}$ as one hot encoding in $\\{0,1\\}^{\sqrt[p]{c}}$. That is, we take
$d=p\cdot\sqrt[p]{c}$.

######  Theorem G.11.

A data-independent BaseConv model needs at least
$\lfloor\log(2p)\rfloor$-layers to solve $\mathrm{MQAR}$ for a JRT prompt
$\bm{u}^{\mathrm{JRT}}\in\\{0,1\\}^{2\cdot 3N\times d}$ for the original input
$\bm{u}\in\\{0,1\\}^{3N\times d}$ in the $p$-hot encoding setting, where
$d=p\cdot\sqrt[p]{c}$.

###  G.2 Recurrent Models and Set Disjointness

In this section, we will provide upper bounds on the class of recurrent models
defined in [7] solving the set disjointness (SD) problem. First, we recall the
definition of recurrent models below.

######  Definition G.12 (Recurrent Models).

A model $\mathcal{M}$ taking an input ${\bm{u}}\in\mathbb{R}^{N\times d}$,
where $N$ is the input length and $d$ is the model dimension, is termed a
_recurrent model_ if its $i$-th state, representing the output at location
$i$, ${\bm{Z}}_{\mathcal{M}}^{i}\in\mathbb{R}^{\tilde{d}}$, with $\tilde{d}$
denoting the state size, is determined exclusively by the preceding elements
of the input ${\bm{u}}[0\ldots i-1]$. The state ${\bm{Z}}_{\mathcal{M}}^{i}$
represents the accumulated information of the model depending on the inputs up
to the $i$-th element, and is distinct from learned parameters that are static
with respect to the input sequence.

Specifically, ${\bm{Z}}_{\mathcal{M}}^{i}({\bm{u}})=\phi({\bm{u}}[0\ldots
i-1])$, indicating that the state is a function of the input history but not
of the entire input sequence simultaneously. Moreover, we can express this as:

${\bm{Z}}_{\mathcal{M}}^{i}({\bm{u}})=f_{\mathcal{M}}^{i}({\bm{Z}}_{\mathcal{M}}^{i-1},{\bm{u}}[i]),$
(16)

for a sequence of functions $\\{f_{\mathcal{M}}^{i}\\}_{i\in[N]}$, where each
function is tailored to evolve the state based on the immediate past state and
the current input.

######  Remark G.13.

Note that definition G.12 excludes models that inherently require the entire
input sequence for computation at any state, such as those based on non-causal
convolutional operations over the full input.

######  Remark G.14.

Given sets $A,B\subseteq\\{0,1\\}^{n}$, the set disjointness (SD) problem
seeks to check whether $A$ and $B$ are disjoint, that is, $A\cap B=\emptyset$.
First, we clarify the format of the input $\bm{u}\in\\{0,1\\}^{N\times(n+1)}$
for the set-disjointness problem with $N=\left\lvert A\right\rvert+\left\lvert
B\right\rvert+1$. The rows of the input $\bm{u}\in\\{0,1\\}^{N\times(n+1)}$
correspond to elements in $A$ and $B$. That is, $\bm{u}[i,0:n-1]\in A\cup
B\cup\\{\bm{0}^{n}\\}$, where $\\{[\bm{0}^{n}::1]\\}$ is a separator element
which separates the contiguously placed (in any arbitrary order) elements of
each set with the last entry of non-separator rows equal to $0$.

######  Theorem G.15.

For any recurrent model $\mathcal{M}$, there exists a function of the input
history
${\bm{Z}}_{\mathcal{M}}^{i}({\bm{u}}^{\mathrm{JRT}})=\phi({\bm{u}}^{\mathrm{JRT}}[0\ldots
i-1])$ that solves the set disjointness problem with
${\bm{Z}}_{\mathcal{M}}^{2N}$ of size $\mathcal{O}(n\cdot\min\\{|A|,|B|\\})$
for the $\mathrm{JRT}$ prompt
${\bm{u}}^{\mathrm{JRT}}\in\\{0,1\\}^{2N\times(n+1)}$ of the input
$\bm{u}\in\\{0,1\\}^{N\times(n+1)}$ for the set-disjointness problem.

###### Proof.

Given a JRT prompt $\bm{u}^{\mathrm{JRT}}\in\\{0,1\\}^{2N\times(n+1)}$
corresponding to the input for the set-disjointness problem, for a recurrent
model $\mathcal{M}$, we define the state ${\bm{Z}}_{\mathcal{M}}^{i}$ in
Algorithm 3.

Algorithm 3 Recurrent Model for Set Disjointness

1: an input $\bm{u}^{\mathrm{JRT}}\in\\{0,1\\}^{2N\times(n+1)}$ for the set-
disjointness problem

2: state size ${\bm{Z}}_{\mathcal{M}}^{2N-1}$.

3: firstSeparator $\leftarrow$ False

4: secondSeparator $\leftarrow$ False

5: smallFirst $\leftarrow$ False

6:for $i\leftarrow$ $0$ to ${2N-1}$ do

7: if $\bm{u}^{\mathrm{JRT}}[i,n]=1$ then

8: if firstSeparator = False then

9: firstSeparator $\leftarrow$ True

10: if $i\leq\lfloor\frac{N}{2}\rfloor$ then

11: smallFirst $\leftarrow$ True

12: else

13: secondSeparator $\leftarrow$ True

14: else

15: if firstSeparator = True then

16: if smallFirst = True then

17: if secondSeparator = False then

18: if $i\geq N$ then

19: Add $\bm{u}^{\mathrm{JRT}}[i,:]$ to ${\bm{Z}}_{\mathcal{M}}^{i}$

20: else

21: if there exists $j$ s.t.
$\bm{u}^{\mathrm{JRT}}[i,:]={\bm{Z}}_{\mathcal{M}}^{i-1}[j,:]$ then

22: ${\bm{Z}}_{\mathcal{M}}^{i-1}[j,n]=1$

23: else

24: if secondSeparator = False then

25: if $i\leq N$ then

26: Add $\bm{u}^{\mathrm{JRT}}[i,:]$ to ${\bm{Z}}_{\mathcal{M}}^{i}$

27: else

28: if there exists $j$ s.t.
$\bm{u}^{\mathrm{JRT}}[i,:]={\bm{Z}}_{\mathcal{M}}^{i-1}[j,:]$ then

29: ${\bm{Z}}_{\mathcal{M}}^{i-1}[j,n]=1$

30:for all $j$ s.t. ${\bm{Z}}_{\mathcal{M}}^{i-1}[j,n]=1$ do

31: return ${\bm{Z}}_{\mathcal{M}}^{i-1}[j,0:n-1]$.

Semantically, we take a JRT input
$\bm{u}^{\mathrm{JRT}}\in\\{0,1\\}^{2N\times(n+1)}$ for the set-disjointness
problem, and find the first separator (lines 7 to 11). If the index $i$ of the
first separator is less than or equal to $\lfloor\frac{N}{2}\rfloor$ (line
10), then we know that the smaller set is placed before the larger set.
Otherwise, the smaller set is placed later (see Figure 4).

$N$$N$$N$$N$$\mathbf{0}^{n}::1$$\mathbf{0}^{n}::1$$\mathbf{0}^{n}::1$$\mathbf{0}^{n}::1$
Figure 4: Placement of the smaller set is determined by when we first
encounter the separator.

Either way, we want to store the smaller set and compare it against the larger
set for intersections. To this end, if the smaller set comes first (line 16),
then we continue until the beginning of the repeated input (line 18) and
collect the smaller set (line 19), which we then use after we encounter the
second separator (lines 21 to 22) to compare against the larger set. If the
smaller set comes second (lines 23 to 29), then after the first separator, we
collect the smaller set (lines 25 to 26) and compare it against the larger set
that comes right after (lines 27 to 29).

For comparison (lines 30 to 31), we use the separator flag at the end. Recall
that non-separator elements of the input have $0$ in the separator flag index,
and thus, so do the elements from the smaller set collected in the state
$\bm{Z}_{\mathcal{M}}$. When comparing against the elements from the larger
set, we simply set the flag to $1$ for an element that is in the intersection
of two sets.

Now, we examine the space requirement for the state $\bm{Z}_{\mathcal{M}}$ of
the model $\mathcal{M}$. Note that we only add an element to
$\bm{Z}_{\mathcal{M}}$ in lines 19 and 26. In both cases, the elements are
from the smaller set, and thus, $|\bm{Z}_{\mathcal{M}}|=\min\\{|A|,|B|\\}$.
Moreover, each element in $A$ and $B$ is of size $n$, and thus, we can
conclude that the model $\mathcal{M}$ with state $\bm{Z}_{\mathcal{M}}$ can
solve the set-disjointness problem with JRT input in
$\mathcal{O}(n\cdot\min\\{|A|,|B|\\})$. ∎

###  G.3 Based Solving SD

In this section, we will show that Based can solve the set disjointness
problem with $\mathrm{JRT}$ inputs. Specifically, this section implements
Algorithm 3 in the Based architecture. Recall here that the Based model
combines two layer types: BaseConv (see definition G.1) and LinearAttention
defined below.

######  Definition G.16 (Linear Attention with Kernels).

Given an input sequence ${\bm{u}}\in\mathbb{R}^{N\times d},$ where $N$ is the
sequence length and $d$ is the model dimension, kernel projections101010By
kernel projections of a matrix ${\bm{u}}\in\mathbb{R}^{m\times n},$ we mean
applying some kernel map $\phi:\mathbb{R}^{N\times d}\to\mathbb{R}^{N\times
f}$ to each row of ${\bm{u}}$.
$\texttt{Projection}_{q},\texttt{Projection}_{k}\in\mathbb{R}^{d\times f}$,
$\texttt{Projection}_{v}\in\mathbb{R}^{d\times d}$, where $f$ is the feature
dimension, the LinearAttention layer computes the following:

$\bm{z}^{\texttt{LinearAttention}}:=\left({{\bm{Q}}}\
{{\bm{K}}}^{\top}\right){\bm{V}}\in\mathbb{R}^{N\times d},$ (17)

where
${\bm{Q}}:=\texttt{Projection}_{q}({\bm{u}}),{\bm{K}}:=\texttt{Projection}_{k}({\bm{u}}),{\bm{V}}:=\texttt{Projection}_{v}({\bm{u}})$.

####  G.3.1 SD with LinearAttention

We first show that with appropriate placement of the two sets, we can solve
the set disjointness problem using a class of kernel maps defined below.

######  Definition G.17 ($IP$-Kernel).

We define the IP-Kernel to be the kernel map
$\phi_{\epsilon,f}:\mathbb{R}^{d}\to\mathbb{R}^{f}$ that takes elements from
$[c]$ to $\mathbb{R}^{f}$ so that, for any $x,y\in[c]$, we have

${\langle\phi_{\epsilon,f}(x),\phi_{\epsilon,f}(y)\rangle}=1\text{ if
}x=y\text{ and
}\left\lvert\langle\phi_{\epsilon,f}(x),\phi_{\epsilon,f}(y)\rangle\right\rvert\leq\epsilon\text{
otherwise}.$

That is, an IP-kernel projects elements from the universal set $[c]$ so that
the inner products are approximately orthogonal. Note that the feature
dimension $f$ is dependent on the tolerance $\epsilon$.

We now show that if there exists an IP kernel with small enough $\epsilon$,
then it can be used to solve the set-disjointness problem with a Linear
Attention layer followed by an MLP layer.

######  Proposition G.18.

Given an input $\bm{u}\in\mathbb{R}^{N\times d}$ encoding the input $(A,B)$ to
the set-disjointness problem (SD) on sets $A,B\subseteq[c]$, there exists a
Linear Attention (+ MLP) layer with state space $O(df)$ that solves the set
disjointness problem for ${\bm{u}}\in\mathbb{R}^{N\times d}$ with the IP
kernel $\phi_{\epsilon,f}$ applied on ${\bm{Q}},{\bm{K}}$ for
$\epsilon=\frac{1}{3|A|}$.111111Our notion of ‘solves’ is a bit non-standard
so we clarify it here. If ${\bm{z}}\in\mathbb{R}^{N\times d}$ is the output
then it encodes the result as follows. If the $i$th element in $B$ appears in
$A$ then ${\bm{z}}[|A|+i,:]$ has all entries in $\left[\frac{1}{3},1\right]$,
otherwise it is $\bm{0}^{d}$. If we want a single value as an answer (since SD
has a Boolean output) we can apply $O(1)$ BaseConv layers on ${\bm{z}}$ to sum
up all the values in the last $|B|$ rows of ${\bm{z}}$. Then if $A\cap
B\neq\emptyset$ then this value is at least $\frac{d}{3}$, otherwise it is
$0$.

###### Proof.

We first define the keys and queries along with the values for the Linear
Attention layer as follows:

${\bm{Q}}[i,:]={\bm{K}}[i,:]=\phi_{\epsilon,f}({\bm{u}}[i,:])\text{ and
}{\bm{V}}[i,j]:=\begin{cases}1&\text{if }i<|A|\\\
0&\text{otherwise}.\end{cases}$

Note that ${\bm{Q}},{\bm{K}}\in\mathbb{R}^{N\times f}$ and
${\bm{V}}\in\mathbb{R}^{N\times d}$.

$\displaystyle\left({{\bm{Q}}}\ {{\bm{K}}}^{\top}\right)[i,j]$
$\displaystyle:={{\bm{Q}}}[i,:]{{\bm{K}}}^{\top}[:,j]$
$\displaystyle=\langle{{\bm{Q}}}[i,:],{{\bm{K}}}[j,:]\rangle$
$\displaystyle=\langle\phi_{\epsilon,f}({\bm{u}}[i,:]),\phi_{\epsilon,f}({\bm{u}}[j,:])\rangle$

Next, the key-query product yields the following

$\displaystyle\bm{z}^{\texttt{LinearAttention}}[i,j]$
$\displaystyle:=\left({{\bm{Q}}}\ {{\bm{K}}}^{\top}\right)[i,:]{\bm{V}}[:,j]$
$\displaystyle=\sum_{k=0}^{N-1}\left({{\bm{Q}}}\
{{\bm{K}}}^{\top}\right)[i,k]\cdot{\bm{V}}[k,j]$
$\displaystyle=\sum_{k=0}^{N-1}\langle\phi({\bm{u}}[i,:]),\phi({\bm{u}}[k,:])\rangle\cdot{\bm{V}}[k,j]$
$\displaystyle=\sum_{k<|A|}\langle\phi_{\epsilon,f}({\bm{u}}[i,:]),\phi_{\epsilon,f}({\bm{u}}[k,:])\rangle$
$\displaystyle=:\rho_{i},$

where the second-last equality follows from the definition of ${\bm{V}}$ and
we can specify $\rho_{i}$ as follows:

$\displaystyle\rho_{i}=1\pm\epsilon\cdot|A|\text{ if there exists
}k\in[0\cdots|A|-1]\text{ s.t. }{\bm{u}}[k,:]\equiv{\bm{u}}[i,:],\text{ and
otherwise, }\rho_{i}\leq\epsilon|A|.$ (18)

For the MLP layer, we define the following parameters (see Definition G.3 for
notation):

${\bm{W}}^{1}={\bm{I}}_{d\times
d},\quad{\bm{B}}^{1}_{\mathrm{MLP}}:=-\frac{1}{3}\bm{1}_{N\times
d},\quad{\bm{W}}^{2}_{\mathrm{MLP}}={\bm{I}}_{d\times
d},\quad{\bm{B}}^{2}_{\mathrm{MLP}}=\bm{0}_{N\times d}$

Next, we note that for $0\leq\ell<N$ and $0\leq j<d$:

$\displaystyle\bm{y}[\ell,j]$
$\displaystyle:=\left(\bm{z}^{\texttt{LinearAttention}}{\bm{W}}^{1}_{\mathrm{MLP}}+{\bm{B}}^{1}_{\mathrm{MLP}}\right)[\ell,j]$
$\displaystyle=\left(\bm{z}^{\texttt{LinearAttention}}-\frac{1}{3}\bm{1}_{|B|\times
d}\right)[\ell,j]$ $\displaystyle=\left(\rho_{\ell}-\frac{1}{3}\right).$

We now use the fact that $\epsilon\leq\frac{1}{3|A|}$ to get bounds on the
above. To this end, for $0\leq\ell<N$, due to equation 18, if there exists
$k\in[0\cdots|A|-1]\text{ such that }{\bm{u}}[k,:]\equiv{\bm{u}}[\ell,:]$, we
have

$\bm{y}[\ell,j]=\left(\rho_{\ell}-\frac{1}{3}\right):=\left(\left(1\pm\epsilon\cdot|A|\right)-\frac{1}{3}\right)\in\left[\frac{2}{3},\frac{4}{3}\right]-\frac{1}{3}=\left[\frac{1}{3},1\right]$

Otherwise, if there is no match, then we have

$\bm{y}[\ell,j]=\left(\rho_{\ell}-\frac{1}{3}\right)\leq\epsilon\cdot|A|-\frac{1}{3}\leq\frac{1}{3}-\frac{1}{3}\leq
0.$

We then get the final output as

$\bm{z}:=\mathrm{ReLU}(\bm{y}){\bm{W}}^{2}_{\mathrm{MLP}}+{\bm{B}}^{2}_{\mathrm{MLP}}=\mathrm{ReLU}(\bm{y}),$

which reduces to

$\bm{z}[\ell,j]\in\left[\frac{1}{3},1\right]\text{ if there exists
}k\in[0\cdots|A|-1]\text{ such that }{\bm{u}}[k,:]\equiv{\bm{u}}[i,:],\text{
and $0$ otherwise}.$

Therefore, the last $|B|$ rows of the output $\bm{z}$ will have non-zero
values if and only if $A\cap B\neq\phi$. Finally, the claim on $O(df)$ space
follows from the well-known recurrent view of LinearAttention (see equation
2).121212To incorporate the MLP part, note that as soon as each row of
${\bm{z}}^{\texttt{LinearAttention}}$ is generated, we can generate the output
of the corresponding row in
$\texttt{MLP}({\bm{z}}^{\texttt{LinearAttention}})$ with $O(d)$ space by
noting that MLP operates independently on each row of its input. ∎

####  G.3.2 Realization of IP Kernels

In this section, we will provide some instances of realizing the IP kernels
from Definition G.17.

##### Exponential Kernels.

The first IP-kernel that we define is the exponential kernel
$\phi^{\exp}:\mathbb{R}^{d}\to\mathbb{R}^{f}$ such that for any $x,y\in[c]$,
we have

$\langle\phi(x),\phi(y)\rangle=\exp\left(\langle x,y\rangle\right),$

where $x$ and $y$ are encoding of the corresponding elements of $[c]$ in
$\\{-1,1\\}^{d},d=O(\log(c))$ with large enough distance131313Specifically, we
will need to use well-known construction of Binary codes with constant rate
and constant relative distance [81].. If $x=y$, we have

$\langle\phi(x),\phi(y)\rangle=\langle\phi(x),\phi(x)\rangle$
$=\exp\left(\langle
x,x\rangle\right)=\exp\left(\sum_{i\in[d]}x_{i}^{2}\right)=\exp\left(\sum_{i\in[d]}1\right)=\exp(d).$

Next, if $x\neq y$, we instead have

$0<\langle\phi(x),\phi(y)\rangle=\exp\left(\langle
x,y\rangle\right)\leq\exp\left(\gamma\cdot d\right)$

for some $\gamma<1$ as the code has constant relative distance. Here, we want
the match $\exp(d)$ to be large enough. That is, we want

$\frac{\exp(d)}{\exp(\gamma\cdot d)}\gg c$

So, we want to pick $d$ large enough so that

$(1-\gamma)\cdot d\gg\ln{c}.$

##### Data-Dependent Kernels.

Here, we define the kernel $\phi$ based on the smaller set $A$. We start by
letting $d:=|A|+\log{c}$ so that we define the embeddings as

$\displaystyle\phi:[c]$ $\displaystyle\to\mathbb{R}^{|A|+\log{c}}$ (19)
$\displaystyle A\ni a$
$\displaystyle\mapsto\begin{bmatrix}{\bm{e}}_{a}&\bm{0}^{\log{c}}\end{bmatrix}$
$\displaystyle A\not\ni b$
$\displaystyle\mapsto\begin{bmatrix}\bm{0}^{|A|}&{\bm{B}}_{b}\end{bmatrix}$

where ${\bm{e}}_{a}\in\\{0,1\\}^{|A|}$ is the $1$-hot encoding of the element
$a$ in $A$ and ${\bm{B}}_{b}$ is the natural binary encoding in $[c]$ on the
element $b$. Using this kernel $\phi$, we achieve orthogonality:

$\langle\phi(x),\phi(y)\rangle=\delta_{xy}.$

That is, we have the tolerance $\epsilon=0$ with feature dimension
$f=|A|+\log_{2}{c}$.

##### Randomized Kernels.

We can also define a random kernel map

$\phi:[c]\to\frac{1}{\sqrt{f}}[-1,1]^{f}.$ (20)

That is, for each $x\in[c]$, we pick a random vector in $\\{-1,1\\}^{f}$ and
normalize it by dividing by $\sqrt{f}$. Here, it is easy to see that for every
$x\in[c]$, we have

$\langle\phi(x),\phi(x)\rangle=\frac{1}{f}\sum_{i\in[f]}1=1.$

Now, for every $x\neq y$, we can apply known concentration inequalities on
Rademacher random variables to get

$\text{Pr\/}\left[\langle\phi(x),\phi(y)\rangle>\frac{t}{\sqrt{f}}\right]\leq
e^{\frac{-t^{2}}{2}}.$

We then pick $t=O(\sqrt{\log{c}})$ so that over all $c^{2}$ pairs, we have

$\text{Pr\/}\left[\langle\phi(x),\phi(y)\rangle>\frac{O(\sqrt{\log{c}})}{\sqrt{f}}\right]<\frac{1}{100c^{2}}.$

Then with a union bound on all $c^{2}$ pairs, with high probability, we get
that $\phi$ has $\epsilon=\frac{t}{\sqrt{f}}$. We then want the threshold to
satisfy the following:

$t/\sqrt{f}<\frac{1}{3|A|}\implies f=\Omega(|A|^{2}\log{c}).$

That is, for $\epsilon=\frac{1}{3|A|}$,
$f=\Theta(\min\\{|A|,|B|\\}^{2}\log{c})$ suffices.

######  Remark G.19 (Practical Justification).

Empirically, prior works shows a variety of kernels that are competitive with
softmax attention quality while using a small amount of space. For instance,
Zhang et al. [76] show that either training MLP projections to mimic softmax
attention weights or using a $2^{\mathrm{nd}}$-order Taylor approximation to
the softmax-exponential function are two effective kernel function choices.
The $2^{\mathrm{nd}}$-order polynomial is only a high fidelity approximation
within a small band of real values, however empirically results in Arora et
al. [7] suggest that the normalized query-key dot products often fall within
this range, resulting in competitive quality with softmax attention. Arora et
al. [7], Chen et al. [82], and others further suggest that combining efficient
sparse plus low-rank attentions (e.g., linear attention plus dense, local
sliding window attention) further diminishes quality gaps versus full
attention.

####  G.3.3 Shifts with BaseConv

Next, we will show that we can use BaseConv layers to move the smaller set to
the start of the sequence. First, based on whether the smaller set is at the
start or not, we need to define separate convolution kernels based on the
input. To this end, we use the following BaseConv model to derive these
kernels.

######  Lemma G.20.

There exists
$\left(2N,O(1),(n+1),\left(2N+\frac{N}{2}\right),(n+1)\right)-\text{{BaseConv}}$
model that takes in a $\mathrm{JRT}$ prompt
${\bm{u}}^{\mathrm{JRT}}\in\mathbb{R}^{2N\times(n+1)}$ of the input
$\bm{u}\in\mathbb{R}^{N\times(n+1)}$ for the set-disjointness (SD) problem
$(A,B)\subseteq\\{0,1\\}^{n}$ and outputs the kernel
${\bm{h}}_{\mathrm{shift}}$ that shifts the input ${\bm{u}}^{\mathrm{JRT}}$ to
get the smaller set at the start of the sequence, where

${\bm{h}}_{\mathrm{shift}}(X):=\begin{cases}X^{|A|+1}&\text{if }|A|\geq|B|\\\
1&\text{otherwise.}\end{cases}$ (21)

###### Proof.

Following the proof of Proposition G.18, we know that it suffices to find the
location of the separator to determine the location of the smaller set. More
specifically, if the separator is within $\left[0,\frac{N}{2}-1\right]$ row
index range, then we know that the smaller set is at the start, and the kernel
being generated is the identity. Otherwise, we generate the kernel $X^{|A|+1}$
which will be used in the proof of Proposition G.21.

We first increase the inner dimension of the JRT input
${\bm{u}}^{\mathrm{JRT}}\in\mathbb{R}^{2N\times(n+1)}$ to
${\bm{u}}^{JRT}_{\mathrm{inner}}\in\mathbb{R}^{\left(2N+\frac{N}{2}\right)\times(n+1)}$
so that we introduce a zero-block between the first seperator and the start of
set $B$. That is, we have

${\bm{u}}^{\mathrm{JRT}}_{\mathrm{inner}}[i,:]=\begin{cases}{\bm{u}}^{\mathrm{JRT}}[i,:]&\text{if
}i<\frac{N}{2}\\\ \bm{0}^{n+1}&\text{if }\frac{N}{2}\leq i<N\\\
{\bm{u}}^{\mathrm{JRT}}[i-\frac{N}{2},:]&\text{if }i\geq N.\end{cases}$

We can achieve this by simply using the remembering primitive from [7,
Definition F.15, Proposition F.13] using a
$\left(\left(2N+\frac{N}{2}\right),8,(n+1),\left(2N+\frac{N}{2}\right),(n+1)\right)-\text{{BaseConv}}$
to remember ${\bm{u}}^{\mathrm{JRT}}[\frac{N}{2}:2N-1,:]$ while applying the
identity kernel to preserve ${\bm{u}}^{\mathrm{JRT}}[0:\frac{N}{2}-1,:]$.

We again apply the remembering primitive from [7, Definition F.15, Proposition
F.13] to get

${\bm{Y}}\leftarrow{\tt
remember}({\bm{u}}^{\mathrm{JRT}}_{\mathrm{inner}},0,N,f),$

using
$\left(\left(2N+\frac{N}{2}\right),8,(n+1),\left(2N+\frac{N}{2}\right),(n+1)\right)-\text{{BaseConv}}$,
where $f$ is applied over
${\bm{x}}:={\bm{u}}^{\mathrm{JRT}}_{\mathrm{inner}}[0:N-1,:]$, the first $N$
rows of ${\bm{u}}^{\mathrm{JRT}}_{\mathrm{inner}}$. That is, we want to
remember the last $\left(N+\frac{N}{2}\right)$ rows of
${\bm{u}}^{\mathrm{JRT}}_{\mathrm{inner}}$. We define $f:=f_{2}\circ f_{1}$,
where $f_{1}$ is the cumulative sum of the first $N$ rows computed using
$\left(N,O(1),(n+1),N,(n+1)\right)-\text{{BaseConv}}$ followed by $f_{2}$
which is the shifting down by $N-1$ using
$\left(N,3,(n+1),N,(n+1)\right)-\text{{BaseConv}}$ [7, Propositions F.41 and
F.38]. That is, for $i\in[0:N-1]$, we have

$\displaystyle f_{1}({\bm{x}})[i,:]$
$\displaystyle=\sum_{k=0}^{i}{\bm{x}}[k,:];$ $\displaystyle
f_{2}(f_{1}({\bm{x}}))[i,:]$ $\displaystyle=f_{1}({\bm{x}})[i-(N-1),:].$

For the $n$th column, we know that for $0\leq i<N$:

${\bm{x}}[i,n]={\bm{u}}^{\mathrm{JRT}}_{\mathrm{inner}}[i,n]={\bm{u}}^{\mathrm{JRT}}[i,n]=\begin{cases}1&\text{if
}|A|\leq|B|\text{ and }i=|A|\\\ 0&\text{otherwise.}\end{cases}$

This is because if $|A|\leq|B|$, the separator is within
$\left[0,\frac{N}{2}-1\right]$ and its $n$th bit is $1$, where
$|A|=:i_{s}\in\left[0,\frac{N}{2}-1\right]$ to be the location of the
separator. We then get

$\displaystyle f_{1}({\bm{x}})[i,n]$ $\displaystyle=\begin{cases}1&\text{if
}|A|\leq|B|\text{ and }i\geq i_{s}\\\ 0&\text{otherwise.}\end{cases}$
$\displaystyle f_{2}(f_{1}({\bm{x}}))[i,n]$
$\displaystyle=\begin{cases}1&\text{if }|A|\leq|B|\text{ and }i=0\\\
0&\text{otherwise.}\end{cases}$

We can thus characterize the $n$th column of the output
${\bm{Y}}\in\mathbb{R}^{\left(2N+\frac{N}{2}\right)\times(n+1)}$ as follows:

${\bm{Y}}[i,n]=\begin{cases}1&\text{if }|A|\leq|B|\text{ and }i=0\\\
0&\text{if }|A|>|B|\text{ and }i=0\text{ or }1\leq i<N\\\
{\bm{u}}^{\mathrm{JRT}}[i+\frac{N}{2},n]&\text{if }i\geq N.\end{cases}$

We now remember ${\bm{Y}}[0:\frac{N}{2}-1,:]$ while shifting down
${\bm{Y}}[\frac{N}{2}:2N+\frac{N}{2}-1,:]$ by $\frac{N}{2}-1$ [7, Proposition
F.13 and F.38] to get ${\bm{Y}}^{\prime}$ such that:

$\displaystyle{\bm{Y}}^{\prime}[i,:]$
$\displaystyle=\begin{cases}{\bm{Y}}[i,:]&\text{if }i<\frac{N}{2}\\\
{\bm{Y}}[i-\frac{N}{2},:]&\text{if }i\geq\frac{N}{2}\end{cases}$
$\displaystyle=\begin{cases}{\bm{Y}}[i,:]&\text{if }i<\frac{N}{2}\\\
{\bm{u}}^{\mathrm{JRT}}[i,:]&\text{if }\frac{N}{2}\leq i<2N-1\\\
\bm{0}^{n}&\text{otherwise.}\end{cases}$

Focusing on the $n$th column, we see that we get for $0\leq i<N$:

$\displaystyle{\bm{Y}}^{\prime}[i,n]$ $\displaystyle=\begin{cases}1&\text{if
}|A|\leq|B|\text{ and }i=0\text{ or }|A|>|B|\text{ and }i=|A|\\\
0&\text{otherwise}\end{cases}.$

Or equivalently

$\begin{aligned}
{\bm{Y}}^{\prime}[0:N-1,n]&=\begin{cases}{\bm{e}}_{0}&\text{if }|A|\leq|B|\\\
{\bm{e}}_{|A|}&\text{if }|A|>|B|.\end{cases}\end{aligned},$

which is exactly what we need as the shift kernel ${\bm{h}}_{\mathrm{shift}}$.
A schematic representation of this process is provided in Figure 5. The final
claim on the overall parameters follows from the fact that we can ‘stack’
BaseConv layers with the same internal dimension [7].

$\frac{N}{2}-1$$N-1$$1$$1$$|A|\ \leq\ |B|$$|A|\ \geq\
|B|$$0$$0$$N-1$$1$$1$$|A|\ \leq\ |B|$$|A|\ \geq\ |B|$$0$$0$$N-1$cumulative_sum
→ shift_downremember$N-1$$1$$1$$|A|\ \leq\ |B|$$|A|\ \geq\
|B|$$0$$0$$\frac{N}{2}-1$$N-1$$1$$1$$|A|\ \leq\ |B|$$|A|\ \geq\ |B|$$0$$0$
Figure 5: Schema for getting input-dependent shift kernels for the set
disjointness (SD) problem.

∎

We now use the kernels from Lemma G.20 to do the appropriate shift.

######  Proposition G.21.

Given a $\mathrm{JRT}$ prompt
${\bm{u}}^{\mathrm{JRT}}\in\mathbb{R}^{2N\times(n+1)}$ of the input
$\bm{u}\in\mathbb{R}^{N\times(n+1)}$ for the set-disjointness (SD) problem
$(A,B)\subseteq\\{0,1\\}^{n}$, there exist $O(1)$ input-dependent BaseConv
layers that can rearrange the input so that the smaller set out of $A$ and $B$
is placed at the start of the sequence.

###### Proof.

The input $\bm{u}\in\\{0,1\\}^{N\times(n+1)}$ is formatted as in Remark G.14.
In the first case where $A$ is the smaller set, we do not need to change the
input. Let ${\bm{s}}:=[\bm{0}^{n}::1]$ be the separator, then we want:

${\bm{u}}^{\mathrm{JRT}}\equiv\begin{bmatrix}\longleftarrow
A\longrightarrow&{\bm{s}}&\longleftarrow
B\longrightarrow&\vline&\longleftarrow
A\longrightarrow&{\bm{s}}&\longleftarrow B\longrightarrow\end{bmatrix}$

Otherwise, if $|B|\leq|A|$, we want to shift the input so that $|B|$ comes at
the start of the input sequence in the $\mathrm{JRT}$ prompt
${\bm{u}}^{\mathrm{JRT}}$. To this end, we want to add a separator between
after the first copy of the input ends. For this purpose, we can keep the
first copy as is, and operate on the duplicate $A$ by shifting it down by $1$
and adding a separator at the start of the second input. We thus apply the
remember(${\bm{u}}^{\mathrm{JRT}}_{{\tt shift\\_up}},N,N+|A|,f$) primitive [7,
Definition F.15] with $8$ layers of BaseConv where $f$ is any function that
maps $(A,{\bm{s}})\mapsto({\bm{s}},A)$, so that we get

${\bm{u}}^{\mathrm{JRT}}\equiv\begin{bmatrix}\longleftarrow
A\longrightarrow&{\bm{s}}&\longleftarrow
B\longrightarrow&\vline&{\bm{s}}&\longleftarrow
A\longrightarrow&\longleftarrow B\longrightarrow\end{bmatrix}$

Next, we shift up using the shift_up(${\bm{u}}^{\mathrm{JRT}},|A|+1$)
primitive [23, Proposition C.5] for BaseConv with $3$ layers by implementing
the kernel ${\bm{h}}_{\mathrm{shift}}$ from Lemma G.20. We then get

${\bm{u}}^{\mathrm{JRT}}_{{\tt shift\\_up}}\equiv\begin{bmatrix}\longleftarrow
B\longrightarrow&{\bm{s}}&\longleftarrow
A\longrightarrow&\vline&\longleftarrow
B\longrightarrow&\bm{0}^{|A|+1}\end{bmatrix}$

That is, in both cases, the final output has the smaller set out of $A$ and
$B$ at the start of the sequence.

To complete the proof we note that we can do the above in one single model
(that uses data dependent convolutions): (1) We add the extra separator after
the second set in ${\bm{u}}^{\mathrm{JRT}}$ and (2) we do the using the
convolution operator in BaseConv where we use the convolution kernel computed
from Lemma G.20.141414We also need to only keep the first $N$ rows of the
matrix, which we can obtain by zeroing out all the remaining rows using
another BaseConv layer. ∎

Finally, we can combine Propositions G.18 and G.21 to claim that Based can
solve $\mathrm{SD}$ with JRT-prompting in space $O(\min\\{|A|,|B|\\})$.

######  Theorem G.22.

Given a $\mathrm{JRT}$ prompt
${\bm{u}}^{\mathrm{JRT}}\in\mathbb{R}^{2N\times(n+1)}$ of the input
$\bm{u}\in\mathbb{R}^{N\times(n+1)}$ for the set-disjointness (SD) problem
$(A,B)$, there exists a (data dependent) Based model (BaseConv \+ MLP +
LinearAttention \+ MLP)151515This matches the architecture in our experiments.
that solves the SD problem with space $O(\min\\{|A|,|B|\\}\cdot n)$.

###### Proof.

First, we use the BaseConv layers from Proposition G.21 to get the smaller set
of $A$ and $B$ in ${\bm{u}}^{\mathrm{JRT}}$ to the start of the sequence in
$\bm{z}^{\texttt{BaseConv}}$. Next, we reduce $\bm{z}^{\texttt{BaseConv}}$
using an MLP layer to get $\bm{z}^{\texttt{BaseConv}}[0:N-1,:]$ as the input
to the LinearAttention (+MLP) layer in Proposition G.18 so that we solve the
SD problem for the original input $\bm{u}$. Finally, for the LinearAttention
layer, we can use the data-dependent IP kernels from equation 19 to get
$f=O(\min\\{|A|,|B|\\}))$, which yields the claimed space usage since we have
$d=n$. ∎

######  Remark G.23.

We note that we can use the random kernels from equation 20 in Theorem G.22 to
get space usage of $O\left((\min\\{|A|,|B|\\})^{2}\cdot n\right)$ without
using data-dependent IP kernels.

###  G.4 $\mathrm{GAR}$ and $\mathrm{SD}$

In this section, we introduce the general associative recall $\mathrm{GAR}$
problem. Recall that the query in the AR problem comes at the end, and thus,
the query is compared with all the keys in the input. On the other hand, in
MQAR, a query at position $i$ is only compared with keys at positions $j<i$.
Moreover, the number of keys and queries in the input are the same for MQAR.
Instead, we introduce the following alternate generalization of AR that has
all the queries at the end with the number of queries different from the
number of keys.

######  Definition G.24 ($\mathrm{GAR}$).

We are given an input sequence

$\bm{u}[0\cdots
N-1]\triangleq({\bm{k}}_{0},{\bm{v}}_{0}),\ldots,({\bm{k}}_{n-1},{\bm{v}}_{n-1});{\bm{q}}_{0},\ldots,{\bm{q}}_{m-1},$
(22)

where $K:=\\{{\bm{k}}_{i}\\}_{i=0}^{n-1},V:=\\{{\bm{v}}_{i}\\}_{i=0}^{n-1},$
and $Q:=\\{{\bm{q}}_{i}\\}_{i=0}^{m-1}$, with each
${\bm{k}}_{i},{\bm{v}}_{i},{\bm{q}}_{i}\in C$ is a token drawn from a
vocabulary of size $c=\left\lvert C\right\rvert$, and we have $N=2n+m$.

Our goal in the general associative recall ($\mathrm{GAR}$) problem is to
check, for each $\bm{q}_{i}\in Q$, whether there exists $\bm{k}_{j}\in K$ such
that $\bm{q}_{i}\equiv\bm{k_{j}}$; if so, output the corresponding value
${\bm{v}}_{j}$, and otherwise, output ${\tt Null}$.

We will first show that SD reduces to $\mathrm{GAR}$.

######  Proposition G.25.

Any algorithm $\mathcal{A}$ solving $\mathrm{GAR}$ can also solve
$\mathrm{SD}$.

###### Proof.

Given an input to the set-disjointness problem $(A,B)$ with
$A:=\\{A_{0},\ldots,A_{|A|-1}\\},B:=\\{B_{0},\ldots,B_{|B|-1}\\}$, we can
construct the following input to the $\mathrm{GAR}$ problem:

$\bm{u}:=(A_{0},A_{0}),\ldots,(A_{|A|-1},A_{|A|-1});B_{0},\ldots,B_{|B|-1}.$

Now, we run algorithm $\mathcal{A}$ on $\bm{u}$, and if for all $q\in Q$, we
get ${\tt Null}$, then we know $A\cap B=\emptyset$, and otherwise, $A\cap
B\neq\emptyset$. This solves the set disjointness (SD) problem. ∎

What we have shown is that $\mathrm{GAR}$ is much more general compared to
$\mathrm{SD}$. However, we can also show that we can solve $\mathrm{GAR}$
under certain conditions if we had access to an algorithm solving
$\mathrm{SD}$.

######  Proposition G.26.

Let $\mathcal{A}_{\mathrm{SD}}$ be an algorithm solving the set disjointness
($\mathrm{SD}$) problem. Then, for a vocabulary $\mathcal{C}$ with
$\left\lvert\mathcal{C}\right\rvert=c$ with values from $[c]$ and represented
as $v_{j}\in\\{0,1\\}^{d}$ where $d=\lceil\log_{2}(c+1)\rceil$ with at most
one match for each query, we can solve the $\mathrm{GAR}$ problem (definition
G.24) with $d$ calls to $\mathcal{A}_{\mathrm{SD}}$.

###### Proof.

Given an input
$({\bm{k}}_{0},{\bm{v}}_{0}),\ldots,({\bm{k}}_{n-1},{\bm{v}}_{n-1});{\bm{q}}_{0},\ldots,{\bm{q}}_{m-1}$
to $\mathrm{GAR}$, for each call $\ell\in[d]$ to algorithm
$\mathcal{A}_{\mathrm{SD}}$, we construct the inputs to algorithm
$\mathcal{A}$ by taking $A:=Q,B:=K_{\ell}$ with $K_{\ell}$ defined as follows:

$k_{j}\in K_{\ell}\iff v_{j}[\ell]=1.$ (23)

That is, we include $k_{j}\in K_{\ell}$ iff the $\ell$’th bit of $v_{j}$ is 1.

We now claim that we can solve the MQAR problem given $Q\cap K_{\ell}$ for all
$\ell\in[d]$. To see this, note that if a query $q\in Q$ is not in $K$, then
$q\notin Q\cap K_{\ell}$ for every $\ell\in[d]$. We thus output ${\tt Null}$
for these queries.

Otherwise, if $q\in Q\cap K$, then there exists a non-empty set of calls
$L\subseteq[d]$ such that $q\in Q\cap K_{\ell}$ for all $\ell\in L$. We can
then extract the $\ell$’th bit of $v_{j}$, where $q=k_{j}$. That is, for
$q=k_{j}$, we use equation 23 to get

$v_{j}[\ell]=\begin{cases}1&\text{if }\ell\in L\\\
0&\text{otherwise.}\end{cases}$

This is exactly the value corresponding to the unique matching key $k_{j}$ for
the query $q$. ∎

####  G.4.1 Lower Bound for $\mathrm{GAR}$ via $\mathrm{SD}$

In this section, we present a lower bound for solving $\mathrm{GAR}$. For this
purpose, we require the following two-way randomized communication
complexity161616Here, in contrast to one-way randomized communication protocol
in section G.1.1, both Alice and Bob are allowed to send messages to each
other. lower bound for set-disjointness ($\mathrm{SD}$).

######  Theorem G.27 ([83]171717[83] provides a lower bound of $n$ for
$|A|=|B|$. However, we can extend it to Theorem G.27 by reducing the
$\min\\{|A|,|B|\\}$ subset to the equal sized set by picking a hard
distribution where both sets are of size $\min\\{|A|,|B|\\}$ and then adding
”extra” elements to only one of them to get a larger set (i.e., one can
increase the universe size by these extra elements to get the desired lower
bound). ).

The two-way randomized communication complexity of the set disjointness
problem with sets $A,B\subseteq[n]$ is $\Omega(\min\\{|A|,|B|\\})$ bits for
$n\geq o(\min\\{|A|,|B|\\})$.

######  Definition G.28 ($\mathrm{JR-}\text{$p$}$ Prompts).

For any model $\mathcal{M}$ with input $\bm{u}\in\mathbb{R}^{N\times d}$, a
$\mathrm{JR-}\text{$p$}$ prompt for input $\bm{u}$ is the $p$-times repeated
input $\bm{u}^{\mathrm{JR-}\text{$p$}}\in{\mathbb{R}^{pN\times d}}$ given by

$\bm{u}^{\mathrm{JR-}\text{$p$}}[i,:]:=\bm{u}[i\mod{N},:]$

######  Proposition G.29.

Given a $\mathrm{JR-}\text{$p$}$ prompt
${\bm{u}}^{\mathrm{JR-}\text{$p$}}\in\\{0,1\\}^{pN\times d}$ for input
${\bm{u}}\in\\{0,1\\}^{N\times d}$ to the $\mathrm{GAR}$ problem, any
recurrent model $\mathcal{M}_{\mathrm{GAR}}$ (definition G.12) solving
$\mathrm{GAR}$ requires
$\max_{i}\left\lvert{\bm{Z}}_{\mathcal{M}_{\mathrm{GAR}}}^{i}\right\rvert$ to
be at least $\Omega\left(\frac{\min\\{|A|,|B|\\}}{p}\right)$-bits.

###### Proof.

We first take the input $\bm{u}\in\\{0,1\\}^{N\times d}$ to the $\mathrm{GAR}$
problem and design a two-way communication protocol for solving
${\mathrm{GAR}}$ given access to the reccurrent model
$\mathcal{M}_{\mathrm{GAR}}$. To this end, Alice with their access of key-
value part generates her part of the input:

${\bm{u}}_{\mathrm{Alice}}:=(k_{0},v_{0}),\ldots,(k_{n-1},v_{n-1})$ (24)

of the input for $\mathrm{GAR}$ (without the queries), and Bob with their
access of the query part generates the following;

${\bm{u}}_{\mathrm{Bob}}:=q_{0},\ldots,q_{m-1}$ (25)

of the input for $\mathrm{GAR}$ (without the key-value pairs) as in equation
22. That is, the concatenation
${\bm{u}}_{\mathrm{Alice}}::{\bm{u}}_{\mathrm{Bob}}\equiv{\bm{u}}$ in equation
22. We then have

$\underbrace{{\bm{u}}_{\mathrm{Alice}}::{\bm{u}}_{\mathrm{Bob}}::\cdots::{\bm{u}}_{\mathrm{Alice}}::{\bm{u}}_{\mathrm{Bob}}}_{p-\mathrm{times}}\equiv{\bm{u}}^{\mathrm{JR-}\text{$p$}},$
(26)

the corresponding $\mathrm{JR-}\text{$p$}$ prompt for the input ${\bm{u}}$ to
the $\mathrm{GAR}$ problem. We now claim that the following protocol
(algorithm 4) is equivalent to running the recurrent model
$\mathcal{M}_{\mathrm{GAR}}$ on the $\mathrm{JR-}\text{$p$}$ prompt
${\bm{u}}^{\mathrm{JR-}\text{$p$}}$:

Algorithm 4 Communication Protocol for $\mathrm{GAR}$

1: A recurrent model $\mathcal{M}_{\mathrm{GAR}}$ solving $\mathrm{GAR}$ along
with the inputs ${\bm{u}}_{\mathrm{Alice}},{\bm{u}}_{\mathrm{Bob}}$ from 24
and 25.

2: $\mathcal{M}_{\mathrm{GAR}}({\bm{u}}^{\mathrm{JR-}\text{$p$}})$.

3:for $i\leftarrow$ $0$ to $p-1$ do

4: for $j\leftarrow$ $0$ to $2n-1$ do

5: ${\bm{Z}}_{\mathcal{M}_{\mathrm{GAR}}}^{i\cdot N+j}\leftarrow
f_{\mathcal{M}}^{i\cdot N+j}({\bm{Z}}_{\mathcal{M}_{\mathrm{GAR}}}^{i\cdot
N+j-1},{\bm{u}}_{\mathrm{Alice}}[j])$ $\triangleright$
${\bm{Z}}_{\mathcal{M}_{\mathrm{GAR}}}^{0}\leftarrow{\bm{u}}_{\mathrm{Alice}}[0,:]$

6: Alice sends ${\bm{Z}}_{\mathcal{M}_{\mathrm{GAR}}}^{i\cdot N+2n-1}$ to Bob

7: for $j\leftarrow$ $0$ to $m-1$ do

8: ${\bm{Z}}_{\mathcal{M}_{\mathrm{GAR}}}^{i\cdot N+2n+j}\leftarrow
f_{\mathcal{M}}^{i\cdot N+j}({\bm{Z}}_{\mathcal{M}_{\mathrm{GAR}}}^{i\cdot
N+2n+j-1},{\bm{u}}_{\mathrm{Bob}}[j])$

9: Bob sends ${\bm{Z}}_{\mathcal{M}_{\mathrm{GAR}}}^{i\cdot N+m-1}$ to Alice

The equivalency of this protocol with running the model
$\mathcal{M}_{\mathrm{GAR}}$ follows from equation 26.

Next, consider an instance $\bm{u}^{\mathrm{SD}}:=(A,B)$ of the set-
disjointness problem with $A,B\subseteq[n]$ and $|A|+|B|=N$, where
$A:=\\{A_{0},\ldots,A_{|A|-1}\\},B:=\\{B_{0},\ldots,B_{|B|-1}\\}$. Due to
proposition G.25, we know that we can generate an equivalent input $\bm{u}$
for $\mathrm{GAR}$ given an input $\bm{u}^{\mathrm{SD}}$ to the $\mathrm{SD}$
problem, whence we can generate inputs for Alice and Bob as in equation 24 and
equation 25. Applying algorithm 4 then solves the $\mathrm{GAR}$ problem for
$\bm{u}$, and consequently, the SD problem for $\bm{u}^{\mathrm{SD}}$. Here,
the total number of bits that are communicated in this protocol is

$T_{\mathrm{bits}}:=\sum_{i=0}^{p-1}\left\lvert{\bm{Z}}_{\mathcal{M}_{\mathrm{GAR}}}^{i\cdot
N+2n-1}\right\rvert+\left\lvert{\bm{Z}}_{\mathcal{M}_{\mathrm{GAR}}}^{i\cdot
N+m-1}\right\rvert.$

Now, if $T_{\mathrm{bits}}$ is $o(\min\\{|A|,|B|\\})$ bits, we have shown that
a two-way communication protocol exists for solving the set-disjointness
($\mathrm{SD}$) that uses $o(\min\\{|A|,|B|\\})$ communication complexity.
However, this contradicts theorem G.27. Thus, we have
$T_{\mathrm{bits}}\geq\Omega\left(\min\\{|A|,|B|\\}\right).$

Finally, note that we have

$\displaystyle p\cdot
2\max_{k}\left\lvert{\bm{Z}}_{\mathcal{M}_{\mathrm{GAR}}}^{k}\right\rvert$
$\displaystyle=\sum_{i=0}^{p-1}2\max_{k}\left\lvert{\bm{Z}}_{\mathcal{M}_{\mathrm{GAR}}}^{k}\right\rvert$
$\displaystyle\geq\sum_{i=0}^{p-1}\left\lvert{\bm{Z}}_{\mathcal{M}_{\mathrm{GAR}}}^{i\cdot
N+2n-1}\right\rvert+\left\lvert{\bm{Z}}_{\mathcal{M}_{\mathrm{GAR}}}^{i\cdot
N+m-1}\right\rvert$ $\displaystyle\geq\Omega\left(\min\\{|A|,|B|\\}\right).$
$\displaystyle\implies\max_{k}\left\lvert{\bm{Z}}_{\mathcal{M}_{\mathrm{GAR}}}^{k}\right\rvert$
$\displaystyle\geq\Omega\left(\frac{\min\\{|A|,|B|\\}}{2p}\right).$

This concludes the proof. ∎

Table 15: JRT-RNN Training Settings. For hybridizing the three layer types – gated convolutions, sliding window, and linear attention – we use linear attention at layers $\\{2,7,12,17,22,27,32\\}$ and sliding window at layers $\\{3,8,13,18,23,28,33\\}$, with gated convolution layers elsewhere. We did not tune the layer orderings and proportions. | 356M | 1.3B
---|---|---
Optimizer | Adam
Optimizer momentum | $\beta_{1},\beta_{2}=0.9,0.95$
Optimizer eps | $1e-8$
Precision | BFloat16
Encoder region length | 1024
Masked language modeling probability | 15%
MLM loss scale | 0.25
NTP loss scale | 1.00
Warmup | 1%
Learning rate decay | Cosine
Learning rate (min, base) | 8e-5, 8e-4
Global batch size | 256
Weight decay | 0.1
Num Layers | 26 | 36
Hidden Size | 1024 | 1792
MLP Activation | SwiGLU
MLP Width | 2
Num. Linear Attn Layers | 5 | 7
Num. Linear Attn Heads | 16
Taylor Feature Dimension | 16
Linear Attn Positional Encodings | None
Num. Sliding Window Layers | 5 | 7
Sliding Window Size | 64 | 16
Sliding Window Heads | 16
Sliding Window Positional Encodings | Rotary
Num. BaseConv Layers | 17 | 22
BaseConv Projection Expansion Factor | 4
BaseConv Filter Size | 3
BaseConv Activation | SiLU
Table 16: Based Training Settings. For hybridizing the three layer types – gated convolutions, sliding window, and linear attention – we use linear attention at layers $\\{2,7,12,17,22,27,32\\}$ and sliding window at layers $\\{3,8,13,18,23,28,33\\}$, with gated convolution layers elsewhere. We did not tune the layer orderings and proportions. | 363M | 1.4B
---|---|---
Optimizer | Adam
Optimizer momentum | $\beta_{1},\beta_{2}=0.9,0.95$
Optimizer eps | $1e-8$
Precision | BFloat16
Warmup | 1%
Learning rate decay | Cosine
Learning rate (min, base) | 8e-5, 8e-4
Global batch size | 256
Weight decay | 0.1
Num Layers | 27 | 36
Hidden Size | 1024 | 1792
MLP Activation | SwiGLU
MLP Width | 2
Num. Linear Attn Layers | 5 | 7
Num. Linear Attn Heads | 16
Taylor Feature Dimension | 16
Linear Attn Positional Encodings | None
Num. Sliding Window Layers | 5 | 7
Sliding Window Size | 128
Sliding Window Heads | 16
Sliding Window Positional Encodings | Rotary
Num. BaseConv Layers | 17 | 22
BaseConv Projection Expansion Factor | 4
BaseConv Filter Size | 3
BaseConv Activation | SiLU
Table 17: Mamba Training Settings | 358M | 1.3B
---|---|---
Optimizer | Adam
Optimizer momentum | $\beta_{1},\beta_{2}=0.9,0.95$
Optimizer eps | $1e-8$
Precision | BFloat16
Warmup | 1%
Learning rate decay | Cosine
Learning rate (min, base) | 8e-5, 8e-4
Global batch size | 256
Weight decay | 0.1
Num Layers | 46
Hidden Size | 1024 | 2048
RMSNorm | True
Norm Epsilon | $1e-5$
Dt State | $16$
Dt (Min, Max) | $(0.001,0.1)$
Dt Init. Strategy | Random
Dt Init. Floor | $1e-4$
Dt Scale | $1.0$
Dt Softplus | True
Projection Expansion Factor | 2
Short Conv Filter Size | 4
Table 18: Attention Training Settings | 360M | 1.3B
---|---|---
Optimizer | Adam
Optimizer momentum | $\beta_{1},\beta_{2}=0.9,0.95$
Optimizer eps | $1e-8$
Precision | BFloat16
Warmup | 1%
Learning rate decay | Cosine
Learning rate (min, base) | 8e-5, 8e-4
Global batch size | 256
Weight decay | 0.1
Num Layers | 24 | 36
Hidden Size | 1024 | 1680
Num Heads | 16 | 24
RMSNorm | True
MLP Bias | False
Flash Attn | True
Rotary Emb. Fraction | 0.5
MLP Activation | SwiGLU
MLP Width | 4