Revisiting SSM (II): Some Leftover Questions about HiPPO
Continuing from where we left off, in the previous post Revisiting SSM (I): Linear Systems and the HiPPO Matrix] we discussed in detail the derivation of the HiPPO approximation framework and its HiPPO matrix. The underlying principle is to dynamically approximate a function that keeps getting updated in real time, using an orthogonal function basis; the dynamics of the projection coefficients then turn out to be exactly a linear system, and if we choose orthogonal polynomials as the basis, the core matrix of that linear system can be solved analytically — this is the matrix called the HiPPO matrix.
Of course, the previous post focused on deriving the HiPPO matrix and did not go further into analyzing its properties. Questions such as "how to discretize it for use on real data" and "whether other bases besides polynomials can also be solved analytically" were not discussed in detail either. In this post we will fill in some of these gaps.
Discretization Schemes
Assuming the reader has already read and understood the previous post, we won't repeat too much of the setup here. In the previous post, we derived two classes of linear ODE systems:
\begin{align} &\text{HiPPO-LegT:}\quad x'(t) = Ax(t) + Bu(t) \label{eq:legt-ode}\\[5pt] &\text{HiPPO-LegS:}\quad x'(t) = \frac{A}{t}x(t) + \frac{B}{t}u(t) \label{eq:legs-ode}\end{align}
where $A,B$ is a constant matrix independent of time $t$, and the HiPPO matrix mainly refers to the matrix $A$. In this section, we discuss how to discretize these two ODEs. more
Transforming the Input
In real-world scenarios, the input data points form a discrete sequence $u_0,u_1,u_2,\cdots,u_k,\cdots$ — for example a streaming audio signal or text embeddings — and we want to use the ODE system above to memorize these discrete points in real time. To do this, we first define
\begin{equation}u(t) = u_k,\quad \text{if} t\in[k\epsilon, (k + 1)\epsilon)\end{equation}
where $\epsilon$ is the discretization step size. This definition means that on the interval $[k\epsilon, (k + 1)\epsilon)$, $u(t)$ is a constant function equal to $u_k$. Clearly, $u(t)$ defined this way loses no information from the original sequence $u_k$, so memorizing $u(t)$ is equivalent to memorizing the sequence $u_k$.
Going from $u_k$ to $u(t)$ turns the input signal back into a function on a continuous interval, which is convenient for later operations such as integration. In addition, keeping it constant within each discretized interval also simplifies the resulting discretized scheme.
The LegT Version
Let's start with the LegT-type ODE $\eqref{eq:legt-ode}$, integrating both sides over the interval:
\begin{equation}x(t+\epsilon) - x(t) = A\int_t^{t+\epsilon} x(s)ds + B\int_t^{t+\epsilon}u(s)ds\end{equation}
where $t=k\epsilon$. By the definition of $u(t)$, it is constantly equal to $u_k$ over the interval $[t, t + \epsilon)$, so the integral of $u(s)$ can be computed directly:
\begin{equation}x(t+\epsilon) - x(t) = A\int_t^{t+\epsilon} x(s)ds + \epsilon B u_k\end{equation}
What comes next depends on how we approximate the integral of $x(s)$. If we assume that $x(s)$ is approximately constant and equal to $x(t)$ on the interval $[t, t + \epsilon)$, we get the forward Euler scheme:
\begin{equation}x(t+\epsilon) - x(t) = \epsilon A x(t) + \epsilon B u_k \quad\Rightarrow\quad x(t+\epsilon) = (I + \epsilon A)x(t) + \epsilon B u_k\end{equation}
If instead we assume that $x(s)$ is approximately constant and equal to $x(t+\epsilon)$ on the interval $[t, t + \epsilon)$, we get the backward Euler scheme:
\begin{equation}x(t+\epsilon) - x(t) = \epsilon A x(t+\epsilon) + \epsilon B u_k \quad\Rightarrow\quad x(t+\epsilon) = (I - \epsilon A)^{-1}(x(t) + \epsilon B u_k)\end{equation}
Forward and backward Euler have the same theoretical order of accuracy, but the backward scheme is usually more numerically stable. To be more accurate, if we assume that $x(s)$ is approximately constant and equal to $\frac{1}{2}[x(t) + x(t+\epsilon)]$ on the interval $[t, t + \epsilon)$, we get the bilinear form:
\begin{equation}\begin{gathered} x(t+\epsilon) - x(t) = \frac{1}{2}\epsilon A [x(t) + x(t+\epsilon)] + \epsilon B u_k \\ \Downarrow \\ x(t+\epsilon) = (I - \epsilon A/2)^{-1}[(I + \epsilon A/2) x(t) + \epsilon B u_k] \end{gathered}\end{equation}
which is equivalent to taking a half-step of forward Euler followed by a half-step of backward Euler. More generally, we could assume that $x(s)$ is approximately constant and equal to $\alpha x(t) + (1 - \alpha) x(t+\epsilon)$ on the interval $[t, t + \epsilon)$, where $\alpha\in[0,1]$, but we won't expand on this further here. In fact, we don't need to approximate at all: combining equation $\eqref{eq:legt-ode}$ with the fact that $u(s)$ is constant, equal to $u_k$, on the interval $[t,t+\epsilon)$, we can solve exactly using the "variation of parameters]" method, giving
\begin{equation}x(t+\epsilon) = e^{\epsilon A} x(t) + A^{-1} (e^{\epsilon A} - I) B u_k\label{eq:legt-ode-sol}\end{equation}
Here the matrix exponential is defined via its series expansion; see An Appreciation of the Identity det(exp(A)) = exp(Tr(A))] for reference.
The LegS Version
Now let's turn to the LegS-type ODE, whose treatment is basically the same as LegT's, with similar results. First, integrate both sides of equation $\eqref{eq:legs-ode}$ to get
\begin{equation}x(t+\epsilon) - x(t) = A\int_t^{t+\epsilon} \frac{x(s)}{s}ds + B\int_t^{t+\epsilon}\frac{u(s)}{s}ds\end{equation}
By the definition of $u(t)$, the $u(s)$ in the second integral is constantly equal to $u_k$ over $[t,t+\epsilon)$, so it amounts to the integral of $1/s$, which can be computed directly to give $\ln\frac{t+\epsilon}{t}$. Of course, replacing it directly with the first-order approximation $\frac{\epsilon}{t}$ is also fine, since the transformation from $u_k$ to $u(t)$ already has a lot of freedom built in, so this small extra error doesn't matter much. As for the first integral, we use the more accurate midpoint approximation directly, giving
\begin{equation}\begin{gathered} x(t+\epsilon) - x(t) = \frac{1}{2}\epsilon A\left(\frac{x(t)}{t}+\frac{x(t+\epsilon)}{t+\epsilon}\right) + \frac{\epsilon}{t} B u_k \\[5pt] \Downarrow \\[5pt] x(t+\epsilon) = \left(I - \frac{\epsilon A}{2(t+\epsilon)}\right)^{-1}\left[\left(I + \frac{\epsilon A}{2t}\right)x(t) + \frac{\epsilon}{t} B u_k\right] \end{gathered}\label{eq:legs-ode-bilinear}\end{equation}
In fact, equation $\eqref{eq:legs-ode}$ can also be solved exactly — we just need to notice that it is equivalent to
\begin{equation}Ax(t) + Bu(t) = t x'(t) = \frac{d}{d\ln t} x(t)\end{equation}
which means that with the change of variables $\tau = \ln t$, the LegS-type ODE can be converted into a LegT-type ODE:
\begin{equation}\frac{d}{d\tau} x(e^{\tau}) = Ax(e^{\tau}) + Bu(e^{\tau})\end{equation}
Using equation $\eqref{eq:legt-ode-sol}$ (and noting that, because of the change of variables, the time interval changes from $\epsilon$ to $\ln(t+\epsilon) - \ln t$) we get
\begin{equation}x(t+\epsilon) = e^{(\ln(t+\epsilon) - \ln t) A} x(t) + A^{-1} \big(e^{(\ln(t+\epsilon) - \ln t) A} - I\big) B u_k\label{eq:legs-ode-sol}\end{equation}
However, although the equation above is an exact solution, it is not as convenient to use as the equally exact equation $\eqref{eq:legt-ode-sol}$, because the exponential-matrix part of equation $\eqref{eq:legt-ode-sol}$ is $e^{\epsilon A}$, which is independent of time $t$ and so only needs to be computed once. But in the equation above, $t$ appears inside the matrix exponential, meaning that during iteration the matrix exponential must be recomputed repeatedly — not computationally friendly. So for the LegS-type ODE we generally just use equation $\eqref{eq:legs-ode-bilinear}$ for discretization.
Nice Properties
Next, LegS is our main focus. It's not hard to guess why: based on the derivation's assumptions, it is the only ODE system we've solved so far that is capable of memorizing the entire history, which is crucial for many scenarios such as multi-turn dialogue. Beyond that, it has several other nice and practically useful properties.
Scale Equivariance
For instance, the LegS discretization scheme $\eqref{eq:legs-ode-bilinear}$ is independent of the step size. Simply substitute $t=k\epsilon$ into it and write $x(k\epsilon)=x_k$, and we find
\begin{equation} x_{k+1} = \left(I - \frac{A}{2(k + 1)}\right)^{-1}\left[\left(I + \frac{A}{2k}\right)x_k + \frac{1}{k} B u_k\right]\end{equation}
The step size $\epsilon$ automatically cancels out, which naturally removes one hyperparameter we'd otherwise need to tune — clearly good news for practitioners doing model training. Note that step-size independence is an intrinsic property of the LegS-type ODE and has nothing to do with the specific discretization method used; for instance, the exact solution $\eqref{eq:legs-ode-sol}$ is likewise independent of the step size:
\begin{equation}x_{k+1} = e^{(\ln(k+1) - \ln k) A} x_k + A^{-1} \big(e^{(\ln(k+1) - \ln k) A} - I\big) B u_k\label{eq:legs-ode-sol-2}\end{equation}
The underlying reason is that the LegS-type ODE satisfies timescale equivariance: if we set $t=\lambda\tau$ and substitute into the LegS-type ODE, we get
\begin{equation}Ax(\alpha\tau) + Bu(\alpha\tau) = (\alpha\tau)\times \frac{d}{d(\alpha\tau)} x(\alpha\tau) = \tau \frac{d}{d\tau}x(\alpha\tau)\end{equation}
This means that when we replace $u(t)$ with $u(\alpha t)$, the form of the LegS ODE doesn't change, while the corresponding solution has $x(t)$ replaced by $x(\alpha t)$. A direct consequence of this property is: when we choose a larger step size, the recursive scheme doesn't need to change, because the resulting $x_k$ automatically scales up with the step size as well — this is the essential reason why LegS-type ODE discretization is independent of the step size.
Long-Tail Decay
Another nice property of the LegS-type ODE is that its memory of historical signals decays polynomially, which is slower than the exponential decay of ordinary RNNs, so in theory it can remember a much longer history and is less prone to vanishing gradients. To understand this, we start from the exact solution $\eqref{eq:legs-ode-sol-2}$. From equation $\eqref{eq:legs-ode-sol-2}$ we can see that at each recursive step, the decay effect on the historical information is described by the matrix exponential $e^{(\ln(k+1) - \ln k) A}$, so going from step $m$ to step $n$, the total decay effect is
\begin{equation}\prod_{k=m}^{n-1} e^{(\ln(k+1) - \ln k) A} = e^{(\ln n - \ln m) A}\end{equation}
Recall the form of $A$ in HiPPO-LegS:
\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}
From this definition we can see that $A$ is a lower triangular matrix, with diagonal entries $-1,-2,-3,\cdots$. We know that the diagonal entries of a triangular matrix are exactly its eigenvalues (see Triangular matrix]), so an $d\times d$-sized $A$ matrix has $d$ distinct eigenvalues $-1,-2,\cdots,-d$. This shows that the matrix $A$ is diagonalizable, i.e., there exists an invertible matrix $P$ such that $A = P^{-1}\Lambda P$, where $\Lambda = \text{diag}(-1,-2,\cdots,-d)$, and thus we have
\begin{equation}\begin{aligned} e^{(\ln n - \ln m) A} =&\, e^{(\ln n - \ln m) P^{-1}\Lambda P} \\ =&\, P^{-1} e^{(\ln n - \ln m) \Lambda}P \\ =&\, P^{-1}\,\text{diag}(e^{-(\ln n - \ln m)},e^{-2(\ln n - \ln m)},\cdots,e^{-d(\ln n - \ln m)})\,P \\ =&\, P^{-1}\,\text{diag}\Big(\frac{m}{n},\frac{m^2}{n^2},\cdots,\frac{m^d}{n^d}\Big)\,P \\ \end{aligned}\end{equation}
We can see that the final decay function is a linear combination of powers $1,2,\cdots,d$ of $1/n$, so the memory decay of the LegS-type ODE with respect to history is at most polynomial — a much longer tail than exponential decay, and hence theoretically better memory capacity.
Computational Efficiency
Finally, we point out that the $A$ matrix of HiPPO-LegS is computationally efficient. Specifically, a naive matrix multiplication of an $d\times d$ matrix by an $d\times 1$-dimensional column vector requires $d^2$ multiplications, but multiplying the LegS $A$ matrix by a vector can be reduced to $\mathcal{O}(d)$ operations. Moreover, we can further show that the discretized $\eqref{eq:legs-ode-bilinear}$ can also be computed in $\mathcal{O}(d)$.
To see this, let's first rewrite the HiPPO-LegS $A$ matrix 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}
For a vector $v = [v_0,v_1,\cdots,v_{d-1}]$, we have
\begin{equation}\begin{aligned} (Av)_n = \sum_{k=0}^n A_{n,k}v_k =&\, \sum_{k=0}^n \left(n\delta_{n,k} - \sqrt{2n+1}\sqrt{2k+1}\right)v_k \\ =&\, n v_n -\sqrt{2n+1}\sum_{k=0}^n \sqrt{2k+1}v_k \end{aligned}\end{equation}
This involves three operations: in the first term, $n v_n$ is the elementwise product of vector $[0,1,2,\cdots,d-1]$ and $v$; in the second term, $\sqrt{2k+1}v_k$ is the elementwise product of vector $[1,\sqrt{3},\sqrt{5},\cdots,\sqrt{2d-1}]$ and $v$; then $\sum\limits_{k=0}^n$ is a $\text{cumsum}$ operation (a cumulative sum), and finally multiplying by $\sqrt{2n+1}$ is another elementwise product with the vector $[1,\sqrt{3},\sqrt{5},\cdots,\sqrt{2d-1}]$. Each step can be done in $\mathcal{O}(d)$, so the total complexity is $\mathcal{O}(d)$.
Now let's look at $\eqref{eq:legs-ode-bilinear}$, which involves two "matrix-vector" multiplications: the first is $(I+\lambda A)v$, where $\lambda$ is an arbitrary real number — we've already shown that $Av$ is computationally efficient, so naturally so is $(I+\lambda A)v$; the second is $(I-\lambda A)^{-1}v$, which we will now show is also computationally efficient. To see this, note that computing $z=(I-\lambda A)^{-1}v$ is equivalent to solving the equation $v = (I-\lambda A)z$. Using the expression for $Av$ given above, we get
\begin{equation}v_n = z_n - \lambda \left(n z_n - \sqrt{2n+1}\sum_{k=0}^n \sqrt{2k+1}z_k\right)\end{equation}
Writing $S_n = \sum\limits_{k=0}^n \sqrt{2k+1}z_k$, we have $z_n = \frac{S_n - S_{n-1}}{\sqrt{2n+1}}$; substituting this into the above equation gives
\begin{equation}v_n = \frac{S_n - S_{n-1}}{\sqrt{2n+1}} - \lambda \left(n \frac{S_n - S_{n-1}}{\sqrt{2n+1}} - \sqrt{2n+1}S_n\right)\end{equation}
which rearranges to
\begin{equation}S_n = \frac{1 - \lambda n}{1+\lambda n + \lambda}S_{n-1} + \frac{\sqrt{2n+1}}{1+\lambda n + \lambda}v_n\end{equation}
This is a scalar recurrence that can be computed either fully sequentially, or in parallel using prefix-sum algorithms (see here]), with a computational complexity of $\mathcal{O}(d)$ or $\mathcal{O}(d\log d)$ — in any case, more efficient than $\mathcal{O}(d^2)$.
The Fourier Basis
Finally, let's close with a derivation using the Fourier basis. In the previous post, we used the Fourier series to introduce the linear system, but only derived the result for the sliding-window form; for the Legendre polynomial basis we derived both the sliding-window and full-interval versions (i.e., LegT and LegS). So can the Fourier basis be used to derive a version comparable to LegS? And what difficulties would we run into? Let's explore this below.
Again, we won't repeat the setup — following the notation of the previous section, the coefficients under the Fourier basis are
\begin{equation}c_n(T) = \int_0^1 u(t_{\leq T}(s)) e^{-2i\pi n s}ds\end{equation}
As with LegS, in order to memorize the signal over the entire interval $[0,T]$, we need a mapping of the form $[0,1]\mapsto [0,T]$; for this, we choose the simplest option $t_{\leq T}(s)=sT$. Substituting this in and differentiating both sides gives
\begin{equation}\frac{d}{dT}c_n(T) = \int_0^1 u'(sT) s e^{-2i\pi n s}ds\end{equation}
Integrating by parts gives
\begin{equation}\begin{aligned} \frac{d}{dT}c_n(T) =&\, \frac{1}{T}\int_0^1 s e^{-2i\pi n s}d u(sT) \\ =&\, \frac{1}{T} u(sT) s e^{-2i\pi n s}\big|_{s=0}^{s=1} - \frac{1}{T}\int_0^1 u(sT) d(s e^{-2i\pi n s})\\ =&\, \frac{1}{T} u(T) - \frac{1}{T}\int_0^1 u(sT) e^{-2i\pi n s} ds + \frac{2i\pi n}{T}\int_0^1 u(sT) s e^{-2i\pi n s} ds\\ =&\, \frac{1}{T} u(T) - \frac{1}{T}c_n(T) + \frac{2i\pi n}{T}\int_0^1 u(sT) s e^{-2i\pi n s} ds\\ \end{aligned}\end{equation}
In the previous post, we mentioned that one of the key reasons HiPPO chose Legendre polynomials as its basis is that $(s+1)p_n'(t)$ can be decomposed as a linear combination of $p_0(t),p_1(t),\cdots,p_n(t)$, whereas the Fourier basis's $s e^{-2i\pi n s}$ cannot achieve this. But in fact, if we allow some error, this claim doesn't quite hold, because we can likewise expand $s$ as a Fourier series:
\begin{equation}s = \frac{1}{2} + \frac{i}{2\pi}\sum_{k\neq 0} \frac{1}{k} e^{2i\pi k s}\end{equation}
This sum has infinitely many terms; truncating it to a finite number of terms introduces some error, but let's not worry about that for now and just substitute it in directly, giving
\begin{equation}\begin{aligned} &\,\frac{2i\pi n}{T}\int_0^1 u(sT) s e^{-2i\pi n s} ds \\ =&\, \frac{2i\pi n}{T}\int_0^1 u(sT) \left(\frac{1}{2} + \frac{i}{2\pi}\sum_{k\neq 0} \frac{1}{k} e^{2i\pi k s}\right) e^{-2i\pi n s} ds \\ =&\, \frac{i\pi n}{T}\int_0^1 u(sT) e^{-2i\pi n s} ds - \frac{1}{T}\sum_{k\neq 0} \frac{n}{k}\int_0^1 u(sT) e^{-2i\pi (n - k) s} ds \\ =&\, \frac{i\pi n}{T}c_n(T) - \frac{1}{T}\sum_{k\neq 0} \frac{n}{k}c_{n-k}(T) \\ =&\, \frac{i\pi n}{T}c_n(T) - \frac{1}{T}\sum_{k\neq n} \frac{n}{n - k}c_k(T) \\ \end{aligned}\end{equation}
This gives us
\begin{equation} \frac{d}{dT}c_n(T) = \frac{1}{T} u(T) + \frac{i\pi n - 1}{T}c_n(T) - \frac{1}{T}\sum_{k\neq n} \frac{n}{n - k}c_k(T)\end{equation}
and so we can write
\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}-\frac{n}{n-k}, &k \neq n \\ i\pi n - 1, &k = n\end{array}\right.\\[8pt] B_n =&\, 1 \end{aligned}\end{equation}
In practice, we only need to truncate $|n|,|k|\leq N$ to obtain an $(2N+1)\times (2N+1)$ matrix. The error introduced by truncation really doesn't matter, because when deriving HiPPO-LegT we likewise introduced a finite-series approximation, and we didn't worry about the error there either — or, put the other way around, for a specific task we would choose an appropriate scale (i.e., the size of $N$), and part of what "appropriate" means here is precisely that the error from truncation is negligible for that task.
For most people, this derivation using the Fourier basis is probably easier to follow, since Legendre polynomials are unfamiliar to many readers, especially the several identities used in deriving LegT and LegS, whereas most readers have at least some familiarity with Fourier series. That said, in terms of the end result, this Fourier-basis version is probably less practical than LegS: first, it introduces complex numbers, which adds implementation complexity; second, the $A$ matrix it derives is not a relatively sparse lower triangular matrix like LegS's, making theoretical analysis considerably more complicated. So, it's best to treat this as an exercise for deepening one's understanding of HiPPO, rather than something to use in practice.
Summary
In this post, we followed up on some leftover questions from the previous article's introduction to HiPPO, including how to discretize the ODEs, some of the nice properties of the LegS-type ODE, and a derivation of a result for memorizing the entire historical interval using the Fourier basis (i.e., a Fourier version of LegS) — all in the hope of achieving a more complete understanding of HiPPO.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.