JetViT: 面向高效高分辨率视觉Transformer的后训练注意力搜索
计算机视觉与模式识别
2026-05-27 v1 人工智能
摘要
我们提出了JetViT，一种新型的混合架构视觉Transformer（ViT）模型家族，它在匹配最先进的全注意力视觉基础模型精度的同时，在高分辨率图像上实现了显著更高的推理效率。我们方法的核心是后训练注意力搜索，这是一种后训练加速框架，通过识别并用线性注意力或窗口注意力模块替换冗余的全注意力模块，将预训练的全注意力ViT转换为高效的混合注意力变体。通过继承基础模型的MLP和注意力权重，后训练注意力搜索通过三个关键步骤高效地探索架构设计空间：(1) 优化线性注意力模块设计；(2) 寻找线性注意力和窗口注意力模块的最佳组合；(3) 识别并保留关键的全注意力模块。我们在两个具有代表性的高分辨率视觉基础模型DINOv3和DepthAnythingV2上评估了JetViT。在NVIDIA H100 GPU上，JetViT在不牺牲精度的情况下实现了高达1.79倍的吞吐量提升和高达44.81%的延迟降低。我们将很快发布我们的代码和加速后的ViT模型。
引用
@article{arxiv.2605.26636,
title = {JetViT: Efficient High-Resolution Vision Transformer with Post-Training Attention Search},
author = {Dongyun Zou and Zhuoyang Zhang and Junyu Chen and Wenkun He and Qinhe Peng and Hanrong Ye and Yao Lu and Hongxu Yin and Yu Wang and Song Han and Han Cai},
journal= {arXiv preprint arXiv:2605.26636},
year = {2026}
}
备注
Accepted to CVPR 2026 Findings