English

Enhanced Structured State Space Models via Grouped FIR Filtering and Attention Sink Mechanisms

Computation and Language 2024-08-02 v1 Machine Learning

Abstract

Structured State Space Models (SSMs) have emerged as compelling alternatives to Transformer architectures, offering linear-time complexity and superior performance in various sequence modeling tasks. Despite their advantages, SSMs like the original Mamba-2 face training difficulties due to the sensitivities introduced by the extended series of recurrent matrix multiplications. In this paper, we propose an advanced architecture that mitigates these challenges by decomposing A-multiplications into multiple groups and optimizing positional encoding through Grouped Finite Impulse Response (FIR) filtering. This new structure, denoted as Grouped FIR-enhanced SSM (GFSSM), employs semiseparable matrices for efficient computation. Furthermore, inspired by the "attention sink" phenomenon identified in streaming language models, we incorporate a similar mechanism to enhance the stability and performance of our model over extended sequences. Our approach further bridges the gap between SSMs and Transformer architectures, offering a viable path forward for scalable and high-performing sequence modeling.

Keywords

Cite

@article{arxiv.2408.00244,
  title  = {Enhanced Structured State Space Models via Grouped FIR Filtering and Attention Sink Mechanisms},
  author = {Tian Meng and Yang Tao and Wuliang Yin},
  journal= {arXiv preprint arXiv:2408.00244},
  year   = {2024}
}