ai papers
Foundations of Mamba Modeling (2) - Efficiently Modeling Long Sequences with Structured State Spaces (S4)
Junyoung Park · 2024-02-22 · 18 min
Before we begin
HiPPO emphasized the importance of the structure of matrix for establishing long-term range in an SSM, and LSSL generalized the construction across all continuously defined . In the previous post, I explained LSSL as one of the foundations of Mamba modeling and included brief derivations of the relevant equations. To be honest, I still do not think I fully understand every one of these fundamentals, and trying to understand Mamba through this route can feel excessively roundabout. That may sound unconfident, but ultimately I do not believe that understanding Mamba modeling requires internalizing state-space modeling all the way down to the bone. My personal view is that top-down, rather than bottom-up, is the right direction.
Many blog posts also explain Mamba clearly in one pass with the help of visualizations. Even so, the reason I chose to write a long review is the same one I gave in the previous post: the attempt to understand Mamba’s foundations has value of its own. If we understand Mamba only at a surface level, we cannot explain precisely why it can address problems in the Transformer architecture—and we should not assume that an architecture is unconditionally good merely because it is new.
What comes easily is lost just as easily. A sound understanding of Mamba’s characteristics should reveal its fundamental advantages and disadvantages, which can then provide a foundation for future research and model development. This post covers S4 modeling, which seeks to resolve LSSL’s computational cost and numerical instability.
Returning to sequence modeling
The central questions in sequence modeling are how efficiently a model can process sequences of a useful length, and whether the amount of accessible context grows with the sequence length. A Transformer sacrifices computation to obtain attention information over the full sequence in a single operation and can extend its reach by increasing the token-embedding length, but ultimately runs into computational limits. RNNs and CNNs have their own characteristics, yet no modeling approach is universal because each comes with trade-offs. Papers on LMU, HiPPO, and related methods proposed state-space modeling (SSM) in response.
To summarize SSM
In one sentence, SSM seeks to construct a latent space with a Linear Time-Invariant system. The matrices in an LTI system’s differential equation are time invariant. At every observation unit obtained by discretizing a continuous signal, the differential equation therefore retains the same form and can create the same latent space regardless of sequence length. Unlike an RNN that depends on gates, if matrix has a particular structure, it can extract a long-range-dependent latent—the feature vector most important to a deep-learning model’s predictions. Rather than computing attention token by token with a Transformer, or forming an implicit latent from a stream of data through an RNN, we can predict a genuinely time-invariant latent space. No matter how much time passes—even when the observation range has moved far from its beginning—the hidden space is implemented by the same function, thereby supporting long-range dependency. An SSM can consequently process longer sequences than a Transformer with a fixed attention context and can theoretically handle inputs of unlimited length.
Let us use the paper’s figure to build intuition. In the “Continuous State Space” section on the left, the SSM input produces an output modeled through the state of a linear system composed of matrices and . That state is called the hidden state, or latent state . If acts as memory that continuously stores the input-to-output transformation , then fixed, theoretically invariant matrices and can continually retain preceding information. This is the “Long-Range Dependencies” shown in the center, and the paper that demonstrated the effectiveness of a structured matrix for this purpose was “HiPPO: Recurrent Memory with Optimal Polynomial Projections.”
An SSM fundamentally has a recurrent system. Its structure maps input to output, and the change in system state resulting from an earlier input and output affects subsequent inputs and outputs. This can be implemented inductively as recurrent operations, or simplified into a convolution. Yet the matrix operations that necessarily grow for high-dimensional data—for example, operations on or —still create a problem of high computational cost.
What if we use a “well”-structured SSM? (S4)
This brings us to the objective of “Efficiently Modeling Long Sequences with Structured State Spaces” (S4). Methods had been proposed to reduce the high computational cost of conventional LSSL—a plain SSM—but all lacked numerical stability or rigor. To improve numerical stability while taking advantage of algorithms applicable to existing “well”-structured matrices, S4 restructures matrix , the basis of the SSM. It exploits the fact that can be decomposed into a component whose rank is far lower than that of the original and a normal-matrix component. Rank is the number of independent columns; a low-rank matrix can be represented in fewer dimensions after dependent columns are removed. A normal matrix is a special square matrix whose properties—including commutation and diagonalizability—make it faster to compute with than a non-normal matrix.
Whereas the conventional SSM used a coefficient-space approach—the function representing the latent was expressed as a weighted sum of coefficients on predefined orthogonal functions—S4 lifts the computation into the frequency domain. Convolution in the time domain becomes multiplication in the frequency domain.
This lets the low-rank matrix be handled through the Woodbury identity and the normal matrix through a Cauchy kernel, reducing computation from to and memory from to .
Conventional SSMs and their representation
We use matrix that is modeled to preserve long contexts, such as a HiPPO matrix. This constructs a system expressed by the following simultaneous differential equations:
A conventional computing system must convert the continuous system into a discrete one with discretized inputs. Applying that conversion produces the discrete SSM used in practice:
For the proof and more detail, see the previous post on LSSL. The important implication of the discrete SSM is that, because an SSM ultimately has a recurrent computational structure, it shares the computational characteristics of an RNN. The discretized matrix serves as the transition matrix for hidden state . If we expand the hidden states and outputs above while assuming , however, the operation can also be expressed with a convolution kernel:
If the convolutional filter were known, an FFT—a fast convolution algorithm—could improve computation speed. But computing the filter itself requires matrix multiplication, and generalizing the operation to a non-normal matrix remains difficult. The paper’s task is therefore to compute this filter efficiently.
Method: Diagonalization
Matrix diagonalization uses a matrix’s eigenvalues and eigenvectors to transform it into a matrix whose diagonal entries are the eigenvalues. If a matrix is diagonalizable, for example, its eigenvalues and corresponding eigenvectors let us write . Here and are
Under these conditions, we can rewrite the original equations as follows. SSMs generally simplify with , so we will write .
The two systems may look different, but multiplying every expression on the left side of the transformed system by shows that it is equivalent to an SSM with . The input-to-output relationship remains the same; only the system latent changes according to eigenvector matrix of matrix . If we can construct the diagonal matrix that plays the role of in the transformed system, we can efficiently reduce the cost of repeated multiplication by in the convolution above. This is called a Vandermonde product; we will simply take that as given.
Unfortunately, the HiPPO matrix cannot be diagonalized stably.
The reason is that when the HiPPO matrix is diagonalized, the entries of the eigenvector matrix become far too large. Put simply, numerically stable computation—agreement with the true mathematical value—requires matrix entries not to grow excessively during matrix operations. Diagonalizing the HiPPO matrix shown earlier produces entries . As the state size grows, entries as large as become unmanageable and can, for example, break the output computation .
Method: Normal + Low-Rank Matrix
The preceding section attempted to make computation easier through diagonalization, but showed that applying it to a basic HiPPO matrix requires additional work. The ideal case is for the target matrix to be diagonalizable by a special matrix such as a unitary matrix. In linear algebra, matrices satisfying this condition are called normal matrices. As you may have guessed, a HiPPO matrix is not normal, which is why its eigenvector entries grow during diagonalization.
Fortunately, the authors discovered that although HiPPO matrix is not normal, it can be represented as the sum of a normal matrix and a low-rank matrix. A new problem appears, however: computing the convolution filter requires taking powers of this sum (Normal + Low Rank), which is itself slow and needs optimization.
To address the problem, the kernel filter is computed with three additional algorithms. The following figure summarizes how they relate. We have not yet covered enough detail to understand it, so let us proceed in order.
To follow the algorithm, we first need to accept that matrix can be represented as NPLR (Normal Plus Low-Rank) or DPLR (Diagonal Plus Low-Rank). Appendix C.1 of the paper demonstrates this empirically for every HiPPO matrix. There is no real need to understand the equations in this part; it can simply be accepted as a fact.
The method then departs from the conventional calculation of the convolution kernel. Instead of directly computing , it uses , the Discrete Fourier Transform (DFT) of . The DFT converts a signal on a discrete time axis—here, think of the filter entries as samples on a continuous time axis—into a spectrum on a discrete frequency axis. The algorithms for the DFT and its inverse, the IDFT, are known as the Fast Fourier Transform (FFT), which runs in .
We will cover the details below, but constructing a truncated SSM and computing in the spectral domain replaces repeated multiplications by during filter computation with a single matrix operation. This introduces the troublesome term —which apparently still requires exponentiation—but reparameterization avoids repeated computation, saves memory, and improves speed. The structure assumed above then shows that computing the spectral kernel is equivalent to a Cauchy kernel, allowing an efficient algorithm. In brief, the sequence is as follows.
- Every HiPPO matrix can be represented as NPLR (or DPLR). This structure reduces the discretization to computation.
- The truncated-SSM generating function of is equivalent to the DFT. We can therefore compute through transforms into and out of the frequency domain, replacing repeated powers of with a single operation.
- This operation has the same structure as a Cauchy-kernel computation, so an efficient algorithm is available.
- The inverse can be simplified with the Woodbury identity.
A more detailed explanation of the method
Let us examine the efficiency improvements in detail. Personally, I find this harder than diffusion papers. Still, let us press on.
Every HiPPO matrix can be represented as NPLR, or Normal Plus Low-Rank. Four HiPPO matrices exist; considering only the most common case, LegS, gives
Adding to every entry produces
Separating the diagonal entries gives , where is a skew-symmetric matrix, which is normal. The matrix that adds the same value to every entry of the original matrix also has rank . See the paper for proofs covering every HiPPO matrix. The implication is that every HiPPO state matrix can be separated into a diagonal part and a low-rank part. We will write this as follows. The paper uses the conjugate symbol () in its proof; because that symbol is less familiar to me, I will replace it with the transpose symbol ().
Using this substitution to express the discrete system matrices gives
Computing poses no serious problem, but contains a major obstacle: the matrix inverse. The computation grows dramatically as the matrix dimension increases. The inverse of a DPLR matrix can be simplified as above with the Woodbury identity. A diagonal matrix is easy to invert, and the attached low-rank term can be handled without inverting the entire expression. Because the Woodbury identity appears repeatedly in the derivation below, it is useful to keep it in mind: it can be applied efficiently whenever a matrix has DPLR structure.
Woodbury's identity applies as follows. For three matrices whose elements form a commutative ring, —where is a commutative ring whose elements commute under multiplication—
Using the resulting expressions, the discrete system can be defined for the DPLR matrices and with computation:
At this point the authors treat the former row vector as a column vector so that its shape matches the other parameters and . Following them, from here onward I will write the in the original system equation as .
We now understand the discretized system matrices. The most important task is to compute the convolution filter of the recurrent system quickly. DPLR improves the efficiency of matrix discretization by enabling the Woodbury identity, but it does not help much with the repeated multiplications needed to compute the convolution filter itself. Exploiting DPLR requires an inverse, so we can instead move to the spectral domain with a generating function of the coefficients. Suppose, for example, that we have an infinite convolution-filter signal:
What we actually have is a finite filter of length , truncated from that ideal convolution filter:
A discrete signal of length can be decomposed into components with frequencies from through . If a variable represents a frequency unit, the signal can be replaced by a set of coefficients of a function of ; this is called the -transform. In general, is a complex number (Real + Imaginary). In frequency units, Euler’s angular conversion lets us express it over .
This transform is the DFT (Discrete Fourier Transform), commonly used to convert a signal on a discrete time axis into one on a discrete frequency axis. Setting the leading factor to gives
We can substitute the closed forms of discretized matrices to express this in terms of :
Substituting directly into the expression above yields
Because we showed that the system matrix is DPLR (Diagonal Plus Low-Rank), we can rewrite this as
Now we can use the Woodbury identity introduced earlier. Defining the following quantity to simplify the expression,
leads to the final expansion
One question concerns the earlier definition . Ordinarily, would have to be computed repeatedly. By reparameterizing the trainable parameter so that it is initialized directly as , however, we can reduce the associated computational cost.
We have reached the final step. To summarize, the operation that computes has been replaced by the generating function , and the result was expanded by exploiting the diagonal component of . The resulting expression matches a Cauchy kernel exactly, and efficient algorithms exist for computing Cauchy kernels. A Cauchy matrix or kernel has the structure
Looking at the part of the final expression—where the original is considered at each element of unit —we see that it can be computed with the cost of a Cauchy matrix-vector multiplication. Ordinarily, computing a length- convolution kernel for hidden states costs . Allowing a small approximation error reduces this to . The proof is mostly unnecessary: computing means obtaining for every in , which is precisely the form of a Cauchy kernel. The condition always holds.
That completes the algorithm. Let us revisit the brief outline from earlier.
Beginning from the fact that is DPLR, we transformed into the Fourier domain to compute the generating function rather than directly. We reparameterized in the resulting expression. We then wrote every discretized matrix in the closed form of the state matrix and expanded it with the Woodbury identity, obtaining
One slight difference between my derivation and the paper’s actual algorithm is that I have treated throughout. Because the structure is the same, I do not believe this creates a serious problem. Applying an inverse Fourier transform to the efficiently computed gives . At last, we have understood every algorithm behind the paper’s core idea.
Deep-layer architecture built with S4
The preceding derivation shows which parameters S4 needs. is initialized as some form of HiPPO matrix, and according to that form it consists of the diagonal and vector components and . Since a diagonal matrix need not store parameters outside its diagonal entries, S4 has a total of trainable parameters for state dimension . S4 itself is a linear mapping—it emits a sequence of the same length—but stacking several S4 layers and adding nonlinearities turns it into a deep layer.
In conclusion
I think it is best to end here rather than separately present the experimental results. I expected them to demonstrate efficiency while retaining long-range effectiveness, and that is exactly what the paper shows. What I found most important was seeing how one can continually identify improvements over the preceding paper that established the modeling foundation for SSM—for example, decomposing a structure to simplify matrix multiplication and proving the result mathematically—and produce a sequence of contributions through that process.