Paper: Window Attention is Bugged: How not to Interpolate Position Embeddings (ICLR-2024)

Reviewer comment (Theory):
The whole idea is based on the visualization of the position embedding, but whether the visualization can represent position information is a question. The draft should further explain this issue to make the paper clear.

Author response:
**W1. The Motivation.**  
These are good points, but we would like to emphasize that the link between position embeddings and attention is broadly treated as common knowledge in the transformer literature (we provide a derivation of this in A below), and as mentioned above, the bug is not that the position embeddings are tiled, but that _not accounting for that tiling when interpolating_ causes an issue (see B below). We elaborate on both points below.

**_A.) It’s not clear that the position embeddings we visualize and their actual effect on the network after passing through multiple layers are the same._**  
They aren’t exactly the same, but they _are_ inherently linked. Note that analyzing the values of position embedding to infer its effect on the rest of the network is a common practice in the ViT literature (e.g., [1,2]). Thus, proving this connection is out of the scope of this work, but we think this is an important discussion, so we’ve described our perspective here.

The amount of processing applied to the features after adding the position embedding is minimal in the early layers. Mathematically, the implementation of Hiera / ViT proceeds like this for each token $x_i$ after the patchification step with position encoding for that token $p_i$: $x'_i = x_i + p_i$

This $x'_i$ gets immediately passed into the first block of the network and thus the first attention layer. For the attention matrix (ignoring heads, as they don’t affect the result), it constructs $q_i$ and $k_i$ as $q_i = x'_i W_q^T$ and $k_i = x'_i W_k^T$. The (i, j)th element of the attention matrix pre-softmax is then

$$q_i k_j^T=x'_i W_q^T W_k {x'}_j^T$$

For brevity, let $W_{qk} = W_q^T W_k$. Then, if we substitute back $x_i$, we get:

$$q_i k_j^T = \left(x_i + p_i\right)W_{qk} \left(x_j + p_j\right)^T$$

When we expand the above expression, we get 4 terms---a term without posemb, 2 cross terms, and the following explicit position embedding term:

$$ q_i k_j^T = … + p_i W_{qk} p_j^T $$

Note that the same $W_{qk}$ is being used for each position here, so the position embeddings are the only parts of this term that vary spatially. Thus, any repeated pattern we see in the actual values of $p_i$ and $p_j$ (what we visualizing in Fig. 2, with similarity plotted in Fig. 5) is very closely reflected in this first layer of the network, and because of skip connections, also in subsequent layers. This is consistent with the commonly held belief that position embeddings have more effect at the start of the network than at the end.

Again, we view this to be out of scope for this work (as this kind of analysis is common), but we’ve added this discussion to the appendix as Appendix E.1 to further solidify our claims. We also have a similar derivation for the bug in Appendix E.2 and an explanation of how absolute win fixes the bug in Appendix E.3. Please take a look if you’re interested.

Reviewer comment (Theory):
The proposed method utilizes the "bug" of the learned window position embedding, which raises the question of whether the learned embedding is indeed a bug.

Author response:
**W2. The fix.**  
Again, we would like to point out that the bug is in a mismatch between the size of windows in the position embedding and the size of window attention. Relative position embeddings can recover the lost accuracy, but this is at the cost of a large speed penalty that absolute win avoids. Furthermore, absolute win leans into the _tiling behavior_ of window attention, not the bug. As explained in Appendix E.3 and shown visually in Fig. 1, it explicitly fixes the bug. We elaborate on these points below.

**_A.) Hiera removed relative position embeddings, adding them back could fix it._**  
This is very true, and is actually what the original ViTDet does. Because of the bug, the original ViTDet requires relpos to get good accuracy. Which is fine for accuracy, but relpos is significantly slower than not having relpos. If we are just careful about how we define position embeddings (abs win in our paper), we can completely avoid having to add relpos and take the speed penalty _and_ get a further accuracy boost. The same would likely be true for Hiera, adding back the relpos might alleviate some of the accuracy drop due to the bug, but it would make the model significantly slower and just like ViTDet, the accuracy would likely still be worse.

Regardless, this is exactly the kind of situation we want to help avoid with this paper. By showing that this bug is a possibility, we hope that the community can avoid future issues like this when designing new models and architectures.

**_B.) If the method is proposed based on the "bug" not fixing it, is it still a "bug"?_**  
We think there may be some confusion here. We _do_ fix the bug, as very clearly shown empirically through our comparisons to the original model (and we further elaborate in the updated Appendix E.3). The bug is as shown in Fig. 1, an induced misalignment between the size of windows in the position embedding and the size of windows in window attention. By explicitly windowing our position embedding (absolute win), we can exactly control these two windows sizes to perfectly align at any image resolution, which means that there’s no mismatch and thus no bug.

Reviewer comment (Theory):
The whole problem (bug) is found and discussed based on the "learned" feature embedding, as shown in Figs. 1 and 2 in the draft. This example plays a foundational role in the study, but why it is a bug and what the fix is need clarification. It is unclear why the visualization can be considered position information. Since some learned position embeddings go through or are fused with semantic features, it is important to confirm whether the defined position embedding is really position embedding. The learned block could also contain semantic features, which could be the reason behind the observed phenomenon.

Author response:
**Clarification: What the bug is.**  
We believe there might be some confusion as to what the bug actually is. We do not claim that the _windowed position embeddings_ (of e.g., Hiera) is the bug. Instead, the bug is _not accounting for these windows_ when interpolating the position embedding. This causes a mismatch between the windows of the position embeddings and the windows in attention that results in poor performance. It is this _mismatch_ that’s the bug (see Fig. 1), and we fix it by carefully accounting for the windows in each case. We have further elaborated on this in Appendix E.2 in the updated draft of the paper.

Reviewer comment (Experiments):
Assuming it is position embedding, how can one confirm the bug is from window-attention + absolute position? Could the learned embedding be the main reason instead? In other words, if pre-defined absolute position information is used with window-attention, would the bug still exist? The performance may drop, but the claim was on exposing the bug.

Author response:
**_B.) Assuming position embeddings are the source of the bug, it’s not clear that this is because of window attention—it could be learned position embeddings instead._**  
The reviewer provides a potential counterexample experiment, e.g., if we train a Hiera model with a fixed absolute position embedding (e.g., sinusoidal), would we still have this bug? The answer is (likely), no. In fact, we already have a similar ablation in Tab. 4 (“global only”). This is because the bug is specifically in the mismatched window sizes between these learned position embeddings and the actual windows used when running the model on bigger images. If the position embedding is explicitly not windowed, the model will probably learn a way around this (but with likely lower accuracy, as the reviewer mentioned).

But, that doesn’t mean you can just ignore the bug and use fixed embeddings. Specifically, the ViTDet models actually do use fixed sinusoidal embeddings (from the pretrained MAE models). And yet they still suffer from the bug, not because the position embedding is windowed, but because they specifically add windows which now cease to align with the pre-existing position embedding. So in our method, we manually add windows to the position embedding by tiling them to fix it.

This is the core of the bug, which is described in Fig. 1—a mismatch between windows. The primary goal of this work is to alert the community of this bug so that they can avoid it in their own designs. We hope the reviewer can agree that we provide a significant amount of empirical verification that the bug does exist and that our method can correct it in both Hiera and ViTDet.

Reviewer comment (Writing):
A further question arises: if the method is proposed based on the "bug" but not actually fixing it, is it still a "bug"?

Author response:
**_B.) If the method is proposed based on the "bug" not fixing it, is it still a "bug"?_**  
We think there may be some confusion here. We _do_ fix the bug, as very clearly shown empirically through our comparisons to the original model (and we further elaborate in the updated Appendix E.3). The bug is as shown in Fig. 1, an induced misalignment between the size of windows in the position embedding and the size of windows in window attention. By explicitly windowing our position embedding (absolute win), we can exactly control these two windows sizes to perfectly align at any image resolution, which means that there’s no mismatch and thus no bug.

Reviewer comment (Evaluation):
The gains from fixing this bug seem somewhat task/architecture dependent (smaller gains on ViTDet).

Author response:
This is a good observation, and there’s a good reason for it. For ViTDet, the bug actually does also cause a 1 mAP drop (similar to HieraDet), but this drop is _hidden_ behind ViTDet’s choice to add extra relative position embeddings. In the original ViTDet paper [1], the authors make a comment (page 11 on the arxiv version) that this addition improved AP by ~1 point. What we show in this work is that by fixing the bug using a more principled approach, we can get that same 1 mAP back (and even more) without those expensive relative position embeddings (thereby increasing inference speed by ~40%). So, in reality, the bug itself has a similar effect in most cases—it’s just that ViTDet’s original design obscured that effect.

[1] Exploring Plain Vision Transformer Backbones for Object Detection. ECCV 2022.

Reviewer comment (Experiments):
Unclear if the bug manifests in other architectures with window attention and absolute embeddings.

Author response:
This is something we also don’t know for sure. But our goal is not necessarily to expose an issue in _every_ method out there. Instead, using two comprehensive recent examples (Hiera and ViTDet), we want to make the fact that this issue exists abundantly clear to the community so that they can fix it or avoid causing it in their own work—pushing the SotA is just a bonus.

Reviewer comment (Experiments):
Does not explore more powerful global position embedding strategies for ViTDet.

Author response:
The scope of this work is primarily about the bug and fixing it. More study into whether relpos is best for the global embedding for ViTDet (other than the existing Tab. 4 ablation) would have been interesting, but we didn’t want to stray too far from the original ViTDet model in practice.

Reviewer comment (Experiments):
Does relative position encoding exhibit similar phenomena as absolute position encoding?

Author response:
This is a good question and a worthy line of inquiry for future work. Without such an investigation, we can only speculate, but relpos is likely a little more resilient to this issue than absolute embeddings. Fundamentally, relpos recomputes and resamples the position embedding depending on the _relative_ position between two tokens. Since there’s a layer of abstraction between the actual position embeddings and the positions of the tokens themselves, it can be recomputed correctly for different image sizes. However, this abstraction is also what makes relpos slow—so, there’s a tradeoff there.

Reviewer comment (Experiments):
The proposed fix of "absolute win" embeddings is intuitive, but the paper does not provide much analysis into why tiling the window embeddings specifically works better than other potential solutions. Some ablation studies on the design choices could add more clarity.

Author response:
In the updated draft, we’ve backed up this intuition with mathematical justification (see Appendix E), which should hopefully show why this is a natural choice to fix the bug. We also have ablations in Tab. 2 that show two different methods to fixing the bug: tiling (window only) and using a position embedding small enough such that all tokens in a window map to the same location in the embedding (global only). Both work to fix the issue, but because Hiera has both window attention layers and global attention layers, we need to add both embeddings to retain the accuracy of the original model at the original resolution.

Reviewer comment (Novelty):
The simplicity of the proposed fix somewhat limits the depth of technical novelty and contribution. Providing more design motivation and analysis for "absolute win" would strengthen this aspect.

Author response:
The design for absolute win was simple to attempt to match the existing position embeddings as much as possible while still fixing the bug—specifically so that the fix could be used as a simple drop in replacement. A more complicated approach would have impeded this goal. As for motivation and analysis, we have provided theoretical justifications in Appendix E.2 and E.3 of the updated draft.

Reviewer comment (Experiments):
Tiling the window position embeddings makes sense based on the observed tiling behavior, but other arrangements could also resolve the misalignment issue when interpolating. The design space is not explored in depth.

Author response:
The goal of this work is to identify this issue and propose a fix so that the broader community can avoid the bug when designing their own models. To that end, we chose the simplest option that deviates the least from the original design of the model. We agree that there likely are other designs that could also fix the bug. However, we believe this kind of analysis is better saved for follow-up work. We have included additional discussion about this in Appendix E.3.

Reviewer comment (Theory):
While shown to work empirically, the approach lacks theoretical analysis into why tiling embeddings avoids the misalignment issue during interpolation. A more formal understanding could improve the method.

Author response:
Thanks for the suggestion. We’ve added a 2 page Appendix E to the updated draft of the paper, which attempts to more rigorously derive position embedding’s effect on attention, how a mismatch between periodicity in the position embedding and window attention causes this misalignment issue, and how absolute win fixes it. Along with this derivation come findings that the bug is most detrimental when the finetuning resolution is slightly different than the pretraining one, which explains the empirical behavior of the original model, like in Tab. 8. Please take a look at the updated draft if interested.