source/authored/appendix/constructions.tex

\section{Omitted Proofs in \Cref{sec merge}}
\label{app:constructions}


\subsection{Proof of \Cref{th lb selective copying}}\label{app proof lb copy}



\begin{theorem}[Restatement of \Cref{th lb selective copying}]
Consider the task of selective copying. There is a distribution $D$ over $\cV^L$ such that any pure state space model that can solve the task for some $\vec{\x}$ drawn from $D$ with probability $90\%$ must have $\sum_{i=1}^k\log(\card{\mathcal{S}_i}) \ge N \log M$, where $\mathcal{S}_i$ is the state space of the $i$th SSM layer, furthermore, any pure Transformer model that can solve the task with probability $90\%$ must have $\sum_{i=1}^kW_i \ge \Omega(L)$.
\end{theorem}

\begin{proof}[Proof of \Cref{th lb selective copying}]
We write down the task in the form of a function composition $F(u(\vec{\x}),v(\vec{\x}))$. Let $u(\vec{\x}) = \vec{\x}_{L-N+1: L}$ and $v(\vec{\x}):=\argmax_{1 \le i \le L} x_i \in \cN$ and $F(u,v)$ be $u_{L+1-v}$. 

We first show that $F(u,v)$ satisfies \Cref{asp ssm lb}. We take $v^{(i)} = i, i \in [N]$, which implies $F(u,v^{(i)}) = u_i, i \in [N]$. Thus, $G(u)=(F(u,v^{(1)}),\dots,F(u,v^{(N)})) = u$ is an injection. Taking $m=q=N$, by \Cref{th SSM lb general}, 
we know that when a pure SSM gives a sequence of input of the form $(u,v)$, where $u$ is drawn uniformly from $\cV$ and $v$ is drawn uniformly from $\cN$, then if the SSM wants to solve the task with probability at least $7/8$, it needs 
\begin{align*}
\sum_{i=1}^k \log(\card{\mathcal{S}_i)} \ge \Omega(N \log(\card{M})),   
\end{align*}
when a pure SSM is given $(u,v)$ such that $u$ is uniformly drawn from $\cV^m$ and $v$ is drawn uniformly from $\cN$. To simulate such a distribution over $(u,v)$ using a long context $\vec{\x}$, we select the first $L-1$ tokens from $\cV$ uniformly at random, while selecting the last token as a random token from $\{2,\dots,N\}$. Notice that such a distribution $D_S$ is sufficient to simulate the required distribution of $(u,v)$, since the effective part of $u$ is $u_{1:N-1}$ and $u$ is $u_{1:N-1}$ is independent on $v$.



We next show that $F(u(x),v(x))$ is $L/2$ sensitive. Let $\vec{\x},\vec{x'} \in \cV^L$ be any input context such that $\vec{\x}_{L/2:L} = \vec{x'}_{L/2:L}$ and $\vec{\x}_i \not\in \cN$, for every $i=L/2,\dots,L$. This implies $v(\vec{\x})$ and $v(\vec{x'})$ only depends on the first $L/2$ coordinates. Thus, $F(u(\vec{\x}),v(\vec{\x}))$ is $L/2$ sensitive. By \Cref{th Transformer lb}, we know that any pure Transformer model that can compute $F(u(\vec{\x}),v(\vec{\x}))$ with a constant probability must have $\sum_{i=1}^k W_i \ge \Omega(L).$ In particular, to construct a distribution $D_T$ over $\vec{\x}$ that makes a pure Transformer fail, we draw $\vec{\x}$ such that $\vec{\x}_{L/2:L}$ is chosen uniformly from $\cM$ and $\vec{\x}_{L/2:L}$ is chosen uniformly from $\cV$.

To conclude the proof of \Cref{th lb selective copying}, we will choose a distribution $D:= D_S/2 + D_T/2$. Under this distribution, a pure SSM with $\sum_{i=1}^k \log(\card{\mathcal{S}_i)} < \Omega(N \log(\card{M}))$ has a constant failure probability when $\vec{\x}$ is drawn from $D_S$ and a Transformer with working memory less than $o(L)$ has a constant failure probability when $\vec{\x}$ drawn from $D_T$.
 \end{proof}







\subsection{Proof of \Cref{th hybrid copy}}\label{app proof hybrid copy}

\begin{theorem}[Restatement of \Cref{th hybrid copy}]
Consider the task of selective copying. There is a two-layer hybrid model that is a combination of a mamba layer and an attention layer that can solve the selective copying task for \emph{every} input sequence $\vec{\x} \in \cV^L$. Furthermore, the hybrid model has an embedding dimension $d = O( \max(\log \card{\cV},\log L))$ such that the Mamba layer has $O(\card{\cV})$ state spaces, while the attention layer has dimension $d$ and a sliding window size of $O(N)$.
\end{theorem}


\begin{proof}[Proof of \Cref{th hybrid copy}]
We first construct the two-layer hybrid model and show the correctness of the construction.

\noindent \textbf{Token Embeddings.}
We begin by constructing an embedding function that maps each $\x \in \cV$ to a vector in $\R^{d'}$ for some $d>0$.
Let $d' = \log \card{\cV}$. We define $\psi': \cV \to \{\pm1\}^{d'}$ to be the binary encoding of the vocabulary $\cV$. For each $\x \in \cV$, we embed the token $\x$ as
\[\psi(\x) = \begin{pmatrix}
    \psi'(\x) & \One_{\{\x \in \calN\}}\psi'(\x) & \mathbf{0}
\end{pmatrix}^\top.\]
Here $\mathbf{0} \in \R^{\ell}$ for some $\ell \le O(\log(\card{\cV})+\log L)$ is a zero vector. For a given input context $\vec{\x}$, the embedded context has the form of 
\begin{equation*}
    \begin{pmatrix}
        \psi'(\x_1) & \psi'(\x_2) & \dots & \psi'(\x_{L-1}) & \psi'(\x_L) \\
        \One_{\{\x_1 \in \calN\}}\psi'(\x_1) & \One_{\{\x_2 \in \calN\}}\psi'(\x_2) & \dots & \One_{\{\x_{L-1} \in \calN\}}\psi'(\x_{L-1}) & \One_{\{\x_L \in \calN\}}\psi'(\x_{L}) \\
        \mathbf{0} & \mathbf{0} & \dots & \mathbf{0} & \mathbf{0}
    \end{pmatrix}.
\end{equation*}
\noindent \textbf{Position Encoding.}
To allow the attention layer to access the position of the input tokens, we next add a position encoding to each input token. Define $\phi': [L] \to \{\pm 1\}^{\log L}$ be the binary encoding of numbers in $L$. For each $i \in [L]$, we encode the position $i$ using a function $\phi$ defined as follows
\begin{align*}
     \phi(i) =  \phi'(L+1-i). 
\end{align*}
After adding the position encoding, the input context has the following form. 
\begin{align*}
    \Phi(\vec{\x}):=\begin{pmatrix}
        \psi'(\x_1) &  \dots & \psi'(\x_{L-1}) & \psi'(\x_L) \\
        \One_{\{\x_1 \in \calN\}}\psi'(\x_1)  & \dots & \One_{\{\x_{L-1} \in \calN\}}\psi'(\x_{L-1}) & \One_{\{\x_L \in \calN\}}\psi'(\x_{L}) \\
        \mathbf{0} &  \dots & \mathbf{0} & \mathbf{0} \\
        \phi(1)  & \dots & \phi(L-1) & \phi(L)
    \end{pmatrix}
\end{align*}
By construction, each column of the input context is in $\R^d$, for some $d = O(\log(\card{\cV}+\log L)$. 
Given the embedding of the input context, we now construct the SSM layer and the attention layer of the hybrid model.


\noindent \textbf{Mamba Layer.} 
We define the weights of the Mamba layer as follows. 
Let $\W_A,\W_B,\W_C$ be as follows, 
\begin{align*}
    \W_B 
    \begin{pmatrix}\psi'(\x_i)\\\One_{\{\x_i \in \calN\}}\psi'(\x_i)\\\mathbf{0}\\\phi'(i)\end{pmatrix} 
    &= \One_{\{\x_i \in \calN\}}\psi'(\x_i)
    , \quad
    \W_C \bv = \begin{pmatrix}\mathbf{0} \\ \mathbf{0} \\ \bv \\ \mathbf{0}\end{pmatrix}, 
\end{align*}




$\W_A = I$, and $\Delta(\x) = \Ind_{\{\x \in \cN\}}$. We claim the following guarantee for the constructed Mamba layer.
\begin{claim}\label{cl mamba2}
    Let $\vec{\x} \in \cV^L$ be a sequence of input contexts. Given the embedded context $\Phi(\vec{\x})$, the SSM layer with parameter $\W_A,\W_B,\W_C,\Delta(\x)$, has the following output
      \begin{align*}
        \begin{pmatrix}
        \psi'(\x_1) & \psi'(\x_2) & \dots & \psi'(\x_{L-1}) & \psi'(\x_L) \\
        \One_{\{\x_1 \in \calN\}}\psi'(\x_1) & \One_{\{\x_2 \in \calN\}}\psi'(\x_2) & \dots & \One_{\{\x_{L-1} \in \calN\}}\psi'(\x_{L-1}) & \One_{\{\x_L \in \calN\}}\psi'(\x_{L}) \\
        \phi({L+1-n_1}) & \phi({L+1-n_2}) & \dots & \phi({L+1-n_{L-1}}) & \phi({L+1-n_L}) \\
        \mathbf{0} & \mathbf{0} & \dots & \mathbf{0} & \mathbf{0} \\
        \phi(1) & \phi(2) & \dots & \phi(L-1) & \phi(L)
    \end{pmatrix},
    \end{align*}
where $n_i = \x_{\argmax_{1 \leq j \leq i} \x_j \in \calN}$
\end{claim}

\begin{proof}[Proof of \Cref{cl mamba}]

By our choice of $\Delta(\x_t)$, if $\x_t \in \N$, then $H_t = \W_B\Phi(\x_t)$ and if $\x_t \not \in \N$, then $H_t = H_{t-1}$. Notice that if we set $H_0$ to be a zero vector, then by induction, for each $t$, $H_t$ is a sparse vector, with the only non-zero component $\phi(L+1-n_i) = \psi'(n_i)$. Using an MLP layer to combine the output with the input, we know that the input context 
\begin{align*}
        \begin{pmatrix}
        \psi'(\x_1) & \psi'(\x_2) & \dots & \psi'(\x_{L-1}) & \psi'(\x_L) \\
        \One_{\{\x_1 \in \calN\}}\psi'(\x_1) & \One_{\{\x_2 \in \calN\}}\psi'(\x_2) & \dots & \One_{\{\x_{L-1} \in \calN\}}\psi'(\x_{L-1}) & \One_{\{\x_L \in \calN\}}\psi'(\x_{L}) \\
        \phi(L+1-{n_1}) & \phi(L+1-{n_2}) & \dots & \phi(L+1-{n_{L-1}}) & \phi(L+1-{n_L}) \\
        \mathbf{0} & \mathbf{0} & \dots & \mathbf{0} & \mathbf{0} \\
        \phi(1) & \phi(2) & \dots & \phi(L-1) & \phi(L)
    \end{pmatrix}.
    \end{align*}
\end{proof}

\noindent \textbf{Attention Layer.} Based on the output of the Mamba layer, we will now construct an attention layer that can solve the copy task. Let $\W_q,\W_k,\W_v$ be the weight matrices of the attention layer.
\begin{align*}
    \W_q \begin{pmatrix}\psi'(\x_i)
    % \\\One_{\{\x_i \in \calN\}}
    \\\One_{\{\x_i \in \calN\}}\psi'(L+1-\x_i) 
    \\\phi'({n_i})\\\mathbf{0}\\\phi'(i)\end{pmatrix} &= M\phi({n_i}) ,\quad
    \W_k \begin{pmatrix}\psi'(\x_i)
    % \\\One_{\{\x_i \in \calN\}}
    \\\One_{\{\x_i \in \calN\}}\psi'(L+1-\x_i) \\\phi'({n_i})\\\mathbf{0}\\\phi'(i)\end{pmatrix} = \phi(i) \\
    \W_v \begin{pmatrix}\psi'(\x_i)
    % \\\One_{\{\x_i \in \calN\}}
    \\\One_{\{\x_i \in \calN\}}\psi'(L+1-\x_i) \\\phi'({n_i})\\\mathbf{0}\\\phi'(i)\end{pmatrix} &= \begin{pmatrix}\mathbf{0}\\0\\\mathbf{0}\\\psi'(\x_i)\\\mathbf{0}\end{pmatrix} \\
\end{align*}
We summarize the performance of the attention layer as the following claim.
\begin{claim}\label{cl tf copy}
    Let $\SSM(\vec{\x})$ be the output of the Mamba layer constructed above. By applying an attention mechanism with parameters $\W_q,\W_k,\W_v$ with a window size of $N$, the last output vector is a sparse vector with the only non-zero part $\psi(\x_{L+1-n_L})$.
\end{claim}
\begin{proof}[Proof of \Cref{cl tf copy}]
    Notice that the last output vector is defined as 
    \begin{align*}
        \sum_{i=L+1-N}^L \frac{\exp(M\phi(n_L)\phi(i))}{\sum_{i=L+1-N}^L \exp(M\phi(n_L)\phi(i))}\begin{pmatrix}\mathbf{0}\\0
        % \\\mathbf{0}
        \\\mathbf{0}\\\psi'(\x_i)\\\mathbf{0}\end{pmatrix} = \begin{pmatrix}\mathbf{0}\\0
        % \\\mathbf{0}
        \\\mathbf{0}\\\psi'(\x_{L+1-n_L})\\\mathbf{0}\end{pmatrix}
    \end{align*}
\end{proof}
In fact, if we use a full attention, then after passing the attention layer, the context now has the form
\begin{align*}
        \begin{pmatrix}
        \psi'(\x_1) & \psi'(\x_2) & \dots & \psi'(\x_{L-1}) & \psi'(\x_L) \\
        \One_{\{\x_1 \in \calN\}}\psi'(\x_1) & \One_{\{\x_2 \in \calN\}}\psi'(\x_2) & \dots & \One_{\{\x_{L-1} \in \calN\}}\psi'(\x_{L-1}) & \One_{\{\x_L \in \calN\}}\psi'(\x_{L}) \\
        \phi({n_1}) & \phi({n_2}) & \dots & \phi({n_{L-1}}) & \phi({n_L}) \\
        \psi'(x_{L+1-n_1}) & \psi'(x_{L+1-n_2}) & \dots & \psi'(x_{L+1-n_{L-1}}) & \psi'(x_{L+1-n_{L}}) \\
        \phi(1) & \phi(2) & \dots & \phi(L-1) & \phi(L)
    \end{pmatrix}.
    \end{align*}
This implies that by applying the natural decoding that copies only the row with $\psi'(\x_{L+1-n_1})$. The final sequence (up to arbitrarily small error) has the form
\textbf{\begin{equation*}
    \begin{bmatrix}
        \psi'(\x_{L+1-n_1}) & \psi'(\x_{L+1-n_2}) & \dots & \psi'(\x_{L+1-n_{L-1}}) & \psi'(\x_{L+1-n_L})
    \end{bmatrix}
\end{equation*}}

To conclude the proof of \Cref{th hybrid copy}, it remains to count the number of parameters and the working memory of the constructed hybrid model. 
\end{proof}

\subsection{Proof of \Cref{th lb associate recall}}\label{app proof lb recall}


\begin{theorem}[Restatement of \Cref{th lb associate recall}]
    Consider the task of associative recall with decoding. 
    There is a distribution $D$ over $\cV^L$ such that any pure state space model that can solve the task for some $\vec{\x}$ drawn from $D$ with probability $90\%$ must have $\sum_{i=1}^k\log(\card{\mathcal{S}_i}) \ge \Omega (W \log W)$, where $\mathcal{S}_i$ is the state space of the $i$th SSM layer and $W = \card{\cM}$, furthermore, any pure Transformer model that can solve the task with probability $90\%$ must have $\sum_{i=1}^kW_i \ge \Omega(L)$.
\end{theorem}

\begin{proof}[Proof of \Cref{th lb associate recall}]
    To prove the hardness of the task, we will consider a constraint version of the problem. To do this, we partition equally partition the vocabulary $\cM$ into $\cM_1 = \{\alpha_1,\dots,\alpha_{W/2}\} ,\cM_2 = \{\beta_{1},\dots,\beta_{W/2}\}$. For every possible input $\vec{\x}$, we restrict $\vec{\x}$ of the following form. $\vec{\x} = (\alpha_{i1},\beta_{i1},b_{i1},\dots,\alpha_{ik},\beta_{ik},b_{ik})\in \cV^L$. Here, for $j \in [l]$, $b_{ij}$ either does not appear or $b_{ij} \in \{0,1\}$. Now, we partition $\vec{\x}$ into two parts. Let $\vec{b}=(b_{i1},\dots,b_{ik})$ be the 0-1 subsequence of $\vec{\x}$ and $\vec{w} = (\alpha_{i1},\beta_{i1},b_{i1},\dots,\alpha_{ik},\beta_{ik},b_{ik})$ be subsequence of $\vec{\x}$ with every token in $\cM$. Given this partition, we write the problem as a function composition. For each $\alpha \in \cM_1$, let $\beta(\alpha) \in \cM_2$ be the next token of the last appearance of $\alpha$ in $\vec{w}$. Let $u = (\beta(\alpha))_{\alpha \in \cM_1} \in \cM_2^k$ and let $v \in \cM_1$ be the token corresponding to binary representation $\vec{b}$ of elements in $\cM_1$. Then $F(u,v) = \beta(v)$. We notice that by choosing $G(u)=(\beta(\alpha_1),\dots,\beta(\alpha_{W/2}))$ is an injection. 
    Taking $m=q=W/2$, by \Cref{th SSM lb general}, 
we know that when a pure SSM gives a sequence of input of the form $(u,v)$, where $u$ is uniformly drawn from $\cM_2^{M/2}$ and $v$ is drawn uniformly from $\cM_1$, then if the SSM wants to solve the task with probability at least $7/8$, it needs 
\begin{align*}
\sum_{i=1}^k \log(\card{\mathcal{S}_i)} \ge \Omega(W \log(\card{W})),   
\end{align*}

To simulate such a distribution over $(u,v)$ using a distribution $D_S$ over a long context $\vec{\x}$, we select each $(\alpha_{ij},\beta_{ij})$ uniformly at random and after selecting $\vec{w}$, we randomly select $\vec{b}$ and append it to $\vec{w}$.


    On the other hand, 
    by sampling $\vec{w}$ uniformly from, we know that with probability at least $99\%$, for each $\alpha \in \cM_1$, the last appearance of $\alpha$ must be at the last $\Tilde{O}(W)$ positions of $\vec{w}$. By \Cref{th Transformer lb}, we know that any pure Transformer model that can compute $F(u(\vec{\x}),v(\vec{\x}))$ with a constant probability must have $\sum_{i=1}^k W_i \ge \Omega(L)$ if we randomly selecting $\vec{b}$ and appending it before $\vec{w}$. we denote the resulting distribution by $D_T$.
    
    
To conclude the proof of \Cref{th lb associate recall}, we will choose a distribution $D:= D_S/2 + D_T/2$. Under this distribution, a pure SSM with $\sum_{i=1}^k \log(\card{\mathcal{S}_i)} < \Omega(W \log(W))$ has a constant failure probability when $\vec{\x}$ is drawn from $D_S$ and a Transformer with working memory less than $o(L)$ has a constant failure probability when $\vec{\x}$ drawn from $D_T$.

    
\end{proof}

\subsection{Proof of \Cref{th hybrid recall}}\label{app proof hybrid recall}


\begin{theorem}{Restatement of \Cref{th hybrid recall}}
Consider the task of associative recall with decoding. There is a three-layer hybrid model that is a combination of a mamba layer and two attention layers that can solve the associative recall with decoding task with probability $99\%$ for an input sequence $\vec{\x} \in \cV^L$ drawn from a uniform distribution. Furthermore, the hybrid model has an embedding dimension $d = O( \max(\log \card{\cV},\log L))$ such that the Mamba layer has $O(\card{\cV})$ state spaces, while the attention layer has dimension $d$ and a sliding window size of $\Tilde{O}(\card{\cV})$.    
\end{theorem}


\begin{proof}[Proof of \Cref{th hybrid recall}]
    We first construct the three-layer hybrid model and show the correctness of the construction. 

\noindent \textbf{Token Embeddings}
Let $\cB = \{0,1\}$ be the set containing the two bits in the vocabulary. We begin by constructing an embedding function that maps each $\x \in \cV$ to a vector in $\R^{d'}$ for some $d>0$.
Let $d' = \log \card{\cV}$. We define $\psi': \cV \to \{\pm 1\}^{d'}$ to be the binary encoding of the vocabulary $\cV$. For each $\x \in \cV$, we embed the token $\x$ as
\[\psi(\x) = \begin{pmatrix}
    \psi'(\x) & \mathbf{0} &\One_{\{\x \in \cB\}}\psi'(\x) & \mathbf{0}
\end{pmatrix}^\top.\]
Here $\mathbf{0} \in \R^{\ell}$ for some $\ell \le O(\log(\card{\cV})+\log L)$ is a zero vector. Also, $\One_{\{\x \in \cB\}}$. For a given input context $\vec{\x}$, the embedded context has the form of 
\begin{equation*}
    \begin{pmatrix}
        \psi'(\x_1) & \psi'(\x_2) & \dots & \psi'(\x_{L-1}) & \psi'(\x_L) \\
        \mathbf{0} & \mathbf{0} & \dots & \mathbf{0} & \mathbf{0} \\
        % \One_{\{\x_1 \in \cB\}} & \One_{\{\x_2 \in \cB\}} & \dots & \One_{\{\x_{L-1} \in \cB\}} & \One_{\{\x_L \in \cB\}}\\
        \One_{\{\x_1 \in \cB\}}\psi'(\x_1) & \One_{\{\x_2 \in \cB\}}\psi'(\x_2) & \dots & \One_{\{\x_{L-1} \in \cB\}}\psi'(\x_{L-1}) & \One_{\{\x_L \in \cB\}}\psi'(\x_{L}) \\
        \mathbf{0} & \mathbf{0} & \dots & \mathbf{0} & \mathbf{0}
    \end{pmatrix}.
\end{equation*}
\noindent \textbf{Position Encoding.}
To allow the attention layer to access the position of the input tokens, we next add a position encoding to each input token. Define $\phi: [L] \to \{\pm 1\}^{\log L}$ be the binary encoding of numbers in $L$. After adding the position encoding, the input context has the following form. 
\begin{align*}
    \Phi(\vec{\x}):=\begin{pmatrix}
        \psi'(\x_1) &  \dots & \psi'(\x_{L-1}) & \psi'(\x_L) \\
        \mathbf{0} &  \dots & \mathbf{0} & \mathbf{0} \\
         % \One_{\{\x_1 \in \cB\}} & \dots & \One_{\{\x_{L-1} \in \cB\}} & \One_{\{\x_L \in \cB\}}\\
        \One_{\{\x_1 \in \cB\}}\psi'(\x_1)  & \dots & \One_{\{\x_{L-1} \in \cB\}}\psi'(\x_{L-1}) & \One_{\{\x_L \in \cB\}}\psi'(\x_{L}) \\
        \mathbf{0} &  \dots & \mathbf{0} & \mathbf{0} \\
        \phi(1)  & \dots & \phi(L-1) & \phi(L)
    \end{pmatrix}
\end{align*}
By construction, each column of the input context is in $\R^d$, for some $d = O(\log(\card{\cV}+\log L)$. 
Given the embedding of the input context, we now construct the SSM layer and the attention layer of the hybrid model.

\noindent \textbf{Mamba Layer.} 
We will need a state dimension $d_s$ equal to the number of bits per bit sequence. Let $\W_A = I - S$ be the block diagonal matrix, where $S$ is the permutation matrix satisfying
\begin{align*}
    S z = (z_2,\dots,z_{d_s-1},z_{d_s},0), \forall z \in \R^{d_s}.
\end{align*}

We define $\W_B$ and $\W_C$ as follows,

\begin{align*}
    \W_B \begin{pmatrix}\psi'(\x_i)
    \\
    \mathbf{0} \\
    \One_{\{\x_i \in \cB\}}\psi'(\x_i)\\\mathbf{0}\\\phi'(i)\end{pmatrix} = 
    \begin{pmatrix}\mathbf{0} \\ \One_{\{\x_i \in \cB\}} \psi'(x_i)\end{pmatrix} 
    , \quad\W_C 
    \bv
   = \begin{pmatrix}\mathbf{0}\\ \mathbf{0} \\\mathbf{0}\\ \bv \\\mathbf{0}\end{pmatrix}, 
\end{align*}


 and $\Delta(\x) = \Ind_{\{\x \in \cB\}}$. We claim the following guarantee for the constructed Mamba layer.
\begin{claim}\label{cl mamba}
    Let $\vec{\x} \in \cV^L$ be a sequence of input contexts. Given the embedded context $\Phi(\vec{\x})$, the SSM layer with parameters $\W_A,\W_B,\W_C,\Delta(\x)$, has the following output
      \begin{align*}
        \begin{pmatrix}
        \psi'(\x_1) & \psi'(\x_2) & \dots & \psi'(\x_{L-1}) & \psi'(\x_L) \\
        \mathbf{0} & \mathbf{0} & \dots & \mathbf{0} & \mathbf{0} \\
        \One_{\{\x_1 \in \cB\}}\psi'(\x_1) & \One_{\{\x_2 \in \cB\}}\psi'(\x_2) & \dots & \One_{\{\x_{L-1} \in \cB\}}\psi'(\x_{L-1}) & \One_{\{\x_L \in \cB\}}\psi'(\x_{L}) \\
        \phi({n_1}) & \phi({n_2}) & \dots & \phi({n_{L-1}}) & \phi({n_L}) \\
        \mathbf{0} & \mathbf{0} & \dots & \mathbf{0} & \mathbf{0} \\
        \phi(1) & \phi(2) & \dots & \phi(L-1) & \phi(L)
    \end{pmatrix},
    \end{align*}
where $n_i = v(\vec{\x}_{1:i})$ corresponds to token in $\cM$ that matches the binary representation of the $0-1$ subsequence of $\vec{\x}_{1:i}$.
\end{claim}

\begin{proof}[Proof of \Cref{cl mamba}]
We initialize $H_0 = 0$ and do induction over $H_t$. By our construction, only the third block of $H_t$ is non-zero. So, for the convenience of notation, we use $H_t$ to denote the third block of the state. Assuming that in time step $t$, $H_t=\phi(n_t)$, we prove this for time step $t+1$. Write $\phi(n_t) = (z_1,\dots,z_{d'})$. If $\x_{t+1} \not\in \{0,1\},$ then $\Delta(\x_{t+1}) =0,$ which implies $H_{t+1} = H_t$ and $n_t = n_{t+1}$. If $\x_{t+1} \in \{0,1\},$ then 
$H_{t+1} = SH_t + \x_{t+1} = \phi(n_{t+1})$. Thus, the third block of the matrix is always $\phi(n_t)$. This implies, after using an MLP layer to combine the output with the input sequence, we know that in the input context has the form of 
      \begin{align*}
        \begin{pmatrix}
        \psi'(\x_1) & \psi'(\x_2) & \dots & \psi'(\x_{L-1}) & \psi'(\x_L) \\
        \mathbf{0} & \mathbf{0} & \dots & \mathbf{0} & \mathbf{0} \\
        \One_{\{\x_1 \in \cB\}}\psi'(\x_1) & \One_{\{\x_2 \in \cB\}}\psi'(\x_2) & \dots & \One_{\{\x_{L-1} \in \cB\}}\psi'(\x_{L-1}) & \One_{\{\x_L \in \cB\}}\psi'(\x_{L}) \\
        \phi({n_1}) & \phi({n_2}) & \dots & \phi({n_{L-1}}) & \phi({n_L}) \\
        \mathbf{0} & \mathbf{0} & \dots & \mathbf{0} & \mathbf{0} \\
        \phi(1) & \phi(2) & \dots & \phi(L-1) & \phi(L)
    \end{pmatrix},
    \end{align*}

\end{proof} 




\noindent \textbf{Transformer Block.} Based on the output of the Mamba layer, we will now construct a Transformer block that can solve the recall task. The Transformer block contains two layers. The first layer contains two heads, while the second layer contains only one head. 

We start with the construction of the first layer. The first layer maps each token $\x_t$ to $(\x_{t-1},\x_{t})^\top$. We denote by $\Attn^{(1)},\Attn^{(2)}$ the two heads of the first attention layer $\AT^{(1)}$. 
For the first head, we define $\W_q^{(1)}=\W_k^{(1)}=0$, $(B_i)_j = -\infty \Ind(j \neq i-1)$. That is to say $\Attn^{(1)}$ is used to select the previous element for each position. We set $\W_v^{(1)}$ such that 
\begin{align*}
    \W_v^{(1)} \begin{pmatrix}\psi'(\x_i)
    \\
    \mathbf{0} \\
    \One_{\{\x_i \in \cB\}}\psi'(\x_i)\\\mathbf{0}\\\phi'(i)\end{pmatrix}  = \begin{pmatrix}\mathbf{0}\\ \psi'(\x_i) \\
    \mathbf{0}\\\mathbf{0}\\\mathbf{0}\end{pmatrix}
\end{align*}
For the second head, we set $\W_q^{(1)}=\W_k^{(1)}=0$, with $(B_i)_j = -\infty \Ind(j \neq i).$ And we set $\W_v^{(2)} = I$. This implies given any input sequence $\vec{\x}$, we have 
\begin{align*}
    \AT^{(1)}(\SSM(\vec{\x})) = \begin{pmatrix}
        \psi'(\x_1) & \psi'(\x_2) & \dots & \psi'(\x_{L-1}) & \psi'(\x_L) \\
        \psi'({\x_0}) & \psi'(\x_1) & \dots & \psi'(\x_{L-2}) & \psi'(\x_{L-1}) \\
        \One_{\{\x_1 \in \cB\}}\psi'(\x_1) & \One_{\{\x_2 \in \cB\}}\psi'(\x_2) & \dots & \One_{\{\x_{L-1} \in \cB\}}\psi'(\x_{L-1}) & \One_{\{\x_L \in \cB\}}\psi'(\x_{L}) \\
        \phi({n_1}) & \phi({n_2}) & \dots & \phi({n_{L-1}}) & \phi({n_L}) \\
        \mathbf{0} & \mathbf{0} & \dots & \mathbf{0} & \mathbf{0} \\
        \phi(1) & \phi(2) & \dots & \phi(L-1) & \phi(L)
    \end{pmatrix}.
\end{align*}
That is to say, the first attention layer maps each column of the input into a query $\x_i^q = \phi(n_i)$, a key $\x_i^k = \psi'(\x_{i-1})$, and a value $\x_i^v = \psi'(\x_i)$.
We next design the second layer $\AT^{(2)}$ that contains only a single head to perform the attention mechanism using the output of $\AT^{(1)}$.
Let $\W_q,\W_k,\W_v$ be the weight matrices of the attention layer.
\begin{align*}
    \W_q \begin{pmatrix}\psi'(\x_i)\\\psi'({\x_{i-1}})\\\One_{\{\x_i \in \cB\}}\psi'(\x_i) \\\phi({n_i})\\\mathbf{0}\\\phi(i)\end{pmatrix} = M\phi({n_i}) ,\quad
    \W_k \begin{pmatrix}\psi'(\x_i)\\\psi'({\x_{i-1}})\\\One_{\{\x_i \in \cB\}}\psi'(\x_i) \\\phi({n_i})\\\mathbf{0}\\\phi(i)\end{pmatrix}= \psi'(\x_{i-1}) ,\quad
    \W_v \begin{pmatrix}\psi'(\x_i)\\\psi'({\x_{i-1}})\\\One_{\{\x_i \in \cB\}}\psi'(\x_i) \\\phi({n_i})\\\mathbf{0}\\\phi(i)\end{pmatrix} = \begin{pmatrix}\mathbf{0}\\\mathbf{0}\\\mathbf{0}\\\mathbf{0}\\\psi'(\x_i)\\\mathbf{0}\end{pmatrix} \\
\end{align*}
We also add a bias $B$ for each position, so that the argmax is achieved at the last recall token. This implies the last output vector is 
    \begin{align*}
        \sum_{i=1}^L \frac{\exp(M(\phi(n_L)\phi(\x_{i-1}) +B_{Li}))}{\sum_{i=1}^L \exp((M\phi(n_L)\phi(\x_{i-1})+B_{Li}))}\begin{pmatrix}\mathbf{0}\\\mathbf{0}\\\mathbf{0}\\\mathbf{0}\\\psi'(\x_i)\\\mathbf{0}\end{pmatrix} = \begin{pmatrix}\mathbf{0}\\\mathbf{0}\\\mathbf{0}\\\mathbf{0}\\\psi'(\x^*_{i+1})\\\mathbf{0}\end{pmatrix}
    \end{align*}
This implies that the hybrid model outputs the correct recall token. We remark that when tokens in $\cM$ are drawn uniformly, with probability $99\%$, each token in $\cM$ appears among the last $\Tilde{O}(W)$ tokens in $\vec{\x}$. This implies that instead of using a window of size $L$, a window of size $\Tilde{O}(W)$ is enough to get the same output.

\end{proof}

\subsection{Construction Implementations}

\label{app constructions}
Both of these constructions are implemented in the code repository. We show here the input embedding and the output of these different constructions. Selective copying's construction can be found in \Cref{fig:construction_selective_copy} and Associative Recall with Decoding's construction can be found in \Cref{fig:construction_decode_recall}. On interesting aspect of these constructions comes from their similarity and dissimilarity to the structures in learned models. Typically, learned models on selective copying learn to output the correct token at each position in the context, while the construction only provides the correct token in the last position. In contrast, the associative recall with decoding construction outputs the correct token at each position in the context, similar to learned models. This difference can be understood in when a task uses fixed positional differences. The selective copying construction uses a fixed mechanism to look-up a distance away from the last token, while learned models learn a more general relative positioning. Associative recall with decoding does not use relative positions, leading to a construction that more readily works at every token position. 

\begin{figure}
    \centering
    \includegraphics[width=0.49\linewidth]{fig/exps/constructions/selective_copy_input.pdf}
    \includegraphics[width=0.49\linewidth]{fig/exps/constructions/selective_copy_output.pdf}
    \caption{An example of the input/embedding and the output for selective copy. The aspects of the construction are kept in relatively similar positions in the implementation. Dark purple is -1, cyan is 0, and yellow is 1.}
    \label{fig:construction_selective_copy}
\end{figure}

\begin{figure}
    \centering
    \includegraphics[width=0.49\linewidth]{fig/exps/constructions/decode_recall_input.pdf}
    \includegraphics[width=0.49\linewidth]{fig/exps/constructions/decode_recall_output.pdf}
    \caption{An example of the input/embedding and the output for associative recall with decoding. The aspects of the construction are kept in relatively similar positions in the implementation. Dark purple is -1, cyan is 0, and yellow is 1.}
    \label{fig:construction_decode_recall}
\end{figure}