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 ψ:RRNout×Nin\psi:\mathbb{R}\rightarrow\mathbb{R}^{N_{out}\times N_{in}}. Continuous function ψ\psi 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, yy is the desired function and zz is its first derivative. Because directly solving a second-order ODE to obtain the desired function is difficult, define first derivative yy^\prime as temporary variable zz, converting the second-order ODE into a first-order system over zz and yy.

This transformed ODE system is called a Hamiltonian system. Solving it implements a continuous function dependent on the input at each time,

H(y,z,t)=α2y2+12z2+i=1m1wilog(cosh(wiyi+(Vu(t))i+bi)) H(y, z, t) = \frac{\alpha}{2} \parallel y \parallel^2 + \frac{1}{2}\parallel z \parallel^2 + \sum_{i=1}^m\frac{1}{w_i} \log (\cosh (w_iy_i + (Vu(t))_i + b_i))

giving the expression above through the Euclidean norms \parallel\cdot\parallel of vectors y,zy,z. 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 <A,B,C,D><A,B,C,D>.

m˙=Am+Buy=Cm+Du \begin{aligned} \dot{m} =& Am + Bu \newline y =& Cm + Du \end{aligned}

Using u(s),y(s)u(s),y(s), the Laplace transforms of the input and output, we define transfer function G(s)G(s) of a SISO system. The difficulty is that this transfer function is infinite-dimensional and cannot cover every continuous delay θ\theta.

G(s)=y(s)u(s)=eθs G(s) = \frac{y(s)}{u(s)} = e^{-\theta s}

To obtain a finite, causal state-space realization, transfer function G(s)G(s) must be expressed as a polynomial in ss. 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 i,j[0,d1]i,j\in[0,d-1] satisfy the following relations.

For details and derivations, see the earlier LMU paper or the derivation below.

Ai,j=(2i+1)θ{1i<j(1)ij+1ijBi=(2i+1)(1)iθCi=(1)il=0i(il)(i+lj)(1)lD=0 \begin{aligned} A_{i,j} =& \frac{(2i+1)}{\theta}\begin{cases} -1 & i < j \newline (-1)^{i-j+1} & i \ge j \end{cases}\newline B_i =& \frac{(2i+1)(-1)^i}{\theta} \newline C_i =& (-1)^i \sum_{l=0}^i {i \choose l}{i+l \choose j}(-1)^l \newline D =& 0 \end{aligned}

The third matrix, CC, contains the central idea. Expanding CC yields Legendre polynomials. Since D=0D=0, the current state mtm_t can determine the accuracy of shifted input u(tθ)u(t-\theta). A signal shifted in phase by θ\theta^\prime, for example, can be approximated with a shifted Legendre polynomial as follows.

Ci(θ)=(1)il=0i(il)(i+1j)(θθ)l, 0θθu(tθ)C(θ)mt \begin{aligned} C_i(\theta^\prime) = (-1)^i \sum_{l=0}^i {i \choose l}{i+1 \choose j}&\left(-\frac{\theta^\prime}{\theta}\right)^l,~0 \le \theta^\prime \le \theta \newline u(t-\theta^\prime) \approx& C(\theta^\prime)^\top m_t \end{aligned}

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 μ\mu, which measures the difference between a true and predicted function, map continuous function ff on each interval into subspace gg, 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 A(t)A(t) and B(t)B(t), 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 θ\theta.

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 u(t)y(t)u(t)\rightarrow y(t) through implicit function x(t)x(t).

As in LMU, AA controls the evolution of the system's implicit function x(t)x(t), while B,C,DB,C,D perform projection.

x˙(t)=Ax(t)+Bu(t)y(t)=Cx(t)+Du(t) \begin{aligned} \dot{x}(t) =& Ax(t) + Bu(t) \newline y(t) =& Cx(t) + Du(t) \end{aligned}

Setting Δt\Delta t 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 Δt\Delta t and state matrix AA, it can also inherit the desirable properties of RNNs and ODEs. Matrix AA 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 dtdt is chosen in a continuous-time differential equation. Controlling step size or time scale Δt\Delta t 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 x˙(t)=f(t,x(t))\dot{x}(t)=f(t,x(t)) has an equivalent integral equation x(t)=x(t0)+t0tf(s,x(s))dsx(t)=x(t_0)+\int_{t_0}^t f(s,x(s))ds. Its integral solution can be found iteratively by substituting an approximation of xx into f(s,x(s))f(s,x(s)). Given initial condition x0(t)=x(t0)x_0(t)=x(t_0), for example,

xi+1(t)=x0(t)+t0tf(s,xi(t))ds x_{i+1} (t) = x_0 (t) + \int_{t_0}^t f(s, x_{i}(t))ds

we obtain the approximation above. This is called Picard iteration.

Discretization

During discretization, we cannot integrate the function directly, so we must obtain x(ti)x(t_i) at separate discrete times tit_i. If the integral equation has an exact closed form, simple downsampling gives x(t0),x(t1),x(t_0),x(t_1),\cdots. Even without a closed form, Picard iteration can be applied to the integral equation on each interval,

x(tk+1)=x(tk)+tktk+1f(s,x(s))ds x(t_{k+1}) = x(t_k) + \int_{t_k}^{t_{k+1}} f(s, x(s)) ds

to sample function values at each tkt_k. Another method is the generalized bilinear transform (GBT), which applies to the linear ODE of interest here. Given

x˙(t)=Ax(t)+Bu(t)y(t)=Cx(t)+Du(t) \begin{aligned} \dot{x}(t) =& Ax(t) + Bu(t) \newline y(t) =& Cx(t) + Du(t) \end{aligned}

the GBT update is given below, where Δt\Delta t is the step size.

x(t+Δt)=(IαΔtA)1(I+(1α)ΔtA)x(t)+Δt(IαΔtA)1Bu(t) x(t+\Delta t) = (I-\alpha \Delta t \cdot A)^{-1}(I+(1-\alpha)\Delta t \cdot A)x(t) +\Delta t(I-\alpha \Delta t \cdot A)^{-1}B \cdot u(t)

The equation looks complicated, but special cases make it easier to understand. Substituting α=0\alpha=0 gives

x(t+Δt)=x(t)+Δt(Ax(t)+Bu(t))=x(t)+Δtx˙(t) \begin{aligned} x(t+\Delta t) =& x(t) + \Delta t \cdot (Ax(t) + Bu(t)) \newline =& x(t) + \Delta t \cdot \dot{x}(t) \end{aligned}

which is the standard Euler method. Parameter α\alpha determines the location at which the derivative is evaluated. When α=1\alpha=1, 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 t+Δtt+\Delta t rather than tt.

x(t+Δt)=(IΔtA)1x(t)+Δt(IΔtA)1Bx˙(t)x(t+\Delta t) = (I-\Delta t A)^{-1}x(t) + \Delta t (I - \Delta t A)^{-1} B \cdot \dot{x}(t)

Using α=12\alpha=\frac{1}{2} averages derivatives at the two positions. For a complex solution with high curvature, this produces a more stable prediction at the same Δt\Delta t. It is called the bilinear method.

x(t+Δt)=(IΔt/2A)1(I+Δt/2A)x(t)+Δt(IΔt/2A)1Bx˙(t)x(t+\Delta t) = (I-\Delta t / 2A)^{-1}(I+\Delta t / 2A) x(t) + \Delta t (I - \Delta t / 2A)^{-1} B\cdot\dot{x}(t)

Let the matrices used by the bilinear method be Aˉ,Bˉ\bar{A},\bar{B}. Discretizing the system with them yields the following discrete-time state-space model.

xt=Aˉxt1+Bˉutyt=Cxt+Dut\begin{aligned} x_t =& \bar{A}x_{t-1} + \bar{B}u_t \newline y_t =& Cx_t + Du_t \end{aligned}

Timescale factor

Dependency decreases as sequence length grows. If Δt\Delta t is the time interval, dependency is inversely proportional to it. Most ODE-based RNNs fix Δt\Delta t, but the gating mechanism of a classical RNN has an effect similar to learning it. From a CNN perspective, Δt\Delta t controls convolution-kernel size. Interpreting either CNNs or RNNs through ODEs therefore turns the problem into one of optimizing time interval Δt\Delta t.

Continuous-time memory

Given input function u(t)u(t) and fixed probability measure (metric) ω(t)\omega(t), suppose there are NN basis functions. At each time tt, all previous inputs u(τ)τ<tu(\tau)\vert_{\tau<t} can be expressed as a combination of these bases, yielding coefficient vector x(t)RNx(t)\in\mathbb{R}^N by projecting the function. The optimal solution at each time depends on distance metric ω(t)\omega(t). Representing u(t)u(t) by coefficients x(t)x(t) 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 ω=I{[0,1]}\omega=\mathbb{I}\{[0,1]\} assigns equal importance to every time step, while exponentially decaying measure ω(t)=exp(t)\omega(t)=\exp(-t) 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 AA has a closed-form solution for a chosen measure, enabling long-range dependency modeling.

The HiPPO appendix derives matrix AA 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 μ\mu 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 iith polynomial must have degree ii.

Pi,Pjμ=Pi(x)Pj(x)dμ(x)=0  (ij), deg(Pi)=i\langle P_i, P_j \rangle_\mu = \int P_i(x) P_j(x) d\mu (x) = 0~~(i \neq j),~\deg (P_i) = i

Under these conditions, the optimal approximation to target function ff is

i=0N1ciPi(x)/Piμ2, where ci=f,Piμ=f(x)Pi(x)dμ(x)\sum_{i=0}^{N-1} c_i P_i(x) / \parallel P_i \parallel_\mu^2,~\text{where }c_i = \langle f,P_i \rangle_\mu = \int f(x)P_i (x) d\mu(x)

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 L2L^2 inner product over [1,1][-1,1] gives a Kronecker delta scaled by 22n+1\frac{2}{2n+1}. Another well-known property is the boundary condition Pn(1)=1,Pn(1)=(1)nP_n(1)=1, P_n(-1)=(-1)^n.

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 [1,1][-1,1]. Rescaling the function domain so it holds on [0,t][0,t] gives

(2n+1)0tPn(2xt1)Pm(2xt1)1tdx=2n+12PnPmωlegdx\begin{aligned} (2n+1)\int_0^t P_n \left( \frac{2x}{t} - 1 \right) P_m \left( \frac{2x}{t}-1 \right) \frac{1}{t} dx = \frac{2n+1}{2}\int P_n P_m \omega_\text{leg} dx \end{aligned}

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.

(2n+1)1/2Pn(2xt1)(2n+1)^{1/2} P_n \left(\frac{2x}{t} - 1\right)

Translated Legendre

Translated Legendre is the Legendre measure with window size θ\theta at current time tt.

ω(t,x)=1θI[tθ,t]pn(t,x)=(2n+1)1/2Pn(2(xt)θ+1)gn(t,x)=λnpn(t,x)\begin{aligned} \omega(t, x) =& \frac{1}{\theta} \mathbb{I}_{[t-\theta, t]} \newline p_n(t, x) =& (2n+1)^{1/2}P_n\left(\frac{2(x-t)}{\theta} + 1\right) \newline g_n(t, x) =& \lambda_n p_n (t, x) \end{aligned}

The concept of tilting introduces a function χ\chi, used when we do not want to rely solely on orthogonal polynomials. If combined function pn(x)χ(x)p_n(x)\chi(x) replaces pn(t,x)p_n(t,x), it becomes orthogonal under ω/χ2\omega/\chi^2 at each time step, since the product contains two orthogonal polynomials and χ2\chi^2. To obtain a normalized measure and orthonormal basis, define

ζ(t)=ωχ2=ω(t)(x)(χ(t)(x))2dx\zeta(t) = \int \frac{\omega}{\chi^2} = \int \frac{\omega^{(t)}(x)}{(\chi^{(t)}(x))^2}dx

which becomes the normalization constant. Normalized measure ν(t)\nu^{(t)} therefore has density ω(t)(x)ζ(t)(χ(t)(x))2\frac{\omega^{(t)}(x)}{\zeta(t)\cdot(\chi^{(t)}(x))^2}. This construction makes the tilted orthogonal polynomials orthonormal. Checking orthogonality with the equation above shows that Legendre orthogonality is normalized under the original measure ω\omega. 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

gn(t,t)=λn(2n+1)1/2gn(t,tθ)=λn(1)n(2n+1)1/2\begin{aligned} g_n(t, t) =& \lambda_n (2n+1)^{1/2} \newline g_n(t,t-\theta) =& \lambda_n (-1)^n (2n+1)^{1/2} \end{aligned}

as shown above.

Projection and Coefficients

This may feel like an enormous detour merely to derive AA, but the equations are essential to understanding HiPPO fully. Using the tilted measure derived above, projection onto the measure yields the coefficients

cn(t)=ζ(t)1/2λnfpn(t)ω(t)χ(t)c_n(t) = \zeta(t)^{-1/2} \lambda_n \int fp_n^{(t)} \frac{\omega^{(t)}}{\chi^{(t)}}

To build an end-to-end model from this expression and make the network represent both past values of ff and the current function during online prediction, c(t)c(t) 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.

ddtcn(t)=ζ(t)1/2λnf(x)(tpn(t,x))ωχ(t,x)dx+f(x)(ζ1/2λnpn(t,x))(tωχ(t,x))dx\begin{aligned} \frac{d}{dt} c_n(t) &= \zeta(t)^{-1/2} \lambda_n \int f(x) \left(\frac{\partial}{\partial t}p_n (t, x) \right) \frac{\omega}{\chi} (t, x) dx \newline &+\int f(x) \left( \zeta^{-1/2}\lambda_n p_n(t, x) \right)\left(\frac{\partial}{\partial t} \frac{\omega}{\chi} (t, x)\right) dx \end{aligned}

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 nnth Legendre polynomial is a linear combination of Legendre polynomials through degree n1n-1. What a remarkable world of Legendre polynomials.

Fortunately, the derivative of λnpn(t,x)\lambda_n p_n(t,x) can therefore be expressed simply in terms of the gg functions.

tgn(t,x)=λn(2n+1)1/22θ(λn11(2n1)1/2gn1+λn31(2n1)1/2gn3+)\frac{\partial}{\partial t} g_n (t, x) = -\lambda_n (2n+1)^{1/2} \frac{2}{\theta} \left( \lambda_{n-1}^{-1} (2n-1)^{1/2}g_{n-1} + \lambda_{n-3}^{-1} (2n-1)^{1/2} g_{n-3} + \cdots \right)

The measure's partial derivative is the derivative of a rectangular function.

tω(t,x)=1θδt1θδtθ\frac{\partial}{\partial t} \omega (t, x) = \frac{1}{\theta}\delta_t - \frac{1}{\theta} \delta_{t-\theta}

With all components ready, we can substitute them into the differential equation for coefficient dynamics derived above.

ddtcn(t)=λnθ(2n+1)1/2k=0N1Mnk(2k+1)1/2ck(t)λk+(2n+1)1/2λnθf(t)\frac{d}{dt}c_n(t) = -\frac{\lambda_n}{\theta} (2n+1)^{1/2} \sum_{k=0}^{N-1} M_{nk} (2k+1)^{1/2} \frac{c_k(t)}{\lambda_k} + (2n+1)^{1/2} \frac{\lambda_n}{\theta} f(t)

where MnkM_{nk} equals 11 when knk\le n and (1)nk(-1)^{n-k} when k>nk>n. Choosing λn=(2n+1)1/2(1)n\lambda_n=(2n+1)^{1/2}(-1)^n gives

ddtc(t)=1θAc(t)+1θBf(t)\frac{d}{dt} c(t) = -\frac{1}{\theta} Ac(t) + \frac{1}{\theta} B f(t)

with

Ank=(2n+1){(1)nkif k<n1if kn,  Bn=(2n+1)(1)nA_{nk} = (2n+1)\begin{cases} (-1)^{n-k}& \text{if }k < n \newline 1 & \text{if }k \ge n \end{cases},~~B_n = (2n+1)(-1)^n

which recovers exactly the LMU introduced earlier.

Interpreting LSSL

Returning to LSSL, assume fixed state-space representation A,B,C,DA,B,C,D. At its simplest, LSSL maps an input sequence to an output sequence. It is defined by parameter matrices A,B,C,DA,B,C,D and the Δt\Delta t 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 xt1x_{t-1}, which contains the input context preceding each time step. The discretized LSSL computes current state xtx_t and output yty_t.

xt=Aˉxt1+Bˉutyt=Cxt+Dut\begin{aligned} x_t =& \bar{A}x_{t-1} + \bar{B}u_t \newline y_t =& Cx_t + Du_t \end{aligned}

It therefore behaves like an RNN and even supports gated recurrence. For example, one-dimensional gated recurrence (1σ(z))xt1+σ(z)ut(1-\sigma(z))x_{t-1}+\sigma(z)u_t is equivalent to discretizing x˙(t)=x(t)+u(t)\dot{x}(t)=-x(t)+u(t) with the backward Euler method. Here zz may be any expression. From the sigmoid's properties and the GBT introduced earlier, setting Δt=exp(z)\Delta t=\exp(z) shows that this gate is backward Euler with A=1,B=1A=-1,B=1. 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 x˙(t)=x(t)+u(t)\dot{x}(t)=-x(t)+u(t), 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 00. The linear state-system output becomes

yk=C(Aˉ)kBˉu0+C(Aˉ)k1Bˉu1++CAˉBˉuk1+Bˉuk+Duky_k = C(\bar{A})^k\bar{B}u_0 + C(\bar{A})^{k-1}\bar{B}u_1 + \cdots + C\bar{A} \bar{B}u_{k-1} + \bar{B}u_k + Du_k

which can be written as a discrete-time convolution.

y=K_L(Aˉ,Bˉ,C)u+DuK_L(Aˉ,Bˉ,C)=(CAiB)_i[L]RL\begin{aligned} &y = \mathcal{K}\_L (\bar{A}, \bar{B}, C) \ast u + Du \newline &\mathcal{K}\_L (\bar{A}, \bar{B}, C) = (CA^iB)\_{i \in [L]} \in \mathbb{R}^L \end{aligned}

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 yy is the convolution of input uu with the system's impulse-response function hh.

y(t)=h(τ)u(tτ)dτy(t) = \int h(\tau)u(t-\tau) d\tau

Conversely, if a convolutional filter is a rational function of degree NN, it can be represented by a state-space model of size NN. Prior work shows that any convolutional filter hh can be expressed as a finite-degree rational function. Consider the HiPPO matrix above. In Translated Legendre, AA was defined under a measure with a uniform probability distribution over interval θ\theta. Fixing dtdt in an ordinary LSSL, the first equation

x˙(t)=Ax(t)+Bu(t)\dot{x}(t) = Ax(t) + Bu(t)

stores historical elements, while the second equation

y(t)=Cx(t)+Du(t)y(t) = Cx(t) + Du(t)

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 LL, it is a parameterized vector-to-vector function RLRL\mathbb{R}^L\rightarrow\mathbb{R}^L. If the LSSL is ψ\psi,

ψ(A,B,C,D,Δt), ARN×N, BRN×1, CR1×N, DR1×1\psi(\cdot \vert A, B, C, D, \Delta t),~A \in \mathbb{R}^{N \times N},~B \in \mathbb{R}^{N \times 1},~C \in \mathbb{R}^{1 \times N},~D \in \mathbb{R}^{1 \times 1}

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 HH, as in a Transformer input (L×HL\times H), the model learns HH 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 Aˉ\bar{A}.

xt=Aˉxt1+Bˉutx_t = \bar{A}x_{t-1} + \bar{B}u_t

Training with gradient descent therefore cannot avoid vanishing gradients. Randomly initializing and learning AA will not produce the desired performance.

Frameworks such as HiPPO instead describe how to remember past functions according to measure ω\omega, 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 AA only for a few measures, including Legendre, and does not establish that a different AA can exist for every general measure.

LSSL extends this result to an arbitrary measure ω\omega and proves that a differential equation can be found for low-recurrence-width AA.

Efficient Algorithms for LSSLs

Although AA and Δt\Delta t 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,

x(t+Δt)=(IαΔtA)1(I+(1α)ΔtA)x(t)+Δt(IαΔtA)1Bu(t)x(t+\Delta t) = (I-\alpha \Delta t \cdot A)^{-1}(I+(1-\alpha)\Delta t \cdot A)x(t) +\Delta t(I-\alpha \Delta t \cdot A)^{-1}B \cdot u(t)

while the latter must multiply Aˉ\bar{A} as many as LL times, where LL is the feature length.

KL(Aˉ,Bˉ,C)=(CB,CAB,,CAL1B)\mathcal{K}_L (\bar{A}, \bar{B}, C) = (CB, CAB, \ldots, CA^{L-1}B)

This introduces another condition for learning AA efficiently. Every existing fixed-LSSL matrix AA is proven to be 3-quasiseparable. If learned AA also preserves quasiseparability, MVM and Krylov-function computation require fewer operations.

Evaluations and Demonstrations

The experiments confirm that learning AA under these conditions outperforms the earlier fixed HiPPO-system AA.

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...