点击下方名片，关注「集智书童」公众号
本文主要解决了什么问题
视觉Transformer（ViT）缺乏弹性推理能力 ：传统ViT在推理时使用固定的计算预算，无法根据输入图像的复杂度动态调整计算资源，限制了其在异构硬件平台上的高效部署。
现有嵌套Transformer模型的效率低下 ：尽管基于Matryoshka架构的嵌套模型实现了多子网络推理，但它们对所有输入分配相同的计算资源，未能根据输入复杂度进行自适应调整，造成资源浪费。
缺乏有效的图像级路由机制 ：现有路由机制要么计算开销大，要么表示能力不足，无法准确评估图像的全局复杂性，难以实现高效的输入感知推理。
本文的核心创新是什么
ThinkingViT架构设计 ：提出了一种基于ViT的嵌套架构，支持多阶段渐进式推理，通过逐步激活更多的注意力头实现动态计算调整。
Token回收机制 ：在不同推理阶段之间循环利用Token Embedding，使得后续推理建立在前一阶段的基础上，提升预测准确性，同时避免重复从头计算。
基于熵的“啊哈时刻”机制 ：通过预测输出的香农熵衡量模型确定性，决定是否提前终止推理，实现无需独立路由机制的输入自适应计算分配。
弹性推理能力 ：用户可通过调整熵阈值灵活控制模型在精度与计算效率之间的权衡，适用于不同部署场景。
结果相较于以前的方法有哪些提升
准确率提升 ：
- • 在 相同吞吐量 下，ThinkingViT比嵌套Baseline模型在ImageNet-1K上的准确率提升了 高达2.0个百分点 。
- • 在 相同GMACs （计算量）条件下，准确率提升 高达2.9个百分点 。
计算效率优化 ：通过熵控制机制，简单样本可以提前终止推理，减少不必要的计算开销，而复杂样本则分配更多计算资源，实现更精准预测。
鲁棒性增强 ：在ImageNet-V2、ImageNet-ReaL、ImageNet-R等分布偏移数据集上均优于现有嵌套模型，且参数量和计算量显著更低。
局限性总结
在均匀困难输入上性能受限 ：当输入样本普遍复杂时（如ImageNet-A），模型频繁进入更深的推理阶段，导致早期终止机制失效，效率优势下降。
依赖预测阶段的质量 ：早期预测的准确性和确定性评估直接影响整体效率，若初始阶段预测质量不高，可能影响后续推理路径的优化。
思考阶段和头数扩展的配置需调优 ：最佳的阶段数、注意力头扩展步长等超参数依赖具体任务与部署目标，需进行针对性调优。
导读
视觉Transformer（ViT）能够提供最先进的性能，但其固定的计算预算阻碍了其在异构硬件上的可扩展部署。基于Matryoshka的架构通过在单个模型中嵌入嵌套子网络来缓解这一问题，从而实现可扩展推理。然而，这些模型对所有输入分配相同量的计算资源，无论其复杂度如何，这导致了效率低下。为解决此问题，作者提出了ThinkingViT，一种嵌套ViT架构，该架构采用渐进式思考阶段，根据输入难度动态调整推理计算量。ThinkingViT通过激活最重要的注意力头的小子集来启动推理，如果预测达到足够的确定性，则提前终止。否则，它会激活额外的注意力头并重新评估输入。ThinkingViT的核心是作者的Token回收机制，该机制将后续推理阶段的条件建立在先前阶段的嵌入上，从而实现渐进式改进。由于其 Backbone 网络保留的设计，ThinkingViT也可作为普通ViT的插件升级。实验表明，在相同的吞吐量和相等的GMACs下，ThinkingViT在ImageNet-1K上的准确率比嵌套 Baseline 提高了高达2.0个百分点（p.p.），在GMACs相同时提高了高达2.9 p.p.。
1 引言
动机：视觉Transformer（ViT）（Dosovitskiy等人，2021）在众多图像识别任务中取得了最先进的结果（Liu等人，2022；Hatamizadeh和Kautz，2024），然而其非弹性推理流程导致计算成本统一且往往过高。这种缺乏灵活性在现实世界部署中构成了重大限制，因为设备在计算能力和延迟约束方面差异很大（Xu等人，2025）。从高通量服务器到移动和嵌入式平台，现代应用需要弹性推理（Cai等人，2024），即模型能够调整其计算占用以匹配目标硬件的能力。缺乏这种适应性，ViT难以满足可扩展和广泛部署所需的性能与效率权衡。
现有工作的局限性：基于Matryoshka（即嵌套）Transformer架构的最新进展为弹性推理提供了一条引人入胜的途径（Kusupati等人，2022；Haberer等人，2024；Devvrit等人，2024；Zhang等人，2024）。这些方法将多个嵌套子网络嵌入单个Transformer主干中，从而在推理时实现动态子网络选择。这种策略使模型能够在不同的延迟和硬件约束下运行，无需任何额外的调整或训练后修改，同时保持统一的参数集。然而，这些方法通常为每个输入分配固定的计算预算，错失了区分简单和复杂样本的机会，导致效率低下，尤其是在资源紧张的情况下。
基于图像的路由：为了使嵌套模型具有适应性，作者必须根据输入图像的难度将其路由到适当的子网络。虽然路由的概念在专家混合（MoE）模型（Zhou等人，2022）中出现过，但作者不能直接使用它们，因为这些路由器在通过轻量级多层感知器（MLP）引导 Token Embedding 时工作在 Token Level ，难以捕捉图像的全局复杂性。在实践中，准确估计图像难度需要与完整规模的分类器相当的表示能力，而这类紧凑的路由器无法满足这一要求（Ong等人，2025；Ding等人，2025）。此外，将大量计算资源分配给一个仅将输入转发到另一个模型的独立路由器会引入冗余。理想情况下，路由过程本身应贡献有用的表示，这些表示可以在 Pipeline 中循环使用以增强最终预测，而不是在做出决策后被丢弃。这引出了一个关键的设计问题：作者如何在嵌套的视觉Transformer（Vision Transformer）中实现输入感知的计算分配，而不依赖于一个独立的昂贵路由机制？
ThinkingViT：受基于思考的架构的最新进展（Guo等人，2025；Shao等人，2024）的启发，作者引入了ThinkingViT，这是一个基于基础ViT架构（Dosovitskiy等人，2021）的ViT，它通过嵌套的渐进式思考阶段根据输入复杂性动态调整其推理工作。ThinkingViT首先激活嵌套网络的一部分，包括最重要的注意力头（例如，50%）以生成初始预测以及一个确定性分数。如果确定性足够高（“啊哈”时刻（Guo等人，2025），推理将提前终止。否则，处理后的嵌入与原始输入嵌入融合，并通过嵌套网络的一个更大的子集，使用一个增加的注意力头集（例如，100%）执行更彻底的重新评估。这种 Token 融合机制使模型能够避免从头开始思考；相反，它循环利用在第一次通过中获得的知识，用更高的精度改进其预测；见图1。作者提出的渐进式思考机制使模型在处理模糊输入时不仅能够思考更多，而且能够更强大，确保更具挑战性的示例获得更大的表示能力，即更多的注意力头。在此 Pipeline 中，预测确定性作为控制弹性推理的核心机制：模型对高确定性案例提前终止，而对更不确定的输入逐步加深计算。通过调整确定性阈值，模型自然地在性能和效率之间取得平衡，支持弹性推理而无需任何独立的路由机制。此外，将所有阶段嵌套在单个统一的模型中避免了参数重复，并实现了更高效的训练和更高的精度。
此外，与语言模型中反复使用同一网络进行迭代优化的方法（Guo等人，2025年）不同，在视觉模型中，渐进式扩展是必要的，因为简单地将ViT自身的输出重新输入到同一网络的朴素迭代链无法实现持续改进，并且很快会达到饱和。在某些情况下，它们甚至会导致性能下降，如表1所示。深度扩展变体表现出类似的局限性。尽管参数数量大幅增加，但它们的性能也会达到平台期，并且在添加过多层时可能会下降（Shen等人，2022年）。这些发现突出了ThinkingViT采用的渐进式扩展策略的必要性，该策略显著提高了准确率。
贡献：
-
- 作者介绍了ThinkingViT，一种基于思考的视觉Transformer（ViT），它通过根据图像难度动态调整计算，为嵌套Transformer带来了输入适应性。
-
- ThinkingViT通过逐步激活更大数量的注意力头进行多轮推理，允许模型根据预测的确定性对易于输入的样本提前停止，同时为需要更丰富表示的复杂样本分配更大的计算能力。
-
- ThinkingViT引入了Token回收机制，将每一轮推理的条件建立在上一轮推理产生的特征之上，从而提升了整体准确率。
-
- ThinkingViT在相同吞吐量下，准确率比 Baseline 模型高2.0个百分点，在相同GMACs下，准确率比 Baseline 模型高2.9个百分点，参见图2。
2 相关工作
嵌套模型：超越利用神经架构搜索在预训练模型中识别最优子网络的方法（Cai等人，2019），研究界一直致力于开发嵌套架构。可缩减网络首次展示了单个模型可通过共享权重和宽度特定归一化在多个宽度下运行（Yu等人，2018；Yu和Huang，2019）。许多后续研究基于这一理念。例如，MatFormer（Devvrit等人，2024）基于Matryoshka表示学习（Kusupati等人，2022），在MLP的隐藏层中引入了多个嵌套子集。DynaBERT（Hou等人，2020）切分了多头注意力（MHA）和MLP层，但由于依赖知识蒸馏，并未跨层切分嵌入。SortedNet（Valipour等人，2023）将嵌套扩展至MLP、归一化层（NORM）、MHA和嵌入维度，尽管其保留了固定数量的注意力头。HydraViT（Haberer等人，2024）和Slicing ViT（Zhang等人，2024）支持跨嵌入、NORM、MLP和MHA的切分，并支持动态头数。然而，与其他方法类似，它对每个输入应用固定的计算预算，限制了其根据输入复杂度自适应的能力，并在资源受限时导致效率低下。
路由：路由机制在MoE框架（Zhou等人，2022）中变得流行，并随后扩展到嵌套设计，如MoNE（Jain等人，2024）、MoD（Raposo等人，2024）和AMoD（Gadhikar等人，2025）。Flextron引入了一个替代损失预测器来指导token路由（Cai等人，2024）。这些方法通常依赖于轻量级MLPs，这些MLPs缺乏在图像 Level 路由的表示能力，而图像 Level 的决策往往需要完整分类器的推理能力。Selective Query（Kag等人，2023）、OCCAM（Ding等人，2025）和RouteLLM（Ong等人，2025）等基于集成的方法使用完整模型作为门控，执行输入感知路由。然而，这些方法在预训练模型（例如CNN或Transformer）的固定集合上运行，并且无法在路由器和被路由模型之间实现任何知识迁移。相比之下，ThinkingViT在各个阶段之间循环token，允许每一轮在先前推理的基础上进行改进预测。
思考：大语言模型（LLM）的最新进展引入了自适应推理深度机制（El-Kishky等人，2025；Openai：Jaech等人，2024）。这些方法通常使用强化学习根据输入复杂度动态调整推理步数（Guo等人，2025；Shao等人，2024）。然而，过长的推理长度会增加推理成本，而准确率提升并不成比例（Kumar等人，2025）。受此启发，近期研究致力于在保留深度推理优势的同时避免冗余计算，通过引入测试时自适应或基于置信度的提前停止（Muennighoff等人，2025；Zeng等人，2025）。然而，这些模型在每一轮中重复使用相同的网络，而ThinkingViT在重新思考时提升模型容量，这对于视觉任务至关重要，见表1。
3 思维视觉Transformer
在本节中，作者介绍了ThinkingViT的设计组件。首先，作者描述了如何在基础ViT Backbone 网络内部构建有序子网络的层次结构。接下来，作者解释了递归推理循环以及ThinkingViT如何在连续轮次中融合知识。最后，作者展示了ThinkingViT如何通过基于熵的"aha!"标准实现弹性推理，该标准在达到足够确定性时停止思考。
3.1 嵌套视觉Transformer
ThinkingViT是在标准ViT架构的基础上构建的。设
为一个具有
个注意力头和
维嵌入维度的ViT模型。ViT首先将输入图像
分割为
个不重叠的块，每个块通过基于CNN的块嵌入函数投影到一个
维的嵌入向量
。在添加位置编码后，得到的序列由
个标准Transformer块进行处理：
诱导嵌套子网络：受Haberer等人（2024）的启发，ThinkingViT在此架构中诱导出
个嵌套子网络。每个子网络用
表示，并使用每个token的前
个嵌入值以及全模型的前
个注意力头构建。为构建这种嵌套结构，ThinkingViT根据这些索引对ViT的所有组件（包括嵌入层、注意力模块、MLP块和归一化层）进行切片。这生成了嵌入维度和注意力能力逐步增加的包含式子网络：
3.2 通过Token循环嵌套Vision Transformers
受推理模型近期进展的启发（Guo等人，2025；Shao等人，2024），ThinkingViT通过采用"Token回收"机制，执行多轮渐进式精炼。尽管在实践中作者发现两个子网络足以实现优异性能（如图3所示），但此处仍呈现通用多轮公式，以增强清晰度和更广泛的适用性。
在构建嵌套架构后，ThinkingViT首先使用最小子网络
，其嵌入维度为
，并具有
个注意力头。该子网络处理输入图像以预测类别并生成 Token Embedding
。随后，为了细化其预测，模型通过激活一个更大的子网络
来扩展计算能力，该子网络具有
个注意力头，其中
且
。之后，它通过一个投影层融合生成的 Token Embedding
与新的输入嵌入
，并通过可学习的参数
进行缩放，从而确定从上一步“回忆”的嵌入的权重：
这种Token回收机制允许模型重用先前的表征，而不是从头开始重新处理。该过程迭代进行，与用于推理的重复使用相同网络的语言模型（Guo等人，2025年）不同，ThinkingViT通过逐步激活更大的注意力头子集来增加模型容量，并在每一步增强表征能力。渐进式扩展在视觉任务中至关重要，因为在表1中所示，朴素递归会迅速达到平台期。附录D和附录C详细比较了回收和融合策略。
3.3 联合训练所有子网络
在训练过程中，ThinkingViT执行所有
个思考轮次，并在所有阶段最小化加权分类损失
。
其中
是真实标签，
控制每个子网络对全局目标的贡献。对于较少的子网络，可以同时优化所有子网络，因为梯度计算图仍然保持相对紧凑。然而，随着
的增加，联合训练所有子网络变得计算密集。在这种情况下，作者采用诸如三明治规则（
和 Huang, 2019）和随机子网络采样（Haberer 等人，2024）等策略来降低开销。
3.4 "啊哈时刻"
通过在模型内部引入嵌套迭代子网络，ThinkingViT构建了一个包含
个嵌套思考步骤的层次结构，其中每个步骤在之前步骤的基础上进行预测的细化。然而，不同输入具有不同的复杂度，对于某些输入，模型仅需
次迭代即可变得自信，其中
（例如，在单轮思考后）。当模型识别到已达到足够置信度时，便会触发"啊哈!"时刻，使其根据输入复杂度分配适当数量的思考轮次（Guo等人，2025）。令
表示在第
次迭代中产生的softmax输出。在
轮之后，模型使用香农熵来衡量其置信度：
其中
为类别数量。如果熵
低于预设阈值
，推理过程提前终止，模型接受当前预测。否则，ThinkingViT 激活更大的子网络以进一步优化结果。该简单指标在 ImageNet-1K 上表现与更复杂的标准相当（Jitkrittum 等人，2023）。作者的设计还允许替代路由模块（Kag 等人，2023；Ding 等人，2025）作为即插即用替换方案集成，无需更改架构或训练。
3.5 弹性推理
ThinkingViT通过其自身的确定性来决定何时停止，从而支持弹性推理。熵阈值
控制了效率与准确率之间的权衡。低阈值会执行更严格的置信度要求，导致大多数输入需要经过更多思考步骤，这会增加计算成本但提高准确率。相反，高阈值允许更多输入更早退出，减少延迟和计算成本，但可能以牺牲潜在准确率为代价。这种灵活性使用户能够根据特定的资源与性能需求来调整模型。
3.6 范味 ViT 同构架构的优势
如先前研究所示 (Devvrit 等人，2024；Shukla 等人，2024；Valipour 等人，2023)，基础 ViT 的均匀层结构，具有一致的嵌入维度、token 数量和注意力头数，使其非常适合切片操作，并允许跨层进行直接选择。相比之下，像 Swin Transformer 这样的层次化模型 (Liu 等人，2022) 由于基于窗口的划分，会变化这些维度，从而增加了切片的复杂性。在大语言模型 (LLMs) 中也观察到了对同构设计的类似偏好 (Cai 等人，2024，2025)。此外，基础 ViT 通过一个简单的投影层使 ThinkingViT 的 Token 回收成为可能，而层次化模型则需要更复杂的机制来对齐不匹配的维度，即 token 数量和嵌入维度。
4 评估
在本节中，作者首先在ImageNet-1K（Russakovsky等人，2015年）及其变体上评估ThinkingViT，并将其与SOTA嵌套 Baseline 进行比较。然后，作者分析ThinkingViT如何利用熵作为信号，根据输入难度分配计算资源。最后，作者强调ThinkingViT的一个关键局限性：当面对均匀具有挑战性的输入时，ThinkingViT的二进制路由效果降低。
4.1 设置
实现细节：作者在ImageNet-1K（Russakovsky等人，2015）上以
的分辨率运行所有实验，使用在timm（Wightman，2019）中实现的模型，并按照Touvron等人（2021）的设置进行训练，所有模型均使用预训练的DeiT-Tiny预训练权重初始化。作者在配备NVIDIA A100 GPU的集群上进行训练，每轮训练约需2张GPU耗时10分钟。在原型开发阶段，作者进行了约100次训练，每次持续300轮。
Baseline 模型：作者将ThinkingViT与多个 Baseline 模型进行比较：MatFormer（Devvrit等人，2024年），该模型仅切割MLP的隐藏层，同时保持MHA和嵌入维度不变；DynaBERT（Hou等人，2020年），该模型切割MHA和MLP的隐藏层，同时保持嵌入维度不变；SortedNet（Valipour等人，2023年），该模型切割所有嵌入、NORM、MHA和MLP，同时保持头数固定；以及HydraViT（Haberer等人，2024年），该模型切割所有嵌入、NORM、MHA和MLP，同时还会改变头数。类似于HydraViT，ThinkingViT在头数、MHA、嵌入、NORM和MLP层上采用切割策略。然而，与上述方法不同，这些方法无论输入难度如何都遵循静态推理路径，ThinkingViT引入了输入自适应计算。
4.2 注意力扩展中的权衡
ThinkingViT通过逐步激活注意力头扩展计算能力。如表2所示，从3个头扩展到6个头（3H6H）获得最佳第一轮精度，而3H9H在ImageNet-1K上实现了最高第二轮精度。值得注意的是，通过使用损失加权，作者可以根据计算与精度之间的期望权衡来调整模型，使其更关注第一轮或第二轮。作者在表3中也探索了3阶段变体，证明ThinkingViT自然扩展到更深层次的思考结构，并获得了更高的最终精度。通常，思考阶段数和注意力头扩展步长是超参数，其最佳配置取决于数据集和目标效率目标。在图3中，作者评估了这种权衡，发现3H6H配置在ImageNet-1K上提供了精度与计算之间的最佳平衡。2H3H6H变体在高性能模型中跨越了最宽的GMACs范围，而与3H6H相比，精度仅略有下降。在高端，3H6H12H实现了最高精度，比3H6H至少高0.82个百分点（p.p.），同时将计算成本从约5 GMACs提高到至少约12 GMACs。这种显著的成本增加导致整体效率低于3H6H设置，突显了将ThinkingViT的思考策略与特定部署目标对齐的重要性。由于作者专注于在精度和GMACs之间提供最佳权衡的配置，因此后续实验采用3H6H。更详细的信息请参见附录I。
4.3 比较ThinkingViT和 Baseline 模型
为评估鲁棒性，作者在ImageNet-V2（Recht等人，2019）和ImageNetReaL（Beyer等人，2020；Russakovsky等人，2015）上评估了ThinkingViT，如图4所示，在ImageNet-R（Hendrycks等人，2020）上评估了ThinkingViT，如图7所示。在附录H中提供了ImageNet-A（Hendrycks等人，2019）和ImageNetSketch（Wang等人，2019）的额外结果。虽然ThinkingViT与 Baseline 之间的性能差距在ImageNet-R、-A和-Sketch上缩小（见第4.5节），但ThinkingViT在所有这些基准上都优于 Baseline 。此外，在ImageNet-ReaL、-Sketch和-R上，尽管ThinkingViT使用的参数显著更少（22.1M vs. 86.6M）且GMACs更低（5.85 vs. 17.56），但它仍然超过了DeiT-Base的准确率，这突出了Token Recycling的有效性。作者还在附录B中将ThinkingViT与Early Exit进行了比较，并在附录F中将ThinkingViT与几种基于token的剪枝方法进行了比较。
4.4 分析熵作为深度思考的信号
在每一阶段，ThinkingViT使用输出熵来评估确定性，并决定是否需要进一步计算。图5展示了在ImageNet-A（Hendrycks等人，2019年）、ImageNet-V2（Recht等人，2019年）和ImageNet-R（Hendrycks等人，2020年）上第一轮推理后的熵，为理解ThinkingViT
如何在不同分布偏移下估计不确定性提供了洞察。
在ImageNet-V2等较容易的数据集上，熵呈现左偏分布，反映了早期预测的自信。相比之下，ImageNet-A和-R等较难的数据集产生右偏分布，触发ThinkingViT进行更深层次的推理。在图8中，作者展示了在变化熵阈值下3H和6H子模型的负载分布。随着熵的降低，更少的样本进入第二阶段，而早期退出的样本很少被误分类。这证实了熵能够有效地阻止对简单输入的计算，同时不牺牲准确性。此外，图6展示了具有不同熵 Level 的图像，以说明ThinkingViT如何调整其计算。例如，清晰的单物体图片等简单图像具有低熵，并在一轮思考后停止。而更难的图像，包括包含多个物体、遮挡视图或光照条件差的图像，则表现出更高的熵，并触发第二轮思考以提高准确性。
4.5 局限性：在均匀困难输入上的性能限制
尽管ThinkingViT具有诸多优势，但它存在一个关键局限性：对于简单输入的早期预测效果高度依赖于每个阶段的预测质量。在每个决策点，模型会评估其确定性，以决定是否进入更深的阶段。当输入难度变化时，这种方法效果良好，允许对简单示例进行早期预测（见图2）。然而，当大多数输入都比较困难时，模型频繁地进入更深的阶段，这略微降低了早期预测的效率提升。在图7中，作者在ImageNetA数据集上展示了这一点，该数据集由ResNet-50（He等人，2016）错误分类的样本组成，在ImageNet-1K上的得分为76%。使用3H的ThinkingViT达到了73.6%，这表明几乎所有ImageNet-A样本都过于复杂，难以进行早期分类，因此触发了第二阶段。这解释了为什么ThinkingViT在严重偏向的数据集上的性能略有下降。有关局限性的进一步讨论，请参见附录J。
参考
[1]. ThinkingViT: Matryoshka Thinking Vision Transformer for Elastic Inference