DORA: Dynamic Online Reinforcement Agent for Token Pruning in Vision Transformers

Kaixuan He; Song Chen; Yi Kang

cs.CVSep 28, 2026

DORA: Dynamic Online Reinforcement Agent for Token Pruning in Vision Transformers

Authors: Kaixuan He, Song Chen, Yi Kang

Organizations: University of Science and Technology of China, Hefei, China · Institute of Artificial Intelligence, Hefei Comprehensive National Science Center, Hefei, China

Abstract

Vision Transformers (ViTs) incur quadratic self-attention cost in the number of tokens. Most token-reduction methods adapt token identities within a prescribed layer-wise compression schedule, or search a static mask offline, and thus limit online adaptation of when and how much to prune. We propose DORA (Dynamic Online Reinforcement Agent), which learns an input-adaptive pruning policy itself for frozen ViTs. At each eligible block, a hierarchical actor decides whether to prune, how many tokens to remove, and which tokens to remove from each image's evolving representation. Because early deletions change the states observed by later decisions, DORA formulates pruning as a finite-horizon Markov decision process. Complete-prefix shadow evaluations convert final-prediction fidelity into localized per-step credit, while closed-loop accuracy feedback adjusts the fidelity penalty toward a shared accuracy-drop target. A privileged critic and all shadow computations are training-only. Deployment retains the frozen backbone and a lightweight actor that applies hard deletion and packed variable-length FlashAttention, converting token reduction into measured speedups. On ImageNet-1K with DeiT-Base, DORA reduces FLOPs by 38.4% relative to the uncompressed backbone within one percentage point of accuracy loss. Averaged across four ViT-type backbones at matched accuracy, DORA uses 13.2% fewer FLOPs and achieves 32.4% higher throughput than the corresponding per-backbone baseline means. Under zero-shot transfer to ImageNet-A, these gains widen to 20.3% and 45.6%, respectively.

Figures & tables

Figure 1: Joint efficiency and input-dependent behavior. Left: GFLOPs reduction and measured throughput gain at accuracy-matched operating points across four ViT-type backbones; DORA (star) achieves the strongest joint efficiency at comparable accuracy. Right: representative DeiT-B trajectories showing that the learned policy applies different pruning intensities and retains different spatial evidence across images.

Figure 2: Training–deployment separation in DORA. Training uses a hierarchical actor with complete-prefix rewards, accuracy feedback, and a privileged critic; deployment retains only the frozen ViT and deterministic actor with hard pruning and variable-length FlashAttention.

Backbone	Method	ImageNet-1K			ImageNet-A
Backbone	Method	Top-1 Acc (%)	GFLOPs	Throughput (images/s)	Top-1 Acc (%)	GFLOPs	Throughput (images/s)
DeiT-B	Native	81.8	35.2	3230	27.9	35.2	3230
	ToMe	80.9	24.1	3163	26.6	29.6	2433
	DiffRate	80.8	22.1	3512	–	–	–
	PiToMe	80.8	24.1	3821	26.7	23.6	3844
	MCTF	80.8	21.9	2122	26.7	26.6	1647

Table 1: Accuracy-matched comparison with seven token-reduction baselines on ImageNet-1K and ImageNet-A. Baseline operating points match DORA within 0.1 pp Top-1. Avg. baselines is the unweighted mean over available baselines; “–” denotes no matching point. Bold marks the best efficiency value in each backbone block and the DORA row.

Figure 3: Learned adaptivity in where and how much to prune. (a) Intervention locations learned for each backbone. (b) Min–max deletion budgets over gate-active validation images at each intervention block; filled points denote collapsed single-budget decisions. (c) The resulting per-image GFLOPs-reduction ranges and means.

Setting	Top-1(%)	GFLOPs	Native drop (pp)
Full DORA	80.79	21.66	1.00
Terminal-only reward	80.81	22.95	0.98
Fixed coefficient ( $a{=}30$ )	78.65	17.78	3.14

Table 2: Training-mechanism ablations on DeiT-B.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Backbone	$d$	$L$	$N_{0}$	$M$	$\mathcal{B}_{0}$
DeiT-B/16	768	12	196	96	$\{2,4,\ldots,192\}$
ViT-B/16	768	12	196	96	$\{2,4,\ldots,192\}$
ViT-L/16	1024	24	196	96	$\{2,4,\ldots,192\}$
DINOv2-B/14	768	12	256	126	$\{2,4,\ldots,252\}$

Appendix

Table 3: Backbone dimensions and initial budget grids.

M

is the budget-head output dimension. Token counts exclude CLS.

Setting	Value
Rollout / total training	256 images / 299,776 images (1,171 updates)
Actor / critic epochs	4 / 4 per update
Gate / controller / critic minibatch	512 / 64 / 128 transitions
Optimizer / gradient clipping	Separate Adam optimizers / max norm 0.5
Gate / critic learning rate	$5\times 10^{-5}$ / $10^{-4}$
Controller learning rates	Shared: $2\times 10^{-5}$ ; budget: $5\times 10^{-5}$ ; selector: $2\times 10^{-5}$

Appendix

Table 4: Shared optimization settings for DORA training.

Backbone	DORA selector	Random	Similarity	Attention-WPR
DeiT-B/16	80.8	$79.1\pm 0.1$	80.5	77.3
ViT-B/16	83.8	$82.0\pm 0.2$	83.1	81.6
ViT-L/16	85.1	$83.3\pm 0.0$	84.2	84.2
DINOv2-B/14	83.7	$82.1\pm 0.1$	83.3	82.2

Appendix

Table 5: Selector substitution on ImageNet-1K (Top-1, %). Random reports mean

\pm

standard deviation over five sampling seeds. Values are rounded to one decimal place.

Run	Top-1 (%)	GFLOPs/image	Native drop (pp)
Original seed	80.79	21.66	1.00
Additional seed 1	80.83	21.90	0.96
Additional seed 2	80.75	21.57	1.04
Additional seed 3	80.90	21.93	0.89
Mean $\pm$ std.	$80.82\pm 0.07$	$21.76\pm 0.18$	$0.97\pm 0.07$

Appendix

Table 6: Four-seed DeiT-B results on ImageNet-1K validation. Checkpoints are selected on the development split. The final row reports mean

\pm

sample standard deviation across the four training runs.

Explore similar work

DORA: Dynamic Online Reinforcement Agent for Token Merging in Vision Transformers
May 12, 2026Kaixuan He, Song Chen, Yi KangVision Transformer Soft-Token Representations
University of Science and Technology of China, Hefei, China · Institute of Artificial Intelligence, Hefei Comprehensive National Science Center
ASAP: Attention Sink Anchored Pruning
May 21, 2026Jaehyuk Lee, Hanyoung Kim, Yanggee Kim +1Visual Token Pruning Vision Transformer
Department of Mathematics Korea University Seoul, Republic of Korea · Korea University Seoul, Republic of Korea
RAPID: Layer-Wise Redundancy-Aware Pruning and Importance-Driven Token Merging for Efficient ViT
Jun 6, 2026Kyumin Choi, Ikbeom JangVisual Token Pruning Self-Supervised Vision Transformers
Hankuk University of Foreign Studies Seoul, South Korea

May 12, 2026, cs.CVJ/K move · Enter open · S save

DORA: Dynamic Online Reinforcement Agent for Token Merging in Vision Transformers

Kaixuan He, Song Chen, Yi Kang

University of Science and Technology of China, Hefei, China · Institute of Artificial Intelligence, Hefei Comprehensive National Science Center

Vision Transformer Soft-Token Representations Fast Inference Offline Reinforcement Learning

Open paper AskPDF arXiv

Vision Transformers (ViTs) incur significant computational overhead due to the quadratic complexity of self-attention relative to the token sequence length. While existing token reduction methods mitigate this issue, they predominantly rely on fixed heuristic metrics, predefined ratios, or static offline masks, which lack the adaptability to capture input-dependent redundancy during inference. In this paper, we propose DORA (Dynamic Online Reinforcement Agent), the first reinforcement learning (RL)-driven online inference framework for dynamic token merging in ViTs. We formulate the merging process as a sequential Markov Decision Process (MDP), where a lightweight RL agent determines the merging strategy for each Transformer block based on the current feature state and layer-specific context. To balance computational efficiency and feature fidelity, the agent is optimized via a dense reward function incorporating a non-linear distillation-based penalty. We implement an asymmetric Actor-Critic architecture that utilizes a high-capacity Critic for stable offline training while retaining a minimal Actor head for low-computation online inference. Evaluations across multiple ViT scales (Tiny to Large) demonstrate that DORA improves the accuracy-efficiency Pareto front compared to current baselines. Under strict negligible accuracy-drop constraints (<= 0.05%), DORA achieves up to a 12.66% token merging rate, and delivers up to a 569.7% relative improvement over the most efficient baseline. On ImageNet-1K, under aligned accuracy constraints, DORA achieves up to a 76% relative improvement in computational savings compared to state-of-the-art methods. Furthermore, on out-of-distribution (OOD) benchmarks such as ImageNet-A and ImageNet-C, DORA attains a relative efficiency advantage of over 430%.

Abstract

Explore similar work

DORA: Dynamic Online Reinforcement Agent for Token Merging in Vision Transformers

ASAP: Attention Sink Anchored Pruning

RAPID: Layer-Wise Redundancy-Aware Pruning and Importance-Driven Token Merging for Efficient ViT