An explicit operator explains end-to-end computation in the modern neural networks used for sequence and language modeling

1. Problem Statement

• The Limitation of Transformers: While Transformers are dominant in sequence modeling, they suffer from quadratic computational complexity ($O(N^2)$) regarding sequence length due to the "matrix attention" mechanism, which computes pairwise interactions between all tokens.
• The Promise and Opacity of SSMs: State Space Models (SSMs), such as the Structured State Space Sequence model (S4), offer linear scaling ($O(N)$) and have achieved state-of-the-art performance. However, their internal mechanics remain largely opaque.
  • SSMs rely on linear dynamics (mathematically solvable) combined with a nonlinear decoder (necessary for complex tasks).
  • The interleaving of linear recurrence and nonlinear activation makes end-to-end analytical explainability extremely difficult. Current interpretability methods are largely empirical (post-hoc) rather than derived from first principles.
• The Gap: There is no exact mathematical description that characterizes the complete input-output map of modern SSMs, specifically how the nonlinear decoder interacts with the linear recurrent state to perform classification.

2. Methodology

The authors establish a rigorous mathematical framework linking SSMs to nonlinear oscillator networks.

• Mathematical Correspondence:

  • They map the diagonal linear time-invariant (LTI) implementation of S4 (specifically S4D) to an exactly solvable network of coupled nonlinear oscillators arranged on a 1D ring.
  • Using a nonlinear coordinate transformation ($x_i = e^{i\psi_i}$), the complex nonlinear oscillator dynamics are transformed into a linear system governed by a matrix exponential: $x(t) = e^{Kt}x(0)$.
  • This reveals that S4D is mathematically equivalent to a system where the state evolves via a diagonal operator in a specific eigenbasis (the Fourier basis).

• Traveling Wave Interpretation:

  • In this mapped framework, the latent state of the SSM is interpreted as traveling waves propagating across the ring topology.
  • Input sequences excite specific eigenmodes (waves) with distinct spatial frequencies and temporal decay rates.

• Carleman Embedding for Nonlinearity:

  • To handle the nonlinear activation function (e.g., GELU) which breaks the linearity, the authors employ Carleman embedding.
  • This technique lifts the nonlinear system into an infinite-dimensional linear space by representing the nonlinearity as a hierarchy of polynomial interactions (monomials) of the state variables.
  • They derive an exact operator expression for the full forward pass, decomposing the computation into:
    1. Linear terms: Direct wave excitation by inputs.
    2. Quadratic/Higher-order terms: Interactions between waves (e.g., $\mu_i \mu_j$) induced by the nonlinearity.

3. Key Contributions

1. Exact Operator Expression: The paper derives a closed-form mathematical operator (Eq. 32) that describes the entire end-to-end computation of an SSM, from input $u_k$ to output $O$, without approximation errors (until truncation of the Carleman series).
2. Physical Interpretation of SSMs: It provides a concrete physical interpretation of SSMs as traveling wave systems. The "black box" of the SSM is revealed to be a mechanism where inputs generate waves that travel, interfere, and interact.
3. Mechanism of Nonlinearity: The study demonstrates that while linear dynamics store information about the past via wave amplitudes, the nonlinear decoder is essential for classification because it induces wave-wave interactions. These interactions extract information based on relative phases and cross-frequency couplings, which linear dynamics alone cannot capture.
4. Generalizability: The framework applies to all diagonal SSMs (including S4D-Lin, S4D-Inv, etc.) and suggests that the specific eigenvalue structure of S4D ensures uniform decay rates, preventing mode domination and allowing flexible processing.

4. Key Results

• Synthetic Data (Sinusoids): On a simple dataset distinguishing 15 Hz vs. 20 Hz sinusoids in noise, the authors showed that classification relies on a small subset of resonant modes. The linear dynamics alone (first-order Carleman term) were sufficient to separate these classes by accumulating energy in specific traveling waves.
• Real-World Data (SelfRegulationSCP1):
  • On a complex real-world sequence classification task, the linear component alone explained 83% of the model's performance.
  • Adding second-order interactions (quadratic terms representing wave-wave coupling) increased the explained performance to 94%.
  • The full GELU activation (infinite order) explained 100%.
  • Conclusion: Real-world sequence processing relies heavily on nonlinear wave interactions (cross-frequency coupling) rather than just the amplitude of individual waves. The nonlinear readout effectively "mixes" the wave modes to separate classes.
• Efficiency of Approximation: Truncating the Carleman expansion to just the second order captures the vast majority of the model's computational power, suggesting that SSMs can be approximated efficiently by low-order polynomial interactions of traveling waves.

5. Significance and Implications

• Theoretical Breakthrough: This is the first work to provide a first-principles, analytical description of the forward pass in modern SSMs. It moves beyond empirical observation to mathematical derivation.
• Interpretability: It shifts the paradigm from "post-hoc" analysis to "prospective" understanding. Researchers can now predict how specific input features propagate through the system as waves and how they interact to produce outputs.
• Controllability: By understanding the computation as explicit mathematical operators, future work could design controllable language models where outputs are steered by manipulating specific eigenmodes or wave interactions, offering new guarantees for safety and reliability.
• Architecture Design: The framework suggests design principles for future SSMs:
  • Tuning eigenvalue spectra to create specific interference patterns.
  • Designing connectivity to optimize wave propagation for longer time dependencies.
  • Extending the analysis to time-varying architectures like Mamba.

In summary, the paper demystifies the "black box" of SSMs by revealing them as nonlinear oscillator networks where computation is performed through the generation, propagation, and nonlinear interaction of traveling waves.