Memba for Language Tasks

This directory includes the code for reproducing the language task results in our paper.

Setup

Install dependencies. Our code is tested on CUDA 11.7, PyTorch 2.0.0, and Python 3.8.

pip install torch==2.0.0 torchvision==0.15.1
pip install -r requirements.txt

Note: causal_conv1d and mamba_ssm require CUDA toolkit for compilation. If the build fails due to pip's build isolation installing a different PyTorch, use:

pip install causal_conv1d mamba_ssm --no-build-isolation

Datasets

Download the commonsense 170k fine-tuning dataset from LLM-Adapters, then place it as:

language/commonsense_reasoning/commonsense_170k.json

Fine-tuning

Full fine-tuning

python finetune.py \
    --base_model state-spaces/mamba-130m-hf \
    --data_path commonsense_170k.json \
    --adapter_name full \
    --output_dir ./results/full \
    --learning_rate 5e-5

Memba (LoRA on in+out proj with LIM)

# Mamba-130M
CUDA_VISIBLE_DEVICES=0 python finetune.py \
    --base_model state-spaces/mamba-130m-hf \
    --data_path commonsense_170k.json \
    --adapter_name lora_in_out_proj_lif \
    --output_dir ./results/memba_130m \
    --learning_rate 1e-3

# Mamba-790M
CUDA_VISIBLE_DEVICES=0 python finetune.py \
    --base_model state-spaces/mamba-790m-hf \
    --data_path commonsense_170k.json \
    --adapter_name lora_in_out_proj_lif \
    --output_dir ./results/memba_790m \
    --learning_rate 5e-4

# Mamba-1.4B
CUDA_VISIBLE_DEVICES=0 python finetune.py \
    --base_model state-spaces/mamba-1.4b-hf \
    --data_path commonsense_170k.json \
    --adapter_name lora_in_out_proj_lif \
    --output_dir ./results/memba_1.4b \
    --learning_rate 1e-4

Supported Memba adapter names:

• lora_in_out_proj_lif — LoRA on in_proj + out_proj with LIM (recommended, Table 2)
• lora_in_proj_lif — LoRA on in_proj with LIM
• lora_out_proj_lif — LoRA on out_proj with LIM

Evaluation

We use lm-evaluation-harness for evaluation.

TASKS='boolq,social_iqa,hellaswag,piqa,arc_easy,arc_challenge,winogrande,openbookqa'

# Zero-shot evaluation (without fine-tuning)
lm_eval --model hf \
    --model_args pretrained="state-spaces/mamba-130m-hf,trust_remote_code=True" \
    --tasks $TASKS

# Memba evaluation
python lm_harness_eval.py --model MambaPEFT \
    --model_args pretrained="state-spaces/mamba-130m-hf,peft_weights=./results/memba_130m,trust_remote_code=True" \
    --output_path results/memba_130m \
    --tasks $TASKS

Acknowledgements

Our code is based on Mamba, LLM-Adapters, PEFT, and Transformers.