$μ$Match: Foundation Models for Semi-supervised Learning and Domain Adaptation in EM

µMatch 기술 요약: 전자현미경(EM)에서의 준지도 학습 및 도메인 적응을 위한 파운데이션 모델

문제 정의
전자현미경(EM)은 나노미터 규모의 세포 초구조를 분석하는 데 필수적인 모달리티입니다. 그러나 광학 현미경이나 컴퓨터 단층 촬영과 같은 다른 이미징 모달리티에 비해 비전 파운데이션 모델(VFM)을 EM에 적용하는 것은 제한적이었습니다. 이러한 한계는 세그멘테이션 작업의 높은 다양성(예: 핵, 미토콘드리아, 신경 돌기)과 포괄적으로 주석이 달린 데이터의 부족에서 기인합니다. 결과적으로, EM 세그멘테이션은 여전히 주로 지도 학습(supervised learning)에 의존하고 있으며, 이는 광범위한 수동 주석 작업을 필요로 하여 초구조 분석의 확장을 저해합니다. 특정 소기관(예: 미토콘드리아)을 위한 초기 파운데이션 모델은 존재하지만, 다양한 EM 소기관을 위한 종합적인 파운데이션 모델은 부재하며, 기존 방법들은 EM 도메인을 위한 준지도 학습(SSL) 및 도메인 적응(DA) 프레임워크에 VFM을 체계적으로 통합하지 못했습니다.

방법론
저자들은 EM에서 준지도 학습과 도메인 적응을 가능하게 하기 위해 VFM을 활용하는 통합 프레임워크인 µMatch를 제안합니다. 이 방법론은 세 가지 핵심 구성 요소로 이루어집니다:

1. 파운데이션 모델 초기화: 프레임워크는 SAM, SAM2, µSAM(광학 현미경 및 EM을 위한 도메인 특화 변형), DINOv2/v3를 포함한 다양한 VFM의 가중치를 사용하여 세그멘테이션 네트워크(특히 UNETR 아키텍처)의 인코더를 초기화합니다. 이는 처음부터 학습된 표준 UNet과 비교됩니다.
2. 학생-교사 학습(Student-Teacher Learning): µMatch는 EM에 적합하게 조정된 세 가지 최첨단 학생-교사 방법을 구현합니다:
   • Mean Teacher: 학생과 교사 모두에게 약한 증강(weak augmentations)을 사용하며, 학생 가중치의 지수 이동 평균(EMA)을 통해 교사를 업데이트합니다.
   • FixMatch: 교사에게는 약한 증강 뷰를, 학생에게는 강한 증강 뷰를 입력하며, 가중치를 공유하고 의사 라벨링(pseudo-labeling)을 위해 신뢰도 임계값을 사용합니다.
   • UniMatch v2: 교사에게 하나의 약한 뷰를, 학생에게 두 개의 강한 뷰를 사용하며, 특징 맵(feature maps)에 상보적 드롭아웃(complementary dropout)을 적용하여 견고성을 높입니다.
   • 학습 전략: 프레임에서 유의미한 교사 예측을 보장하기 위해, 무감독 손실(Dice loss)을 적용하기 전 1,000회 반복의 지도 학습 웜업(warm-up) 단계를 수행하며, 이때 임계값($t_c$)을 초과하는 의사 라벨링된 픽셀에 대해 손실을 계산합니다.
3. 학습 설정: 프레임워크는 세 가지 설정에서 평가됩니다:
   • 지도 학습(Supervised Learning): 라벨이 있는 데이터만 사용하여 학습합니다.
   • 준지도 학습(Semi-Supervised Learning, SSL): 동일한 도메인 내에서 적은 양의 라벨 데이터와 더 많은 양의 라벨 없는 데이터를 함께 사용하여 학습합니다.
   • 도메인 적응(Domain Adaptation, DA): 라벨이 완전히 지정된 소스 도메인에서 사전 학습한 후, 비지도 DA(UDA, 타겟 라벨 없음) 또는 준지도 DA(SSDA, 소량의 타겟 라벨 있음)를 사용하여 타겟 도메인에서 학습합니다.

주요 기여

• VFM의 체계적 평가: 본 논문은 핵, 미토콘드리아, 신경 돌기 세그멘테이션에 대해 여러 파운데이션 모델(SAM, SAM2, µSAM, DINOv2/v3)을 지도 학습, SSL 및 DA 설정에서 종합적으로 평가합니다.
• SSL/DA 방법의 구현: 저자들은 EM 데이터의 고유한 도전 과제(예: 경계 및 어피니티 출력에 대한 특정 증강 전략)를 해결하면서 Mean Teacher, FixMatch, UniMatch v2를 EM 작업에 맞게 조정하고 평가했습니다.
• 절제 연구(Ablation Studies): 신뢰도 임계값 전략(고정형 vs 동적 스케줄링)의 영향과 라벨링된 데이터 양의 변화에 따른 효과를 조사합니다.
• 다양한 데이터셋 벤치마킹: Drosophila, 마우스, 인간을 포함한 다양한 생물종과 이미징 조건을 다루는 Seymour(핵), MitoEM(미토콘드리아), CREMI(신경 돌기) 및 다양한 MitoEM v2 조건을 포함한 광범위한 데이터셋에서 실험을 수행합니다.

결과

• 지도 학습: VFM으로 초기화하는 것이 처음부터 학습하는 것보다 성능을 크게 향 향상시킵니다. 구체적으로, SAM 기반 가중치(SAM, SAM2, µSAM)로 초기화된 UNETR 모델이 가장 좋은 결과를 얻었습니다. 핵 세그멘테이션의 경우, 최적의 UNETR(DINOv3)은 처음부터 학습된 UNet의 mSA 0.097에 비해 0.456을 달고했습니다. 미토콘드리아의 경우, µSAM-lm 초기화가 UNet 베이스라인보다 11% 개선된 결과를 보였습니다. DINOv2/v3 백본은 일반적으로 SAM 기반 모델보다 성능이 낮았습니다.
• 준지도 학습(SSL): SSL은 모든 작업과 데이터셋에서 일관되게 성능을 향상시켰습니다.
  • 핵(Nuclei): µSAM-em 백본을 사용한 UniMatch v2가 가장 높은 mSA(0.574)를 달성하였으며, 이는 지도 학습 베이스라인(0.375)에 비해 상당한 개선입니다.
  • 미토콘드리아(Mitochondria): SAM 초기화를 사용한 Mean Teacher가 가장 좋은 결과(mSA 0.503)를 냈으며, 이는 지도 학습 모델보다 약 7.6%, UNet 베이스라인보다 약 16% 개선된 수치입니다.
  • 신경 돌기(Neurites): SAM 초기화를 사용한 UniMatch v2가 가장 높은 CREMI 점수(0.435)를 달성했습니다.
  • 관찰 사항: FixMatch는 이러한 EM 작업에서 지도 학습 베이스라인에 비해 성능 저하를 지속적으로 초래했습니다.
• 도메인 적응(DA): 서로 크게 다른 데이터셋 간(예: Seymour에서 FAFB/MICRONS로) 이동할 때, 순수 SSL을 타겟 도메인에 직접 학습시키는 것이 DA 접근 방식(UDA 및 SSDA)보다 일반적으로 더 우수한 성능을 보였습니다. 이는 현재 EM에서 큰 도메인 변화가 표준 DA 전략의 효과를 제한하고 있음을 시사합니다. 그러나 분포 변화가 작은 경우(예: 특정 MitoEM v2 조건)에는 SSDA가 개선 효과를 보였습니다. 사전 학습된 모델(MitoNet, µSAM)의 제로샷(Zero-shot) 성능은 일반적으로 낮았으며, 이는 미세 조정(fine-tuning)의 필요성을 강조합니다.

의의 및 주장
본 논문은 µMatch가 준지도 학습 및 도메인 적응 프레임워크 내에서 파운데이션 모델을 효과적으로 활용함으로써 EM의 주석 작업 노력을 실질적으로 줄이는 경로를 구축한다고 주장합니다. 저자들은 VFM이 지도 학습 설정에서 엄청난 상승을 제공하는 한편, 이를 학생-교사 SSL 프레임워크에 통합하면 강력한 베이스라인보다 일관되고 유의미한 개선을 얻을 수 있음을 강조합니다.

연구의 결론은 다음과 같습니다:

1. SAM 기반 모델이 EM 세그멘테이션 작업에 가장 효과적인 초기화 방법인 반면, DINO 변형은 이 문맥에서 이점을 제공하지 못합니다.
2. 학생-교사 SSL은 EM에 매우 유익하며, Mean Teacher는 성능과 계산 비용 사이의 견고한 절충안을 제공하고, UniMatch v2는 계산 자원이 허용될 때 최상의 결과를 제공합니다.
3. EM에서의 도메인 적응은 큰 도메인 변화로 인해 현재 제한적이며, 소스 도메인으로부터의 DA보다 타겟 도메인에서의 SSL이 종종 더 우수한 성능을 보입니다. 그러나 이 방법론은 도메인 변화가 작은 시나리오에서는 유망합니다.
4. 이 프레임워크는 사용자가 주석을 단 데이터에 대해 미세 조정을 용이하게 하기 위해 기존 현미경 도구(예: µSAM)에 통합되도록 설계되었으며, 이를 통해 초구조 분석의 병목 현상인 수동 주석 문제를 해결합니다.