File size: 7,716 Bytes
4ca4e4c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
\section{Function Compositions and Limitations of Pure Models}\label{sec lb}
In this work, we define the following family of tasks that we term \emph{function-composition} tasks, which expose the limitations of pure models. 


\begin{definition}[Function Composition]\label{def function composition}
Let $\cV$ be a vocabulary of tokens. Let $m,n \in \Z^+$. Consider a function $F(u,v) : \cV^m \times \cV^n \to \mathcal{Y}$. Let $u(\vec{\x}): \cV^L \to \cV^m$ and $v(\vec{\x}): \cV^L \to \cV^n$ be two functions that map a long sequence of tokens to parameters needed for computing $F$.
The goal for model $M$ is to compute $M(\vec{\x}) = F(u(\vec{\x}),v(\vec{\x}))$.
\end{definition}

Computing deep sequential function compositions has been used as a technique for understanding the limitations of Transformer-based models empirically or through communication complexity \cite{chen2024theoretical,dziri2023faith}. Though the family of tasks we consider here shares a similar flavor to prior works, we need only consider a function of composition with depth 1. Intuitively, it is convenient to think about $u(\vec{\x})$ as a subsequence of $\vec{\x}$ that contains essential information that one should look at (of length $m$ for $m \ll L$ but moderately long, i.e., the width of the necessary context), while $v(\vec{\x})$ can be thought as a small parameter that controls the result of $F(u,v)$.  

Many \emph{long context tasks} naturally fall into such function composition categories. For example, in a natural question answering task, the input context is usually very long, but the question (which must be learned from the context) only sparsely depends on part of the context. Transformers often struggle retrieving the information without consuming almost the whole sequence into memory, while after retrieving information, a pure SSM requires an extremely large state space to perform the rest of the computation.
We start by showing that for a broad class of very simple $u,v$, \emph{pure Transformers and pure state space models cannot compute $F(u,v)$ without sufficient scale}.




\subsection{Limitations of SSMs}
To make the above intuition formal,
we first provide conditions under which $F$ is hard to compute by an SSM. 

\begin{assumption}\label{asp ssm lb}
    Consider any function $F$ that satisfies \Cref{def function composition}. We say the function $F$ is \emph{hard to compute by an SSM} if it satisfies the following property: There exists a set $Q=\{v^{(i)}\}_{i=1}^q \subseteq \cV^n$ such that $G(u):=(F(u,v^{(1)}),\dots,F(u,v^{(q)}))$ is an injection.
\end{assumption}
Our first result shows that if $F$ satisfies \Cref{asp ssm lb}, then a $k$-layer SSM either requires $k$ to be $\Omega(m)$ or needs to have one layer with number of states exponential in $m$. That is, to compute $F$, \textit{the size of an SSM must grow linearly with respect to the hidden parameter $m$}. Formally, %we present the following \Cref{th SSM lb general}.
\begin{theorem}\label{th SSM lb general}
    Let $F$ be a function defined as in \Cref{def function composition} that satisfies \Cref{asp ssm lb}. There is a distribution $D$ over the input $(u,v)$ such that any model $M$ that is a composition of $k$ state space layers $\SSM_i$, with state space $\mathcal{S}_i, i \in [k]$ that can compute $F$ with probability $1/2$ must satisfy $\sum_{i=1}^k \log(\card{\mathcal{S}_i}) \ge \Omega(m\log(\card{\cV})-q\log(\card{\mathcal{Y}}))$.
\end{theorem}
\begin{remark}
    The distribution $D$ considered in \Cref{th SSM lb general} is defined over $(u,v)$ instead of the actual distribution over the input context $\vec{\x}$. For concrete tasks that satisfy \Cref{asp ssm lb}, we construct distributions over $\vec{\x}$ to simulate $D$.
\end{remark}

To prove \Cref{th SSM lb general}, we first prove a structural result for a pure multi-layer SSM. Roughly speaking, if a model is a sequence of multiple layers of state space models, then we can view them as single-layer SSMs. We defer the proof of \Cref{lm group ssm} to \Cref{app proof group ssm}.


\begin{lemma}\label{lm group ssm}
Consider a $k$-layer model $M$ defined that is a composition of $k$ state space layers $\SSM_j, j \in [k]$. There is a model $\SSM'$ that only consists of a single layer SSM, which behaves the same as $M$. In particular, denote the state space of $\SSM_j$ as $\cS_j$, $j \in [k]$, and denote by the state space of $\SSM'$ as $\cS'$, then $\card{\cS'} \le \prod_{j=1}^k \card{\cS_j}$.
\end{lemma}

Given \Cref{lm group ssm} and Yao's min-max principle, we only need to consider deterministic single-layer SSMs. The main technical difficulty of the proof is that, unlike in \citet{jelassi2024repeat,zhan2025overcoming}, which prove hardness against specific tasks, we have little knowledge of the structure of $F$ and cannot compute the probability of failure directly. 
We use information theoretic arguments:  
at a high level, for a fixed sample prefix and $m$ different random control parameters $v$, there need to be $\Omega(m \log |\cV|)$ bits to store all of the necessary information from the prefix to use $v$ correctly.
The full proof of \Cref{th SSM lb general} is in \Cref{app proof SSM lb general}.





\subsection{Limitations of Transformers}
Next we study the limitations of using a Transformer to solve this problem under a memory constraint. We consider \emph{sliding window attention}, a dominant design choice for large context models. The size of the sliding window characterizes the working memory of a Transformer based model. Our hardness result is developed in \Cref{ass lb Transformer} and \Cref{th Transformer lb}. Roughly speaking, when the underlying function $F$ is \emph{locally sensitive}, which implies predicting a token at a position requires information very far from the current position, any pure Transformer model must either be very deep or have one layer that is very dense. We remark that \Cref{ass lb Transformer} is very natural. In the context of function composition, $F(u,v)$, though the length of the essential context $u,v$ may be small, the control parameter $v$ might be sensitive and depend on a long range of the input context, making using a standard Transformer costly. For example, if $v(\vec{x})$ is the last token in the sequence that satisfies some property, then any Transformer must maintain a very long window size in order to compute the control parameter.
The proof of \Cref{th Transformer lb} is in \Cref{app proof transform lb}. 


\begin{assumption}\label{ass lb Transformer}
    Let $F$ be a function that satisfies \Cref{def function composition}. We say the function $F(u(\vec{\x}),v(\vec{\x}))$ is hard to compute by a Transformer if it is $R$-local sensitive. That is, there are two sequences $\vec{\x},\vec{\x'}$ such that $\vec{\x}_{L-R+1:L} = \vec{\x'}_{L-R+1:L}$ but $F(\vec{\x}) \neq F(\vec{\x'})$.
\end{assumption}

\begin{theorem}\label{th Transformer lb}
    Let $F$ be a function that satisfies \Cref{ass lb Transformer}. There is a distribution $D$ over the input such that any model $M$ that is a composition of $k$ Transformer layers $\TF_1,\dots,\TF_k$ that can compute $F$ with probability $2/3$ must satisfy $\sum_{i=1}^k W_i \ge R$.
\end{theorem}




\begin{figure}
    \centering
    \includegraphics[width=0.6\linewidth]{fig/construction_structure.pdf}
    \caption{The construction's style follows taking an input $x$ and implementing 2 functions $u, v$ with an SSM. Typically, $u$ is a truncation of the input, and $v$ is a control parameter (represented in purple). Lastly, a Transformer combines these by implementing $F$ to perform the complete task (represented in red).}
    \label{fig:general_construction}
\end{figure}