Revisiting SSM (III): Efficient Computation of HiPPO (S4)
In the previous two posts, Revisiting SSM (I): Linear Systems and the HiPPO Matrix and Revisiting SSM (II): Some Remaining Issues with HiPPO, we introduced the ideas and derivations behind HiPPO — using an orthogonal function basis to approximate a continuously updated function in real time, whose fitting-coefficient dynamics turn out to be expressible as a linear ODE system, and for a particular choice of basis and approximation scheme, the key matrix of this linear system can be computed exactly. We also discussed the discretization of HiPPO and some related properties, laying the theoretical groundwork for the subsequent line of SSM work.
Next, we'll introduce a follow-up application of HiPPO, Efficiently Modeling Long Sequences with Structured State Spaces (S4 for short). It uses the results derived for HiPPO as a basic tool for sequence modeling, and explores efficient computation and training from a fresh angle, ultimately validating its effectiveness on a number of long-sequence modeling tasks. It's fair to call it one of the landmark works in the revival of SSMs and, more broadly, RNNs.
Basic Framework
The sequence modeling framework used by S4 is the following linear ODE system:
\begin{equation}\begin{aligned} x'(t) =&\, A x(t) + B u(t) \\ y(t) =&\, C^* x(t) + D u(t) \end{aligned}\end{equation}more
Here $u,y,D\in\mathbb{R};x\in\mathbb{R}^d;A\in\mathbb{R}^{d\times d};B,C\in\mathbb{R}^{d\times 1}$, and ${}^*$ denotes the conjugate transpose — for a real matrix, this is simply the transpose. Since a full model typically also includes a residual connection, the last term $D u(t)$ can be absorbed into the residual, so we can just assume $D=0$ to slightly simplify the form, without reducing the model's expressive power.
This system has a similarity invariance property: if $\tilde{A}$ is a similarity transform of $A$, i.e. $A = P^{-1}\tilde{A}P$, then substituting and rearranging gives
\begin{equation}\begin{aligned} Px'(t) =&\, \tilde{A} Px(t) + PB u(t) \\ y(t) =&\, ((P^{-1})^* C)^* P x(t) \end{aligned}\end{equation}
Treating $Px(t)$ as a whole to replace the original $x(t)$, the new system changes as $(A,B,C)\to(\tilde{A},PB,(P^{-1})^*C)$, but the output remains entirely unchanged. This means that if there exists some similarity matrix $\tilde{A}$ of $A$ that makes the computation simpler, we can shift the analysis entirely to $\tilde{A}$ without changing the result — this is the core idea behind the whole series of analyses that follows.
In particular, S4 chooses the matrix $A$ to be the HiPPO-LegS matrix, i.e.
\begin{equation}A_{n,k} = -\left\{\begin{array}{l}\sqrt{(2n+1)(2k+1)}, &k < n \\ n+1, &k = n \\ 0, &k > n\end{array}\right.\end{equation}
What's special about this choice is that, as we derived earlier, the ODE satisfied by LegS takes the form $x'(t) =\frac{A}{t} x(t) + \frac{B}{t} u(t)$, while the ODE for LegT takes the form $x'(t) = A x(t) + B u(t)$. So here we're pairing the LegT-style ODE with the LegS-style $A$ matrix, and the first question we should ask is: what effect does this combination have? For instance, is its memory of history still complete and equally weighted, the way LegS's is?
Exponential Decay
The answer is no — the ODE system chosen by S4 decays exponentially with respect to history. We can understand this from two angles.
The first angle starts from the transformation discussed in Revisiting SSM (II): Some Remaining Issues with HiPPO, where the LegS-type ODE can be equivalently rewritten as:
\begin{equation}Ax(t) + Bu(t) = t x'(t) = \frac{d}{d\ln t} x(t)\end{equation}
So setting $\tau=\ln t$ turns the LegS-type ODE into a LegT-type ODE with time variable $\tau$ — which is exactly the ODE used by S4. We know that LegS treats every point in history equally, but this is under the premise that the input is $u(t)=u(e^{\tau})$; the ODE used by S4, however, effectively changes the input directly to $u(\tau)$. If we now perform uniform discretization on $\tau$, the resulting weights are no longer equal — assuming $t\in[0,T]$, in terms of probability density this is $dt/T=\rho(\tau)d\tau$, i.e. $\rho(\tau)=e^{\tau}/T$, i.e. the weight is an exponential function of $\tau$, so more recent history gets a larger weight.
The second angle requires a bit more linear algebra. Also in Revisiting SSM (II): Some Remaining Issues with HiPPO, we noted that the HiPPO-LegS matrix $A$ is, in principle, diagonalizable, with eigenvalues $[-1,-2,-3,\cdots]$. So there exists an invertible matrix $P$ such that $A = P^{-1}\Lambda P$, where $\Lambda = \text{diag}(-1,-2,\cdots,-d)$. By similarity invariance, the original system is equivalent to the new system
\begin{equation}\begin{aligned} x'(t) =&\, \Lambda x(t) + PB u(t) \\ y(t) =&\, C^* P^{-1} x(t) \end{aligned}\end{equation}
After discretization (using forward Euler as an example):
\begin{equation}x(t+\epsilon) = (I + \epsilon\Lambda) Px(t) + \epsilon P B u(t)\end{equation}
Here $I + \epsilon\Lambda$ is a diagonal matrix with all entries less than 1 in magnitude, meaning that with each iteration, historical information is multiplied by a factor less than 1. Stacking up many steps, this produces an exponential decay effect.
Discretization Scheme
Although exponential decay may seem less elegant than LegS's equal treatment of all history, there's really no free lunch: for a fixed-size memory state $x(t)$, as the memory window grows, LegS's policy of treating every point in history equally ends up blurring every point in history, which is actually a poor fit for scenarios where "recent matters more, distant matters less." Furthermore, the right-hand side of the S4-type ODE has no explicit dependence on time $t$, which also helps training efficiency.
Now that we have a clear picture of the memory properties of the S4-type ODE, we can move to the next step. To handle actual discrete sequences, we first need to discretize. In the previous post we gave two high-precision discretization schemes. One is the bilinear form
\begin{equation}x_{k+1} = (I - \epsilon A/2)^{-1}[(I + \epsilon A/2) x_k + \epsilon B u_k] \end{equation}
which has second-order accuracy — this is the scheme S4 adopts, and the one we'll be examining in this post. The other is based on exactly solving the ODE for a constant input, giving
\begin{equation}x_{k+1} = e^{\epsilon A} x_k + A^{-1} (e^{\epsilon A} - I) B u_k\end{equation}
The author's later works, including Mamba, use this format instead, and it generally requires assuming that $A$ is a diagonal matrix, since computing the matrix exponential for the LegS matrix $A$ is not particularly convenient.
Now let's define:
\begin{equation}\bar{A}=(I - \epsilon A/2)^{-1}(I + \epsilon A/2),\quad\bar{B}=\epsilon(I - \epsilon A/2)^{-1}B,\quad\bar{C}=C\end{equation}
which gives us the linear RNN:
\begin{equation}\begin{aligned} x_{k+1} =&\, \bar{A} x_k + \bar{B} u_k \\ y_{k+1} =&\, \bar{C}^* x_{k+1} \\ \end{aligned}\label{eq:s4-r}\end{equation}
where $\epsilon > 0$ is the discretization step size, a manually chosen hyperparameter.
Convolution
In the previous post we also mentioned that the HiPPO-LegS matrix $A$ is efficient to compute with, specifically in that multiplying $A$ or $\bar{A}$ with a vector $x$ admits an efficient algorithm with complexity $\mathcal{O}(d)$ instead of the general $\mathcal{O}(d^2)$. But this only means that Equation $\eqref{eq:s4-r}$ is more efficient than a generic RNN when computed recursively — for efficient training, pure recursion isn't enough, and we need to look into parallel computation methods.
There are two approaches to parallelizing a linear RNN. One is to treat it as a Prefix Sum problem, as introduced in Google's New Paper Tries to "Revive" RNNs: Can RNNs Shine Again?, and compute it directly with Associative Scan algorithms such as Upper/Lower, Odd/Even, or Ladner-Fischer; see the paper Prefix Sums and Their Applications. The other is to convert it into a convolution between a matrix sequence and a vector sequence, and accelerate it using the Fast Fourier Transform (FFT) — this is S4's approach. Either way, they face a common bottleneck: computing the power matrix $\bar{A}^k$.
Specifically, we generally set the initial state $x_0$ to zero, which lets us write:
\begin{equation}\begin{aligned} y_1 =&\, \bar{C}^*\bar{B}u_0\\ y_2 =&\, \bar{C}^*(\bar{A}x_0 + \bar{B}u_1) = \bar{C}^*\bar{A}\bar{B}u_0 + \bar{C}^*\bar{B}u_1\\ y_3 =&\, \bar{C}^*(\bar{A}x_1 + \bar{B}u_2) = \bar{C}^*\bar{A}^2 Bu_0 + \bar{C}^*\bar{A}Bu_1 + \bar{C}^*\bar{B}u_2\\[5pt] \vdots \\ y_L =&\, \bar{C}^*(\bar{A} x_{L-1}+\bar{B}u_{L-1}) = \sum_{k=0}^{L-1} \bar{C}^*\bar{A}^k \bar{B}u_{L-k} = \bar{K}_{< L} * u_{< L} \end{aligned}\end{equation}
where $*$ denotes the convolution operation, and
\begin{equation}\bar{K}_k = \bar{C}^*\bar{A}^k\bar{B},\quad \bar{K}_{< L} = \big(\bar{K}_0,\bar{K}_1,\cdots,\bar{K}_{L-1}\big),\quad u_{< L} = (u_0,u_1,\cdots,u_{L-1})\end{equation}
Note that, under our current conventions, both $\bar{C}^*\bar{A}^k \bar{B}$ and $u_k$ are scalars, so $\bar{K}_{< L},u_{< L}\in\mathbb{R}^L$ holds. As we know, convolution can be computed via the (discrete) Fourier transform by converting it into multiplication in the frequency domain and then transforming back; this has complexity $\mathcal{O}(L\log L)$, where $L$ is the sequence length. Although this complexity looks larger than the $\mathcal{O}(L)$ of direct recursion, the Fourier transform can be parallelized, so in practice it's actually faster.
So the question now is how to efficiently compute the convolution kernel $\bar{K}_{< L}$, which requires computing the power matrix $\bar{A}^k$ — and computing this directly by definition is still quite expensive. Of course, if we only needed to compute $\bar{A}^k$, that wouldn't be a problem, since $A$ is a constant matrix, and given $\epsilon$, $\bar{A}$ is also a constant matrix — no matter how hard its powers are to compute, we can just compute and cache them ahead of time. However, $\bar{A}^k$ is only an intermediate step; we still need to compute $\bar{C}^*\bar{A}^k\bar{B}$, and since S4 treats $\bar{C},\bar{B}$ as a trainable parameter, $\bar{C}^*\bar{A}^k\bar{B}$ can't be precomputed — precomputing $\bar{A}^k$ alone isn't efficient enough.
Generating Functions
Before going further, let's take a detour to introduce the concept of a generating function, which is one of the basic building blocks for the efficient computation that follows. For readers not so familiar with convolution and the discrete Fourier transform, this can also serve as a bit of background, giving a rough sense of the basic principle behind using the Fourier transform to accelerate convolution.
For a given sequence $a = (a_0,a_1,a_2,\cdots)$, its generating function treats each component as a coefficient of a power series:
\begin{equation}\mathcal{G}(z|a) = \sum_{k=0}^{\infty} a_k z^k\end{equation}
If we have two sequences $a = (a_0,a_1,a_2,\cdots)$ and $b = (b_0,b_1,b_2,\cdots)$, the product of their generating functions is:
\begin{equation}\mathcal{G}(z|a)\mathcal{G}(z|b) = \left(\sum_{k=0}^{\infty} a_k z^k\right)\left(\sum_{l=0}^{\infty} b_l z^l\right) = \sum_{k=0}^{\infty}\sum_{l=0}^{\infty}a_k b_l z^{k+l} = \sum_{l=0}^{\infty}\left(\sum_{k=0}^l a_k b_{l-k}\right) z^l \end{equation}
Notice something? The coefficient of the $l$-th term of $\mathcal{G}(z|a)\mathcal{G}(z|b)$ (i.e., the coefficient of $z^{l-1}$) is exactly the convolution of $a_{< l}=(a_0,\cdots,a_{l-1})$ and $b_{< l}=(b_0,\cdots,b_{l-1})$. If we have a way to quickly compute generating functions and quickly extract a particular coefficient from a generating function, we can convert convolution into: form the generating functions, multiply them as ordinary functions, then extract the relevant coefficient.
The Discrete Fourier Transform (DFT) is precisely this kind of approach to constructing generating functions. First, note that if we only need to convolve at most the first $L$ terms of $a,b$, then the upper limit of the summation in the generating function need not go to positive infinity — changing the upper limit to $L-1$ works just as well. Motivated by this, instead of computing the generating function for all $z$, the DFT evaluates it only at specific values of $z=e^{-2i\pi l/L},l=0,1,2,\dots,L-1$:
\begin{equation}\hat{a}_l = \sum_{k=0}^{L-1} a_k \left(e^{-2i\pi l/L}\right)^k = \sum_{k=0}^{L-1} a_k e^{-2i\pi kl/L}\end{equation}
The inverse transform for extracting coefficients (Inverse DFT, IDFT) is
\begin{equation}a_k = \frac{1}{L}\sum_{l=0}^{L-1} \hat{a}_l e^{2i\pi kl/L}\end{equation}
Both DFT and IDFT can be efficiently computed via the Fast Fourier Transform (FFT), and most numerical computing frameworks already have built-in functions for this, so efficiency isn't an issue. But note: if we want to use the DFT to compute a convolution, we need a small adjustment. Because $e^{-2i\pi l/L}$ is periodic, we can't distinguish $e^{-2i\pi l/L}$ from $e^{-2i\pi (l+L)/L}$, and when we multiply the DFTs of two sequences each summed over $L$ terms, the result will produce a term with exponent $l \geq L$ of $e^{-2i\pi kl/L}$ that gets mixed together with the $e^{-2i\pi k(l-L)/L}$ term — so applying IDFT would actually give us the sum of two coefficients, which is not the correct convolution result.
The fix is to change the $L$ in $e^{-2i\pi l/L}$ to $2L$ (while keeping the summation over $L$ terms), i.e. increasing the period so that the product stays within a single period. That is, redefine the DFT as
\begin{equation}\hat{a}_l = \sum_{k=0}^{L-1} a_k e^{-i\pi kl/L}\end{equation}
However, off-the-shelf FFT functions generally don't support adjusting the period independently — by default the period equals the array length — so the equivalent workaround is to pad $(a_0,a_1,\cdots,a_{L-1})$ with $L$ zeros, perform an ordinary DFT, take the product, apply IDFT, and finally keep only the first $L$ results.
From Powers to Inverses
For the convolution kernel $\bar{K}$, we have
\begin{equation}\mathcal{G}(z|\bar{K}) = \sum_{k=0}^{\infty} \bar{C}^*\bar{A}^k \bar{B}z^k = \bar{C}^*\left(I - z\bar{A}\right)^{-1}\bar{B}\label{eq:k-gen}\end{equation}
We can see that the generating function not only accelerates the convolution computation, it also converts the original power-matrix computation $\bar{A}^k$ into a matrix-inverse computation $\left(I - z\bar{A}\right)^{-1}$.
What kind of matrix $\bar{A}$ has an easy-to-compute $\left(I - z\bar{A}\right)^{-1}$? First, a diagonal matrix is certainly fine: if $\bar{A}$ is diagonal, then $I - z\bar{A}$ is also diagonal, and the inverse of a diagonal matrix is obtained simply by inverting each diagonal entry. Second, if $\bar{A}$ can be diagonalized as $\bar{\Lambda}$, i.e. $\bar{A}=P^{-1}\bar{\Lambda} P$, then $\left(I - z\bar{A}\right)^{-1}$ is equally easy to compute, since
\begin{equation}\left(I - z\bar{A}\right)^{-1} = \left(P^{-1}(I - z\bar{\Lambda})P\right)^{-1} = P^{-1}\left(I - z\bar{\Lambda}\right)^{-1} P\end{equation}
So can $\bar{A}$ be diagonalized? This depends on whether $A$ can be diagonalized. If $A=P^{-1}\Lambda P$, then by similarity invariance we can shift the entire computation to the new system $A=\Lambda$, and by definition the new $\bar{A}$ becomes:
\begin{equation}\begin{aligned} \bar{A}=&\,(I - \epsilon A/2)^{-1}(I + \epsilon A/2) \\ =&\,(I - \epsilon\Lambda/2)^{-1}(I + \epsilon\Lambda/2) \end{aligned}\end{equation}
which is clearly a diagonal matrix.
So can $A$ be diagonalized? In theory, yes; in practice, no. It's theoretically possible because, in theory, almost every matrix can be diagonalized over the complex numbers, and in the previous post we already gave the eigenvalues of LegS's $A$, $[-1,-2,-3,\cdots]$ — so we even know what the diagonalized matrix looks like. But it's practically infeasible because it's numerically very hard: numerical computation has to take precision, memory, and time into account, and if any one of these exceeds acceptable limits, a theoretically viable algorithm becomes practically unusable.
For the $A$ matrix, the practical difficulty mainly comes from the fact that the matrix $P$ needed to diagonalize $A$ suffers from numerical instability — ultimately a consequence of the limited precision of floating-point computation. On this point, the original paper simply gives the analytic solution for the matrix $P$ without explanation and then verifies it, which is clearly not very helpful for readers trying to understand it. Below, I'll offer another way to think about this, starting from the computation of eigenvectors.
Eigenvectors
Diagonalizing $A$ is equivalent to diagonalizing $-A$, and since the eigenvalues of $A$ are all negative, for simplicity let's instead consider diagonalizing $-A$, which has $d$ distinct eigenvalues $\lambda=1,2,\cdots,d$. The matrix needed to diagonalize it is exactly the stack of its eigenvectors, so finding $P$ essentially amounts to finding eigenvectors. And for a matrix with known eigenvalues, the direct method for solving eigenvectors is to solve the equation $-Av=\lambda v$.
In the "Efficient Computation" section of the previous post, we already gave the result for the $n$-th component of $Av$:
\begin{equation}(Av)_n = n v_n -\sqrt{2n+1}\sum_{k=0}^n \sqrt{2k+1}v_k \end{equation}
So $-Av=\lambda v$ implies
\begin{equation}\sqrt{2n+1}\sum_{k=0}^n \sqrt{2k+1}v_k - n v_n = \lambda v_n\end{equation}
Writing $S_n = \sum\limits_{k=0}^n \sqrt{2k+1}v_k$, we get $\sqrt{2n+1}v_n=S_n - S_{n-1}$, which after a little rearranging gives
\begin{equation}S_{n-1} = \frac{\lambda - n - 1}{\lambda + n}S_n\end{equation}
Note that $-Av=\lambda v$ is an underdetermined equation, so we have some freedom to adjust things flexibly (i.e., the eigenvector is not unique). Since the largest value of $n$ is $d-1$, we can set $S_{d-1}=1$ and then recursively work backwards, until $\lambda - n - 1=0$ gives us $S_{\lambda - 1} = 0$, after which $S_n = 0$ holds for all $\forall n < \lambda - 1$. For $n > \lambda - 1$, we have
\begin{equation}S_n = (-1)^{d-n-1}\frac{(d-\lambda)! (n+\lambda)!}{(d+\lambda-1)! (n-\lambda + 1)!}\end{equation}
Since our goal here is to demonstrate the numerical instability of $P$, it suffices to examine a single eigenvector. Take $n=\lambda=d/3$ (if $d$ is not a multiple of 3, simply round to the nearest integer — the conclusion is unaffected), giving
\begin{equation}|S_{d/3}| = \frac{\left(\frac{2d}{3}\right)! \left(\frac{2d}{3}\right)!}{\left(\frac{4d}{3}-1\right)!}\sim \mathcal{O}(\sqrt{d}\,2^{-4d/3})\end{equation}
The final $\sim$ can be obtained from Stirling's formula. This result shows that for the eigenvalue $d/3$, going from $S_{d-1}$ to $S_{d/3}$ there's an exponential decay (or, going the other way, an exponential blow-up). This means the eigenvector's components also decay similarly from $v_{d-1}$ to $v_{d/3}$, and within the finite precision of floating-point numbers, it's very hard to handle such an eigenvector accurately. So directly diagonalizing the matrix $P$ of $A$ suffers from numerical instability.
Diagonal Plus Low-Rank
Besides diagonal matrices, when $\bar{A}$ admits a low-rank decomposition, this can likewise reduce the difficulty of computing $\left(I - z\bar{A}\right)^{-1}$. This is because of the following Woodbury identity:
\begin{equation}(I - UV^*)^{-1} = \sum_{k=0}^{\infty} (UV)^k = I + U\left(\sum_{k=0}^{\infty}(V^* U)^k\right)V = I + U(I - V^* U)^{-1} V^*\end{equation}
Here $U,V\in\mathbb{R}^{d\times r}$, and the derivation makes use of $(UV^*)^k = U(V^* U)^{k-1}V$. If $d \gg r$, then in theory the cost of computing $(I - V^* U)^{-1}$ is much lower than that of $(I - UV^*)^{-1}$, so this speeds things up. In particular, if $r=1$, then $(I - V^* U)^{-1}$ is just the reciprocal of a scalar, which is the simplest case of all.
However, we know that $A$ is a lower-triangular matrix with no zero diagonal entries, so it must be full rank. Combined with the conclusion of the previous section, this means $A$ is neither low-rank nor practically diagonalizable — so neither approach applies. Is there another way? Yes! Using the Woodbury identity above, we can derive a more general version of it:
\begin{equation}\begin{aligned} (M - UV^*)^{-1} =&\, (M(I - (M^{-1}U)V^*))^{-1} = (I - (M^{-1}U)V^*)^{-1}M^{-1} \\ =&\, (I + M^{-1}U(I - V^*M^{-1}U)^{-1} V^*)M^{-1} \\ =&\, M^{-1} + M^{-1}U(I - V^*M^{-1}U)^{-1} V^*M^{-1} \\ \end{aligned}\end{equation}
This result tells us that if the inverse of $M$ is easy to compute, then adding or subtracting a low-rank matrix from it keeps the inverse easy to compute. And which matrices have easy-to-compute inverses? We're back to the answer from the previous section — diagonal matrices. So we should try to force either $A$ or $\bar{A}$ into a "diagonal + low-rank" form.
In fact, if we look closely, the matrix $A$ already has hints of a "diagonal + low-rank" structure. In the previous post, we rewrote the definition of $A$ equivalently as:
\begin{equation}A_{n,k} = \left\{\begin{array}{l}n\delta_{n,k} - \sqrt{2n+1}\sqrt{2k+1}, &k \leq n \\ 0, &k > n\end{array}\right.\end{equation}
where $n\delta_{n,k}$ is essentially the diagonal matrix $\text{diag}(0,1,2,\cdots)$, and $\sqrt{2n+1}\sqrt{2k+1}$ can be rewritten in low-rank matrix form as $v v^*$, where $v=[1,\sqrt{3},\sqrt{5},\cdots]^*\in\mathbb{R}^{d\times 1}$. In other words, if it weren't for the restriction $k > n, A_{n,k}=0$, then $A$ itself would already be a diagonal matrix minus a low-rank matrix.
The Masterstroke
While the lower-triangular restriction means this pattern no longer directly applies, we can still make full use of the $v v^*$ structure that's already there to help construct a new diagonalizable matrix. This trick is, frankly, quite clever — a real masterstroke, and it's hard not to admire the original authors for it. Specifically, consider $A+\frac{1}{2}v v^*$:
\begin{equation}\left(A + \frac{1}{2}v v^*\right)_{n,k} = \left\{\begin{array}{l}n\delta_{n,k} - \frac{1}{2}\sqrt{2n+1}\sqrt{2k+1}, &k \leq n \\ \frac{1}{2}\sqrt{2n+1}\sqrt{2k+1}, &k > n\end{array}\right.\end{equation}
The diagonal entries of this new matrix are exactly $-\frac{1}{2}I$. Adding $\frac{1}{2}I$ to it gives
\begin{equation}\left(A + \frac{1}{2}v v^*+\frac{1}{2}I\right)_{n,k} = \left\{\begin{array}{} - \frac{1}{2}\sqrt{2n+1}\sqrt{2k+1}, &k < n \\ 0, &k=n \\ \frac{1}{2}\sqrt{2n+1}\sqrt{2k+1}, &k > n\end{array}\right.\end{equation}
And here's the key point: this is a skew-symmetric matrix, so it's guaranteed to be diagonalizable (over the complex numbers)! We've thus decomposed $A$ into the sum of a diagonalizable matrix and a low-rank matrix! One might object: wasn't $A$ already guaranteed to be diagonalizable in theory, yet still had numerical stability issues — so shouldn't we worry about the numerical stability of diagonalizing this skew-symmetric matrix too? Here's the crux of it: a skew-symmetric matrix isn't just guaranteed to be diagonalizable — it's guaranteed to be diagonalizable by an orthogonal matrix (a unitary matrix, in the complex case)! Unitary matrices generally have excellent numerical stability, so there's nothing to worry about here — which is exactly why we don't diagonalize $A$ directly but instead take this detour through the skew-symmetric matrix.
We now have that there exists a diagonal matrix $\Lambda$ and a unitary matrix $U$ such that $A + \frac{1}{2}v v^*+\frac{1}{2}I = U^*\Lambda U$, from which
\begin{equation}A = U^*\Lambda U - \frac{1}{2}I - \frac{1}{2}v v^* = U^*\left(\Lambda - \frac{1}{2}I - \frac{1}{2}(Uv)(Uv)^*\right) U\end{equation}
Stripping away the scaffolding, we find that the final conclusion can be summarized as: "$A$ is isomorphic to a diagonal matrix minus a rank-1 matrix." That is, there exists a unitary matrix $U$, diagonal matrix $\Lambda$, and column vector $u,v$, such that:
\begin{equation}A = U^*\left(\Lambda - uv^*\right) U\end{equation}
Note that a "diagonal + low-rank" matrix can be multiplied by a vector efficiently, e.g.:
\begin{equation}\left(\Lambda - uv^*\right)x = \Lambda x - u(v^*x)\end{equation}
$\Lambda x$ amounts to treating $\Lambda$ as a vector and multiplying it elementwise with $x$, while $u(v^*x)$ is the inner product of $v$ with $x$ followed by multiplying the resulting scalar with the vector $u$ — both of which can be done in $\mathcal{O}(d)$.
The Final Push
With $A=U^*\left(\Lambda - uv^*\right) U$ in hand, and once again invoking similarity invariance, all our remaining computation can be shifted into $A=\Lambda - uv^*$, so from here on we take $A=\Lambda - uv^*$. First, for $\bar{A}$:
\begin{equation}\bar{A}=\big(I - \epsilon (\Lambda - uv^*)/2\big)^{-1}\big(I + \epsilon (\Lambda - uv^*)/2\big)\end{equation}
Noticing that $I - \epsilon (\Lambda - uv^*)/2= \frac{\epsilon}{2}(D + uv^*)$, where $D=\frac{2}{\epsilon}I - \Lambda$ is diagonal, we can apply the Woodbury identity to obtain:
\begin{equation}\big(I - \epsilon (\Lambda - uv^*)/2\big)^{-1} =\frac{2}{\epsilon}(D + uv^*)^{-1} = \frac{2}{\epsilon}\left[D^{-1} - D^{-1}u(I + v^*D^{-1}u)^{-1} v^*D^{-1}\right]\end{equation}
Look closely and you'll see this is again "diagonal + low-rank"; multiplying by $\big(I + \epsilon (\Lambda - uv^*)/2\big)$ then completes the computation of $\bar{A}$, and the final result is the product of two "diagonal + low-rank" matrices, which means it, too, is efficient to compute — a result that can be used during recursive inference.
Finally, we need the convolution kernel required for parallel training, which we've already converted into the generating function $\eqref{eq:k-gen}$, and we can now finish computing it. First, by an operation similar to "putting over a common denominator," we can show that:
\begin{equation}\begin{aligned} \mathcal{G}(z|\bar{K}) = \bar{C}^* \left(I - \bar{A}z\right)^{-1}\bar{B} =&\, \bar{C}^* \left(I - (I - \epsilon A/2)^{-1}(I + \epsilon A/2)z\right)^{-1}\bar{B} \\ =&\, \bar{C}^* \left[(I - \epsilon A/2)^{-1}\big((I - \epsilon A/2)-(I + \epsilon A/2)z\big)\right]^{-1}\bar{B} \\ =&\, \bar{C}^* \big[(I - \epsilon A/2)-(I + \epsilon A/2)z\big]^{-1}(I - \epsilon A/2)\bar{B} \\ =&\, \bar{C}^* \big[(I - \epsilon A/2)-(I + \epsilon A/2)z\big]^{-1}B\epsilon \\ =&\, \bar{C}^* \big[(1-z)I - (1+z)\epsilon A / 2\big]^{-1}B\epsilon \\ =&\, \frac{2}{1+z}\bar{C}^* \left[\frac{2}{\epsilon}\frac{1-z}{1+z}I - A\right]^{-1}B \\ \end{aligned}\end{equation}
Substituting $A=\Lambda - uv^*$ gives
\begin{equation}\mathcal{G}(z|\bar{K}) = \frac{2}{1+z}\bar{C}^* \left[\frac{2}{\epsilon}\frac{1-z}{1+z}I - (\Lambda - uv^*)\right]^{-1}B=\frac{2}{1+z}\bar{C}^* (R_z + uv^*)^{-1}B\end{equation}
Here $R_z = \frac{2}{\epsilon}\frac{1-z}{1+z}I - \Lambda$ is diagonal, so we can once again apply the Woodbury identity to complete the computation:
\begin{equation}\mathcal{G}(z|\bar{K}) = \frac{2}{1+z}\bar{C}^* \left[R_z^{-1} - R_z^{-1}u(I + v^*R_z^{-1}u)^{-1} v^*R_z^{-1}\right]B\end{equation}
This is a scalar function of $z$. But there's a subtlety worth noting: what the Fourier transform actually needs is the "truncated generating function":
\begin{equation}\mathcal{G}_L(z|\bar{K}) = \sum_{k=0}^{L-1} \bar{C}^*\bar{A}^k \bar{B}z^k = \bar{C}^*(I - z^L\bar{A}^L)\left(I - z\bar{A}\right)^{-1}\bar{B}\end{equation}
which amounts to replacing $\bar{C}^*$ in $\mathcal{G}(z|\bar{K})$ with $\bar{C}^*(I - z^L\bar{A}^L)$, where $L$ is the maximum training length chosen ahead of time. From here, we just need to substitute $z=e^{-2i\pi l/L},l=0,1,2,\dots,L-1$ into the computation, obtaining the DFT of $\bar{K}$, then apply IDFT to get $\bar{K}$. This process could be further accelerated by converting it to a Cauchy kernel problem, but I don't think that's particularly essential, so I won't go into it here. Lastly, there's one more trick: for $z=e^{-2i\pi l/L}$ we have $z^L=1$, which just amounts to replacing $\bar{C}^*$ with $\bar{C}^*(I - \bar{A}^L)$. Since S4 treats $\bar{C}$ as a trainable parameter, we can directly treat $\bar{C}^*(I - \bar{A}^L)$ as the trainable parameter and later solve for $\bar{C}$ from it for use in inference, which lets us avoid ever computing $\bar{A}^L$ during training.
At this point, one might think we could equally well substitute $z=e^{-i\pi l/L}$ directly to compute the DFT of $\bar{K}$ used for the convolution, rather than the roundabout process of first taking the IDFT to get $\bar{K}$, padding with zeros, and then taking the DFT again. But the problem is that $z^L=(-1)^l$ is not a fixed value here, so we can't treat $\bar{C}^*(I - z^L\bar{A}^L)$ as a single trainable parameter — this would require computing $\bar{A}^L$ during training, which is fairly expensive (though, of course, if $\bar{A}$ is completely fixed during training, it could be precomputed — it depends on the situation).
Wrapping Up
After this rather lengthy discussion, we've now worked through the key mathematical details of S4, and I hope this has been useful to readers interested in understanding it. As we've seen, S4 further supplements and completes HiPPO, and its crucial contribution is showing that $A$ is equivalent to a "diagonal + low-rank" matrix form, laying the groundwork for the rest of the analysis. This matters because $A$ was originally defined piecewise rather than as a matrix expression, and such a definition is not amenable to applying standard linear-algebra tools for general analysis.
Since the HiPPO derivation was carried out assuming that $u(t)$ is a one-dimensional function, up to this point S4's $u_k$ has also remained a scalar. So how does S4 handle vector-sequence inputs? Rather brute-force: it applies the aforementioned linear RNN independently to each component, with each RNN using a different set of $\epsilon,B,C$ parameters, and then concatenates the results — an approach that persists even in the author's more recent work, Mamba. There is, of course, a simplified alternative: handling vector inputs directly within a single RNN, just by changing $B,C$ to a matrix accordingly — this is S5 (whose author is no longer Albert Gu). This approach can be understood as simply borrowing S4's linear RNN form and HiPPO's matrix $A$, while discarding all the other fine details of HiPPO, and it too achieves solid results.
Ironically, S4 proposed a wealth of ingenious mathematical techniques to simplify and accelerate the computation of $A$, yet starting with Diagonal State Spaces are as Effective as Structured State Spaces, the original author's subsequent work — including Mamba — largely abandoned all of this and simply assumed $A$ to be a diagonal matrix from the outset, at which point the RNN component becomes essentially the same as the LRU introduced in Google's New Paper Tries to "Revive" RNNs: Can RNNs Shine Again?. So, from the perspective of the current state of the art in SSMs and linear RNNs, the HiPPO/S4 line of work is, in a sense, already "outdated." Many articles explaining Mamba that start from HiPPO and S4 are, with hindsight, perhaps not strictly necessary.
That said, for me personally, spending so much time studying HiPPO and S4 wasn't simply about understanding or using the latest SSM and RNN models — it was about learning the assumptions and derivations behind HiPPO, understanding how linear systems handle memory and where their bottlenecks lie, and building up a richer set of ideas for constructing new models and methods in the future. Beyond that, the many elegant mathematical techniques found in HiPPO and S4 are genuinely a pleasure to study, and make for excellent practice for sharpening one's mathematical skills.
Summary
This post introduced S4, the follow-up work to HiPPO. Its key contribution is the "diagonal matrix + low-rank matrix" decomposition, which enables efficient, parallel computation of the HiPPO matrix. We've focused mainly on walking through and deriving the more difficult mathematical details involved.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.