用强化学习动态合并视觉变换器的令牌,提升效率且几乎不丢精度。
DORA: Dynamic Online Reinforcement Agent for Token Merging in Vision Transformers

- 设计在线强化学习代理,按输入内容实时决定令牌合并策略。
- 在保持精度损失小于0.05%时,最高实现12.66%的令牌合并率。
- 适合追求高效推理的视觉模型部署,尤其对OOD数据表现优异。
视觉变换器(ViTs)因自注意力机制与令牌序列长度呈二次复杂度,导致计算开销大。现有令牌压缩方法多依赖固定启发式规则、预设比例或静态离线掩码,缺乏对输入相关冗余的适应能力。本文提出DORA(Dynamic Online Reinforcement Agent),首个基于强化学习的在线令牌合并框架。将合并过程建模为马尔可夫决策过程,轻量级强化学习代理根据当前特征状态和层特定上下文,动态决定每层的合并策略。通过融合非线性知识蒸馏惩罚的密集奖励函数,在计算效率与特征保真间取得平衡。采用非对称演员-评论家架构,高容量评论家用于稳定离线训练,演员头极小以支持低计算量在线推理。在多个ViT规模(Tiny到Large)上评估显示,相比基线方法,DORA显著提升精度-效率权衡。在严格精度损失约束(≤0.05%)下,最高达12.66%的令牌合并率,并实现最高效基线569.7%的相对提升。在ImageNet-1K上,对齐精度约束下,相较先进方法,计算节省提升高达76%。在分布外(OOD)基准如ImageNet-A和ImageNet-C上,效率优势超过430%。
原文摘要 · Abstract (English)
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%.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。