Summary
This study proposes a training scheme called stochastic dropout training, thereby presenting its resulting models, called HydraViT. From the observation that ViT with a smaller version can be understood as a subset of a larger version, the authors propose to sample a subnetwork of a larger ViT. Although this scheme trains the subnetwork of ViT, during inference, the size of ViT can be flexibly chosen considering hardware specs. Experiments demonstrate the validity of the proposed method.
Strengths
I would like to admit good motivation for choosing the subnetwork of ViT for training by slicing its dimension. Further, incorporating order importance can be useful for other purposes during training and inference. The flexible choice of dimension during inference can be useful for several practical scenarios, as described in the manuscript.
Weaknesses
When we change the size of ViT, such as hidden dim, head size, and embedding dim, can we guarantee consistent behavior of ViT? In other words, if several dimensions of ViT change, ViT might exhibit inconsistent behavior. For example, the self-attention operation is defined as $softmax(\frac{Q K}{\sqrt{d_k}})V$ and calibrates its scale using $d_k$, the dimension of the quary and key. I would like to know how the authors incorporated different dimensions into the self-attention layers.
The concept of training subnetwork might be similar to Dropout. I would say that though Dropout applies stochastic dropping without order-awareness, the schochastic dropout training proposed in this study incorporates the ordering, starting from the first elements. In this regard, the proposed method can be understood as a special case of Dropout. I am not saying that the contribution is weak due to similarity; however, I would like to know whether the authors adopted a $1/p$ scaling similar to Dropout. Specifically, the Dropout scheme drops elements during training and uses all elements during the test phase. To have consistent output from Dropout, we apply $1/p$ scaling during training. For $p=0.5$, we amplify the output of Dropout by twice. By adopting this practice, we obtain a consistent output scale. In consideration of this, the training subnetwork of ViT may require scale calibration similar to Dropout. Similarly, though the size of ViT can be flexibly chosen during inference, this may require scale calibration with respect to the size of ViT to have consistent behavior of ViT.
Questions
See weaknesses. Although I would like to admit the motivation and novelty of the proposed method, changing demensions of ViT may cause inconsistent behavior of ViT in self-attention or may require scale calibration for an intermediate feature map similar to Dropout. The proposed method may seem plausible at first glance, but technically, we should be careful when changing the dimension of ViT. I would like to know whether the authors addressed these issues.
Limitations
The authors discussed several limitations as occasion arises.