[논문 리뷰] HUMUS-Net: Hybrid unrolled multi-scale network architecture for accelerated MRI reconstruction

Zalan Fabian, Tinaz, Berk외 1명|arXiv (Cornell University)|2022. 03. 15.

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.

연구 동기 및 목표

• 가장 긴 거리 의존성을 모델링하는 purely convolutional 모델의 한계를 해결하여 가속 MRI 재구성 개선을 목표로 한다.
• 고해상도 세부 정보를 보존하면서 다중 해상도 Transformer 기반 특징 추출을 활용하는 하이브리드 아키텍처를 제안한다.
• 이미지 도메인 디노이저가 고해상도와 저해상도 특징을 융합하는 언롤링 재구성 프레임워크를 도입한다.
• 대규모 MRI 데이터셋에서 최첨단 성능을 보여주고 설계 선택을 검증하는 애블레이션 연구를 제공한다.

제안 방법

• HUMUS-Net을 k-공간 MRI 재구성용 언롤링 반복의 cascade로 도입한다.
• HUMUS-Block은 3x3 합성곱 계층을 통해 고해상도 특징을 추출하고 합성곱 다운샘플링 경로를 통해 저해상도 특징을 추출한다.
• 새로운 다중 해상도 하이브리드 특징 추출기가 Swin-transformer 블록과 점프 연결을 사용하여 이미지 표현과 토큰 표현에서 인코더-디코더 처리를 수행한다.
• 저해상도 특징은 Transformer-합성곱 하이브리드 구성을 통해 장거리 의존성을 포착하면서 계산량을 현실화한다.
• 재구성은 고해상도 특징과 깊은 저해상도 특징을 잔차 이미지 학습 모듈로 결합하여: x_out = x_in + f_R(F_H, F_D).
• 포워드 모델은 데이터 일관성과 코일 민감도 확장/축소를 사용하며 인접 슬라이스 재구성을 추가로 사용해 컨텍스트를 개선할 수 있다.

실험 결과

연구 질문

• RQ1 하이브리드 합성곱-Transformer 아키텍처가 순수 CNN 기반 또는 Transformer 기반 모델과 비교하여 가속 MRI 재구성 품질을 향상시키는가?
• RQ2 다중 해상도 처리 및 언롤링된 k-공간 최적화가 fastMRI 같은 대규모 공개 MRI 데이터셋에서 최첨단 결과를 제공하는가?
• RQ3 고해상도 이미지 특징과 저해상도 Transformer 기반 특징을 함께 처리하는 것이 재구성 정확도에 미치는 영향은 얼마나 큰가?
• RQ4 인접 슬라이스 재구성(ASR)이 데이터셋 간에 일관된 이점을 제공하는가?
• RQ5 단일 해상도 SwinIR vs MUST vs 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 및 언롤된 변형보다 우수하며, 스탠포드 3D MRI에서 HUMUS-Net은 SSIM=0.9449±0.0064를 달성한다.
• ASR은 인접 슬라이스의 컨텍스트를 활용해 재구성을 추가로 개선하지만 토큰 수를 변경하지 않는다.
• 전반적으로 HUMUS-Net은 fastMRI 무릎에서 새로운 최첨단 SSIM을 달성하고 스탠포드 2D/3D 데이터셋에서 견고함을 입증한다.