\section{Preliminaries and Notations}\label{sec notation}
We provide necessary preliminaries and notation. A complete list of preliminaries is deferred to \Cref{app notations}.

We consider sequence-to-sequence token prediction problems. Let $\cV$ be some vocabulary of tokens, $V = \card{\cV}$, and $\vec{\x}=(\x_i)_{i=1}^L$ be an input sequence. A language model $M$ a 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 $\circ$ is function composition and each $F_i$ is a sequence-to-sequence map called a \emph{layer}.  
We will consider types of layers.


\noindent \textbf{Transformer Layer.}
Consider an embedded input sequence $\vec{\x}=(\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*}
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. 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 formalism of a SSM layer as in \citet{jelassi2024repeat}.
A state space $\mathcal{S}$ is a 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 model} (GSSM) is a layer defined by an update rule $u: \mathcal{S} \times \cV \rightarrow \mathcal{S}$ and an 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 = S_i(\x_1, \ldots, \x_i)$ and the output token is denoted by $r_i=R_i(\x_1, \ldots, \x_i)$. The state and output are defined recursively:
\begin{enumerate}[leftmargin=*, nosep]
    \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{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. \emph{Input-dependent memory}, also called \textit{working memory}, is the size of the intermediate state of the model, applicable to SSMs. \emph{Input-independent memory} is used to characterize the number of parameters in the model.