\section{Complete Preliminaries and Notations}\label{app notations}

\begin{table}[t]
    \centering
    \begin{tabular}{|c|c|}
    \hline
        Symbol & Meaning \\
    \hline \hline
        $\phi, \psi, \Phi$ & Token and Positional Embeddings \\ 
        $u,v$ & Control parameters of the task \\
        $F$ & The target task \\
        $H$ & (Relative) Entropy \\
        $I$ & Mutual Information \\
        $\W_q,\W_k,\W_v,\W_o$ & Transformer parameters \\
        $\W_A,\W_B,\W_C,\Delta$ & SSM parameters \\
        $\x$ & The input sequence \\
        $\cV$ & Vocabulary space \\
        $\cN, \cM$ & Number/Vocabulary components of $\cV$\\
        $Y$ & Target space, typically $\cV^n$  \\
        $d$ & The token dimension \\
        $d_s$ & The state dimension \\
    \hline
    \end{tabular}
    \caption{Notation.}
    \label{tab:notation}
\end{table}






We consider sequence to sequence token prediction problem. Let $\cV$ be some vocabulary of tokens and $V = |\cV|$ and $\vec{\x}=(\x_i)_{i=1}^L$ be and input sequence. A language model $M$ is sequence to sequence map $M: \cV^L \to \cV^m$ of the form $M(\vec{\x}) = F_N\circ F_{N-1} \circ \dots \circ F_1 (\vec{\x})$, where each $F_i$ is a sequence to sequence map called layer.  









We will consider several different layers in this paper.

\noindent \textbf{Transformer Layer.}
Consider an input $\x_1,\dots,\x_L$ such that $\x_i \in \R^d$. An attention head $\Attn$ is defined by matrices $\W_k,\W_q,\W_v\in \R^{d \times d}$ such that $\Attn(\vec{\x})_j = \sum_{i=1}^n \alpha_{ji} \W_v\x_i$, where
\begin{align*}
    \alpha_{ji} := \frac{\exp\left((\W_q \x_j) \cdot (\W_k \x_i)\right)}{\sum_{i=1}^n \exp\left((\W_q \x_j) \cdot (\W_k \x_i)\right)}.
\end{align*}
\begin{remark}
    We remark that for some practical applications, a bias term $B$ will also be added to the computation of attention.
\end{remark}

An attention layer $\AT$ is defined by $H$ attention heads $\Attn_1,\dots,\Attn_H$ and a projection matrix $\W_o \in \R^{d \times dH}$. Denote by $\mathbf{O}:= (\Attn_1(\vec{\x})^\top,\dots,\Attn_H(\vec{\x})^\top)^\top \in \R^{dH \times L}$ the concatenation of the outputs of the $H$ attention heads, so the attention layer $\AT(\vec{\x})$ outputs $\W_o \mathbf{O} \in \R^{d \times L}$.

A Transformer layer $\TF$ is defined by an attention layer $\AT$ and an MLP layer $\MLP$. In particular, an MLP layer is defined as $f(\x): \R^d \to \R^d$ as $f(\x) = \mathbf{U}_2 \sigma(\mathbf{U}_1 \x)$, where $\mathbf{U}_1,\mathbf{U}_2$ are matrices and $\sigma$ is an activation function applied coordinate-wise. Specifically, $\TF(\vec{\x}) = \MLP(\AT(\vec{\x}))$.

\noindent \textbf{State Space Model Layer.}
We use a similar definition of SSM layer as \cite{jelassi2024repeat}.
A state space $\mathcal{S}$ is some finite set. We denote by $\mathrm{mem}(\mathcal{S})$ the number of bits required to encode the states of $\mathcal{S}$, namely $\mathrm{mem}(\mathcal{S}) = \log(|\mathcal{S}|)$. A \textit{generalized state space layer} (GSSM) is a sequence model defined by an update rule $u: \mathcal{S} \times \cV \rightarrow \mathcal{S}$ and some output function $r: \mathcal{S} \rightarrow \cV$. Let $s_0 \in \mathcal{S}$ be some initial state. Given some sequence $\x_1, \ldots, \x_L$, the state of the model at iteration $i$ is denoted by $S_i(\x_1, \ldots, \x_i)$ and the output token is denoted by $R_i(\x_1, \ldots, \x_i)$. The state and output are defined recursively:

\begin{enumerate}
    \item $S_0(\emptyset) = s_0$,
    \item $S_i(\x_1, \ldots, \x_i) = u(S_{i-1}(\x_1, \ldots, \x_{i-1}), \x_i)$,
    \item $R_i(\x_1, \ldots, \x_i) = r(S_i(\x_1, \ldots, \x_i))$.
\end{enumerate}

\noindent \textbf{Mamba Layer.}
% In this work, we will be interested in a specific family of state space models, Mamba. Mamba is defined as follows.
Our SSM layers in our constructions are defined as follows.
Let $\MB\^i$ be a Mamba layer. Let $d$ be the token embedding dimension size and $d_s$ be the state dimension size. 
% be equivalent to the original Mamba formulation. 
Let 
% \begin{align*}
%     \W_A(\x_t) \in \R^{d \times d},
%     \W_B(\x_t) \in \R^{d},
%     \W_C(\x_t) \in \R^{d},
%     \Delta(\x_t) \in \R
% \end{align*}
\begin{align*}
    \W_A \in \R^{d_s \times d_s},
    \W_B \in \R^{d_s \times d},
    \W_C \in \R^{d \times d_s},
    \Delta(\x_t) \in \R
\end{align*}
be constants and a function that returns the shift in time $\Delta(\x_t)$ based on the current token. We update the hidden state matrix as follows:
\begin{align*}
    H_t &= (I - \Delta(\x_t)\W_A) H_{t-1}  + \Delta(\x_t) \W_B \x_t  \\
    y_t &= \W_C H_t
\end{align*}

This follows from the first-order approximation of the S6, where $\exp(-\Delta A) \approx I - \Delta A$. This is useful for construction purposes to avoid large constants $M \gg 1$ to push $\exp$ close to $0$. Equivalent constructions can be found by instead placing a large constant $M$ is specific locations with the original $\exp$.

We also omit the typical per-token dependencies that exist for $\W_B$ and $\W_C$ in standard Mamba, as constant functions were sufficient for our constructions.





In this work, we will consider two kinds of models used in practice, encoder-based models and decoder-based models. We next formally defined the two models as follows. 

\noindent \textbf{Encoder-Based Model.} An Encoder-Based Model $M$ can be thought of as a sequence-to-sequence map. That is to say, given a sequence of input tokens $\vec{x} = (\x_i)_{i=1}^L$, the model $M$ maps it to another sequence $\vec{y} = (y_i)_{i=1}^L$, each token $y_i$ corresponds to a token $\x_i$ in the sequence. Specifically, let $M = F_N \circ F_{N-1} \circ \cdots \circ F_1(\vec{\x})$ be a language model with $N$ layers. Each layer $F_t$ has an input $h^{(t-1)} \in \R^{d\times L}$ and an output $h^{(t)} \in \R^{d \times L}$, in particular $h\^0$ is the embedding of the input sequence $\vec{x}$. Furthermore, for every $t \in [N]$ and $i \in [L]$, $h\^t_i$ is a function of the whole vector $h^{(t-1)}$.

\noindent \textbf{Decoder-Based Model(Autoregressive Model).}
A Decoder-Based Model $M$ can be thought of as an autoregressive(generative) model. Roughly speaking, for each input sequence $\vec{\x}$, we consider recursively generating $\x_{L+i} = M(\x_1,\dots,\x_{L+i-1})$ and denote by $M(\vec{\x}) = \x_{L+1},\x_{L+2},\dots$ the final output of the model $M$. For a multi-layer model $M = F_N \circ F_{N-1} \circ \cdots \circ F_1(\vec{\x})$, the inference stage of $M$ contains two stages. In the first stage, 
each layer $F_t$ has an input $h^{(t-1)} \in \R^{d\times L}$ and an output $h^{(t)} \in \R^{d \times t}$, in particular $h^{(0)}$ is the embedding of the input sequence $\vec{\x}$. However, unless the encoder-based model, $h^{(t)}_i$ only depends on $h^{(t-1)}_1,\dots,h^{(t-1)}_i$. We will decode $h^{(N)}_L$ as $\x_{L+1}$ and enter the second stage, where we generate $\x_{L+2}$ by consuming $(\x_1,\dots,\x_{L+1})$.


\noindent \textbf{Memory Budget.} In this work, we will compare the behavior of different models according to their memory budget. In particular, we will consider two types of budgets: input-dependent memory and input-independent memory. By input-dependent memory, we mean the space needed for storing the input as well as the intermediate results. By input-independent memory, we mean the number of parameters of the model.