Revisiting SSM (I): Linear Systems and the HiPPO Matrix
A few days ago I read several articles introducing SSMs (State Space Models), and realized that I had actually never properly studied SSMs myself. So I decided to sit down and learn about them properly, and started this new series to record what I learn along the way.
The concept of SSM has a long history, but here we specifically mean the SSM used in deep learning, whose opening work is generally considered to be S4 from 2021 — not too old — while the newest and hottest SSM variant is probably last year's Mamba. Of course, when we talk about SSM we may also be referring more broadly to all linear RNN models, in which case RWKV, RetNet, and the LRU we introduced earlier in Google's New Work Tries to "Revive" RNNs: Can RNNs Shine Again? can all be grouped into this category. Many SSM variants aim to become competitors to the Transformer, and although I don't think a complete replacement is really possible, the elegant mathematical properties of SSM itself are well worth studying.
Although we say that SSM originates from S4, before S4 there was an extremely powerful foundational work on SSM: HiPPO: Recurrent Memory with Optimal Polynomial Projections (abbreviated HiPPO), so this post starts from HiPPO. more
Basic Form
A quick digression first: the first author of HiPPO, S4, and Mamba — the representative SSM works mentioned above — is all the same person, Albert Gu. He has many other SSM-related works, and it's no exaggeration to say that these works form the foundation of the entire SSM edifice. Regardless of how the future of SSM turns out, this kind of relentless dedication to a single research topic deserves our sincere admiration.
Back to the main topic. For readers who already have some prior knowledge of SSM, you probably know that SSM modeling uses a 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}\label{eq:ode}\end{equation}
where $u(t)\in\mathbb{R}^{d_i}, x(t)\in\mathbb{R}^{d}, y(t)\in\mathbb{R}^{d_o}, A\in\mathbb{R}^{d\times d}, B\in\mathbb{R}^{d\times d_i}, C\in\mathbb{R}^{d_o\times d}, D\in\mathbb{R}^{d_o\times d_i}
$. Of course, we can also discretize it, turning it into a linear RNN model — we'll expand on that in a later post. Whether discretized or not, the key word is "linear," which immediately raises a natural question: why a linear system? Is a linear system enough?
We can answer this question from two angles: a linear system is both sufficiently simple and sufficiently complex. Simple means that, in theory, linearization is often the most basic approximation of a complex system, so a linear system is usually an unavoidable starting point. Complex means that even such a simple system can fit remarkably complicated functions. To see this, we only need to consider a simple example with $\mathbb{R}^4$:
\begin{equation} x'(t) =\begin{pmatrix} 1 & 0 & 0 & 0 \\ 0 & -1 & 0 & 0 \\ 0 & 0 & 0 & 1 \\ 0 & 0 & -1 & 0 \end{pmatrix}x(t)\end{equation}
The general solution of this example is $x(t) = (e^t, e^{-t}, \sin t, \cos t)$. What does this mean? It means that as long as $d$ is large enough, this linear system can fit sufficiently complex functions through combinations of exponential and trigonometric functions, and we know that even the highly expressive Fourier series is nothing more than a combination of trigonometric functions — adding exponential functions on top clearly makes it even more expressive. So it's reasonable to expect that linear systems also have sufficiently strong fitting power.
Of course, these explanations are, in a sense, all after the fact. The result given by HiPPO is more fundamental: when we try to use an orthogonal basis to approximate a dynamically updating function, the result is exactly the linear system above. This means HiPPO not only tells us that a linear system can approximate sufficiently complex functions, but also tells us exactly how to approximate it, and even how good that approximation is.
Finite Compression
Going forward, we will only consider the special case where $d_i=1$, with $d_i > 1$ simply being the parallel generalization for $d_i=1$. In this case, the output of $u(t)$ is a scalar. Further, to start off simply, let's assume $t\in[0, 1]$. HiPPO's goal is: to store the information of the segment $u(t)$ using a finite-dimensional vector.
This looks like an impossible requirement at first glance, since $t\in[0,1]$ means that $u(t)$ could correspond to a vector made up of infinitely many points, and compressing that into a finite-dimensional vector could cause severe distortion. However, if we make some assumptions about $u(t)$ and allow some loss, this compression becomes achievable — and in fact most readers have already tried something like this. For example, when $u(t)$ is differentiable up to order $n+1$ at some point, its corresponding order-$n$ Taylor expansion is often a good approximation of $u(t)$, so we can just store the $n+1$ coefficients of the expansion as an approximate representation of $u(t)$, which successfully compresses $u(t)$ into an $n+1$-dimensional vector.
Of course, for real-world data, a condition like "differentiable up to order $n+1$" is extremely demanding, so we usually prefer expansion in an orthogonal function basis under square-integrability, such as the Fourier series, whose coefficients are computed as
\begin{equation}c_n = \int_0^1 u(t) e^{-2i\pi n t}dt \label{eq:fourier-coef-1}\end{equation}
If we take a sufficiently large integer $N$ and keep only the coefficients of $|n|\leq N$, then we compress $u(t)$ into an $2N + 1$-dimensional vector.
Now let's raise the difficulty. We just said $t\in[0,1]$, which is a static interval, but in practice $u(t)$ represents a continuously acquired signal, so new data keeps coming in. For instance, suppose we've now approximated the data over the interval $[0,1]$, and immediately the data of $[1,2]$ arrives — we need to update the approximation to try to memorize the whole interval $[0,2]$, followed by $[0,3]$, $[0,4]$, and so on. We call this "online function approximation." The Fourier coefficient formula above, $\eqref{eq:fourier-coef-1}$, applies only to the interval $[0,1]$, so we need to generalize it.
To do this, let $t\in[0,T]$, where $s\mapsto t_{\leq T}(s)$ is a mapping from $[0,1]$ to $[0,T]$. Then, when we treat $u(t_{\leq T}(s))$ as a function of $s$, its domain becomes $[0,1]$, and we can reuse equation $\eqref{eq:fourier-coef-1}$:
\begin{equation}c_n(T) = \int_0^1 u(t_{\leq T}(s)) e^{-2i\pi n s}ds \label{eq:fourier-coef-2}\end{equation}
Here we've already marked the coefficient as $(T)$, to indicate that this coefficient will change as $T$ changes.
The First Signs of Linearity
There are infinitely many functions that can map $[0,1]$ to $[0,T]$, and the final result differs depending on the choice of $t_{\leq T}(s)$. Some fairly intuitive and relatively simple choices are as follows:
1. $t_{\leq T}(s) = sT$, i.e., mapping $[0,1]$ uniformly onto $[0,T]$;
2. Note that $t_{\leq T}(s)$ doesn't have to be surjective, so something like $t_{\leq T}(s)=s + T - 1$ is also allowed, meaning we only keep the information of the nearest window $[T-1,T]$ and discard the earlier part. More generally we could have $t_{\leq T}(s)=sw + T - w$, where $w$ is a constant, meaning that information before $T-w$ is discarded;
3. We can also choose a non-uniform mapping, such as $t_{\leq T}(s) = T\sqrt{s}$, which is likewise a surjection from $[0,1]$ to $[0,T]$, but when $s=1/4$, it maps to $T/2$ — meaning that although we pay attention to the full history, we place more weight on the information near time $T$.
Now, taking $t_{\leq T}(s)=sw + T - w$ as an example, substituting into equation $\eqref{eq:fourier-coef-2}$ gives
\begin{equation}c_n(T) = \int_0^1 u(sw + T - w) e^{-2i\pi n s}ds\end{equation}
Now let's take the derivative of both sides with respect to $T$:
\begin{equation}\begin{aligned} \frac{d}{dT}c_n(T) =&\, \int_0^1 u'(sw + T - w) e^{-2i\pi n s}ds \\ =&\, \left.\frac{1}{w} u(sw + T - w) e^{-2i\pi n s}\right|_{s=0}^{s=1} + \frac{2i\pi n}{w}\int_0^1 u(sw + T - w) e^{-2i\pi n s}ds \\ =&\, \frac{1}{w} u(T) - \frac{1}{w} u(T-w) + \frac{2i\pi n}{w} c_n(T) \\ \end{aligned}\label{eq:fourier-dc}\end{equation}
where the second equality uses integration by parts. Since we only keep the coefficients of $|n|\leq N$, according to the Fourier series formula, the following can be considered a good approximation of $u(sw + T - w)$:
\begin{equation}u(sw + T - w) \approx \sum_{k=-N}^{k=N} c_k(T) e^{2i\pi k s}\end{equation}
so $u(T - w) = u(sw + T - w)|_{s=0}\approx \sum\limits_{k=-N}^{k=N} c_k(T)$, and substituting into equation $\eqref{eq:fourier-dc}$ gives:
\begin{equation}\frac{d}{dT}c_n(T) \approx \frac{1}{w} u(T) - \frac{1}{w} \sum_{k=-N}^{k=N} c_k(T) + \frac{2i\pi n}{w} c_n(T)\end{equation}
Renaming $T$ as $t$, then stacking all the $c_n(t)$ together and denoting them as $x(t) = (c_{-N},c_{-(N-1)},\cdots,c_0,\cdots,c_{N-1},c_N)$, and without distinguishing between $\approx$ and $=$, we arrive at
\begin{equation}x'(t) = Ax(t) + Bu(t),\quad A_{n,k} = \left\{\begin{array}{l}(2i\pi n - 1)/w, &k=n \\ -1/w,&k\neq n\end{array}\right.,\quad B_n = 1/w\end{equation}
This is exactly the linear ODE system shown in equation $\eqref{eq:ode}$. That is, when we try to use a Fourier series to memorize the state of a real-time function within its nearest window, the result naturally leads to a linear ODE system.
The General Framework
Of course, so far we've only chosen a particular $t_{\leq T}(s)$; a different choice of $t_{\leq T}(s)$ won't necessarily give such a simple result. Moreover, the Fourier series result is stated over the complex domain — it's possible to go further and make it purely real, but the form becomes more complicated. So we need to generalize the process from the previous section into a general framework, in order to get a more general and simpler, purely real-valued conclusion.
Let $t\in[a,b]$, and suppose we have a target function $u(t)$ and a function basis $\{g_n(t)\}_{n=0}^N$. We want to approximate the former with a linear combination of the latter, with the goal of minimizing the $L_2$ distance:
\begin{equation}\mathop{\text{argmin}}_{c_1,\cdots,c_N}\int_a^b \left[u(t) - \sum_{n=0}^N c_n g_n(t)\right]^2 dt\end{equation}
Here we're mainly working within the real numbers, so we can just square the bracketed expression directly without needing to take the modulus. A more general objective function could also include a weighting function $\rho(t)$, but we won't consider that here, since HiPPO's main conclusions actually don't consider this weighting function either.
Expanding the objective function, we get
\begin{equation}\int_a^b u^2(t) dt - 2\sum_{n=0}^N c_n \int_a^b u(t) g_n(t)dt + \sum_{m=0}^N\sum_{n=0}^N c_m c_n \int_a^b g_m(t) g_n(t) dt\end{equation}
Here we only consider an orthonormal function basis, defined by $\int_a^b g_m(t) g_n(t) dt = \delta_{m,n}$, where $\delta_{m,n}$ is the Kronecker delta. In this case, the expression above simplifies to
\begin{equation}\int_a^b u^2(t) dt - 2\sum_{n=0}^N c_n \int_a^b u(t) g_n(t)dt + \sum_{n=0}^N c_n^2 \end{equation}
This is just a quadratic function of $c_n$, and its minimum has an analytical solution:
\begin{equation}c^*_n = \int_a^b u(t) g_n(t)dt\end{equation}
This is also called the inner product of $u(t)$ and $g_n(t)$ — it's the natural generalization of the inner product from finite-dimensional vector spaces to function spaces. For simplicity, when there is no risk of confusion, we'll take $c_n$ to mean $c^*_n$ by default.
The subsequent treatment is the same as in the previous section: we want to consider the approximation of $u(t)$ for a general $t\in[0, T]$, so we find a mapping $s\mapsto t_{\leq T}(s)$ from $[a,b]$ to $[0,T]$, and then compute the coefficient
\begin{equation}c_n(T) = \int_a^b u(t_{\leq T}(s)) g_n(s) ds\end{equation}
Again, we take the derivative of both sides with respect to $T$, and use integration by parts:
\begin{equation}\scriptsize\begin{aligned} \frac{d}{dT}c_n(T) =&\, \int_a^b u'(t_{\leq T}(s)) \frac{\partial t_{\leq T}(s)}{\partial T} g_n(s) ds = \int_a^b \left(\frac{\partial t_{\leq T}(s)}{\partial T}\left/\frac{\partial t_{\leq T}(s)}{\partial s}\right.\right) g_n(s) d u(t_{\leq T}(s)) \\ =&\,\left.u(t_{\leq T}(s))\left(\frac{\partial t_{\leq T}(s)}{\partial T}\left/\frac{\partial t_{\leq T}(s)}{\partial s}\right.\right) g_n(s)\right|_{s=a}^{s=b} - \int_a^b u(t_{\leq T}(s)) \,d\left[\left(\frac{\partial t_{\leq T}(s)}{\partial T}\left/\frac{\partial t_{\leq T}(s)}{\partial s}\right.\right) g_n(s)\right] \end{aligned}\label{eq:hippo-base}\end{equation}
Enter Legendre
The subsequent calculation now depends on the specific forms of $g_n(t)$ and $t_{\leq T}(s)$. HiPPO's full name is High-order Polynomial Projection Operators, and the first "P" is indeed the initial of "Polynomial" — so the key to HiPPO is choosing polynomials as the basis. Now, following Fourier, let's bring in another great mathematician — Legendre. The function basis we're about to choose is precisely the "Legendre polynomials" named after him.
The Legendre polynomial $p_n(t)$ is a degree-$n$ function of $t$, with domain $[-1,1]$, satisfying
\begin{equation}\int_{-1}^1 p_m(t) p_n(t) dt = \frac{2}{2n+1}\delta_{m,n}\end{equation}
So the $p_n(t)$ are only mutually orthogonal, not yet normalized (i.e., the integral of the square isn't 1) — $g_n(t)=\sqrt{\frac{2n+1}{2}} p_n(t)$ is the orthonormal basis.
When we apply Gram–Schmidt orthogonalization to the function basis $\{1,t,t^2,\cdots, t^n\}$, the result is exactly the Legendre polynomials. Compared to the Fourier basis, the advantage of Legendre polynomials is that they're defined purely within the real domain, and their polynomial form helps simplify part of the derivation of $t_{\leq T}(s)$, as we'll see later. Legendre polynomials have many different definitions and properties, which we won't go through exhaustively here — interested readers can check the linked Wikipedia article.
Next, we'll use two recursive formulas to derive an identity. These two recursive formulas are
\begin{align} p_{n+1}'(t) - p_{n-1}'(t) = (2n+1)p_n(t) \label{eq:leg-r1}\\[5pt] p_{n+1}'(t) = (n + 1)p_n(t) + t p_n'(t) \label{eq:leg-r2}\\ \end{align}
Iterating the first formula, $\eqref{eq:leg-r1}$, gives:
\begin{equation}\begin{aligned} p_{n+1}'(t) =&\, (2n+1)p_n(t) + (2n-3)p_{n-2}(t) + (2n-7)p_{n-4}(t) + \cdots \\ =&\, \sum_{k=0}^n (2k+1) \chi_{n-k} p_k(t) \end{aligned}\label{eq:leg-dot}\end{equation}
where, when $k$ is even, $\chi_k=1$, and otherwise $\chi_k=0$. Substituting into the second formula $\eqref{eq:leg-r2}$ gives
\begin{equation}t p_n'(t) = n p_n(t) + (2n-3)p_{n-2}(t) + (2n-7)p_{n-4}(t) + \cdots\end{equation}
and thus
\begin{equation}\begin{aligned} (t+1) p_n'(t) =&\, n p_n(t) + (2n-1)p_{n-1}(t) + (2n-3)p_{n-2}(t) + \cdots\\ =&\,-(n+1) p_n(t) + \sum_{k=0}^n (2k + 1) p_k(t) \end{aligned}\label{eq:leg-dot-t1}\end{equation}
These are the identities we'll need shortly. In addition, the Legendre polynomials satisfy $p_n(1)=1,p_n(-1)=(-1)^n$, a boundary value we'll also use later.
Just as an $n$-dimensional space has more than one orthogonal basis, there is also more than one family of orthogonal polynomials besides Legendre — for instance Chebyshev polynomials, and if we allow a weighted target function (i.e., $\rho(t)\not\equiv 1$), there are also Laguerre polynomials, among others. All of these are mentioned in the original paper, but HiPPO's main conclusions are still built around the Legendre polynomial expansion, so we won't go further into the rest here.
The Nearby Window
With the groundwork done, we can now substitute the specific form of $t_{\leq T}(s)$ to carry out the calculation. The process is much like the Fourier series example, except the basis function is now the orthonormal basis constructed from Legendre polynomials, $g_n(t)=\sqrt{\frac{2n+1}{2}} p_n(t)$. As our first example, let's again consider keeping only the information in the nearest window, where $t_{\leq T}(s) = (s + 1)w / 2 + T - w$ maps $[-1,1]$ onto $[T-w,T]$. The original paper calls this case "LegT (Translated Legendre)."
Substituting directly into equation $\eqref{eq:hippo-base}$, we immediately get
\begin{equation}\small\frac{d}{dT}c_n(T) = \frac{\sqrt{2(2n+1)}}{w}\left[u(T) - (-1)^n u(T-w)\right] - \frac{2}{w}\int_{-1}^1 u((s + 1)w / 2 + T - w) g_n'(s) ds\end{equation}
Let's first handle the $u(T-w)$ term. Using the same idea as in the Fourier series case, we truncate $n\leq N$ as an approximation of $u((s + 1)w / 2 + T - w)$:
\begin{equation}u((s + 1)w / 2 + T - w)\approx \sum_{k=0}^N c_k(T)g_k(s)\end{equation}
which gives $u(T-w)\approx \sum\limits_{k=0}^N c_k(T)g_k(-1) = \sum\limits_{k=0}^N (-1)^k c_k(T) \sqrt{\frac{2k+1}{2}}$. Next, using equation $\eqref{eq:leg-dot}$, we get
\begin{equation}\begin{aligned} &\,\int_{-1}^1 u((s + 1)w / 2 + T - w) g_n'(s) ds \\ =&\,\int_{-1}^1 u((s + 1)w / 2 + T - w) \sqrt{\frac{2n+1}{2}} p_n'(s) ds \\ =&\, \int_{-1}^1 u((s + 1)w / 2 + T - w)\sqrt{\frac{2n+1}{2}}\left[\sum_{k=0}^{n-1} (2k+1) \chi_{n-1-k} p_k(s)\right]ds \\ =&\, \int_{-1}^1 u((s + 1)w / 2 + T - w)\sqrt{\frac{2n+1}{2}}\left[\sum_{k=0}^{n-1} \sqrt{2(2k+1)} \chi_{n-1-k} g_k(s)\right]ds \\ =&\, \sqrt{2n+1}\sum_{k=0}^{n-1} \sqrt{2k+1} \chi_{n-1-k} c_k(T) \end{aligned}\end{equation}
Combining these results together, we get
\begin{equation}\begin{aligned} \frac{d}{dT}c_n(T) \approx &\, \frac{\sqrt{2(2n+1)}}{w}u(T) - \frac{\sqrt{2(2n+1)}}{w} (-1)^n \overbrace{\sum\limits_{k=0}^N (-1)^k c_k(T) \sqrt{\frac{2k+1}{2}}}^{u(T-w)} \\ &\quad- \frac{2}{w}\overbrace{\sqrt{2n+1}\sum_{k=0}^{n-1} \sqrt{2k+1} \chi_{n-1-k} c_k(T)}^{\int_{-1}^1 u((s + 1)w / 2 + T - w) g_n'(s) ds} \\[12pt] = &\, \frac{\sqrt{2(2n+1)}}{w}u(T) - \frac{\sqrt{2n+1}}{w} \sum\limits_{k=0}^N (-1)^{n-k} c_k(T) \sqrt{2k+1} \\ &\quad- \frac{2}{w}\sqrt{2n+1}\sum_{k=0}^{n-1} \sqrt{2k+1} \chi_{n-1-k} c_k(T) \\[12pt] = &\, \frac{\sqrt{2(2n+1)}}{w}u(T) - \frac{\sqrt{2n+1}}{w} \sum\limits_{k=n}^N (-1)^{n-k} c_k(T) \sqrt{2k+1} \\ &\quad- \frac{\sqrt{2n+1}}{w}\sum_{k=0}^{n-1} \sqrt{2k+1} \underbrace{\left(2\chi_{n-1-k} + (-1)^{n-k}\right)}_{\equiv 1}c_k(T) \\ \end{aligned}\label{eq:leg-t}\end{equation}
Again, renaming $T$ back as $t$, and stacking all the $c_n(t)$ together and denoting them as $x(t) = (c_0,c_1,\cdots,c_N)$, we can write, based on the equation above,
\begin{equation}\begin{aligned} x'(t) =&\, Ax(t) + Bu(t)\\[8pt] \quad A_{n,k} =&\, -\frac{1}{w}\left\{\begin{array}{l}\sqrt{(2n+1)(2k+1)}, &k < n \\ (-1)^{n-k}\sqrt{(2n+1)(2k+1)}, &k \geq n\end{array}\right.\\[8pt] B_n =&\, \frac{1}{w}\sqrt{2(2n+1)} \end{aligned}\label{eq:leg-t-hippo-1}\end{equation}
We can also introduce a scaling factor for each $c_n(T)$, to make the above result more general. For example, let $c_n(T) = \lambda_n \tilde{c}_n(T)$; substituting into equation $\eqref{eq:leg-t}$ and simplifying gives
\begin{equation}\begin{aligned} \frac{d}{dt}\tilde{c}_n(T) \approx &\, \frac{\sqrt{2(2n+1)}}{w\lambda_n}u(T) - \frac{\sqrt{2n+1}}{w} \sum\limits_{k=n}^N (-1)^{n-k} \tilde{c}_k(T) \frac{\lambda_k\sqrt{2k+1}}{\lambda_n} \\ &\quad- \frac{\sqrt{2n+1}}{w}\sum_{k=0}^{n-1} \frac{\lambda_k\sqrt{2k+1}}{\lambda_n} \tilde{c}_k(T) \\ \end{aligned}\end{equation}
If we take $\lambda_n = \sqrt{2}$, then $A$ stays unchanged and $B_n = \frac{1}{w}\sqrt{2n+1}$, which matches the result of the original paper. If instead we take $\lambda_n = \frac{2}{\sqrt{2n+1}}$, we get the result from Legendre Memory Units:
\begin{equation}\begin{aligned} x'(t) =&\, Ax(t) + Bu(t)\\[8pt] \quad A_{n,k} =&\, -\frac{1}{w}\left\{\begin{array}{l}2n+1, &k < n \\ (-1)^{n-k}(2n+1), &k \geq n\end{array}\right.\\[8pt] B_n =&\, \frac{1}{w}(2n+1) \end{aligned}\label{eq:leg-t-hippo-2}\end{equation}
These forms are theoretically equivalent, but may have different numerical stability. For example, generally speaking, when the behavior of $u(t)$ isn't particularly bad, we can expect that the larger $n$ is, the relatively smaller the value of $|c_n|$ tends to be. So if we use $c_n$ directly, the scales of different components of the $x(t)$ vector won't be well matched, which tends to cause numerical stability issues in actual computation. Taking $\lambda_n = \frac{2}{\sqrt{2n+1}}$ and instead using $\tilde{c}_n$ means that components with small values get appropriately amplified, which may help alleviate the multi-scale issue and make the numerical computation more stable.
The Full Interval
Now let's move on to another example: $t_{\leq T}(s) = (s + 1)T / 2$, which uniformly maps $[-1,1]$ onto $[0,T]$, meaning we don't discard any historical information and treat all history equally. The original paper calls this case "LegS (Scaled Legendre)."
Similarly, substituting into equation $\eqref{eq:hippo-base}$ gives
\begin{equation}\frac{d}{dT}c_n(T) = \frac{\sqrt{2(2n+1)}}{T}u(T) - \frac{1}{T}\int_{-1}^1 u((s + 1)T / 2) \left[g_n(s) + (s+1) g_n'(s)\right] ds\end{equation}
Using formula $\eqref{eq:leg-dot-t1}$ gives
\begin{equation}\begin{aligned} &\,\int_{-1}^1 u((s + 1)T / 2) \left[g_n(s) + (s+1) g_n'(s)\right] ds \\ =&\,c_n(T) + \int_{-1}^1 u((s + 1)T / 2) (s+1) g_n'(s) ds \\ =&\, c_n(T) + \int_{-1}^1 u((s + 1)T / 2)(s+1) \sqrt{\frac{2n+1}{2}} p_n'(s) \\ =&\, c_n(T) + \int_{-1}^1 u((s + 1)T / 2)\sqrt{\frac{2n+1}{2}}\left[-(n+1) p_n(s) + \sum_{k=0}^n (2k + 1) p_k(s)\right] ds \\ =&\, c_n(T) + \int_{-1}^1 u((s + 1)T / 2)\left[-(n+1) g_n(s) + \sum_{k=0}^n \sqrt{(2n+1)(2k + 1)} g_k(s)\right] ds \\ =&\, -n c_n(T) + \sum_{k=0}^n \sqrt{(2n+1)(2k + 1)} c_k(T) \\ \end{aligned}\end{equation}
and thus
\begin{equation}\frac{d}{dT}c_n(T) = \frac{\sqrt{2(2n+1)}}{T}u(T) - \frac{1}{T}\left(-n c_n(T) + \sum_{k=0}^n \sqrt{(2n+1)(2k + 1)} c_k(T)\right)\label{eq:leg-s}\end{equation}
Renaming $T$ back as $t$, and stacking all the $c_n(t)$ together and denoting them as $x(t) = (c_0,c_1,\cdots,c_N)$, we can write, based on the equation above,
\begin{equation}\begin{aligned} x'(t) =&\, \frac{A}{t}x(t) + \frac{B}{t}u(t)\\[8pt] \quad A_{n,k} =&\, -\left\{\begin{array}{l}\sqrt{(2n+1)(2k+1)}, &k < n \\ n+1, &k = n \\ 0, &k > n\end{array}\right.\\[8pt] B_n =&\, \sqrt{2(2n+1)} \end{aligned}\label{eq:leg-s-hippo}\end{equation}
It's also possible to introduce a scaling factor to generalize the result: let $c_n(T) = \lambda_n \tilde{c}_n(T)$; substituting into equation $\eqref{eq:leg-t}$ and simplifying gives
\begin{equation}\frac{d}{dT}\tilde{c}_n(T) = \frac{\sqrt{2(2n+1)}}{T\lambda_n}u(T) - \frac{1}{T}\left(-n \tilde{c}_n(T) + \sum_{k=0}^n \frac{\sqrt{(2n+1)(2k + 1)}\lambda_k}{\lambda_n} \tilde{c}_k(T)\right)\end{equation}
Taking $\lambda_n=\sqrt{2}$ leaves $A$ unchanged and turns $B$ into $B_n = \sqrt{2n+1}$, which matches the result of the original paper. If instead we take $\lambda_n=\sqrt{\frac{2}{2n+1}}$, we can, just as in the LegT result in the previous section, remove the square root:
\begin{equation}\begin{aligned} x'(t) =&\, \frac{A}{t}x(t) + \frac{B}{t}u(t)\\[8pt] \quad A_{n,k} =&\, -\left\{\begin{array}{l}2n+1, &k < n \\ n+1, &k = n \\ 0, &k > n\end{array}\right.\\[8pt] B_n =&\, 2n+1 \end{aligned}\label{eq:leg-s-hippo-2}\end{equation}
But the original paper doesn't consider this case, for reasons unknown.
Further Thoughts
Looking back at the whole derivation of Leg-S, we can see that a key step is expanding $(s+1) g_n'(s)$ as a linear combination of $g_0(s),g_1(s),\cdots,g_n(s)$. For orthogonal polynomials, $(s+1) g_n'(s)$ is a degree-$n$ polynomial, so this expansion is guaranteed to hold exactly. But in the case of Fourier series, $g_n(s)$ is an exponential function, and a similar expansion isn't achievable there, at least not exactly — so it's fair to say that the fundamental purpose of choosing orthogonal polynomials as the basis is to simplify the subsequent derivation.
It's especially worth pointing out that HiPPO is a bottom-up framework: it doesn't assume from the outset that the system must be linear, but rather derives, starting from orthogonal-basis approximation, that the dynamics of the coefficients satisfy a linear ODE system. This gives us confidence that, as long as we accept the assumptions made, the capacity of the linear ODE system is indeed sufficient — we don't need to worry that the linearity of the system limits what it can express.
Of course, HiPPO's assumptions for each solution, and their physical meaning, are also quite clear. So, for SSMs that reuse the HiPPO matrix, exactly how they store history and how much history they can store can be read directly off the underlying HiPPO assumptions. For example, LegT only keeps the information within the nearest window of size $w$ — if you use LegT's HiPPO matrix, this is analogous to Sliding Window Attention. LegS, in theory, can capture the entire history, but there's a resolution issue: the dimension of $x(t)$ represents the order of the fit, which is a fixed value, and using a fixed-order function basis to fit a function is always going to be more accurate over a smaller interval and less accurate over a larger one. This is much like having to stand further back in order to see a large picture all at once, at the cost of seeing less detail.
Models such as RWKV and LRU don't reuse the HiPPO matrix, opting instead for a trainable matrix, which in principle offers more possibilities for breaking through this bottleneck. But from the analysis above, we can roughly sense that, across different matrices, the linear ODE just corresponds to a different choice of function basis, and fundamentally it may still just be the coefficient dynamics of a finite-order function-basis approximation. If that's the case, resolution and memory length still can't both be had for free — if you want to remember longer inputs while keeping performance unchanged, the only option is to increase the overall size of the model (i.e., equivalent to increasing hidden_size). This is probably a characteristic shared by all linear systems.
Summary
This post has, as simply as possible, reproduced the main derivation of HiPPO: Recurrent Memory with Optimal Polynomial Projections (abbreviated HiPPO). Through an appropriate memory assumption, HiPPO derives a linear ODE system in a bottom-up manner, and works out the corresponding analytical solution (the HiPPO matrix) for the Legendre polynomial case. Its results have since been used by many subsequent SSMs (State Space Models), making it a genuinely foundational piece of work for SSM.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.