MambaPEFT: Exploring Parameter-Efficient Fine-Tuning for Mamba

TL;DR

This paper investigates parameter-efficient fine-tuning (PEFT) for Mamba, a State Space Model–based alternative to Transformers with linear-time sequence processing. It adapts existing Transformer PEFTs to Mamba, redesigns them for SSM architecture, and introduces new Mamba-specific PEFT methods (Affix-tuning, Partial LoRA, Partial-tuning, Additional-scan) alongside a two-stage Hybrid PEFT Search to identify effective method combinations. Experiments on image and language tasks show PEFT yields larger performance gains for Mamba than for Transformers, with LoRA variants excelling under limited data and Additional-scan proving strong with larger datasets, while Affix-tuning benefits large models. The proposed framework enables efficient, scalable fine-tuning by balancing method diversity, parameter count, and hyperparameters, and the authors plan to release code to promote reproducibility and broader adoption.

Abstract

An ecosystem of Transformer-based models has been established by building large models with extensive data. Parameter-efficient fine-tuning (PEFT) is a crucial technology for deploying these models to downstream tasks with minimal cost while achieving effective performance. Recently, Mamba, a State Space Model (SSM)-based model, has attracted attention as a potential alternative to Transformers. While many large-scale Mamba-based models have been proposed, efficiently adapting pre-trained Mamba-based models to downstream tasks remains unexplored. In this paper, we conduct an exploratory analysis of PEFT methods for Mamba. We investigate the effectiveness of existing PEFT methods for Transformers when applied to Mamba. We also modify these methods to better align with the Mamba architecture. Additionally, we propose new Mamba-specific PEFT methods that leverage the distinctive structure of Mamba. Our experiments indicate that PEFT performs more effectively for Mamba than Transformers. Lastly, we demonstrate how to effectively combine multiple PEFT methods and provide a framework that outperforms previous works. To ensure reproducibility, we will release the code after publication.

Figures (5)

• Figure 1: An overview of our proposed MambaPEFT. We investigate, improve, and propose 20 variations of seven PEFT methods for Mamba and search for the best combination.
• Figure 2: (a) Prefix-tuning designed for Transformer. It can't be applied to Mamba. (b) The proposed Affix-tuning, which we re-design for Mamba. It discards prefixes after SSM. This design allows us to insert affix tokens at arbitrary locations. (c) Additional-scan that we design for Mamba. In this method, we add a learnable dimension to the hidden state in SSM.
• Figure 6: Ablation studies for vision tasks. (a) shows the relationship between the number of tokens and performance in Affix-tuning with projection. (b) shows the relationship between the number of tokens and performance in Prompt-tuning. (c) shows the relationship between the rank of LoRA and performance in Mamba. (d) shows the relationship between the rank of LoRA and performance in ViT. Performance consistently decreases as the rank increases, contrary to the case in Mamba.
• Figure 11: Scatter plot of the evaluations for the methods in \ref{['app:detail_methods']}, showing the relationship between the proportion of trainable parameters and performance
• Figure 12: This ablation study investigates how optimal hyperparameters differ per PEFT method and the size of base models in language tasks. (a) Optimal learning rate per PEFT method. (b) The relationship between the ratio of trainable parameters and performance on Mamba 370M base model. (c) Suitable learning rate per base model size for Additional-scan. (d) Optimal additional dimension per base model size for Additional-scan.