To overcome the quadratic complexity of standard self-attention [1]}, we alternately use axial attention modules [2]}, [3]} and standard attention modules as illustrated in Fig. REF (top-left). The axial attention module can be implemented easily by adjusting the tensor shape for row-wise/column-wise and then processing them with dot product-based self-attention respectively.
To improve the spatial relation, we also provide Relative Position Encoding (RPE) [4]} for each axial-attention module. For the input feature \(\mathbf {X}\in \mathbb {R}^{h\times w\times c}\) , we suppose that \(\mathbf {x}_{ri,rj}, \mathbf {x}_{ci,cj}\in \mathbb {R}^{c}\)  mean feature vectors of rows \(i,j\)  and columns \(i,j\)  of \(\mathbf {X}\) .
Then the row and column-wise RPE based axial attention scores \(\mathbf {A}^{row}, \mathbf {A}^{col}\)  can be written as
\(\begin{split}\mathbf {A}_{i,j}^{row}&=\mathbf {x}_{ri}\mathbf {W}_{rq}\mathbf {W}_{rk}^T\mathbf {x}_{rj}^T+R_{i,j}^{row},\\\mathbf {A}_{i,j}^{col}&=\mathbf {x}_{ci}\mathbf {W}_{cq}\mathbf {W}_{ck}^T\mathbf {x}_{cj}^T+R_{i,j}^{col},\end{split}\) 
