GroupMixFormer：推动视觉 Transformer 的新进展

GroupMixFormer: 突破传统自注意力的界限

近年来,Vision Transformer (ViT) 凭借其强大的长程依赖建模能力,在计算机视觉领域掀起了一场革命。然而,传统的多头自注意力机制(MHSA)仅能在单一粒度上捕捉 token 间的相关性,这在一定程度上限制了模型的表征能力。为了突破这一瓶颈,来自香港大学和香港中文大学的研究团队提出了一种创新的注意力机制 - Group-Mix Attention (GMA),并基于此开发出了强大的视觉骨干网络 GroupMixFormer。

Group-Mix Attention: 多粒度相关性建模的关键

Group-Mix Attention (GMA) 是 GroupMixFormer 的核心创新。与传统的自注意力不同,GMA 能够同时捕捉 token-to-token、token-to-group 和 group-to-group 三种不同粒度的相关性,从而大幅提升模型的表征能力。具体而言,GMA 的工作原理如下:

1. 将 Query、Key 和 Value 均匀分割成多个段。
2. 对分割后的段进行不同的组聚合操作,生成组代理(group proxies)。
3. 基于 token 和组代理的混合计算注意力图。
4. 利用生成的注意力图重新组合 Value 中的 token 和组。

通过这种方式,GMA 能够在不同的组大小下捕捉多粒度的相关性,从而使模型获得更全面、更丰富的特征表示。

GroupMixFormer: 强大而高效的视觉骨干网络

基于 GMA,研究团队设计了 GroupMixFormer 这一新型视觉骨干网络。与现有模型相比,GroupMixFormer 在参数量更少的情况下,在多个视觉任务上取得了优异的性能:

• 图像分类: GroupMixFormer-L (70.3M 参数,384^2 输入)在 ImageNet-1K 上达到了 86.2% 的 Top-1 准确率,无需额外数据。
• 语义分割: GroupMixFormer-B (45.8M 参数)在 ADE20K 数据集上达到了 51.2% 的 mIoU。

这些出色的结果充分证明了 GroupMixFormer 的强大性能和高效性。

GroupMixFormer 的实现与应用

为了方便研究人员和开发者使用 GroupMixFormer,项目团队在 GitHub 上开源了完整的代码实现。以下是使用 GroupMixFormer 的主要步骤:

环境配置

首先,需要创建并激活一个 conda 虚拟环境:

conda create -n groupmixformer python=3.8 -y
conda activate groupmixformer

然后安装 PyTorch、torchvision 等依赖库:

pip install torch==1.8.0+cu111 torchvision==0.9.0+cu111 -f https://download.pytorch.org/whl/torch_stable.html
pip install timm==0.4.12 tensorboardX six tensorboard ipdb yacs tqdm fvcore

数据准备

对于图像分类任务,需要下载并解压 ImageNet 数据集,目录结构如下:

│path/to/imagenet/
├──train/
│  ├── n01440764
│  │   ├── n01440764_10026.JPEG
│  │   ├── n01440764_10027.JPEG
│  │   ├── ......
│  ├── ......
├──val/
│  ├── n01440764
│  │   ├── ILSVRC2012_val_00000293.JPEG
│  │   ├── ILSVRC2012_val_00002138.JPEG
│  │   ├── ......
│  ├── ......

模型训练

以 GroupMixFormer-Small 为例,在单节点 8 GPU 上训练 300 轮的命令如下:

python3 -m torch.distributed.launch --nproc_per_node 8 --nnodes 1 --use_env train.py \
  --data-path <Your data path> \
  --batch-size 64 \
  --output <Your target output path> \
  --cfg ./configs/groupmixformer_small.yaml \
  --model-type groupmixformer \
  --model-file groupmixformer.py \
  --tag groupmixformer_small

模型评估

使用预训练权重对 GroupMixFormer-Small 进行评估的命令如下:

CUDA_VISIBLE_DEVICES=1 OMP_NUM_THREADS=1 python3 -m torch.distributed.launch --nproc_per_node 1 --nnodes 1 --use_env test.py \
  --data-path <Your data path> \
  --batch-size 64 \
  --output <Your target output path> \
  --cfg ./configs/groupmixformer_small.yaml \
  --model-type groupmixformer \
  --model-file groupmixformer.py \
  --tag groupmixformer_small

GroupMixFormer 模型族

研究团队提供了一系列预训练的 GroupMixFormer 模型,适用于不同的应用场景:

这些预训练模型可以直接用于图像分类任务,也可以作为其他下游任务(如目标检测、语义分割)的骨干网络。

GroupMixFormer 的影响与展望

GroupMixFormer 的提出为视觉 Transformer 的发展开辟了新的方向。通过引入多粒度的相关性建模,GroupMixFormer 不仅提升了模型性能,还为解释模型内部工作机制提供了新的视角。未来,Group-Mix Attention 机制有望在更多领域发挥作用,如自然语言处理、多模态学习等。

此外,GroupMixFormer 的高效性也为部署大规模视觉模型提供了新的可能。随着硬件能力的不断提升,我们可以期待看到基于 GroupMixFormer 的更大规模模型,以及在更多实际应用中的落地。

总的来说,GroupMixFormer 代表了视觉 Transformer 领域的一个重要进展,为提升计算机视觉系统的性能和效率提供了新的思路。我们期待看到更多基于 GroupMixFormer 的创新应用,以及它在推动人工智能技术发展中发挥的重要作用。