无 Softmax 的线性 Transformer

SOFT:线性复杂度的无 Softmax Transformer,
Jiachen Lu, Jinghan Yao, Junge Zhang, Xiatian Zhu, Hang Xu, Weiguo Gao, Chunjing Xu, Tao Xiang, Li Zhang
NeurIPS 2021

无 Softmax 的线性 Transformer,
Jiachen Lu, Junge Zhang, Xiatian Zhu, Jianfeng Feng, Tao Xiang, Li Zhang
IJCV 2024

新特性

我们提出了一种具有更强泛化能力的归一化无 Softmax 自注意力机制。
SOFT 现在可应用于更多视觉任务（目标检测和语义分割）。

新闻

[2024/02/12] 我们的期刊扩展版无 Softmax 的线性 Transformer被 IJCV 接收。
[2022/07/05] SOFT 现已可用于下游任务！SOFT 应用了一种高效的归一化方法。请参考 SOFT-Norm

环境要求

timm==0.3.2
torch>=1.7.0 及匹配的 torchvision
cuda>=10.2

在 cuda < 10.2 上编译可能会失败。
我们已在 cuda 10.2 和 cuda 11.2 上成功编译。

数据准备

从 http://image-net.org/ 下载并解压 ImageNet 训练集和验证集图像。目录结构是 torchvision datasets.ImageFolder 的标准布局，训练数据和验证数据分别位于 train/ 文件夹和 val 文件夹中：

/path/to/imagenet/
  train/
    class1/
      img1.jpeg
    class2/
      img2.jpeg
  val/
    class1/
      img3.jpeg
    class/2
      img4.jpeg

安装

git clone https://github.com/fudan-zvg/SOFT.git
python -m pip install -e SOFT

主要结果

ImageNet-1K 图像分类

模型	分辨率	参数量	FLOPs	Top-1 %	配置	预训练模型
SOFT-Tiny	224	13M	1.9G	79.3	SOFT_Tiny.yaml, SOFT_Tiny_cuda.yaml	SOFT_Tiny, SOFT_Tiny_cuda
SOFT-Small	224	24M	3.3G	82.2	SOFT_Small.yaml, SOFT_Small_cuda.yaml
SOFT-Medium	224	45M	7.2G	82.9	SOFT_Meidum.yaml, SOFT_Meidum_cuda.yaml
SOFT-Large	224	64M	11.0G	83.1	SOFT_Large.yaml, SOFT_Large_cuda.yaml
SOFT-Huge	224	87M	16.3G	83.3	SOFT_Huge.yaml, SOFT_Huge_cuda.yaml
SOFT-Tiny-Norm	224	13M	1.9G	79.4	SOFT_Tiny_norm.yaml	SOFT_Tiny_norm
SOFT-Small-Norm	224	24M	3.3G	82.4	SOFT_Small_norm.yaml	SOFT_Small_norm
SOFT-Medium-Norm	224	45M	7.2G	83.1	SOFT_Meidum_norm.yaml	SOFT_Medium_norm
SOFT-Large-Norm	224	64M	11.0G	83.3	SOFT_Large_norm.yaml	SOFT_Large_norm
SOFT-Huge-Norm	224	87M	16.3G	83.4	SOFT_Huge_norm.yaml

COCO目标检测（2017验证集）

骨干网络	方法	学习率调度	边界框mAP	掩码mAP	参数量
SOFT-Tiny-Norm	RetinaNet	1x	40.0	-	23M
SOFT-Tiny-Norm	Mask R-CNN	1x	41.2	38.2	33M
SOFT-Small-Norm	RetinaNet	1x	42.8	-	34M
SOFT-Small-Norm	Mask R-CNN	1x	43.8	40.1	44M
SOFT-Medium-Norm	RetinaNet	1x	44.3	-	55M
SOFT-Medium-Norm	Mask R-CNN	1x	46.6	42.0	65M
SOFT-Large-Norm	RetinaNet	1x	45.3	-	74M
SOFT-Large-Norm	Mask R-CNN	1x	47.0	42.2	84M

ADE20K语义分割（验证集）

骨干网络	方法	裁剪大小	学习率调度	mIoU	参数量
SOFT-Small-Norm	UperNet	512x512	1x	46.2	54M
SOFT-Medium-Norm	UperNet	512x512	1x	48.0	76M

入门指南

训练

我们有两种高斯核的实现：PyTorch版本和使用cuda实现的精确高斯函数形式。包含cuda的配置文件是cuda实现。两种实现都能达到相同的性能。在运行cuda版本之前，请先安装SOFT。

./dist_train.sh ${GPU数量} --data ${数据路径} --config ${配置文件}
# 例如，使用8个GPU在ImageNet训练集上训练SOFT-Tiny
./dist_train.sh 8 --data ${数据路径} --config config/SOFT_Tiny.yaml

测试

./dist_train.sh ${GPU数量} --data ${数据路径} --config ${配置文件} --eval_checkpoint ${检查点文件} --eval

# 例如，使用8个GPU在ImageNet验证集上测试SOFT-Tiny

./dist_train.sh 8 --data ${数据路径} --config config/SOFT_Tiny.yaml --eval_checkpoint ${检查点文件} --eval

参考文献

@inproceedings{SOFT,
    title={SOFT: Softmax-free Transformer with Linear Complexity}, 
    author={Lu, Jiachen and Yao, Jinghan and Zhang, Junge and Zhu, Xiatian and Xu, Hang and Gao, Weiguo and Xu, Chunjing and Xiang, Tao and Zhang, Li},
    booktitle={NeurIPS},
    year={2021}
}

@article{Softmax,
    title={Softmax-free Linear Transformers}, 
    author={Lu, Jiachen and Zhang, Li and Zhang, Junge and Zhu, Xiatian and Feng, Jianfeng and Xiang, Tao},
    journal={International Journal of Computer Vision},
    year={2024}
}