[논문 리뷰] Self-supervised Learning for Gastrointestinal Pathologies Endoscopy Image Classification with Triplet Loss

Tai Nguyen-D-P|arXiv (Cornell University)|2023. 03. 03.

Gastrointestinal Bleeding Diagnosis and Treatment인용 수 4

한 줄 요약

이 논문은 자가학습 학습 프레임워크인 SsLTF를 제안한다. 이는 표현 학습을 위해 Barlow Twins를 사용하고, 내부 클래스 분산을 줄이고 클래스 간 유사도를 감소시키기 위해 트리플릿 손실을, 클래스 불균형 문제를 해결하기 위해 포칼 손실을 적용한다. 간단한 ResNet-50 백본을 사용하여 Hyper-Kvasir 데이터셋에서 F1 스코어 0.824를 기록하며 최신 기술 수준의 성능을 달성하였으며, 더 복잡한 모델들보다 뛰어난 성능을 보였다.

ABSTRACT

Recently, the amount of GI tract datasets is introduced more and more by gathering from contests and challenges. The most common task needs to solve that is to classify images from the GI tract into various classes. However, the contributions of the existing approaches exhibit lots of limitations. In this paper, we aim to develop a computer-aided diagnosis system to classify the pathological findings in endoscopy images, the system can classify some common pathologies including polyps, esophagitis, and ulcerative -- colitis. To evaluate the proposed work, we use the public dataset which is Hyper--Kvasir instead of gathering the data. The key idea of our system is to develop self-supervised learning based on the Barlow Twins framework with a downstream task which is an endoscopy image classification integrated with triplet loss and focal loss functions. The self-supervision framework and focal loss function are used to overcome class-imbalanced data, while the triplet loss function is to tackle the domain-specific properties in endoscopy images which are inter/intra-class problems. An extensive experimental study on the pathological finding images in the Hyper--Kvasir dataset has shown that our proposed system is in general better than the compared methods, whereas using a simple neural network model. This means the proposed system can be used efficiently and capable of accurately for the classification of pathology images in the GI tract.

연구 동기 및 목표

• 내시경 영상에서 다수의 위장병변(폴립, 식도염, 궤양성 대장염 등)을 분류하기 위한 컴퓨터 지원 진단 시스템을 개발한다.
• 모델의 일반화 능력을 저해하는 내시경 영상에서의 클래스 불균형과 내/외부 클래스 변동성 문제를 해결한다.
• 대규모 비공개 데이터셋에 의존하지 않고 자가학습 사전 훈련과 손실 함수 통합을 통해 분류 성능을 향상시킨다.
• 간단한 모델에 철저히 설계된 손실 함수를 적용할 경우 더 복잡한 아키텍처보다 의료 영상 분류 과제에서 뛰어난 성능을 낼 수 있음을 보여준다.

제안 방법

• 시스템은 자가학습 사전 훈련을 위해 unlabeled 내시경 영상에서 강력한 특징 표현을 학습하기 위해 Barlow Twins 프레임워크를 사용한다.
• 하류 분류 헤드는 트리플릿 손실과 포칼 손실의 조합을 통해 미세조정되며, 이는 분류 능력을 향상시키고 클래스 불균형 문제를 완화한다.
• 트리플릿 손실은 내부 클래스 분산을 최소화하고 클래스 간 간격을 최대화하여, 유사한 외관을 가진 병변의 특징 분리 능력을 향상시킨다.
• 포칼 손실은 학습 중에 쉬운 음성 예측과 다수 클래스의 영향을 줄이고, 어려운 예제에 집중함으로써 학습을 효과적으로 조정한다.
• 모델는 6개 클래스로 구성된 Hyper-Kvasir 데이터셋에서 훈련된다: 폴립, 식도염-grade-a, 식도염-grade-b-d, 궤양성 대장염-grade-1, 궤양성 대장염-grade-2, 궤양성 대장염-grade-3.
• 최종 분류 파이프라인은 병합된 손실 목적함수를 통해 엔드 투 엔드로 훈련된 ResNet-50 백본을 사용한다.

실험 결과

연구 질문

• RQ1 Barlow Twins를 활용한 자가학습 사전 훈련이 위장내 내시경 영상 분류를 위한 표현 표현을 향상시킬 수 있는가?
• RQ2 내시경 영상에서 높은 시각적 유사성을 가지는 병변에 대해 트리플릿 손실이 내부 클래스 분산을 줄이고 클래스 간 분리 능력을 향상시키는 데 얼마나 효과적인가?
• RQ3 의료 영상 데이터셋에서 클래스 불균형으로 인한 성능 저하를 포칼 손실이 어느 정도 완화할 수 있는가?
• RQ4 고도로 발전된 손실 함수와 자가학습 사전 훈련을 결합한 경량 모델인 ResNet-50이 더 큰 아키텍처인 DenseNet-161과 EfficientNet보다 성능을 뛰어넘을 수 있는가?

주요 결과

• 제안된 SsLTF 방법은 Hyper-Kvasir 데이터셋에서 F1 스코어 0.824를 기록하여, DenseNet-161과 EfficientNet 등 더 큰 모델을 사용한 비교 모델들을 모두 능가하는 성능을 보였다.
• SsLTF의 매크로 평균 F1 스코어(0.824)는 Borgli 등 [3] (0.6747), He 등 [11] (0.6701), Gjestang [7] (0.5800)보다 높았으며, 이는 더 단순한 ResNet-50 아키텍처를 사용한 것에도 불구하고 성과를 냈다.
• 트리플릿 손실의 통합은 특히 다양한 정도의 궤양성 대장염과 식도염과 같이 높은 시각적 유사성을 가지는 병변에 대해 특징의 분류 능력을 크게 향상시켰다.
• 포칼 손실은 다수 클래스인 폴립의 총 손실에 미치는 영향을 효과적으로 줄여, 궤양성 대장염-grade-3와 같은 소수 클래스의 재현율을 향상시켰다.
• Barlow Twins를 활용한 자가학습 사전 훈련 단계는 특징 품질을 향상시켜, 제한된 레이블 데이터 조건에서도 더 나은 하류 성능을 가능하게 하였다.
• 모델의 단순성에도 불구하고 SsLTF는 경쟁적인 성능을 달성하였으며, 이는 이 설정에서 손실 함수 설계와 사전 훈련이 모델 깊이보다 더 큰 영향을 미친다는 것을 보여주었다.