fastiSSM: Fast inference of state space model with online model approximation in frequency-domain.

Abstract

State space model (SSM) is a promising structure for long sequence time-series forecasting (LSTF) and natural language processing (NLP) applications. Although it has demonstrated superior sequence modeling efficiency and ability to compete with and even outperform the traditional models, the large size of SSM leads to slow speed and excessive memory consumption for inference, prohibiting it from being used in applications with speed/memory/energy constraints, especially on devices with limited resources. In this work, we propose fastiSSM, which accelerates the inference of general SSMs through online model approximation. To overcome the computational efficiency and approximation accuracy challenges in existing model approximation methods for SSMs, we introduce two new techniques in fastiSSM. First, fastiSSM has a parameter-sharing structure in its continuous-time models. With this structure, only one model approximation is needed to generate the compact recurrences for fast inference, compared with multiple ones performed in traditional methods. This brings high model approximation efficiency for fastiSSM. Second, we introduce the frequency-domain based model approximation technique in fastiSSM by exploiting the linear time-invariant (LTI) property of SSM. It results in robust compact models and better error control in model approximation, and finally leads to fast and accurate inference. Comprehensive experiments are performed by comparing fastiSSM with both original recurrence and compact recurrence based inference methods on multiple LSTF and audio generation tasks. Experimental results confirm that fastiSSM is able to approximate SSMs with online computing efficiency and accelerate the inference effectively with marginal error. It outperforms the comparison methods in both model approximation efficiency and inference performance, making it a better choice to accelerate the inference of large-scale SSMs. The fastiSSM code is available as open source on GitHub.