Paper: Win-Win: Training High-Resolution Vision Transformers from Two Windows (ICLR-2024)

Reviewer comment (Experiments):
The evaluation appears to be somewhat constrained. High-resolution Vision Transformers (ViTs) possess a wide range of potential applications, including object detection, semantic segmentation, instance segmentation, and serving as a backbone in diffusion models. However, the authors have restricted their analysis to just two representative applications and a limited set of datasets. This limitation raises questions about the general applicability of the proposed method. Specifically, I am keen to ascertain the viability of applying this method to object detection and diffusion models.

Author response:
Thanks for your suggestion. In this paper we focus on pixelwise task. This is because for other tasks, hierarchical backbones already offer a viable alternative to train at high resolution. For instance they have been shown to work for tasks like object detection. However these hierarchical backbones have inherent limitations when considering pixelwise tasks, in particular because maintaining the input resolution is desirable. 

We do understand and agree with the request for more tasks. To provide that, we experimented with the task of monocular depth estimation on the NYU v2 dataset. While it was not specifically requested, we believe it was the best candidate because a) it is pixelwise as discussed, and b) it provides reasonably high-resolution images, which is also necessary for our method to be relevant.

Below, we report preliminary results with the delta-1 (higher is better) metric on this dataset for full-resolution training, ViT-Det, random tokens and our approach. These results confirm the previous finding of the paper. We perform on par with full-resolution training while being 3 to 4 times faster to train. These results are mentioned in Section 4.1 and detailed in Appendix D and Table 10 of the revised paper.

| NYU Depth (delta 1)                      |         |
|------------------------------------------|---------|
| Full-res ViT      (1200 tokens)          |   87.7  |
| Full-res ViT-Det  (1200 tokens)          |   86.8  |
| Random tokens (392 tokens)               |   80.8  |
| **Win-Win (2-win. 14x14) (392 tokens)** |   87.3  |

Reviewer comment (Experiments):
The ablation study regarding window selection is currently lacking completeness. Notably, an essential baseline - random and non-contiguous window selection - has been omitted for semantic segmentation tasks. I am curious whether Win-Win works better than random masking in single-frame applications, though Table 8 demonstrated its effectiveness for optical flow prediction.

Author response:
Great point. Thanks. We initially did not run this ablation as masking random regions is not compatible with typical convolutional heads that are run on top (e.g. mentioned at the end of page 2).
On semantic segmentation, when running our method with 968 randomly sampled patches, i.e., with a similar number of patches as Win-Win, we obtain a mIoU performance of 59.9, compared to 63.6 with our approach. 

To verify that this is not simply due to the fact that training the convolutional head with a larger receptive field is better, we also try with a simple linear head. In this case, we obtain 59.9 for randomly selected patches, 61.6 with full-resolution training and 62.1 with our proposed Win-Win approach. 

These results further highlight the benefit of our approach with structured masking.
Intuitively, using structured masking allows to model both local and global interactions between patches, while training significantly faster than the full-resolution counterpart. We added this baseline to Table 4 in the revised paper, and the results with a linear head in Appendix C and Table 9.

|                                    | conv. head | linear head |
|------------------------------------|------------|-------------|
| Random tokens (968 tokens)         | 59.9       | 59.9        |
| Win-Win (968 tokens, 2 win. 22x22) | **63.6**   | **62.1**    |
| Full-resolution (3600 tokens)      | 63.2       | 61.6        |

Reviewer comment (Novelty):
The paper brings to mind GridMask, a well-known data augmentation technique widely applicable in various vision-related tasks. It would be beneficial if the author could incorporate a discussion on GridMask, highlighting the similarities and differences with the Win-Win method. Furthermore, the paper lacks any discussion regarding inference acceleration techniques for Vision Transformers (ViTs) through token pruning, such as DynamicViT, A-ViT, IA-RED^2, and SparseViT. Incorporating a brief comparison or analysis of Win-Win in the context of these methods would enhance the completeness of the paper.

Author response:
Thanks a lot for bringing the GridMask reference which we had missed. In GridMask, the idea is to color some squared areas in black, to act as a data augmentation for semantic tasks: the model should implicitly leverage context to implicitly fill-in the missing information, similar to MAE for pre-training. 

The small gain we observe for our approach compared to full-resolution training may be explained by the fact that it also performs a similar augmentation. We thus add and discuss this insight in the revised paper (at the end of page 6).

We also note that our method has key differences with it: 
* a) We mainly aim at efficient training despite the quadratic complexity of attention: we completely remove patches, similarly to MAE, rather than simply changing their color. We do not wish to claim quantitative performance gains from using it as a data augmentation tool, as it is less clear what to expect from this augmentation when considering more geometric tasks.
* b) An other difference is that we use a structured masking where large areas are completely removed: it is thus harder to implicitly fill-in the blank in these regions
* c) We extend it to the case of binocular task with a smart masking strategy in the second image, guided by flow. We show that in this case, naively masking both images is sub-optimal. 

Regarding the token pruning strategies, we mention them in the second point of the second paragraph of the introduction, where we added the suggested references, thanks. They are indeed relevant, but these token pruning strategies are most of the time not applicable to pixelwise prediction tasks and rather to image-level or object-level tasks, for which reducing the number of tokens is not harmful. In contrast, with pixelwise prediction tasks, it is natural to keep the input resolution intact.

Reviewer comment (Novelty):
Lack of Novelty: Masking out image patches is not a new approach in the literature, and it is simple data augmentation tuning to mask out most of them. I highly recommend the authors to do in-depth exploration and analysis of this method.

Author response:
Masking out image patches has been extensively studied in the context of pre-training image representation, following the seminal work of MAE [He et al, 2022] or as a form of data augmentation e.g with MSN [Assran et al, 2022]. These types of masking are quite different from ours. 
* a) In our case, the targets are aligned with the *visible* patches, contrary to MAE where they are aligned with the *masked* patches, or to MSN where the targets are *invariant* to the masking. This is because these methods tackle pre-training and not downstream tasks.
* b) Our proposed method goes beyond a naive random masking of patches: we use a **structured masking strategy** where we mask out all tokens except two windows. This allows to model both local interaction and global interaction that are both capital in downstream performances. In the case of multiple input views, we show that guiding the window selection using relevant signal is key to preserve performance, which does not happen with random masking. 
* c) Our purpose is not to improve quantitative performance, which is what typical `simple data augmentation tuning' would achieve. It is to enable efficient training at high resolutions; we believe that this claim is well substantiated (see e.g. Figure 1). The masking is not a data augmentation here, it is a way to control the computational budget.

We actually did compare to random tokens for optical flow in Table 8, and we now also provide a comparison on semantic segmentation where with 968 random tokens, we obtain 59.9 mIoU instead of 63.6 with our structured masking strategy. This loss of performance with random masking is not only due to the fact that it is harder to train convolutional head on top, as we still obtain a significant gain up to 62.1 compared to 59.9 with random masking.

Our results demonstrate that a structured masking strategy consistently outperforms random masking, for both monocular and binocular tasks. Extending a masking strategy to binocular tasks like optical flow is also novel.

Reviewer comment (Experiments):
The paper presents semantic segmentation result for a single train and test resolution (1280x720). Does the performance hold for the solution exceeding 1280x720?

Author response:
Yes, it would work at higher resolution. As an example, our optical flow experiments on the Spring benchmark are performed on **1920x1080** images. 

We have rephrased the beginning of the last paragraph of Section 4.2 where the results on Spring are presented to insist more on how critical it becomes to reduce training and inference costs when dealing with such large resolutions.

Reviewer comment (Evaluation):
Win-Win is better than ViT-Det by a mere 0.3%, suggesting a marginal enhancement. Can 0.3% deemed as a substantial improvement?

Author response:
It is indeed a rather limited improvement, but an improvement nonetheless. Most importantly, we do not seek improved performance and do not wish to claim it, we have made sure that the paper is clear in that respect. We do wish to claim that Win-Win is efficient: it is **3 times faster** to train than ViT-Det (Figure 1) and **4 times faster** to train than ViT. It is also conceptually simple, and does not restrict nor modify the vanilla transformer architecture. Additionally, our contribution of using structured masking is orthogonal to attention approximation and thus open new ways of tackling the problem. We believe this makes the overall contribution significant.

Reviewer comment (Evaluation):
For results in Table 1, row 1 (968 tokens) compared with row 2 (988), using non-square window has a large drop of 1 mIOU, do authors have any insights on that? And for row 5 (1021 tokens), why a 0.6 mIOU drop compared to 1009 tokens, is it because of minimum window size?

Author response:
This could be a possible explanation. However, the main message from Table 1 is that, to decide the size of the two windows, choosing two squared windows of the same size is simplest and efficient. One explanation could be that this allows to maximize the number of training instances and the size for local interactions.

Reviewer comment (Evaluation):
Table 3 shows larger variance when using different window strategy compared to Table 1, any insights for this?

Author response:
The metrics are significantly different. For optical flow, it is an average error in pixels, compared to the standard mean intersection-over-union for semantic segmentation, where changes in scores have an effect on the metric only if this changes the relative ordering of the score classes.

Reviewer comment (Experiments):
From Table 3, the best setting is 2 win. 10x10, 4 win. 7x7, then as stated ‘using the best settings found previously’ to evaluate on MPI-Sintel, however, the setting used has different window size. How did the authors choose the window size on MPI-Sintel? Does this mean each time, a search of window size is needed when applying to different data?

Author response:
Good point, thanks. We did indeed increase the number of tokens, which differs from the setting in the ablations on the synthetic dataset and the sentence was indeed lacking clarity. We did keep the rest of the configuration: when increasing the number of tokens, we keep the number and shapes of the windows and thus simply increase their sizes. We report below the ablation when increasing the number of tokens with around 200, 300, 400 and 500 tokens. Note that 500 tokens corresponding roughly to one quarter of the total number of patches for the full-resolution images for MPI-Sintel, a similar ratio as for semantic segmentation. This number of tokens is also close to the one used by CroCo-Flow (384x320 crops, which represents 480 tokens). We added these results to the main paper in Section 4.2 and Table 4 left.

| window settings                   | (#ntokens img1/img2) | EPE val |
|-----------------------------------|----------------------|---------|
| 2 win. 10x10 -> 4 win. 7x7        | (200/196)            | 2.00    |
| 2 win. 12x12 -> 4 win. 9x8        | (288/288)            | 1.96    |
| 2 win. 14x14 -> 4 win. 10x10      | (392/400)            | 1.83    |
| **2 win. 16x16 -> 4 win. 11x11    | (512/484)            | **1.67**|