Lag Operator SSMs: A Geometric Framework for Structured State Space Modeling

TL;DR

This work presents a direct discrete-time framework for Structured State Space Models built around a lag operator that geometrically tracks how a warped projection basis evolves between timesteps. By projecting the history onto time-warped orthonormal bases and updating via a backward lag, the authors derive the entire discrete recurrence from a single inner product, bypassing the traditional continuous-time ODE–discretization pipeline. They demonstrate that an exponential warp recovers the HiPPO-LegS system, providing a principled geometric foundation for HiPPO and enabling modular memory design through the warp function. Numerical experiments show exact HiPPO equivalence in matrix form and faithful replication of memory dynamics, validating the framework as a flexible, interpretable building block for long-range sequence modeling and potential multi-resolution memory schemes.

Abstract

Structured State Space Models (SSMs), which are at the heart of the recently popular Mamba architecture, are powerful tools for sequence modeling. However, their theoretical foundation relies on a complex, multi-stage process of continuous-time modeling and subsequent discretization, which can obscure intuition. We introduce a direct, first-principles framework for constructing discrete-time SSMs that is both flexible and modular. Our approach is based on a novel lag operator, which geometrically derives the discrete-time recurrence by measuring how the system's basis functions "slide" and change from one timestep to the next. The resulting state matrices are computed via a single inner product involving this operator, offering a modular design space for creating novel SSMs by flexibly combining different basis functions and time-warping schemes. To validate our approach, we demonstrate that a specific instance exactly recovers the recurrence of the influential HiPPO model. Numerical simulations confirm our derivation, providing new theoretical tools for designing flexible and robust sequence models.

Figures (4)

• Figure 1: Illustration of a time warping function $\sigma_t(s)$. The infinite time interval $T_t = (-\infty, t]$ (s-axis) is mapped to the finite canonical interval $Z = (0, 1]$ (z-axis), concentrating recent history (near $t$) and compressing older history (towards $-\infty$).
• Figure 2: Conceptualization of the extended signal $\tilde{u}_{t{+}1}(s)$. The new input $u_{t{+}1}$ is mixed over the interval boundary to the prior signal approximation $\hat{u}_t(s)$.
• Figure 3: Example basis functions lag-shifted forward and backward using the lag operator $\bm{A}_\Delta$. The backward shift compresses toward old history, while the forward shift expands with amplitude shrinkage due to stabilization.
• Figure 4: Visualization of the final compressed memory state $\bm{c}_T$ (reconstructed via eq. \ref{['eq:recon-u-w-projection']}) for our Lag Operator model and the HiPPO-LegS baseline. The near-perfect alignment demonstrates that both models accurately capture recent history while forgetting the distant past, a behavior governed by the plotted exponential measure ($\omega_t(s)$), which dictates the memory weighting.

Theorems & Definitions (5)

• Proposition 1: Continuous Generator $\bm{A}_{\text{gen}}$
• Proposition 2: Continuous Generator for $\bm{B}_{\text{gen}}$
• Proposition 3: Proposition \ref{['prop:a-gen']} Restated: Continuous Generator $\bm{A}_{\text{gen}}$
• Proposition 4: Proposition \ref{['prop:b-gen']} Restated: Continuous Input Generator
• proof