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

TL;DR

This paper introduces GFSSM, an improved Structured State Space Model that uses grouped FIR filtering and attention sink mechanisms to enhance training stability and performance in sequence modeling tasks.

Contribution

The paper proposes GFSSM, combining grouped FIR filtering and attention sink mechanisms to address training challenges and improve SSM performance.

Findings

Enhanced training stability over long sequences

Improved sequence modeling performance

Bridging SSMs and Transformer architectures

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.