A Unified Implicit Attention Formulation for Gated-Linear Recurrent Sequence Models

Itamar Zimerman¹ *, Ameen Ali¹ * and Lior Wolf¹
itamarzimm@gmail.com, ameenali023@gmail.com, liorwolf@gmail.com
¹ Tel Aviv University, (*) equal contribution

This repository provides the official implementation for A Unified Implicit Attention Formulation for Gated-Linear Recurrent Sequence.

The purpose of this repository is to provide tools for the explainability and interpretability of modern sub-quadratic architectures, based on implicit attention representation.

Supported Models:

• Mamba
• Griffin
• RWKV
• Vision Mamba

Usage:

We provide the following Jupyter notebooks ('I' denotes installation instructions.):

• RWKV Notebook , I
• Griffin Notebook, I
• Mamba Notebook, I
• Vision Mamba (Coming Soon!)
• • Heatmaps Extraction (Coming Soon!)
  • Segmentation (Coming Soon!)

Citation

If you use this codebase, or otherwise found our work valuable, please cite:

@misc{zimerman2024unified,
      title={A Unified Implicit Attention Formulation for Gated-Linear Recurrent Sequence Models}, 
      author={Itamar Zimerman and Ameen Ali and Lior Wolf},
      year={2024},
      eprint={2405.16504},
      archivePrefix={arXiv},
      primaryClass={cs.LG}
}

Acknowledgement:

This repository is heavily based on Transformers and Mamba. Thanks for their wonderful works.