Summary
The authors propose to use example packing to train ViTs, where training examples of various lengths are packed into a single sequence. This requires a few straightforward architecture changes, including modified attention masking, pooling, and positional embeddings. This scheme allows for some interesting ideas, such as variable image resolutions during training, variable token drop rates, and adaptive inference time computation.
Extensive experiments in the paper show that NaViT results in more efficient training and finetuning (in terms of TPU hours) and a better compute-performance tradeoff at inference time. Further analysis shows that mixed resolution training is beneficial to model performance, that the time-varying token drop ratio allowed by the model can improve results when the correct schedule is used, and that the proposed factorized positional embeddings can offer very good generalization to aspect ratios and resolutions unseen during training. Finally, additional experiments show promising behavior with respect to calibration, fairness, object detection, and video classification.
Strengths
Overall, this is a very comprehensive paper. The number of experiments and the aspects of the proposed models performance evaluated is impressive. The authors touch on a large number of relevant properties of the model, including calibration stability, positional embedding ablations, and out of distribution performance. Additionally, results are quite positive, showing promising results on image classification, object detection, video classification, and training efficiency. Finally, the appendix is thorough, and gives enough details to accurately reproduce all experiments from what I can see.
Weaknesses
According to appendix section B.1, classification experiments seem to be "compute-matched." Are there any experiments analyzing the asymptotic performance between ViT and NaViT, or can the authors discuss this? The experiments in the paper seem inconclusive on this point.
The question of compute matching also applies to downstream experiments. For example, in the fairness analysis, are compute matched ViT and NaViT being compared? If so then the conclusion from Appendix H that "native image resolution improves the performance of fairness signal annotation" may not hold.
NaViT seems to benefit from training at the original aspect ratio, and at variable resolutions. However, could either of these benefits be achieved through scale and crop data augmentation? What data augmentation is used for ViT? Additionally, I was under the impression that image stretching or shearing could be useful as a data augmentation. Could the authors please discuss this?
Where are the contrastive results? I might have overlooked something, but I can't find zero-shot imagenet or COCO image-text retrieval results as discussed on line 154.
Minor typos:
- Figure 2 typo: "image 2" is repeated twice in the "data processing" part
- Figure 9 is before Figure 7 and 8
- It seems that a lot of (if not all) appendix section references are incorrect. Examples of these typos occur in line 152, line 158, line 339, and line 120.
Questions
See weaknesses section, thank you.
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Limitations
I think the "limitations" section from the appendix could be more forthcoming and candid. It mostly just says "NaViT" is a great idea, but we didn't get to apply it to all the tasks we wanted to. It could benefit from a more honest discussion of the method's limitations, such as compute overhead from packing or additional architectural changes needed to support example packing.