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

TL;DR

This paper establishes a mathematical link between state space models used in sequence modeling and nonlinear oscillator networks, providing an exact analytical expression for their computation.

Contribution

It introduces a novel mathematical correspondence that offers an exact operator expression for the S4D model, enhancing interpretability of neural sequence models.

Findings

Derived an exact operator expression for S4D's forward pass

Revealed how nonlinear decoders enable sequence classification

Provided a physical interpretation of SSM architectures as oscillator networks

Abstract

We establish a mathematical correspondence between state space models, a state-of-the-art architecture for capturing long-range dependencies in data, and an exactly solvable nonlinear oscillator network. As a specific example of this general correspondence, we analyze the diagonal linear time-invariant implementation of the Structured State Space Sequence model (S4). The correspondence embeds S4D, a specific implementation of S4, into a ring network topology, in which recent inputs are encoded, as waves of activity traveling over the one-dimensional spatial layout of the network. We then derive an exact operator expression for the full forward pass of S4D, yielding an analytical characterization of its complete input-output map. This expression reveals that the nonlinear decoder in the system induces interactions between these information-carrying waves that enable classifying real-world sequences. These results generalize across modern SSM architectures, and show that they admit an exact mathematical description with a clear physical interpretation. These insights enable a new level of interpretability for these systems in terms of nonlinear oscillator networks.