Win-Win: Training High-Resolution Vision Transformers from Two Windows
作者: Vincent Leroy, Jerome Revaud, Thomas Lucas, Philippe Weinzaepfel
分类: cs.CV
发布日期: 2023-10-01 (更新: 2024-03-22)
备注: ICLR 2024
💡 一句话要点
提出Win-Win方法以高效训练高分辨率视觉变换器
🎯 匹配领域: 支柱三：空间感知与语义 (Perception & Semantics)
关键词: 视觉变换器 高分辨率图像 语义分割 深度预测 光流估计 训练效率 局部特征学习
📋 核心要点
- 现有方法在高分辨率像素级任务中训练成本高，且在测试时性能下降明显。
- 提出通过在训练中屏蔽大部分高分辨率输入，仅保留随机窗口，学习局部和全局交互。
- 在语义分割、单目深度预测和双目光流任务中，验证了该方法的有效性和推理速度提升。
📝 摘要（中文）
变换器已成为最先进视觉架构的标准，在图像级和密集像素级任务中表现出色。然而，训练高分辨率视觉变换器的成本极高。现有方法通常依赖于分层架构、快速近似注意力或低分辨率裁剪，这些方法在测试时会导致性能显著下降。本文提出了一种新策略，通过在训练过程中屏蔽大部分高分辨率输入，仅保留N个随机窗口，使模型能够学习窗口内的局部交互和不同窗口间的全局交互。该方法在测试时可以直接处理高分辨率输入，且训练速度比全分辨率网络快4倍。我们在三个高分辨率密集预测任务中验证了该策略的有效性。首先，在语义分割任务中，使用两个窗口的简单设置表现最佳，因此命名为Win-Win。其次，在单目深度预测任务中也确认了这一结果。最后，在双目光流任务中，该方法在Spring基准测试中达到了最先进的性能，推理速度比最佳竞争者快一个数量级。
🔬 方法详解
问题定义：本文旨在解决高分辨率视觉变换器训练成本高和测试性能下降的问题。现有方法如低分辨率裁剪导致在高分辨率测试时性能显著下降，且需要复杂的后处理。
核心思路：提出通过在训练过程中屏蔽大部分高分辨率输入，仅保留N个随机窗口，使模型能够专注于局部和全局交互。这种设计使得模型在测试时能够直接处理高分辨率输入，无需额外技巧。
技术框架：整体架构包括输入高分辨率图像，随机选择窗口进行训练，利用相对位置嵌入（如旋转嵌入）来增强模型的学习能力。训练过程中，模型在每个窗口内学习局部特征，同时也学习不同窗口间的全局特征。
关键创新：最重要的创新在于通过窗口化输入来减少计算量，同时保持模型在高分辨率输入上的有效性。这与现有方法的根本区别在于不再依赖低分辨率输入或复杂的后处理。
关键设计：在参数设置上，选择合适的窗口数量N和相对位置嵌入方式，损失函数设计上确保模型能够有效学习局部和全局特征。网络结构上，采用标准的视觉变换器架构，结合窗口化策略进行优化。
🖼️ 关键图片
📊 实验亮点
在实验中，使用Win-Win方法的语义分割任务表现最佳，两个窗口的设置显著提升了性能。此外，在单目深度预测和双目光流任务中也取得了优异结果，尤其在Spring基准测试中，推理速度比最佳竞争者快一个数量级，展现了该方法的实用性和高效性。
🎯 应用场景
该研究具有广泛的应用潜力，尤其在需要高分辨率图像处理的领域，如自动驾驶、医学影像分析和视频监控等。通过提高训练效率和推理速度，能够推动相关技术的实际应用和发展，降低计算资源的需求。
📄 摘要（原文）
Transformers have become the standard in state-of-the-art vision architectures, achieving impressive performance on both image-level and dense pixelwise tasks. However, training vision transformers for high-resolution pixelwise tasks has a prohibitive cost. Typical solutions boil down to hierarchical architectures, fast and approximate attention, or training on low-resolution crops. This latter solution does not constrain architectural choices, but it leads to a clear performance drop when testing at resolutions significantly higher than that used for training, thus requiring ad-hoc and slow post-processing schemes. In this paper, we propose a novel strategy for efficient training and inference of high-resolution vision transformers. The key principle is to mask out most of the high-resolution inputs during training, keeping only N random windows. This allows the model to learn local interactions between tokens inside each window, and global interactions between tokens from different windows. As a result, the model can directly process the high-resolution input at test time without any special trick. We show that this strategy is effective when using relative positional embedding such as rotary embeddings. It is 4 times faster to train than a full-resolution network, and it is straightforward to use at test time compared to existing approaches. We apply this strategy to three dense prediction tasks with high-resolution data. First, we show on the task of semantic segmentation that a simple setting with 2 windows performs best, hence the name of our method: Win-Win. Second, we confirm this result on the task of monocular depth prediction. Third, we further extend it to the binocular task of optical flow, reaching state-of-the-art performance on the Spring benchmark that contains Full-HD images with an order of magnitude faster inference than the best competitor.