NN LAB 🧠

Recurrent neural networks (RNNs) historically suffered from either vanishing gradients (fading ancient memory) or exploding gradients. The HiPPO (High-order Polynomial Projection Operators; Gu et al., NeurIPS 2020) framework solved this fundamental challenge by casting online memory as optimal continuous function approximation.

1. Continuous History Projection onto Legendre Polynomials

Consider a continuous 1D input signal $u(t)$ observed up to current time $t$. We seek to maintain a degree-$N$ polynomial approximation $g^{(t)}(\tau)$ of the history $u(\tau)$ for $\tau \in [0, t]$ with respect to a time-varying probability measure $\mu^{(t)}(\tau) = \frac{1}{t} \mathbb{I}_{[0, t]}(\tau)$:

Expressing $g^{(t)}$ in the orthogonal basis of shifted Legendre polynomials $\{P_n\}_{n=0}^{N-1}$, the projection coefficients$c(t) \in \mathbb{R}^N$ satisfy a differential equation that guarantees optimal online coefficient updates without storing past values:

2. Canonical HiPPO Transition Matrix

Converting this time-varying system into a time-invariant state-space representation yields the canonical continuous HiPPO-LegS matrix:

Why this prevents memory fading: Standard RNN transition matrices $W_h$ compress history through repeated unconstrained matrix multiplications, causing geometric decay $\sim \lambda^t$. In contrast, HiPPO projects continuous history onto orthogonal Legendre bases, ensuring that ancient signals retain non-zero projection coefficients over unbounded sequence horizons. In Mamba, the continuous transition matrix is initialized via structured diagonal decomposition with negative real poles ($A_n = -n - 1$) derived from the HiPPO spectrum.

State Space Models describe continuous dynamical systems governed by ordinary differential equations (ODEs):

Digital computing architectures operate on discrete sequences $\{x_0, x_1, \dots, x_{T-1}\}$ sampled at intervals parameterized by step size $\Delta$. The Zero-Order Hold (ZOH) assumption postulates that the input $x(t)$ remains piecewise-constant over each sampling interval $[t\Delta, (t+1)\Delta)$.

1. Derivation of the Discrete Transition Matrices

Solving the linear continuous ODE over $\tau \in [0, \Delta]$ using the matrix exponential integrating factor:

Because $x(t)$ is constant over the interval, we factor it out of the integral:

Defining discrete transition matrices $\bar{A}$ and $\bar{B}$:

For diagonal $A = \text{diag}(a_1, \dots, a_N)$, each dimension evaluates independently as:

2. Proof of Spectral Stability & Contraction Mapping

Let $\lambda_i(A)$ denote the continuous eigenvalues of $A$. By HiPPO construction, continuous poles lie strictly in the open left half-plane:$\text{Re}(\lambda_i(A)) < 0$.

Under the exponential mapping, the eigenvalues of discrete operator $\bar{A}$ are:

Because all discrete eigenvalues have magnitude strictly less than 1, $\bar{A}$ is a strict contraction mapping on $\mathbb{R}^N$. Consequently, the latent state remains strictly bounded for arbitrarily long sequence lengths ($\|h_t\| < \infty$ as $t \to \infty$), preventing gradient explosions during training and unbounded state drift during autoregressive inference.

While Transformers revolutionized deep learning by parallelizing training over time, their quadratic attention matrix $O(T^2)$ creates severe memory and computational bottlenecks for long sequences. Conversely, standard RNNs have linear complexity $O(T)$ but failed because sequential recurrent steps cannot parallelize across modern GPU architectures.

1. The Modern GPU Memory Hierarchy Dilemma

Modern accelerators (e.g. Nvidia H100 SXM) have two primary memory tiers:

• High-Bandwidth Memory (HBM): Massive capacity (80–96 GB) but comparatively limited bandwidth (~3.35 TB/s).
• On-Chip Static RAM (SRAM): Ultra-fast bandwidth (~33 TB/s, $10\times$ faster) but minimal capacity (~228 KB per SM, ~50 MB total).

Standard RNNs require reading and writing state vectors $h_t$ between HBM and SRAM at every sequential token step $t$. The arithmetic intensity is negligible ($< 1$ FLOP/byte), stalling tensor execution units while waiting on HBM memory transfers.

2. The Associative Property of Linear Recurrence

Mamba makes the critical realization that linear state space recurrence is an associative operator. Consider adjacent state transitions:

Defining the binary tuple operator $\circ$ on elements $(p_i, q_i) = (\bar{A}_i, \bar{B}_i x_i)$:

Because matrix multiplication and addition are associative, $((a \circ b) \circ c) = (a \circ (b \circ c))$. This enables computing the entire sequence recurrence using the Blelloch Parallel Prefix Scan in $O(\log T)$ parallel time steps across GPU threads.

3. Kernel Fusion in On-Chip SRAM

Mamba implements an optimized Triton/CUDA kernel that fuses discretization, causal 1D convolution, and parallel scan directly in on-chip SRAM:

In Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality (Dao & Gu, 2024), the authors establish an exact theoretical equivalence between linear attention and structured state space models.

1. The 1-Semiseparable Matrix Connection

Unrolling an unforced SSM over sequence length $T$ reveals that the transformation from input sequence $X \in \mathbb{R}^{T \times 1}$ to output sequence $Y \in \mathbb{R}^{T \times 1}$ can be expressed as multiplication by a lower-triangular structured matrix $M \in \mathbb{R}^{T \times T}$:

This matrix is 1-semiseparable: every submatrix below the main diagonal has rank at most 1 (or rank $N$ in multi-channel SSMs). Linear attention computes a special case of this transformation where $\bar{A}_k = I$ (no decay):

2. Bridging Tensor Cores and Recurrence

In Mamba-1, general diagonal $A$ required custom associative scan kernels that underutilized GPU Tensor Cores (specialized matrix multiply units). Mamba-2 restricts $A$ to scalar-times-identity structure ($A = \alpha I$). This enables decomposing the semiseparable matrix into block matrix multiplications:

• Intra-chunk computation (e.g. 64 tokens): Evaluated using dense matrix multiplications ($Q K^T V$ style) on Tensor Cores with peak compute efficiency.
• Inter-chunk computation: Carried across chunk boundaries via recurrent state propagation in linear time.

The Result: Mamba-2 matches the expressive capacity of selective SSMs while training up to $8\times$ faster than Mamba-1 by computing the core recurrence directly on Tensor Cores, unifying Attention and SSM architectures into a single foundational framework.