where \(\mathbf {W}_{rq},\mathbf {W}_{rk},\mathbf {W}_{cq},\mathbf {W}_{ck}\)  are trainable parameters for query and key in row and column; \(R_{i,j}^{row}\)  is the trainable RPE value between row \(i\)  and \(j\) , and \(R_{i,j}^{col}\)  means the RPE value between columns \(i,j\) . Then, the attention scores are processed by the softmax operation.
To stabilize the training, we advocate the pre-norm trick in [1]}.
Compared with the \(O(n^2)\)  complexity of the standard self-attention, the axial attention only has \(O(2n^{\frac{3}{2}})\) , which allows us can handle more attention layers for a better model capacity.
Additionally, we also remain some vanilla attention modules in TSR for learning better global correlation.
