RWKV再下一城！Restore-RWKV：图像恢复全新模型

点击下方卡片，关注“CVer”公众号

AI/CV重磅干货，第一时间送达

点击进入—>【Mamba/多模态/扩散】交流群

添加微信号：CVer111，小助手会拉你进群！

扫描下方二维码，加入CVer学术星球！可以获得最新顶会/顶刊上的论文idea和CV从入门到精通资料，及最前沿应用！发论文/搞科研/涨薪，强烈推荐！

作者：夫未（已授权CVer转载）

https://zhuanlan.zhihu.com/p/709498751

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

代码：https://github.com/Yaziwel/Restore-RWKV

论文：https://arxiv.org/abs/2407.11087

单位：北航，清华，北京协和医院，字节跳动

全局建模能力在图像恢复任务中至关重要，因为图像恢复的一个核心在于从图像本身寻找有用的token信息来弥补当前token的信息缺失。然而，RWKV/Mamba/xLSTM等模型起源于NLP领域，因此它们的原始算子主要用于处理1D序列。由此，其后续视觉相关的模型的研究重点在于如何将其1D序列建模能力扩展到2D图像上。今天介绍的是Restore-RWKV是第一篇将RWKV应用于图像恢复任务的论文，其依然沿用了1D—>2D拓展的思路讲故事。

动机

由于图像恢复任务是pixel-level的像素值回归，那么其每个pixel都是一个token，而医学图像通常分辨率较高（例如CT 512×512），如果使用self-attention进行全局建模，那么消耗的计算资源是无法忍受的。所以现有的很多方法都采样efficient版本的self-attention，核心思路是缩小attention的计算范围，最典型的是swin transformer的window attention，但是这样做实质上丧失了全局建模能力。当然Mamba也是一种选择，其SSM具有线性复杂度的长距离序列建模能力。然而目前使用的SSM是单向的，其感受野范围局限于从序列的第一个token到当前token这个范围，虽然通过不同方向的扫描可以最终保证其全局建模能力，但这样做效果仍然是欠缺的。我们举例可视化了不同方法的有效感受野如下：

可以看到，只有基于Mamba的MambaIR和本文提出的基于RWKV的Restore-RWKV实现了全局感受野，而且Restore-RWKV的有效感受野更好。那么这是如何实现的呢？

方法

原始的RWKV有两个核心创新，一是WKV attention用于捕获长距离信息，二是token shift用于捕获局部信息。当然，它们都是用于处理1D序列的，为了实现2D图像的建模，本文分别对两种机制进行改进，并提出reccurent WKV attention (Re-WKV)和omnidirectional token shift (Omni-Shift)，并将其应用于图像恢复网络中，网络结构如下：

接下来分别介绍①Re-WKV和②Omni-Shift：

①Re-WKV

Re-WKV的目的是以线性复杂度实现2D图像全局建模。它由两部分组成，首先其沿用Vision-RWKV论文中提出的bidirectional WKV attention (Bi-WKV)算子，将原始单向WKV attention的有限视野扩充到全局视野；然后其通过循环attention的方式，从不同扫描方向上进一步促进2D图像全局token之间的交互。

②Omni-Shift

Omni-Shift作用是捕获局部信息，原始的RWKV的token shift是单向的（将上一个token信息转移到当前token），Vision-RWKV的改进版本是四向（将上下左右四个token信息转移到当前token），但是很显然这对2D图像来说都不是最好的token shift策略，2D图像的局部性应该是全向的，邻近区域的各个方向的token都是相关的，所以最好的token shift应该就是卷积。于是本文采用depth-wise卷积，并且考虑到不同context范围的token重要性不一样，在训练阶段仿照RepVGG利用了结构重整化，使用多个卷积分支，每个分支负责学习一定context范围的token shift，在测试时，这些分支合并成一个卷积：

Restore-RWKV在使用Re-WKV+Omn-Shift组合的改进之后，相比于原始RWKV的单向的Uni-WKV+Uni-Shift组合，以及Vision-RWKV的Bi-WKV和Quad-Shift组合，其模型的感受野大大提升：

实验结果

最后，论文在MRI image super-resolution、CT image denoising、PET image synthesis以及all-in-one medical image restoration任务上验证了其有效性。可视化对比如下：

总结与思考

本质是寻求将1D的序列模型拓展到2D。token shift使用卷积是完美的2D拓展操作。但是WKV attention部分仍然需要进行扫描操作，比较丑陋，如果能够通过改写 CUDA，将其改为 2D WKV attention 操作，那就完美了。

另外RWKV或者Mamba为什么会在图像恢复上work？我的理解是①利用卷积保证局部信息的利用(卷积十分重要！！！)；②WKV attention或者SSM提供了实现长距离依赖的途径，真正有用的信息依然可以根据这个途径传播到当前token对其进行信息补充。所以RWKV和Mamba相当于是在CNN和Transformer之间进行取舍之后寻找的一个平衡点。

何恺明在MIT授课的课件PPT下载

在CVer公众号后台回复：何恺明，即可下载本课程的所有566页课件PPT！赶紧学起来！

ECCV 2024 论文和代码下载

在CVer公众号后台回复：ECCV2024，即可下载ECCV 2024论文和代码开源的论文合集

CVPR 2024 论文和代码下载

在CVer公众号后台回复：CVPR2024，即可下载CVPR 2024论文和代码开源的论文合集

Mamba、多模态和扩散模型交流群成立

扫描下方二维码，或者添加微信号：CVer111，即可添加CVer小助手微信，便可申请加入CVer-Mamba、多模态学习或者扩散模型微信交流群。另外其他垂直方向已涵盖：目标检测、图像分割、目标跟踪、人脸检测&识别、OCR、姿态估计、超分辨率、SLAM、医疗影像、Re-ID、GAN、NAS、深度估计、自动驾驶、强化学习、车道线检测、模型剪枝&压缩、去噪、去雾、去雨、风格迁移、遥感图像、行为识别、视频理解、图像融合、图像检索、论文投稿&交流、PyTorch、TensorFlow和Transformer、NeRF、3DGS、Mamba等。
一定要备注：研究方向+地点+学校/公司+昵称（如Mamba、多模态学习或者扩散模型+上海+上交+卡卡），根据格式备注，可更快被通过且邀请进群

▲扫码或加微信号: CVer111，进交流群
CVer计算机视觉（知识星球）来了！想要了解最新最快最好的CV/DL/AI论文速递、优质实战项目、AI行业前沿、从入门到精通学习教程等资料，欢迎扫描下方二维码，加入CVer计算机视觉（知识星球），已汇集上万人！

▲扫码加入星球学习

▲点击上方卡片，关注CVer公众号
整理不易，请赞和在看