PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Kaixuan He, Song Chen, Yi Kang

arXiv:2609.34325v1cs.CV

TL;DR

Vision Transformers face quadratic self-attention cost, while many token-reduction methods retain a shared pruning schedule or search a static mask. DORA learns an input-adaptive sequential pruning policy for frozen backbones, achieving lower FLOPs and higher throughput at matched accuracy across four ViT-type backbones.

  • Problem

    Many token-reduction methods adapt token identities within a prescribed layer-wise schedule or search a static mask, limiting online adaptation of when and how much to prune.

    Source: §1 · p. 1

  • Method

    DORA uses a hierarchical actor to choose whether, how much, and which tokens to prune, with complete-prefix shadow evaluations and closed-loop accuracy feedback during training.

    Source: §1 · p. 2, §3.3 · p. 5, §3.4 · p. 5

  • Results

    At matched ImageNet-1K accuracy, DORA averages 13.2% fewer FLOPs and 32.4% higher throughput than per-backbone baseline means across four ViT-type backbones.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Across the evaluated frozen-backbone methods, DORA improves the accuracy-efficiency trade-off, and hard deletion with packed variable-length FlashAttention converts reduced computation into measured acceleration.

    Source: §5 · p. 10

    Evaluation is limited to image classification on four ViT-type backbones, with runtime benefits measured on GPU using variable-length FlashAttention; training incurs extra shadow-evaluation and critic costs.

    Source: §5 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

DORA learns when, how much, and which tokens to prune online for each image, addressing the limits of shared compression schedules and offline masks. Its training assigns localized credit to sequential pruning decisions and adjusts fidelity penalties toward a shared accuracy-drop target.

Motivation

  • Existing methods generally adapt token identities within a prescribed layer-wise schedule, although input-dependent redundancy means a shared budget can retain unnecessary tokens or remove useful detail.DORA therefore makes the compression budget itself input-adaptive, based on each image’s evolving representations.Source: §1 · p. 1
  • Because early deletions change later representations and individually expendable tokens may be unsafe to remove together, DORA formulates pruning as a finite-horizon MDP.This lets the objective account for long-term effects of discrete actions rather than relying only on local attention or similarity proxies.Source: §1 · p. 1

Main results

  • At matched ImageNet-1K accuracy, DORA averages 13.2% fewer FLOPs and 32.4% higher throughput across four ViT-type backbones; on zero-shot ImageNet-A, the gains are 20.3% and 45.6%.Figure 1 compares matched-accuracy efficiency across backbones and shows representative input-dependent pruning trajectories.Source: §1 · p. 2, Figure 1 · p. 2

DORA

  • A hierarchical actor decides where to prune, how many tokens to remove, and which tokens to remove, using rollouts that physically delete tokens so later decisions observe the altered state.The policy learns pruning for frozen ViTs rather than searching for a static mask.Source: §1 · p. 2

Training and deployment

  • Complete-prefix shadow evaluations localize prediction-fidelity credit, while closed-loop accuracy feedback adjusts the fidelity penalty toward a shared accuracy-drop target across backbones.Deployment retains a deterministic actor and frozen backbone; hard deletion with packed variable-length FlashAttention turns token reduction into measured speedups.Source: §1 · p. 2

2 RELATED WORK

Prior token-reduction work scores, selects, or aggregates tokens, but often leaves some part of the layer schedule, token count, or mask shared across inputs. Some importance and merging mechanisms also require changes that are incompatible with standard fused attention.

Token pruning and merging

  • Pruning methods score and discard tokens, while merging methods aggregate them; despite image-dependent token choices, their per-layer budgets may remain prescribed and shared across inputs.Examples include attention- or graph-based pruning and similarity-based or graph-based merging.Source: §2 · pp. 2–3

Adaptive token reduction

  • Adaptive methods still fix some schedule element: DynamicViT uses predefined token counts and fine-tunes the backbone, DiffRate shares learned rates across images, and RL4EViT retains predefined stages.V-Pruner searches offline for a mask shared across inputs; RL4EViT uses attention masking rather than physically removing tokens.Source: §2 · p. 3

Execution with fused attention

  • Methods that rank tokens using intermediate attention probabilities or add token-size terms to attention logits cannot use the standard fused operator directly without changes.Representation Shift restores compatibility using a training-free importance score derived from representation changes.Source: §2 · p. 3

3 METHODOLOGY

DORA treats token pruning as sequential control over a frozen ViT: a hierarchical actor adapts pruning decisions to each image’s evolving representation. Complete-prefix evaluations localize fidelity costs, accuracy feedback regulates their penalty, and deployment uses only the frozen backbone and deterministic actor.

3.1 OVERVIEW AND FORMULATION

  • DORA models one pruning decision per block as a finite-horizon MDP whose state includes surviving keys and pruning history, and whose actions specify whether, how much, and which tokens to prune.Deletion physically shortens the sequence after attention; the CLS token is preserved, and removed tokens do not enter that block’s MLP or later blocks.Source: §3.1 · p. 4

3.2 HIERARCHICAL POLICY

  • A hierarchical actor gates pruning, chooses a removal budget, and selects tokens; its shared encoder uses surviving-token context, and selector scores depend on the requested removal fraction.Training samples decisions, while deployment uses a gate threshold of 0.5, the largest budget logit, and the highest selector scores.Source: §3.2 · p. 4

3.3 LOCALIZED REWARDS FROM COMPLETE-PREFIX EVALUATION

  • Complete-prefix shadow evaluations replay each sampled pruning prefix and finish the backbone without further pruning, measuring fidelity against unpruned logits; successive fidelity-cost differences attribute change to each block’s action.The reward combines incremental fidelity cost with token-layer compression, which accounts for computation saved in downstream blocks; the fidelity increments telescope to the final divergence, with no separate terminal reward.Source: §3.3 · p. 5

3.4 CLOSED-LOOP ACCURACY FEEDBACK

  • DORA periodically measures the deterministic policy’s paired Top-1 drop on a held-out training shard and adjusts the fidelity coefficient toward a one-percentage-point target.The target regulates training but is not a hard constraint on unseen validation data.Source: §3.4 · p. 5

3.5 OPTIMIZATION AND INFERENCE

  • A privileged critic and shadow evaluations are training-only; optimization uses GAE and separate clipped PPO objectives, while inference retains only the frozen backbone and deterministic actor.At inference, packed sequences use variable-length FlashAttention; reported throughput includes actor execution, selection, packing, and compaction.Source: §3.5 · pp. 5–6

4 EXPERIMENTS

DORA is evaluated on four frozen ViT backbones against seven token-reduction baselines under accuracy-matched protocols. It improves average compute and throughput on ImageNet-1K and ImageNet-A, while learning backbone- and input-adaptive pruning behavior and using closed-loop feedback to regulate accuracy drop.

4.1 EXPERIMENTAL SETUP

  • The evaluation uses four frozen pretrained backbones, seven token-reduction baselines, a shared BF16 inference contract, and end-to-end throughput measurements on one RTX 4090 at batch size 320.Checkpoint selection minimizes GFLOPs subject to a one-percentage-point Native accuracy-drop constraint; baseline operating points are matched to DORA’s realized accuracy.Source: §4.1 · p. 6

4.2 ACCURACY–COMPUTE–THROUGHPUT TRADE-OFFS

  • At matched accuracy, DORA averages 13.2% fewer FLOPs and 32.4% higher throughput on ImageNet-1K, widening to 20.3% and 45.6% on zero-shot ImageNet-A.Table 1 covers accuracy-matched comparisons with seven baselines on both datasets.Source: §4.2 · pp. 6–7, Table 1 · p. 7

4.3 DORA LEARNS BACKBONE- AND INPUT-ADAPTIVE PRUNING

  • Figure 3 shows backbone-specific learned intervention blocks and image-varying deletion budgets, with mean GFLOPs reductions of 38.3%, 32.8%, 39.2%, and 34.8% for DeiT-B, ViT-B, ViT-L, and DINOv2-B.DINOv2-B block 3 is conditional, activating on 66.5% of validation images.Source: Figure 3 · p. 8, §4.3 · p. 8

4.4 ACCURACY FEEDBACK CALIBRATES FIDELITY ACROSS BACKBONES

  • Closed-loop feedback moves recent Top-1 drops toward a shared 1-pp target; on DeiT-B, fixed a=30 fails the development constraint, while target sweeps preserve the highest measured matched-accuracy throughput.The fixed-coefficient terminal checkpoint reaches 78.65% Top-1 with a 3.14-pp Native drop; targets of 0.5, 1.0, and 1.5 pp yield distinct accuracy–efficiency regimes.Source: §4.4 · pp. 8–9

4.5 ABLATIONS AND STABILITY

  • Terminal-only reward selects 80.81% Top-1 at 22.95 GFLOPs versus 80.79% and 21.66 GFLOPs with localized credit; four runs cluster at 80.82±0.07% and 21.76±0.18 GFLOPs.Table 2 reports training-mechanism ablations on DeiT-B.Source: §4.5 · p. 9, Table 2 · p. 9

5 CONCLUSION

DORA frames pruning as sequential control and uses localized fidelity credit with accuracy feedback. Across evaluated frozen-backbone methods, it improves the accuracy-efficiency trade-off, with scope limited to image classification on four ViT-type backbones.

A.1 GATE, SHARED ENCODER, BUDGET HEAD, AND SELECTOR

  • A hierarchical actor separates whether, how much, and which tokens to prune, using complete-prefix shadow evaluations to assign localized credit for coupled decisions.The backbone remains frozen, and deployment uses a deterministic actor.Source: §5 · p. 9, §A.1 · p. 13

5 CONCLUSION

  • DORA consistently improves the accuracy-efficiency trade-off among evaluated frozen-backbone methods, while hard deletion with packed variable-length FlashAttention yields measured acceleration.The evaluation targets GPU execution with variable-length FlashAttention.Source: §5 · p. 10
  • The evaluation is limited to image classification on four ViT-type backbones, and training incurs additional cost from shadow evaluations and the privileged critic.Runtime benefits are measured on GPU execution with variable-length FlashAttention.Source: §5 · p. 10

A.2 BACKBONE-SPECIFIC BUDGET GRIDS

  • Candidate removal counts come from a backbone-specific budget grid, and no pruning is represented by gate-off rather than a zero budget.The grid preserves at least four visual tokens; token counts exclude CLS.Source: §A.2 · p. 13, Table 3 · p. 13

A.3 TRAINING-TIME SAMPLING AND DETERMINISTIC INFERENCE

During training, DORA samples pruning actions, while inference applies deterministic budget and token-selection decisions. The sampled order is retained for PPO likelihood evaluation.

A.3 TRAINING-TIME SAMPLING AND DETERMINISTIC INFERENCE

  • Conditional on an active gate, training samples a feasible removal count and an ordered sequence of token indices without replacement.The selector scores tokens once for each pruning decision.Source: §A.3 · p. 13
  • PPO retains the sampled order to evaluate the same action’s likelihood under the old and new policies, although deletion depends only on the selected set.Source: §A.3 · p. 14
  • At inference, the budget selects the highest-logit feasible count, and the selector removes the highest-scoring visual tokens while preserving CLS and survivor order.Deletion occurs after the attention residual and before the MLP.Source: §A.3 · p. 14

B COMPLETE-PREFIX SHADOW EVALUATION AND REWARD DERIVATION

The reward combines token-layer compression with changes in prediction fidelity measured on complete-prefix shadows. A telescoping reward assigns fidelity changes to individual actions, and feedback is evaluated periodically on held-out training shards.

B.1 TOKEN-LAYER COMPRESSION

  • The token-layer compression surrogate weights each deletion by the number of later block outputs in which the token is absent.Deleted tokens still participate in current-block attention because deletion follows the attention residual and precedes the MLP.Source: §B.1 · p. 14

B.2 SHADOW EVALUATION PROCEDURE

  • A complete-prefix shadow replays rollout deletions through block l, then completes the remaining blocks without further pruning to produce logits.The unpruned Native forward is the empty prefix; gate-off steps add no pruning, and shadows are gradient-free.Source: §B.2 · p. 14

B.3 TELESCOPING REWARD DERIVATION

  • The signed fidelity increment Dl − Dl−1 measures the current action’s marginal effect on final-prediction fidelity, avoiding repeated penalties for accumulated divergence.With the undiscounted objective, adjacent fidelity terms cancel and no additional terminal reward is added.Source: §B.3 · pp. 14–15

C.1 FEEDBACK SPLIT AND CADENCE

  • Feedback uses eight fixed shards of 2,048 images, checking one shard every eight PPO updates after rollout collection and before the policy update.The feedback split is disjoint from the PPO rollout pool and development split; development and validation data are not used for feedback.Source: §C.1 · p. 15

C.2 COEFFICIENT INITIALIZATION AND UPDATE

The section specifies how DORA regulates its fidelity coefficient from paired accuracy feedback and how its critic and PPO updates support the pruning policy. The paired drop compares Native and DORA predictions on the same shard, while the critic supplies training-time value baselines for the decisions.

C.2 COEFFICIENT INITIALIZATION AND UPDATE

  • The fidelity coefficient starts at 30 for all four backbones, and feedback adjusts it toward the selected accuracy-drop target.The default target is one percentage point; the coefficient remains fixed during rollout collection and its associated PPO update, with the next value applied to the next rollout.Source: §C.2 · p. 15

C.3 PAIRED ACCURACY-DROP COMPUTATION

  • The paired drop is the fraction of Native-correct/DORA-wrong predictions minus Native-wrong/DORA-correct predictions on the same shard.It retains its sign, and 100 times the drop gives percentage points.Source: §C.3 · p. 15

D.1 CRITIC INPUTS AND ARCHITECTURE

  • The training-only critic uses pre-action tokens, Native logits, and scalar state to predict value baselines for gate, budget, and selector decisions.The selector value additionally conditions on the chosen budget; no value head observes selected token identities, post-action states, or ground-truth labels.Source: §D.1 · pp. 15–16

D.2 VALUE TARGETS AND PPO OPTIMIZATION

  • GAE uses γ = 1 and λGAE = 0.95, with gate, budget, and selector optimized through separate clipped PPO objectives.All three decisions share a return, while their own value baselines are used; budget and selector value heads train only on gate-on transitions.Source: §D.2 · p. 16

D.3 PCGRAD

  • Symmetric PCGrad projects conflicting budget and selector gradients on shared controller parameters before summing them.Their private heads retain original gradients, and the resulting shared-controller gradient is clipped to norm 0.5; gate and critic updates do not use PCGrad.Source: §D.3 · p. 16

D.4 TRAINING HYPERPARAMETERS

DORA uses shared optimization settings across all four backbones.

D.4 TRAINING HYPERPARAMETERS

  • The optimization settings in Table 4 are shared across all four backbones.Table 4 presents the shared DORA training settings.Source: §D.4 · p. 16, Table 4 · p. 17

E.1 DATA SPLITS

The appendices specify the data, checkpoint, baseline, compute-accounting, and timing protocols used to evaluate DORA and its component substitutions. The selector study finds higher Top-1 accuracy than three substitutes under DORA's learned gate and budget decisions.

E.1 DATA SPLITS

  • PPO rollouts, accuracy feedback, and checkpoint selection use disjoint ImageNet-1K training subsets; final accuracy uses all 50,000 validation and 7,500 ImageNet-A images.ImageNet-A predictions use its official 200-class mapping; preprocessing uses deterministic resizing and center cropping without stochastic augmentation.Source: §E.1 · p. 17

E.2 CHECKPOINT SELECTION

  • The default checkpoint minimizes development-set GFLOPs subject to at most a one-percentage-point Native-relative Top-1 drop, while baseline points are matched within 0.1 percentage points without interpolation.Validation accuracy and throughput are not used for DORA checkpoint selection, and the selected checkpoint transfers to ImageNet-A unchanged.Source: §E.2 · p. 17

E.3 BASELINE IMPLEMENTATIONS

  • All methods share frozen backbone weights, classifier, and preprocessing within each backbone setting; the evaluated baselines include paper-based reimplementation and offline-search variants.V-Pruner's exported mask is fixed across images and applied before the first block; offline search is excluded from deployment measurements.Source: §E.3 · p. 17

E.4 GFLOPS ACCOUNTING

  • GFLOPs count two FLOPs per multiply–accumulate and include the complete online forward, including token-reduction computations.Counts follow each image's realized token trajectory before averaging; training-only critic and shadow computations are excluded.Source: §E.4 · p. 18

E.5 RUNTIME AND THROUGHPUT MEASUREMENT

  • Throughput measurements use an RTX 4090 with batch size 320, and timed forwards include online token selection, packing, and compaction.ImageNet-1K timing uses five repeats over 8,000 GPU-resident images; ImageNet-A uses a fixed 7,360-image subset, and median latency and throughput are reported.Source: §E.5 · p. 18

F.1 SELECTOR SUBSTITUTION STUDY

  • The learned selector preserves higher Top-1 accuracy than random, similarity, and Attention-WPR substitutes on every backbone, conditional on DORA's learned gate and budget decisions.The substitutions reuse recorded gate and budget decisions and are not evaluations of the complete baseline methods.Source: §F.1 · pp. 18–19, Table 5 · p. 18

F.2 ACCURACY-FEEDBACK TARGET SWEEP

  • DeiT-B feedback-target runs use drops of 0.5, 1.0, and 1.5 percentage points, changing only the target across separate training runs.Each run selects the lowest-GFLOPs checkpoint meeting its development-set accuracy constraint, and baselines are matched to realized validation accuracy within 0.1 percentage points.Source: §F.2 · p. 19

F.3 FOUR-SEED STABILITY

Across four DeiT-B training runs, selected operating points remain closely clustered.

F.3 FOUR-SEED STABILITY

  • 80.82 ± 0.07% Top-1 and 21.76 ± 0.18 GFLOPs per image across four DeiT-B training runs.Each checkpoint is selected independently under the one-percentage-point development-set constraint; reported Native drops are measured on ImageNet-1K validation.Source: §F.3 · p. 19, Table 6 · p. 19
Loading 2609.34325v1…