Restore-RWKV: Efficient and Effective Medical Image Restora-tion with RWKV

1.Innovations

1. 模型创新
   • 提出 Restore-RWKV，首次将 RWKV 模型应用于医学图像恢复领域，为医学图像恢复提供了一种高效且有效的骨干模型，填补了 RWKV 在该领域应用的空白。
2. 注意力机制创新
   • 提出 Recurrent WKV (Re-WKV) 注意力机制，该机制通过整合双向注意力实现全局感受野，并利用循环注意力更好地适应二维图像的依赖关系建模，以线性计算复杂度有效捕获高分辨率医学图像中的全局依赖。
   • 与原始 RWKV 中的单向 WKV (Uni-WKV) 注意力机制相比，Re-WKV 解决了其在处理二维图像时感受野有限的问题，能够更好地捕捉图像中的长距离依赖关系，从而提升模型在医学图像恢复任务中的性能。
3. Token Shift 层创新
   • 开发 Omnidirectional Token Shift (Omni-Shift) 层，与以往仅从有限方向移动 token 的机制不同，Omni-Shift 通过卷积从所有方向移动和融合相邻 token，增强了局部依赖关系的捕获能力。
   • 在训练时采用多分支结构学习不同上下文范围的准确 token 移动，测试时利用结构重参数化策略保证效率，从而提升模型性能，使模型能够更有效地利用局部上下文信息，进一步提高医学图像恢复的质量。

2.Methodology

2.1 Architecture

2.2 Recurrent WKV

          Δdir（·）分别代表 H-Scan 和 V-Scan 操作，H-Scan 和V-Scan 交替出现

2.3 Omnidirectional Token Shift

Training：

Testing:

       与单分支相比，多分支设计不可避免地会引入额外的参数和计算成本。为了缓解这个问题，Omni-Shift 采用了 RepVGG 中使用的方法，即在测试期间进行结构重参数化。由于较小核的卷积权重可以通过对较小核的权重进行零填充来与较大核的卷积权重整合，因此在测试期间，Omni-Shift 的多分支设计可以合并为一个使用卷积核大小为 5 的单分支（如图 4 (b) 所示）。这种重参数化策略保证了 Omni-Shift 的准确性和测试时的效率。

3.Experiments