If you have any questions, please don't hesitate to contact me ([email protected]).

You can also refer to this github repository: QuadMamba.

• 2024.10 QuadMamba is accepted by NeurIPS.
• 2024.11 Release the raw code of QuadMamba.
• 2025.06 Fixed the bugs in the code.
• 2025.06 Improve the QuadMamba scheme.

Quadtree-based selective scan with prediction modules. Image tokens are partitioned into bi-level window quadrants from coarse to fine. A fully differentiable partition mask is then applied to generate the 1D sequence with negligible computational overhead.

QuadMamba is based on the old version of VMamba; The installation can be referred to VMamba.

Environment Setup:

VMamba recommends setting up a conda environment and installing dependencies via pip. Use the following commands to set up your environment: Also, We recommend using the pytorch>=2.0, cuda>=11.8. But lower version of pytorch and CUDA are also supported.

Create and activate a new conda environment

conda create -n quadmamba
conda activate quadmamba

Install Dependencies

pip install -r requirements.txt
cd kernels/selective_scan && pip install .

Dependencies for Detection and Segmentation (optional)

pip install mmengine==0.10.1 mmcv==2.1.0 opencv-python-headless ftfy regex
pip install mmdet==3.3.0 mmsegmentation==1.2.2 mmpretrain==1.2.0

Classification

To train QuadMamba models for classification on ImageNet, use the following commands for different configurations:

python -m torch.distributed.launch --nnodes=1 --node_rank=0 --nproc_per_node=8 --master_addr="127.0.0.1" --master_port=29501 main.py --cfg </path/to/config> --batch-size 128 --data-path </path/of/dataset> --output /tmp

If you only want to test the performance (together with params and flops):

python -m torch.distributed.launch --nnodes=1 --node_rank=0 --nproc_per_node=1 --master_addr="127.0.0.1" --master_port=29501 main.py --cfg </path/to/config> --batch-size 128 --data-path </path/of/dataset> --output /tmp --pretrained </path/of/checkpoint>

To train with mmdetection or mmsegmentation:

bash ./tools/dist_train.sh </path/to/config> 8

If you find it useful, please consider to cite the paper, thanks!

@article{xie2024quadmamba,
    title={QuadMamba: Learning Quadtree-based Selective Scan for Visual State Space Model},
    author={Xie, Fei and Zhang, Weijia and Wang, Zhongdao and Ma, Chao},
    journal={Advances in Neural Information Processing Systems},
    year={2024}
}

This project is based on VMamba, LocalMamba, and DynamicViT. Thanks for their great work!