[论文解读] U-Net vs Transformer: Is U-Net Outdated in Medical Image Registration?

Xi Jia, Joseph W. Bartlett+4 more|arXiv (Cornell University)|Aug 7, 2022

Brain Tumor Detection and Classification被引用 5

一句话总结

该论文提出 LKU-Net，一种通过引入大感受野卷积模块改进的 U-Net 架构，以增强感受野，在 3D 医学图像配准任务中表现出优于基于 Transformer 的先进方法（如 TransMorph）的性能。尽管参数量仅为 TransMorph 的 1.12%，乘加操作量仅为其 10.8%，LKU-Net 在基于图谱和跨被试者的配准任务中均实现了更优性能，挑战了 U-Net 在该领域已过时的观点。

ABSTRACT

Due to their extreme long-range modeling capability, vision transformer-based networks have become increasingly popular in deformable image registration. We believe, however, that the receptive field of a 5-layer convolutional U-Net is sufficient to capture accurate deformations without needing long-range dependencies. The purpose of this study is therefore to investigate whether U-Net-based methods are outdated compared to modern transformer-based approaches when applied to medical image registration. For this, we propose a large kernel U-Net (LKU-Net) by embedding a parallel convolutional block to a vanilla U-Net in order to enhance the effective receptive field. On the public 3D IXI brain dataset for atlas-based registration, we show that the performance of the vanilla U-Net is already comparable with that of state-of-the-art transformer-based networks (such as TransMorph), and that the proposed LKU-Net outperforms TransMorph by using only 1.12% of its parameters and 10.8% of its mult-adds operations. We further evaluate LKU-Net on a MICCAI Learn2Reg 2021 challenge dataset for inter-subject registration, our LKU-Net also outperforms TransMorph on this dataset and ranks first on the public leaderboard as of the submission of this work. With only modest modifications to the vanilla U-Net, we show that U-Net can outperform transformer-based architectures on inter-subject and atlas-based 3D medical image registration. Code is available at https://github.com/xi-jia/LKU-Net.

研究动机与目标

• 探究基于 U-Net 的方法在医学图像配准任务中是否已落后于现代基于 Transformer 的架构。
• 评估由 Transformer 建模的长程依赖关系是否对精确的 3D 医学图像配准是必要的。
• 设计一种改进的 U-Net 架构，在不依赖注意力机制的前提下增强感受野。
• 通过轻量化、高效的架构在公开的 3D 医学图像配准基准上实现最先进性能。
• 证明单一、精心设计的 U-Net 变体可超越复杂且参数量庞大的 Transformer 模型。

提出的方法

• 提出 LKU-Net，一种采用大感受野卷积模块的 U-Net 架构，将并行的大感受野卷积模块集成到标准 U-Net 的编码器和解码器路径中，以扩展有效感受野。
• 在大感受野卷积模块中使用大卷积核（如 15×15×15），在不增加网络深度的前提下捕获更大区域的空间上下文信息。
• 保留标准 U-Net 的跳跃连接架构，以维持特征融合与梯度流动。
• 采用标准无监督损失函数，包含一个变形层、数据保真项和正则化项，实现网络端到端训练。
• 在 IXI 和 Learn2Reg 2021 数据集上，采用与 TransMorph 相同的数据预处理、训练协议和评估指标，确保公平比较。
• 使用相同的网络架构，通过与不通过雅可比行列式正则化，评估微分与非微分配准性能。

实验结果

研究问题

• RQ1 基于 Transformer 的网络在医学图像配准中的性能提升，是否主要源于其长程建模能力？
• RQ2 在不使用注意力机制的前提下，改进的 U-Net 架构能否实现最先进性能？
• RQ3 3D 脑图像中平均位移仅为 2.1–2.3 个体素，是否意味着长程依赖关系对配准并非必要？
• RQ4 轻量化 U-Net 变体能否在准确率与效率方面超越更大、更复杂的 Transformer 模型（如 TransMorph）？
• RQ5 在现代 3D 医学图像配准任务中，U-Net 架构是否已过时，还是可通过架构改进保持竞争力？

主要发现

• 在 3D IXI 脑部数据集上，LKU-Net 的 Dice 分数达到 0.757，优于 TransMorph（0.746），且参数量仅为后者的 1.12%。
• LKU-Net 将计算成本降低至 TransMorph 的 10.8%，同时保持或提升了配准精度。
• 在 MICCAI Learn2Reg 2021 任务 3 数据集上，LKU-Net 的 Dice 分数达到 0.8861，在公开验证排行榜中排名第一，优于 TransMorph-LC 和挑战优胜模型 LapIRN。
• 在图谱到被试者配准任务中，8 通道、卷积核大小为 15 的 LKU-Net 变体取得 0.757 的 Dice 分数，超越所有其他方法（包括 TransMorph 和 nnFormer）。
• 在微分配准中，LKU-Net-diff 8,5 的 Dice 分数达到 0.753，显著优于 TransMorph-diff（0.599），并接近 TransMorph-bspl（0.752）。
• 可视化对比显示，LKU-Net 生成的形变更平滑，且在复杂解剖结构区域的形变图像更逼真。