Reg4Pru: Regularisation Through Random Token Routing for Token Pruning

Technical Summary: Reg4Pru – Regularisation Through Random Token Routing for Token Pruning

Problem Statement
Vision Transformers (ViTs) have become foundational for modern vision models due to their scalability and generalisation capabilities. However, their reliance on self-attention mechanisms results in computational complexity that scales quadratically ($O(N^2)$) with the number of tokens $N$. This bottleneck is particularly acute in medical imaging segmentation, where high-resolution inputs are required to preserve critical diagnostic details.

While token pruning methods (removing redundant tokens) and reactivation strategies (reintroducing pruned tokens at inference for dense pixel-wise predictions) offer a path to efficiency, the authors observe a significant performance degradation in deep layers. Specifically, when pruned tokens are reactivated during inference, they exhibit drastically different representations compared to their state during training. This train-test distribution shift leads to instability and poor dense prediction performance, counter-intuitively causing performance to decline as model depth increases.

Methodology
The paper proposes Reg4Pru, a training-time regularisation technique designed to mitigate the instability associated with token pruning and reactivation. The core mechanism involves random token routing during the training phase to simulate the effects of pruning without permanently discarding tokens.

1. Random Token Routing: During the forward pass of training, the model selects a random routing range $[l, n]$ between intermediate transformer blocks. A random subset of tokens ($\tau_r$) is chosen to skip the computation of blocks from $l$ to $n$, while the remaining tokens ($\tau_k$) are processed normally.
2. Gradient Flow and Representation: The routed tokens bypass specific blocks, effectively cutting their gradients with respect to those layers. However, by randomising both the subset of routed tokens and the routing span $[l, n]$ at every step, different tokens contribute gradients to different blocks across iterations. This forces the model to learn depth-invariant representations, ensuring that tokens can be reactivated at any depth during inference without suffering from representation mismatch.
3. Inference Strategy: At inference, a separate policy network predicts which tokens to prune based on a hierarchical keep-ratio schedule. Pruned tokens are temporarily removed to save computation and are reactivated at the final block to generate dense segmentation masks.
4. Loss Functions: The training objective combines a standard ratio-based loss (to enforce token budgets) and, optionally, an informed policy loss ($L_{pol}$) that uses ground truth segmentation masks to guide pruning decisions. The routing mechanism acts as a regulariser alongside these losses.

Key Contributions

• Depth-wise Instability Analysis: The authors identify and characterise a specific instability in token pruning for dense prediction, where performance degrades in deeper layers due to the reactivation of tokens that have not been exposed to the full depth of the network during training.
• Reg4Pru Framework: They introduce a training-time token routing strategy with random bounds. Unlike previous routing methods that use fixed bounds or learned routing policies, Reg4Pru randomises the routing range to approximate the feature distribution across all pruning depths, thereby regularising the model.
• Empirical Validation: The method is validated on the FIVES retinal blood vessel segmentation dataset, demonstrating that random routing significantly stabilises performance compared to fixed routing or no routing.

Results
Experiments were conducted on the FIVES dataset (2048$\times$2048 images, resized to quadrants for throughput testing) using a DINOv2-small initialized EoMT architecture.

• Performance Gains: Compared to a pruned model trained without routing, Reg4Pru improved the Average Precision (AP) by an absolute 46% (from 25.60% to 71.67% for the $L_{ratio}$ configuration).
• Efficiency: The Reg4Pru configuration achieved a 29% relative speedup in wall-clock time compared to the non-pruned baseline while maintaining high segmentation quality.
• Stability: Depth-wise analysis (Figure 5) showed that while non-routed pruned models suffered a sharp decline in AP across deeper blocks, Reg4Pru maintained stable performance. Fixed-bound routing mitigated some instability but failed to prevent performance drops after the fixed range ended, confirming the necessity of randomisation.
• Comparison: Reg4Pru outperformed Token Merging (ToMe) methods in terms of AP when comparing similar throughput configurations, particularly when ToMe was applied over both MHSA and MLP layers.

Significance and Claims
The paper claims that Reg4Pru is a valuable regulariser for token reduction strategies in dense prediction tasks. The authors argue that the instability observed in pruning is a result of the train-test distribution shift regarding token representations at different depths. By randomising the routing range during training, Reg4Pru forces the model to learn representations that are robust to being skipped at various depths, effectively bridging the gap between training and inference.

The authors remain modest regarding the generalisability of their findings, noting that experiments were conducted on a single high-resolution medical dataset. They suggest that while the method is effective for segmentation, further evaluation is required to assess its applicability to other dense prediction tasks (e.g., depth estimation) and its potential as a general regularisation mechanism when combined with techniques like dropout or stochastic depth. They also acknowledge that while they explored informed policies, ratio-based policies currently yield superior performance in their specific setup.