Reg4Pru: Regularisation Through Random Token Routing for Token Pruning

TL;DR

Reg4Pru addresses the quadratic cost $O(N^2)$ of attention in Transformers for dense prediction by regularising token pruning during training. It introduces train-time token routing with random bounds to emulate pruning and reduce depth-wise instability, plus an informed-policy extension for multi-budget deployment. The approach yields large gains on FIVES, with an absolute AP improvement around 46 percentage points over the same pruned model without routing and about a 29% wall-clock speedup relative to a non-pruned baseline, demonstrating practical throughput gains without sacrificing accuracy. This work suggests Reg4Pru as a general regulariser for token reduction strategies and a step toward robust, adaptable dense-prediction models.

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.

Figures (5)

• Figure 1: Top: Output heatmaps for pruned models with and without routing. Bottom: Average precision (AP) vs. throughput (H100, batch size 1, $1024^2$ input). Key: (a) Reg4Pru, (b) Pruned (Fixed Routing), (c) Pruned (No Routing), (d) No Pruning, (e) ToMe bolya2023ToMeDiff over MHSA+MLP. Pruning keep ratios range from $0.9$ to $0.5$; ToMe merge ratios range from $0.7$ to $0.4$. All metrics are computed from the final block.
• Figure 2: Selected Pruning Policy Frequency - Frequency with which each token is selected across blocks $[3 \mathinner{{\ldotp}{\ldotp}} 11]$ (zero-indexed), overlaid on the input image (also shown in \ref{['fig:qualitative_logit_routing_ablation']} row 1). Tokens shown as fully transparent are selected at every block. Remaining frequencies are visualised using the viridis colourmap, where darker values indicate less frequent selection and lighter values indicate more frequent selection. Examples shown for models trained with the ratio loss $\mathcal{L}_{\text{ratio}}$, and $\mathcal{L}_{\text{pol}}$ where random routing (Reg4Pru) is used or no routing is used.
• Figure 3: Qualitative segmentation results - Final block segmentation logits after combining the weighted combination of logits across query masks. Each row contains a distinct input from the FIVES dataset FIVES. While each column highlights no routing, fixed bounds routing and randomised routing Reg4Pru along with a baseline without any pruning or routing.
• Figure 4: AP vs. convolutional decode block for each method. Results are from the same forward pass per checkpoint, with all efficient methods achieving ${\sim}100$ images/s throughput. Key: (a) Reg4Pru, (b) Pruned (Fixed Routing), (c) Pruned (No Routing), (d) No Pruning, (e) ToMe bolya2023ToMeDiff over MHSA+MLP, (f) ToMe bolya2023ToMeDiff over MHSA, (g) Pruned with constant $\rho{=}0.5$ (No Routing).
• Figure 5: Block-wise Qualitative segmentation results - Block-wise segmentation logits after combining the weighted combination of logits across query masks. The top row shows the outputs after training with no routing. The middle row shows the fixed routing outputs, and the final row shows the model trained with Reg4Pru.Each row contains a distinct input from the FIVES dataset FIVES. Notably, the instability in row 1 is mitigated in rows 2 and 3. Row 3 also is able to produce less noisy outputs, especially around the optic disc. Further, the false-positive logits around the retinal field, introduced by pruning, are removed in both the fixed and random routing cases.