Learning to Dissipate Energy in Oscillatory State-Space Models
By: Jared Boyer, T. Konstantin Rusch, Daniela Rus
Potential Business Impact:
Helps computers remember long stories better.
State-space models (SSMs) are a class of networks for sequence learning that benefit from fixed state size and linear complexity with respect to sequence length, contrasting the quadratic scaling of typical attention mechanisms. Inspired from observations in neuroscience, Linear Oscillatory State-Space models (LinOSS) are a recently proposed class of SSMs constructed from layers of discretized forced harmonic oscillators. Although these models perform competitively, leveraging fast parallel scans over diagonal recurrent matrices and achieving state-of-the-art performance on tasks with sequence length up to 50k, LinOSS models rely on rigid energy dissipation ("forgetting") mechanisms that are inherently coupled to the timescale of state evolution. As forgetting is a crucial mechanism for long-range reasoning, we demonstrate the representational limitations of these models and introduce Damped Linear Oscillatory State-Space models (D-LinOSS), a more general class of oscillatory SSMs that learn to dissipate latent state energy on multiple timescales. We analyze the spectral distribution of the model's recurrent matrices and prove that the SSM layers exhibit stable dynamics under simple, flexible parameterizations. D-LinOSS consistently outperforms previous LinOSS methods on long-range learning tasks, without introducing additional complexity, and simultaneously reduces the hyperparameter search space by 50%.
Similar Papers
The Curious Case of In-Training Compression of State Space Models
Machine Learning (CS)
Shrinks computer models during learning for speed.
Deep Learning-based Approaches for State Space Models: A Selective Review
Machine Learning (Stat)
Helps computers understand changing information better.
Regularity and Stability Properties of Selective SSMs with Discontinuous Gating
Machine Learning (CS)
Makes AI remember things better and more reliably.