source/authored/sections/upper_bound.tex

\section{Hybrid Model for Function Composition}\label{sec merge}
Using our hardness results, if a function $F$ satisfies \Cref{asp ssm lb} and \Cref{ass lb Transformer} simultaneously, then any pure SSM must need a large number of parameters to solve the problem, making a key bottleneck for training and deploying the model, while any pure Transformer must need a working memory nearly in $\Omega(L)$ due to data redundant, making inference over long-context data a key bottleneck.


To achieve the best of both worlds, i.e., a model with a small number of parameters and relatively small working memory, we consider hybrid models that combine state-space and Transformer layers. The intuitive motivation for hybrid use is that \emph{a state space model can implicitly act as an encoder that summarizes the information from the long context $\vec{\x}$ and passes the compressed information to a Transformer}. Since the Transformer itself does not have the ability to select the correct positions to look at unless it has a deep depth or a large window size, using more comprehensive information can reduce the space requirement of the Transformer. 

To avoid making the notation messy, we think of $\vec{v}(\vec{\x})$ as a single token so that given $(\vec{u},\vec{v})$, a Transformer solves the problem by answering a query $q_v$. Suppose we have a state space model $\SSM_u$ parameterized by $(S^u,R^u)$ that maps $\vec{\x} \in \cV^L$ to a sequence $\vec{a} \in \cV^L$ such that $\vec{a}_{L-m+1:L} = \vec{u}(\vec{\x})$ and an encoder based model $\SSM_v$ parameterized by $S^v,R^v$ that maps $\vec{\x} \in \cV^L$ to a sequence $\vec{b} \in \cV^L$ such that $\vec{b}_L = \vec{v}(\vec{\x})$. We can build a merged $\SSM$ by combining the two state space models in a black box way. That is,  
\begin{align*}
    & S(\x_1,\dots,\x_i) = (S^u (\x_1,\dots,\x_i), S^v (\x_1,\dots,\x_i)), \\
    & R(S(\x_1,\dots,\x_i)) = \begin{pmatrix}
        R^u(S^u (\x_1,\dots,\x_i)) \\
        R^v(S^v (\x_1,\dots,\x_i)
    \end{pmatrix}.
\end{align*}
In matrix form, $\SSM$ maps $\vec{\x}$ to the following matrix
\begin{align*}
    \begin{pmatrix}
        r^u_1 & r^u_2 &,\dots, & r^u_L \\
        r^v_1 & r^v_2 &,\dots, & r^v_L
    \end{pmatrix}.
\end{align*}
In particular, if we look at the last $m$ columns of the output of the state-space model, then we have
\begin{align*}
    \begin{pmatrix}
        u_1 & u_2 &,\dots, & u_m \\
        r^v_{L-m+1} & r^v_{L-m+2} &,\dots, & v
    \end{pmatrix}.
\end{align*}
Suppose we have a Transformer $\TF$ parameterized by $(\W_q^{\TF},\W_k^{\TF},\W_v^{\TF})$ such that given $(\vec{u},\vec{v})$, it can compute $F(\vec{u},\vec{v})$. We consider the following attention layer $\W_q(r^u_i, r^v_i)^\top = \W_q^{\TF}r^v_i, \W_k(r^u_i, r^v_i)^\top = \W_k^{\TF}r^u_i, \W_v(r^u_i, r^v_i)^\top = \W_v^{\TF}r^u_i$.
The output of $\TF\circ \SSM$ is exactly $F(\vec{u},\vec{v})$. Furthermore, such a construction \emph{preserves the model size and the working memory of both state space models and Transformers}. We next use this idea to show that for several natural synthetic tasks, we can construct small scale hybrid models achieving good working memory efficiency.