SWSformer: Sub-Window Shuffle Transformer for Image Restoration

Transformers have demonstrated superior performance over conventional methods in various tasks due to their ability to capture long-range dependencies and adaptively generate weights. However, their computational complexity increases quadratically with the number of tokens, limiting their applicability to high-resolution image tasks. Recent image restoration methods mitigate this by adopting mechanisms that sacrifice the ability to fully capture long-range dependencies. In contrast, Shuffle Transformer's window shuffle self-attention (WS-SA) is a mechanism that allows not only suppressing the computational complexity but also not restricting long-range dependency capture. Nevertheless, WS-SA has a problem that it causes a very sparse SA with a small receptive field for high-resolution images. In this work, we propose sub-window shuffle self-attention (SWS-SA), which is a mechanism to expand the receptive field without changing the computational complexity of WS-SA. In SWS-SA, the windows before shuffling are further partitioned into sub-windows, and shuffle is applied to them on a sub-window basis. After that, average pooling per sub-window is applied only to queries. Furthermore, we propose a Sub-Window Shuffle Transformer (SWSformer) that takes SWS-SA as its main mechanism and incorporates effective mechanisms proposed in related works. SWSformer achieves the same or better performance than state-of-the-art methods on denoising and deblurring tasks.