Summary
The paper presents a framework to speed up the pre-training of Vision Transformers (ViTs) in a self-supervised contrastive learning setup. The proposed method incorporates randomized token dropout and flexible patch scaling. The authors leverage this framework to analyze estimated gradient errors and its downstream performance. Additionally, they propose to determine an optimal dynamic acceleration schedule during training. Experimental findings demonstrate improvements in the convergence rate of the MoCo-v3 model across IN-100 and IN-1k datasets.
Strengths
- The proposed acceleration framework brings noticeable improvements in the pre-training convergence of MoCo-v3 on IN-100 and IN-1k.
- The framework incorporates various sequence compression strategies. The authors investigate how these strategies affect gradient estimation errors and analyze their impact on downstream performance.
- The acceleration framework includes a dynamic scheduler that adapts during training. It is validated across different training budgets and supported by ablation studies that highlight the importance of the individual contribution of token dropout and patch-scaling.
Weaknesses
- The proposed framework consists of two components: (1) randomized token dropout and (2) flexible patch scaling. Both ideas exist and have generally been studied for efficient pre-training of ViTs. For instance, the idea of dropping tokens in ViTs has been explored in various forms since their introduction. Recent research has focused on more sophisticated and targeted ways of applying token dropout in ViTs [1, 2, 3]. As a result, the main contributions of this work may be a bit limited.
- In section 4.3 - 'Since there are many more linear operations than quadratic ones, the time complexity of linear operations dominates. Thus, for the sake of simplicity, we consider the time complexity of the ViT architecture to be linear in the sequence length'. This statement seems oversimplified. While it's true that there are more distinct linear operations, this doesn't automatically mean they dominate the time complexity. I am not convinced that having more linear operations would negate the impact of the quadratic operation, even with a sequence length of ~200. This assertion requires further clarification and possibly revision.
- The concepts of randomized token dropout and patch scaling could potentially benefit ViTs in other self-supervised learning (SSL) approaches such as distillation-based SSL (DINO [4], iBOT [5]). I think there will be more contribution if the authors explored pre-training of ViTs in a broader SSL setting than being restricted to just MoCo-v3.
- I am skeptical about labeling the described approach as truly dynamic acceleration. Although it employs cost-adjusted MSE to compare strategies efficiently, the acceleration schedule is predetermined based on the analysis of intermediate checkpoints from a pre-trained model (as mentioned in section 5.1 under Dynamic Acceleration). A genuinely dynamic acceleration approach would involve real-time adjustments during training. Therefore, the method could be better described as an optimized static schedule that varies across different training stages rather than a fully adaptive, dynamic approach.
[1] Marin, Dmitrii, et al. "Token pooling in vision transformers." arXiv preprint arXiv:2110.03860 (2021).
[2] Ryoo, Michael S., et al. "Tokenlearner: What can 8 learned tokens do for images and videos?." arXiv preprint arXiv:2106.11297 (2021).
[3] Wang, Yulin, et al. "Not all images are worth 16x16 words: Dynamic transformers for efficient image recognition." Advances in neural information processing systems 34 (2021): 11960-11973.
[4] Caron, Mathilde, et al. "Emerging properties in self-supervised vision transformers." Proceedings of the IEEE/CVF international conference on computer vision. 2021.
[5] Zhou, Jinghao, et al. "ibot: Image bert pre-training with online tokenizer." arXiv preprint arXiv:2111.07832 (2021).
Questions
- Line 210: 'The variance is a function of their computational cost.' Could you provide references or further clarification on how variance is a direct function of computational cost?
- Line 268 - 'The only modification was the use of a non-symmetric loss.' Is there a specific reason for this, considering that MoCo-v3 uses symmetric loss?
- Line 271 - 'Non-symmetric version produces more diverse batches.' Could you elaborate on what diverse batches mean and how beneficial they are exactly?
- See weaknesses.