让视觉Transformer根据输入难易动态调整计算量,提升效率与精度。
ThinkingViT: Matryoshka Thinking Vision Transformer for Elastic Inference
- 通过分阶段激活注意力头实现动态推理,输入越复杂计算越多。
- 在ImageNet-1K上相同算力下准确率提升2.0个百分点,算力相同时提升2.9个百分点。
- 支持即插即用升级现有ViT模型,适用于分割、Swin等架构。
ViT虽性能优异,但固定计算预算限制了其在异构硬件上的可扩展部署。现有嵌套式Transformer通过在单模型中嵌入子网络实现可扩展推理,但对所有输入分配相同算力,造成效率损失。为此,本文提出ThinkingViT,一种采用渐进式思考阶段的嵌套视觉Transformer,能根据输入难度动态调整推理计算。模型首先激活少量关键注意力头进行初步预测,若置信度达标则提前终止;否则在同一带宽内逐步激活更多注意力头并重新前向传播,直至达到预设置信度或耗尽最大容量。为提升后续阶段性能,引入令牌复用机制,融合前一阶段的嵌入表示。实验表明,ThinkingViT在ImageNet-1K上,相同吞吐量下准确率比基线最高提升2.0百分点,相同GMACs下提升2.9百分点。其保持骨干结构的设计使其可作为下游任务(如语义分割)的即插即用升级模块,并有效迁移至Swin Transformer等其他架构。源代码已开源。
原文摘要 · Abstract (English)
ViTs deliver SOTA performance, yet their fixed computational budget prevents scalable deployment across heterogeneous hardware. Recent Matryoshka-style Transformer architectures mitigate this by embedding nested subnetworks within a single model to enable scalable inference. However, these models allocate the same amount of compute to all inputs, regardless of their complexity, which leads to inefficiencies. To address this, we introduce ThinkingViT, a nested ViT architecture that employs progressive thinking stages to dynamically adjust inference computation based on input difficulty. ThinkingViT first activates a small subset of the most important attention heads to produce an initial prediction. If the prediction confidence exceeds a predefined threshold, inference terminates early. Otherwise, within the same backbone, it activates a larger subset of attention heads and conducts a new forward pass. This process continues iteratively until the model reaches the predefined confidence level or exhausts its maximum capacity. To boost the performance of subsequent rounds, we introduce a Token Recycling approach that fuses the input embeddings with the embeddings from the previous stage. Experiments show that ThinkingViT surpasses nested baselines by up to 2.0 percentage points (p.p.) in accuracy at the same throughput and by up to 2.9 p.p. at equal GMACs on ImageNet-1K. We show that the backbone-preserving design of ThinkingViT allows it to serve as a plug-in upgrade for ViTs in downstream tasks such as semantic segmentation. We also demonstrate that ThinkingViT transfers effectively to other architectures such as Swin Transformers. The source code is available at https://github.com/ds-kiel/ThinkingViT.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。