用强化学习动态剪枝视觉变压器，提速近一半且精度几乎不变
Reinforcement Learning-based Token Pruning in Vision Transformers: A Markov Game Approach
- 将剪枝建模为马尔可夫博弈，多智能体协同决策每张补丁
- ImageNet-1k上推理速度提升44%，准确率仅降0.4%
- 首次实现跨层自适应剪枝，适合部署高效率ViT模型的开发者
视觉变压器（ViTs）的计算开销随令牌数量呈二次增长，亟需高效的令牌剪枝策略。现有策略多为手工设计，缺乏对输入变化的自适应能力，且未考虑多层间剪枝的序列特性。本文首次（据我们所知）利用强化学习实现数据自适应的剪枝策略。将剪枝建模为序列决策问题，构建马尔可夫博弈框架，采用多智能体近端策略优化（MAPPO），每个智能体独立决策一个令牌的去留。设计奖励函数，促使智能体在效率与精度间协同竞争。在ImageNet-1k数据集上，该方法最高可提升44%的推理速度，同时仅造成0.4%的精度损失。源代码已公开于https://github.com/daashuai/rl4evit。
原文摘要 · Abstract (English)
Vision Transformers (ViTs) have computational costs scaling quadratically with the number of tokens, calling for effective token pruning policies. Most existing policies are handcrafted, lacking adaptivity to varying inputs. Moreover, they fail to consider the sequential nature of token pruning across multiple layers. In this work, for the first time (as far as we know), we exploit Reinforcement Learning (RL) to data-adaptively learn a pruning policy. Formulating token pruning as a sequential decision-making problem, we model it as a Markov Game and utilize Multi-Agent Proximal Policy Optimization (MAPPO) where each agent makes an individualized pruning decision for a single token. We also develop reward functions that enable simultaneous collaboration and competition of these agents to balance efficiency and accuracy. On the well-known ImageNet-1k dataset, our method improves the inference speed by up to 44% while incurring only a negligible accuracy drop of 0.4%. The source code is available at https://github.com/daashuai/rl4evit.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。