Parameter-Efficient Fine-Tuning of State Space Models

Authors: Kevin Galim, Wonjun Kang, Yuchen Zeng, Hyung Il Koo, Kangwook Lee
Primary area: transfer learning, meta learning, and lifelong learning
Keywords: parameter-efficient fine-tuning, state space model, mamba, lora

Abstract:
Deep State Space Models (SSMs), such as Mamba `(Gu \& Dao, 2023)`, have emerged as powerful tools for language modeling, offering high performance with efficient inference and linear scaling in sequence length. However, the application of parameter-efficient fine-tuning (PEFT) methods to SSM-based models remains largely unexplored. This paper aims to systematically study two key questions: (i) How do existing PEFT methods perform on SSM-based models? (ii) Which modules are most effective for fine-tuning? We conduct an empirical benchmark of four basic PEFT methods on SSM-based models. Our findings reveal that prompt-based methods (e.g., prefix-tuning) are no longer effective, an empirical result further supported by theoretical analysis. In contrast, LoRA remains effective for SSM-based models. We further investigate the optimal application of LoRA within these models, demonstrating both theoretically and experimentally that applying LoRA to linear projection matrices without modifying SSM modules yields the best results, as LoRA is not effective at tuning SSM modules. To further improve performance, we introduce LoRA with Selective Dimension tuning (SDLoRA), which selectively updates certain channels and states on SSM modules while applying LoRA to linear projection matrices. Extensive experimental results show that this approach outperforms standard LoRA.

Decision: reject
Reviewer scores: 5, 6, 5, 3

Reviews and discussion:
<thread>
creation date:
1730656975301

thread_type:
official_review

<thread_object>
creator: reviewer_atck

<content>
summary:
This paper presents a systematic study on the application of parameter-efficient fine-tuning (PEFT) methods to Deep State Space Models (SSMs). The paper reveals that prompt-based methods are less effective for SSMs, while LoRA remains effective. The authors further propose SDLoRA, which selectively updates certain channels and states on SSM modules while applying LoRA to linear projection matrices, demonstrating improved performance over standard LoRA.

strengths:
The paper systematically analyzes how existing PEFT methods perform on SSM-based models and which modules are most effective for fine-tuning, revealing that prompt-based methods like prefix-tuning are no longer effective. Additionally, it shows that applying LoRA to linear projection matrices without modifying SSM modules yields the best results. The paper further introduces SDLoRA, a novel approach to parameter-efficient fine-tuning (PEFT) for SSMs. This method's innovativeness lies in its selective dimension updating strategy within SSM modules. The document also references various datasets used for evaluating the proposed methods, and the clarity of the paper is generally good.

weaknesses:
1. While the paper presents a novel approach for parameter-efficient fine-tuning SSMs, the innovation seems to be more incremental than groundbreaking. 
2. The paper lacks a detailed analysis of the computational overhead from the selective dimension tuning process. This is crucial to understanding the trade-offs between parameter reduction and computational efficiency in SDLoRA.
3. The paper would benefit from a  detailed analysis of SDLORA with hybrid models that combine SSMs and Transformers, as these models are becoming increasingly popular and have shown promise in various domains.

questions:
1. Could the authors comment on how SDLoRA benefit to hybrid models that combine SSMs and Transformers? 
2. On Mamba-130M, the authors use GLUE and DART benchmarks, while on Mamba-1.4B, they use SAMSum and Spider benchmarks. Could the authors elaborate on the considerations behind this benchmark selection strategy?
3. In Table 4, SDLoRA outperforms Full Fine-Tuning. Was this outcome expected, and if so, could the authors provide insights into why this might be the case? Additionally, have the authors considered conducting experiments on more challenging visual tasks, such as Imagenet-1k, to further validate the effectiveness of SDLoRA?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730679814316

thread_type:
official_review

<thread_object>
creator: reviewer_1uqc

<content>
summary:
The paper aims to explore the application of LoRA to SSMs. It explores different adapter types as well as different ways of how to apply them to the SSM block. The paper additionally includes theoretical results to justify their choices.

strengths:
The paper is looking at something which definitely needs to be studied, as it is not a foregone conclusion that adapters, prompt tuning, etc... will have the same benefits for SSMs as they do for transformers. A detailed study understanding the different tradeoffs brought on by the SSM architecture should be explored.

The paper has a good mix of theoretical and empirical analysis.

weaknesses:
The paper, to me, does not have enough substance. It needs either more detailed theoretical results, which result in some innovation or needs more empirical results which help the community understand how different adapters perform with transformers vs SSMs.

Here are some specifics:
1. Mamba2 is not included in any of the results, the paper should be updated to look at this architecture as well
2. The theoretical analysis is largely only for the S4 block, its not clear to me the conclusions would extend to S6, and aren't convincing to me for that reason.
3. The empirical results are lacking. The SDLoRA module is slightly different application of the standard LoRA block, essentially targeting different layers within the block instead of new a design. Also only two adapters are compared within this work. To have a true empirical study of this, many more experiments need to be conducted, even in the presence of theoretical results.

I think the paper might be more interesting if it were to do something like the following:
- Compare many adapters in a standardized for Transformers, Mamba1, and Mamba2
- Isolate differences in which adapters perform well for each class of model
- See if these insights give rise to a new adapter design specifically for this model class
- Draw theoretical results to try to explain why certain adapters perform differently than in Transformers

The current structure doesn't feel as though it is contributing much

questions:
Have you looked at all into Hybrid Architectures?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730723128984

thread_type:
official_review

<thread_object>
creator: reviewer_mtue

<content>
summary:
The paper explores parameter-efficient fine-tuning (PEFT) methods for deep State Space Models (SSMs), especially in the context of language modeling tasks. It investigates the effectiveness of various PEFT methods, such as prompt-based prefix-tuning and Low-Rank Adaptation (LoRA), applied to SSM architectures like Mamba. A new variant called SDLoRA (Selective Dimension LoRA) is proposed in this paper to selectively update channels and states in the SSM modules, aiming to enhance the fine-tuning performance while reducing parameters. The results indicate that SDLoRA outperforms conventional LoRA when fine-tuning SSM-based models.

strengths:
(1) Efficiency and Scalability: By focusing on selective parameter updates, the SDLoRA method enhances computational efficiency, which is crucial for large-scale models. Experimental results show that SDLoRA consistently outperforms traditional LoRA across several benchmarks, proving its efficacy in SSM architectures.
(2) Adaptability: The proposed SDLoRA method demonstrates adaptability across multiple tasks, including NLP tasks and vision tasks

weaknesses:
(1) Limited Applicability Beyond SSMs: The focus on SSMs means SDLoRA may not generalize well to non-SSM architectures or hybrid models such as Transformer models or Transformer-SSM combinations. Its broader applicability to other architectures remains untested.
(2) Parameter Selection: The dimension selection process in SDLoRA relies on parameter magnitudes, which may not be optimal and could benefit from a more sophisticated selection algorithm. And what if the magnitude of each channel changes during the fine-tuning stage?

questions:
See the weaknesses
</content>
</thread_object>

</thread>
<thread>
creation date:
1730790960165

thread_type:
official_review

<thread_object>
creator: reviewer_n23z

<content>
summary:
This paper presents the first study on the performance of PEFT methods applied to SSM-based models. Specifically, prompt-based and parameter-based methods are involved. With theoretical analysis and extensive experiments, LoRA tends to achieve better performance. To further improve the performance, this paper introduces SDLoRA, which selectively updates certain channels and states on SSM modules while applying LoRA to linear projection matrices.

strengths:
1.	Reasonable theoretical analysis and comprehensive experiments.
2.	The introduced SDLoRA is novel and effective.
3.	Through extensive experiments, the findings in this paper is useful and inspired.

weaknesses:
1.	The speed of SDLoRA is nor reported.
2.	Experimental results on larger datasets are needed.

questions:
1.	During the process of selective dimension tuning, the authors select the target channels and states based on magnitude, any other metrics have been tried?
2.	Will SDLoRA's training speed be slower compared to vanilla LoRA? How much slower will it be?
3.	What is the accuracy of SDLoRA on a larger data set, such as ImageNet?
4.	Some other advanced parameter-efficient tuning method like DoRA [1] can be adapted to Mamba? or the proposed SDLoRA can be adapted to Jamba?

[1] DoRA: Weight-Decomposed Low-Rank Adaptation
</content>
</thread_object>

</thread>
<thread>
creation date:
1732279187304

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
To AC and All Reviewers

comment:
We would like to thank all reviewers for their comments and helpful feedback. We are particularly encouraged that the reviewers have found that (i) our paper is clear (`R-AtCk`), useful and inspired (`R-n23z`), and studying what definitely needs to be studied (`R-1UQC`), (ii) our paper provides theoretical (`R-n23z`, `R-1UQC`), empirical (`R-1UQC`), systematic (`R-AtCk`), and detailed (`R-1UQC`) analysis with comprehensive (`R-n23z`) and extensive (`R-n23z`) experiments across various datasets (`R-AtCk`), and (iii) our method is novel (`R-n23z`, `R-AtCk`), innovative (`R-AtCk`), effective (`R-n23z`), efficient (`R-mTUE`), scalable (`R-mTUE`) and adaptable (`R-mTUE`).

In response to the feedback, we have addressed each concern, added experimental results, and will update our paper accordingly. Below, we summarize the major updates in our rebuttal. 


# Major Update 1: Expanded Model Coverage — Jamba & Mamba-II.
**Key Insight: SDLoRA outperforms LoRA on Jamba [1] and Mamba-II [2] models.**


In response to the reviewers' feedback, we expanded our analysis beyond the deep S4 model and Mamba presented in the original paper. Specifically, we incorporated the Transformer-SSM hybrid model Jamba (Jamba-Tiny-319M and Jamba-Mini-52B) and Mamba-II (Mamba-II-130M and Mamba-II-1.4B). 

> New Experiment Results on Jamba.

We froze the Transformer layers, tuning only the Mamba layers, while adhering to the same experimental settings used for Mamba. To accommodate the Jamba-Tiny 52B model on a single 80GB GPU, we quantized all non-Mamba layers to 4-bit precision, following an approach similar to QLoRA, and reduced the batch size.

**Results.** The performance comparison between LoRA and SDLoRA is shown in the table below ((!) indicates still training).  SDLoRA outperforms LoRA on nine out of eleven tasks, demonstrating that SDLoRA's strong performance on Mamba effectively transfers to hybrid models as well.


|  | Jamba-Tiny-319M |  |  |  |  |  |  |  |  |  |  | Jamba-Mini-52B |  |  |  |  |  |  |  |
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|  |  |  | DART |  |  | SAMSum |  |  |  | Spider |  |  |  | DART |  |  | SAMSum |  |  |
|  | Params (%) |  | BLEU ($\uparrow$) | METEOR  ($\uparrow$) |  | R1 ($\uparrow$) | R2  ($\uparrow$) | RL  ($\uparrow$) |  | Acc. ($\uparrow$) |  | Params (%) |  | BLEU ($\uparrow$) | METEOR  ($\uparrow$) |  | R1   ($\uparrow$) | R2   ($\uparrow$) | RL   ($\uparrow$) |
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| LoRA | 0.05030 |  | 27.17 | 65.03 |  | 37.13 | 16.43 | 30.90 |  | 35.49 |  | 0.004951 |  | 52.86 | 73 |  | 55.31 | 31.71 | 46.47 |
|  | 0.05690 |  | **39.02** | **67.90** |  | 40.80 | 18.54 | 33.87 |  | 44.07 |  | 0.005629 |  | 52.65 | 72.81 |  | 55.12 | 31.63 | 46.64 |
|  | 0.06153 |  | 23.18 | 65.05 |  | 39.15 | 17.70 | 32.79 |  | 37.67 |  | 0.006051 |  | 52.63 | 72.94 |  | 56.36 | 33.48 | 47.91 |
|  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |  |
| SDLoRA | 0.05536 |  | 31.49 | 67.18 |  | 41.11 | 18.48 | 33.84 |  | 48.58 |  | 0.005484 |  | 51.86 (!) | 72.42 (!) |  | 56.08 | 32.79 | 47.61 |
|  | 0.05540 |  | 31.43 | 67.86 |  | 41.69 | 19.17 | 34.47 |  | **50.40** |  | 0.005488 |  | 52.79 (!) | **73.07** (!) |  | **56.53** | **33.5** | **47.96** |
|  | 0.05549 |  | 33.03 | 67.80 |  | **42.18** | **19.19** | **34.95** |  | 49.60 |  | 0.005497 |  | **53.11** | 72.95 |  | 56.14 | 33.08 | 47.56 |


> New Experiment Results on Mamba-II. 

For Mamba-II, however, applying SDLoRA is not straightforward because Mamba-II further constrains $A$ such that all (non-zero) entries must have the same value. Therefore, our original dimension selection approach cannot be directly applied here. We consider a naive extension of SDLoRA by selecting dimensions in the projection matrices for input mapping vector $B$ and the projection matrices for output mapping vector $C$ using their respective magnitude, and fine-tune the selected dimensions and all elements of state transition matrix $A$. 

**Results**: The table below compares the performance of LoRA and SDLoRA on Mamba-II. The results demonstrate that SDLoRA consistently outperforms LoRA on Mamba-II models.


|  |  | DART (Mamba-II-130M) |  |  |  | SAMSum (Mamba-II-1.3B) |  |  |  |
|---|---|---|---|---|---|---|---|---|---|
|  |  | Params (%) | BLEU ($\uparrow$) | METEOR  ($\uparrow$) |  | Params (%) | R1   ($\uparrow$) | R2   ($\uparrow$) | RL   ($\uparrow$) |
| LoRA |  | 0.3354 | 48.09 | 68.71 |  | 0.1614 | 49.73 | 26.14 | 41.53 |
| SDLoRA |  | 0.3393 | **48.93** | **70.60** |  | 0.1767 | **50.72** | **27.21** | **42.54** |
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper presents a systematic investigation into parameter-efficient fine-tuning (PEFT) methods for Deep State Space Models (SSMs), with a particular focus on language modeling tasks. They introduce SDLoRA, a novel approach that selectively updates certain channels and states on SSM modules while applying LoRA to linear projection matrices, showing improved performance over standard LoRA.

One of the main concerns raised by the reviewers was the limited applicability of the proposed SDLoRA method beyond SSMs. While the authors have demonstrated the effectiveness of SDLoRA on hybrid models like Jamba, more comprehensive testing on a broader range of architectures would be beneficial. The comparison with existing PEFT methods was noted as insufficient by some reviewers. While the authors defended their approach, a more comprehensive comparison could have been made to better position the proposed method against existing benchmarks. The complexity of the proposed model, particularly with components like the selective dimension tuning, raised concerns about its practicality, which was not fully addressed in the rebuttal.

The reviewers expressed mixed opinions, with some leaning towards acceptance due to the novelty and potential of the work, especially in the context of SSMs, while others remained skeptical due to the perceived incremental innovation and the lack of broader architectural testing. The paper shows promise in advancing the field of parameter-efficient fine-tuning for SSMs but would benefit from further refinement, particularly in addressing the concerns about broader applicability and providing more robust comparisons with existing methods.

This paper received an average score that, while close to the acceptance threshold, is not competitive among this year's submissions. Given the balance of strengths and weaknesses, the final recommendation is to reject this submission in its current form.