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 RWKV in the NLP field has attracted much attention as it can 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 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. Extensive experiments demonstrate that Restore-RWKV achieves superior performance across various medical image restoration tasks, including MRI image super-resolution, CT image denoising, PET image synthesis, and all-in-one medical image restoration. Code is available at: \href{https://github.com/Yaziwel/Restore-RWKV.git}{https://github.com/Yaziwel/Restore-RWKV}.


翻译:Transformer模型已经彻底改变了医学图像复原领域,但其二次复杂度在处理高分辨率医学图像时仍存在应用限制。近期,RWKV在自然语言处理领域的出现因其能够高效处理长序列而备受关注。为了利用其先进设计,我们提出了Restore-RWKV,这是首个基于RWKV的医学图像复原模型。由于原始RWKV模型专为一维序列设计,我们针对二维图像的空间关系建模进行了两项必要改进。首先,我们提出了一种循环WKV(Re-WKV)注意力机制,该机制能以线性计算复杂度捕获全局依赖关系。Re-WKV融合了双向注意力作为全局感受野的基础,并通过循环注意力有效建模来自不同扫描方向的二维依赖关系。其次,我们开发了全向令牌移位(Omni-Shift)层,该层通过全方位、宽上下文范围的令牌移位来增强局部依赖关系。这些改进使得所提出的Restore-RWKV成为医学图像复原领域高效且有效的模型。大量实验表明,Restore-RWKV在多种医学图像复原任务中均取得卓越性能,包括MRI图像超分辨率、CT图像去噪、PET图像合成以及一体化医学图像复原。代码发布于:\href{https://github.com/Yaziwel/Restore-RWKV.git}{https://github.com/Yaziwel/Restore-RWKV}。

0
下载
关闭预览

相关内容

FlowQA: Grasping Flow in History for Conversational Machine Comprehension
专知会员服务
34+阅读 · 2019年10月18日
Stabilizing Transformers for Reinforcement Learning
专知会员服务
60+阅读 · 2019年10月17日
《DeepGCNs: Making GCNs Go as Deep as CNNs》
专知会员服务
32+阅读 · 2019年10月17日
Keras François Chollet 《Deep Learning with Python 》, 386页pdf
专知会员服务
164+阅读 · 2019年10月12日
【SIGGRAPH2019】TensorFlow 2.0深度学习计算机图形学应用
专知会员服务
41+阅读 · 2019年10月9日
Transferring Knowledge across Learning Processes
CreateAMind
29+阅读 · 2019年5月18日
强化学习的Unsupervised Meta-Learning
CreateAMind
18+阅读 · 2019年1月7日
Unsupervised Learning via Meta-Learning
CreateAMind
44+阅读 · 2019年1月3日
meta learning 17年:MAML SNAIL
CreateAMind
11+阅读 · 2019年1月2日
A Technical Overview of AI & ML in 2018 & Trends for 2019
待字闺中
18+阅读 · 2018年12月24日
STRCF for Visual Object Tracking
统计学习与视觉计算组
15+阅读 · 2018年5月29日
Hierarchical Imitation - Reinforcement Learning
CreateAMind
19+阅读 · 2018年5月25日
Focal Loss for Dense Object Detection
统计学习与视觉计算组
12+阅读 · 2018年3月15日
IJCAI | Cascade Dynamics Modeling with Attention-based RNN
KingsGarden
13+阅读 · 2017年7月16日
From Softmax to Sparsemax-ICML16(1)
KingsGarden
74+阅读 · 2016年11月26日
国家自然科学基金
2+阅读 · 2017年12月31日
国家自然科学基金
13+阅读 · 2017年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
3+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
VIP会员
最新内容
《美陆军条例:陆军指挥政策(2026版)》
专知会员服务
7+阅读 · 今天8:10
《军用自主人工智能系统的治理与安全》
专知会员服务
5+阅读 · 今天8:02
《系统簇式多域作战规划范畴论框架》
专知会员服务
9+阅读 · 4月20日
高效视频扩散模型:进展与挑战
专知会员服务
4+阅读 · 4月20日
乌克兰前线的五项创新
专知会员服务
8+阅读 · 4月20日
 军事通信系统与设备的技术演进综述
专知会员服务
7+阅读 · 4月20日
《北约标准:医疗评估手册》174页
专知会员服务
6+阅读 · 4月20日
相关资讯
Transferring Knowledge across Learning Processes
CreateAMind
29+阅读 · 2019年5月18日
强化学习的Unsupervised Meta-Learning
CreateAMind
18+阅读 · 2019年1月7日
Unsupervised Learning via Meta-Learning
CreateAMind
44+阅读 · 2019年1月3日
meta learning 17年:MAML SNAIL
CreateAMind
11+阅读 · 2019年1月2日
A Technical Overview of AI & ML in 2018 & Trends for 2019
待字闺中
18+阅读 · 2018年12月24日
STRCF for Visual Object Tracking
统计学习与视觉计算组
15+阅读 · 2018年5月29日
Hierarchical Imitation - Reinforcement Learning
CreateAMind
19+阅读 · 2018年5月25日
Focal Loss for Dense Object Detection
统计学习与视觉计算组
12+阅读 · 2018年3月15日
IJCAI | Cascade Dynamics Modeling with Attention-based RNN
KingsGarden
13+阅读 · 2017年7月16日
From Softmax to Sparsemax-ICML16(1)
KingsGarden
74+阅读 · 2016年11月26日
相关基金
国家自然科学基金
2+阅读 · 2017年12月31日
国家自然科学基金
13+阅读 · 2017年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2015年12月31日
国家自然科学基金
2+阅读 · 2015年12月31日
国家自然科学基金
3+阅读 · 2015年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
国家自然科学基金
0+阅读 · 2014年12月31日
Top
微信扫码咨询专知VIP会员