기억의 기록(나의 기기)

[Computer Vision] A semi-supervised multi-task learning framework for cancer classification with weak annotation in whole-slide images 논문 리뷰 본문

논문 리뷰

[Computer Vision] A semi-supervised multi-task learning framework for cancer classification with weak annotation in whole-slide images 논문 리뷰

황경하 2025. 2. 14. 22:43
반응형

[Computer Vision] A semi-supervised multi-task learning framework for cancer classification with weak annotation in whole-slide images

안녕하세요, 데이터 과학 전공하고 있는 황경하입니다.

오늘은 지난 포스팅에 이어 딥러닝이 병리학적 데이터에 적용된 논문을 살펴보겠습니다. 본 논문은 다른 논문과 달리 Supervised-Learning 방식이 아닌 Semi-Supervised Learning 방식으로 진행하는 Multi-task Learning이라는 점을 눈여겨볼 필요가 있습니다.

 

*리뷰할 논문: https://www.sciencedirect.com/science/article/pii/S1361841522002808?casa_token=cxsD2ADHN3IAAAAA:EPAGtIFuSJKosWJFgixL3M70uMZ5rhXwt1LfU72s4J_NBtgeLbPD8YBRie6tXaQr-VQqysxtvX0 (IF: 10.7)

* 해당 글에 첨부된 사진들은 모두 위 논문에서 발췌한 것임을 미리 밝힙니다.

* 해당 글에 설명은 논문을 기반으로 한 주관적 해석이므로 틀린 부분이 있다면, 댓글로 지적 부탁드립니다!


# Abstract

본 논문은 딥러닝을 이용한 암 영역 검출(CRD) 및 하위 유형 분류에 대한 연구입니다.

 

기존 연구에는 두 가지 주요 한계가 존재합니다.

  • 일반적으로 많은 노동력을 줄이려는 원래 의도와 모순되는 정확한 주석이 포함된 대규모 데이터 세트가 필요
  • 하위 유형 분류 작업의 경우 비암성 영역이 WSI 내의 암성 영역과 동일하게 취급되어 학습 과정에서 하위 유형 분류 모델에 혼란을 줌.

여기서 후자가 이 논문의 핵심입니다. 기존 연구의 경우, 이 문제를 해결하기 위해 하나의 모델로 CRD를 수행한 후, 그 결과에 대해 다시 다른 모델로 하위 유형의 암 유형을 분류하였습니다. 그러나, 이 방식의 경우 앞서 수행하는 모델이 CRD를 잘못 검출하면 하위 유형 분류까지 영향을 주게 되고, 암 영역과 하위 유형 간의 상관관계를 무시하게 된다는 단점이 존재합니다.

 

이에, 본 연구에서는 암 분류를 위한 반지도 학습 다중 작업 학습(MTL) 프레임워크를 제안합니다.

 

** CRD와 하위 유형 분류를 동시에 학습하여 두 작업 간의 상호 작용을 고려할 수 있다는 점이 포인트 **

 

MTL 프레임워크는 백본 특징 추출기, 두 개의 작업별 분류기 및 가중치 제어 메커니즘으로 구성됩니다. 

  • 백본 특징 추출기는 두 개의 작업별 분류기에서 공유되므로 CRD 및 하위 유형 분류 작업의 상호 작용을 포착할 수 있습니다. 
  • 가중치 제어 메커니즘은 이러한 두 작업의 순차적 관계를 유지하고 MTL 프레임워크에서 하위 유형 분류 작업에서 CRD 작업으로의 오류 역전파를 보장합니다.
  • 최소 점 기반(min-point) 주석 전략에 의해 생성된 소량의 주석만 포함하는 데이터 세트를 사용하여 전체 프레임워크를 반지도 설정으로 학습합니다.

# Introduction

기존에 병리학적 이미지를 컴퓨터로 처리하려 할 때 생긴 문제점은 아래와 같습니다. 

  • 학습 데이터의 라벨링에 너무 많은 시간이 걸린다. → 노동력이 너무 많이 필요하다.
  • WSI 이미지가 너무 커서 컴퓨터가 처리하기 힘들어, 패치 단위로 사용하게 됩니다. 다만, 하나의 세포 이미지를 패치화 시켜서 라벨링 하다 보니, 모든 패치가 다 같은 레이블을 가지게 됩니다. 하지만, 비암성 영역의 패치의 경우에는 오히려 암의 하위 유형 분류에 악영향을 줄 수 있게 됩니다.

첫 번째 문제에 대해서는 기존 연구들은 Supervised Learning 방식으로 학습하기에 WSI 전체 또는 특정 영역에 대한 주석을 필요로 했습니다. 하지만, 본 논문은 Semi-Supervised Learning (SSL) 방식으로 학습하기에 최소 점 주석 방법으로 몇 개의 점만으로 레이블을 진행해 주석 작업 시간을 크게 축소시킨다는 장점이 있습니다.

 본 논문에서는 최소 점 주석 방법을 제안합니다. (이후 내용에서 설명 예정)

 

두 번째 문제에 대해서 기존 연구에서는 CRD 분류 모델을 잘 학습시키고, 이를 이용해 비암성 영역을 제외한 후 암성 영역만 가지고 하위 유형 분류 작업을 수행했었습니다. 이는 정확도를 크게 올릴 수 있으나, CRD 모델이 잘못 예측하는 경우 하위 유형 분류도 잘못 예측이 됩니다. 또한, CRD와 하위 유형 분류의 상관관계에 대해서도 무시됩니다.

본 논문에서는 두 작업을 통합된 프레임워크에 포함하고 함께 훈련합니다.

 

** 추가 설명: WSI 이미지에서 CRD 작업의 경우 각 패치는 서로 다른 레이블을 가질 수 있음. Ex) 이 영역의 패치는 암성 영역이지만, 이 영역의 패치는 비암성 영역. 그러나, 하위 유형 분류의 경우에는 모든 패치가 같은 레이블을 가짐. 따라서, 위에서 비암성 영역으로 분류했더라도, 하위 유형 분류에서는 레이블이 있게 되어 혼란을 줌. **

 

또한, TCGA 데이터셋을 사용하되, 반지도 학습 방법을 채택하여 수행합니다. 그 결과, CRD 작업에서는 supervised learning 방식으로 학습된 모델과 성능(AUROC)이 비슷했으며, 하위 유형 분류 작업에서는 기존 연구 방식보다 최대 10% 더 뛰어났습니다. 마지막으로, 신장, 폐, 유방암의 세 가지 다른 암 유형의 공용 데이터 세트와 하나의 외부 유효성 검사 RCC 데이터 세트에서 프레임워크의 효과와 일반화 능력을 검증했습니다.


# Related Work

Semi-supervised learning: 제한된 양의 레이블링 된 데이터로 모델을 훈련한 다음 모델을 사용하여 레이블링 되지 않은 이미지에 대한 유사 레이블을 생성하여 추가 모델 훈련에 사용합니다.

 

Multi-task learning: MTL은 복잡한 작업을 하위 작업으로 분해하고 각 하위 작업을 기계 학습 모델을 통해 병렬적으로 학습하는 공동 학습 방식이며, 모델 성능은 각 하위 작업의 영향을 받을 수 있습니다.

  • 하드 매개변수 공유:  서로 다른 작업 앞에 계층에서 백본 네트워크를 공유하는 것
  • 소프트 매개변수 공유: 각 작업에 대해 서로 다른 백본 네트워크를 사용하고 해당 매개변수에 제약을 적용

논문에서는 MTL을 위해서는 CRD는 필수적인 전제 조건이며, 이는 CRD와 하위 유형 분류가 순차적으로 배열되어야 함을 강조합니다.


# Method

## Min-point annotation

Fig 2. Min-point annotation 그림

 

(a)는 병리학자가 직접 배율을 키웠다 줄였다 하며 레이블링 한 것으로 노동력이 매우 많이 들어가는 작업입니다. (b)는 낮음에서 중간 배율로 고정하여 대략적인 레이블링 한 것으로 소요 시간이 줄어들긴 하지만 그래도 많은 노동력이 필요합니다. (c)는 본 논문에서 제안하는 최소 점 주석 방법으로 레이블링 한 결과입니다.

 

최소 점 주석 방법의 규칙:

  • 암성 및 비암성 영역 모두에 동일한 수의 점을 표시 (암성을 5개 표시  비암성도 5개 표시, 본 논문에서는 이를 5로 설정)
  • 전체 이미지 내에서 점을 고르게 분포 (편향되지 않게)
  • 공백, 가장자리, 심하게 염색된, 손상된 및 기타 비정상적인 영역에 점을 표시하지 않음

주석은 P1, P2 (숙련된 병리학자), T1~T4 (의대생)에게 맡겼으며 OpenHI 소프트웨어로 처리하였습니다. 병리학자에게는 최소 점 주석 방법을 수행하게 하고, 시간 비교를 위해 몇 장의 이미지에 대해서만 기존 방식(a와 b 방식)으로 주석 처리를 맡겼습니다.

 

Table 1. 레이블링 시간 비교. (단위: second)

 

 그 결과, Complete 영역 주석과 Rough 영역 주석에 비해 각각 약 76%와 95%의 주석 시간을 절약할 수 있었습니다.


## Semi-supervised multi-task learning framework

Multi-task Learning은 훈련과 테스트 단계를 나누어 진행합니다.

 

훈련 단계:

  • 최소 점 주석이 있는 WSI에서 레이블링 된 및 레이블링 되지 않은 이미지 패치 추출.
  • 특별히 설계된 반지도 다중 작업 학습 프레임워크를 사용하여 다중 헤드 CNN 훈련. (이후 Algorithm1으로 설명.)

테스트 단계:

  • 테스트하고자 하는 WSI 이미지에서 이미지 패치 세트 획득.
  • 잘 훈련된 다중 헤드 CNN을 사용하여 각 이미지 패치의 CRD 및 하위 유형 레이블 예측.
  • CRD 분기에서 암성 영역으로 예측된 관련 패치의 예측된 하위 유형 레이블을 다수결 투표로 집계하여 테스트 WSI의 최종 하위 유형 레이블 획득.

## Semi-supervised multi-task loss

논문에서 학습 시 사용한 loss는 두 가지 loss를 결합한 형태로 완성됩니다.

 

1) 레이블링 된 데이터 (labeled patches)

l_ai: CRD Loss, l_bi: 하위 유형 분류 loss 각 패치에 대하여 loss들을 계산하고 더한 후, 가중치를 줍니다.

 

** 하위 유형 분류 loss에는 y_i가 곱해지게 되는데, 이는 해당 패치가 암성 영역인지, 비암성 영역인지를 뜻하는 것으로 만약 비암성 영역이라면, 이는 훈련에 방해를 주므로 0이 곱해져 영향을 없앰으로써 컨트롤합니다. **  

 

2) 레이블링 되지 않은 데이터 (unlabeld patches)

 

l_cj: CRD Loss, l_dj: 하위 유형 분류 loss. λc, λd: 각각 CRD 손실과 하위 유형 분류 손실의 가중치 스칼라. p_j: CRD 분류기가 u_j를 암성 영역으로 예측한 확률 (위와 동일 효과로 이 역시 가중치 역할을 합니다.)

 

** 단, 이 경우에는 암성 영역인지 비암성 영역인지 불명확하므로 예측 확률(0에서 1 사이 값)인 p_j를 곱하여 소프트 가중치를 주어 영향을 낮추기만 합니다. **

 

최종 loss는 이 두 가지를 결합하여 사용합니다. 이를 그림으로 나타내면 아래와 같습니다.

 

이를 단계별로 살펴본다면, 논문에서 Algorithm 1으로 설명한 것과 동일합니다.

 

Algorithm 1:

  • X: 라벨링 된 패치, U: 라벨링 되지 않은 패치, Y_labeld: 라벨링 된 패치의 CRD 레이블,  Z: 모든 패치들의 하위 유형 분류 레이블.
  • 가중치 초기화
  • 이다음부터는 반복 과정.
  • 미니 배치를 통해 데이터 샘플링.
  • U에 K 증강을 적용하여 U_k 구하기. (K 증강은 증강 작업을 K번 반복한 것으로 표현)
  • CNN 모델을 통해 U_k 데이터의 CRD 레이블 예측 후 전처리를 통해 y_unlabeld 구하기.
  • Y는 {labeld, unlabeld}로 함께 구성 후, X_m, U_m, Y_m, Z_m을 mixup 하여 X’, U’ 생성 후 loss 계산.
  • 가중치 업데이트.

# Training strategy

  • 모델이 학습 초기에는 라벨링 되지 않은 데이터에 대해서 학습이 어려우므로 𝜆𝑐가 훈련 과정에 걸쳐 선형적으로 최댓값까지 증가합니다. 단, 너무 큰 값이면 모델 성능에 큰 영향을 주므로 𝜆𝑐의 최댓값을 40으로 설정하여 제한합니다.
  • CRD 분류기가 수렴할 때까지 즉, 두 훈련 epoch에서 𝑙𝑎와 𝑙𝑐의 값 변화가 해당 경험적 임계값보다 모두 작을 때까지 𝜆𝑑를 0으로 설정하여 CRD 분류기 학습에 unlabeld data의 영향을 줄여줍니다. 그 후, 𝜆𝑑는 상수로 설정합니다.
  • 학습은 Adam Optimizer를 사용합니다.
  • Early Stopping 기법 사용

# Dataset

TCGA-RCC(명세포, 유두상, 신장 세포암종 데이터셋), TCGA-LU(선암, 편평 세포암 데이터셋), TCGA-BR (관상피암, 소엽암 데이터셋)으로 총 3가지 TCGA 데이터셋을 사용합니다. 또한, 일반화 능력 검증을 위해 인근 병원에서 얻은 LH-RCC 데이터셋도 활용합니다.

 

데이터는 훈련:검증:테스트 = 6:2:2 비율로 나누어 사용하게 됩니다.

 

하위 유형 최종 예측은 암성, 비암성 여부에 관계없이 다수결을 통해 집결.

 

* 다수결은 여러 패치들의 결과를 통해 집계. (예를 들어, {a1, a2, a3} 패치에 대해 모델의 예측 결과가 {1, 2, 1}이 나왔다면 1로 예측합니다.)


# Comparison methods

논문에서 제안된 모델: ResSSMTL-MD ResMTL-CD (MTL: Multi task learning)

 

ResMTL-CD ResSSMTL-MD 완전 지도 학습 버전이라고 생각하시면 됩니다. 결과를 보면, ResMTL-CD ResF-C 완전 지도 학습 모델이지만, 논문에서 제안한 모델과 0.1~0.4% 차이밖에 나지 않는 걸로 보아 Semi-Supervised Learning 방식이 높은 성능을 보임을 알 수 있습니다.


# Results

 

하위 유형 분류의 경우 TCGA-RCC, TCGA-BR, TCGA-LU 데이터에서 논문에서 제안된 모델 ( 마지막 ) 가장 좋은 성능을 보임을 알 수 있습니다.

 

 

제안된 모델 (g) 그림을 보면, 암 영역 검출에서도 마찬가지로, 기존 supervised learning 방식과 매우 유사한 결과를 보여주며 모두 실제 레이블(a) 굉장히 유사한 결과를 보여줌을 알 수 있습니다.

 

 

그리고, 훈련된 모델이 처음보는 데이터에 대해서도 얼마나 적용되는지 일반화 성능을 평가하기 위해 LH-RCC데이터셋을 적용한 결과를 비교하였습니다. Fine-Tuning 전과 후로 나누어 진행하며, Fine-Tuning 성능 향상의 결과를 보여주었습니다. 이는 논문에서 제안된 방식이 새로운 데이터에 적용되어도 빠르게 적용되며 높은 성능을 보일 있음을 의미합니다.


# Conclusion

논문에서는 기존의 레이블링 데이터셋이 부족하다는 한계점과 upstream task ( 영역 검출) downstream task (하위 유형 분류) 방식으로 진행하는 연구의 한계점을 지적하였습니다.

 

이에 min point annotation 방식을 채택하여 semi-supervised learning 방식으로 학습을 진행하고, multi-task learning 방식으로 가지 태스크를 동시에 진행하여 해결하였습니다. 논문에서 제안된 모델은 실제 기존에 연구된 supervised learning 방식의 모델보다 유사하거나 높은 성능을 기록하였습니다. 또한, 일반화 성능도 높아 추후 새로운 데이터에 적용 나은 성능을 보일 있다는 장점을 가지고 있습니다.

 

* 코드는 공개되어 있지 않습니다. (논문 github 없음 + github 존재하지 않음, paper with code에도 나와있지 않음)

 

마치며...

오늘은 병리학적 데이터에서 Semi-Supervised Learning 방식과 Multi-Task Learning 방식을 이용한 딥러닝 논문을 리뷰해 보았습니다. 의료 데이터에서는 항상 Annotation이 많은 노동력을 필요로 한다는 단점이 존재합니다. 이에, 많은 데이터셋이 공개되어 사용되었으면 하는 바람이 있습니다. (그래야 많은 연구가 진행될 테니까요.)

 

제 글이 이 논문을 읽으시려는 분들에게 조금이나마 도움이 되었으면 합니다. 오늘도 읽어주셔서 감사합니다 :)  

 

 

반응형