[논문 리뷰] U-Net vs Transformer: Is U-Net Outdated in Medical Image Registration?

Xi Jia, Joseph W. Bartlett외 4명|arXiv (Cornell University)|2022. 08. 07.

Brain Tumor Detection and Classification인용 수 5

한 줄 요약

이 논문은 대용량 커널 컨볼루션 블록을 통합한 수정된 U-Net인 LKU-Net을 제안하여 수용 영역를 향상시키며, 3D 의료 영상 정렬에서 TransMorph와 같은 최신 트랜스포머 기반 방법들을 능가함을 입증한다. TransMorph의 파rameter의 1.12%와 Multiply-Add 연산의 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 기반 방법이 오래되었는지 여부를 조사하기 위해.
• 정확한 3D 의료 영상 정렬을 위해 트랜스포머가 모델링하는 장거리 의존성은 필수적인가를 평가하기 위해.
• 주의 메커니즘에 의존하지 않고 수용 영역를 향상시키는 수정된 U-Net 아키텍처를 설계하기 위해.
• 경량적이고 효율적인 아키텍처를 사용하여 공개 3D 의료 영상 정렬 벤치마크에서 최신 기술 수준의 성능을 달성하기 위해.
• 단일의 잘 설계된 U-Net 변종이 복잡하고 파라미터가 많은 트랜스포머 모델을 능가할 수 있음을 입증하기 위해.

제안 방법

• 기본 U-Net의 인코더 및 디코더 경로에 병렬 대용량 커널 컨볼루션 블록을 통합하여 효과적 수용 영역를 확장한 대용량 커널 U-Net인 LKU-Net을 제안한다.
• 깊이를 증가시키지 않고도 더 넓은 영역의 공간적 맥락을 포괄하기 위해 대용량 커널 컨볼루션(예: 커널 크기 15×15×15)을 대용량 커널 블록에 사용한다.
• 특징 융합과 기울기 흐름을 유지하기 위해 표준 U-Net의 스킵 커넥션 아키텍처를 유지한다.
• 네트워크를 엔드 투 엔드로 훈련하기 위해 표준 비지도 학습 손실, 워핑 레이어, 데이터 적합성 항목 및 정규화를 사용한다.
• IXI 및 Learn2Reg 2021 데이터셋에서 공정한 비교를 위해 TransMorph와 동일한 데이터 전처리, 훈련 프rotocol 및 평가 지표를 적용한다.
• 자기 동역학성 및 비자기 동역학성 정렬을 동일한 네트워크 아키텍처를 사용하여 재구성된 자코비안 행렬 행렬식 정규화 여부에 따라 평가한다.

실험 결과

연구 질문

• RQ1 트랜스포머 기반 네트워크의 성능 향상이 주로 장거리 모델링 능력 덕분인가?
• RQ2 주의 메커니즘을 사용하지 않고도 수정된 U-Net 아키텍처가 최신 기술 수준의 성능을 달성할 수 있는가?
• RQ3 3D 뇌 영상에서 평균 이동 거리가 2.1–2.3 벽소자일 경우, 장거리 의존성이 정렬에 필수적인가?
• RQ4 경량 U-Net 변종이 정확성과 효율성 측면에서 TransMorph와 같은 더 크고 복잡한 트랜스포머 기반 모델을 능가할 수 있는가?
• RQ5 현대 3D 의료 영상 정렬 분야에서 U-Net 아키텍처는 오래되었는가, 아니면 아키텍처적 개선을 통해 경쟁력을 유지할 수 있는가?

주요 결과

• 3D IXI 뇌 데이터셋에서 LKU-Net은 딱도스 점수 0.757을 기록하여 TransMorph(0.746)를 능가하며, 파라미터는 TransMorph의 1.12%만 사용한다.
• LKU-Net은 TransMorph의 Multiply-Add 연산량의 10.8%로 감소시키면서도 정확도를 유지하거나 향상시킨다.
• MICCAI Learn2Reg 2021 Task 3 데이터셋에서 LKU-Net은 딱도스 점수 0.8861을 기록하여 공개 검증 리더보드에서 1위를 차지하며 TransMorph-LC와 도전 우승한 LapIRN을 모두 능가한다.
• 8채널 및 커널 크기 15를 가진 LKU-Net 변종은 애틀라스-서브젝트 정렬에서 딱도스 점수 0.757을 기록하여 TransMorph 및 nnFormer를 포함한 모든 다른 방법을 능가한다.
• 자기 동역학성 정렬에서 LKU-Net-diff 8,5는 딱도스 점수 0.753을 기록하여 TransMorph-diff(0.599)를 크게 능가하고 TransMorph-bspl(0.752)와 유사한 성능을 달성한다.
• 시각적 비교 결과, LKU-Net은 더 매끄러운 변형과 더 현실적인 변형된 영상을 생성하며, 특히 복잡한 해부학적 구조를 가진 영역에서 뚜렷한 우수성을 보였다.