spb/artificial-neural-networks-book Public
Artificial Neural Networks — Methods, Equations and Graphical Representations: a complete book, every method with rigorous equations, pseudocode and native TikZ figures.
TeX 100%
1% ============================================================================2% Artificial Neural Networks — Methods, Equations and Graphical3% Representations4% Author : Simon-Pierre Boucher — contact@spboucher.ai5% Chapter 7 : Modern Transformer Variants (chapters/07-modern-transformers.tex)6% ============================================================================7\chapter{Modern Transformer Variants: ViT, Mixture of Experts and State-Space Models}8\label{chap:modern}910The Transformer of the previous chapter is less a single architecture than a11\emph{substrate}: a stack of residually connected token-mixing and12channel-mixing operations that can be re-instantiated for new modalities,13scaled to extreme parameter counts, or re-engineered to escape its quadratic14cost. This chapter develops the three directions along which that substrate15has evolved most consequentially. First, the \emph{Vision Transformer}16transplants the architecture, essentially unchanged, from token sequences to17images by an embedding trick \cite{dosovitskiy2021}. Second, the18\emph{mixture of experts} decouples parameter count from per-token19computation through sparse conditional routing \cite{shazeer2017}. Third,20\emph{efficient attention} and \emph{selective state-space models} attack21the $O(n^{2}d)$ bottleneck of self-attention itself, the latter replacing22attention with a linear dynamical system whose inference cost is constant23in sequence length \cite{gu2023}. We close with the empirical scaling laws24that govern how all of these models improve with size, data and compute.2526% ============================================================================27\section{The Vision Transformer}28\label{sec:mt-vit}2930% ----------------------------------------------------------------------------31\subsection{From pixels to tokens}3233Self-attention consumes a \emph{sequence} of vectors; an image34$\mathsf{x} \in \R^{H \times W \times C}$ is not one. The Vision35Transformer (ViT) of Dosovitskiy et al.\ \cite{dosovitskiy2021} resolves36the mismatch with a deliberately simple embedding: partition the image into37$N$ non-overlapping square patches of side $P$ (typically $P = 16$),38flatten each patch into a vector, and treat the result as a sequence of39tokens,40\begin{equation}41 \vect{x}_p^{i} \in \R^{P^{2}C},42 \qquad43 i = 1, \dots, N,44 \qquad45 N = \frac{HW}{P^{2}} .46 \label{eq:mt-patches}47\end{equation}48Each flattened patch is mapped to the model dimension $D$ by a single49learned linear projection $\mat{E} \in \R^{(P^{2}C) \times D}$ — the visual50analogue of a word-embedding table. A learnable classification token51$\vect{x}_{\mathrm{class}} \in \R^{D}$ is prepended, and learned positional52embeddings $\mat{E}_{\mathrm{pos}} \in \R^{(N+1) \times D}$ are added to53restore the spatial arrangement that flattening destroyed:54\begin{equation}55 \mat{Z}_056 =57 \bigl[\,\vect{x}_{\mathrm{class}};\;58 \vect{x}_p^{1}\mat{E};\;59 \vect{x}_p^{2}\mat{E};\;60 \dots;\;61 \vect{x}_p^{N}\mat{E}\,\bigr]62 + \mat{E}_{\mathrm{pos}} .63 \label{eq:mt-embed}64\end{equation}6566The sequence $\mat{Z}_0$ then passes through $L$ standard pre-norm67Transformer encoder blocks — multi-head self-attention (MSA) and a68position-wise MLP, each wrapped in layer normalization and a residual69connection \cite{vaswani2017}:70\begin{align}71 \mat{Z}'_{\ell} &= \operatorname{MSA}\bigl(\operatorname{LN}(\mat{Z}_{\ell-1})\bigr)72 + \mat{Z}_{\ell-1},73 \label{eq:mt-msa}\\74 \mat{Z}_{\ell} &= \operatorname{MLP}\bigl(\operatorname{LN}(\mat{Z}'_{\ell})\bigr)75 + \mat{Z}'_{\ell},76 \qquad \ell = 1, \dots, L .77 \label{eq:mt-mlp}78\end{align}79Classification reads out only the final state of the class token,80$\vect{z}_{L}^{0}$:81\begin{equation}82 \hat{\vect{y}}83 =84 \softmax\!\bigl(\operatorname{LN}(\vect{z}_{L}^{0})\,\mat{W}_{\mathrm{head}}\bigr),85 \qquad86 \mat{W}_{\mathrm{head}} \in \R^{D \times K},87 \label{eq:mt-head}88\end{equation}89for $K$ classes. Figure~\ref{fig:mt-vit} traces the full pipeline, and90Algorithm~\ref{alg:mt-vit} states the forward pass together with one91supervised training step.9293\begin{figure}[t]94 \centering95 \begin{tikzpicture}[font=\small]96 % ---- input image as 2x2 patch grid ----97 \fill[cinput!15] (0,0) rectangle (0.8,0.8);98 \fill[cinput!30] (0.8,0) rectangle (1.6,0.8);99 \fill[cinput!45] (0,0.8) rectangle (0.8,1.6);100 \fill[cinput!60] (0.8,0.8) rectangle (1.6,1.6);101 \draw[black!70] (0,0) rectangle (1.6,1.6);102 \draw[black!70] (0.8,0) -- (0.8,1.6);103 \draw[black!70] (0,0.8) -- (1.6,0.8);104 \node[etiquette, anchor=north] at (0.8,-0.15) {input image, $N$ patches};105106 % ---- flatten arrow ----107 \draw[fleche] (1.8,0.8) -- (2.9,0.8)108 node[midway, above, etiquette, align=center] {flatten +\\ project $\mat{E}$};109110 % ---- token row ----111 \node[mem, minimum width=8mm, minimum height=8mm] (cls) at (3.65,0.8) {$\vect{x}_{\mathrm{cls}}$};112 \node[blochidden, minimum width=8mm, minimum height=8mm] (t1) at (4.72,0.8) {$\vect{x}_p^{1}\mat{E}$};113 \node[blochidden, minimum width=8mm, minimum height=8mm] (t2) at (5.79,0.8) {$\vect{x}_p^{2}\mat{E}$};114 \node[blochidden, minimum width=8mm, minimum height=8mm] (t3) at (6.86,0.8) {$\vect{x}_p^{3}\mat{E}$};115 \node[blochidden, minimum width=8mm, minimum height=8mm] (t4) at (7.93,0.8) {$\vect{x}_p^{4}\mat{E}$};116 \node[etiquette, anchor=north] at (5.79,0.25)117 {$+\;\mat{E}_{\mathrm{pos}}$ (learned positional embeddings)};118 \node[etiquette, anchor=south] at (3.65,1.35) {\texttt{[CLS]}};119120 % ---- encoder ----121 \draw[fleche] (8.5,0.8) -- (9.2,0.8);122 \node[blochidden, minimum width=1.9cm, minimum height=1.7cm, align=center]123 (enc) at (10.35,0.8) {Transformer\\ encoder $\times L$};124125 % ---- head ----126 \draw[fleche] (enc.east) -- ++(0.65,0)127 node[midway, above, etiquette] {$\vect{z}_L^0$};128 \node[blocoutput, minimum width=1.3cm, minimum height=8mm]129 (head) at (12.7,0.8) {MLP head};130 \draw[fleche] (head.east) -- ++(0.6,0) node[right] {$\hat{\vect{y}}$};131 \end{tikzpicture}132 \caption{The Vision Transformer \cite{dosovitskiy2021}. The image is cut133 into $N$ non-overlapping patches (here $N = 4$ for legibility), each134 flattened and linearly projected to $\R^{D}$ as in135 \eqref{eq:mt-embed}. A learnable \texttt{[CLS]} token (violet) is136 prepended, learned positional embeddings are added, and the resulting137 sequence flows through $L$ standard encoder blocks,138 eqs.~\eqref{eq:mt-msa}--\eqref{eq:mt-mlp}; the classification139 head \eqref{eq:mt-head} reads only the final \texttt{[CLS]} state.}140 \label{fig:mt-vit}141\end{figure}142143\begin{algorithm}[t]144\caption{Vision Transformer: forward pass and one training step}145\label{alg:mt-vit}146\begin{algorithmic}[1]147\Require image $\mathsf{x}$, label $y$, patch size $P$, depth $L$,148 parameters $\theta = \{\mat{E}, \mat{E}_{\mathrm{pos}},149 \vect{x}_{\mathrm{class}}, \text{encoder blocks},150 \mat{W}_{\mathrm{head}}\}$151\Function{ViTForward}{$\mathsf{x}$}152 \State split $\mathsf{x}$ into $N = HW/P^{2}$ patches; flatten each to153 $\vect{x}_p^{i} \in \R^{P^{2}C}$154 \State $\mat{Z}_0 \gets [\vect{x}_{\mathrm{class}};\,155 \vect{x}_p^{1}\mat{E};\, \dots;\, \vect{x}_p^{N}\mat{E}]156 + \mat{E}_{\mathrm{pos}}$157 \For{$\ell = 1, \dots, L$}158 \State $\mat{Z}'_{\ell} \gets159 \operatorname{MSA}(\operatorname{LN}(\mat{Z}_{\ell-1}))160 + \mat{Z}_{\ell-1}$161 \State $\mat{Z}_{\ell} \gets162 \operatorname{MLP}(\operatorname{LN}(\mat{Z}'_{\ell}))163 + \mat{Z}'_{\ell}$164 \EndFor165 \State \Return $\hat{\vect{y}} =166 \softmax(\operatorname{LN}(\vect{z}_{L}^{0})\,\mat{W}_{\mathrm{head}})$167\EndFunction168\State $\hat{\vect{y}} \gets \Call{ViTForward}{\mathsf{x}}$169\State $\Loss \gets -\log \hat{y}_{y}$170 \Comment{cross-entropy on the true class}171\State compute $\nabla_{\theta} \Loss$ by backpropagation172\State update $\theta$ with AdamW (weight decay, warmup, cosine decay)173\end{algorithmic}174\end{algorithm}175176% ----------------------------------------------------------------------------177\subsection{What is lost and what is gained}178179\begin{remark}[Inductive bias versus data]180\label{rem:mt-inductive}181A convolutional layer hard-wires locality and translation equivariance;182ViT's attention layers assume neither — any patch may interact with any183other from the first layer onward. The architecture therefore184\emph{underperforms} CNNs of similar size when trained on modest datasets,185but \emph{overtakes} them once pre-training data reaches tens or hundreds186of millions of images \cite{dosovitskiy2021}: given enough data, the model187learns locality where it is useful instead of being confined to it, and the188global receptive field is immediate — no stacking of layers is needed to189relate distant patches.190\end{remark}191192Two costs follow from the embedding. The attention cost is quadratic in193the number of patches, so halving $P$ quadruples $N$ and multiplies the194attention cost by sixteen; hierarchical variants confine attention to195local windows to recover linear complexity in image size. And discarding196all tokens but the class token in \eqref{eq:mt-head} is a design choice,197not a necessity: dense prediction tasks (detection, segmentation) instead198read out the full sequence $\mat{Z}_L$.199200% ============================================================================201\section{Mixture of Experts: Scaling by Sparsity}202\label{sec:mt-moe}203204% ----------------------------------------------------------------------------205\subsection{Conditional computation}206207In a dense Transformer every parameter participates in every token's208forward pass: doubling the parameters doubles the per-token compute. The209sparsely-gated mixture of experts (MoE) of Shazeer et al.\210\cite{shazeer2017} severs that link.211212\begin{definition}[Sparse mixture of experts]213\label{def:mt-moe}214Let $E_1, \dots, E_{E}$ be $E$ \emph{expert} networks (in Transformers,215independent FFN blocks) with identical input and output dimensions, and let216$G : \R^{d} \to \R^{E}$ be a \emph{router} producing a sparse weight vector217with at most $k \ll E$ non-zero entries. The layer computes218$\vect{y} = \sum_{i=1}^{E} G(\vect{x})_i \, E_i(\vect{x})$, evaluating only219the experts for which $G(\vect{x})_i \neq 0$. Parameter count grows with220$E$; per-token compute grows only with $k$.221\end{definition}222223The simplest router is a linear map followed by a softmax,224\begin{equation}225 G_{\mathrm{dense}}(\vect{x}) = \softmax\bigl(\mat{W}_g\transp \vect{x}\bigr),226 \qquad \mat{W}_g \in \R^{d \times E},227 \label{eq:mt-gate}228\end{equation}229but \eqref{eq:mt-gate} is dense — every expert receives every token.230Sparsity is obtained by keeping only the $k$ largest logits, after adding231tunable Gaussian noise that encourages exploration and load dispersion232during training:233\begin{equation}234 H(\vect{x})_i235 =236 \bigl(\mat{W}_g\transp \vect{x}\bigr)_i237 + \varepsilon_i \cdot \operatorname{softplus}\!\bigl(238 (\mat{W}_{\mathrm{noise}}\transp \vect{x})_i\bigr),239 \qquad \varepsilon_i \sim \mathcal{N}(0, 1),240 \label{eq:mt-noisy}241\end{equation}242\begin{equation}243 \operatorname{TopK}(\vect{h}, k)_i244 =245 \begin{cases}246 h_i & \text{if } h_i \text{ is among the } k \text{ largest entries of } \vect{h},\\247 -\infty & \text{otherwise},248 \end{cases}249 \label{eq:mt-topk}250\end{equation}251so that the router output and the layer output are252\begin{equation}253 G(\vect{x}) = \softmax\bigl(\operatorname{TopK}(H(\vect{x}), k)\bigr),254 \qquad255 \vect{y} = \sum_{i \,\in\, \mathcal{S}(\vect{x})} G(\vect{x})_i \, E_i(\vect{x}),256 \label{eq:mt-combine}257\end{equation}258where $\mathcal{S}(\vect{x})$ is the selected index set. The $-\infty$259entries vanish under the softmax, so gradients flow only to the selected260experts and to the router itself. Figure~\ref{fig:mt-moe} shows one token's261route through the layer.262263\begin{figure}[t]264 \centering265 \begin{tikzpicture}[font=\small]266 \node[blocinput, minimum width=1.4cm] (tok) at (0,0) {token $\vect{x}$};267 \node[gate, minimum width=1.6cm, minimum height=1.0cm, align=center]268 (router) at (2.7,0) {router\\ $G(\vect{x})$};269 \draw[fleche] (tok) -- (router);270271 % experts272 \node[blochidden, minimum width=1.7cm] (e1) at (6.3, 2.25) {$E_1$ (FFN)};273 \node[blochidden, minimum width=1.7cm] (e2) at (6.3, 0.75) {$E_2$ (FFN)};274 \node[blochidden, minimum width=1.7cm] (e3) at (6.3,-0.75) {$E_3$ (FFN)};275 \node[blochidden, minimum width=1.7cm] (e4) at (6.3,-2.25) {$E_4$ (FFN)};276277 % dispatch: top-2 = experts 2 and 3 solid, others dashed faded278 \draw[flechep, black!30] (router.east) -- (e1.west);279 \draw[fleche, cgate!80!black] (router.east) -- (e2.west)280 node[midway, above, sloped, etiquette] {$g_2$};281 \draw[fleche, cgate!80!black] (router.east) -- (e3.west)282 node[midway, below, sloped, etiquette] {$g_3$};283 \draw[flechep, black!30] (router.east) -- (e4.west);284285 % combine286 \node[op] (sum) at (9.6,0) {$+$};287 \draw[flechep, black!30] (e1.east) -- (sum);288 \draw[fleche] (e2.east) -- (sum)289 node[pos=0.32, above, sloped, etiquette] {$g_2 E_2(\vect{x})$};290 \draw[fleche] (e3.east) -- (sum)291 node[pos=0.32, below, sloped, etiquette] {$g_3 E_3(\vect{x})$};292 \draw[flechep, black!30] (e4.east) -- (sum);293 \draw[fleche] (sum.east) -- ++(1.0,0) node[right] {$\vect{y}$};294 \end{tikzpicture}295 \caption{A sparsely-gated mixture-of-experts layer with $E = 4$ experts296 and top-$k$ routing, $k = 2$ \cite{shazeer2017}. For this token the297 router \eqref{eq:mt-combine} selects experts $E_2$ and $E_3$ (solid298 green arrows); the non-selected experts (dashed grey) receive neither299 the token nor any gradient. The output is the gate-weighted sum of the300 two active experts.}301 \label{fig:mt-moe}302\end{figure}303304% ----------------------------------------------------------------------------305\subsection{Load balancing and capacity}306307Left to itself, the router collapses: a few experts win early, receive308more gradient, and win forever. Training therefore adds an auxiliary loss309that pushes the dispatch distribution toward uniformity. With $f_i$ the310fraction of tokens in a batch whose first choice is expert $i$, and $P_i$311the mean router probability assigned to expert $i$ over the batch,312\begin{equation}313 \Loss_{\mathrm{aux}}314 =315 \alpha \, E \sum_{i=1}^{E} f_i \, P_i ,316 \label{eq:mt-aux}317\end{equation}318which is minimized when both distributions are uniform319($f_i = P_i = 1/E$ gives $\Loss_{\mathrm{aux}} = \alpha$); the320coefficient $\alpha \approx 10^{-2}$ trades balance against task loss.321The product form makes \eqref{eq:mt-aux} differentiable through $P_i$322even though the counts $f_i$ are not.323324\begin{remark}[Capacity factor and overflow]325\label{rem:mt-capacity}326On parallel hardware every expert is allocated a fixed buffer of327$\lceil C \cdot kT/E \rceil$ token slots per batch of $T$ tokens, where328$C \geq 1$ is the \emph{capacity factor}. Tokens routed to a full expert329\emph{overflow}: they skip the expert and pass through the residual330connection unchanged. Algorithm~\ref{alg:mt-moe} makes this explicit.331Top-1 routing (the Switch simplification) and top-2 routing are the332dominant regimes; with $E = 8$ and $k = 2$, a model can hold $47$ billion333parameters while activating only ${\sim}13$ billion per token.334\end{remark}335336\begin{algorithm}[t]337\caption{MoE layer: noisy top-$k$ routing with capacity factor}338\label{alg:mt-moe}339\begin{algorithmic}[1]340\Require batch of $T$ token vectors $\{\vect{x}_t\}$, experts341 $E_1, \dots, E_E$, router weights $\mat{W}_g,342 \mat{W}_{\mathrm{noise}}$, top-$k$, capacity factor $C$343\State $\mathrm{cap} \gets \lceil C \cdot kT/E \rceil$;\quad344 $\mathrm{load}_i \gets 0$ for $i = 1, \dots, E$345\For{$t = 1, \dots, T$}346 \State $\vect{h}_t \gets$ noisy logits by \eqref{eq:mt-noisy}347 \State $\mathcal{S}_t \gets$ indices of the $k$ largest entries of348 $\vect{h}_t$349 \State $\vect{g}_t \gets \softmax\bigl(\operatorname{TopK}(\vect{h}_t,350 k)\bigr)$ \Comment{eq.~\eqref{eq:mt-topk}}351 \State $\vect{y}_t \gets \vect{0}$352 \For{$i \in \mathcal{S}_t$}353 \If{$\mathrm{load}_i < \mathrm{cap}$}354 \State $\vect{y}_t \gets \vect{y}_t + g_{t,i} \, E_i(\vect{x}_t)$;355 \quad $\mathrm{load}_i \gets \mathrm{load}_i + 1$356 \EndIf \Comment{overflowed tokens rely on the residual path}357 \EndFor358\EndFor359\State add $\Loss_{\mathrm{aux}}$ of \eqref{eq:mt-aux} to the task loss360\end{algorithmic}361\end{algorithm}362363% ============================================================================364\section{Efficient Attention}365\label{sec:mt-efficient}366367Self-attention over $n$ tokens of width $d$ costs368\begin{equation}369 \underbrace{O(n^{2} d)}_{\text{time}}370 \qquad \text{and} \qquad371 \underbrace{O(n^{2})}_{\text{memory for } \mat{Q}\mat{K}\transp},372 \label{eq:mt-complexity}373\end{equation}374which at $n = 10^{5}$ tokens makes the attention matrix alone prohibitive.375Three families of remedies exist: \emph{restrict} which pairs may376interact, \emph{approximate} the softmax kernel, or \emph{reorganize} the377exact computation around the memory hierarchy.378379\paragraph{Sliding-window (local) attention.} Each token attends only to380the $W$ preceding tokens. The additive mask381\begin{equation}382 M_{ij} =383 \begin{cases}384 0 & \text{if } 0 \leq i - j < W,\\385 -\infty & \text{otherwise},386 \end{cases}387 \label{eq:mt-window}388\end{equation}389reduces the cost to $O(nWd)$, and stacking $L$ such layers still yields an390effective receptive field of $L \cdot W$ positions — exactly as stacked391small convolutions enlarge a CNN's receptive field.392393\paragraph{Linear (kernelized) attention.} If the exponential kernel of394the softmax is replaced — or approximated — by an inner product of feature395maps, $\exp(\vect{q}\transp\vect{k}) \approx396\phi(\vect{q})\transp\phi(\vect{k})$ with397$\phi : \R^{d} \to \R^{r}$, the attention output factorizes, and the398multiplication order can be changed:399\begin{equation}400 \operatorname{Attn}(\mat{Q}, \mat{K}, \mat{V})401 \approx402 \phi(\mat{Q}) \,403 \bigl(\phi(\mat{K})\transp \mat{V}\bigr),404 \label{eq:mt-linear}405\end{equation}406where the bracketed product is $r \times d$ — independent of $n$ — so the407total cost is $O(nrd)$, linear in sequence length. The price is an408approximation, made unbiased by suitable random-feature constructions for409$\phi$.410411\paragraph{Exact IO-aware attention.} A complementary line accelerates412\emph{exact} attention by observing that the bottleneck on modern413accelerators is memory traffic, not arithmetic: tiling414$\mat{Q}, \mat{K}, \mat{V}$ into on-chip blocks, computing the softmax415incrementally with a running maximum and normalizer, and never416materializing the $n \times n$ matrix reduces memory from $O(n^{2})$ to417$O(n)$ with the output unchanged. Table~\ref{tab:mt-complexity} compares418the regimes.419420\begin{table}[t]421 \centering422 \caption{Cost of one attention (or mixing) layer over $n$ tokens of423 width $d$; $W$ is the window size of \eqref{eq:mt-window}, $r$ the424 feature dimension of the kernel map in \eqref{eq:mt-linear}, and $N$425 the state size of the SSM in \eqref{eq:mt-selective}.}426 \label{tab:mt-complexity}427 \begin{tabular}{lccc}428 \toprule429 Mechanism & Time & Memory & Exact?\\430 \midrule431 Full softmax attention \cite{vaswani2017} & $O(n^{2}d)$ & $O(n^{2})$ & yes\\432 Sliding window, eq.~\eqref{eq:mt-window} & $O(nWd)$ & $O(nW)$ & restricted\\433 Linear / kernelized, eq.~\eqref{eq:mt-linear} & $O(nrd)$ & $O(nr)$ & approximate\\434 Tiled exact (IO-aware) & $O(n^{2}d)$ & $O(n)$ & yes\\435 Selective SSM, eq.~\eqref{eq:mt-selective} & $O(nNd)$ & $O(Nd)$ & different model\\436 \bottomrule437 \end{tabular}438\end{table}439440% ============================================================================441\section{State-Space Models and Mamba}442\label{sec:mt-ssm}443444% ----------------------------------------------------------------------------445\subsection{The linear dynamical view of sequence modeling}446447\begin{definition}[State-space model]448\label{def:mt-ssm}449A (continuous-time, linear) state-space model maps an input signal450$u(t) \in \R$ to an output $y(t) \in \R$ through a hidden state451$\vect{h}(t) \in \R^{N}$ obeying452\begin{equation}453 \vect{h}'(t) = \mat{A}\,\vect{h}(t) + \vect{b}\,u(t),454 \qquad455 y(t) = \vect{c}\transp \vect{h}(t),456 \label{eq:mt-ssm-cont}457\end{equation}458with parameters $\mat{A} \in \R^{N \times N}$ and459$\vect{b}, \vect{c} \in \R^{N}$. In deep SSMs each channel of a460$d$-dimensional sequence carries its own scalar system, and the layer is461wrapped in the usual residual and normalization scaffolding.462\end{definition}463464To operate on sampled sequences $u_1, u_2, \dots$, the continuous system465\eqref{eq:mt-ssm-cont} is discretized with step size $\Delta$ by the466zero-order hold, exact when $u(t)$ is piecewise constant between samples:467\begin{equation}468 \bar{\mat{A}} = \exp(\Delta \mat{A}),469 \qquad470 \bar{\vect{b}}471 = (\Delta \mat{A})^{-1}\bigl(\exp(\Delta \mat{A}) - \mat{I}\bigr)\,472 \Delta \vect{b},473 \label{eq:mt-zoh}474\end{equation}475yielding the linear recurrence476\begin{equation}477 \vect{h}_t = \bar{\mat{A}}\,\vect{h}_{t-1} + \bar{\vect{b}}\,u_t,478 \qquad479 y_t = \vect{c}\transp \vect{h}_t .480 \label{eq:mt-ssm-disc}481\end{equation}482483\begin{property}[An LTI recurrence unrolls into a convolution]484\label{prop:mt-conv}485If $\bar{\mat{A}}, \bar{\vect{b}}, \vect{c}$ do not depend on $t$ (a486linear \emph{time-invariant} system), unrolling \eqref{eq:mt-ssm-disc}487from $\vect{h}_0 = \vect{0}$ gives488$y_t = \sum_{j=0}^{t-1} \vect{c}\transp \bar{\mat{A}}^{\,j}489\bar{\vect{b}}\, u_{t-j}$, i.e.\ a causal convolution490\begin{equation}491 \vect{y} = \vect{u} * \bar{\vect{K}},492 \qquad493 \bar{\vect{K}}494 = \bigl(\vect{c}\transp\bar{\vect{b}},\;495 \vect{c}\transp\bar{\mat{A}}\bar{\vect{b}},\;496 \vect{c}\transp\bar{\mat{A}}^{2}\bar{\vect{b}},\; \dots\bigr),497 \label{eq:mt-kernel}498\end{equation}499computable for a length-$n$ sequence in $O(n \log n)$ by the FFT. The500same model therefore trains \emph{in parallel} as a convolution and runs501inference \emph{recurrently} with $O(1)$ memory per step — a duality that502attention does not possess.503\end{property}504505% ----------------------------------------------------------------------------506\subsection{Selectivity: making the dynamics depend on the input}507508An LTI system applies the same dynamics to every token: it cannot decide,509based on \emph{content}, what to store and what to forget. The selective510SSM of Mamba \cite{gu2023} breaks time invariance by making the step511size and the input/output projections functions of the current input512$\vect{u}_t \in \R^{d}$:513\begin{equation}514 \Delta_t = \operatorname{softplus}\bigl(\mat{W}_{\Delta}\vect{u}_t\bigr),515 \qquad516 \vect{b}_t = \mat{W}_B \vect{u}_t,517 \qquad518 \vect{c}_t = \mat{W}_C \vect{u}_t,519 \qquad520 \vect{h}_t = \bar{\mat{A}}_t\,\vect{h}_{t-1}521 + \bar{\vect{b}}_t\,u_t,522 \label{eq:mt-selective}523\end{equation}524with $\bar{\mat{A}}_t, \bar{\vect{b}}_t$ obtained from \eqref{eq:mt-zoh}525using $\Delta_t$ and $\vect{b}_t$. A large $\Delta_t$ resets the state526toward the current input (\emph{attend}); $\Delta_t \to 0$ leaves the527state untouched (\emph{ignore}) — a content-dependent gate reminiscent of528the LSTM's, embedded in a principled continuous-time model. Input529dependence destroys the convolutional form \eqref{eq:mt-kernel}, so530training uses a \emph{hardware-aware parallel scan}: the recurrence531\eqref{eq:mt-selective} is associative in the pairs532$(\bar{\mat{A}}_t, \bar{\vect{b}}_t u_t)$, so $n$ steps reduce in533$O(\log n)$ parallel depth with states kept in on-chip memory.534Algorithm~\ref{alg:mt-scan} gives the sequential form, which is also the535constant-memory inference procedure; Figure~\ref{fig:mt-ssm} contrasts536the two mixing mechanisms.537538\begin{algorithm}[t]539\caption{Selective scan (sequential form; inference-time recurrence)}540\label{alg:mt-scan}541\begin{algorithmic}[1]542\Require sequence $\vect{u}_1, \dots, \vect{u}_n$, parameters $\mat{A},543 \mat{W}_{\Delta}, \mat{W}_B, \mat{W}_C$544\State $\vect{h}_0 \gets \vect{0}$545\For{$t = 1, \dots, n$}546 \State $\Delta_t \gets \operatorname{softplus}(\mat{W}_{\Delta}\vect{u}_t)$;547 \quad $\vect{b}_t \gets \mat{W}_B \vect{u}_t$;548 \quad $\vect{c}_t \gets \mat{W}_C \vect{u}_t$549 \State $\bar{\mat{A}}_t \gets \exp(\Delta_t \mat{A})$;\quad550 $\bar{\vect{b}}_t \gets (\Delta_t\mat{A})^{-1}551 (\exp(\Delta_t \mat{A}) - \mat{I})\,\Delta_t \vect{b}_t$552 \Comment{ZOH, eq.~\eqref{eq:mt-zoh}}553 \State $\vect{h}_t \gets \bar{\mat{A}}_t \vect{h}_{t-1}554 + \bar{\vect{b}}_t u_t$;\qquad555 $y_t \gets \vect{c}_t\transp \vect{h}_t$556\EndFor557\State \Return $y_1, \dots, y_n$558 \Comment{training uses an associative parallel scan instead}559\end{algorithmic}560\end{algorithm}561562\begin{figure}[t]563 \centering564 \begin{tikzpicture}[font=\small]565 % ================= panel (a): SSM recurrence =================566 \begin{scope}567 \node[mem, minimum width=1.0cm] (h1) at (0,0) {$\vect{h}_1$};568 \node[mem, minimum width=1.0cm] (h2) at (1.8,0) {$\vect{h}_2$};569 \node[mem, minimum width=1.0cm] (h3) at (3.6,0) {$\vect{h}_3$};570 \node[mem, minimum width=1.0cm] (h4) at (5.4,0) {$\vect{h}_4$};571 \draw[fleche, cmem!80!black] (h1) -- (h2)572 node[midway, above, etiquette] {$\bar{\mat{A}}_2$};573 \draw[fleche, cmem!80!black] (h2) -- (h3)574 node[midway, above, etiquette] {$\bar{\mat{A}}_3$};575 \draw[fleche, cmem!80!black] (h3) -- (h4)576 node[midway, above, etiquette] {$\bar{\mat{A}}_4$};577 \foreach \i in {1,...,4} {578 \node[ninput] (u\i) at ({(\i-1)*1.8}, -1.6) {$u_{\i}$};579 \node[noutput] (y\i) at ({(\i-1)*1.6*1.125}, 1.6) {$y_{\i}$};580 \draw[fleche] (u\i) -- (h\i);581 \draw[fleche] (h\i) -- (y\i);582 }583 \node[etiquette, anchor=east] at (-0.15,-0.95) {$\bar{\vect{b}}_t$};584 \node[etiquette, anchor=east] at (-0.15,0.95) {$\vect{c}_t\transp$};585 \node[align=center, font=\small] at (2.7,-2.75)586 {(a) selective SSM: $O(n)$ time,\\ state of fixed size $N$};587 \end{scope}588 % ================= panel (b): attention all-pairs =================589 \begin{scope}[xshift=8.6cm]590 \foreach \i in {1,...,4}591 \node[nhidden] (t\i) at ({(\i-1)*1.8}, 0) {$\vect{x}_{\i}$};592 % causal all-pairs arcs above593 \draw[fleche, chidden!80!black] (t1) to[bend left=45] (t2);594 \draw[fleche, chidden!80!black] (t2) to[bend left=45] (t3);595 \draw[fleche, chidden!80!black] (t3) to[bend left=45] (t4);596 \draw[fleche, chidden!80!black] (t1) to[bend left=55] (t3);597 \draw[fleche, chidden!80!black] (t2) to[bend left=55] (t4);598 \draw[fleche, chidden!80!black] (t1) to[bend left=65] (t4);599 \node[align=center, font=\small] at (2.7,-2.75)600 {(b) causal self-attention:\\ $O(n^{2})$ pairwise interactions};601 \end{scope}602 \end{tikzpicture}603 \caption{Two mechanisms for mixing information along a sequence.604 (a)~The selective state-space recurrence \eqref{eq:mt-selective}:605 inputs $u_t$ enter a fixed-size state $\vect{h}_t$ (violet) that is606 carried forward by input-dependent transitions $\bar{\mat{A}}_t$; the607 cost is linear in length and the inference memory constant.608 (b)~Causal self-attention: every token interacts directly with every609 earlier token, an immediate global view at quadratic cost610 \cite{vaswani2017,gu2023}.}611 \label{fig:mt-ssm}612\end{figure}613614\begin{remark}[Where each mechanism wins]615\label{rem:mt-duality}616Attention retrieves \emph{exactly}: any past token can be recalled617verbatim, which underlies in-context learning and copying. An SSM618compresses the past into a fixed-size state — retrieval is lossy, but the619cost is $O(n)$ in time and $O(1)$ in inference memory620(Table~\ref{tab:mt-complexity}). Selective SSMs match Transformers of621roughly twice their size on language modeling \cite{gu2023}, and hybrid622stacks interleaving the two layer types are increasingly common, spending623quadratic attention only where exact retrieval pays for itself.624\end{remark}625626% ============================================================================627\section{Scaling Laws}628\label{sec:mt-scaling}629630All the architectures of this chapter obey strikingly regular631\emph{scaling laws}: over many orders of magnitude, the test cross-entropy632of an autoregressive Transformer falls as a power law in the parameter633count $N_{\mathrm{par}}$ and the dataset size $D$ (in tokens),634\begin{equation}635 \Loss(N_{\mathrm{par}}) =636 \Bigl(\frac{N_c}{N_{\mathrm{par}}}\Bigr)^{\alpha_N},637 \qquad638 \Loss(D) = \Bigl(\frac{D_c}{D}\Bigr)^{\alpha_D},639 \qquad640 \alpha_N \approx 0.076,\;641 \alpha_D \approx 0.095,642 \label{eq:mt-kaplan}643\end{equation}644when the other factor is not binding. A refined joint parametrization645separates an irreducible entropy of text $E_0$ from two reducible terms,646\begin{equation}647 \Loss(N_{\mathrm{par}}, D)648 =649 E_0 + \frac{A}{N_{\mathrm{par}}^{\alpha}} + \frac{B}{D^{\beta}},650 \qquad651 \alpha \approx 0.34,\;652 \beta \approx 0.28 .653 \label{eq:mt-chinchilla}654\end{equation}655656\begin{remark}[Compute-optimal training]657\label{rem:mt-chinchilla}658Training cost is approximately $C \approx 6 N_{\mathrm{par}} D$ FLOPs.659Minimizing \eqref{eq:mt-chinchilla} subject to fixed $C$ gives660$N_{\mathrm{par}}^{\mathrm{opt}} \propto C^{a}$ and661$D^{\mathrm{opt}} \propto C^{b}$ with $a \approx b \approx 0.5$:662parameters and tokens should be scaled \emph{in equal proportion},663roughly twenty tokens per parameter — a sharp correction to earlier664practice, which grew models far faster than their training sets.665Production systems now deliberately train \emph{past} this optimum on666smaller models, accepting extra training compute to reduce the inference667cost that dominates a deployed model's lifetime. The MoE construction of668Section~\ref{sec:mt-moe} bends these laws favourably by growing669$N_{\mathrm{par}}$ without growing per-token compute, and selective SSMs670(Section~\ref{sec:mt-ssm}) do so by removing the $O(n^{2})$ cost of671context length itself.672\end{remark}673674Together, the three levers of this chapter — new modalities through675embeddings, more parameters through sparsity, longer contexts through676sub-quadratic mixing — define the current design space of large-scale677neural networks, all resting on the same residual substrate introduced in678the previous chapter.679