MambaLRP: Explaining Selective State Space Sequence Models
Authors: Farnoush Rezaei Jafari, Grégoire Montavon, Klaus-Robert Müller, Oliver Eberle
NeurIPS 2024 | Conference PDF | Archive PDF | Plain Text | LLM Run Details
| Reproducibility Variable | Result | LLM Response |
|---|---|---|
| Research Type | Experimental | To evaluate our proposed approach, we benchmark its effectiveness against various methods previously proposed in the literature for interpreting neural networks. We empirically evaluate our proposed methodology using Mamba-130M, Mamba-1.4B, and Mamba-2.8B language models [33], which are trained on diverse text datasets. The training details can be found in Appendix C.1. For the vision experiments, we use the Vim-S model [83]. Moreover, we perform several ablation studies to further investigate our proposed method. |
| Researcher Affiliation | Collaboration | 1Machine Learning Group, Technische Universit at Berlin, 10587 Berlin, Germany 2BIFOLD Berlin Institute for the Foundations of Learning and Data, 10587 Berlin, Germany 3Department of Mathematics and Computer Science, Freie Universit at Berlin, Arnimallee 14, 14195 Berlin, Germany 4Department of Artificial Intelligence, Korea University, Seoul 136-713, South Korea 5Max Planck Institute for Informatics, Stuhlsatzenhausweg 4, 66123 Saarbr ucken, Germany 6Google Deep Mind, Berlin, Germany |
| Pseudocode | Yes | Algorithm 1: Mamba LRP in Si LU activation layer" and "Algorithm 2: Mamba LRP in Mamba block" (Appendix C.2) |
| Open Source Code | Yes | Our code is publicly available.1 https://github.com/Farnoush RJ/Mamba LRP |
| Open Datasets | Yes | Datasets In this study, we perform experiments on four text classification datasets, namely SST-2 [67], Medical BIOS [28], Emotion [60], and SNLI [16]. The SST-2 dataset encompasses around 70K English movie reviews, categorized into binary classes, representing positive and negative sentiments. The Medical BIOS dataset consists of short biographies (10K) with five specific medical occupations as targets. The SNLI corpus (version 1.0) comprises 570k English sentence pairs, with the labels entailment, contradiction, and neutral, used for the natural language inference (NLI) task. The Emotion dataset (20K) is a collection of English tweets, each labeled with one of six basic emotions. For the vision experiments, we use Image Net dataset [23] with 1.3M images and 1K classes. |
| Dataset Splits | Yes | Table 5: Data statistics. Dataset Train Test Validation SST-2 68K 2K 1K Med-BIOS 8K 1K 1K Emotion 16K 2K 2K SNLI 550K 10K 10K Image Net 1.3M 50K 100K" and "We employed early stopping and ended training as soon as the validation loss ceased to improve. |
| Hardware Specification | Yes | All baseline methods are evaluated on a single A100-40GB GPU with a batch size of 1. |
| Software Dependencies | No | The paper mentions software like PyTorch and the Captum library, and specific optimizers (Adam W), but does not provide version numbers for these software dependencies, which are necessary for full reproducibility. |
| Experiment Setup | Yes | Training details During training, we used a batch size of 32. To train the Mamba-1.4B and Mamba2.8B models on the SNLI dataset, a batch size of 64 is used. We employed the {Eleuther AI/gpt-neox20b}4 tokenizer. The models parameters were optimized using Adam W optimizer with a learning rate set at 7e 5. Additionally, we used a linear learning rate scheduler with an initial factor of 0.5. All models were trained for a maximum of 10 epochs. We employed early stopping and ended training as soon as the validation loss ceased to improve. |