JetViT: Efficient High-Resolution Vision Transformer with Post-Training Attention Search

📄 arXiv: 2605.26636v1 📥 PDF

作者: Dongyun Zou, Zhuoyang Zhang, Junyu Chen, Wenkun He, Qinhe Peng, Hanrong Ye, Yao Lu, Hongxu Yin, Yu Wang, Song Han, Han Cai

分类: cs.CV, cs.AI

发布日期: 2026-05-26

备注: Accepted to CVPR 2026 Findings

💡 一句话要点

JetViT：通过后训练注意力搜索实现高效高分辨率视觉Transformer

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱三：空间感知与语义 (Perception & Semantics) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 视觉Transformer 高分辨率图像 模型加速 后训练 注意力搜索

📋 核心要点

  1. 现有ViT模型在高分辨率图像上计算复杂度高，推理效率低，难以满足实际应用需求。
  2. JetViT提出后训练注意力搜索框架，通过将冗余全注意力块替换为更高效的线性或窗口注意力块，实现模型加速。
  3. 实验表明，JetViT在保持精度的同时，显著提高了高分辨率图像上的推理吞吐量，降低了延迟。

📝 摘要（中文）

本文介绍了一种新型混合架构的视觉Transformer (ViT) 模型家族JetViT，它在保持与最先进的全注意力视觉基础模型相当的精度的同时，显著提高了高分辨率图像上的推理效率。我们的方法核心是后训练注意力搜索，这是一种后训练加速框架，通过识别和替换冗余的全注意力块为线性或窗口注意力块，将预训练的全注意力ViT转换为高效的混合注意力变体。通过继承基础模型的MLP和注意力权重，后训练注意力搜索通过三个关键步骤有效地探索架构设计空间：(1) 优化线性注意力块设计；(2) 找到线性注意力和窗口注意力块的最佳组合；(3) 识别并保留关键的全注意力块。我们在两个具有代表性的高分辨率视觉基础模型DINOv3和DepthAnythingV2上评估了JetViT。在NVIDIA H100 GPU上，JetViT在不牺牲精度的情况下，实现了高达1.79倍的吞吐量和高达44.81%的延迟降低。我们将很快发布我们的代码和加速的ViT模型。

🔬 方法详解

问题定义：现有Vision Transformer (ViT) 模型在处理高分辨率图像时，由于全注意力机制的计算复杂度与图像尺寸呈平方关系，导致计算成本过高，推理速度慢，难以部署到资源受限的设备上。因此，如何降低ViT模型在高分辨率图像上的计算复杂度，提高推理效率，是本文要解决的关键问题。

核心思路：JetViT的核心思路是利用后训练注意力搜索（Post-Training Attention Search）框架，在预训练好的全注意力ViT模型的基础上，通过识别并替换冗余的全注意力块为计算量更小的线性或窗口注意力块，从而在不损失精度的前提下，显著降低模型的计算复杂度，提高推理效率。这种方法避免了从头训练的巨大开销，并且可以有效利用预训练模型的知识。

技术框架：JetViT的整体框架包括以下几个主要步骤：1) 线性注意力块设计：优化线性注意力块的结构和参数，使其在计算效率和表示能力之间取得平衡。2) 混合注意力块搜索：通过搜索算法，找到线性注意力和窗口注意力块的最佳组合，以最大程度地降低计算复杂度，同时保持模型的精度。3) 关键全注意力块保留：识别并保留对模型性能至关重要的全注意力块，以确保模型在加速的同时不会损失关键信息。整个过程在预训练模型的基础上进行，无需重新训练。

关键创新：JetViT的关键创新在于提出了后训练注意力搜索框架，该框架能够在预训练的全注意力ViT模型上，自动搜索并替换冗余的注意力块，从而在不损失精度的前提下，显著提高模型的推理效率。与以往的注意力机制加速方法相比，JetViT无需重新训练模型，可以有效利用预训练模型的知识，并且能够自动找到最佳的注意力块组合。

关键设计：JetViT的关键设计包括：1) 线性注意力块的结构：采用了高效的线性注意力机制，例如线性Transformer或Performer，以降低计算复杂度。2) 注意力块搜索算法：使用基于梯度或进化算法的搜索策略，自动找到最佳的注意力块组合。3) 全注意力块重要性评估：通过计算注意力权重的信息熵或梯度等指标，评估全注意力块的重要性，并保留重要的块。4) 损失函数：在搜索过程中，使用精度和计算复杂度作为优化目标，以平衡模型的性能和效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

JetViT在NVIDIA H100 GPU上进行了评估，结果表明，在DINOv3和DepthAnythingV2等高分辨率视觉基础模型上，JetViT在不牺牲精度的情况下，实现了高达1.79倍的吞吐量提升和高达44.81%的延迟降低。这些结果表明，JetViT是一种高效的ViT模型加速方法，能够显著提高高分辨率图像上的推理效率。

🎯 应用场景

JetViT具有广泛的应用前景，尤其是在需要处理高分辨率图像的场景中，例如自动驾驶、医学图像分析、遥感图像处理等。通过提高推理效率，JetViT可以降低计算成本，并使得ViT模型能够在资源受限的设备上部署，从而推动视觉Transformer在实际应用中的普及。

📄 摘要（原文）

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.