Reg4Pru: Regularisation Through Random Token Routing for Token Pruning

技术摘要：Reg4Pru —— 通过随机标记路由实现正则化以进行标记剪枝

问题陈述
视觉 Transformer (ViT) 因其可扩展性和泛化能力，已成为现代视觉模型的基础。然而，它们对自注意力机制的依赖导致计算复杂度随标记数量 $N$ 呈二次方级增长 ($O(N^2)$)。这种瓶颈在医学图像分割中尤为突出，因为该领域需要高分辨率输入以保留关键的诊断细节。

虽然标记剪枝方法（移除冗余标记）和重激活策略（在推理时重新引入被剪枝的标记，以进行稠密像素级预测）提供了一条实现高效性的路径，但作者观察到深层网络中存在显著的性能退化。具体而言，当被剪枝的标记在推理阶段被重新激活时，它们表现出的表示与训练期间的状态存在巨大差异。这种训练-测试分布偏移会导致不稳定性并导致稠密预测性能变差，从而产生一种反直觉的现象：随着模型深度的增加，性能反而下降。

方法论
本文提出了 Reg4Pru，一种旨在减轻与标记剪枝和重激活相关联的不稳定性的训练期正则化技术。其核心机制是在训练阶段通过随机标记路由 (Random Token Routing) 来模拟剪枝的效果，而不永久丢弃标记。

1. 随机标记路由： 在训练的前向传播过程中，模型会在中间 Transformer 块之间选择一个随机路由范围 $[l, n]$。一个随机选定的标记子集 ($\tau_r$) 将跳过从 $l$ 到 $n$ 的块计算，而剩余的标记 ($\tau_k$) 则正常进行处理。
2. 梯度流与表示： 被路由的标记会绕过特定的块，从而有效地切断了它们相对于这些层的梯度。然而，通过在每次迭代中随机化路由的标记子集以及路由范围 $[l, n]$，不同的标记会在不同的迭代中为不同的块贡献梯度。这迫使模型学习深度不变的表示，确保标记在推理时的任何深度被重新激活时，都不会因表示失配而遭受影响。
3. 推理策略： 在推理阶段，一个独立的策略网络根据层次化的保留比例方案来预测哪些标记应该被剪枝。被剪枝的标记会被暂时移除以节省计算量，并在最后一个块处被重新激活，以生成稠密的分割掩码。
4. 损失函数： 训练目标结合了标准的基于比例的损失（以强制执行标记预算）以及可选的有信息策略损失 ($L_{pol}$)，后者利用地面真值（ground truth）分割掩码来引导剪枝决策。路由机制作为正则化项与这些损失函数共同发挥作用。

核心贡献

• 深度维度的不稳定性分析： 作者识别并表征了用于稠密预测的标记剪枝中一种特定的不稳定性，即由于重新激活了在训练期间未经历完整网络深度的标记，导致性能在更深层中发生退化。
• Reg4Pru 框架： 他们引入了一种带有随机边界的训练期标记路由策略。与使用固定边界或学习型路由策略的以往路由方法不同，Reg4Pru 通过随机化路由范围来近似所有剪枝深度下的特征分布，从而实现正则化。
• 实证验证： 该方法在 FIVES 视网膜血管分割数据集上进行了验证，结果表明，与固定路由或无路由相比，Reg4Pru 显著稳定了性能。

结果
实验在 FIVES 数据集（2048$\times$2048 图像，为进行吞吐量测试进行了象限缩放）上进行，使用了以 DINOv2-small 为初始化的 EoMT 架构。

• 性能提升： 与未进行路由的剪枝模型相比，Reg4Pru 将平均精度 (AP) 绝对提升了 46%（对于 $L_{ratio}$ 配置，从 25.60% 提升至 71.67%）。
• 效率： Reg4Pru 配置在保持高分割质量的同时，实现了比非剪枝基准模型 29% 的相对加速（墙钟时间）。
• 稳定性： 深度分析（图 5）显示，虽然非路由剪枝模型在较深层块中会出现 AP 急剧下降，但 Reg4Pru 保持了稳定的性能。固定边界路由缓解了部分不稳定性，但未能防止在固定范围结束后出现的性能下降，这证实了随机化的必要性。
• 对比： 在比较相似吞吐量配置下的 AP 时，Reg4Pru 的表现优于标记合并 (ToMe) 方法，特别是在 ToMe 同时应用于 MHSA 和 MLP 层的情况下。

意义与主张
本文主张 Reg4Pru 是标记减少策略在稠密预测任务中的一种有价值的正则化工具。作者认为，观察到的剪枝不稳定性是由于不同深度下标记表示的训练-测试分布偏移造成的。通过在训练期间随机化路由范围，Reg4Pru 迫使模型学习出对在各种深度被跳过的过程具有鲁棒性的表示，从而有效地弥合了训练与推理之间的差距。

作者对研究结果的普适性保持谨慎态度，指出实验是在单个高分辨率医学数据集上进行的。他们指出，虽然该方法对于分割任务非常有效，但仍需进一步评估以确定其在其他稠密预测任务（如深度估计）中的适用性，以及其作为与 Dropout 或随机深度 (stochastic depth) 结合使用的通用正则化机制的潜力。他们同时也承认，尽管他们探索了有信息策略，但在目前的设置中，基于比例的策略目前表现更优。