[논문 리뷰] RaBiT: An Efficient Transformer using Bidirectional Feature Pyramid Network with Reverse Attention for Colon Polyp Segmentation

Nguyen Hoang Thuan, Thi-Oanh Nguyen외 3명|arXiv (Cornell University)|2023. 07. 12.

Radiomics and Machine Learning in Medical Imaging인용 수 7

한 줄 요약

RaBiT는 다중 척도 특징 융합과 경계 정밀 조정을 향상시키기 위해 반대 주의(RaBiFPN)와 벽돌형 최적화된 반대 주의 모듈을 사용하는 경량 트랜스포머 기반 인코더-디코더 네트워크를 제안한다. 고해상도 ETIS-Larib에서 기존 모델 대비 4.5% mIoU 향상과 함께 파라미터 수를 37% 감소시켜 최신 기술 수준의 성능을 달성한다.

ABSTRACT

Automatic and accurate segmentation of colon polyps is essential for early diagnosis of colorectal cancer. Advanced deep learning models have shown promising results in polyp segmentation. However, they still have limitations in representing multi-scale features and generalization capability. To address these issues, this paper introduces RaBiT, an encoder-decoder model that incorporates a lightweight Transformer-based architecture in the encoder to model multiple-level global semantic relationships. The decoder consists of several bidirectional feature pyramid layers with reverse attention modules to better fuse feature maps at various levels and incrementally refine polyp boundaries. We also propose ideas to lighten the reverse attention module and make it more suitable for multi-class segmentation. Extensive experiments on several benchmark datasets show that our method outperforms existing methods across all datasets while maintaining low computational complexity. Moreover, our method demonstrates high generalization capability in cross-dataset experiments, even when the training and test sets have different characteristics.

연구 동기 및 목표

• 기존 모델이 다양한 대장내시경 데이터셋 간의 일반화 능력과 다중 척도 특징을 포괄하지 못하는 한계를 해결하기 위해.
• 새로운 반대 주의 메커니즘을 활용한 반복적 개선을 통해 폴립 경계 세분화를 향상시키기 위해.
• 정확도를 유지하면서도 계산 비용을 줄이기 위해.
• 인코더-디코더 수준 간의 특징 융합을 향상시키기 위해 이중 방향 특징 피라미드 구조를 사용하기 위해.
• 다중 클래스 세분화를 효과적으로 수행하기 위해 다중 클래스 호환 가능한 반대 주의 모듈을 도입하기 위해.

제안 방법

• 인코더는 계층적 전역 의미 관계를 다중 척도에서 포괄하기 위해 경량 MiT-B3 기반 트랜스포머를 사용한다.
• 디코더는 가중치가 부여된 연결과 스킵 커넥션을 사용해 다양한 인코더 레벨의 특징을 융합하는 이중 방향 피라미드(RaBiFPN) 모듈 스택을 활용한다.
• 파라미터와 계산 비용을 줄이면서도 폴립 경계의 반복적 개선을 가능하게 하는 경량 벽돌형 기반 반대 주의(RA) 모듈을 도입한다.
• 비암암성 및 암성 폴립을 더 잘 분류하기 위해 영역 경계 강조 기능을 갖춘 다중 클래스 RA 모듈을 설계한다.
• RA 모듈은 시그모이드 또는 소프트맥스 활성화 함수를 사용하며, 실험 결과 소형 객체 세분화에 소프트맥스가 더 뛰어난 성능을 보였다.
• 엔드 투 엔드로 교차 엔트로피 손실과 딱지 손실을 사용해 학습하며, 추상화 연구를 통해 구성 요소의 효과를 검증한다.

실험 결과

연구 질문

• RQ1 경량 트랜스포머 인코더가 대장 폴립 세분화에서 전역적 맥락과 다중 척도 특징을 효과적으로 모델링할 수 있는가?
• RQ2 반대 주의 기반 이중 방향 피라미드가 표준 FPN 또는 BiFPN에 비해 특징 융합과 경계 정밀 조정을 향상시키는가?
• RQ3 벽돌형 최적화된 반대 주의 모듈이 계산 비용을 줄이면서도 세분화 정확도를 유지하거나 향상시킬 수 있는가?
• RQ4 제안된 방법이 다양한 대장내시경 데이터셋 간의 분포 차이(해상도 및 폴립 특성)를 고려할 때 어떻게 일반화되는가?
• RQ5 다중 클래스 반대 주의 모듈이 암성 및 비암성 폴립 양쪽의 세분화 성능을 향상시키는가?

주요 결과

• ETIS-Larib 고해상도 데이터셋에서 표준 BiFPN 대비 4.5% mIoU 향상으로 더 큰 이미지에서 뛰어난 효과를 입증했다.
• Neo-Small 데이터셋에서 반대 주의 모듈에 소프트맥스 활성화 함수를 적용함으로써 딱지 비율이 6.43% 향상되어 소형 객체 세분화 성능 향상이 확인되었다.
• 제안된 벽돌형 모듈 덕분에 파라미터 수가 81.79M에서 51.72M로 감소하고 GFLOPs는 41.37에서 23.12로 감소하여 계산 비용을 크게 낮췄다.
• RaBiFPN 모듈을 4개 사용할 경우, 모든 벤치마크 데이터셋에서 SegFormer-B3 및 기타 최신 기술 수준의 모델을 능가하는 최고의 성능을 달성했다.
• 다양한 데이터셋 간 일반화 테스트에서 학습 및 테스트 세트의 분포가 다를 경우에도 뛰어난 성능 유지를 보이며 높은 강건성을 입증했다.
• 추상화 연구 결과, RaBiFPN과 벽돌형 모듈이 각각 mIoU 향상에 기여하며, 특히 고해상도 데이터셋에서 가장 큰 성과를 보였다.