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
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
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.
Figure 5
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
N0
M
B0
DeiT-B/16
768
12
196
96
{2,4,…,192}
ViT-B/16
768
12
196
96
{2,4,…,192}
ViT-L/16
1024
24
196
96
{2,4,…,192}
DINOv2-B/14
768
12
256
126
{2,4,…,252}
Appendix
Table 3: Backbone dimensions and initial budget grids. M is the budget-head output dimension. Token counts exclude CLS.
Table 4: Shared optimization settings for DORA training.
Backbone
DORA selector
Random
Similarity
Attention-WPR
DeiT-B/16
80.8
79.1±0.1
80.5
77.3
ViT-B/16
83.8
82.0±0.2
83.1
81.6
ViT-L/16
85.1
83.3±0.0
84.2
84.2
DINOv2-B/14
83.7
82.1±0.1
83.3
82.2
Appendix
Table 5: Selector substitution on ImageNet-1K (Top-1, %). Random reports mean ± 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 ± std.
80.82±0.07
21.76±0.18
0.97±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 ± sample standard deviation across the four training runs.
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%.
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 Transformers (ViTs) face severe computational bottlenecks due to the quadratic complexity of self-attention at high resolutions. Existing token reduction methods rely on local metrics - such as single-layer attention scores - that are inherently vulnerable to the attention sink phenomenon, where uninformative tokens are paradoxically preserved over salient foreground objects. We propose ASAP (Attention Sink Anchored Pruning), a training-free framework that recasts this sink as a feature. Modeling ViT information flow as a Lazy Random Walk, ASAP identifies the sink as a dominant accumulator of probability mass. By computing the diffusion distance to the sink within the cumulative transition matrix, ASAP partitions tokens via Radial Diffusion Clustering and compresses background redundancy through Transition Weight Pooling in a single shot. Extensive experiments across image, video, and vision-language tasks demonstrate ASAP outperforms state-of-the-art methods, accelerating throughput by up to 48% while maintaining - or even exceeding - baseline accuracy.
Jaehyuk Lee, Hanyoung Kim, Yanggee Kim +1
Department of Mathematics Korea University Seoul, Republic of Korea · Korea University Seoul, Republic of Korea
Vision Transformers (ViTs) achieve strong performance but suffer from high computational costs due to quadratic self-attention complexity. Although token reduction techniques such as pruning and merging mitigate this, they typically overlook how representations evolve across network depth. We propose RAPID, a depth-aware token reduction framework that adapts reduction strategies to the layer-wise characteristics of token representations. The primary methodological contribution is a bifurcated strategy: in shallow-to-middle layers, RAPID employs a redundancy-similarity aware pruning metric to eliminate over-represented local patterns. As features transition to global semantic concepts in deeper layers, the framework shifts to an importance-similarity aware merging mechanism. This stage leverages classification (CLS) token attention weights to protect semantically critical tokens while fusing less important but similar neighbors. Empirical validation on ImageNet-1K using ViT and DeiT architectures demonstrates that RAPID establishes a superior accuracy-compression Pareto frontier compared to plug-and-play baselines such as ToMe and ToFu. RAPID is particularly robust in aggressive compression regimes, achieving up to 4.29% higher accuracy than ToMe at extreme reduction rates. Our framework provides a training-free template for optimizing vision models by aligning reduction strategies with hierarchical feature evolution.
Kyumin Choi, Ikbeom Jang
Hankuk University of Foreign Studies Seoul, South Korea