Restore-RWKV：基于RWKV的高效医学图像修复方法
图像与视频处理
2025-01-07 v3 计算机视觉与模式识别
摘要
Transformer已在医学图像修复领域取得显著进展，但其二次复杂度仍限制了在高分辨率医学图像上的应用。近期在自然语言处理领域出现的Receptance Weighted Key Value（RWKV）模型因其高效处理长序列的能力而引起广泛关注。为充分利用其先进设计，我们提出了第一个基于RWKV的医学图像修复模型——Restore-RWKV。由于原始RWKV模型设计用于1D序列，本文对2D医学图像的空间关系建模进行了两个必要的修改。首先，我们提出了循环WKV（Re-WKV）注意力机制，通过线性计算复杂度捕获全局依赖性。Re-WKV将双向注意力作为基本组件，以实现全局感受野，并利用循环注意力有效地从各种扫描方向建模2D依赖性。其次，我们开发了全向token位移（Omni-Shift）层，通过从所有方向以及跨越宽广上下文范围的位移来增强局部依赖性。这些改编使所提出的Restore-RWKV成为医学图像修复高效且有效的模型。即使是仅包含116万参数的轻量级变体，也可实现与现有最先进（SOTA）方法相当甚至更好的结果。广泛的实验表明，所得到的Restore-RWKV在包括PET图像合成、CT图像去噪、MRI图像超分辨率以及全方位医学图像修复在内的一系列医学图像修复任务中实现了SOTA性能。代码已公开：https://github.com/Yaziwel/Restore-RWKV。
引用
@article{arxiv.2407.11087,
title = {Restore-RWKV: Efficient and Effective Medical Image Restoration with RWKV},
author = {Zhiwen Yang and Jiayin Li and Hui Zhang and Dan Zhao and Bingzheng Wei and Yan Xu},
journal= {arXiv preprint arXiv:2407.11087},
year = {2025}
}
备注
This paper introduces the first RWKV-based model for image restoration