ai papers
Foundations of Mamba Modeling (1): Linear State-Space Layers (LSSL)
Junyoung Park · 2024-02-01 · 20 min
The Evolution of DNNs for Sequential Data
Deep-learning models have undergone countless changes to process sequential data. Today, the Transformer architecture dominates LLM and multimodal research. Before Transformers, recurrent neural networks (RNNs) were developed with long-term modules such as LSTM and GRU to capture relationships across distant context, while CNNs (Convolutional Neural Networks), the most basic DNN architecture, were adapted to temporal datasets—for example, 3D convolutions add a time axis to video data so temporal information can be used.
Another, less celebrated line of neural-network development directly models neural differential equations (NDEs) to process more complex continuous data.
Each network had clear strengths and weaknesses. Despite advances in long-term modules, RNN computation and latency still grow in proportion to context length. Reliance on those modules also creates a fundamental difficulty in learning context from complex data.
The vanishing-gradient and exploding-gradient problems are also major causes of catastrophic forgetting in continual learning and in training sequential models such as RNNs.
CNNs optimize local information—features adjacent in their dimensions—quickly and provide useful inductive biases for object- or semantic-centric data such as images and videos with reasonably clear context. But their computation is not inherently structured for sequences, so longer inputs create problems similar to RNNs. Convolution ultimately extracts only local information within a fixed context, leaving computation and time proportional to context length.
NDE modeling provides a mathematical theory for a particular modality or fixed problem, but it is not very efficient. Diffusion modeling is a representative example: diffusion generative models address efficiency by discretizing the score function. This reduces the burden on the implicit model and enables a simple U-Net; other implicit mappings such as consistency modeling can help as well. But these methods work because the differential-equation solution can be obtained numerically on each interval. There is no single architecture that lets neural networks learn every kind of differential equation efficiently in a uniform way.
The ideal direction for model architecture can therefore be summarized as follows:
- It should support parallel computation like convolution for efficiency.
- It should process context through recurrent state inference.
- Like a differential equation, it should apply over a time scale rather than only to discretized signals.
There have been continual attempts to find such a model.
An Introduction to Several Modeling Approaches
CKConv
CKConv (Continuous Kernel Convolution) treats a convolution kernel as a continuous vector-valued function . Continuous function is parameterized and learned by a small MLP. The MLP receives a time step as a scalar value and outputs the convolution kernel at that position as a vector.
UnICORNN
ODE-based RNN modeling that seeks long-range dependencies through time scaling includes work such as UnICORNN. In brief, this RNN uses the Euler method to discretize a second-order ordinary differential equation along the time axis.
In the figure, is the desired function and is its first derivative. Because directly solving a second-order ODE to obtain the desired function is difficult, define first derivative as temporary variable , converting the second-order ODE into a first-order system over and .
This transformed ODE system is called a Hamiltonian system. Solving it implements a continuous function dependent on the input at each time,
giving the expression above through the Euclidean norms of vectors . Discretizing this continuous-signal differential equation with the Euler method produces the desired discrete dynamical system.
LMU
Parallelizing Legendre Memory Unit Training (LMU) attempts to solve an RNN's lack of parallelism with linear recurrent convolution. If we can build a delay system that retrieves states before and after a given input, its output provides the input's recurrent structure. We seek a linear system—natural because neural-network computation is based on tensors and matrices—so the goal is to find an LTI system represented by four matrices .
Using , the Laplace transforms of the input and output, we define transfer function of a SISO system. The difficulty is that this transfer function is infinite-dimensional and cannot cover every continuous delay .
To obtain a finite, causal state-space realization, transfer function must be expressed as a polynomial in . A transfer function is generally a ratio of polynomials with fixed degrees; in a proper-dimensional system, the denominator must have the higher degree to guarantee convergence. We therefore need to approximate the expression above. Implementing it as a linear system requires the matrices introduced earlier, whose elements for satisfy the following relations.
For details and derivations, see the earlier LMU paper or the derivation below.
The third matrix, , contains the central idea. Expanding yields Legendre polynomials. Since , the current state can determine the accuracy of shifted input . A signal shifted in phase by , for example, can be approximated with a shifted Legendre polynomial as follows.
To restate the long explanation: conventional linear state-machine design represents an ideal delay as an LTI system, and LMU learns that system with a nonlinear neural network. It trains the network using the delay system as its solution.
HiPPO
HiPPO is a generalization of LMU.
The HiPPO method works as follows:
In Hilbert space , which measures the difference between a true and predicted function, map continuous function on each interval into subspace , then represent it as an array of coefficients in a suitable vector basis. The continuous-time ODE can then be expressed as the differential equation of an LTI system. Functions and , which govern the system, map importance over sequence memory. This generalizes LMU to continuous-time memorization. Conventional LMU, the Legendre Memory Unit, instead directly discretizes the LTI differential equation of an ideal delay system with fixed sliding-window size .
LSSL Modeling
These earlier modules all approach the architecture and limitations of CNNs and RNNs at the time-step level. Their limitation is that none fundamentally solves the problems of convolutional and recurrent models.
The Linear State-Space Layer (LSSL) aims to design an architecture combining the advantages shown above. Its formulation is not radically different from earlier approaches: LSSL maps a one-dimensional function or sequence through implicit function .
As in LMU, controls the evolution of the system's implicit function , while perform projection.
Setting as a discrete step size lets us interpret LSSL as a recurrent model that evolves state over time with a fixed amount of memory and computation. The two LTI-system equations above can be expressed as continuous convolution, so their discrete-time version can also be parallelized as a convolution, accelerating training. Finally, because an LTI system is itself modeled by a differential equation, LSSL can reproduce every applicable continuous-time-model setting.
The paper's claim is that, as classical control theory already establishes, LSSL can represent any one-dimensional convolution. With an appropriate step size and state matrix , it can also inherit the desirable properties of RNNs and ODEs. Matrix is the learned driver of system evolution and, as earlier work such as HiPPO showed, must capture both continuous-time memory and long-range dependencies.
Continuous-time memorization
Approximating continuous-time memorization is shared theoretical and technical background for HiPPO and LSSL.
Because continuous-time modeling cannot be applied directly, it must be approximated or downsampled into a discrete-time model.
As diffusion modeling demonstrates, generation quality depends on how is chosen in a continuous-time differential equation. Controlling step size or time scale is therefore crucial when discretizing a continuous-time model.
This section summarizes concepts that provide the basis for insights into LSSL's properties.
Approximations of differential equations
Every differential equation has an equivalent integral equation . Its integral solution can be found iteratively by substituting an approximation of into . Given initial condition , for example,
we obtain the approximation above. This is called Picard iteration.
Discretization
During discretization, we cannot integrate the function directly, so we must obtain at separate discrete times . If the integral equation has an exact closed form, simple downsampling gives . Even without a closed form, Picard iteration can be applied to the integral equation on each interval,
to sample function values at each . Another method is the generalized bilinear transform (GBT), which applies to the linear ODE of interest here. Given
the GBT update is given below, where is the step size.
The equation looks complicated, but special cases make it easier to understand. Substituting gives
which is the standard Euler method. Parameter determines the location at which the derivative is evaluated. When , this becomes the backward Euler method. It is still a first-order approximation based on the derivative at a particular point, but that point is rather than .
Using averages derivatives at the two positions. For a complex solution with high curvature, this produces a more stable prediction at the same . It is called the bilinear method.
Let the matrices used by the bilinear method be . Discretizing the system with them yields the following discrete-time state-space model.
Timescale factor
Dependency decreases as sequence length grows. If is the time interval, dependency is inversely proportional to it. Most ODE-based RNNs fix , but the gating mechanism of a classical RNN has an effect similar to learning it. From a CNN perspective, controls convolution-kernel size. Interpreting either CNNs or RNNs through ODEs therefore turns the problem into one of optimizing time interval .
Continuous-time memory
Given input function and fixed probability measure (metric) , suppose there are basis functions. At each time , all previous inputs can be expressed as a combination of these bases, yielding coefficient vector by projecting the function. The optimal solution at each time depends on distance metric . Representing by coefficients in this way is the HiPPO (High-Order Polynomial Projection Operator) introduced earlier.
HiPPO proposes two measures, called LegT and LagT in the paper: uniform measure assigns equal importance to every time step, while exponentially decaying measure assigns greater importance to recent steps. As an aside, HiPPO uses scaled Legendre (LegS), rather than translated Legendre (LegT) with a fixed sliding-window size, to address long-range dependencies and forgetting. Both weight the contents of the window uniformly, but the LegS window grows over time. The important point is that matrix has a closed-form solution for a chosen measure, enabling long-range dependency modeling.
The HiPPO appendix derives matrix for each metric. I will give a simplified summary. The prerequisites are too extensive for a complete proof here, but understanding HiPPO as a whole is necessary to interpret LSSL.
Orthogonal Polynomials
Orthogonal polynomials—polynomials perpendicular to one another—are a basic tool for analyzing functions. Every measure induces a unique sequence of functions corresponding to its orthogonal polynomials. Here, the measure can be understood as the subspace over which integration occurs. Distinct orthogonal polynomials integrate to zero under the measure, and the th polynomial must have degree .
Under these conditions, the optimal approximation to target function is
The Fourier-series basis is perhaps the best-known example; Jacobi, Laguerre, and Hermite polynomials are others. The orthogonal polynomial considered here is the Legendre polynomial, a member of the Jacobi family.
Legendre Polynomials
Legendre polynomials are common in spherical coordinates; they bring back nightmares from engineering mathematics. Their familiar orthogonal relationship is that the inner product over gives a Kronecker delta scaled by . Another well-known property is the boundary condition .
Linearity lets us derive corresponding polynomials on different time scales. This is possible because the measure space in which Legendre polynomials are defined is itself a uniform probability distribution.
The original orthogonality holds on . Rescaling the function domain so it holds on gives
Simply matching the integration interval recovers the Kronecker delta. Thus, when the measure is rescaled, the corresponding Legendre polynomial follows by rescaling the original polynomial.
Translated Legendre
Translated Legendre is the Legendre measure with window size at current time .
The concept of tilting introduces a function , used when we do not want to rely solely on orthogonal polynomials. If combined function replaces , it becomes orthogonal under at each time step, since the product contains two orthogonal polynomials and . To obtain a normalized measure and orthonormal basis, define
which becomes the normalization constant. Normalized measure therefore has density . This construction makes the tilted orthogonal polynomials orthonormal. Checking orthogonality with the equation above shows that Legendre orthogonality is normalized under the original measure . This special formulation is not used for the Legendre projection itself, so the ordinary equation is enough for our purposes. Using the Legendre properties mentioned earlier, shifted and scaled Legendre functions have boundaries
as shown above.
Projection and Coefficients
This may feel like an enormous detour merely to derive , but the equations are essential to understanding HiPPO fully. Using the tilted measure derived above, projection onto the measure yields the coefficients
To build an end-to-end model from this expression and make the network represent both past values of and the current function during online prediction, must be a vector—which is exactly our objective of obtaining a coefficient vector.
The coefficients must always update from the current prediction. They are not fixed, but continuously changing functions, suggesting the following differential equation.
Coefficient dynamics with Translated Legendre
We said tilting can be ignored when projecting Legendre polynomials. Solving the expression above then requires partial derivatives of both the orthogonal polynomials and the measure. Omitting the detailed derivation, the derivative of the th Legendre polynomial is a linear combination of Legendre polynomials through degree . What a remarkable world of Legendre polynomials.
Fortunately, the derivative of can therefore be expressed simply in terms of the functions.
The measure's partial derivative is the derivative of a rectangular function.
With all components ready, we can substitute them into the differential equation for coefficient dynamics derived above.
where equals when and when . Choosing gives
with
which recovers exactly the LMU introduced earlier.
Interpreting LSSL
Returning to LSSL, assume fixed state-space representation . At its simplest, LSSL maps an input sequence to an output sequence. It is defined by parameter matrices and the required for discretization. Let us now see how it can possess the properties of RNNs, CNNs, and neural ODEs.
LSSL to RNN
The recurrent state in LSSL is , which contains the input context preceding each time step. The discretized LSSL computes current state and output .
It therefore behaves like an RNN and even supports gated recurrence. For example, one-dimensional gated recurrence is equivalent to discretizing with the backward Euler method. Here may be any expression. From the sigmoid's properties and the GBT introduced earlier, setting shows that this gate is backward Euler with . One might still ask whether a state layer built from a linear system can express the nonlinearity and complexity of an ordinary deep RNN.
We described this as discretizing , but can reinterpret it through Picard iteration: a deep RNN searches for a function by performing Picard iteration during training. If LSSL uses nonlinear rather than linear recurrence, it can likewise learn nonlinearity. In this sense, RNNs and LSSL have an equivalence relationship.
LSSL to CNN
For simplicity, assume the initial state is . The linear state-system output becomes
which can be written as a discrete-time convolution.
LSSL can therefore be interpreted as a model whose output is computed by convolution, which can be accelerated with an FFT.
In an ordinary continuous state-space system, output is the convolution of input with the system's impulse-response function .
Conversely, if a convolutional filter is a rational function of degree , it can be represented by a state-space model of size . Prior work shows that any convolutional filter can be expressed as a finite-degree rational function. Consider the HiPPO matrix above. In Translated Legendre, was defined under a measure with a uniform probability distribution over interval . Fixing in an ordinary LSSL, the first equation
stores historical elements, while the second equation
extracts meaningful features within the window. LSSL can therefore be viewed as equivalent to learning a convolutional kernel with learnable width.
Deep Linear State-System Layers
In summary, an ordinary LSSL maps an input sequence to an output sequence. For a signal of length , it is a parameterized vector-to-vector function . If the LSSL is ,
it can be written as above. Since one LSSL has both recurrent and convolutional properties, it can serve like an RNN's recurrent unit or a CNN's convolution kernel.
If the input sequence has hidden dimension , as in a Transformer input (), the model learns independent LSSLs and can reproduce the Transformer's multi-head effect.
Stacking LSSLs can approximate varied functions like a conventional DNN and can be combined with techniques such as normalization and residual connections.
LSSL and Continuous-Time Memorization
Compatibility with conventional DNN modeling is not enough by itself; LSSL must also demonstrate an advantage.
Long dependency into LSSLs
In a discretized linear-system ODE, the system evolves through repeated multiplication by discrete parameter .
Training with gradient descent therefore cannot avoid vanishing gradients. Randomly initializing and learning will not produce the desired performance.
Frameworks such as HiPPO instead describe how to remember past functions according to measure , through projection and coefficients. HiPPO's problem is that its manually defined matrix cannot be learned and must remain fixed. It finds structured solution matrices only for a few measures, including Legendre, and does not establish that a different can exist for every general measure.
LSSL extends this result to an arbitrary measure and proves that a differential equation can be found for low-recurrence-width .
Efficient Algorithms for LSSLs
Although and are crucial parameters, a naive LSSL struggles to learn them. Computing MVM (Matrix–Vector Multiplication) and the Krylov function, corresponding respectively to convolution and recurrence, requires matrix inversion in the former case,
while the latter must multiply as many as times, where is the feature length.
This introduces another condition for learning efficiently. Every existing fixed-LSSL matrix is proven to be 3-quasiseparable. If learned also preserves quasiseparability, MVM and Krylov-function computation require fewer operations.
Evaluations and Demonstrations
The experiments confirm that learning under these conditions outperforms the earlier fixed HiPPO-system .
Classification of long audio signals also demonstrates long-range dependency modeling.
The model reaches performance comparable to the previous state of the art in far fewer training epochs.
Conclusion
We examined LSSL, the foundation of Mamba modeling, and covered HiPPO as well because it is essential to understanding LSSL. I do not yet know how many posts this Mamba review will take, but state modeling does not seem to have been organized clearly elsewhere...