SPOT: Sparsification with Attention Dynamics via Token Relevance in Vision Transformers

This repository contains PyTorch implementation for SPOT.

SPOT is a SParsification with attentiOn dynamics via Token relevance framework for the early detection of redundant tokens within ViTs, leveraging comprehensive contextual information.
Our proposed method outperforms existing methods in identifying salient tokens, yielding a significant reduction in computational cost, while maintaining near-identical or even improving accuracy.

Examples

Models

We provide our SPOT models pretrained on ImageNet, allowing to achieve a range of performance levels adaptable to varying computational resource limitations:

Backbone model	URL	Performance
DeiT-T	Google Drive
DeiT-S	Google Drive
LV-ViT-T	Google Drive
LV-ViT-S	Google Drive

Usage

Requirements

A detailed list of required packages is provided in requirements.yml. You can install them using conda:

conda env create -f requirements.yml

Data preparation: download and extract ImageNet images from http://image-net.org/. The directory structure should be

│ILSVRC2012/
├──train/
│  ├── n01440764
│  │   ├── n01440764_10026.JPEG
│  │   ├── n01440764_10027.JPEG
│  │   ├── ......
│  ├── ......
├──val/
│  ├── n01440764
│  │   ├── ILSVRC2012_val_00000293.JPEG
│  │   ├── ILSVRC2012_val_00002138.JPEG
│  │   ├── ......
│  ├── ......

Model preparation: download pretrained models if necessary:

model	url
DeiT-T	link
DeiT-S	link
LV-ViT-T	link
LV-ViT-S	link

Evaluation

To evaluate a pretrained SPOT model on the ImageNet validation set with a single GPU, run: (You can fine-tune models with different computational budgets by adjusting base_rate. )

DeiT-T

python evaluate.py --data_path /path/ImageNet/ --model deit-t --models_path /path/fine_tuned_SPOT --base_rate desired_rho

DeiT-S

python evaluate.py --data_path /path/ImageNet/ --model deit-s --models_path /path/fine_tuned_SPOT --base_rate desired_rho

LV-ViT-T

python evaluate.py --data_path /path/ImageNet/ --model lvvit-t --models_path /path/fine_tuned_SPOT --base_rate desired_rho

LV-ViT-S

python evaluate.py --data_path /path/ImageNet/ --model lvvit-s --models_path /path/fine_tuned_SPOT --base_rate desired_rho

Fine-tuning

To fine-tune models on ImageNet using the SPOT framework, run: (You can fine-tune models with different computational budgets by adjusting base_rate. )

DeiT-T

torchrun --nproc_per_node=8 /root/main.py --data_path /path/ImageNet/ --output_dir /results/SPOT_deit_t --model deit-t --batch_size 64 --base_rate desired_rho

DeiT-S

torchrun --nproc_per_node=8 /root/main.py --data_path /path/ImageNet/ --output_dir /results/SPOT_deit_s --model deit-s --batch_size 64 --base_rate desired_rho

LV-ViT-T

torchrun --nproc_per_node=8 /root/main.py --data_path /path/ImageNet/ --output_dir /results/SPOT_lvvit_t --model lvvit-t --batch_size 64 --base_rate desired_rho

LV-ViT-S

torchrun --nproc_per_node=8 /root/main.py --data_path /path/ImageNet/ --output_dir /results/SPOT_lvvit_s --model lvvit-s --batch_size 64 --base_rate desired_rho