LISM: Long-range Integrative State space Models via Input-Latent State Interactions

Abstract

State space models (SSMs) are an emerging paradigm that achieves linear-time scaling, however, they intrinsically suffer from "curse of memory". The memory of SSMs, including Mamba, decays exponentially as long as the recursive update is stable. In this work, we developed Long-range Integrative State space Models (LISM) to mitigate the curse of memory in SSMs. LISM allows the recursive update to occasionally be unstable, which enables longer memory retention, and incorporates a dynamic input-state interaction mechanism to identify inputs that should be memorized for longer periods. We theoretically establish probability bounds for the overall stability of LISM and derive a probabilistic stability condition to guide parameterization and training. We evaluated LISM's performance on multiple tasks, a synthetic Selective Copy task, two real language tasks WikiText-103 and LRA Retrieval, and an additional biological prediction task. LISM achieved the best performance on the Selective Copying task, WikiText-103, and the biological prediction task compared to SSM-based models and achieved comparable performance on the LRA-Retrieval task with the state-of-the-art. Beyond average performance, LISM also shows strong repeatability across runs. In the two real language tasks, LISM reduced the relative performance variance (coefficient of variation) by an average of 33.3%, achieving a standard deviation of $\pm$0.30% on WikiText-103 compared to $\pm$0.46% for Mamba-2. The implementation of LISM is available at https://github.com/CM-Compbio-Group/LISM_dev.