source/authored/sections/experiments.tex

\section{Experiments}\label{sec experiments}

Our theoretical results show fundamental expressivity differences between pure models and hybrid constructions. We empirically validate three claims related to these results:
\begin{itemize}[leftmargin=*, nosep]
\item \textbf{C1.} The construction for the hybrid model empirically outperforms the pure Transformer and pure SSM baselines, as predicted by our theoretical results,
\item \textbf{C2.} Under standard training approaches, a \emph{learned} (rather than constructed) hybrid \text{also outperforms} pure Transformer and pure SSM baselines, 
\item \textbf{C3.} In further and more realistic settings, including out-of-distribution and length generalization scenarios, learned hybrids continue to outperform pure Transformer and pure SSM baselines.
\end{itemize} 
Here, C1 verifies our basic theoretical claims, while C2 and C3 show that the benefits of hybrids persist in typical scenarios (e.g., the hybrid model is trained in standard ways, the tasks and settings deviate from the exact ones we studied). By validating these claims, {\textbf{\emph{we demonstrate that the fundamental benefits of hybrid models carry over to practical settings}}}---and are not just a theoretical curiosity. 


\subsection{Construction Implementations}

Each of the constructions for \Cref{th hybrid copy} and \Cref{th hybrid recall} have been implemented to test their validity. Both perform their respective tasks on the last position in the context, as desired, validating C1. For details, see Appendix \ref{app constructions}.

\subsection{Learnability Experiments}
To validate C2, we conduct experiments designed to test the capacity of a hybrid to learn the two main tasks we studied, selective copying and associative recall with decoding. Besides these, we also empirically study two additional tasks: multi-key associative recall (MKAR) and needle-in-a-haystack (NH). These tasks also fall into the category of function composition and are standard tasks used to evaluate pure Transformer-based and SSM models. Our goal is to determine whether the hybrid models learned with standard training approaches (rather than explicitly constructed) continue to outperform non-hybrids. As we shall see, our findings confirm that this is the case. 

To simplify comparisons, rather than directly comparing the state sizes or the input-independent memory for the Transformers, we compare models with similar parameter counts, controlled through the embedding dimension of the tokens. Both the state size and the input-independent memory scale similarly as embedding dimension is increased. 

\noindent \textbf{Experiment Details.} These models are comprised of GPTNeoX and Mamba layers for attention and SSM layers. We use RoPE positional encodings. Unless otherwise specified, we are using windowed, causal attention, the transformers have a single head, and the Mamba state dimension expansion is 1. All experiments are trained to convergence using linearly decaying learning rates. The experiments are seq-to-seq, and accuracy if measured over all valid tokens. Transformers use a windowed attention mask.

When listed in figures, layers are read left-to-right. For example, SSM-TF is an $\SSM$ layer followed by a $\TF$ layer. Additional details can be found in Appendix \ref{app experiments}.


\noindent \textbf{Selective Copy.} 
\begin{figure}[t]
    \centering
    \includegraphics[width=0.6\linewidth]{fig/exps/var_copy/layers_dim.png}
    \small
    \begin{tabular}{c c c c c}
    \toprule
    Parameters & Pure TF & Pure SSM & TF$\to$SSM & \textit{SSM$\to$TF} \\
    \midrule
    $\sim1000$  & 0.056 & 0.084 & 0.100 & 0.087 \\
    $\sim2000$  & 0.352 & 0.305 & 0.433 & 0.999 \\
    $\sim6000$  & 0.727 & 0.485 & 0.822 & 1.000 \\
    $\sim12000$ & 0.923 & 0.931 & 0.908 & 1.000 \\
    \bottomrule
    \end{tabular}

    \caption{Results from training small models on Selective Copy, across an increase in the hidden dimension of the models. At 2000 parameters, hybrid models consistently attain perfect accuracy. The pure models, with 6x the parameters, only attain around 0.9 accuracy.}
    \label{fig:exp_var_copy}
    \vspace{-0.2cm}
\end{figure}
\noindent \emph{Setup.} Given the relative simplicity of this task, we can study model expressivity for small models, with up to approximately ten thousand parameters.  

\noindent \emph{Results.} We depict the results in \Cref{fig:exp_var_copy}. As expected, the hybrid performs the task to 90\% accuracy with significantly fewer parameters than either the pure Transformer or pure SSM, by around a factor of $6 \times$. 
Also note that the ``reverse'' hybrid, with the Transformer layer first, performs no differently than the pure models. This is consistent with prior work \cite{park2024can}.

\noindent {\bf Associative Recall with Decoding.} 
\begin{figure}[t]
    \centering
    \includegraphics[width=0.6\linewidth]{fig/exps/binary_recall/layers_dim_3.png}
    \caption{Results from training small models on Associative Recall with Decoding. Even at much smaller scales than the pure models, the hybrid is the only architecture that attains 0.5 accuracy. At the scales tested, none of the pure models performed the task with more than 0.4 accuracy.}
    \label{fig:exp_decode_recall}
\end{figure}
\emph{Setup.} For this task, we expect both the pure Transformer and the pure SSM to struggle, while a hybrid has the expressive power to represent this task. Unlike the other tasks, this experiment utilized three layer models rather than two layer ones. This is due to the construction for \Cref{th hybrid recall}, where three layers were needed. These models required significantly more parameters than the other models, close to 1 million. 

\emph{Results}. For these three layer models, we observe the same behavior: the hybrid excels while the pure models struggle. As seen in \Cref{fig:exp_decode_recall}, at the scales tested, none of the pure models achieved greater than 40\% accuracy, while the hybrid did at much smaller scales, eventually surpassing 50\% accuracy. At a high level, this task is more difficult than the others as it requires the more complex computation of binary values before performing the commonly tested task of associative recall. 


\noindent \textbf{Multi-Key Associative Recall.}
\textit{Setup.} We now turn our attention to a further pair of tasks, the first being MKAR. 

\begin{figure}[t]
    \centering
    \includegraphics[width=0.6\linewidth]{fig/exps/assoc_recall_mk/layers_dim.png}
    \small
    \begin{tabular}{c c c c c}
\toprule
Parameters & Pure TF & Pure SSM & TF$\to$SSM & \textit{SSM$\to$TF} \\
\midrule
$\sim1000$  & 0.124 & 0.158 & 0.131 & 0.144 \\
$\sim2000$  & 0.159 & 0.173 & 0.183 & 0.512 \\
$\sim6000$  & 0.230 & 0.356 & 0.286 & 0.990 \\
$\sim12000$ & 0.668 & 0.517 & 0.524 & 0.989 \\
\bottomrule
\end{tabular}
    \caption{Results from training small models on Multi-Key Associative Recall, across an increase in the hidden dimension. The hybrid consistently outperforms the pure models of the same depth and similar parameter counts. The hybrid models could perform the task to 60\% accuracy with $6 \times$ fewer parameters than any of the pure Transformers.}
    \label{fig:exp_assoc_recall_mk}
\end{figure}


\begin{definition}[MKAR]
    Let $\vec{\x}$ be a sequence of length $L$ sampled from a vocabulary $\cV$, and let $k$ be some small number. Let $K = \vec{\x}_{L-k:L}$. Let $i$ be the last position in the context where $\vec{\x}_{i:i+k} = K$. Performing \textit{MKAR} is outputting $\vec{\x}_{i+k}$.
\end{definition}

In the framework of function-composition, we can take $v$ to be the empty map, $u$ to include enough context to find the key, and $F$ to be the look-up operation. Since $F$ can have any output depending heavily on the context, this is hard for an SSM. Since $u$ could be large, this is also hard for a Transformer. However, since $v$ is empty, function-composition does not immediately indicate if a separation exists between Transformers and hybrids.

\textit{Results.} In \Cref{fig:exp_assoc_recall_mk}, we see that SSMs perform quite poorly, while hybrids and Transformers can perform the task at scale. However, we see a similar separation between hybrids and Transformers present for selective copying. Specifically, hybrids perform the task on average with $6 \times$ fewer parameters than the pure Transformers to an accuracy of $60\%$. 


\begin{figure}[t]
    \centering
    \includegraphics[width=0.6\linewidth]{fig/exps/needle/layers_dim.png}
    \caption{Results from training small models on Needle in a Haystack, across an increase in the hidden dimension of the models with no context windowing. The hybrid and SSM perform this task with fewer parameters than the Transformer, however we still see the hybrid with a slight improvement. This task was expected to be hard for the Transformer and not the SSM.}
    \label{fig:exp_needle}
\end{figure}

\noindent {\bf Needle in a Haystack.} \textit{Setup.} To complement MKAR, we trained the same models on needle-in-a-haystack (NH). 
\begin{definition}[NH]
    Let $\vec{\x}$ be a sequence of length $L$ sampled from a vocabulary $\cV \cup \{M\}$. Let the location of $M$ be denoted by $i^*$. Performing \textit{NH} is outputting $\vec{\x}_{i^*+1}$.
\end{definition}
This task is simple: copy the token(s) after a marker token when requested for at the end of the input sequence. This task is hard for Transformers due to windowing, while SSMs and hybrids should perform this task easily. In the framework of function-composition, $u$ is the empty map, and $F$ is the identity. As such, we do not have the hardness criterion for SSMs; we should therefore not immediately expect hybrids to perform differently than SSMs.

\textit{Results.} We show results in \Cref{fig:exp_needle}  for full-context attention. Even in the situation where the Transformer could possibly learn the task, small token dimensions result in learnability issues. 
SSMs also perform more inconsistently than hybrid models in small parameter regimes. The mechanism behind these separations is not directly characterized by function-composition and is left to future work.


For both of these tasks, we see the same property: on these synthetic tasks and at the scales tested, \textbf{hybrids outperform pure models, even when we use tasks that are outside of the function-composition framework.}






\subsection{Further Experiments}

In contrast to the much smaller expressivity experiments, our next set of experiments are designed to increase the scale of the models to somewhere nearer those of modern LLMs. These models will have around 100 million parameters each, closer to the scale of standard language models. Different properties of these models are tested since accuracy on many of the above tasks already approaches 1, and scaling up difficulty parameters, such as vocabulary size, which only increases the difficulty by a small amount, leads to models with similar behaviors to the expressivity experiments.


\noindent {\bf Associative Recall with Decoding.}
To analyze these different analyses, we use the more difficult task of associative recall with decoding. Empirically, this task proved far more challenging than selective copy, where 2-layer models with the same scale of parameters learned nothing. 

\begin{figure}[t]
    \centering
    \includegraphics[width=0.6\linewidth]{fig/exps/mini_decode_recall.png}
    \caption{The distribution of accuracies across different input sequence lengths. Hybrid models with comparatively similar parameters as their attention/Transformer counterparts perform better at longer lengths consistently.}
    \label{fig:decode_recall_length_gen}
\end{figure}

\noindent {\bf Length Generalization.}  
\textit{Setup.} First, we investigate the length generalization of different models. Each model is trained on sequences of length 20 to 50, and tests on longer sequences as well. Comparing the hybrid model to the Transformer (T\_rope) and SSM (mamba), Figure \ref{fig:decode_recall_length_gen} shows how the hybrid models \textit{consistently} outperforms the pure models. 

\textit{Results.} As expected, the performance drops as sequences grow longer; however, hybrids lose performance at the slowest rate. This means that even though hybrids and Transformers behave within 2\% of each other for short sequences, this separation grows to be around 10\% for longer sequences.

\begin{table}[t]
    \centering
    % \small
    \begin{tabular}{c|c c c}
    \toprule
    Train Proportion & SSM & TF & Hybrid \\
    \midrule
    0.05 & 0.24 & \textbf{0.47} & \textbf{0.47} \\
    0.1  & 0.34 & 0.40 & \textbf{0.47} \\
    0.3  & 0.17 & 0.64 & \textbf{0.74} \\
    0.5  & 0.46 & 0.63 & \textbf{0.77} \\
    0.8  & 0.67 & 0.63 & \textbf{0.83} \\
    0.9  & \textbf{0.86} & 0.61 & 0.80 \\
    \bottomrule
    \end{tabular}

    \caption{Results from training 12-layer models with different proportions of bits for Associative Recall with Decoding. Data are evaluation accuracies for evaluation bit proportions of 0.2. Each architecture tends to improve performance as the training bit proportion increases, with hybrids consistently out-performing the pure models.}
    \label{tab:robustness}
\end{table}


\noindent {\bf OOD Generalization.}
\textit{Setup.} We also tested these models as their sampling distributions are changed. Specifically, we tested Associative Recall with Decoding with differing proportions of bits between test and train time. This task allows us to test these behaviors on larger models, as the other two tasks saturate, only showing 100\% accuracy on all tests. Under these distributions, we can see which of the architectures learns representations that consistently perform well across different varying sampling distributions.

\textit{Results.} The results can be seen in \Cref{tab:robustness}. For almost all training distributions, the hybrid indeed performs the best on a 0.2 proportion test set. However, there are some other notable trends within these results beyond just the hybrid's performance. In particular, the different architectures show varying behavior across training distributions. SSMs tend to improve the most as more training bits are added, while Transformers improve the least. Hybrid models attain the best of both works, acting well with both a high frequency and a low frequency of bits.