Transformers are widely adopted in modern vision models due to their strong ability to scale with dataset size and generalisability. However, this comes with a major drawback: computation scales quadratically to the total number of tokens. Numerous methods have been proposed to mitigate this. For example, we consider token pruning with reactivating tokens from preserved representations, but the increased computational efficiency of this method results in decreased stability from the preserved representations, leading to poorer dense prediction performance at deeper layers. In this work, we introduce Reg4Pru, a training regularisation technique that mitigates token-pruning performance loss for segmentation. We compare our models on the FIVES blood vessel segmentation dataset and find that Reg4Pru improves average precision by an absolute 46% compared to the same model trained without routing. This increase is observed using a configuration that achieves a 29% relative speedup in wall-clock time compared to the non-pruned baseline. These findings indicate that Reg4Pru is a valuable regulariser for token reduction strategies.

翻译：Transformer因其能够随数据集规模扩展的强大能力以及良好的泛化性，在现代视觉模型中被广泛采用。然而，这带来一个主要缺点：计算量随令牌总数呈二次方增长。已有众多方法被提出以缓解此问题。例如，我们考虑通过从保留的表征中重新激活令牌来进行令牌剪枝，但该方法提升的计算效率会导致保留表征的稳定性下降，从而在更深层网络中获得较差的密集预测性能。在本工作中，我们引入了Reg4Pru，一种训练正则化技术，旨在缓解分割任务中令牌剪枝带来的性能损失。我们在FIVES血管分割数据集上比较了我们的模型，发现与未使用路由训练的相同模型相比，Reg4Pru将平均精度绝对提升了46%。这一提升是在一种配置下观察到的，该配置相比未剪枝的基线，在挂钟时间上实现了29%的相对加速。这些发现表明，Reg4Pru是针对令牌缩减策略的一种有价值的正则化器。