高效模型与推理加速 · CVPR2026 ·2026.06.03

ThinkingViT:面向弹性推理的套娃式思考视觉 Transformer

ThinkingViT: Matryoshka Thinking Vision Transformer for Elastic Inference

Ali Hojjat、Janek Haberer、Sören Pirk、Olaf Landsiedel

与你有关的部分

把 Matryoshka 式嵌套 ViT 与按输入难度自适应的渐进式推理结合:先用少量注意力头出预测,置信度不足再在同一骨干内激活更多头重跑,并用词元回收融合上一阶段嵌入,可直接作为现有 ViT 骨干的即插即用加速方案。

原文依据 ThinkingViT first activates a small subset of the most important attention heads to produce an initial prediction.
可迁移到 图像分类语义分割弹性推理

摘要译文

由模型翻译 · 术语统一 · 请对照原文核对关键结论

ViT 虽能取得 SOTA 性能,但其固定的计算预算使其难以在异构硬件上可扩展地部署。近期的套娃式 Transformer 架构通过在单一模型内嵌入嵌套子网络来缓解这一问题,从而实现可扩展推理。然而,这类模型对所有输入分配相同的计算量,不考虑输入的复杂度,从而导致效率低下。为解决这一问题,我们提出 ThinkingViT,一种采用渐进式思考阶段的嵌套 ViT 架构,可根据输入难度动态调整推理计算量。ThinkingViT 首先激活一小部分最重要的注意力头以产生初始预测。若预测置信度超过预设阈值,推理即提前终止;否则,它在同一骨干内激活更大一部分注意力头并执行新的前向计算。该过程迭代进行,直至模型达到预设置信度或耗尽最大容量。为提升后续轮次的性能,我们提出词元回收方法,将输入嵌入与上一阶段的嵌入相融合。实验表明,在 ImageNet-1K 上,ThinkingViT 在相同吞吐量下比嵌套基线最高提升 2.0 个百分点(p.p.)的准确率,在相同 GMACs 下最高提升 2.9 个百分点。我们进一步表明,ThinkingViT 保持骨干不变的设计使其可作为 ViT 在下游任务(如语义分割)中的即插即用升级方案。我们还证明 ThinkingViT 能有效迁移到 Swin Transformer 等其他架构。

英文原文摘要 1658 字符

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.

标签

会议 CVPR2026
页码 41923-41933
Bibkey
来源 CVF Open Access