This repository provides the official implementation for Explaining Modern Gated-Linear RNNs via A Unified Implicit Attention Formulation.

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

Usage:

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

• RWKV Notebook , GriffinInstall
• Griffin Notebook, GriffinInstall
• Mamba Notebook, MambaNLPInstall

Vision

• For the segmentation experiemt:
  1. Download the data gtsegs_ijcv.mat and put it under '/UnifiedImplicitAttnRepr/MambaVision'
  2. Configure Checkpoint path at line 466 in imagenet_seg.eval.py (You can download the checkpoint from vim_s_midclstok_80p5acc.pth and put it under '/UnifiedImplicitAttnRepr/MambaVision')
  3. run 'python ./UnifiedImplicitAttnRepr/MambaVision/imagenet_seg_eval.py'
• For Heatmap Extraction follow the notebook in './UnifiedImplicitAttnRepr/MambaVision/Inference.ipynb'

Citation

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

@inproceedings{
zimerman2025explaining,
title={Explaining Modern Gated-Linear {RNN}s via a Unified Implicit Attention Formulation},
author={Itamar Zimerman and Ameen Ali Ali and Lior Wolf},
booktitle={The Thirteenth International Conference on Learning Representations},
year={2025},
url={https://openreview.net/forum?id=wnT8bfJCDx}
}

Acknowledgement:

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