[논문 리뷰] Capsules for Object Segmentation

Rodney LaLonde, Ulaş Bağcı|arXiv (Cornell University)|2018. 04. 11.

Advanced Neural Network Applications참고 문헌 15인용 수 215

한 줄 요약

이 논문은 Locally-constrained routing과 deconvolutional capsules를 사용하는 SegCaps를 소개합니다. 이는 512x512 이미지에서 U-Net보다 훨씬 적은 파라미터로 객체 분할을 수행하며, CT 스캔의 병리학적 폐 분할에서 경쟁력 있는 dice 점수를 보여줍니다.

ABSTRACT

Convolutional neural networks (CNNs) have shown remarkable results over the last several years for a wide range of computer vision tasks. A new architecture recently introduced by Sabour et al., referred to as a capsule networks with dynamic routing, has shown great initial results for digit recognition and small image classification. The success of capsule networks lies in their ability to preserve more information about the input by replacing max-pooling layers with convolutional strides and dynamic routing, allowing for preservation of part-whole relationships in the data. This preservation of the input is demonstrated by reconstructing the input from the output capsule vectors. Our work expands the use of capsule networks to the task of object segmentation for the first time in the literature. We extend the idea of convolutional capsules with locally-connected routing and propose the concept of deconvolutional capsules. Further, we extend the masked reconstruction to reconstruct the positive input class. The proposed convolutional-deconvolutional capsule network, called SegCaps, shows strong results for the task of object segmentation with substantial decrease in parameter space. As an example application, we applied the proposed SegCaps to segment pathological lungs from low dose CT scans and compared its accuracy and efficiency with other U-Net-based architectures. SegCaps is able to handle large image sizes (512 x 512) as opposed to baseline capsules (typically less than 32 x 32). The proposed SegCaps reduced the number of parameters of U-Net architecture by 95.4% while still providing a better segmentation accuracy.

연구 동기 및 목표

• 문헌상 처음으로 캡슐 네트워크를 객체 분할로 확장한다.
• 로컬 제약 라우팅과 공유 변환을 통해 전통적 캡슐의 메모리 및 파라미터 폭발 문제를 해결한다.
• 대형 영상에서 심층 분할 아키텍처를 가능하게 하기 위해 deconvolutional 캡슐을 도입한다.
• 양성 클래스의 마스킹 재구성을 통한 정규화로 입력 공간의 임베딩을 개선한다.
• 병리학적 폐 분할에서 SegCaps를 시연하고 U-Net 및 Tiramisu와 비교한다.

제안 방법

• 부분-전체 관계를 보존하기 위해 최대풀링을 컨볼루션 스트라이드와 라우팅으로 대체한다.
• k_h x k_w 커널 내에서 라우팅 계수는 계산되는 로컬 제약 동적 라우팅을 도입한다.
• 매 캡슐 타입 내에서 변환 행렬을 공유하여 파라미터를 줄인다.
• 로컬 라우팅으로 라우팅된 전치 합성곱을 사용하는 deconvolutional 캡슐을 추가하여 심층 분할을 가능하게 한다.
• 3-레이어 1x1 컨브 네트로와 마스킹된 MSE를 이용한 양성 클래스의 마스킹 재구성으로 정규화 손실을 사용한다.
• 네겹 교차 검증으로 LIDC-IDRI의 LUNA16 부분집합에서 512x512 CT 슬라이스를 학습 및 평가한다.

실험 결과

연구 질문

• RQ1 캡슐 네트워크가 큰 이미지 크기(512x512)에서 픽셀 수준의 객체 분할에 효과적으로 적용될 수 있는가?
• RQ2 로컬 제약 라우팅과 공유 변환 행렬이 메모리/파라미터를 줄이면서 분할 정확도를 유지하거나 향상시키는가?
• RQ3 deconvolutional 캡슐과 마스킹 재구성이 U-Net 및 관련 아키텍처에 비해 분할 성능을 향상시키는가?
• RQ4 CT 스캔의 병리학적 폐 분할에서 SegCaps의 성능은 최신 방법에 비해 어떤가?

주요 결과

방법	매개변수	분할-0 (%)	분할-1 (%)	분할-2 (%)	분할-3 (%)	평균 (%)
U-Net	31.0 M	98.353	98.432	98.476	98.510	98.449
Tiramisu	2.3 M	98.394	98.358	98.543	98.339	98.410
Baseline Caps	1.7 M	82.287	79.939	95.121	83.608	83.424
SegCaps (R1)	1.4 M	98.471	98.444	98.401	98.362	98.419
SegCaps	1.4 M	98.499	98.523	98.455	98.474	98.479

• SegCaps는 4개 폴드에 걸쳐 평균 dice 점수 98.479%를 달성하여 U-Net과 Tiramisu를 약간 능가했다.
• SegCaps는 1.4M 파라미터를 사용하여 U-Net 대비 파라미터를 95.4% 감소시키고 Tiramisu 대비 38.4% 감소시킨다.
• 기본 3계층 캡슐 네트워크는 현저히 열악한 성능을 보여 깊은 컨볼루션-디컨볼루션 아키텍처의 이점을 강조한다.
• SegCaps는 512x512 입력에서 작동하며, 이전의 캡슐 모델이 (<=32x32)와 같이 작은 입력에 제한된 것과는 다르다.
• 정성적 결과는 SegCaps가 CT 슬라이스에서 U-Net에 비해 분할 누출을 줄임을 보여준다.
• 시각화는 최종 캡슐 벡터가 분할 영역 전반에서 뚜렷한 텍스처 특성을 포착한다는 것을 시사한다.