source/authored/sections/tasks.tex

\noindent \textbf{Selective Copying.}
Our first task is defined as:

\begin{definition}[Selective Copying]
Consider a vocabulary of tokens $\cV = \cN \cup \cM$, where $\cN = \{\#1,\#2,\dots,\#N\}$, $|\cN| = N$ and $|\cM| = M$. The other values in $\cM$ are arbitrary. Provided some sequence $(\x_i)_{i=1}^L$, let $i = \argmax_{1 \leq i \leq L} \x_i \in \cN$,
the goal is to extract the token $\x_{L+1-\x_{i}}$, i.e.
$\x_{L+1} = \x_{L+1-\x_{i}}.$ 
\end{definition}

As a direct application of \Cref{th SSM lb general} and \Cref{th Transformer lb}, the size of a pure SSM that can solve the selective copying task well must scale linearly with respect to $\card{\cN}$, while a pure Transformer must have working memory that scales linearly in the context length $L$. The proof is in \Cref{app proof lb copy}.

\begin{theorem}\label{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{figure}
    \centering
    \includegraphics[width=0.6\linewidth]{fig/selective_copy_construction_structure.pdf}
    \caption{The construction solving selective copy takes an input sequence and finds the most recent number token (as represented in the bottom squares of the output of the SSM). The Transformer can then use these to look back some relative distance to find the correct token to output.}
    \label{fig:var_copy_construction}
\end{figure}

\noindent \textbf{Hybrid Model for Selective Copying}

To sidestep the limitations of pure models, we \textbf{design a two-layer hybrid that provably solves the problem with small input-independent memory \emph{and} input-dependent memory}. Our construction follows the discussion in \Cref{sec merge}. In particular, the SSM component can be realized by a Mamba model.
\begin{theorem}\label{th hybrid copy}
There is a two-layer hybrid with 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 space, while the attention layer has dimension $d$ and a sliding window size of $O(N)$.
\end{theorem}
The proof of \Cref{th hybrid copy} is in \Cref{app proof hybrid copy}. We remark that the number of parameters of our hybrid model is only $\poly\log(\max(\card{\cV},L))$ with working memory $\Tilde{O}(N)$, which is much smaller than $L$, unless $N$ is extremely large. The construction follows a structure where the SSM stores the most recent number token into its state, adding it to the current token. The Transformer can then use this information to copy the token that many positions in the past (Fig.  \ref{fig:var_copy_construction}).


\noindent \textbf{Associative Recall with Decoding}

As a further concrete application of our framework, we introduce another task, for which a hybrid model outperforms both pure state space models and Transformers.
\begin{definition}[Associative Recall with Decoding]
    Consider a vocabulary of tokens $\cV = \cM \cup \{0,1\}$, where $\cM$ is a set of word tokens. Let $\vec{\x} \in \cV^L$ be a sequence of input tokens and let $v(\vec{x}) \in \{0,1\}^{\log(\card{\cV})}$ be the $0-1$ subsequence of $\vec{\x}$. Denote by $\Phi(\vec{\x}) \in \cM$, the token with binary representation $\vec{\x}$.
    Given $v(\vec{\x})$, the goal is to output the next token in $\vec{\x}$ behind the last $\Phi(v(\vec{\x}))$ token. 
\end{definition}

As another implication of \Cref{th SSM lb general} and \Cref{th Transformer lb}, the size of a pure SSM that can solve the associative recall with decoding task well must scale linearly with respect to the number of all possible word tokens, while a pure Transformer must have working memory that scales linearly in the context length $L$. 
The proof of \Cref{th lb associate recall} is in \Cref{app proof lb recall}.


\begin{theorem}\label{th lb associate recall}
    Consider the task of associate 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}$. Any pure Transformer model that can solve the task with probability $90\%$ must have $\sum_{i=1}^kW_i \ge \Omega(L)$.
\end{theorem}



We next show that for very natural distributions (including the one stated in \Cref{th lb associate recall}), the performance of a hybrid model can be much better than that of a pure model. 
In particular, we show that we can construct a hybrid model such that the number of parameters of the model scales with the logarithm of the task size. The intuition here is that for associative recall, the only tokens that matter tend to appear at the end of the sequence, which implies that the window used by the Transformer can be improved to much smaller than $L$. For example, if each token is sampled uniformly from the vocabulary, then with probability $99\%$, within a window of size $\Tilde{O}(\card{\cM})$, we can see all distinct tokens from the vocabulary. In fact, the hard distribution we considered in \Cref{th lb associate recall} also satisfies such a property.
Thus, once we use a state space model to extract the control variable $\vec{x}$, we are able to solve the problem with a small model with a small working memory. We defer the proof of \Cref{th hybrid recall} to \Cref{app proof hybrid recall}.



\begin{theorem}\label{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 selective copying task with probability $99\%$ for an input sequence $\vec{\x} \in \cV^L$ such that tokens in $\vec{\x}\setminus v(\x)$ are drawn from a uniform distribution. 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}

% }