💻 computer science

Putting Registers to Work: Task Registers for Token Pruning in Vision Transformers

本文提出了任务自适应剪枝（Task-Adaptive Pruning, TAP），这是一种利用特定任务寄存器在不同视觉 Transformer 任务中动态排序 Token 并分配剪枝预算的方法，在保持图像分类、语义分割和目标检测方面具有竞争力的性能的同时，实现了显著的吞吐量提升。

原作者： Hongsen Cao, Mona Jaber, Shanxin Yuan, Ahmed Sayed

发布于 2026-08-12

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Hongsen Cao, Mona Jaber, Shanxin Yuan, Ahmed Sayed

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象你拥有一个超级聪明的机器人大脑，它通过观察图片来尝试理解其中的内容。这个大脑的构造就像是一个由无数微小工人组成的庞大团队，每个工人手里都拿着一块拼图碎片（即“标记/token”）。为了做出决策，每个工人都必须与每一个其他工人进行交流。问题在于，如果图片很大，这种“聊天”的数量就会爆炸式增长，导致大脑过载、变慢，并且极其渴望电力。科学家们一直试图通过告诉大脑忽略那些无聊的部分，只保留重要的部分来解决这个问题，这有点像老师告诉班上的学生，只听取那些有真知灼见的学生发言。这被称为“标记剪枝（token pruning）”。但棘手的地方在于，什么才算得上是“聪明”或“重要”，会随着机器人试图完成的任务而改变。一张对于识别猫咪至关重要的图片碎片，对于判断天空是什么颜色来说可能毫无用处。长期以来，科学家们尝试使用一套通用的“忽略名单”来应对所有任务，但事实证明，一种尺寸并不适用于所有情况。

这篇题为《让寄存器发挥作用（Putting Registers to Work）》的论文，正是在解决这个难题。作者是来自伦敦玛丽女王大学的一个团队，他们发现决定丢弃哪些图片碎片的规则对于不同的任务来说是完全不同的。他们发现，一种在物体检测（比如在公园里找狗）中表现极佳的策略，可能会损害机器人理解场景的能力（比如弄清楚那只狗是在草地上还是在人行道上）。为了解决这个问题，他们发明了一个聪明的全新系统，称为任务自适应剪枝（Task-Adaptive Pruning, TAP）。

把机器人的大脑想象成一个繁忙的厨房。通常情况下，厨师（剪枝策略）会使用同样的食谱来决定扔掉哪些食材，无论他是在做汤还是在做蛋糕。作者意识到这是一个错误。相反，他们为厨房提供了一套特殊的“任务寄存器”——可以把它们想象成根据烹饪菜肴而变化的智能、神奇的食谱卡片。当机器人需要识别一只猫时，它会抽出一张“猫卡”，这张卡会告诉大脑究竟应该保留哪些图像部分以及忽略哪些部分。当它需要绘制整个房间的地图时，它会换上一张“房间卡”。这些卡片不仅决定了要扔掉什么，还决定了在脑处理图像的过程中，在每一层（layer）的哪个阶段进行丢弃。

研究人员通过冻结一个标准的机器人大脑，并在不重新训练它的情况下测试不同的“忽略规则”来进行测试。他们发现了三个巨大的惊喜。首先，挑选物体检测中的重要碎片与挑选场景映射中的重要碎片，其最佳方式是截然相反的。其次，机器人对处理最初几层的过程非常敏感，尤其是对于像判断图像中是猫还是狗这类简单的任务。第三，当机器人丢弃一个碎片时，它需要记住这个碎片原本的样子，以便稍后填补空缺。但关键在于：对于寻找物体，最好的办法是假装那个碎片从未存在过，并使用附近的邻居来填补空隙；而对于绘制场景地图，更好的办法是实际记住精确的差异并在稍后将其粘贴回来。

为了解决这一切，TAP 使用一张针对当前任务的单一“活跃”寄存器卡。这张卡会在大脑中穿行，实时做出三个决定：

选择（Selection）： 保留哪些标记，丢弃哪些标记。
预算（Budget）： 在每一步丢弃多少个标记（也许是早期多丢一些，或者后期多丢一些，反之亦然）。
恢复（Recovery）： 当机器人需要绘制详细地图时，要找回多少“被遗忘”的细节。

结果令人印象深刻。通过使用这些特定于任务的卡片，机器人在保持智能的同时变得更快了。在针对物体检测的标准测试（COCO）中，它变得快了 1.32 倍，同时几乎保持了相同的准确度（仅下降了 0.3 点）。在针对场景映射的测试（ADE20K）中，它变得快了 1.30 倍，且质量仅有微小的下降。即使对于简单的图像分类，它也保持了竞争力。

作者认为，这种方法是迈向真正统一机器人大脑的一步——一个可以根据不同的“思考风格”在运行中切换任务的机器人，无论是识别一辆车、描绘一条街道，还是阅读一个标牌，都可以使用同一个核心引擎。他们不仅仅是猜测这行得通，他们通过运行数千次实验证明了这一点：通过冻结大脑来测试单个规则，并展示了他们的新型“寄存器”系统如何始终优于旧有的“一刀切”方法。这提醒我们，在人工智能的世界里，有时最聪明的事情就是知道何时该忽略噪音，而“噪音”的定义完全取决于你想要观察的对象。

技术摘要：让寄存器发挥作用：用于视觉 Transformer 中 Token 修剪的任务寄存器

问题陈述

视觉 Transformer (ViT) 的计算量随 Token 数量呈二次方增长，这促使了通过 Token 修剪来降低推理成本的需求。虽然修剪准则（如注意力分数、空间覆盖率）多是针对图像分类任务开发的，但预训练的 ViT 骨干网络正越来越多地被应用于语义分割和目标检测等密集预测任务。

作者发现了一个关键差距：为一种任务设计的修剪策略并不一定能有效地迁移到其他任务。 通过受控的“探测”实验，论文证明了：

准则排名分歧： 分割和检测任务对修로나准则的排序不同。例如，基于注意力的修剪在早期层会导致分类性能显著下降，而空间覆盖率则更为稳定。相反，使检测误差最小化的准则在分割任务中的表现往往较差。
深度敏感性： 修剪准则的有效性随网络深度变化而显著不同，且在不同任务之间也存在差异。
恢复冲突： 密集预测任务需要对被移除的位置进行特征重建。论文发现，分割和检测任务偏好的恢复终点截然不同：分割任务受益于恢复完整的特征偏移（被移除 Token 与保留 Token 之间的差异），而检测任务在使用仅有的“替代” Token 时表现更好。

现有方法通常采用单一、静态的修剪策略或使用任务无关的选择器，未能考虑到这些任务特定且依赖深度的变化。

方法论：任务自适应修剪 (TAP)

论文提出了 任务自适应修剪 (TAP) 框架，该框架利用 任务寄存器 (Task Registers) 来动态控制 Token 修剪、预算分配和特征恢复。

核心机制：

任务寄存器： 寄存器不再仅仅作为任务无关的特征伪影存储空间（如 Darcet 等人 2024 年的工作所述），TAP 引入了一个每个任务可学习的初始寄存器。在向前传播过程中，只有与当前任务相关的寄存器处于激活状态。
演进状态： 激活的寄存器随图像 Token 穿过骨干网络而演进。其在任何层的状态都会调节三个关键决策：
1. Token 选择： 寄存器通过复用当前块现有的查询 (Query) 和键 (Key) 投影，为每个候选补丁计算得分。Token 被排序，并保留一个子集用于自注意力计算。
2. 精确预算分配： 一个以寄存器状态为条件的共享线性读取器，决定了每一层修剪的精确整数 Token 数量。这确保了能够精确满足全局移除预算，而无需使用速率损失 (rate loss)。
3. 恢复缩放： 对于密集任务，寄存器控制一个缩放因子 ( $\alpha$ )，决定了在读取时恢复多少存储的特征偏移（被移除 Token 与其替代 Token 之间的差异）。

实现变体：

TAP-J (联合微调)： 每个任务拥有独立的骨干网络、任务寄存器和读取器。这实现了完全的专业化。
TAP-F (冻结基座)： 多个任务共享一个冻结的骨干网络。任务特定的寄存器、低秩 (LoRA) 更新和头部进行联合训练。这在保持适应性的同时减少了参数存储。

特征恢复：
被移除的 Token 会根据键空间余弦相似度最高的方式，确定性地匹配到保留的 Token（即“替代” Token）上。特征差异（偏移）被存储下来。在密集读取点，重建计算为 $\hat{x}_i = x_{\text{stand-in}} + \alpha \cdot \delta_i$ ，其中 $\alpha$ 是针对每个任务和深度学习的。分类任务会跳过此恢复步骤。

核心贡献

受控框架： 作者开发了一个框架，通过冻结无修剪检查点并逐层应用无参数减少准则，以隔离分类、分割和检测任务中修剪选择的影响。这揭示了准则排名和恢复偏好的不一致性。
任务寄存器作为控制器： 论文重新定义了寄存器的角色——从被动的存储空间转变为主动的、任务特定的稀疏计算控制器。寄存器协调了 Token 选择、跨深度的预算分配以及恢复缩放。
统一的修剪与恢复： TAP 将 Token 移除和特征恢复视为一个统一的操作。它将丢弃的 Token 与替代 Token 进行匹配，并根据任务寄存器缩放恢复的偏移量，使重建特征保持在编码器之外，以维持稀疏性。

实验结果

作者在 ImageNet-1K（分类）、ADE20K（分割）和 COCO（检测/实例分割）上评估了 TAP，使用的 ViT-B 骨干网络参数为 $\rho = 0.5$ （保留 50% 的 Token）。

性能：
- TAP-J 在所有全量微调方法中实现了最小的精度下降。在 ADE20K 上，它达到了 47.0 mIoU（仅比无修剪情况下降 0.2），吞吐量为 1.30×。在 COCO 上，它达到了 53.7 box AP（下降 0.3），吞吐量为 1.32×。在 ImageNet-1K 上，它达到了 83.2 top-1（下降 0.4）。
- TAP-F（共享骨干）略逊于 TAP-J，但显著优于静态规则和其他基线，在 ADE20K 上达到 45.8 mIoU，在 COCO 上达到 52.5 box AP，且参数开销极小（约 1.2M 任务特定参数）。
效率： TAP-J 在保持竞争性精度的同时，提供了 1.26× 至 1.32× 的编码器加速。控制器开销（寄存器查询、预算读取、簿记）仅占编码器运行时间的约 0.7%。
内存： 与密集基线相比，TAP 减少了 9–18% 的峰值推理内存，其中分割任务由于中间块激活尺寸的减小，收益最为显著。
消融研究： 去除任务寄存器（使用共享寄存器）或将分配/恢复固定为静态值会导致性能显著下降，证实了任务特定及深度自适应控制的必要性。

意义与主张

论文声称，Token 减少不能被视为单一的可迁移规则。随着 ViT 表示从局部结构向高层语义演进，Token 的价值随深度和读取该骨干网络的特定任务而变化。

TAP 的意义在于其能够利用共享骨干网络，使稀疏计算适应不同视觉任务的具体需求。通过让演进的任务寄存器发挥作用，TAP：

解决了不同任务之间修剪准则排名的冲突。
优化了移除预算在网络深度上的分布。
使特征恢复取决于密集与稀疏读取的具体要求。

作者将 TAP 定位为迈向统一视觉系统的一步，这类系统可以通过共享骨干网络服务于多种任务，在预测性能与编码器吞吐量之间取得比静态修剪策略更有效的平衡。论文保持了谦逊的态度，指出 TAP-J 比最佳的分类专用方法 (DiffRate) 低 0.1 个百分点，这可能是因为分类依赖于类 Token (class token) 并跳过了密集恢复，导致 TAP 的其中一项机制在该任务中未被充分利用。未来的工作建议探索图像感知的预算分配和多替代 Token 机制。