Restore-RWKV：使用 RWKV 实现高效且有效的医学图像复原
Restore-RWKV: Efficient and Effective Medical Image Restoration with RWKV
AI总结：
本文提出首个基于 RWKV 的医学图像复原模型 Restore-RWKV，通过 Re-WKV 与 Omni-Shift 适配二维图像，以线性复杂度和轻量参数实现多任务 SOTA。
摘要：
Transformers 已经彻底改变了医学图像复原，但其二次复杂度仍然限制了它们在高分辨率医学图像中的应用。近来，自然语言处理领域出现的 Receptance Weighted Key Value（RWKV）模型因其高效处理长序列的能力而备受关注。为利用其先进设计，我们提出 Restore-RWKV，这是首个用于医学图像复原的基于 RWKV 的模型。由于原始 RWKV 模型是为一维序列设计的，我们针对二维医学图像中的空间关系建模进行了两项必要修改。首先，我们提出一种循环 WKV（Re-WKV）注意力机制，以线性计算复杂度捕获全局依赖关系。Re-WKV 将双向注意力作为获得全局感受野的基础，并结合循环注意力，以有效建模来自多种扫描方向的二维依赖关系。其次，我们开发了一个全方向 token 移位（Omni-Shift）层，通过从所有方向并跨越广泛上下文范围移位 token 来增强局部依赖关系。这些调整使所提出的 Restore-RWKV 成为一个高效且有效的医学图像复原模型。即使是仅拥有 116 万参数的 Restore-RWKV 轻量变体，也能取得与现有最先进（SOTA）方法相当甚至更优的结果。大量实验表明，所得 Restore-RWKV 在一系列医学图像复原任务上达到 SOTA 性能，包括 PET 图像合成、CT 图像去噪、MRI 图像超分辨率以及一体化医学图像复原。代码可在 https://github.com/Yaziwel/Restore-RWKV 获取。
英文摘要：
Transformers have revolutionized medical image restoration, but the quadratic complexity still poses limitations for their application to high-resolution medical images. The recent advent of the Receptance Weighted Key Value (RWKV) model in the natural language processing field has attracted much attention due to its ability to process long sequences efficiently. To leverage its advanced design, we propose Restore-RWKV, the first RWKV-based model for medical image restoration. Since the original RWKV model is designed for 1D sequences, we make two necessary modifications for modeling spatial relations in 2D medical images. First, we present a recurrent WKV (Re-WKV) attention mechanism that captures global dependencies with linear computational complexity. Re-WKV incorporates bidirectional attention as basic for a global receptive field and recurrent attention to effectively model 2D dependencies from various scan directions. Second, we develop an omnidirectional token shift (Omni-Shift) layer that enhances local dependencies by shifting tokens from all directions and across a wide context range. These adaptations make the proposed Restore-RWKV an efficient and effective model for medical image restoration. Even a lightweight variant of Restore-RWKV, with only 1.16 million parameters, achieves comparable or even superior results compared to existing state-of-the-art (SOTA) methods. Extensive experiments demonstrate that the resulting Restore-RWKV achieves SOTA performance across a range of medical image restoration tasks, including PET image synthesis, CT image denoising, MRI image super-resolution, and all-in-one medical image restoration. Code is available at: https://github.com/Yaziwel/Restore-RWKV.
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
