% ============================================================================
%  Artificial Neural Networks — Methods, Equations and Graphical
%  Representations
%  Author  : Simon-Pierre Boucher — contact@spboucher.ai
%  Chapter 7 : Modern Transformer Variants (chapters/07-modern-transformers.tex)
% ============================================================================
\chapter{Modern Transformer Variants: ViT, Mixture of Experts and State-Space Models}
\label{chap:modern}

The Transformer of the previous chapter is less a single architecture than a
\emph{substrate}: a stack of residually connected token-mixing and
channel-mixing operations that can be re-instantiated for new modalities,
scaled to extreme parameter counts, or re-engineered to escape its quadratic
cost. This chapter develops the three directions along which that substrate
has evolved most consequentially. First, the \emph{Vision Transformer}
transplants the architecture, essentially unchanged, from token sequences to
images by an embedding trick \cite{dosovitskiy2021}. Second, the
\emph{mixture of experts} decouples parameter count from per-token
computation through sparse conditional routing \cite{shazeer2017}. Third,
\emph{efficient attention} and \emph{selective state-space models} attack
the $O(n^{2}d)$ bottleneck of self-attention itself, the latter replacing
attention with a linear dynamical system whose inference cost is constant
in sequence length \cite{gu2023}. We close with the empirical scaling laws
that govern how all of these models improve with size, data and compute.

% ============================================================================
\section{The Vision Transformer}
\label{sec:mt-vit}

% ----------------------------------------------------------------------------
\subsection{From pixels to tokens}

Self-attention consumes a \emph{sequence} of vectors; an image
$\mathsf{x} \in \R^{H \times W \times C}$ is not one. The Vision
Transformer (ViT) of Dosovitskiy et al.\ \cite{dosovitskiy2021} resolves
the mismatch with a deliberately simple embedding: partition the image into
$N$ non-overlapping square patches of side $P$ (typically $P = 16$),
flatten each patch into a vector, and treat the result as a sequence of
tokens,
\begin{equation}
  \vect{x}_p^{i} \in \R^{P^{2}C},
  \qquad
  i = 1, \dots, N,
  \qquad
  N = \frac{HW}{P^{2}} .
  \label{eq:mt-patches}
\end{equation}
Each flattened patch is mapped to the model dimension $D$ by a single
learned linear projection $\mat{E} \in \R^{(P^{2}C) \times D}$ — the visual
analogue of a word-embedding table. A learnable classification token
$\vect{x}_{\mathrm{class}} \in \R^{D}$ is prepended, and learned positional
embeddings $\mat{E}_{\mathrm{pos}} \in \R^{(N+1) \times D}$ are added to
restore the spatial arrangement that flattening destroyed:
\begin{equation}
  \mat{Z}_0
  =
  \bigl[\,\vect{x}_{\mathrm{class}};\;
        \vect{x}_p^{1}\mat{E};\;
        \vect{x}_p^{2}\mat{E};\;
        \dots;\;
        \vect{x}_p^{N}\mat{E}\,\bigr]
  + \mat{E}_{\mathrm{pos}} .
  \label{eq:mt-embed}
\end{equation}

The sequence $\mat{Z}_0$ then passes through $L$ standard pre-norm
Transformer encoder blocks — multi-head self-attention (MSA) and a
position-wise MLP, each wrapped in layer normalization and a residual
connection \cite{vaswani2017}:
\begin{align}
  \mat{Z}'_{\ell} &= \operatorname{MSA}\bigl(\operatorname{LN}(\mat{Z}_{\ell-1})\bigr)
                     + \mat{Z}_{\ell-1},
  \label{eq:mt-msa}\\
  \mat{Z}_{\ell}  &= \operatorname{MLP}\bigl(\operatorname{LN}(\mat{Z}'_{\ell})\bigr)
                     + \mat{Z}'_{\ell},
  \qquad \ell = 1, \dots, L .
  \label{eq:mt-mlp}
\end{align}
Classification reads out only the final state of the class token,
$\vect{z}_{L}^{0}$:
\begin{equation}
  \hat{\vect{y}}
  =
  \softmax\!\bigl(\operatorname{LN}(\vect{z}_{L}^{0})\,\mat{W}_{\mathrm{head}}\bigr),
  \qquad
  \mat{W}_{\mathrm{head}} \in \R^{D \times K},
  \label{eq:mt-head}
\end{equation}
for $K$ classes. Figure~\ref{fig:mt-vit} traces the full pipeline, and
Algorithm~\ref{alg:mt-vit} states the forward pass together with one
supervised training step.

\begin{figure}[t]
  \centering
  \begin{tikzpicture}[font=\small]
    % ---- input image as 2x2 patch grid ----
    \fill[cinput!15] (0,0)    rectangle (0.8,0.8);
    \fill[cinput!30] (0.8,0)  rectangle (1.6,0.8);
    \fill[cinput!45] (0,0.8)  rectangle (0.8,1.6);
    \fill[cinput!60] (0.8,0.8) rectangle (1.6,1.6);
    \draw[black!70] (0,0) rectangle (1.6,1.6);
    \draw[black!70] (0.8,0) -- (0.8,1.6);
    \draw[black!70] (0,0.8) -- (1.6,0.8);
    \node[etiquette, anchor=north] at (0.8,-0.15) {input image, $N$ patches};

    % ---- flatten arrow ----
    \draw[fleche] (1.8,0.8) -- (2.9,0.8)
      node[midway, above, etiquette, align=center] {flatten +\\ project $\mat{E}$};

    % ---- token row ----
    \node[mem,        minimum width=8mm, minimum height=8mm] (cls) at (3.65,0.8) {$\vect{x}_{\mathrm{cls}}$};
    \node[blochidden, minimum width=8mm, minimum height=8mm] (t1)  at (4.72,0.8) {$\vect{x}_p^{1}\mat{E}$};
    \node[blochidden, minimum width=8mm, minimum height=8mm] (t2)  at (5.79,0.8) {$\vect{x}_p^{2}\mat{E}$};
    \node[blochidden, minimum width=8mm, minimum height=8mm] (t3)  at (6.86,0.8) {$\vect{x}_p^{3}\mat{E}$};
    \node[blochidden, minimum width=8mm, minimum height=8mm] (t4)  at (7.93,0.8) {$\vect{x}_p^{4}\mat{E}$};
    \node[etiquette, anchor=north] at (5.79,0.25)
      {$+\;\mat{E}_{\mathrm{pos}}$ (learned positional embeddings)};
    \node[etiquette, anchor=south] at (3.65,1.35) {\texttt{[CLS]}};

    % ---- encoder ----
    \draw[fleche] (8.5,0.8) -- (9.2,0.8);
    \node[blochidden, minimum width=1.9cm, minimum height=1.7cm, align=center]
      (enc) at (10.35,0.8) {Transformer\\ encoder $\times L$};

    % ---- head ----
    \draw[fleche] (enc.east) -- ++(0.65,0)
      node[midway, above, etiquette] {$\vect{z}_L^0$};
    \node[blocoutput, minimum width=1.3cm, minimum height=8mm]
      (head) at (12.7,0.8) {MLP head};
    \draw[fleche] (head.east) -- ++(0.6,0) node[right] {$\hat{\vect{y}}$};
  \end{tikzpicture}
  \caption{The Vision Transformer \cite{dosovitskiy2021}. The image is cut
  into $N$ non-overlapping patches (here $N = 4$ for legibility), each
  flattened and linearly projected to $\R^{D}$ as in
  \eqref{eq:mt-embed}. A learnable \texttt{[CLS]} token (violet) is
  prepended, learned positional embeddings are added, and the resulting
  sequence flows through $L$ standard encoder blocks,
  eqs.~\eqref{eq:mt-msa}--\eqref{eq:mt-mlp}; the classification
  head \eqref{eq:mt-head} reads only the final \texttt{[CLS]} state.}
  \label{fig:mt-vit}
\end{figure}

\begin{algorithm}[t]
\caption{Vision Transformer: forward pass and one training step}
\label{alg:mt-vit}
\begin{algorithmic}[1]
\Require image $\mathsf{x}$, label $y$, patch size $P$, depth $L$,
         parameters $\theta = \{\mat{E}, \mat{E}_{\mathrm{pos}},
         \vect{x}_{\mathrm{class}}, \text{encoder blocks},
         \mat{W}_{\mathrm{head}}\}$
\Function{ViTForward}{$\mathsf{x}$}
  \State split $\mathsf{x}$ into $N = HW/P^{2}$ patches; flatten each to
         $\vect{x}_p^{i} \in \R^{P^{2}C}$
  \State $\mat{Z}_0 \gets [\vect{x}_{\mathrm{class}};\,
         \vect{x}_p^{1}\mat{E};\, \dots;\, \vect{x}_p^{N}\mat{E}]
         + \mat{E}_{\mathrm{pos}}$
  \For{$\ell = 1, \dots, L$}
    \State $\mat{Z}'_{\ell} \gets
           \operatorname{MSA}(\operatorname{LN}(\mat{Z}_{\ell-1}))
           + \mat{Z}_{\ell-1}$
    \State $\mat{Z}_{\ell} \gets
           \operatorname{MLP}(\operatorname{LN}(\mat{Z}'_{\ell}))
           + \mat{Z}'_{\ell}$
  \EndFor
  \State \Return $\hat{\vect{y}} =
         \softmax(\operatorname{LN}(\vect{z}_{L}^{0})\,\mat{W}_{\mathrm{head}})$
\EndFunction
\State $\hat{\vect{y}} \gets \Call{ViTForward}{\mathsf{x}}$
\State $\Loss \gets -\log \hat{y}_{y}$
       \Comment{cross-entropy on the true class}
\State compute $\nabla_{\theta} \Loss$ by backpropagation
\State update $\theta$ with AdamW (weight decay, warmup, cosine decay)
\end{algorithmic}
\end{algorithm}

% ----------------------------------------------------------------------------
\subsection{What is lost and what is gained}

\begin{remark}[Inductive bias versus data]
\label{rem:mt-inductive}
A convolutional layer hard-wires locality and translation equivariance;
ViT's attention layers assume neither — any patch may interact with any
other from the first layer onward. The architecture therefore
\emph{underperforms} CNNs of similar size when trained on modest datasets,
but \emph{overtakes} them once pre-training data reaches tens or hundreds
of millions of images \cite{dosovitskiy2021}: given enough data, the model
learns locality where it is useful instead of being confined to it, and the
global receptive field is immediate — no stacking of layers is needed to
relate distant patches.
\end{remark}

Two costs follow from the embedding. The attention cost is quadratic in
the number of patches, so halving $P$ quadruples $N$ and multiplies the
attention cost by sixteen; hierarchical variants confine attention to
local windows to recover linear complexity in image size. And discarding
all tokens but the class token in \eqref{eq:mt-head} is a design choice,
not a necessity: dense prediction tasks (detection, segmentation) instead
read out the full sequence $\mat{Z}_L$.

% ============================================================================
\section{Mixture of Experts: Scaling by Sparsity}
\label{sec:mt-moe}

% ----------------------------------------------------------------------------
\subsection{Conditional computation}

In a dense Transformer every parameter participates in every token's
forward pass: doubling the parameters doubles the per-token compute. The
sparsely-gated mixture of experts (MoE) of Shazeer et al.\
\cite{shazeer2017} severs that link.

\begin{definition}[Sparse mixture of experts]
\label{def:mt-moe}
Let $E_1, \dots, E_{E}$ be $E$ \emph{expert} networks (in Transformers,
independent FFN blocks) with identical input and output dimensions, and let
$G : \R^{d} \to \R^{E}$ be a \emph{router} producing a sparse weight vector
with at most $k \ll E$ non-zero entries. The layer computes
$\vect{y} = \sum_{i=1}^{E} G(\vect{x})_i \, E_i(\vect{x})$, evaluating only
the experts for which $G(\vect{x})_i \neq 0$. Parameter count grows with
$E$; per-token compute grows only with $k$.
\end{definition}

The simplest router is a linear map followed by a softmax,
\begin{equation}
  G_{\mathrm{dense}}(\vect{x}) = \softmax\bigl(\mat{W}_g\transp \vect{x}\bigr),
  \qquad \mat{W}_g \in \R^{d \times E},
  \label{eq:mt-gate}
\end{equation}
but \eqref{eq:mt-gate} is dense — every expert receives every token.
Sparsity is obtained by keeping only the $k$ largest logits, after adding
tunable Gaussian noise that encourages exploration and load dispersion
during training:
\begin{equation}
  H(\vect{x})_i
  =
  \bigl(\mat{W}_g\transp \vect{x}\bigr)_i
  + \varepsilon_i \cdot \operatorname{softplus}\!\bigl(
    (\mat{W}_{\mathrm{noise}}\transp \vect{x})_i\bigr),
  \qquad \varepsilon_i \sim \mathcal{N}(0, 1),
  \label{eq:mt-noisy}
\end{equation}
\begin{equation}
  \operatorname{TopK}(\vect{h}, k)_i
  =
  \begin{cases}
    h_i & \text{if } h_i \text{ is among the } k \text{ largest entries of } \vect{h},\\
    -\infty & \text{otherwise},
  \end{cases}
  \label{eq:mt-topk}
\end{equation}
so that the router output and the layer output are
\begin{equation}
  G(\vect{x}) = \softmax\bigl(\operatorname{TopK}(H(\vect{x}), k)\bigr),
  \qquad
  \vect{y} = \sum_{i \,\in\, \mathcal{S}(\vect{x})} G(\vect{x})_i \, E_i(\vect{x}),
  \label{eq:mt-combine}
\end{equation}
where $\mathcal{S}(\vect{x})$ is the selected index set. The $-\infty$
entries vanish under the softmax, so gradients flow only to the selected
experts and to the router itself. Figure~\ref{fig:mt-moe} shows one token's
route through the layer.

\begin{figure}[t]
  \centering
  \begin{tikzpicture}[font=\small]
    \node[blocinput, minimum width=1.4cm] (tok) at (0,0) {token $\vect{x}$};
    \node[gate, minimum width=1.6cm, minimum height=1.0cm, align=center]
      (router) at (2.7,0) {router\\ $G(\vect{x})$};
    \draw[fleche] (tok) -- (router);

    % experts
    \node[blochidden, minimum width=1.7cm] (e1) at (6.3, 2.25) {$E_1$ (FFN)};
    \node[blochidden, minimum width=1.7cm] (e2) at (6.3, 0.75) {$E_2$ (FFN)};
    \node[blochidden, minimum width=1.7cm] (e3) at (6.3,-0.75) {$E_3$ (FFN)};
    \node[blochidden, minimum width=1.7cm] (e4) at (6.3,-2.25) {$E_4$ (FFN)};

    % dispatch: top-2 = experts 2 and 3 solid, others dashed faded
    \draw[flechep, black!30] (router.east) -- (e1.west);
    \draw[fleche, cgate!80!black] (router.east) -- (e2.west)
      node[midway, above, sloped, etiquette] {$g_2$};
    \draw[fleche, cgate!80!black] (router.east) -- (e3.west)
      node[midway, below, sloped, etiquette] {$g_3$};
    \draw[flechep, black!30] (router.east) -- (e4.west);

    % combine
    \node[op] (sum) at (9.6,0) {$+$};
    \draw[flechep, black!30] (e1.east) -- (sum);
    \draw[fleche] (e2.east) -- (sum)
      node[pos=0.32, above, sloped, etiquette] {$g_2 E_2(\vect{x})$};
    \draw[fleche] (e3.east) -- (sum)
      node[pos=0.32, below, sloped, etiquette] {$g_3 E_3(\vect{x})$};
    \draw[flechep, black!30] (e4.east) -- (sum);
    \draw[fleche] (sum.east) -- ++(1.0,0) node[right] {$\vect{y}$};
  \end{tikzpicture}
  \caption{A sparsely-gated mixture-of-experts layer with $E = 4$ experts
  and top-$k$ routing, $k = 2$ \cite{shazeer2017}. For this token the
  router \eqref{eq:mt-combine} selects experts $E_2$ and $E_3$ (solid
  green arrows); the non-selected experts (dashed grey) receive neither
  the token nor any gradient. The output is the gate-weighted sum of the
  two active experts.}
  \label{fig:mt-moe}
\end{figure}

% ----------------------------------------------------------------------------
\subsection{Load balancing and capacity}

Left to itself, the router collapses: a few experts win early, receive
more gradient, and win forever. Training therefore adds an auxiliary loss
that pushes the dispatch distribution toward uniformity. With $f_i$ the
fraction of tokens in a batch whose first choice is expert $i$, and $P_i$
the mean router probability assigned to expert $i$ over the batch,
\begin{equation}
  \Loss_{\mathrm{aux}}
  =
  \alpha \, E \sum_{i=1}^{E} f_i \, P_i ,
  \label{eq:mt-aux}
\end{equation}
which is minimized when both distributions are uniform
($f_i = P_i = 1/E$ gives $\Loss_{\mathrm{aux}} = \alpha$); the
coefficient $\alpha \approx 10^{-2}$ trades balance against task loss.
The product form makes \eqref{eq:mt-aux} differentiable through $P_i$
even though the counts $f_i$ are not.

\begin{remark}[Capacity factor and overflow]
\label{rem:mt-capacity}
On parallel hardware every expert is allocated a fixed buffer of
$\lceil C \cdot kT/E \rceil$ token slots per batch of $T$ tokens, where
$C \geq 1$ is the \emph{capacity factor}. Tokens routed to a full expert
\emph{overflow}: they skip the expert and pass through the residual
connection unchanged. Algorithm~\ref{alg:mt-moe} makes this explicit.
Top-1 routing (the Switch simplification) and top-2 routing are the
dominant regimes; with $E = 8$ and $k = 2$, a model can hold $47$ billion
parameters while activating only ${\sim}13$ billion per token.
\end{remark}

\begin{algorithm}[t]
\caption{MoE layer: noisy top-$k$ routing with capacity factor}
\label{alg:mt-moe}
\begin{algorithmic}[1]
\Require batch of $T$ token vectors $\{\vect{x}_t\}$, experts
         $E_1, \dots, E_E$, router weights $\mat{W}_g,
         \mat{W}_{\mathrm{noise}}$, top-$k$, capacity factor $C$
\State $\mathrm{cap} \gets \lceil C \cdot kT/E \rceil$;\quad
       $\mathrm{load}_i \gets 0$ for $i = 1, \dots, E$
\For{$t = 1, \dots, T$}
  \State $\vect{h}_t \gets$ noisy logits by \eqref{eq:mt-noisy}
  \State $\mathcal{S}_t \gets$ indices of the $k$ largest entries of
         $\vect{h}_t$
  \State $\vect{g}_t \gets \softmax\bigl(\operatorname{TopK}(\vect{h}_t,
         k)\bigr)$ \Comment{eq.~\eqref{eq:mt-topk}}
  \State $\vect{y}_t \gets \vect{0}$
  \For{$i \in \mathcal{S}_t$}
    \If{$\mathrm{load}_i < \mathrm{cap}$}
      \State $\vect{y}_t \gets \vect{y}_t + g_{t,i} \, E_i(\vect{x}_t)$;
             \quad $\mathrm{load}_i \gets \mathrm{load}_i + 1$
    \EndIf \Comment{overflowed tokens rely on the residual path}
  \EndFor
\EndFor
\State add $\Loss_{\mathrm{aux}}$ of \eqref{eq:mt-aux} to the task loss
\end{algorithmic}
\end{algorithm}

% ============================================================================
\section{Efficient Attention}
\label{sec:mt-efficient}

Self-attention over $n$ tokens of width $d$ costs
\begin{equation}
  \underbrace{O(n^{2} d)}_{\text{time}}
  \qquad \text{and} \qquad
  \underbrace{O(n^{2})}_{\text{memory for } \mat{Q}\mat{K}\transp},
  \label{eq:mt-complexity}
\end{equation}
which at $n = 10^{5}$ tokens makes the attention matrix alone prohibitive.
Three families of remedies exist: \emph{restrict} which pairs may
interact, \emph{approximate} the softmax kernel, or \emph{reorganize} the
exact computation around the memory hierarchy.

\paragraph{Sliding-window (local) attention.} Each token attends only to
the $W$ preceding tokens. The additive mask
\begin{equation}
  M_{ij} =
  \begin{cases}
    0 & \text{if } 0 \leq i - j < W,\\
    -\infty & \text{otherwise},
  \end{cases}
  \label{eq:mt-window}
\end{equation}
reduces the cost to $O(nWd)$, and stacking $L$ such layers still yields an
effective receptive field of $L \cdot W$ positions — exactly as stacked
small convolutions enlarge a CNN's receptive field.

\paragraph{Linear (kernelized) attention.} If the exponential kernel of
the softmax is replaced — or approximated — by an inner product of feature
maps, $\exp(\vect{q}\transp\vect{k}) \approx
\phi(\vect{q})\transp\phi(\vect{k})$ with
$\phi : \R^{d} \to \R^{r}$, the attention output factorizes, and the
multiplication order can be changed:
\begin{equation}
  \operatorname{Attn}(\mat{Q}, \mat{K}, \mat{V})
  \approx
  \phi(\mat{Q}) \,
  \bigl(\phi(\mat{K})\transp \mat{V}\bigr),
  \label{eq:mt-linear}
\end{equation}
where the bracketed product is $r \times d$ — independent of $n$ — so the
total cost is $O(nrd)$, linear in sequence length. The price is an
approximation, made unbiased by suitable random-feature constructions for
$\phi$.

\paragraph{Exact IO-aware attention.} A complementary line accelerates
\emph{exact} attention by observing that the bottleneck on modern
accelerators is memory traffic, not arithmetic: tiling
$\mat{Q}, \mat{K}, \mat{V}$ into on-chip blocks, computing the softmax
incrementally with a running maximum and normalizer, and never
materializing the $n \times n$ matrix reduces memory from $O(n^{2})$ to
$O(n)$ with the output unchanged. Table~\ref{tab:mt-complexity} compares
the regimes.

\begin{table}[t]
  \centering
  \caption{Cost of one attention (or mixing) layer over $n$ tokens of
  width $d$; $W$ is the window size of \eqref{eq:mt-window}, $r$ the
  feature dimension of the kernel map in \eqref{eq:mt-linear}, and $N$
  the state size of the SSM in \eqref{eq:mt-selective}.}
  \label{tab:mt-complexity}
  \begin{tabular}{lccc}
    \toprule
    Mechanism & Time & Memory & Exact?\\
    \midrule
    Full softmax attention \cite{vaswani2017} & $O(n^{2}d)$ & $O(n^{2})$ & yes\\
    Sliding window, eq.~\eqref{eq:mt-window}  & $O(nWd)$    & $O(nW)$    & restricted\\
    Linear / kernelized, eq.~\eqref{eq:mt-linear} & $O(nrd)$ & $O(nr)$ & approximate\\
    Tiled exact (IO-aware) & $O(n^{2}d)$ & $O(n)$ & yes\\
    Selective SSM, eq.~\eqref{eq:mt-selective} & $O(nNd)$ & $O(Nd)$ & different model\\
    \bottomrule
  \end{tabular}
\end{table}

% ============================================================================
\section{State-Space Models and Mamba}
\label{sec:mt-ssm}

% ----------------------------------------------------------------------------
\subsection{The linear dynamical view of sequence modeling}

\begin{definition}[State-space model]
\label{def:mt-ssm}
A (continuous-time, linear) state-space model maps an input signal
$u(t) \in \R$ to an output $y(t) \in \R$ through a hidden state
$\vect{h}(t) \in \R^{N}$ obeying
\begin{equation}
  \vect{h}'(t) = \mat{A}\,\vect{h}(t) + \vect{b}\,u(t),
  \qquad
  y(t) = \vect{c}\transp \vect{h}(t),
  \label{eq:mt-ssm-cont}
\end{equation}
with parameters $\mat{A} \in \R^{N \times N}$ and
$\vect{b}, \vect{c} \in \R^{N}$. In deep SSMs each channel of a
$d$-dimensional sequence carries its own scalar system, and the layer is
wrapped in the usual residual and normalization scaffolding.
\end{definition}

To operate on sampled sequences $u_1, u_2, \dots$, the continuous system
\eqref{eq:mt-ssm-cont} is discretized with step size $\Delta$ by the
zero-order hold, exact when $u(t)$ is piecewise constant between samples:
\begin{equation}
  \bar{\mat{A}} = \exp(\Delta \mat{A}),
  \qquad
  \bar{\vect{b}}
  = (\Delta \mat{A})^{-1}\bigl(\exp(\Delta \mat{A}) - \mat{I}\bigr)\,
    \Delta \vect{b},
  \label{eq:mt-zoh}
\end{equation}
yielding the linear recurrence
\begin{equation}
  \vect{h}_t = \bar{\mat{A}}\,\vect{h}_{t-1} + \bar{\vect{b}}\,u_t,
  \qquad
  y_t = \vect{c}\transp \vect{h}_t .
  \label{eq:mt-ssm-disc}
\end{equation}

\begin{property}[An LTI recurrence unrolls into a convolution]
\label{prop:mt-conv}
If $\bar{\mat{A}}, \bar{\vect{b}}, \vect{c}$ do not depend on $t$ (a
linear \emph{time-invariant} system), unrolling \eqref{eq:mt-ssm-disc}
from $\vect{h}_0 = \vect{0}$ gives
$y_t = \sum_{j=0}^{t-1} \vect{c}\transp \bar{\mat{A}}^{\,j}
\bar{\vect{b}}\, u_{t-j}$, i.e.\ a causal convolution
\begin{equation}
  \vect{y} = \vect{u} * \bar{\vect{K}},
  \qquad
  \bar{\vect{K}}
  = \bigl(\vect{c}\transp\bar{\vect{b}},\;
          \vect{c}\transp\bar{\mat{A}}\bar{\vect{b}},\;
          \vect{c}\transp\bar{\mat{A}}^{2}\bar{\vect{b}},\; \dots\bigr),
  \label{eq:mt-kernel}
\end{equation}
computable for a length-$n$ sequence in $O(n \log n)$ by the FFT. The
same model therefore trains \emph{in parallel} as a convolution and runs
inference \emph{recurrently} with $O(1)$ memory per step — a duality that
attention does not possess.
\end{property}

% ----------------------------------------------------------------------------
\subsection{Selectivity: making the dynamics depend on the input}

An LTI system applies the same dynamics to every token: it cannot decide,
based on \emph{content}, what to store and what to forget. The selective
SSM of Mamba \cite{gu2023} breaks time invariance by making the step
size and the input/output projections functions of the current input
$\vect{u}_t \in \R^{d}$:
\begin{equation}
  \Delta_t = \operatorname{softplus}\bigl(\mat{W}_{\Delta}\vect{u}_t\bigr),
  \qquad
  \vect{b}_t = \mat{W}_B \vect{u}_t,
  \qquad
  \vect{c}_t = \mat{W}_C \vect{u}_t,
  \qquad
  \vect{h}_t = \bar{\mat{A}}_t\,\vect{h}_{t-1}
             + \bar{\vect{b}}_t\,u_t,
  \label{eq:mt-selective}
\end{equation}
with $\bar{\mat{A}}_t, \bar{\vect{b}}_t$ obtained from \eqref{eq:mt-zoh}
using $\Delta_t$ and $\vect{b}_t$. A large $\Delta_t$ resets the state
toward the current input (\emph{attend}); $\Delta_t \to 0$ leaves the
state untouched (\emph{ignore}) — a content-dependent gate reminiscent of
the LSTM's, embedded in a principled continuous-time model. Input
dependence destroys the convolutional form \eqref{eq:mt-kernel}, so
training uses a \emph{hardware-aware parallel scan}: the recurrence
\eqref{eq:mt-selective} is associative in the pairs
$(\bar{\mat{A}}_t, \bar{\vect{b}}_t u_t)$, so $n$ steps reduce in
$O(\log n)$ parallel depth with states kept in on-chip memory.
Algorithm~\ref{alg:mt-scan} gives the sequential form, which is also the
constant-memory inference procedure; Figure~\ref{fig:mt-ssm} contrasts
the two mixing mechanisms.

\begin{algorithm}[t]
\caption{Selective scan (sequential form; inference-time recurrence)}
\label{alg:mt-scan}
\begin{algorithmic}[1]
\Require sequence $\vect{u}_1, \dots, \vect{u}_n$, parameters $\mat{A},
         \mat{W}_{\Delta}, \mat{W}_B, \mat{W}_C$
\State $\vect{h}_0 \gets \vect{0}$
\For{$t = 1, \dots, n$}
  \State $\Delta_t \gets \operatorname{softplus}(\mat{W}_{\Delta}\vect{u}_t)$;
         \quad $\vect{b}_t \gets \mat{W}_B \vect{u}_t$;
         \quad $\vect{c}_t \gets \mat{W}_C \vect{u}_t$
  \State $\bar{\mat{A}}_t \gets \exp(\Delta_t \mat{A})$;\quad
         $\bar{\vect{b}}_t \gets (\Delta_t\mat{A})^{-1}
         (\exp(\Delta_t \mat{A}) - \mat{I})\,\Delta_t \vect{b}_t$
         \Comment{ZOH, eq.~\eqref{eq:mt-zoh}}
  \State $\vect{h}_t \gets \bar{\mat{A}}_t \vect{h}_{t-1}
         + \bar{\vect{b}}_t u_t$;\qquad
         $y_t \gets \vect{c}_t\transp \vect{h}_t$
\EndFor
\State \Return $y_1, \dots, y_n$
       \Comment{training uses an associative parallel scan instead}
\end{algorithmic}
\end{algorithm}

\begin{figure}[t]
  \centering
  \begin{tikzpicture}[font=\small]
    % ================= panel (a): SSM recurrence =================
    \begin{scope}
      \node[mem, minimum width=1.0cm] (h1) at (0,0)   {$\vect{h}_1$};
      \node[mem, minimum width=1.0cm] (h2) at (1.8,0) {$\vect{h}_2$};
      \node[mem, minimum width=1.0cm] (h3) at (3.6,0) {$\vect{h}_3$};
      \node[mem, minimum width=1.0cm] (h4) at (5.4,0) {$\vect{h}_4$};
      \draw[fleche, cmem!80!black] (h1) -- (h2)
        node[midway, above, etiquette] {$\bar{\mat{A}}_2$};
      \draw[fleche, cmem!80!black] (h2) -- (h3)
        node[midway, above, etiquette] {$\bar{\mat{A}}_3$};
      \draw[fleche, cmem!80!black] (h3) -- (h4)
        node[midway, above, etiquette] {$\bar{\mat{A}}_4$};
      \foreach \i in {1,...,4} {
        \node[ninput]  (u\i) at ({(\i-1)*1.8}, -1.6) {$u_{\i}$};
        \node[noutput] (y\i) at ({(\i-1)*1.6*1.125},  1.6) {$y_{\i}$};
        \draw[fleche] (u\i) -- (h\i);
        \draw[fleche] (h\i) -- (y\i);
      }
      \node[etiquette, anchor=east] at (-0.15,-0.95) {$\bar{\vect{b}}_t$};
      \node[etiquette, anchor=east] at (-0.15,0.95) {$\vect{c}_t\transp$};
      \node[align=center, font=\small] at (2.7,-2.75)
        {(a) selective SSM: $O(n)$ time,\\ state of fixed size $N$};
    \end{scope}
    % ================= panel (b): attention all-pairs =================
    \begin{scope}[xshift=8.6cm]
      \foreach \i in {1,...,4}
        \node[nhidden] (t\i) at ({(\i-1)*1.8}, 0) {$\vect{x}_{\i}$};
      % causal all-pairs arcs above
      \draw[fleche, chidden!80!black] (t1) to[bend left=45] (t2);
      \draw[fleche, chidden!80!black] (t2) to[bend left=45] (t3);
      \draw[fleche, chidden!80!black] (t3) to[bend left=45] (t4);
      \draw[fleche, chidden!80!black] (t1) to[bend left=55] (t3);
      \draw[fleche, chidden!80!black] (t2) to[bend left=55] (t4);
      \draw[fleche, chidden!80!black] (t1) to[bend left=65] (t4);
      \node[align=center, font=\small] at (2.7,-2.75)
        {(b) causal self-attention:\\ $O(n^{2})$ pairwise interactions};
    \end{scope}
  \end{tikzpicture}
  \caption{Two mechanisms for mixing information along a sequence.
  (a)~The selective state-space recurrence \eqref{eq:mt-selective}:
  inputs $u_t$ enter a fixed-size state $\vect{h}_t$ (violet) that is
  carried forward by input-dependent transitions $\bar{\mat{A}}_t$; the
  cost is linear in length and the inference memory constant.
  (b)~Causal self-attention: every token interacts directly with every
  earlier token, an immediate global view at quadratic cost
  \cite{vaswani2017,gu2023}.}
  \label{fig:mt-ssm}
\end{figure}

\begin{remark}[Where each mechanism wins]
\label{rem:mt-duality}
Attention retrieves \emph{exactly}: any past token can be recalled
verbatim, which underlies in-context learning and copying. An SSM
compresses the past into a fixed-size state — retrieval is lossy, but the
cost is $O(n)$ in time and $O(1)$ in inference memory
(Table~\ref{tab:mt-complexity}). Selective SSMs match Transformers of
roughly twice their size on language modeling \cite{gu2023}, and hybrid
stacks interleaving the two layer types are increasingly common, spending
quadratic attention only where exact retrieval pays for itself.
\end{remark}

% ============================================================================
\section{Scaling Laws}
\label{sec:mt-scaling}

All the architectures of this chapter obey strikingly regular
\emph{scaling laws}: over many orders of magnitude, the test cross-entropy
of an autoregressive Transformer falls as a power law in the parameter
count $N_{\mathrm{par}}$ and the dataset size $D$ (in tokens),
\begin{equation}
  \Loss(N_{\mathrm{par}}) =
  \Bigl(\frac{N_c}{N_{\mathrm{par}}}\Bigr)^{\alpha_N},
  \qquad
  \Loss(D) = \Bigl(\frac{D_c}{D}\Bigr)^{\alpha_D},
  \qquad
  \alpha_N \approx 0.076,\;
  \alpha_D \approx 0.095,
  \label{eq:mt-kaplan}
\end{equation}
when the other factor is not binding. A refined joint parametrization
separates an irreducible entropy of text $E_0$ from two reducible terms,
\begin{equation}
  \Loss(N_{\mathrm{par}}, D)
  =
  E_0 + \frac{A}{N_{\mathrm{par}}^{\alpha}} + \frac{B}{D^{\beta}},
  \qquad
  \alpha \approx 0.34,\;
  \beta \approx 0.28 .
  \label{eq:mt-chinchilla}
\end{equation}

\begin{remark}[Compute-optimal training]
\label{rem:mt-chinchilla}
Training cost is approximately $C \approx 6 N_{\mathrm{par}} D$ FLOPs.
Minimizing \eqref{eq:mt-chinchilla} subject to fixed $C$ gives
$N_{\mathrm{par}}^{\mathrm{opt}} \propto C^{a}$ and
$D^{\mathrm{opt}} \propto C^{b}$ with $a \approx b \approx 0.5$:
parameters and tokens should be scaled \emph{in equal proportion},
roughly twenty tokens per parameter — a sharp correction to earlier
practice, which grew models far faster than their training sets.
Production systems now deliberately train \emph{past} this optimum on
smaller models, accepting extra training compute to reduce the inference
cost that dominates a deployed model's lifetime. The MoE construction of
Section~\ref{sec:mt-moe} bends these laws favourably by growing
$N_{\mathrm{par}}$ without growing per-token compute, and selective SSMs
(Section~\ref{sec:mt-ssm}) do so by removing the $O(n^{2})$ cost of
context length itself.
\end{remark}

Together, the three levers of this chapter — new modalities through
embeddings, more parameters through sparsity, longer contexts through
sub-quadratic mixing — define the current design space of large-scale
neural networks, all resting on the same residual substrate introduced in
the previous chapter.