OTLesMix: Wasserstein Barycenter and Optimal Transport Map for Synthetic Lesion Generation with Diverse Shapes and Locations

기술 요약: OTLesMix: 합성 병변 생성을 위한 Wasserstein Barycenter 및 Optimal Transport Map

문제 정의
딥러닝은 의료 영상 분할 분야에 혁신을 가져왔으나, 지도 학습(supervised training)은 특히 정밀한 3D 분할 작업에서 전문가의 주석(annotation)에 소요되는 높은 비용과 시간이라는 제약에 직면해 있습니다. 데이터 증강(data augmentation)은 모델 학습을 개선하기 위한 표준적인 기법이지만, 전통적인 방식(공간적/강도 변환)과 고급 생성 모델(GAN, Diffusion Models)은 각각 특정한 한계를 가집니다. 생성 모델은 자체 학습을 위해 대규모의 다양하고 풍부한 데이터셋이 필요하며, 이는 의료 환경에서 가용하지 않을 수 있습니다. 반대로, 기존의 "믹스 기반(mix-based)" 증강 전략(예: MixUp, CutMix, CarveMix, LesionMix)은 병변의 형태와 위치에 충분한 다양성을 생성하는 데 실패하는 경우가 많습니다. 이러한 방식들은 대개 기존 마스크의 단순 무작위 샘플링이나 경직된 기하학적 변환에 의존하며, 그 결과 기존 데이터의 단순 복제나 회전에 불과하여 분할 모델을 강건하게 학습시키는 데 필요한 다양성을 확보하지 못합니다.

방법론: OTLesMix
저자들은 Optimal Transport (OT) 이론을 활용하여 현실적이고 다양한 합성 병변을 생성하는 새로운 영상 합성 방법인 OTLesMix를 제안합니다. 이 방법은 베이스 데이터셋 $D = \{(X_i, Y_i)\}$로부터 새로운 샘플 $(\tilde{X}, \tilde{Y})$를 생성하기 위해 두 가지 주요 단계를 거칩니다:

1. Wasserstein Barycenter를 통한 병변 마스크 생성:

   • 두 개의 베이스 샘플 $(X_i, Y_i)$와 $(X_j, Y_j)$를 무작위로 선택합니다.
   • 그라운드 트루스 마스크 $Y_i$와 $Y_j$ 내 비제로(non-zero) 픽셀의 좌표를 확률 분포로 취급합니다.
   • 새로운 병변 마스크 $\tilde{Y}_{les}$는 이 두 분포의 Wasserstein barycenter를 계산하여 생성됩니다. 이는 소스 마스크들로의 가중 합산된 운송 비용(transport costs)을 최소화하도록 새로운 마스크 좌표를 최적화하는 과정을 포함합니다:
     $$B(\lambda, \alpha_i, \alpha_j, \beta) = \lambda OT(\alpha_i, \beta) + (1-\lambda)OT(\alpha_j, \beta)$$
     여기서 $\lambda \sim U(0,1)$은 무작위 가중치 매개변수입니다.
   • 이 과정은 두 원본 마스크 사이의 형태와 위치를 보간(interpolation)하며, 이를 통해 새로운 병변은 원본 위치들 사이의 어딘가에 존재하게 되고, 형태는 원본들의 부드러운 보간 형태를 띠게 됩니다.
   • 계산 효율성을 보장하기 위해 저자들은 de-biased Sinkhorn divergence와 Sinkhorn 알고리즘을 활용하여, 현대적인 GPU 상에서 초 단위 내에 최대 $10^6$개의 포인트를 처리할 수 있도록 했습니다.

2. Optimal Transport Map을 통한 영상 및 레이블 생성:

   • 새로운 마스크 $\tilde{Y}_{les}$가 정의되면, 합성된 병변의 픽셀 강도(intensity)와 레이블을 결정해야 합니다.
   • 생성된 마스크 $\tilde{Y}_{les}$와 하나의 소스 마스크(예: $Y_j$) 사이의 optimal transport plan이 계산됩니다.
   • 새로운 마스크의 각 픽셀에 대해, 운송 맵(transport map) $T$는 해당 픽셀의 대응하는 위치를 소스 영상 $X_j$에서 찾아냅니다.
   • 강도 합성(Intensity Synthesis): 합성된 픽셀의 강도는 매핑된 위치에서의 소스 영상 강도를 선형 보간하여 얻습니다.
   • 레이블 합성(Label Synthesis): 다중 클래스 작업의 경우, 소스 레이블 맵의 최근접 이웃(nearest-neighbor) 보간을 통해 그라운드 트루스 레이벨이 할당됩니다.
   • 합성된 병변은 베이스 영상( $X_i$, $X_j$, 또는 데이터셋의 제3의 영상) 위에 인페인팅(inpainting)됩니다. 후처리 단계에서는 기존 병변과의 중첩이나 비현실적인 해부학적 위치(예: 두개골, 뇌척수액)를 제거합니다.

주요 기여

• 새로운 합성 프레임워크: Wasserstein barycenter를 통한 형태/위치 보간과 optimal transport map을 통한 강도/레이블 전이를 결합한 독창적인 OTLesMix를 도입했습니다.
• 향상된 다양성: 이 방법은 기존의 믹스 기반 방식들이 가진 한계를 명시적으로 해결하여, 단순히 기존 데이터를 복사하거나 단순 회전시킨 것이 아닌, 다양한 형태와 위치를 가진 병변을 생성합니다.
• 최첨단 성능(State-of-the-Art Performance): OTLesMix는 표준 믹스 기반 증강 기법들에 비해 다운스트림 분할 작업의 성능을 향상시키는 데 탁월한 성능을 입증했습니다.

실험 결과
저자들은 nnUNet 백본을 사용하여 세 가지 뇌 병변 분할 데이터셋에서 OTLesMix를 평가했습니다:

1. BraTS 2020: 교모세포종 및 저등급 신경교종 분할 (2D 슬라이스).
2. ATLAS v.2.0: 뇌졸중 신경 영상 (3D 볼륨).
3. ISLES 2022: 허혈성 뇌졸중 병변 분할 (3D 볼륨).

정량적 결과:

• OTLesMix는 모든 데이터셋에서 최첨단 믹스 기반 방법들(MixUp, CutMix, CarveMix) 및 베이스라인(전통적인 데이터 증강만 적용된 nnUNet)보다 우수한 성능을 보였습니다.
• 베이스라인 대비 개선도: 합성 데이터를 사용하지 않고 훈련된 모델에 비해 Dice 유사도 계수(DSC)가 2.9에서 6.6 포인트 향상되었습니다.
• 경쟁 모델 대비 개선도: OTLesMix는 두 번째로 우수한 방법보다 0.8에서 1.8 포인트 높은 DSC를 달성했습니다.
• 통계적 유의성: 다른 생성적 방법들과 비교했을 때 개선 효과는 통계적으로 유의미했습니다($p \le 0.05$). 가장 큰 p-value는 ATLAS 데이터셋에서 CarveMix와 비교했을 때의 0.012였습니다.

정성적 결과:
시각적 검사를 통해 OTLesMix가 다른 방법들에 비해 형태와 위치 면에서 더 높은 다양성을 가진 병변을 생성함을 확인했습니다. 이 방법은 다중 클래스 분할(BraTS)과 여러 개의 연결된 구성 요소를 가진 병변(ATLAS, ISLES)을 성공적으로 처리하였으며, 보간된 형태의 일관성을 유지했습니다.

의의 및 한계
본 논문은 OTLesMix가 데이터가 부족한 의료 영상 분야에서 강력한 데이터 증강 솔루션을 제공하며, 믹싱 방법의 단순함과 전체 생성 모델 학습의 복잡성 사이의 중간 지점을 제공한다고 주장합니다. Optimal Transport를 활용함으로써, 합성된 샘플이 기하학적 및 강도 측면에서 기초 데이터 분포와 일치하도록 보장합니다.

그러나 저자들은 다음과 같은 특정 한계를 인정합니다:

• 계산 비용: 단일 샘플을 생성하는 데 2D 영상의 경우 약 1초, 3D 볼륨의 경우 해상도에 따라 1~3분이 소요됩니다.
• 현실성 제약: 이 방법은 픽셀 강도의 선형 보간에 의존합니다. 만약 믹싱에 사용된 두 소스 영상 사이에 상당한 도메인 차이(예: 서로 다른 스캐너 또는 획득 프로토콜)가 존재할 경우, 비현실적인 인페인팅된 병변이 생성될 수 있습니다.
• 향후 방향: 저자들은 향end work로서, 합성된 강도의 현실성을 더욱 개선하기 위해 OTLesMix의 마스크 생성 절차를 마스크 조건부 딥러닝 생성 모델과 결합할 것을 제안합니다.