Revisiting Positional Information in Transformers in the era of Fused Attention

We thank the reviewer for their insightful comments.

Regarding novelty, we request the reviewer to take a look at the common response for clarification.

Regarding applicability outside fused attention, we would like to clarify that RoPE and our implementations of RoPE can be used with both fused or unfused attention alike. In fact, we observe significant throughput gains in both cases, as explained in Table 3.

Regarding Table 3, Table 3 shows throughput of the respective models with RPB, Axial RoPE and without any bias. The “No bias” case implies there is no additional computation done in the model to impart positional bias. Thus, it is expected that the models without any bias would be slightly faster, since they don’t perform a chunk of computation as compared to other models with biases.

We will cite the mentioned work in the Related Works section. Thanks for bringing the same to our attention!

We thank the reviewer for their comments and constructive criticism. We would like to clarify the reviewer’s concerns here:

Regarding unclear contribution, we request the reviewer to look at the common response for clarification.

Regarding inconclusive results, shared angles and limited generalization:

Following your suggestion, we are currently performing the experiments with k=8 and shared angles across all heads. We will share an update when we have some numbers regarding the same. At the time of submission, we chose to stick to a reasonable default, since performing experiments with all possible settings is infeasible. We could not include experiments on detection and segmentation owing to limited academic compute.

Following your suggestion, we ran experiments with shared angles for all heads and k=8 for the small variant of all models.

We observed that while NAT and DiNAT do not exhibit any noticeable changes, we observe that Swin-small degrades by 0.3%, while ViT-small improving by 0.1%. From this set of experiments, we conclude that the actual choice of hyperparameters would depend on the model at hand the attention mechanism used in the same. From the above results, we would only hypothesize that since NAT and DiNAT already have an implicit spatial bias built into the attention mechanism, they are not significantly affected by changes in the implementation of RoPE. However, since Swin and ViT do not have any implicit spatial bias, they are more sensitive to changes in RoPE's implementation.

Thanks a lot for your suggestion!

Results from additional ablations:

Table 1: Ablations for small variant of all models with shared PE across all heads and $k_{rope}=8$.

We thank the reviewer for their comments and constructive criticism. We would like to clarify the reviewer’s concerns here:

Regarding weakness (1), we want to clarify Section 3.3.2 is a review of different rotary embedding approaches [1,2]. Our ablations have been limited to k_rope and shared angles in Tables 4 and 5, but following your suggestion, we will include ablations of all design choices mentioned in Section 3.3.2. Thank you for bringing this to our attention!

Regarding weakness (2), since we were limited to image classification due to compute constraints, we report ImageNet-1k accuracy and inference throughput for different positional biases. It is indeed commonly accepted that RoPE has equal or better accuracy than RPB, and thus we focus our efforts on showing why RoPE is more scalable than RPB.

Regarding weakness (3), we would like to point out that even though the gains seem marginal, they cannot be achieved trivially. The gains also support the hypothesis that RoPE consistently outperforms RPB, while being hardware-friendly at the same time.

Regarding weakness (4), we request the reviewer to look at the common response for clarification.

Regarding question (1), we hypothesize that RoPE outperforms RPB (in terms of accuracy) because it simply provides the model a spatial bias, without having the need to learn the same. In contrast to this, RPB requires the model to learn this spatial bias itself, which might not be optimal given the data-hungry nature of transformer-based vision models.

Regarding question (2) yes, we expect the gains to effectively scale to larger models and resolutions. As the model size and resolution increases, RPB will indeed become a larger bottleneck due to the nature of fused attention mechanisms. In addition, RoPE will continue to scale up training efficiently as well, since it is a simple elementwise operation in both the forward and backward pass. Attention biases become a much larger bottleneck in training, since their backward pass is a reduction operation, and typically bound by memory bandwidth. We can see this in comparative training throughputs of RoPE and RPB models. For example, a ViT Base model with RPB has a training throughput of 6071 images/sec, whereas a ViT Base with Axial RoPE has a training throughput of 6596 images/sec, thereby observing a training speedup of ~8.5%.

References:

[1] Crowson, Katherine, et al. "Scalable high-resolution pixel-space image synthesis with hourglass diffusion transformers." Forty-first International Conference on Machine Learning. 2024.

[2] Heo, Byeongho, et al. "Rotary position embedding for vision transformer." arXiv preprint arXiv:2403.13298 (2024).

Following your suggestion, we perform ablations of all design choices in Table 1 -- position co-ordinates and angle generator. We perform the ablations while keeping the settings on Axial and 2D RoPE for the other two design choices (k and sharing of angles across heads). We will add the same in the manuscript.

Thanks a lot for the suggestion!

Results from additional ablations:

Table 1: Ablation on angle generator and position co-ordinates with non-shared angles and $k_{rope}=2$.

Table 1: Ablation on angle generator and position co-ordinates with shared angles and $k_{rope}=1$.