Flow Matching with Optimized Subclass Priors for Medical Image Augmentation

기술 요약: 의료 이미지 증강을 위한 최적화된 서브클래스 사전분포를 활용한 흐름 매칭

문제 제기
의료 영상 데이터셋은 심각한 장꼬리 불균형이 특징이며, 흔한 질환이 지배적인 반면 드물고 임상적으로 중요한 병리 (예: 작은 기흉, 기관지확장증) 는 과소대표됩니다. 이러한 불균형으로 인해 판별형 분류기는 헤드 클래스에 과적합되고 희귀 질환에서는 실패하게 됩니다. 생성형 데이터 증강은 잠재적인 해결책을 제공하지만, 표준 조건부 흐름 매칭 (FM) 과 확산 모델은 거친 질병 레이블로 훈련될 때 두 가지 중첩된 어려움에 직면합니다:

1. 다중 모드 집계: 단일 진단 레이블은 종종 다양한 아형, 획득 프로토콜, 그리고 동반 질환을 집계합니다. 이는 생성기를 지배적인 서브모드로 편향시키는 다중 모드 조건부를 생성하여, 드물지만 임상적으로 유효한 변형을 효과적으로 잘라냅니다.
2. 비효율적인 운송 기하학: 모든 생성 궤적을 공유된 표준 가우스 소스에서 끌어오면 특정 하위 집단이 불균형적으로 긴 운송 경로를 통과하게 됩니다. 이는 회귀 난이도를 과장하고 꼬리 클래스 생성의 충실도를 저하시킵니다.

방법론
저자들은 흐름 매칭 모델의 핵심 훈련 루프를 수정하지 않고도 생성 과정을 정제하기 위해 두 수준에서 정보적 사전분포를 도입하는 오프라인 전략을 제안합니다.

1. 가우스 혼합 모델링 (GMM) 을 통한 서브클래스 유도:

   • 이 방법은 생성 모델의 잠재 공간 내에서 각 거친 레이블 $c$를 일관된 서브모드 (서브클래스 $k$) 로 분할합니다.
   • 고정된 VAE(FLUX1.dev) 를 사용하여 결정론적 잠재 변수를 계산합니다. 각 클래스 $c$에 대해, 클래스 중심 $\mu_c$를 뺀 후 잔여 잠재 변수에 대각 가우스 혼합 모델을 적합시킵니다.
   • 구성 요소의 수 $K_c$는 확장 베이지안 정보 기준 (EBIC) 을 사용하여 선택됩니다. 샘플은 서브클래스에 대한 하드 할당을 받으며, 이는 조건부 분산이 더 낮은 분할을 생성합니다.
   • 이론적으로 저자들은 이러한 정제된 레이블에 조건을 부여하면 베이스 최적 흐름 매칭 위험이 증가할 수 없음을 증명합니다. 개선은 조건부 평균 변위의 서브클래스 간 변동성의 기대값과 같습니다.

2. 서브클래스 의존적 소스 최적화:

   • 공유된 가우스 소스 대신, 이 방법은 각 서브클래스에 대한 특정 소스 분포 $\pi_0(\cdot | c, k)$를 학습합니다. 이러한 소스는 학습 가능한 가우스 $N(\mu_{c,k}, \text{diag}(\sigma^2_{c,k}))$로 매개변수화되어 시작 분포를 재중앙화하고 재스케일링합니다.
   • 이는 전형적인 운송 궤적을 단축하고 서브클래스 내 분산을 줄입니다.
   • 기하학적 제어: 퇴화한 해 (예: 노름 팽창 또는 지나치게 집중된 경로) 를 방지하기 위해 저자들은 다음을 강제하는 오프라인 프록시 목적 함수를 최적화합니다:
     • 방향적 응집성: 정규화된 변위 방향이 학습 가능한 프로토타입과 적절히 정렬되도록 장려합니다.
     • 경로 길이 제한: 클래스 중심까지의 거리의 99 백분위수로 정의된 서브클래스별 제한 ($r_{cap}$) 을 사용하여 이상치에 고정된 페널티를 부과합니다. 이는 모델이 변위 노름을 팽창시켜 손실을 우회하는 것을 방지합니다.
     • 스케일 정규화: 소스 분포의 로그 스케일을 단위 참조 방향으로 정규화합니다.

주요 기여

1. 서브클래스 조건부 공식화: GMM 유도 하드 할당을 통해 다중 모드 조건부를 분할하여 조건부 모호성을 효과적으로 줄이는 장꼬리 의료 이미지 생성을 위한 새로운 접근법.
2. 학습된 서브클래스 의존적 소스: 서브클래스별로 소스 분포를 재중앙화하고 재스케일링하여 조건부 운송을 재구성하는 정보적 사전분포의 도입으로, 서브클래스 내 분산을 감소시킵니다.
3. 기하학적으로 정규화된 목적 함수: 적절한 각도적 응집성과 명시적인 경로 길이 이상치 제어 사이의 균형을 맞추는 실용적인 프록시 목적 함수로, 노름 팽창 단축과 방향적 붕괴를 방지합니다.
4. 실증적 검증: 흉부 X 선 (MIMIC-LT, NIH-LT) 과 CT(CT-RATE) 벤치마크에서 향상된 꼬리 커버리지와 다운스트림 장꼬리 분류 성능을 입증했습니다.

실험 결과
이 방법은 세 가지 장꼬리 의료 영상 벤치마크에서 평가되었습니다:

• 다운스트림 분류: 제안된 방법은 비증강 베이스라인 및 기타 증강 전략 (Vanilla FM, CBFM, CPD) 에 비해 세 가지 데이터셋 전체에서 일관되게 가장 높은 균형 정확도 (bAcc) 와 매크로-F1 점수를 달성했습니다. 특히 X 선과 CT 모드 모두에서 성능을 신뢰성 있게 향상시킨 유일한 전략이었습니다.
• 생성 품질: MIMIC-LT 에서 이 방법은 프레이chet 시작 거리 (FID) 와 시작 비율 점수 (IRS) 에서 최고의 성과를 거두어 충실도와 다양성의 동시 향상을 나타냈습니다. CT-RATE 에서는 FID 가 베이스라인보다 약간 높았지만, 높은 IRS 와 우수한 다운스트림 분류를 유지하여 단순 암기가 아닌 더 나은 분포 커버리지를 시사했습니다.
• 어블레이션 연구: 결과에 따르면 파이프라인의 두 단계 (서브클래스 분류 및 소스 최적화) 모두 필요했습니다. 서브클래스 분류만으로도 성능이 향상되었지만, 조합이 가장 좋은 결과를 낳았습니다. 서브클래스 분류 없이 최적화하는 경우 때로는 상충되는 모드를 평균화하여 성능을 저하시켰습니다.
• 서브클래스 유효성: 분석은 유도된 서브클래스가 획득 메타데이터뿐만 아니라 병리학 관련 구조를 포착함을 보여주었습니다. 교란자 탐지기는 메타데이터가 감지 가능하지만 지배적이지는 않음을 나타냈으며, 학습된 서브클래스의 kNN 순도는 무작위 분할을 크게 상회했습니다.

의의 및 주장
이 논문은 소스 설계와 조건부 정제가 임상적 장꼬리 환경에서 견고한 생성 증강을 위한 효과적인 레버라고 주장합니다. 이질적인 레이블을 분할하고 서브클래스 의존적 소스를 학습함으로써, 이 방법은 조건부 모호성을 줄이고 계산 비용이 거의 들지 않는 운송 궤적을 안정화합니다 (전체 과정은 오프라인이며 데이터셋당 1 GPU 시간 미만이 소요됨). 저자들은 이를 꼬리 클래스 생성의 충실도와 다양성을 향상시켜 희귀 의료 질환에 대한 더 신뢰할 수 있는 다운스트림 분류기로 이어지는 원칙적인 전략으로 위치시킵니다. 이 연구는 운송 기하학 문제로 인해 단순 조건부 생성이 종종 드문 하위 집단에서 실패하며, 이는 구조화된 사전분포를 통해 완화될 수 있음을 강조합니다.