Mamba State-Space Models Are Lyapunov-Stable Learners

TMLR Paper5050 Authors

Abstract:

Mamba state-space models (SSMs) have recently outperformed state-of-the-art (SOTA) Transformer large language models (LLMs) in various tasks and been widely adapted. However, a major concern for stable learning in recurrent-based deep models (such as SSMs) is the sensitivity of their recurrent dynamics. Despite widespread adaptation, the sensitivity of Mamba’s recurrent dynamics under common fine-tuning methods–e.g., mixed-precision fine-tuning (MPFT) and parameter-efficient fine-tuning (PEFT)–remains unexplored. Empirically,
we show that Mamba LLMs are extremely stable to changes introduced by combinations of MPFT and PEFT, in stark contrast to Transformer LLMs, which we demonstrate may drastically diverge from their respective full-precision counterparts under different
combinations of MPFT and PEFT (despite the near-ubiquitous adaptation of these fine-tuning frameworks for attention-based models). The demonstrated robustness of Mamba LLMs are due to their recurrent dynamics, which we prove are guaranteed to be stable using
dynamical systems theory (in particular, Lyapunov stability). We conclude by using MPFT and PEFT to novelly study Mamba LLMs’ in-context learning (ICL) abilities on natural language tasks, thus supplementing other recent work.

Submission Length: Regular submission (no more than 12 pages of main content)

Assigned Action Editor: ~Bamdev_Mishra1

Submission Number: 5050