🤖 machine learning

MergeOver: Post-Training Token Merging for Recursive Vision Transformers

MergeOver 是一种后训练方法，它将 Token 合并技术集成到递归权重共享的视觉 Transformer 中，在无需昂贵重训的情况下，显著降低了边缘设备上的内存占用和延迟，同时保持高精度。

原作者： Junseo Kim, Uraz Odyurt, Amirreza Yousefzadeh

发布于 2026-08-14

📖 1 分钟阅读☕ 轻松阅读

原作者： Junseo Kim, Uraz Odyurt, Amirreza Yousefzadeh

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

想象一下，你正试图教一个超级聪明的机器人，仅通过看照片就能识别出猫、狗和汽车。为了做到这一点，这个机器人使用了一种特殊的“大脑”，叫做“视觉 Transformer”（Vision Transformer）。你可以把这个大脑想象成一个由微型侦探组成的团队，照片中的每一个像素都会分配到一个侦探。像素越多，你就需要越多的侦探，而且他们之间必须进行更多的交流，才能弄清楚自己看到了什么。这对于提高准确率非常有帮助，但对于像智能手表或树莓派（Raspberry Pi）这样的小型计算机来说，这是一个巨大的难题，因为它们的内存或电池功率无法支撑成千上万个侦探同时进行聊天。

为了解决这个问题，科学家们尝试了两种主要的技巧。第一种是“递归权重共享”（recursive weight-sharing），这就像是雇佣了一名超级侦探，让他一遍又一遍地重复同样的工作，通过重复使用笔记而不是为每一步都雇佣一个新团队来节省空间。这种方法虽然节省了空间，但会让侦探的工作变得更慢、更沉重。第二种技巧是“Token 合并”（token merging），即当你意识到有些侦探正在观察相同的东西时，你只需将他们合并成一个大侦探以节省时间。现在的问题是，如果你尝试同时使用这两种技巧，且不需要从头开始重新训练整个机器人，会发生什么？这就像是试图教一个侦探在工作过程中如何与他的克隆体合并，而不破坏他正在遵循的地图。

来自特文特大学（University of Twente）的 Junseo Kim 及其团队撰写的这篇论文提出了一个名为 MergeOver 的新方法来解决这个谜题。想象一下机器人的大脑是一栋多层建筑，侦探们在楼层间移动。每当他们移动到新的一层时，建筑的结构要求他们必须站在一个完美的方格阵列中。但合并侦探通常会破坏这个方格阵列，留下空位，从而导致“电梯”故障。MergeOver 引入了一个聪明的“解合并”（Unmerge）技巧：它暂时将合并后的侦探拆分回原始的网格位置，时间长到刚好能通过电梯，然后一旦到达下一层安全区域，再将他们重新合并。这使得机器人在不需要重新训练的情况下，依然能保持较低的内存占用。

该团队在著名的图像数据集 ImageNet-1K 上进行了测试，使用了从强大的显卡到微小的树莓派 5 等各种类型的计算机。他们发现，MergeOver 是否奏效很大程度上取决于你让机器人一次看多少张照片（即“批处理大小”/batch size）。当机器人一次只看一张照片时，合并与解合并带来的额外工作实际上会让速度变慢。然而，当机器人一次看 16 张照片时，奇迹发生了。在强大的 GPU 上，这种设置减少了 38.4% 的内存需求，并且实际上让机器人快了 21.7%。在微小的树莓派上，对于 16 个批处理的任务，它让机器人快了 17.6%。最棒的是，机器人并没有因此变笨太多；其准确率仅下降了 1.47 个百分点，作者认为，为了获得如此大的速度提升，这是一个很小的代价。

不过，论文也谨慎地指出，这并不是一个能瞬间解决一切问题的魔杖。作者明确反对那种认为可以把这个方法直接套用到任何模型上并期望在任何地方都完美运行的想法。他们发现，对于单张图片任务（批处理大小为 1），由于“解合并”技巧带来的开销对于仅仅处理一张图像来说并不划算，该方法在高性能计算机上反而会降低速度。他们还指出，虽然他们的方法在节省内存和加速批处理方面表现出色，但并不一定能击败所有那些从底层设计之初就是为了快速运行而生的其他高效 AI 模型。研究结果是客观且真实的，但他们认为这只是一个将这些技术结合起来的“基准”（baseline），而非最终的完美解决方案。作者建议，未来的工作可能需要将 MergeOver 与其他技巧（例如更好的软件优化）相结合，以使其更快，尤其是在针对小型设备处理单张图片任务时。

=== 技术摘要：MergeOver

问题陈述

视觉 Transformer (ViT) 已在计算机视觉领域取得了最先进的性能，但由于庞大的参数量以及相对于序列长度的二次方计算复杂度 ( $O(N^2)$ )，其在资源受限的边缘硬件上的部署面临重大障碍。虽然递归权重共享（例如在 Sliced Recursive Transformer, SReT 中）通过在多次传递中复用权重有效地减少了参数量，但它也加剧了计算瓶颈和峰值激活内存 (PAM) 问题。相反，Token Merging (ToMe) 提供了一种后训练机制，用于减少序列长度并缓解这些瓶颈。

然而，将这两种范式进行集成并非易事。直接集成会产生严重的架构冲突：

空间约束： SReT 依赖于层次化卷积池化层，这些层需要严格的 2D 空间网格。标准的 Token Merging 会破坏序列长度的完美正方形属性，导致维度冲突。
分组约束： SReT 的分片组自注意力机制 (SGA) 要求序列长度必须能被特定组大小整除。不受约束的合并会违反这些边界。
追踪约束： SReT 利用空间置换来移动注意力窗口。ToMe 依赖于 token-mass 向量来维持比例注意力；如果该向量没有与特征同步进行置换，注意力机制将会被破坏。

现有的高效架构通常需要重新训练、蒸馏或结构重设计，这为结合递归权重共享与动态 Token 缩减的后训练压缩方法留下了空白。

方法论：MergeOver

作者提出了 MergeOver，一种将 ToMe 集成到 SReT 中且无需重新训练的后训练方法。该方法通过特定机制解决了上述三种不兼容性：

1. Unmerge 追踪栈（解决空间约束）

为了满足阶段间卷积池化层的几何要求，MergeOver 实现了一个 Unmerge 栈。

机制： 在一个阶段内，当 Token 被合并时，相应的“Unmerge”操作（将合并后的特征还原回其原始空间坐标）会被压入一个后进先出 (LIFO) 栈中。
恢复： 在序列进入阶段间卷积池化层之前，栈会被弹出，并按逆序应用 Unmerge 操作。这使序列恢复到原始长度和 2D 网格结构，从而允许卷积核正确运行，同时让 Transformer 块处理压缩后的序列。

2. 约束安全的合并率调整（解决分组约束）

合并率 ( $r$ ) 不能是任意的。它必须满足：

整除性： 结果序列长度 ( $N-r$ ) 必须能被当前及后续递归迭代中使用的组大小 ( $g$ ) 的最小公倍数 (LCM) 整除。
二分限制： ToMe 的双向软匹配 (BSM) 在单次迭代中合并的 Token 数量不能超过总数的一半 ( $r \le \lfloor N/2 \rfloor$ )。
实现： 目标缩减量被调整为最近的“安全”值 ( $r_{safe}$ )，以满足 $(N-r) \mod g = 0$ 且 $N-r \ge g$ 。

3. 同步 Token-Mass 追踪（解决追踪约束）

为了在 SReT 的空间置换期间维持比例注意力：

Token-mass 张量 ( $s$ ) 初始化为全 1，并在各层之间进行更新，以追踪每个 Token 所代表的原始 Patch 数量。
同步： Mass 张量与特征序列并行地通过完全相同的置换和逆置换函数进行路由。这确保了用于注意力计算（公式 2）的权重与其对应的 Token 保持对齐。
加权合并： 通过使用 mass 张量进行加权平均来进行合并，以防止“重型”Token 被低估。

4. Token 缩减调度

论文评估了四种调度策略，以确定最佳的缩减轨迹：

全局调度： 常数型和线性递减型（标准的 ToMe 基准）。
阶段内调度：
- 指数型： 在一个阶段内进行逐渐减小的缩减。
- 单次触发型 (Single-shot)： Token 缩减仅发生在每个阶段的第一个 Block，从而在该阶段随后的递归迭代中保持固定的序列长度。

核心贡献

MergeOver 公式化： 一种无需重新训练即可在层次化递归 Transformer 中实现 Token Merging 的后训练框架。
架构集成： 通过 Unmerge 栈和同步 mass 追踪，解决了 SReT 的空间/注意力约束与 ToMe 合并逻辑之间的特定不兼容问题。
调度分析： 对 Token 缩减策略进行了评估，确定了阶段内单次触发 (stage-wise single-shot) 策略与 SReT 架构的兼容性最高。
跨平台基准测试： 在涵盖 GPU、x86 CPU 和 ARM CPU (Raspberry Pi 5) 平台的 ImageNet-1K 上进行了全面的评估。

实验结果

在 ImageNet-1K 验证集上使用 SReT-Tiny-Distill 基准进行评估：

准确率与效率的权衡

选定配置： 采用阶段内单次触发调度，初始缩减比例为 $\rho_{shot} = 0.25$ 。
准确率： 仅降低了 1.47 个百分点 的 Top-1 准确率（从 77.39% 降至 75.92%）。
FLOPs： 减少了 22.0% 的理论 FLOPs。

硬件性能表现 (Batch Size 16)

GPU (RTX 4060 Ti)：
- 吞吐量： 提升了 21.7%。
- 峰值激活内存 (PAM)： 减少了 38.4%。
- 注：在 Batch Size 为 1 时，由于开销原因吞吐量下降了 21.7%，表明该方法更倾向于在 GPU 上进行批处理执行。
x86 CPU (Intel Core Ultra 9)：
- 延迟： 降低了 30.0%。
- 吞吐量： 提升了 42.9%（在激进的 $\rho_{shot}=0.40$ 配置下观察到，该配置会导致 5.03% 的准确率下降）。
- 内存： 过程驻留集大小 ( $\Delta$ RSS) 减少了 22.0%。
- 注：与 GPU 不同，MergeOver 加速了 x86 上的单流推理（Batch Size 1），使其延迟降低了 17.9%。
ARM CPU (Raspberry Pi 5)：
- 延迟： 在 Batch Size 16 时降低了 17.6%（在 Batch Size 1 时降低 2.4%）。
- 吞吐量： 在 Batch Size 16 时提升了 21.4%。
- 内存： 在 Batch Size 16 时 $\Delta$ RSS 减少了 39.6%。

重要性与主张

论文声称 MergeOver 成功弥合了递归权重共享与 Token Merging 之间的鸿沟，而这种组合此前从未被探索过。其主要意义在于：

后训练可行性： 它证明了无需重新训练、蒸馏或结构重设计，即可实现复杂的架构集成。
恢复效率： 它回收了递归权重共享引入的大部分吞吐量和内存成本，使递归模型在边缘部署中更具可行性。
平台依赖性： 结果强调了 Token 缩减带来的收益高度依赖于处理平台和 Batch Size。虽然单流 GPU 推理受到开销影响，但批处理执行和基于 CPU 的边缘设备（x86 和 ARM）显示出了显著的延迟和内存改进。
建立基准： MergeOver 为将 Token Merging 与层次化递归 Transformer 相结合提供了基础性的基准，为未来的低级算子优化和其他压缩技术（如量化）铺平了道路。

作者保持了谦逊的态度，承认评估仅限于 SReT-Tiny-Distill 模型，并且观察到的性能增益尚未与特定的架构属性或低级算子优化完全隔离。他们将 MergeOver 定位为面向边缘 AI 的多轴压缩研究的一个起点。