SMR : State Memory Replay for Long Sequence Modeling (2024)

Abstract: Despite the promising performance of state space models (SSMs) in long sequence modeling, limitations still exist. Advanced SSMs like S5 and S6 (Mamba) in addressing non-uniform sampling, their recursive structures impede efficient SSM computation via convolution. To overcome compatibility limitations in parallel convolutional computation, this paper proposes a novel non-recursive non-uniform sample processing strategy. Theoretical analysis of SSMs through the lens of Event-Triggered Control (ETC) theory reveals the Non-Stable State (NSS) problem, where deviations from sampling point requirements lead to error transmission and accumulation, causing the divergence of the SSM’s hidden state. Our analysis further reveals that adjustments of input sequences with early memories can mitigate the NSS problem, achieving Sampling Step Adaptation (SSA).Building on this insight, we introduce a simple yet effective plug-and-play mechanism, State Memory Replay (SMR), which utilizes learnable memories to adjust the current state with multi-step information for generalization at sampling points different from those in the training data. This enables SSMs to stably model varying sampling points. Experiments on long-range modeling tasks in autoregressive language modeling and Long Range Arena demonstrate the general effectiveness of the SMR mechanism for a series of SSM models.

## introduction long sequence modeling has attracted extensive interest due to its broad prospects in natural language processing @xcite @xcite . the mainstream architectures for sequence modeling mainly focus on attention-based transformers @xcite . however, the quadratic complexity of softmax attention brings a computational bottleneck * equal contributions. although the introduction of linear attention @xcite reduces the computational complexity, it cannot well approximate the performance of the vanilla transformer. more importantly, purely attention-based architectures cannot capture longrange dependencies well. on the other hand, state space model (ssm)-based architectures @xcite show superior performance on the long range arena (lra) @xcite benchmark for long sequence modeling due to their linear computational complexity and excellent long-range dependency capturing ability. existing ssm-based model architectures, such as s5 @xcite and s6 @xcite , primarily rely on recursive structures to tackle the varying sampling step issue. s5 introduced learnable step sizes for each step to improve the sampling step adaptation (ssa) capability of ssm. s6 (mamba) introduced data-dependent parameter settings, which makes the state propagation of the ssm model more flexible. however, this restricts its inference computation to parallel scanning instead of the original efficient convolution mode, significantly hampering training efficiency and imposing a heavier inference burden when handling long inputs at once. to address the mentioned issues, we aim to propose a method that goes beyond recursive constraints to improve ssa capability. this strategy seeks to enhance the ssm, making it more adaptable and flexible for various parallel convolution computation types, including advanced architectures like s4 @xcite , mega @xcite , spade @xcite , and more. specifically, we leverage the event-triggered control (etc) theory @xcite to provide the first demonstration of the non-stable state (nss) problem in ssms. we show that for a fixed-parameter ssm, varying sampling steps, deviating from the model's sampling point requirements, triggers error propagation and accumulation, ultimately leading to the divergence of the hidden state. our analysis further reveals that adjustments based on early memories of the input sequence can achieve ssa, effectively solving the nss problem. inspired by this finding, we propose a simple yet effective plug-and-play mechanism, state memory replay (smr), it can significantly alleviate the nss problem in ssms by improving ssm the capability of ssa thus bring further sequence modeling capabilities. in particular, smr can achieve better generalization ability at different sampling points, especially when dealing with the stochastic selected sampling points. we conduct experiments on autoregressive language modeling on wikitext-103 @xcite and long sequence modeling on lra. the results show that the smr mechanism can bring better performance to ssm-based model, on both autoregressive language modeling and long sequence modeling tasks. it can also further improve a series of competitive ssm-based models such as s5, spade, mega, and s6, which verifies the generality and effectiveness of the proposed smr mechanism. in summary, our main contributions are three folds: ## preliminaries: state space models the state space model is formally defined by eq.( 1 ) and eq.( 2 ): ## ssa via state memory replay in this section, we aim to reveal the problem of nss in ssm caused by changes in sampling points through etc theory @xcite (section 3.1). we demonstrate that unstable hidden states lead to errors in ssm (section 3.2). furthermore, through the analysis based on etc theory, we propose a simple but effective step-size adaptation mechanism, smr, to enhance the model's ssa capability thus alleviate the nss problem in ssm with fixed step setting (section 3.3). experimental results indicate that the smr mechanism can not only enhance the ssa capabil-ity of ssm with fixed parameters but can also be extended to other ssm-based models, improving their ssa capabilities (section 3.4). ## conclusion in this paper, we investigated the nss issue in ssms for long sequence modeling, we found that when input data deviates from the model's sampling requirements, it leads to error accumulation and hidden state divergence. our analysis further revealed that early memory adjustments in the input sequence can achieve adaptive sampling, effectively solving the nss problem. inspired by this, we proposed a simple yet efficient plug-and-play mechanism, smr. theoretical analysis and experiments demonstrated that smr effectively alleviates nss, enhancing the generalization ability of ssms to diverse sampling points and leading to superior 1 https://github.com/alxndrtl/mamba.py.git sequence modeling performance. we evaluated smr on various ssm-based models, including the convolution-based and recurrence-based ssms, applying it to both autoregressive language modeling (on wikitext-103) and the lra benchmark. the results demonstrate that smr significantly improves the performance of ssm-based models on these tasks, solidifying its effectiveness and broad applicability.