Reg4Pru: Regularisation Through Random Token Routing for Token Pruning
This addresses computational inefficiency in Transformers for vision tasks like blood vessel segmentation, but it is incremental as it builds on existing token pruning methods.
The paper tackled the problem of performance loss in token pruning for Transformers in vision models, specifically for segmentation, and found that Reg4Pru improved average precision by 46% while achieving a 29% speedup.
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.