Restore-RWKV: Efficient and Effective Medical Image Restoration with RWKV

变压器技术已经彻底改变了医学图像恢复的方法，但是二次复杂度仍然限制了它们在高分辨率医学图像上的应用。最近，在NLP领域中，RWKV的出现引起了广泛关注，因为它可以高效地处理长序列。为了利用其先进的设计，我们提出了Restore-RWKV，这是第一个基于RWKV的医学图像恢复模型。由于原始的RWKV模型是为1D序列设计的，我们对其进行了两个必要的修改，以建模2D图像中的空间关系。首先，我们提出了一种递归WKV（Re-WKV）注意机制，以线性计算复杂度捕获全局依赖关系。Re-WKV将双向注意力作为全局感受野的基础，并采用循环注意力来有效地建模来自各种扫描方向的2D依赖关系。其次，我们开发了一个全向令牌移位（Omni-Shift）层，通过从所有方向和跨越广泛的上下文范围移动令牌来增强局部依赖关系。这些改进使得所提出的Restore-RWKV成为一种高效而有效的医学图像恢复模型。广泛的实验表明，Restore-RWKV在各种医学图像恢复任务中均表现出优异的性能，包括MRI图像超分辨率、CT图像去噪、PET图像合成和全能医学图像恢复。代码可在以下网址获得：\href{https://github.com/Yaziwel/Restore-RWKV.git}{https://github.com/Yaziwel/Restore-RWKV}。

论文旨在解决高分辨率医学图像恢复中变压器的二次复杂度限制问题，提出了一种基于RWKV的医学图像恢复模型。

提出了一种适用于2D图像的Re-WKV注意机制，用于捕捉全局依赖关系，并且提出了一种Omnidirectional Token Shift（Omni-Shift）层，用于增强局部依赖关系。这些改进使得提出的Restore-RWKV模型在医学图像恢复中具有高效和有效的性能。

论文在MRI图像超分辨率、CT图像去噪、PET图像合成和全方位医学图像恢复等多个医学图像恢复任务上进行了广泛的实验，并展示了Restore-RWKV模型在这些任务上的卓越性能。代码已经开源。

最近的相关研究包括使用变压器进行医学图像恢复的先前工作，以及在NLP领域中使用RWKV的研究。