| \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.} |
| |
| 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. |
| |
| Let |
| |
| |
| |
| |
| |
| |
| \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. |
