Provable Distillation of Observable Systems into Linear Dynamical Systems

Abstract

We present the first provable method for learning an observable nonlinear dynamical system using a symmetric linear dynamical system with sublinear regret relative to the best linear predictor. Our approach builds on recent work on Observational Spectral Filtering for learning dynamical systems and on recent work on the distillation of linear dynamical systems. We showcase how our theoretical guarantees translate into strong practical models, achieving comparable performance to a parameter-matched Gated Recurrent Unit on MuJoCo environments, and with extensions to these methods, we unlock efficiency gains on the FlashSTU convolutional language model. Building on SpectraLDS, we accelerate inference on the 550M FlashSTU model by an additional 24%, and we accelerate inference on the Spectral Transform Unit by up to 9.4x at large batch sizes. Additionally, we improve the pre-existing STU inference baselines, reducing the hidden dimension of SpectraLDS by 50% and producing a PyTorch implementation of the Continuous FutureFill algorithm.