Flow Equivariant State Space Model
Abstract
Natural sequences often contain signals that move over time, such as translating objects, rotating digits, and drifting storm cells. Selective state space models such as Mamba process such data efficiently, but their recurrent state is updated at fixed spatial coordinates. As a result, evidence from a moving object can be accumulated across inconsistent locations rather than in the object's moving reference frame. To address this problem, we introduce Flow-SSM, a flow-equivariant selective state space model that lifts recurrent memory to candidate motion branches and transports each branch along its corresponding flow before applying the selective update. To handle real-world sequences where broad motion and local deformation often occur at different scales, we further introduce Hierarchical Flow-SSM, which organizes flow-aligned memory in a coarse-to-fine hierarchy rather than a single flat representation. We validate our architecture progressively: demonstrating exact transformation capture on Moving and Rotating MNIST, interpretable flow structures on KTH action videos, and the necessity of our multi-scale design for complex forecasting on the SEVIR weather dataset. Beyond merely improving predictive accuracy, our analyses confirm that the explicit flow alignment and hierarchical communication contribute measurably to the observed gains. These results support transported recurrence as an effective inductive bias for spatiotemporal sequence modeling.