Synthetic Data Generation for Long-Tail Medical Image Classification: A Case Study in Skin Lesions

기술적 요약: 장기 꼬리 의료 이미지 분류를 위한 합성 데이터 생성

문제 제기
장기 꼬리 클래스 분포는 다중 클래스 의료 데이터셋에서 광범위하게 존재하며, 심층 학습 모델에 상당한 도전을 제기합니다. 이러한 시나리오에서 소수의 지배적인 '헤드' 클래스는 대부분의 샘플을 포함하는 반면, 수많은 임상적으로 중요한 '꼬리' 클래스 (종종 희귀하고 고위험 질환에 해당) 는 심각하게 과소 대표됩니다. 이러한 편향된 데이터로 훈련된 모델은 일반적으로 경쟁력 있는 전체 정확도를 달성하지만, 꼬리 클래스에 대한 일반화에는 실패합니다. 손실 함수 설계, 모델 아키텍처 개선, 클래스 재균형으로 분류되는 기존 솔루션들은 오직 미미한 개선만 제공합니다. 이러한 방법들은 종종 정밀한 하이퍼파라미터 튜닝에 의존하며, 서로 다른 모델 간 일반화 능력이 제한적이고, 수동으로 제작된 데이터 증강의 결정론적 특성으로 인해 확장성에 어려움을 겪습니다. 또한, 확산 모델이 자연 이미지 생성에서 성공을 거두었음에도 불구하고, 도메인 특화적이고 장기 꼬리인 의료 영상에 대한 적용, 특히 다양하면서도 임상적으로 의미 있는 샘플 생성과 관련하여서는 아직 충분히 탐구되지 않았습니다.

방법론
저자들은 피부 병변 분류에 구체적으로 적용된 장기 꼬리 의료 이미지 분류 (LTMIC) 를 위한 확산 모델 기반 합성 데이터 증강 파이프라인을 제안합니다. 이 접근법은 세 가지 주요 단계로 구성됩니다:

1. 인페인팅 확산 모델 아키텍처:
   이 방법의 핵심은 새로운 인페인팅 확산 모델입니다. 고도로 전문화된 의료 설명에는 적합하지 않은 표준 텍스트 임베딩 모델과 달리, 저자들은 간단한 완전 연결 클래스 임베딩 네트워크를 훈련합니다. 이 모델은 네 가지 입력을 받습니다: 원본 피부경 이미지 ($x$), 이진 관심 영역 마스크 ($b$), 피부경 이미지의 배경 ($d$), 그리고 클래스 레이블 ($y$) 입니다. 이미지와 배경은 사전 훈련된 변이 오토인코더 (VAE) 를 통해 잠재 공간으로 인코딩됩니다. 이진 마스크는 잠재 차원과 일치하도록 재조정됩니다. 이러한 구성 요소들은 연결되어 $\theta_d$로 매개변수화된 UNet 에 입력되어 전방 확산 과정에서 추가된 노이즈를 추정합니다. $\theta_c$로 매개변수화된 클래스 임베딩 네트워크는 레이블 조건을 제공합니다.

2. 샘플 선택을 통한 미세 조정:
   확산 모델이 각 클래스의 분포를 균등하게 잘 학습하도록 보장하기 위해, 미세 조정 동안 오버샘플링 전략이 사용됩니다. 저자들은 생성 모델을 훈련하는 데 '대략 분포 내' (AID) 샘플 (분포 중심에서 약간 떨어진 샘플) 이 가장 중요하다는 이전 연구의 결과를 활용합니다. 따라서, 각 샘플에 대한 이상 점수를 계산하기 위해 자기지도 학습 의료 이상 탐지 방법이 사용됩니다. 가중치는 샘플이 해당 클래스의 중앙 이상 점수로부터의 거리에 따라 할당되며, AID 영역에 더 가까운 샘플은 더 높은 가중치를 부여받고 모든 클래스 간의 미세 조정 데이터셋을 균형 있게 만들기 위해 반복됩니다.

3. 합성 생성 및 OOD 필터링:
   미세 조정된 모델은 원본 훈련 세트의 배경 이미지와 클래스 레이블을 입력으로 받아 합성 이미지를 생성합니다. 확산 모델의 확률적 특성으로 인해, 초기 합성 데이터셋 ($D_{syn}$) 은 임상적으로 유효하지 않을 수 있는 분포 외 (OOD) 이미지를 포함한 '노이즈' 샘플을 포함합니다. 이를 해결하기 위해 $D_{syn}$에 의료 OOD 필터링 메커니즘을 적용하여 정제된 합성 데이터셋 ($D_{clean}^{syn}$) 을 생성합니다. 최종 훈련 데이터셋은 원본 장기 꼬리 데이터셋 ($D_{LT}$) 과 $D_{clean}^{syn}$을 결합하여 형성됩니다. 하이퍼파라미터 $\gamma$는 각 클래스별로 합성 세트에서 유지되는 정제된 샘플의 비율을 제어합니다.

주요 기여
본 논문은 세 가지 주요 기여를 제시합니다:

1. 새로운 파이프라인: LTMIC 문제를 해결하도록 특별히 설계된 인페인팅 확산 모델과 OOD 필터링을 통합한 합성 데이터 생성 파이프라인.
2. 전문 아키텍처: 범용 텍스트 임베딩 대신 전용 클래스 임베딩 네트워크와 배경 조건화를 활용하여 의료 이미지 생성에 맞춤화된 새로운 인페인팅 확산 모델 아키텍처.
3. 샘플링 전략: 의료 장기 꼬리 데이터셋에서 확산 모델을 미세 조정하기 위한 새로운 OOD 기반 샘플링 전략으로, 생성된 데이터의 품질을 향상시키기 위해 꼬리 클래스에 대해 AID 샘플을 우선시합니다.

실험 결과
이 방법은 8 개 클래스에 걸쳐 25,331 개의 이미지를 포함하며 심각한 불균형 (모반은 12,875 개 샘플에서 혈관 병변은 62 개 샘플까지) 을 보이는 벤치마크인 ISIC2019 피부 병변 분류 데이터셋에서 평가되었습니다. 실험은 데이터 생성 방법의 영향을 격리하기 위해 5 교차 검증과 앙상블이나 외부 데이터 없이 간단한 CNN 기반 분류기를 사용했습니다.

• 전체 성능: 제안된 접근법은 다중 모델 앙상블 (MRE) 및 기타 확산 기반 접근법 (예: CBDM, Dreambooth) 을 포함한 현재 최첨단 (SOTA) 방법들을 크게 능가하는 0.802 의 균형 잡힌 다중 클래스 정확도 (BMA) 를 달성했습니다. 특히, 외부 데이터를 필요로 하지 않고 BMA 에서 MRE 방법보다 약 4% 더 높은 성능을 보였습니다.
• 꼬리 클래스 개선: 가장 적은 샘플을 가진 클래스 (피부섬유종, DF) 에서 가장 큰 향상이 관찰되었습니다. 이 클래스의 경우, 합성 데이터가 없는 기준선 대비 정확도가 28% 이상 향상되었습니다.
• OOD 필터링의 영향: 연구 결과 OOD 필터링이 성능 균형을 맞추는 데 중요하다는 것이 밝혀졌습니다. 필터링 없이 합성 데이터는 꼬리 클래스를 크게 개선시켰지만, 헤드 클래스 성능을 약간 저하시키는 노이즈를 도입했습니다. 필터링 과정은 꼬리 클래스에서의 이득을 유지하면서 헤드 클래스 성능을 회복시켰습니다.
• 하이퍼파라미터 민감도: $\gamma$(정제된 합성 샘플의 유지율) 에 대한 제거 연구는 소량의 합성 데이터로도 성능이 급격히 향상되지만, $\gamma$가 1 에 가까워질수록 (필터링 없음) 성능이 정체되거나 감소함을 보여주었습니다. $\gamma$의 최적 범위는 0.2 에서 0.6 사이로 확인되었습니다.

의의 및 주장
저자들은 그들의 결과가 장기 꼬리 불균형을 완화하고 의료 분류의 견고성을 강화하는 확산 기반 증강의 효과성을 입증한다고 주장합니다. 다양하고 사실적이며 임상적으로 의미 있는 합성 샘플을 생성함으로써, 이 방법은 하이퍼파라미터에 민감한 복잡한 모델 앙상블이나 전문 손실 함수에 의존하지 않고 희귀 질환의 데이터 부족 문제를 해결합니다. 논문은 현재 접근법이 데이터 생성과 분류를 독립적인 과정으로 취급하지만, 향후 연구는 LTMIC 성능을 더욱 향상시키기 위해 이러한 구성 요소들의 공동 최적화에 초점을 맞출 것이라고 결론지었습니다. 이 방법은 어떤 분류 아키텍처와도 통합 가능한 모델 무관한 것으로 제시됩니다.