# Reg4Pru: Regularisation Through Random Token Routing for Token Pruning

arXiv: https://arxiv.org/abs/2602.02163 · HTML: https://mlsanity.com/paper/2602.02163/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Token pruning and merging for ViTs, one of 12 directions in Efficient LLM inference and serving.
- Builds on 35 of its 44 references that are in the map.
- First posted on arXiv 2 Feb 2026.

## Placed in

- Efficient LLM inference and serving / [Token pruning and merging for ViTs](https://mlsanity.com/inference/token-pruning-and-merging-for-vits/)

## Builds on (most cited first, 30 of 35)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [DINOv2: Learning Robust Visual Features without Supervision](https://mlsanity.com/paper/2304.07193/) (2023)
- [End-to-End Object Detection with Transformers](https://mlsanity.com/paper/2005.12872/) (2020)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers](https://mlsanity.com/paper/2105.15203/) (2021)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [Masked-attention Mask Transformer for Universal Image Segmentation](https://mlsanity.com/paper/2112.01527/) (2021)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Root Mean Square Layer Normalization](https://mlsanity.com/paper/1910.07467/) (2019)
- [Token Merging: Your ViT But Faster](https://mlsanity.com/paper/2210.09461/) (2022)
- [Deep Networks with Stochastic Depth](https://mlsanity.com/paper/1603.09382/) (2016)
- [DynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification](https://mlsanity.com/paper/2106.02034/) (2021)
- [Going deeper with Image Transformers](https://mlsanity.com/paper/2103.17239/) (2021)
- [Vision Transformer Adapter for Dense Predictions](https://mlsanity.com/paper/2205.08534/) (2022)
- [Albumentations: fast and flexible image augmentations](https://mlsanity.com/paper/1809.06839/) (2018)
- [Perception Encoder: The best visual embeddings are not at the output of the network](https://mlsanity.com/paper/2504.13181/) (2025)
- [The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits](https://mlsanity.com/paper/2402.17764/) (2024)
- [Mixture-of-Depths: Dynamically allocating compute in transformer-based language models](https://mlsanity.com/paper/2404.02258/) (2024)
- [Rotary Position Embedding for Vision Transformer](https://mlsanity.com/paper/2403.13298/) (2024)
- [Token Merging for Fast Stable Diffusion](https://mlsanity.com/paper/2303.17604/) (2023)
- [Flex Attention: A Programming Model for Generating Optimized Attention Kernels](https://mlsanity.com/paper/2412.05496/) (2024)
- [SPViT: Enabling Faster Vision Transformers via Soft Token Pruning](https://mlsanity.com/paper/2112.13890/) (2021)
- [What Matters in Transformers? Not All Attention is Needed](https://mlsanity.com/paper/2406.15786/) (2024)
- [Your ViT is Secretly an Image Segmentation Model](https://mlsanity.com/paper/2503.19108/) (2025)
- [Dynamic Token Pruning in Plain Vision Transformers for Semantic Segmentation](https://mlsanity.com/paper/2308.01045/) (2023)
- [Revisiting Token Pruning for Object Detection and Instance Segmentation](https://mlsanity.com/paper/2306.07050/) (2023)

## Abstract (arXiv)

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.