🤖 machine learning

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

This paper introduces a novel geometric framework for Structured State Space Models (SSMs) that utilizes a lag operator to directly derive discrete-time recurrences through domain expansion, offering a flexible, modular alternative to traditional continuous-to-discrete modeling that unifies basis functions and time-warping schemes while recovering established models like HiPPO.

Original authors: Sutashu Tomonaga, Kenji Doya, Noboru Murata

Published 2026-07-16

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Sutashu Tomonaga, Kenji Doya, Noboru Murata

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to remember a long story, but your brain has a strict rule: you can only hold a tiny, fixed-size note card in your hand at any one time. As new sentences arrive, you have to decide what to write on the card and what to erase. If you erase too much, you forget the plot; if you write too much, you run out of space. This is the daily struggle of artificial intelligence when it tries to understand long sequences of data, like a novel, a song, or a stock market history.

For a long time, the best way to solve this "memory card" problem was to use a method called Structured State Space Models (SSMs). Think of these models as a super-smart librarian who doesn't just read the story but projects the entire history of the book onto a special, shifting wall. As the story gets longer, the wall stretches, and the librarian uses a complex set of mathematical rules to figure out how to squish the old pages into the corner while making room for the new ones. The most famous version of this, called HiPPO, works incredibly well, but figuring out how it works is like trying to understand a magic trick by watching a magician perform it in slow motion while wearing heavy gloves. The math behind it involves continuous time, differential equations, and a multi-step process that can feel like a tangled knot of wires. It works, but nobody is entirely sure why the magic happens until it's too late.

This is where a new paper by Sutashu Tomonaga, Kenji Doya, and Noboru Murata steps in. They decided to untangle that knot by looking at the problem from a completely different angle. Instead of starting with the complex, continuous-time magic and trying to force it into a digital computer, they started with the digital computer itself and asked, "What if we just look at how the memory wall stretches from one second to the next?"

The authors introduce a new tool called the "Lag Operator." To understand this, imagine you are watching a movie on a screen that keeps getting wider. Every time the movie moves forward one frame, the screen stretches a little bit. The Lag Operator is like a ruler that measures exactly how much the screen stretched and how the characters on it shifted to fit the new size. By using this ruler, the authors can calculate the new memory state directly, without needing the complicated "magic" steps of the old method. They found that if you use a specific way of stretching the screen (an exponential stretch), your new, simple ruler gives you the exact same result as the famous, complex HiPPO librarian.

In their simulations, the team proved that their new, simpler method is mathematically identical to the old, complicated one. When they fed a chaotic, unpredictable signal (from a system called Lorenz63) into their model, the memory it built was so close to the original HiPPO model that the difference was almost invisible—only about 0.000000576 in error. This suggests that their "Lag Operator" isn't just a clever trick; it's a fundamental way to understand how these memory systems work.

The best part is that this new framework is like a set of Lego bricks. Because the math is so clean and direct, you can now swap out different types of "stretching rules" (time-warping) and different types of "projection walls" (basis functions) to create custom memory systems. For example, you could build a model that remembers the distant past very slowly while keeping a sharp focus on what happened just a second ago, all within a single, unified system. The authors show that this approach opens the door to designing smarter, more flexible AI that can handle long stories without getting confused, all while making the math behind the curtain much easier to see and understand.

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

Problem Statement

Structured State Space Models (SSMs), such as those underpinning the Mamba architecture, have emerged as powerful alternatives to Transformers for long-range sequence modeling, offering linear scaling and effective dynamical memory compression. However, the theoretical foundation of existing SSMs, particularly the High-order Polynomial Projection Operators (HiPPO) framework, relies on a complex, multi-stage derivation. This process involves:

Approximating signal history using continuous-time Ordinary Differential Equations (ODEs) with time-varying measures.
Deriving state matrices from these continuous dynamics.
Discretizing the resulting ODEs to obtain a practical recurrence.

This intricate pipeline obscures the geometric intuition behind state updates, making principled extensions and the design of novel SSMs difficult. The authors argue that the direct geometric connection between signal history and the final state update is hidden by the intermediate steps of continuous-time modeling and subsequent discretization.

Methodology

The paper introduces a Lag Operator Framework that constructs discrete-time SSMs directly from first principles, bypassing the intermediate continuous-time ODE stage. The core of this approach is a novel lag operator ( $\sigma_t \circ \sigma^{-1}_{t+1}$ ) that geometrically derives the discrete-time recurrence by measuring how the system's basis functions undergo "domain expansion" from one timestep to the next.

Key Components:

Time Warping and Basis Construction:
- The framework defines a time-warping function $\sigma_t$ that maps the infinite history interval $T_t = (-\infty, t]$ to a canonical interval $Z$ (e.g., $(0, 1]$ ).
- An orthonormal basis $\Phi$ is defined on $Z$ . The basis on the history interval, $\Psi_t$ , is constructed by warping $\Phi$ via $\sigma_t$ .
- An induced inner product with a specific measure density $\omega_t(s) = |\sigma'_t(s)|$ ensures the warped basis remains orthonormal.
Signal Approximation and Online Updating:
- The signal history is approximated by projecting it onto the time-warped basis $\Psi_t$ , yielding coefficients $c_t$ .
- Upon receiving a new input $u_{t+1}$ , the history interval expands to $T_{t+1}$ . The framework defines an interim extended signal and projects it onto the new basis $\Psi_{t+1}$ .
The Lag Operator Derivation:
- The state transition matrix $A_t$ and input matrix $B_t$ are derived directly via a single inner product involving the lag operator.
- Specifically, $A_t$ is defined as the inner product of the new basis function $\psi_{t+1,n}$ and the old basis function $\psi_{t,m}$ , which simplifies to $\langle \phi_n, \phi_m \circ (\sigma_t \circ \sigma^{-1}_{t+1}) \rangle$ .
- This formulation interprets the state transition as a "re-warping" of projections, where the lag operator $\sigma_t \circ \sigma^{-1}_{t+1}$ accounts for the domain expansion.
Modularity:
- The framework decouples the choice of the canonical basis $\Phi$ from the time-warping function $\sigma_t$ . This allows for the flexible construction of SSMs with diverse memory properties (e.g., multi-resolution memory) by simply combining different warping schemes and basis functions.

Key Contributions

Direct Discrete-Time Construction: The paper provides a first-principles framework for constructing discrete-time SSMs that entirely bypasses the intermediate continuous-time ODE and discretization steps required by traditional HiPPO derivations.
Geometric Interpretation: It introduces the lag operator as a geometric mapping that quantifies domain expansion, offering a clear intuition for state transitions as a "re-warping" of the signal's compressed history.
Theoretical Recovery of HiPPO: The authors prove that a specific instance of their framework (using exponential time warping and Legendre polynomials) exactly recovers the recurrence relations of the influential HiPPO-LegS model. This provides a more fundamental geometric foundation for HiPPO itself.
Modular Design Space: The framework enables the creation of novel SSMs by independently selecting warping functions and basis functions, facilitating designs such as multi-resolution memory (tracking long-term trends and fine-grained details simultaneously) and frequency-preserving oscillatory models.

Results

The paper validates the framework through both theoretical proofs and numerical simulations:

Theoretical Validation:
- Proposition 1 & 2: The authors derive continuous-time generators ( $A_{gen}, B_{gen}$ ) from their discrete lag operator framework. They demonstrate that for exponential warping, these generators match the known HiPPO-LegS matrices (up to transposition and a tilting shift), confirming the framework's validity.
- Convergence: The discrete transition matrix $A_\Delta$ is shown to converge to the matrix exponential of the continuous generator as the timestep $\Delta \to 0$ .
Numerical Validation:
- Matrix Equivalence: Simulations using a chaotic Lorenz63 signal confirm that the matrices derived from the lag operator framework are numerically equivalent to the analytical HiPPO-LegS matrices, with negligible Frobenius norm differences.
- Memory Dynamics: The reconstructed signal history from the Lag Operator model aligns almost perfectly with the HiPPO-LegS baseline (Mean Squared Error $\approx 5.76 \times 10^{-7}$ ). Both models successfully capture recent history while exhibiting the expected exponential decay for distant past events.

Significance and Claims

The paper positions this work as a foundational step toward a more flexible and interpretable theory of Structured State Space Models.

Conceptual Clarity: The framework unifies memory control into a single component—the time-warp $\sigma_t$ —replacing the complex, multi-stage derivations of HiPPO with a direct geometric projection. This simplifies the understanding of how memory profiles are engineered.
Design Flexibility: By decoupling the basis from the time-warp, the framework opens a modular design space. The authors claim this allows for the principled construction of SSMs with heterogeneous memory decay rates (e.g., multi-resolution memory) and specialized handling of oscillatory signals, which are difficult to achieve with standard ODE-based derivations.
Interchangeability: The numerical results validate that the proposed framework produces a memory component that behaves identically to the original HiPPO recurrence, establishing it as a robust, interchangeable building block for future sequence modeling architectures.

The authors explicitly state that while they provide proof-of-concept validation and theoretical sketches for novel designs (like Multi-Resolution Lag-SSMs), full implementation and downstream benchmarking on standard tasks are deferred to future work. The current contribution is primarily theoretical, offering a new perspective on the mechanisms underlying state-space models.