Parameter-Efficient Fine-Tuning of State Space Models

2024-10-11

arXiv.org

AI Summary

Existing parameter-efficient fine-tuning (PEFT) methods—particularly LoRA—exhibit poor adaptability to state space models (SSMs), such as Mamba, due to structural mismatches with SSM’s core components. Method: This work is the first to systematically characterize the structural sensitivity of SSMs to mainstream PEFT techniques and proposes Sparse Dimensional Tuning (SDT), the first SSM-aware PEFT method. SDT applies sparse low-rank updates directly to key SSM parameters (e.g., A, B, C matrices), enabling fine-grained, structure-aware optimization. It further integrates with LoRA into a unified SDT+LoRA framework. Contribution/Results: Extensive experiments on language modeling tasks demonstrate that SDT achieves state-of-the-art performance with <0.5% parameter overhead—significantly outperforming LoRA and other baselines—while attaining superior trade-offs between accuracy and efficiency.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Search and Optimization: Learning to Search
• Natural Language Processing: (Large) Language Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Large language models for search
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

Deep State Space Models (SSMs), such as Mamba (Gu&Dao, 2024), have become powerful tools for language modeling, offering high performance and linear scalability with sequence length. However, the application of parameter-efficient fine-tuning (PEFT) methods to SSM-based models remains underexplored. We start by investigating two fundamental questions on existing PEFT methods: (i) How do they perform on SSM-based models? (ii) Which parameters should they target for optimal results? Our analysis shows that LoRA and its variants consistently outperform all other PEFT methods. While LoRA is effective for linear projection matrices, it fails on SSM modules-yet still outperforms other methods applicable to SSMs, indicating their limitations. This underscores the need for a specialized SSM tuning approach. To address this, we propose Sparse Dimension Tuning (SDT), a PEFT method tailored for SSM modules. Combining SDT for SSMs with LoRA for linear projection matrices, we achieve state-of-the-art performance across extensive experiments.

Problem

• Explores parameter-efficient fine-tuning for SSM-based models.
• Identifies limitations of existing PEFT methods on SSM modules.
• Proposes Sparse Dimension Tuning (SDT) for improved SSM performance.

Innovation

• Proposes Sparse Dimension Tuning for SSMs
• Combines SDT with LoRA for linear projections
• Achieves state-of-the-art performance in experiments