7. Shifted and Rotated Transformers.
Swin, SwinV2 & Erwin.
Introduction
Transformers have become the go-to method in all the areas of Machine Learning and scientific computing is not an exception. The Attention mechanism proved to be very effective also in this Area. However, its quadratic complexity with respect to the input size makes standard attention impractical for large spatial meshes.
A natural strategy to mitigate this issue is to restrict attention to local sliding windows, similarly to how convolutions operate. While this significantly reduces computational cost, a windowed-attention transformer without additional mechanisms remains strictly local, limiting its ability to capture global interactions.
The Shifted Window Transformer (Swin) addresses this limitation by shifting the window partition across layers, enabling information exchange between neighboring windows and progressively expanding the receptive field. SwinV2 further refines this architecture with improved stability and scaling properties. More recently, Erwin extends the shifted-window paradigm to irregular domains such as point clouds, broadening its applicability to geometric data and numerical simulations on complex geometries.
Windowed attention
Windowed attention is simply standard self-attention applied locally. Instead of computing attention across all tokens (e.g., pixels in an image, mesh points, or patches), the computation is restricted to tokens that lie inside a predefined spatial window.
Therefore, given an input X to the attention layer, to the attention layer, the attention operation is identical to the classical scaled dot-product attention (that we discussed more broadly in the last blog post):
that this operation is applied independently inside windows that partition the input domain. Each window defines a local subset of tokens, and standard self-attention is computed within that subset.
If you are surprised by the B, it’s because it’s not present in the standard attention but the authors found beneficial to add a relative position bias B (to each head) before the softmax.
Complexity: Given a sequence of N tokens in dimension d, global attention has a computational complexity of
While windowed attention, with windows of size w has a complexity of
This follows from the fact that the input is partitioned into N / w windows, and computing attention within each window costs O(w^2 d). Depending on the choice of the window size w, this can be significantly smaller than the quadratic complexity of global attention.
Windowed attention reduces the cost of attention by restricting the set of tokens to which each token pays attention to a window. Without additional modifications, the resulting model will not be able to capture interactions between points that lie in different windows.
A simple fix, would be to introduce overlapping windows by using a stride smaller than the window size. However, I think, they don’t do that because it will increment the cost of attention that is already high (for example by using a stride of half the size of the window the cost will double). Note that, differently, convolutions are very cheap so using stride does not incur in a high additional cost.
Instead, Swin’s authors propose to shift windows: from one layer to the next, each window is shifted by half of the window size. The image below explains it better than I can!
This enables the model to obtain a global receptive field with sufficient depth. In particular, for a sequence of length N and window size w, the model will have
since each shifted-window layer increases the effective communication range by roughly half a window. This doesn’t look great, and I believe this is why Swin adopts a hierarchical design, where larger patches are progressively used from layer to layer:
This significantly accelerates the emergence of a global receptive field, and I believe it is a key ingredient in the success of this model.
SwinV2’s improvements
After the great success of Swin, researchers at Microsoft decided to scale it up and introduced 3 key innovations that enabled training a model with 3 billion parameters, the largest vision transformer at the time (2022), on images with resolutions up to 1536×1536.
1) Cosine Attention
In the standard attention mechanism we discussed earlier, the similarity between pairs of tokens is computed using a dot product. In SwinV2, the authors observed that, in some blocks and attention heads, the learned attention maps were often dominated by a small number of token pairs. To mitigate this effect, they replaced the dot-product similarity with cosine similarity, defined as follows
The cosine function is naturally normalized, and thus can have milder attention values.
2) Post-norm
A residual post-norm replaced the pre-norm present in the previous Swin transformer (V1). Doing so, the output of each residual block is normalized before merging back into the main branch, and the amplitude of the main branch does not accumulate when the layer goes deeper.
3) Continuous relative position bias
Instead of directly learning the position bias, SwinV2 adopts a small meta network on the relative coordinates:
where B is a small network, by default a 2-layer MLP with a ReLU activation in between. The meta network B generates bias values for arbitrary relative coordinates, and thus can be more naturally (more than interpolation techniques previously used) transferred to fine-tuning tasks with arbitrarily varying window sizes.
How we can do the same on point clouds?
Erwin (a.k.a. Rwin (Rotated Windows Transformer)) extends the main ideas of the Swin Transformer to point clouds. Due to the irregular nature of point clouds, spatial windows are no longer appropriate. Instead, they replace windows with local neighborhoods defined as balls in the underlying metric space:
And, analogously to shifting windows in Swin, they rotate the balls between layers to enable interactions across neighboring regions.
Finally, to better capture the multiscale nature of physical phenomena, they adopt a U-Net–like hierarchical structure.
This, together with a great (and very readable) implementation, makes the model — in their own words — blazingly fast.
Conclusion
In this blog, we explored how Swin Transformers efficiently scale attention to large inputs by restricting self-attention to local windows and shifting these windows across layers to enable cross-window communication. We also discussed the improvements introduced in the second version and how Erwin extends these ideas to point clouds. As a reference I recommend the original papers [1, 2, 3] and a blog on Erwin written by the authors [4] that is really well written.