StainSolver: Accelerating Diffusion-Controlled Neural Style Transfer with DPM-Solver++ and Quality-Driven Stain Normalization

기술 요약: StainSolver

문제 정의
조직 병리 이미지, 특히 헤마톡실린 및 에오신(H&E) 염색된 전체 슬라이드 이미지(WSI)는 스캐너, 시약 배치 및 고정 프로토콜의 차이로 인해 상당한 색상 및 대비 가변성을 겪습니다. 이러한 "염색 가변성"은 특정 실험실에서 학습된 컴퓨터 보조 진단 모델의 성능을 저하시키는 도메인 시프트를 생성합니다. 신경 스타일 전이(NST)와 생성적 적대 신경망(GAN)이 이러한 이미지를 정규화하기 위해 적용되어 왔으나, 이들은 결정적인 한계에 직면해 있습니다. 전통적인 NST는 반복적인 최적화로 인해 임상 규모에서는 계산 비용이 너무 많이 들며, GAN은 병리학적 구조를 환각(hallucination)할 위험이 있습니다.

최근 연구인 StainFuser는 대규모 의사 쌍 데이터셋(SPI-2M)으로 학습된 조건부 잠재 확산 모델을 사용하여 이러한 문제를 해결하였으며, 환각 현상 없이 고품질의 정규화를 달кси했습니다. 그러나 StainFuser는 의사 수치 방법 확산 모델(PNDM) 스케줄러에 의존하며, 이는 계산 비용이 많이 들고(이미지당 약 20단계 필요) 구조적 보존과 염색 충실도를 동시에 정량화하는 전용 평가 지표가 부족합니다.

제안된 프레임워크인 StainSolver는 모델 가중치를 재학습하지 않고 고정된 StainFuser 백본(Stable Diffusion v2.1 기반)을 직접 활용하여 구축되었습니다. 이 모델은 두 가지 별도의 확장을 도입합니다:

1. DPM-Solver++를 통한 추론 가속화:
   저자들은 기본 PNDM 스케줄러를 고차 상미분 방정식(ODE) 솔버인 **DPM-Solver++**로 교체합니다. DPM-Solver++는 로그-SNR 도메인에서 작동하며 2차 적분 기법을 사용합니다. 이를 통해 선형 다단계 방식인 PNDM에 비해 수렴 품질을 유지하거나 개선하면서도 훨씬 적은 단계(5~10단계)로 확산 궤적을 근사할 수 있습니다. 구현 과정에서는 VAE, ControlNet, UNet2D 가중치를 건드리지 않고 호환성을 보장하기 위해 PNDM 특유의 설정 파라미터를 필터링하는 과정을 포함합니다.

2. 염색 정규화 품질 점수 (SNQS):
   소스 이미지에 대한 충실도(fidelity)를 측정하는 일반적인 SSIM 및 PSNR의 부적절함을 해결하기 위해, 저자들은 SNQS를 제안합니다. 이 지표는 CIE-LAB 색 공간에서 작동하여 휘도(L)와 채도(a, b)를 분리합니다. 이는 다음 두 구성 요소의 가중 합으로 이루어집니다:

   • StructSim: 출력물과 소스 이미지 간의 휘도 채널에 대해 계산된 구조적 유사도 지수(SSIM)로, 조직 형태가 보존되었는지 확인합니다.
   • ColorMatch: 출력물과 참조 타겟 이미지 사이의 $a$ 및 $b$ 채널 히스토그램 간의 평균 피어슨 상관계수($\rho$)로, 염색 분포가 정확하게 전이되었는지 확인합니다.
     최종 점수는 $SNQS = \alpha \cdot StructSim + (1-\alpha) \cdot ColorMatch$로 정의되며, 기본값으로 $\alpha=0.5$를 사용합니다.

주요 기여

• 스케줄러 교체: DPM-Solver++가 모델 재학습 없이 StainFuser의 드롭인(drop-in) 대체제로 기능할 수 있음을 입증하였으며, 이를 통해 추론 시간을 획기적으로 단축했습니다.
• 새로운 평가 지표 도입: 구조적 보존과 염색 충실도를 공동으로 정량화하여 기존의 염색 정규화 평가 방법의 공백을 메우는 작업 특화 지표인 SNQS를 도입했습니다.
• 효율성-품질 트레이드오프 분석: 다양한 단계 수에 따른 DPM-Solver++와 PNDM의 수렴 프로파일에 대한 포괄적인 분석을 제공합니다.

실험 결과
실험은 100개의 소스 이미지와 512개의 참조 타겟을 포함하는 SPI-2M 데이터셋에서 수행되었습니다. 결과는 다음과 같습니다:

• 속도 및 품질: 10단계의 DPM-Solver++(DPM++@10)는 베이스라인인 PNDM@20보다 1.92배 빠른 속도(이미지당 4.14초 vs 7.94초)를 달성하는 동시에 SNQS를 3.7% 향상(0.691 vs 0.667)시켰습니다.
• 우수한 수렴성: DPM++@10은 속도와 품질 모두에서 PNDM@20을 능가했습니다. 특히, 단 5단계의 DPM++가 50단계의 PNDM을 앞질렀으며, 7.6배의 속도 향상(2.43초 vs 18.40초)과 더 높은 SNQS(0.704 vs 0.691)를 달성했습니다.
• 구성 요소 분석: 품질 향상은 주로 ColorMatch 구성 요소의 +7.0% 개선에 의해 주도되었으며, 이는 DPM-Solver++가 참조 타겟에 더 가까운 염색 분포를 가진 출력을 생성함을 나타냅니다.
• 강건성: DPM++의 성능 우위는 100개의 테스트 이미지 전체에서 일관되게 유지되었습니다. 또한, SNQS 지표는 가중치 파라미터 $\alpha$의 변화에도 강건함을 보였습니다.
• 타겟 선택: 본 연구는 소스와 유사한 히스토그램을 가진 참조 타겟을 선택하는 것이 구조적 보존에는 영향을 주지 않으면서 ColorMatch를 크게 향상(+14.9%)시킨다는 점을 확인하였으며, 이는 사전 선택 전략이 추론 비용 증가 없이 품질을 더욱 높일 수 있음을 시사합니다.

의의
본 논문은 StainSolver가 확산 기반 방법의 계산적 병목 현상을 줄임으로써 기가픽셀 단위 WSI의 실시간 정규화를 가능하게 한다고 주장합니다. 스케줄러를 모델 가중치로부터 분리함으로써, 이 접근 방식은 기존 StainFuser 배포 모델에 즉각적이고 실용적인 업그레이드를 제공합니다. 또한, SNQS의 도입은 일반적인 이미지 충실도를 넘어 염색 전이와 형태학적 보존을 구체적으로 측정하는 데 있어 보다 진단적으로 유의미한 평가 프레임워크를 제공합니다. 저자들은 이러한 기여가 의료 영상 파이프라인에서 고차 ODE 솔버의 잠재력을 입증하며, 향후 염색 정규화 연구를 위한 분해 가능한 지표를 제공한다고 결론짓습니다.