Multi-Agent Reinforcement Learning for Efficient Vision Transformer with Dynamic Token Selection

1. To address the proposed token selection by using RL, a straightforward approach is to optimize it with a single-agent reinforcement learning method. The input features can be regarded as the input state and the decision of whether to retain or discard each token can be regarded as the action. However, the authors didn’t explore single-agent RL methods or discuss the differences. It would be better for the authors to explain the reason for not exploring single-agent RL methods or analyzing the differences between single-agent and multi-agent RL methods.
2. Except for classification results on ImageNet-1K on ViT-B, there are no other datasets (CIFAR-10/100) or vision models (ViT-T/ViT-L/Swin). It would be better for the authors to validate RL4DViT on more datasets or vision models or explain the reasons for choosing only one experimental setting.
3. There are several outstanding token pruning methods that the authors didn’t compare or mention [1][2]. I think the authors should discuss the differences or the advantages of RL4DViT with more baseline methods.
4. Should the rewards of discarded tokens be the same? I think the tokens discarded at earlier stages should have higher rewards, which can cause more computation cost reduction. This is just my suggestion that may be helpful for further improvement and will not affect my ratings.

[1] Kong, Zhenglun, et al. "Spvit: Enabling faster vision transformers via latency-aware soft token pruning." European conference on computer vision. Cham: Springer Nature Switzerland, 2022.

[2] Bolya, Daniel, et al. "Token Merging: Your ViT But Faster." The Eleventh International Conference on Learning Representations.

1. Reference format error: \citep{} should be used for references w.r.t. the ICLR submission style.

2. Lacking proper references: In the Introduction section Lines 80-81, the authors state that previous works on dynamic token pruning favour Gumbel-Softmax since RL-based methods converge slowly. Could the authors cite which existing work(s) claims this? Besides, in the Introduction section Lines 87-88, MAPPO is mentioned as a representative RL-based method, but without reference.

3. Missing quite a lot of important token reduction works: In both the Introduction and Related Work sections, only a few fundamental yet outdated token reduction methods are cited. EViT [1] that proposes an efficient token selection strategy based on the [CLS] attention should be mentioned and compared as a strong baseline. ATS [2] that utilizes a learnable scoring function for estimating the importance of each token should be mentioned and compared as a strong baseline as well. In addition to [1,2], many following token pruning methods [3,4,5] and token merging methods [6,7,8] should be included and possibly compared in this paper.

4. Insufficient experiments:

   4.1. Lacking backbones: The proposed RL4DViT is only adopted and validated on DeiT-B [9]. However, its performance on other backbones is unclear. To demonstrate its generalizability on different model sizes, experiments on DeiT-S and/or DeiT-T should be conducted. To demonstrate its generalizability on different ViT architectures, experiments on LV-ViT [10] or Swin-Transformer [11] should be conducted.

   4.2. Lacking runtime comparisons: Although this paper provides theoretical computational complexities (i.e., GFLOPs), these complexities do not indeed reflect the model's efficiency. Some methods with low GFLOPs may result in an even longer inference time since some operations (e.g., tensor reshaping, and in-memory selection) do not count toward the theoretical complexity [12]. Following the latest common practice, I suggest the authors report the real inference time.

5. Lacking motivations on using multi-agent RL: When utilizing MAPPO in RL4DViT, the authors adopt the parameter-sharing schema for agent policies and value functions. Thus, it arouses an intuitive question that whether current MAPPO can be replaced by single-agent PPO. This question is not well addressed from the perspective of both the Introduction and Method parts. In addition, in the experiments, owing to the utilization of MAPPO, both PPO and IPPO [13] should be included in the baseline methods to illustrate the advantages of the multi-agent framework and the centralized critique respectively.

6. Trivial performance gain: While DynamicViT [14] achieves 81.3% top-1 accuracy on DeiT-B with 11.2GFLOPs, the proposed MAPPO-DeiT-B only achieves 81.38% with 11.6GFLOPs. Such performance gain is trivial and does not demonstrate the superiority of using RL in token selection. Nonetheless, DynamicViT is a 2021 work and has been surpassed by many following works in both accuracy and efficiency.

7. Lacking in-depth analysis: Following Weakness 6, given that this paper focuses on the token selection part, the authors should justify why the MARL-based selection is better, with comparisons to other token selection strategies outlined in [1,2,6,7,8]. However, this paper lacks a quantitative analysis of the benefits of using MARL. And the qualitative analysis in Figure 2 does not clearly demonstrate its advantage over existing token selection methods.

[1] Liang, Youwei, et al. "Not all patches are what you need: Expediting vision transformers via token reorganizations." ICLR, 2022.

[2] Fayyaz, Mohsen, et al. "Adaptive token sampling for efficient vision transformers." ECCV, 2022.

[3] Xu, Yifan, et al. "Evo-vit: Slow-fast token evolution for dynamic vision transformer." AAAI, 2022.

[4] Kong, Zhenglun, et al. "Spvit: Enabling faster vision transformers via latency-aware soft token pruning." ECCV, 2022.

[5] Kong, Zhenglun, et al. "Peeling the onion: Hierarchical reduction of data redundancy for efficient vision transformer training." AAAI, 2023.

[6] Bolya, Daniel, et al. "Token merging: Your vit but faster." ICLR, 2023.

[7] Kim, Minchul, et al. "Token fusion: Bridging the gap between token pruning and token merging." WACV, 2024.

[8] Xu, Xuwei, et al. "GTP-ViT: Efficient Vision Transformers via Graph-based Token Propagation." WACV. 2024.

[9] Touvron, Hugo, et al. "Training data-efficient image transformers & distillation through attention." ICML, 2021.

[10] Jiang, Zi-Hang, et al. "All tokens matter: Token labeling for training better vision transformers." NeurIPS, 2021.

[11] Liu, Ze, et al. "Swin transformer: Hierarchical vision transformer using shifted windows." ICCV, 2021.

[12] Haurum, Joakim Bruslund, et al. "Which tokens to use? investigating token reduction in vision transformers." ICCV, 2023.

[13] Witt, De, et al. "Is independent learning all you need in the starcraft multi-agent challenge?." arXiv preprint arXiv:2011.09533 (2020).

[14] Rao, Yongming, et al. "Dynamicvit: Efficient vision transformers with dynamic token sparsification." NeurIPS, 2021.

• Several issues with spacing and formatting throughout the paper.
• Typos in Figure 1. “specified”, “zeroed-out tokens” and many other typos throughout the manuscript.
• Experiments are only conducted on ImageNet which calls into question the scalability / applicability of the RL4DViT approach on other computer vision tasks and datasets.
• Little motivation provided for the design decisions in formulating the Multi-Agent MDP problem. For example, why is each token an agent in the environment?
• Too many unnecessary RL implementation details provided in the methods section. Equation 1 and 2 are simply the actor and critic losses for PPO training and doesn’t feel necessary.
• What does it mean for an agent to be alive in the reward function definition? Also it is not clear why these agents are both competitive and cooperative? It is not clear why the agents’ actions are not independent of one another.
• Instead of framing this as a MAMDP, could this just be a single agent that predicts the full binary mask?
• Can Figure 2 also show what the token selection is for the baseline methods so there is a qualitative comparison?
• In Table 1, it looks like the base Deit-B model has a higher GLOPs and better Top-1 Acc on compared to the proposed approach.

• I think the most significant weakness of this paper is its presentation.
  • This paper doesn't give much insight into the advantage of modeling token selection as a MARL process, nor does it explain how the token selection procedure can be transformed into an RL formulation. The paper only lists the computing procedure, which makes the methodology hard to follow.
  • Additionally, there are many repetitive rows in both tables. The paper does not give any explanation of the meanings of these rows. This makes me quite confused.
• I am also confused about the performance of the proposed method.
  • Is the performance of the proposed method better or worse than DynamicVit? The paper only provides two tables, but doesn't present any explanation for the meaning of the figures in the table.