分享自:

基于RWKV的精确高效体医学图像分割方法U-RWKV

期刊:IEEE Transactions on Image ProcessingDOI:10.1109/tip.2026.3654389

本文提出了一种面向三维医学图像分割的高效且精确的新型网络架构 U-RWKv。该研究由 Hongyu Cai、Yifan Wang、Liu Wang、Jian Zhao 和 Zhejun Kuang 共同完成,作者均来自长春大学计算机科学与技术学院及吉林省人体健康状态辨识功能与增强重点实验室。论文发表于 *IEEE Transactions on Image Processing*,第 35 卷,2026 年,页码 1025 至 1038,数字对象标识符为 10.1109/TIP.2026.3654389。

从学术背景看,三维医学图像分割在临床诊断、术前规划和疾病进展监测中具有重要作用。传统的卷积神经网络(CNN)受限于局部感受野,难以捕获长距离上下文信息;基于 Transformer 的方法虽然具备强大的全局建模能力,但其自注意力机制的计算复杂度随输入分辨率呈二次增长,在处理高分辨率三维医学图像时效率较低。近期提出的 Mamba 模型基于状态空间模型(SSM)实现了线性复杂度的序列建模,但现有方法通常将三维数据展平为一维序列并进行单向扫描,容易破坏三维空间结构的完整性。为解决上述问题,本文探索了一种基于 RWKV(Receptance Weighted Key Value)的新型分割模型。RWKV 通过线性注意力机制 WKV 和 token shift 操作捕获局部上下文,在自然语言处理中已表现出与 Transformer 和 Mamba 相当的性能,同时计算效率更高,但其在三维医学图像分割中的应用尚未被充分探索。因此,本研究旨在构建一个端到端的三维医学图像分割网络,在保证高精度的同时实现高效计算。

在方法设计上,本文提出了 U-RWKV 网络,其整体采用四层编码器—解码器的 U 形结构,并在对应层之间加入跳跃连接。输入三维图像首先经过 patch embedding 操作被划分为不重叠的小块,然后进入由 TSE-R(Tri-directional Spatial Enhancement RWKV)模块构建的编码器—解码器网络。编码器逐层提取特征并进行下采样,解码器通过转置卷积逐层上采样恢复分辨率。最终通过基于卷积的 seg head 生成分割结果。TSE-R 模块是本文的核心创新。该模块首先将输入特征从三维形状重新排列为一维序列,随后进行空间移位(spatial shift)操作。空间移位操作将序列恢复为三维形式,并沿深度、高度和宽度三个轴分别进行移位,再将不同通道段拼接,从而在不引入线性插值的前提下扩大局部感受野并增强相邻块之间的交互。接着,移位后的特征经过三个线性层生成 R、K、V,其中 K 和 V 被送入三方向 WKV(tri-directional WKV)模块。三方向 WKV 将序列还原为原始三维空间,然后沿三个不同方向分别展平为一维序列,并分别进行 WKV 注意力计算,最后通过可学习权重的加权融合得到包含三维空间结构信息的全局注意力输出。该输出经过 sigmoid 调制的 R 进行控制,再经过层归一化得到空间混合结果。随后,特征进入通道混合部分,通过空间移位和两个线性层生成 Rc 和 Kc,Kc 经线性变换映射到隐藏维度后使用平方 ReLU 激活,再映射回原通道维度并与 sigmoid 处理后的 Rc 相乘,最后经过层归一化完成通道混合。在 TSE-R 模块末端,还加入了一个三维卷积模块以保留体素间的局部相关性。

在实验验证方面,研究在三个公开的三维医学图像分割数据集上进行了广泛测试,包括 Synapse 腹部多器官数据集、VerSe2019 脊椎数据集和 BraTS2021 脑肿瘤数据集,另外还使用了 WORD 腹部多器官数据集和 ASOCA 冠状动脉数据集。在 Synapse 数据集上,当裁剪尺寸为 96×96×96 时,U-RWKV 的平均 Dice 达到 85.61%,HD95 为 12.74 mm;当裁剪尺寸增大到 64×192×192 时,平均 Dice 提升至 87.21%,HD95 降至 8.60 mm,优于 nnU-Net 和 nnFormer 等先进方法。在 VerSe 数据集上,U-RWKV 对腰椎 L1–L5 的平均 Dice 为 86.63%,HD95 为 3.47 mm,显著优于 SegMamba 等方法。在 BraTS2021 数据集上,U-RWKV 在整体肿瘤和肿瘤核心的 Dice 分别为 93.53% 和 90.61%,平均 HD95 为 4.64 mm,与针对脑肿瘤专门设计的 CKD-TransBTS 性能相当。在 WORD 数据集的 16 类腹部器官分割中,U-RWKV 的平均 Dice 为 78.59%,HD95 为 28.34 mm,优于所有对比模型。在 ASOCA 冠状动脉分割中,U-RWKV 的 Dice 为 84.89%,HD95 为 18.32 mm,同样取得最佳性能。

在计算效率方面,U-RWKV 仅使用 9.33 M 参数,相比 nnFormer 的 150.05 M 减少了 16.08 倍。在裁剪尺寸为 96×96×96 时,其 FLOPs 仅为 55.29 G,而 nnU-Net 为 476.57 G,nnFormer 为 213.36 G。在较大裁剪尺寸下,U-RWKV 的 FLOPs 也仅为 147.44 G,远低于同类方法。这表明 U-RWKV 在保持高精度的同时具有显著的效率优势。

消融实验进一步验证了各模块的有效性。仅使用三维卷积模块时 Dice 仅为 70.08%,去除空间混合部分后降至 81.84%,去除通道混合部分后为 77.34%,去除空间移位操作后为 83.96%。扫描方向的消融结果显示,单向扫描和双向扫描的 Dice 均低于三方向扫描,说明三方向 WKV 对保留三维空间结构信息至关重要。有效感受野(ERF)可视化分析表明,加入三方向 WKV 和空间移位后,模型的感受野显著扩大且分布更加均匀,从而增强了全局上下文建模能力。在有限标注和噪声监督的鲁棒性实验中,U-RWKV 在仅使用 10%、30%、50% 标注数据时均优于其他方法,且在标签腐蚀、随机翻转和切片缺失等噪声条件下性能下降最小,表现出较强的鲁棒性。

本研究的结论是,基于 RWKV 架构构建的 U-RWKV 模型成功解决了三维医学图像分割中全局上下文建模、计算效率和空间结构保持三个关键问题。TSE-R 模块完全基于 RWKV 架构实现,其计算复杂度与输入数据呈线性关系,使得 U-RWKV 整体计算复杂度也保持线性。三方向 WKV 扫描和加权融合机制有效避免了序列化建模对三维空间结构的破坏。该方法在多个数据集上均优于基于 CNN、Transformer 和 Mamba 的代表性方法,同时参数量大幅减少,为三维分割模型在边缘设备上的高效部署奠定了基础。本文的主要亮点包括:首次将 RWKV 引入三维医学图像分割任务;设计了融合空间移位和三方向 WKV 的 TSE-R 模块;在精度和效率之间取得了优异平衡;并在低标注和噪声监督条件下展现出较强的鲁棒性。此外,研究还分析了失败案例,指出在输入数据包含严重噪声或异常解剖结构时模型性能可能下降,未来工作将致力于设计更鲁棒的模块以应对此类挑战。

上述解读依据用户上传的学术文献,如有不准确或可能侵权之处请联系本站站长:admin@fmread.com