# Reinforcement Learning-based Token Pruning in Vision Transformers: A Markov Game Approach

Source: [Reinforcement Learning-based Token Pruning in Vision Transformers: A Markov Game Approach](<https://www.sota2.com/research/paper/68fa77ad1e209131cea74a6f>)

Collection index: [llms.txt](<https://www.sota2.com/research/llms.txt>)

Data updated: 2025-10-26T04:24:32.092Z

Authors: Chenglong Lu, Shen Liang, Xuewei Wang, Wei Wang

Published: 2025-03-30T14:34:28.000Z

## Abstract

Vision Transformers (ViTs) have computational costs scaling quadratically with the number of tokens, calling for effective token pruning policies. Most existing policies are handcrafted, lacking adaptivity to varying inputs. Moreover, they fail to consider the sequential nature of token pruning across multiple layers. In this work, for the first time (as far as we know), we exploit Reinforcement Learning (RL) to data-adaptively learn a pruning policy. Formulating token pruning as a sequential decision-making problem, we model it as a Markov Game and utilize Multi-Agent Proximal Policy Optimization (MAPPO) where each agent makes an individualized pruning decision for a single token. We also develop reward functions that enable simultaneous collaboration and competition of these agents to balance efficiency and accuracy. On the well-known ImageNet-1k dataset, our method improves the inference speed by up to 44% while incurring only a negligible accuracy drop of 0.4%. The source code is available at https://github.com/daashuai/rl4evit.

## Recorded benchmark results

A research result applies to its stated dataset, split, metric and evaluation configuration. Compare like-for-like settings, preserve the metric's unit and direction, and distinguish paper-reported results from independent reproduction. The best result recorded here is not a claim of exhaustive, current global coverage.

## Original sources and implementation

[Original paper](<https://arxiv.org/abs/2503.23459>)

[PDF](<https://arxiv.org/pdf/2503.23459.pdf>)
