🐍 The Hidden Attention of Mamba Models 🐍

Official PyTorch Implementation of "The Hidden Attention of Mamba Models"

The Mamba layer offers an efficient state space model (SSM) that is highly effective in modeling multiple domains including long-range sequences and images. SSMs are viewed as dual models, in which one trains in parallel on the entire sequence using convolutions, and deploys in an autoregressive manner. We add a third view and show that such models can be viewed as attention-driven models. This new perspective enables us to compare the underlying mechanisms to that of the self-attention layers in transformers and allows us to peer inside the inner workings of the Mamba model with explainability methods.
You can access the paper through : The Hidden Attention of Mamba Models

Set Up Environment

• Python 3.10.13

  • conda create -n your_env_name python=3.10.13

• Activate Env

  • conda activate your_env_name

• CUDA TOOLKIT 11.8

  • conda install nvidia/label/cuda-11.8.0::cuda-toolkit

• torch 2.1.1 + cu118

  • pip install torch==2.1.1 torchvision==0.16.1 torchaudio==2.1.1 --index-url https://download.pytorch.org/whl/cu118

• Requirements: vim_requirements.txt

  • pip install -r vim/vim_requirements.txt

• Install jupyter

  • pip install jupyter

• Install causal_conv1d and mamba from our source

  • cd causal-conv1d
  • pip install --editable .
  • cd ..
  • pip install --editable mamba-1p1p1

Pre-Trained Weights

We have used the official weights provided by Vim, which can be downloaded from here:

Notes:

• In all of our experiments, we have worked with Vim-small.

Vision-Mamba Explainability Notebook:

Follow the instructions in vim/vmamba_xai.ipynb notebook, in order to apply a single-image inference for the 3 introduced methods in the paper.

To-Do

For the segmentation experiment, please check out our follow-up work.

• XAI - Single Image Inference Notebook
• XAI - Segmentation Experimnts

Citation

if you find our work useful, please consider citing us:

@misc{ali2024hidden,
      title={The Hidden Attention of Mamba Models}, 
      author={Ameen Ali and Itamar Zimerman and Lior Wolf},
      year={2024},
      eprint={2403.01590},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

Acknowledgement

This repository is heavily based on Vim, Mamba and Transformer-Explainability. Thanks for their wonderful works.