[논문 리뷰] UNeXt: MLP-based Rapid Medical Image Segmentation Network

Jeya Maria Jose Valanarasu, Vishal M. Patel|arXiv (Cornell University)|2022. 03. 09.

Advanced Neural Network Applications인용 수 54

한 줄 요약

UNeXt는 의학 영상 분할을 위한 컨볼루션 스템에 이어 토큰화된 MLP 블록을 사용하는 모델로, TransUNet 및 UNe t 변종에 비해 매개변수가 훨씬 적고 계산량이 크게 낮아 최첨단 성능을 달성하며 현장 진료용 CPU 추론을 더 빠르게 가능하게 한다.

ABSTRACT

UNet and its latest extensions like TransUNet have been the leading medical image segmentation methods in recent years. However, these networks cannot be effectively adopted for rapid image segmentation in point-of-care applications as they are parameter-heavy, computationally complex and slow to use. To this end, we propose UNeXt which is a Convolutional multilayer perceptron (MLP) based network for image segmentation. We design UNeXt in an effective way with an early convolutional stage and a MLP stage in the latent stage. We propose a tokenized MLP block where we efficiently tokenize and project the convolutional features and use MLPs to model the representation. To further boost the performance, we propose shifting the channels of the inputs while feeding in to MLPs so as to focus on learning local dependencies. Using tokenized MLPs in latent space reduces the number of parameters and computational complexity while being able to result in a better representation to help segmentation. The network also consists of skip connections between various levels of encoder and decoder. We test UNeXt on multiple medical image segmentation datasets and show that we reduce the number of parameters by 72x, decrease the computational complexity by 68x, and improve the inference speed by 10x while also obtaining better segmentation performance over the state-of-the-art medical image segmentation architectures. Code is available at https://github.com/jeya-maria-jose/UNeXt-pytorch

연구 동기 및 목표

• 제한된 컴퓨트 자원으로 현장 진료 영상 분할을 촉진한다.
• 컨볼루션과 토큰화된 MLP 구성 요소를 결합한 경량 인코더-디코더 아키텍처를 개발한다.
• 잠재 표현을 효율적으로 모델링하기 위해 축 방향 시프트를 갖춘 토큰화된 MLP 블록을 도입한다.
• 분할 정확도를 유지하거나 향상시키면서 매개변수와 FLOPs의 큰 감소를 보여준다.

제안 방법

• 초기 컨볼루션 스테이지에 이어 토큰화된 MLP 스테이지가 이어지는 2단계 아키텍처.
• 토큰화된 MLP 블록은 컨볼루션 특성을 토큰으로 투사하고 지역 의존성 모델링을 위해 시프트된 MLP를 적용한다.
• 축 방향 시프트(W 및 H)는 토큰화 이전에 지역성을 유도하며, Tok-MLP 블록에서 깊이별 합성곱과 GELU 활성화를 사용한다.
• 토큰화된 MLP 블록 내에서 잔차 연결과 계층 정규화가 사용된다.
• 인코더와 디코더 사이의 스킵 연결은 UNet을 반영하고, 디코더는 토큰화된 MLP 블록 뒤에 컨볼루션 블록을 사용한다.

실험 결과

연구 질문

• RQ1 컨볼루션 스템과 토큰화된 MLP를 잠재 공간에서 결합하여 매개변수 및 계산량을 줄이면서 분할 정확도를 유지할 수 있는가?
• RQ2 Tok-MLP 블록의 축 방향 시프트가 경쟁력 있는 의료 영상 분할에 충분한 지역성을 제공하는가?
• RQ3 유닛(UNet), UNet++, ResUNet, MedT, TransUNet에 비해 UNeXt의 정확도, 매개변수, FLOPs, CPU 추론 시간은 어떤 차이가 있는가?

주요 결과

Networks	Params (M)	Inference Speed (ms)	GFLOPs	ISIC F1	ISIC IoU	BUSI F1	BUSI IoU
UNet	31.13	223	55.84	84.03 ± 0.87	74.55 ± 0.96	76.35 ± 0.89	63.85 ± 1.12
UNet++	9.16	173	34.65	84.96 ± 0.71	75.12 ± 0.65	77.54 ± 0.74	64.33 ± 0.75
ResUNet	62.74	333	94.56	85.60 ± 0.68	75.62 ± 1.11	78.25 ± 0.74	64.89 ± 0.83
MedT	1.60	751	21.24	87.35 ± 0.18	79.54 ± 0.26	76.93 ± 0.11	63.89 ± 0.55
TransUNet	105.32	246	38.52	88.91 ± 0.63	80.51 ± 0.72	79.30 ± 0.37	66.92 ± 0.75
UNeXt	1.47	25	0.57	89.70 ± 0.96	81.70 ± 1.53	79.37 ± 0.57	66.95 ± 1.22

• UNeXt는 ISIC 및 BUSI 데이터셋에서 경쟁력 있거나 우수한 F1 및 IoU 점수를 달성한다.
• UNeXt는 1.47M 매개변수와 0.57 GFLOPs를 사용하여 TransUNet(105.32M, 38.52 GFLOPs)보다 훨씬 적다.
• UNeXt은 ISIC에서 89.70 F1 및 81.70 IoU, BUSI에서 79.37 IoU를 달성하며 CPU 추론 시간은 25 ms이다.
• 전체 기준선에서 UNeXt는 정확도와 효율성의 최적 균형을 제공하며, 주의(attention) 기반 모델 대비 계산 및 매개변수 수에서 우수하다.
• 교차실험(ablation) 결과 Tok-MLP의 시프트 및 컨볼루션과 MLP 스테이지의 결합이 복잡도 증가를 최소화하면서 가장 높은 성능을 보인다.