AiluRus: A Scalable ViT Framework for Dense Prediction

This repository provides an MMSegmentation implementation of AiluRus (Adaptive resolution with spatial-awaRe clustering), accepted at NeurIPS 2023.

Jin Li, Yaoming Wang, Xiaopeng Zhang, Bowen Shi, Dongsheng Jiang, Chenglin Li, Wenrui Dai, Hongkai Xiong, and Qi Tian.
AiluRus: A Scalable ViT Framework for Dense Prediction, NeurIPS 2023.

Overview

High-resolution dense prediction creates long ViT token sequences, making self-attention expensive. AiluRus observes that boundaries and detailed regions generally need more spatial resolution than homogeneous interiors. It therefore assigns resolution adaptively instead of retaining every image token throughout the network.

At an intermediate ViT layer, AiluRus:

1. estimates each token's local density using a spatially weighted feature distance;
2. ranks tokens using density-peak clustering (DPC) and selects representative tokens;
3. merges tokens assigned to the same representative, shortening the sequence processed by later transformer blocks; and
4. unfolds the reduced prediction to the original token grid before producing the final dense output.

The procedure is parameter-free and can be applied during inference or fine-tuning. The paper reports that Segmenter ViT-L can increase FPS by 48% without fine-tuning while maintaining performance. In its fine-tuning study, one setting reduces training time by 52%, increases FPS by 146% (2.46x), and loses 0.09 mIoU. These are paper results measured under the authors' hardware and configurations; throughput on other systems will differ.

Repository scope

This snapshot is based on MMSegmentation v1.2.1 and currently provides the Segmenter ViT-B/16 semantic-segmentation path on ADE20K. The paper additionally evaluates ViT-L, Cityscapes, Pascal Context, object detection, instance segmentation, classification, and video generation; their complete configurations and checkpoints are not included here.

The implementation includes spatial-aware DPC, representative-token averaging, and output-resolution recovery. The attention re-weighting formulation described in Section 3.2 of the paper is not exposed as a separate attention module in this snapshot. Accordingly, do not assume that every number in the paper is reproduced by the single configuration included here.

Installation

Requirements:

• Python 3.8+
• PyTorch compatible with your CUDA or CPU environment
• MMCV >=2.0.0rc4,<2.2.0
• MMEngine >=0.5.0,<1.0.0

Create an environment and install PyTorch first using the command appropriate for your platform from the PyTorch installation guide. Then:

python -m pip install -U openmim
mim install "mmcv>=2.0.0rc4,<2.2.0"
python -m pip install "mmengine>=0.5.0,<1.0.0"
python -m pip install -e .

Verify the installation:

python -c "import mmseg; print(mmseg.__version__)"

This code base tracks MMSegmentation 1.2.1 and intentionally retains its package version. Avoid installing a second mmsegmentation package in the same environment.

Data preparation

Download ADE20K and arrange it following the MMSegmentation convention:

data/ade/
├── ADEChallengeData2016/
│   ├── annotations/
│   │   ├── training/
│   │   └── validation/
│   └── images/
│       ├── training/
│       └── validation/
└── release_test/

See the MMSegmentation ADE20K documentation for the dataset terms and complete preparation instructions. Dataset files are not distributed by this repository.

Usage

The provided configuration is:

configs/segmenter/ailurus/
└── segmenter_ailurus_vit-b_mask_8xb1-160k_ade20k-512x512.py

The base Segmenter checkpoint is downloaded from OpenMMLab by the config when needed. AiluRus has no additional learnable parameters, so that checkpoint can also initialize the AiluRus model for direct evaluation or fine-tuning.

Train on one GPU:

python tools/train.py \
  configs/segmenter/ailurus/segmenter_ailurus_vit-b_mask_8xb1-160k_ade20k-512x512.py \
  --work-dir work_dirs/ailurus_vit_b

Train on eight GPUs:

bash tools/dist_train.sh \
  configs/segmenter/ailurus/segmenter_ailurus_vit-b_mask_8xb1-160k_ade20k-512x512.py \
  8 \
  --work-dir work_dirs/ailurus_vit_b

Evaluate a checkpoint:

python tools/test.py \
  configs/segmenter/ailurus/segmenter_ailurus_vit-b_mask_8xb1-160k_ade20k-512x512.py \
  work_dirs/ailurus_vit_b/iter_160000.pth

Measure throughput:

python tools/analysis_tools/benchmark.py \
  configs/segmenter/ailurus/segmenter_ailurus_vit-b_mask_8xb1-160k_ade20k-512x512.py \
  work_dirs/ailurus_vit_b/iter_160000.pth \
  --repeat-times 3

Use the same input size, batch size, warm-up procedure, precision, and hardware when comparing throughput. FPS values from different environments are not directly comparable.

AiluRus parameters

The backbone configuration accepts the following AiluRus-specific fields:

Field	Paper symbol	Meaning	Included default
start_loc	location	Zero-based transformer-layer index after which clustering occurs	2
start_num	M	Number of representative tokens retained; currently must be a perfect square	484
num_neighbours	lambda	Number of spatial neighbors used by the spatial constraint	50
dpc_k	k	Nearest feature neighbors used to estimate local density	1
spatial_alpha	alpha	Spatial-weight strength in the range [0, 1]	0.9

Earlier clustering and fewer representatives usually improve throughput but can reduce accuracy. The paper's ablation on ADE20K reports alpha=0.9, lambda=50, and k=1 as its best tested values. start_num must not exceed the number of patch tokens produced by the configured input size.

Example override without editing the config:

python tools/test.py CONFIG CHECKPOINT \
  --cfg-options model.backbone.start_loc=2 model.backbone.start_num=400

Testing

Install the test dependencies and run the AiluRus tests:

python -m pip install -r requirements/tests.txt
python -m pytest tests/test_models/test_backbones/test_ailurus.py -q

Run the broader MMSegmentation test suite before publishing a release:

python -m pytest tests -q

Acknowledgements and license

This implementation is built on MMSegmentation and uses the Segmenter model. Please follow the licenses and citation requirements of those projects and of the datasets/checkpoints you use.

The source code in this repository is distributed under the Apache License 2.0. See NOTICE for attribution and modification information. The model weights and datasets referenced by the configuration are separate artifacts and may have different terms.

Citation

If this repository is useful in your research, please cite:

@inproceedings{li2023ailurus,
  title={AiluRus: A Scalable ViT Framework for Dense Prediction},
  author={Li, Jin and Wang, Yaoming and Zhang, Xiaopeng and Shi, Bowen and
          Jiang, Dongsheng and Li, Chenglin and Dai, Wenrui and Xiong, Hongkai
          and Tian, Qi},
  booktitle={Advances in Neural Information Processing Systems},
  volume={36},
  year={2023}
}

Please also cite MMSegmentation and Segmenter when appropriate.