Reg4Pru: Regularisation Through Random Token Routing for Token Pruning

TL;DR

Reg4Pru is a novel regularisation method that enhances token pruning in vision transformers, significantly improving segmentation accuracy and efficiency without sacrificing dense prediction performance.

Contribution

This paper introduces Reg4Pru, a regularisation technique that mitigates token pruning performance loss, enabling more efficient and accurate vision transformer models for segmentation tasks.

Findings

Reg4Pru improves average precision by 46% on blood vessel segmentation.

It achieves a 29% relative speedup in wall-clock time.

Reg4Pru acts as an effective regulariser for token reduction strategies.

Abstract

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.