JetViT通过后训练注意力搜索，让高分辨率ViT模型提速近1.8倍且不丢精度。
JetViT: Efficient High-Resolution Vision Transformer with Post-Training Attention Search
- 后训练阶段用线性/窗口注意力替换冗余全注意力块，实现高效架构优化
- 在H100上达到1.79倍吞吐量提升，延迟降低44.81%，精度不变
- 适合需要高分辨率图像推理效率的视觉任务，如医学影像、遥感分析
我们提出JetViT，一种新型混合架构视觉Transformer系列，在保持最先进的全注意力视觉基础模型精度的同时，显著提升高分辨率图像的推理效率。核心是后训练注意力搜索框架，通过识别并替换冗余的全注意力模块为线性或窗口注意力模块，将预训练的全注意力ViT转换为高效的混合注意力变体。该方法继承基础模型的MLP和注意力权重，通过三个关键步骤高效探索架构空间：(1) 优化线性注意力模块设计；(2) 找到最优的线性与窗口注意力模块组合；(3) 识别并保留关键全注意力模块。我们在两个代表性高分辨率视觉基础模型DINOv3和DepthAnythingV2上进行评估。在NVIDIA H100 GPU上，JetViT实现最高1.79倍的吞吐量提升和最高44.81%的延迟降低，且精度无损失。代码和加速后的ViT模型即将开源。
原文摘要 · Abstract (English)
We introduce JetViT, a novel family of hybrid-architecture Vision Transformer (ViT) models that match the accuracy of state-of-the-art full-attention vision foundation models while achieving substantially higher inference efficiency on high-resolution images. At the core of our approach is Post-Training Attention Search, a post-training acceleration framework that converts pre-trained full-attention ViTs into efficient hybrid-attention variants by identifying and replacing redundant full-attention blocks with linear or window-attention blocks. By inheriting the MLP and attention weights from the base model, Post-Training Attention Search efficiently explores the architectural design space through three key steps: (1) optimizing the linear-attention block design; (2) finding the best combination of linear-attention and window-attention blocks; and (3) identifying and preserving critical full-attention blocks. We evaluate JetViT on two representative high-resolution vision foundation models, DINOv3 and DepthAnythingV2. On the NVIDIA H100 GPU, JetViT achieves up to 1.79x higher throughput and up to 44.81% lower latency without sacrificing accuracy. We will release our code and accelerated ViT models soon.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。