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

📄 arXiv: 2605.11683v1 📥 PDF

作者: Kaixuan He, Song Chen, Yi Kang

分类: cs.CV

发布日期: 2026-05-12

备注: Preprint. Under review

💡 一句话要点

提出DORA：一种基于强化学习的ViT动态Token融合在线推理方法

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: Vision Transformer Token融合 强化学习 在线推理 动态网络 模型加速 计算效率

📋 核心要点

  1. 现有ViT Token缩减方法依赖固定规则，无法根据输入动态调整，导致效率和精度难以兼顾。
  2. DORA将Token融合建模为MDP，利用强化学习Agent在线决策，根据特征状态动态调整融合策略。
  3. 实验表明，DORA在精度几乎不变的情况下，显著提升了ViT的计算效率，并在OOD数据上表现出色。

📝 摘要（中文）

Vision Transformer (ViT) 由于自注意力机制的二次复杂度，导致计算开销巨大。现有Token缩减方法主要依赖于固定的启发式指标、预定义的比例或静态的离线掩码，缺乏捕获推理过程中输入相关冗余的自适应性。本文提出DORA (Dynamic Online Reinforcement Agent)，这是第一个基于强化学习(RL)的在线推理框架，用于ViT中的动态Token融合。我们将融合过程建模为序列马尔可夫决策过程(MDP)，其中轻量级RL Agent基于当前特征状态和特定层的上下文来确定每个Transformer块的融合策略。为了平衡计算效率和特征保真度，Agent通过包含基于非线性蒸馏的惩罚的密集奖励函数进行优化。我们实现了一种非对称Actor-Critic架构，该架构利用高容量Critic进行稳定的离线训练，同时保留最小的Actor Head用于低计算的在线推理。在多个ViT规模（Tiny到Large）上的评估表明，与当前基线相比，DORA提高了精度-效率帕累托前沿。在严格的、可忽略的精度下降约束（<= 0.05%）下，DORA实现了高达12.66%的Token融合率，并且比最有效的基线提高了高达569.7%。在ImageNet-1K上，在对齐的精度约束下，与最先进的方法相比，DORA在计算节省方面实现了高达76%的相对改进。此外，在诸如ImageNet-A和ImageNet-C之类的分布外(OOD)基准测试中，DORA获得了超过430%的相对效率优势。

🔬 方法详解

问题定义：Vision Transformer由于自注意力机制的特性，其计算复杂度与Token序列长度呈平方关系，导致计算开销巨大。现有的Token缩减方法，如基于启发式规则或静态掩码的方法，无法根据输入图像的内容动态地调整Token的数量，从而限制了模型在不同场景下的效率和精度。

核心思路：DORA的核心思路是将Token融合过程视为一个序列决策问题，并利用强化学习Agent来动态地决定哪些Token应该被融合。通过这种方式，模型可以根据输入图像的特征自适应地调整Token的数量，从而在保持精度的同时提高计算效率。这种在线决策的方式避免了离线训练的局限性，能够更好地适应不同的输入。

技术框架：DORA采用Actor-Critic架构的强化学习框架。整体流程如下：首先，ViT模型提取图像特征。然后，对于每个Transformer块，Actor网络根据当前特征状态和层特定的上下文信息，输出一个Token融合策略。Critic网络评估该策略的价值，并用于指导Actor网络的训练。最终，根据Actor网络输出的策略，对Token进行融合，并输入到下一个Transformer块。该过程在ViT的每个Transformer块中重复进行。

关键创新：DORA的关键创新在于将Token融合问题建模为一个马尔可夫决策过程，并使用强化学习Agent进行在线决策。与传统的基于固定规则或离线训练的方法相比，DORA能够根据输入图像的特征动态地调整Token的数量，从而在保持精度的同时提高计算效率。此外，DORA采用非对称的Actor-Critic架构，使用高容量的Critic进行离线训练，同时保持轻量级的Actor Head用于在线推理，从而降低了在线推理的计算开销。

关键设计：DORA的关键设计包括：1) 密集奖励函数，用于平衡计算效率和特征保真度，其中包含基于非线性蒸馏的惩罚项；2) 非对称Actor-Critic架构，使用高容量Critic进行稳定的离线训练，同时保持轻量级的Actor Head用于低计算的在线推理；3) 将Token融合过程建模为马尔可夫决策过程，并使用强化学习Agent进行在线决策。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

DORA在多个ViT模型（Tiny到Large）上进行了评估，结果表明，与现有基线相比，DORA显著提高了精度-效率帕累托前沿。在精度下降小于0.05%的约束下，DORA实现了高达12.66%的Token融合率，并且比最有效的基线提高了高达569.7%。在ImageNet-1K上，DORA在计算节省方面实现了高达76%的相对改进。在OOD基准测试中，DORA获得了超过430%的相对效率优势。

🎯 应用场景

DORA具有广泛的应用前景，尤其是在资源受限的场景下，如移动设备、嵌入式系统和边缘计算等。通过动态地减少Token数量，DORA可以显著降低ViT模型的计算开销，使其能够在这些平台上高效运行。此外，DORA还可以应用于实时图像处理、视频分析和自动驾驶等领域，提高模型的响应速度和能源效率。

📄 摘要（原文）

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%.