ICLR2024

Window Attention is Bugged: How not to Interpolate Position Embeddings

Daniel Bolya, Chaitanya Ryali, Judy Hoffman, Christoph Feichtenhofer

被引用 21 次

摘要

Window attention, position embeddings, and high resolution finetuning are core concepts in the modern transformer era of computer vision. However, we find that naïvely combining these near ubiquitous components can have a detrimental effect on performance. The issue is simple: interpolating position embeddings while using window attention is wrong. We study two state-of-the-art methods that have these three components, namely Hiera and ViTDet, and find that both do indeed suffer from this bug. To fix it, we introduce a simple absolute window position embedding strategy, which solves the bug outright in Hiera and allows us to increase both speed and performance of the model in ViTDet. We finally combine the two to obtain HieraDet, which achieves 61.7 box mAP on COCO, making it state-of-the-art for models that only use ImageNet-1k pretraining. This all stems from what is essentially a 3 line bug fix, which we name "absolute win". INTRODUCTION Transformer-based architectures dominate many tasks throughout computer vision (Zhai et al., 2021; Li et al., 2022a; Kirillov et al., 2023) . But despite their ubiquity, these architectures are relatively new, and thus many best practices have yet to be set in stone. In this paper, we focus on a relatively unassuming operation in modern vision transformers: interpolating position embeddings. Absolute position embeddings added at the beginning of a transformer allow the model to distinguish between tokens based on location, an important detail for most vision tasks. While natural language processing (NLP) can sometimes get away with omitting a position embedding (Haviv et al., 2022) , those in computer vision often find themselves adding additional position embeddings to deal with more difficult tasks (e.g., detection in Li et al. (2022a)). Several recent works even design custom position embeddings to inject directly into attention (Liu et al., 2021b; Graham et al., 2021; Li et al., 2022b), so that location information is right where it is needed. However, these additional position embeddings can be costly: techniques like attention bias (Liu et al., 2021b; Graham et al., 2021) or relative position embeddings (Li et al., 2022b;a) are added directly into the attention matrix. Not only are these operations slow, but they also cannot benefit from recent innovations such as Flash Attention (Dao et al., 2022; Dao, 2023) that speed up transformers by not constructing the attention matrix. Ideally, we would like to avoid relative position embeddings and just use simple and fast absolute position embeddings like the original ViT (Dosovitskiy et al., 2020). So, why do most architectures not use absolute positioning embeddings? One potential reason becomes apparent when we study Hiera (Ryali et al., 2023) , a modern hierarchical vision transformer that only uses absolute position embeddings. Hiera is as powerful and more efficient than other state-of-the-art vision architectures, while being composed entirely of simple ViT blocks. Instead of adding more position bias architecturally (e.g., with relative position embeddings), it learns good spatial biases through a strong pretext task (i.e., MAE (He et al., 2022)). This makes it the perfect case study for modern architectural design paired with simple absolute position embeddings. And immediately, an issue presents itself: Hiera does not interpolate well. When finetuning Hiera on images that are even slightly larger than what it was trained on, the accuracy of the resulting model plummets. This includes mediocre results on detection, where ViT outperforms Hiera (Ryali et al., 2023) . The culprit, we discover, is the interaction between window attention and absolute position embeddings. That is, having window attention and absolute position embeddings in the same model * Work done during an internship at Meta.