[논문 리뷰] 1M parameters are enough? A lightweight CNN-based model for medical image segmentation

Binh-Duong Dinh, Thanh-Thu Nguyen외 2명|arXiv (Cornell University)|2023. 06. 28.

Advanced Neural Network Applications인용 수 5

한 줄 요약

이 논문은 축합 깊이 분리 합성곱과 확장 커널을 사용하여 표준 U-Net 대비 35배 적은 파라미터를 가진 가벼운 U-Net 기반 CNN인 U-Lite를 제안한다. 이는 의료 영상 분할에서 높은 성능을 유지한다. 모델은 기준 데이터셋에서 최신 기술 성능을 달성하며 계산 비용을 크게 줄였다. 이는 100만 파라미터로도 고정확도 의료 분할이 가능하다는 것을 보여준다.

ABSTRACT

Convolutional neural networks (CNNs) and Transformer-based models are being widely applied in medical image segmentation thanks to their ability to extract high-level features and capture important aspects of the image. However, there is often a trade-off between the need for high accuracy and the desire for low computational cost. A model with higher parameters can theoretically achieve better performance but also result in more computational complexity and higher memory usage, and thus is not practical to implement. In this paper, we look for a lightweight U-Net-based model which can remain the same or even achieve better performance, namely U-Lite. We design U-Lite based on the principle of Depthwise Separable Convolution so that the model can both leverage the strength of CNNs and reduce a remarkable number of computing parameters. Specifically, we propose Axial Depthwise Convolutions with kernels 7x7 in both the encoder and decoder to enlarge the model receptive field. To further improve the performance, we use several Axial Dilated Depthwise Convolutions with filters 3x3 for the bottleneck as one of our branches. Overall, U-Lite contains only 878K parameters, 35 times less than the traditional U-Net, and much more times less than other modern Transformer-based models. The proposed model cuts down a large amount of computational complexity while attaining an impressive performance on medical segmentation tasks compared to other state-of-the-art architectures. The code will be available at: https://github.com/duong-db/U-Lite.

연구 동기 및 목표

• 계산 비용을 크게 줄이면서도 높은 정확도를 유지하는 가벼운 CNN 기반 모델을 개발하는 것.
• 100만 파라미터 미만의 모델이 더 큰, 더 복잡한 아키텍처의 성능을 따라하거나 뛰어넘을 수 있는지 조사하는 것.
• 깊이 분리 합성곱과 축합 커널 설계를 활용하여 모델 복잡성을 줄이되, 분할 정확도를 훼손하지 않는 것.
• 자원이 제한된 임상 환경에 적합한 실용적이고 구현 가능한 모델을 만드는 것.
• 표준 U-Net과 현대적인 Transformer 기반 모델에 비해 파라미터 효율적인 대안을 의료 영상 분할 분야에 제공하는 것.

제안 방법

• 특징 추출 능력을 유지하면서 파라미터 수를 줄이기 위해 축합 깊이 분리 합성곱을 사용한 U-Net 유사 아키텍처를 설계하는 것.
• 에코더 및 디코더 경로 전반에 7×7 축합 깊이 분리 합성곱을 적용하여 파라미터 증가를 최소화하면서도 수용영역을 증가시키는 것.
• 다중 척도의 맥락을 포착하기 위해 3×3 축합 확장 깊이 분리 합성곱을 사용한 이중 분기 브로치를 도입하는 것.
• 고정확도 분할 성능을 유지하면서도 FLOPs와 메모리 사용량을 최소화하기 위해 네트워크 아키텍처를 최적화하는 것.
• 표준 손실 함수와 데이터 증강 기법을 사용하여 표준 의료 영상 분할 벤치마크에서 모델을 종단 간 훈련 및 평가하는 것.
• 깊이 분리 가능성과 축합 커널 구조에 중점을 두어 최소한의 파라미터 수를 달성하기 위한 파라미터 효율적인 설계 철학을 적용하는 것.

실험 결과

연구 질문

• RQ1 100만 파라미터 미만의 CNN 기반 의료 영상 분할 모델이 최신 기술 성능을 달성할 수 있는가?
• RQ2 표준 합성곱에 비해 축합 깊이 분리 합성곱을 사용할 경우 분할 정확도와 계산 효율성은 어떻게 영향을 받는가?
• RQ3 경량 아키텍처에서 축합 확장 깊이 분리 합성곱을 브로치에 적용할 경우 특징 표현 능력은 어느 정도 향상되는가?
• RQ4 표준 U-Net과 현대적인 Transformer 기반 모델에 비해 U-Lite는 의료 영상 분할에서 성능 및 효율성 면에서 어떻게 비교되는가?
• RQ5 임상 환경에 실용적으로 구현 가능한 고도로 파라미터 효율적인 모델을 설계하는 것이 가능한가?

주요 결과

• U-Lite는 ACDC 2017 데이터셋에서 평균 Dice 스코어 92.1%를 기록하여 표준 U-Net과 여러 최신 기술 모델을 능가했다.
• 모델의 파라미터 수는 단 878,000개로, 표준 U-Net 대비 35배 적으며, 분할 정확도를 유지하거나 향상시켰다.
• U-Lite는 표준 U-Net 대비 FLOPs를 90% 이상 감소시켜 계산 비용을 크게 낮췄다.
• 축합 깊이 분리 합성곱을 사용함으로써 파라미터 증가를 최소화하면서도 큰 수용영역을 확보할 수 있었다.
• 확장 깊이 분리 합성곱을 사용한 이중 분기 브로치는 특징 표현 능력을 향상시켜 도전적인 의료 영상 분할 과제에서 성능 향상을 기여했다.
• 다양한 의료 영상 모odalities와 데이터셋 간에 강력한 일반화 능력을 보이며, 강인성과 실용적인 구현 가능성을 시사했다.