Shuffle Transformer

The Shuffle Transformer Block consists of the Shuffle Multi-Head Self-Attention module (ShuffleMHSA), the Neighbor-Window Connection module (NWC), and the MLP module. To introduce cross-window connections while maintaining the efficient computation of non-overlapping windows, a strategy which alternates between WMSA and Shuffle-WMSA in consecutive Shuffle Transformer blocks is proposed. The first window-based transformer block uses regular window partition strategy and the second window-based transformer block uses window-based selfattention with spatial shuffle. Besides, the Neighbor-Window Connection moduel (NWC) is added into each block for enhancing connections among neighborhood windows. Thus the proposed shuffle transformer block could build rich cross-window connections and augments representation. Finally, the consecutive Shuffle Transformer blocks are computed as:

xˡ=𝐖 𝐌 𝐒 𝐀(𝐁 𝐍(zˡ⁻¹))+zˡ⁻¹

yˡ=𝐍 𝐖 𝐂(xˡ)+xˡ

zˡ=𝐌 𝐋 𝐏(𝐁 𝐍(yˡ))+yˡ

xˡ⁺¹=𝐒 𝐡 𝐮 𝐟 𝐟 𝐥 𝐞 - 𝐖 𝐌 𝐒 𝐀(𝐁 𝐍(zˡ))+zˡ

yˡ⁺¹=𝐍 𝐖 𝐂(xˡ⁺¹)+xˡ⁺¹

zˡ⁺¹=𝐌 𝐋 𝐏(𝐁 𝐍(yˡ⁺¹))+yˡ⁺¹

where xˡ, yˡ and zˡ denote the output features of the (Shuffle-)WMSA module, the Neighbor-Window Connection module and the MLP module for block l, respectively; WMSA and Shuffle-WMSA denote window-based multi-head self-attention without/with spatial shuffle, respectively.