Win-Win: Training High-Resolution Vision Transformers from Two Windows

We sincerely thank the reviewer for the relevant remarks and points raised. We answer in the following about the applicability, ablations studies, missing references, and about the relationship to token pruning. The paper was updated accordingly.

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.

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.

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.

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.

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.

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.

We gladly thank the reviewer for the constructive feedback. We answer the points regarding the lack of comparisons with other ViT architectures, and the novelty.

Lack of Comparative Analysis: It does not provide a comprehensive comparative analysis with a wide range of existing methods for training high-resolution vision transformers.

Most works on fast vision transformers do not tackle pixelwise prediction tasks. In particular, many of them rely on pooling or downsampling tokens, which definitely harm pixelwise predictions. For such problems, maintaining the input resolution throughout is an intuitive solution. We provide a solution to do that while keeping the computational budget constrained. We do welcome any specific method suggestion to compare with.

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.

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.We gladly thank the reviewer for the constructive feedback. We answer the points regarding the lack of comparisons with other ViT architectures, and the novelty.

Lack of Comparative Analysis: It does not provide a comprehensive comparative analysis with a wide range of existing methods for training high-resolution vision transformers.

Most works on fast vision transformers do not tackle pixelwise prediction tasks. In particular, many of them rely on pooling or downsampling tokens, which definitely harm pixelwise predictions. For such problems, maintaining the input resolution throughout is an intuitive solution. We provide a solution to do that while keeping the computational budget constrained. We do welcome any specific method suggestion to compare with.

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.

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.

Questions:

For results in Table 1, row 1 (968 tokens) compare 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, it is because of min window size?

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.

Table 3 shows larger variance when using different window strategy compared to table 1, any insights for this?

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.

From Table 3, the best setting is 2 win. 10x10, 4 win. 7x7, then as state ‘using the best settings found previously), to eval on MPI-Sintel’, however, the setting used have different window size, how did the author choose the window size on MPI-Sintel? Does this mean each time, a search of window size is needed when applying to different data?

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.

In table 5, compared to CroCoFlow, which is also the paper’s base method, it has a performance drop, authors states using different backbone (vit-l v.s. vit-base), the comparison of using the same backbone is missing.

The results from Table 5 are from the official leaderboard. Results for Croco-Flow are with a ViT-Large backbone while we limited the scope of this paper to ViT-Base backbones to control our computational budget. We cannot have the results of CroCo-Flow with a ViT-Base backbone as this benchmarks uses an official leaderboard.

However, looking at values reported in the CroCo-Flow paper, we observe a difference of 0.08 EPE on MPI-Sintel clean between the Large and the Base backbone on the validation set, and both backbones are on par on MPI-Sintel final.

Our Win-Win method is only 0.06 EPE below CroCo-Flow (Large) backbone on the official leaderboard on the clean rendering pass, and is performing 0.10 better on the final pass. Additionally, note that CroCo-Flow (with the large overlap of 0.9 between tiles at inference time used for the submission to the official leaderboard) is 20 times slower at inference when using the same backbone size!

To summarize, our main contribution is not about improving the final performance. Rather, we provide a simple method to train fast at high resolution and with a straightforward inference, compared to training on crops and using tiling at inference, or training on full-resolution images (see Figure 1 right). Our approach does maintain the task performance.

We were glad to read the positive and constructive remarks. We answer below the questions raised, concerning the number of tokens for comparisons, non-square windows, metrics variance, choosing the number of visible tokens, and comparisons to CroCo-Flow.

Weaknesses:

While the extensive experiments show the two window strategy is the best one with its simple and good performance, some analysis of such strategy/experiment results is missing. For example, in Table 1 right, why using 1021 tokens will drop 0.6 performance to using 1009 is not clear.

Thanks for your question. Our paper shows that the number of tokens is not what matters the most, and that the structure of the masking has more impact. In Table 1 right, we aim at comparing different rectangle shapes for the two windows. We make this comparison while keeping the number of tokens approximately constant. This table mainly shows that having the two windows of the same size performs better than having two windows of different sizes, which is also the simplest strategy. One hypothesis is that this window scheme with the same size allows to better model local interactions compared to an asymmetric one.

The results from Table 3 indicates the select of window strategy affects a lot for optical flow task, which suggests difficulties of generating to other data/task (search of window is needed).

The average endpoint error (EPE) metric for optical flow corresponds to the difference in pixels between the ground-truth and estimated motions. In other words, the maximal difference range is less than 1 pixel, which remains rather small.

Thanks for your constructive and positive feedback. We now address the concerns raised, in particular regarding the scalability to higher image resolutions, and the substantiality of some improvements.

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

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.

Win-Win is better than ViT-Det by a mere 0.3%, suggesting a marginal enhancement. Can 0.3% deemed as a substantial improvement?

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.