[论文解读] HUMUS-Net: Hybrid unrolled multi-scale network architecture for accelerated MRI reconstruction

Zalan Fabian, Tinaz, Berk+1 more|arXiv (Cornell University)|Mar 15, 2022

Advanced X-ray and CT Imaging被引用 45

一句话总结

HUMUS-Net 将卷积和基于 Transformer 的模块结合在一个展开的多尺度架构中，以提升加速 MRI 重建，在 fastMRI 膝部数据上达到最先进的结果。

ABSTRACT

In accelerated MRI reconstruction, the anatomy of a patient is recovered from a set of under-sampled and noisy measurements. Deep learning approaches have been proven to be successful in solving this ill-posed inverse problem and are capable of producing very high quality reconstructions. However, current architectures heavily rely on convolutions, that are content-independent and have difficulties modeling long-range dependencies in images. Recently, Transformers, the workhorse of contemporary natural language processing, have emerged as powerful building blocks for a multitude of vision tasks. These models split input images into non-overlapping patches, embed the patches into lower-dimensional tokens and utilize a self-attention mechanism that does not suffer from the aforementioned weaknesses of convolutional architectures. However, Transformers incur extremely high compute and memory cost when 1) the input image resolution is high and 2) when the image needs to be split into a large number of patches to preserve fine detail information, both of which are typical in low-level vision problems such as MRI reconstruction, having a compounding effect. To tackle these challenges, we propose HUMUS-Net, a hybrid architecture that combines the beneficial implicit bias and efficiency of convolutions with the power of Transformer blocks in an unrolled and multi-scale network. HUMUS-Net extracts high-resolution features via convolutional blocks and refines low-resolution features via a novel Transformer-based multi-scale feature extractor. Features from both levels are then synthesized into a high-resolution output reconstruction. Our network establishes new state of the art on the largest publicly available MRI dataset, the fastMRI dataset. We further demonstrate the performance of HUMUS-Net on two other popular MRI datasets and perform fine-grained ablation studies to validate our design.

研究动机与目标

• 通过解决纯卷积模型在建模长程依赖方面的局限性，推动加速 MRI 重建的改进。
• 提出一种混合架构，在利用多尺度 Transformer 基于特征提取的同时，保留高分辨率细节。
• 引入一个展开式重建框架，带有在像素域的去噪器，融合高分辨率和低分辨率特征。
• 在大规模 MRI 数据集上展示最先进的性能，并提供消融研究以验证设计选择。

提出的方法

• 引入 HUMUS-Net，一串用于 k 空间 MRI 重建的展开迭代级联。
• HUMUS-Block 通过一个 3x3 卷积层提取高分辨率特征，通过卷积下采样路径提取低分辨率特征。
• 一个新颖的多尺度混合特征提取器 MUST 在图像表示和标记表示中使用 Swin-Transformer 块进行编码器-解码器处理，并带跳连。
• 低分辨率特征在 Transformer-卷积混合架构中处理，以捕获长程依赖，同时保持计算可行。
• 重建通过一个残差图像学习模块合成高分辨率和深层低分辨率特征：x_out = x_in + f_R(F_H, F_D)。
• 前向模型使用数据一致性和线圈灵敏度扩展/缩减，且可选的相邻切片重建以改善上下文。

实验结果

研究问题

• RQ1 相较于仅基于 CNN 或 Transformer 的模型，混合卷积-Transformer 架构是否能提升加速 MRI 的重建质量？
• RQ2 多尺度处理与展开的 k 空间优化是否在像 fastMRI 这样的大型公开 MRI 数据集上实现最先进的结果？
• RQ3 同时处理高分辨率图像特征与低分辨率 Transformer 基于特征对重建保真度的影响是什么？
• RQ4 相邻切片重建（ASR) 是否在不同数据集上提供一致的增益？
• RQ5 消融研究（单尺度 SwinIR 与 MUST 与 HUMUS-Block）如何揭示多尺度混合处理的重要性？

主要发现

• HUMUS-Net 在 fastMRI 膝部测试数据上实现了 SSIM=0.8936，PSNR=37.0，NMSE=0.0086，优于若干基线。
• 一个变体 HUMUS-Net-L 达到 SSIM=0.8951，PSNR=37.4，NMSE=0.0080，嵌入维度较大。
• 在 fastMRI 验证中，HUMUS-Net 达到 SSIM=0.8934（对比 E2E-VarNet 0.8908）。
• 消融研究显示多尺度混合处理（MUST）优于单尺度 SwinIR 与展开变体，HUMUS-Net 在斯坦福 3D MRI 上达到 SSIM=0.9449±0.0064。
• ASR 通过利用相邻切片的上下文在不改变标记数量的情况下进一步提高重建。
• 总体而言，HUMUS-Net 在 fastMRI 膝部数据上取得新的最先进 SSIM，并在 Stanford 2D/3D 数据集上表现出鲁棒性。