source/authored/appendix/missing_proof_lb.tex

\section{Omitted Proofs in \Cref{sec lb}}\label{app proof lb}


\subsection{Proof of \Cref{lm group ssm}}\label{app proof group ssm}

\begin{lemma}[Restatement of \Cref{lm group ssm}]
Consider a $k$-layer auto-regressive SSM $M$ defined as $\SSM_1 \to \SSM_2 \to \dots \to \SSM_k$, where for $j \in [k]$, $\SSM_j$ is an SSM layer. There is a model $\SSM'$ that only consists 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}

\begin{proof}[Proof of \Cref{lm group ssm}]
    Let $s\^j_0$, $u\^j$, $r\^j$, $R\^j_i$, and $S\^j_i$ be the SSM parameters that define $\SSM_j$. The output of this $k$-layer model is computed as follows:
    \begin{align*}
        S\^j_0(\emptyset) &= s\^j_0 \quad \text{ for } j \in [k] \\
        S\^1_i(\x_1,\dots,\x_i) &= u\^1(S\^1_{i-1}(\x_1,\dots,\x_{i-1}), \x_1) \\
        S\^j_i(R\^{j-1}_1,\dots,R\^{j-1}_i) &= u\^j(S\^j_{i-1}(R\^{j-1}_1,\dots,R\^{j-1}_{i-1}), R\^{j-1}_i) \quad \text{ for } 2 \le j \le k \\
        R\^j_i(R\^{j-1}_1,\dots,R\^{j-1}_i) &= r\^j(S\^j_i(R\^{j-1}_1,\dots,R\^{j-1}_i))
    \end{align*}
    The output of each layer is fed into the next layer to update its state.
    %By Lemma D.1 in \cite{jelassi2024repeat}, we know that for every $i \in [k]$, there is a decoding map $G_i: \mathcal{S} \to \cV$ for every input sequence $\vec{x}=(x_1,\dots,x_L)$, we can write $SSM_i(\vec{x}) = G_i ( S^i(\vec{x}) )$, where $S^i(\vec{x})$ is the state of $SSM_i$ after seeing the whole input $\vec{x}$. 
    We define $\cS'$ to be $\{(s_1,\dots,s_k) \mid s_i \in\cS^i, i \in [k] \}$, so $\card{\cS'} = \prod_{i=1}^k \card{\cS^i}$. Let $\SSM'$ have state
    \[S'_i(\x_1,\dots,\x_i) = \lp(S\^1_i(\x_1,\dots,\x_i),\ S\^2_i(R\^1_1,\dots,R\^1_i),\ \dots\ ,S\^j_i(R\^{j-1}_1,\dots,R\^{j-1}_i)\rp).\]
    We remark that the definition above is well defined, since for every $j \in [k]$, $(R\^j_1,\dots,R^j_i)$ only depends on $(\x_1,\dots,\x_i)$, thus is available at time $i$. The output function $r'$ is defined to be the output of the final layer of the multilayer SSM $(S'_i)_k$,
    \[r'(S'_i(\x_1, \dots, \x_i)) = R\^k_i((S'_i)_k)\]
    All other parameters are simply their application across the different components of the compound state.
    \begin{align*}
        S'_0(\emptyset) &= (s\^1_0, \dots, s\^j_0) \\
        u'(S'_{i-1}(\x_1, \dots, \x_{i-1}), \x_i) &= (S\^1_i, \dots, S\^j_i)
    \end{align*}
    % For each $i \in [k]$, we define $z_{ij}$ to be the $j$th input of $SSM_i$. We construct a single layer of a state space model as follows.
    % \begin{align*}
    %     & S'_j(x_1,\dots,x_j) = (S^1(z_{11},\dots,z_{1j}),\dots,S^k(z_{k1},\dots,z_{kj})) \\
    %     & R'_j(x_1,\dots,x_j) = r^k(S^k(z_{k1},\dots,z_{kj}))
    % \end{align*}
    % We remark that the definition above is well-defined, because for every $i \in [k]$, $(z_{i1},\dots,z_{ij})$ only depends on $(x_1,\dots,x_j)$. Furthermore, the output of $SSM'$ is the same as $M$ for every input sequence by definition. 
    This construction computes the same result as the original multilayer model with the same size of state.
\end{proof}


\subsection{Proof of \Cref{th SSM lb general}} \label{app proof SSM lb general}

\begin{theorem}[Restatement of \Cref{th SSM lb general}]
   Let $F$ be a function defined \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{proof}[Proof of \Cref{th SSM lb general}]
    By \Cref{lm group ssm}, we only need to show the hardness against a single layer of state space model. By Yao's min-max principle, it is sufficient to show that we cannot construct any deterministic state-space model that can solve the problem with a good probability when the input is drawn from some distribution $D$.
    We consider the following distribution $D$ over a sequence of length $m+n$ tokens. For the first $m$ tokens, we sample each token uniformly from $\cV$, representing $u$ is drawn uniformly at random. For the last $n$ tokens, we will draw $v \sim Q$ independent on $u$. We next lower bound the size of $\card{\mathcal{S}}$ of any SSM that can give the correct output when the input $(u,v)$ is drawn from $D$. Denote by $Y_i = F(u,v^{(i)})$ for $i \in [q]$ and $Y=(Y_1,\dots,Y_q)$. Furthermore, let $s = S_m(u)$ the random state of the model after reading $u$. By \Cref{asp ssm lb}, we know that $H(s\mid u) = 0$.
    Since 
    \begin{align*}
        I(u;s) = H(u) - H(u \mid s) = H(s) - H(s \mid u),
    \end{align*}
    we have
    \begin{align*}
        H(u) - H(u \mid s) = H(s) \le \log(\card{\mathcal{S}}),
    \end{align*}
    because the support of $s$ has size at most $\card{\mathcal{S}}$. We next upper bound $H(u \mid s)$. By the symmetry of mutual information, we have
    \begin{align*}
        H(u \mid s) & = H(Y \mid s) + H(u \mid Y, s) - H(Y \mid u, s) \\
        & = H(Y \mid s) \le \sum_{i=1}^q H(Y_i \mid s),
    \end{align*}
    where the second equation holds by $H(u \mid Y, s) - H(Y \mid u, s) = 0$. It remains to upper bound $H(Y_i \mid s)$. Let $\err_i:=\Pr_{u\sim \cV^m}(R(S_{m+n}(u,v^{(i)}))\neq Y_i)$. By Fano's inequality, we have 
    \begin{align*}
        H(Y_i \mid s)\le H_2(\err_i) + \err_i \log(\card{\mathcal{Y}}). 
    \end{align*}
    Thus,
    \begin{align*}
        H(u\mid s) &\le q \frac{1}{q}\sum_{i=1}^qH(Y_i \mid s)\\
        &\le q \frac{1}{q}\sum_{i=1}^q(H_2(\err_i) + \err_i \log(\card{\mathcal{Y}}))\\
        &\le q (H_2(\err)+\err \log(\card{\mathcal{Y}}))
    \end{align*}
    This implies, if we set up $\err<1/8$, then we have 
    \begin{align*}
        \log(\card{\mathcal{S}}) \ge m\log(\card{\cV}) - q(H_2(1/8)+\log(\card{\mathcal{Y}})/8).
    \end{align*}

%\mnote{fix the last inequality by considering the range of $q$.}
To conclude the proof of \Cref{th SSM lb general}, we discuss the possible choice of $q$ and $\card{\mathcal{Y}}$ that satisfy \Cref{asp ssm lb}. Since $G(u):=(F(u,v^{(1)}),\dots,F(u,v^{(q)}))$ is an injection, we know that the smallest possible choice of $q, \card{\mathcal{Y}}$ satisfies $q\log(\card{\mathcal{V}}) = O(m\log(\card{\cV}))$. This implies $\log(\card{\mathcal{S}}) \ge \Omega(m\log(\card{\cV}))$


\end{proof}



\subsection{Proof of \Cref{th Transformer lb}}\label{app proof transform lb}


\begin{theorem}[Restatement of \Cref{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 blocks $\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{proof}[Proof of \Cref{th Transformer lb}]
Consider any sliding window Transformer $M$ with window size $W_i$ for the $i$-th layer. Denote by $W = \sum_{i=1}^k W_i$ the effective window size of a Transformer.
Notice that the output of $M$ is a deterministic function of $(\x_{L-W+1},\dots,\x_L)$. Thus, we construct a distribution $D$ that draws $\vec{\x}$ and $\vec{\x'}$ uniformly. For any input drawn from $D$, $M$ has the same output. However, by definition, $F(\vec{\x}) \neq F(\vec{\x'})$, which implies that $M$ fails to output correctly with probability $1/2$.






% We first construct two distributions over the input $D_0$ and $D_1$. $D_0$ is a uniform distribution drawn from $\cV^L$. $D_1$ first draws an input sequence $\vec{x}$ from $D_0$ and chooses a random position $i$ and modifies $x_i$ to be a random number token. 

% Notice that for any fixed Transformer $M$ and for any input sequence $\vec{x}$, if we random draw a position $i$ from $L$, then the probability that this position is stored by any head of the model is at most $\sum_{i=1}^{M}\ell_i /L$. Since the output of any Transformer is uniquely defined by the stored key/value pairs, if we draw a sequence $\vec{x}$ from the mixture $D_0/2+D_1/2$, then with probability $3/4$, model $M$ will fail to output $\perp$ or the correct letter token if $\sum_{i=1}^{M}\ell_i /L<1/100$. 
\end{proof}