ProCreations's picture
Reproduction logbook (paper-82EJxJzG6r)
4ca4e4c verified
Raw
History Blame Contribute Delete
5.64 kB
\section{Introduction}
Transformers are the workhorse architecture for modern language models. While highly expressive and capable, Transformer-based models suffer from high complexity, particularly for inference time processing of long sequence inputs. As a result, developing alternative non-Transformer architectures has become among the most important problems in LLM development. Structured state space models (SSMs) like Mamba \cite{gu2024mamba} are among the most promising such alternatives. Such models trade off complexity for expressivity \cite{jelassi2024repeat}, achieving higher throughput---but typically lower performance---compared to Transformer-based models.
A natural question is whether we can sidestep this tradeoff and produce a model architecture that offers the best of both worlds. \emph{Hybrid sequence models} seek to achieve this objective. These models, which mix layers from Transformer architectures (e.g., attention layers) with SSM layers, ideally outperform either Transformer-only or SSM-only models. In a short time, hybrids that can empirically do so on particular tasks have been scaled up from tiny sizes to as large as 50 billion parameters. For example, Nvidia's Nemotron-H hybrid model family \cite{blakeman2025nemotron} offers both better downstream evaluation performance \emph{and} higher throughput (due to the presence of lower-complexity Mamba layers) than Transformer-only baselines.
Despite these empirical successes, \textbf{\emph{we have no principled understanding}} of why hybrid models can outperform models made up of a single type of layer. Similarly, we do not yet know for what basic tasks we should expect hybrids to behave in this way. This paper takes the first steps towards providing a \textbf{\emph{fundamental theory addressing architectural tradeoffs for hybrid models}}. It does so by first showing that on a family of core tasks where \emph{pure} (i.e., standard Transformer-only and SSM-only) models provably suffer from limitations (in terms complexity and memory). In contrast, we build constructions of hybrid models that \emph{do not} have the same limitations on representative tasks---including key tasks like associative recall and selective copying---thus exhibiting provable benefits for hybrids.
Concretely, we evaluate the performance of a model by analyzing its \emph{input-independent memory (model size) and input-dependent memory (working memory)}. We focus on a \emph{function-composition} family of tasks (Fig.~\ref{fig:example}) that combine both a long-context control variable and a local context-addressable lookup; such tasks naturally model real-world data. For these, (i) under an injectivity condition, we prove
%memory lower bounds showing
that, for this family, any pure SSM requires large internal state (or many layers); as a result, their size scales linearly with respect to the hidden dimension of the problem to solve the problem.
Likewise, (ii), under a local-sensitivity condition, any sliding-window Transformer (which includes full-window attention) requires a large window scaling linearly with respect to the length of the input context. Together, this pair of results indicates that for a broad class of tasks, pure SSM-based models and pure Transformer-based models \textbf{\emph{fail to achieve good expressivity and inference efficiency simultaneously}}.
We study two representative synthetic tasks in this family, namely selective copying %\cite{jelassi2024repeat}
and a variant of associative recall \cite{arora2023zoology}.
For these tasks, we construct \textbf{\emph{provably successful shallow hybrid models}} whose size scales with the logarithm of the size of the tasks while using only sublinear memory.
Empirically, we validate our theoretical results and investigate hybrid versus non-hybrid performance in further settings and on additional tasks, such multi-key associative recall (MKAR) and needle-in-a-haystack (NH). We find that for selective copying and MKAR, \textbf{\emph{hybrids can perform the task with similar or better quality than the pure models with {6 times fewer parameters}}}. For associative recall with decoding, at the tested scales, the pure models \emph{never} match the performance possible with the hybrid model.
On top of measuring performance at fixed model sizes, we observe that hybrid models exhibit stronger length generalization and out-of-distribution (OOD) robustness. We see that when trained on the same distribution of short examples, hybrids consistently out-perform pure Transformers by around 10\% accuracy for long sequences.
For out-of-distribution testing, the hybrid model sometimes attains over 15\% higher performance than either the Transformer or the SSM with around the same number of parameters.
\noindent \textbf{Roadmap of the paper.}
In \Cref{sec notation}, we provide necessary preliminaries and notations. In \Cref{sec lb}, we introduce a family of tasks formulated as computing a function composition and provide conditions under which non-hybrid models fail to solve the tasks efficiently. Next, in \Cref{sec merge}, we focus on two specific tasks (of varying difficulty) that are within this family and construct hybrid models that outperform non-hybrids. In \Cref{sec experiments}, we conduct experiments to show the benefits of hybrid models empirically.
\begin{figure}[t]
\centering
\includegraphics[width=0.6\linewidth]{fig/example.png}
\caption{Example function composition task. The answer to a learned question only depends on a part of the long context input.}
\label{fig:example}
\vspace{-1em}
\end{figure}