ID-RWKV: Image Deraining RWKV

Abstract

Although Transformer has achieved considerable results in image deraining tasks, the quadratic complexity of self-attention in this structure limits its ability to process high-resolution rainy images. The Receptance Weighted Key Value (RWKV) in the field of natural language processing (NLP) solves the expensive computational cost of self-attention in Transformers by using linear computational complexity to learn long-range dependencies of features. Based on advanced design of RWKV, we propose an image deraining RWKV (ID-RWKV) that gradually obtains fine rainless images by constructing multiple deraining stages. At each stage, a U-shaped rain removal network (U-RRNet) is constructed to encode and decode image features, and output the rain removal results for the current stage. Each layer of U-RRNet consists of Fourier enhancement module (FFM) and local-global RWKV Block (LGRB). FEM is constructed to extract and enhance the features in the frequency domain. LGRB is designed based on RWKV structure to improve the performance and efficiency of deraining models. To better learn local and global contextual information, we propose a LG-WKV attention mechanism to enhance local and global dependencies. To reduce the loss of background information and increase the stability of the network, we construct a deep-shallow feature fusion module (DSFFM) to supplement shallow features. Following extensive experimentation, we demonstrate that our method not only outperforms the current state-of-the-art methods, but also requires fewer parameters and less computation than the Transformer-based method.