[논문 리뷰] Federated attention consistent learning models for prostate cancer diagnosis and Gleason grading

Fei Kong, Xiyue Wang외 11명|arXiv (Cornell University)|2023. 02. 13.

Privacy-Preserving Technologies in Data인용 수 5

한 줄 요약

이 논문은 전립선암 진단 및 글리슨 등급 분류에서 모델의 일반화 능력과 기밀성을 향상시키기 위해 로컬 클라이언트 모델과 중앙 서버 모델 간의 주의 일관성을 강제하는 연합 대비 학습(Federated Contrastive Learning, FCL) 프레임워크를 제안한다. 이 방법은 다수의 기관에서 수집한 19,635장의 웨이브슬라이드 이미지(WSI)를 바탕으로 진단에서 97% AUC, 글리슨 등급 분류에서 0.84 카파를 기록하며, 외부 데이터셋에서도 뛰어난 성능을 보이며 주의 히트맵을 통한 해석 가능성까지 향상시킨다.

ABSTRACT

Artificial intelligence (AI) holds significant promise in transforming medical imaging, enhancing diagnostics, and refining treatment strategies. However, the reliance on extensive multicenter datasets for training AI models poses challenges due to privacy concerns. Federated learning provides a solution by facilitating collaborative model training across multiple centers without sharing raw data. This study introduces a federated attention-consistent learning (FACL) framework to address challenges associated with large-scale pathological images and data heterogeneity. FACL enhances model generalization by maximizing attention consistency between local clients and the server model. To ensure privacy and validate robustness, we incorporated differential privacy by introducing noise during parameter transfer. We assessed the effectiveness of FACL in cancer diagnosis and Gleason grading tasks using 19,461 whole-slide images of prostate cancer from multiple centers. In the diagnosis task, FACL achieved an area under the curve (AUC) of 0.9718, outperforming seven centers with an average AUC of 0.9499 when categories are relatively balanced. For the Gleason grading task, FACL attained a Kappa score of 0.8463, surpassing the average Kappa score of 0.7379 from six centers. In conclusion, FACL offers a robust, accurate, and cost-effective AI training model for prostate cancer pathology while maintaining effective data safeguards.

연구 동기 및 목표

• 민감하고 공유 불가능한 데이터를 가진 다수의 의료 기관에서 전립선암 진단 및 글리슨 등급 분류를 위한 강력한 AI 모델을 훈련하는 데 도전하는 것.
• 로컬 클라이언트 모델과 중앙 서버 모델 간의 주의 일관성을 강제하여 연합 학습에서 모델 일반화 능력을 향상시키는 것.
• 모델 업데이트에 노이즈를 주입하여 차별적 프라이버시를 통합함으로써 연합 학습에서 기밀성 위험을 완화하는 것.
• 공개 및 비공개 웨이브슬라이드 이미지 컬렉션을 포함한 외부 데이터셋에서 모델 성능과 강건성을 검증하는 것.
• 웨이브슬라이드 이미지(WSI)에서 진단적으로 관련 있는 영역을 강조하는 주의 히트맵을 시각화하여 모델의 해석 가능성 향상

제안 방법

• FCL 프레임워크는 로컬 클라이언트 표현을 글로벌 서버 모델과 대비함으로써, 클라이언트 간 및 서버와의 주의 일관성을 극대화한다.
• 훈련 중에 로컬 모델의 주의 맵을 글로벌 모델의 주의 맵과 일치시킴으로써 주의 일관성을 강제하여 특징 정렬과 일반화 능력을 향상시킨다.
• 차별적 프라이버시는 모델 기울기의 조정 과정에 캘리브레이션된 노이즈를 추가하여 적용되며, 소유자 식별 및 모델 역설계 공격을 방지한다.
• 프레임워크는 클라이언트가 원시 데이터를 업로드하지 않고 업데이트된 모델 파라미터만 업로드하는 주요 서버 아키텍처를 채택하여 데이터 기밀성과 소유권을 유지한다.
• 주의 히트맵은 224×224 패치 크기와 90% 겹침을 사용하며, 정규화된 주의 점수는 빨간색 = 높은 주의를 나타내는 RGB 색상 맵으로 시각화된다.
• 모델 성능 평가에는 진단에 대해 AUC, 글리슨 등급 분류에 대해 Kappa를 사용하여 다수의 클라이언트 및 외부 데이터셋에서 평가한다.

실험 결과

연구 질문

• RQ1 주의 일관성을 갖춘 연합 대비 학습이 다수 기관에서 전립선암 진단 및 글리슨 등급 분류 성능을 향상시키는가?
• RQ2 비독립 동일 분포(non-IID) 데이터 분포가 존재하는 환경에서 로컬 모델과 글로벌 모델 간의 주의 일관성은 모델 일반화에 어떤 영향을 미치는가?
• RQ3 차별적 프라이버시는 연합 의료 AI에서 데이터 泄露를 방지하면서도 모델 유틸리티를 어느 정도 유지하는가?
• RQ4 FCL 프레임워크는 다른 기관에서 온 외부의, 볼 수 없는 데이터셋에서도 강건한 성능을 유지하는가?
• RQ5 모델이 생성한 주의 히트맵은 웨이브슬라이드 이미지에서 임상적으로 관련 있는 영역을 신뢰성 있게 강조하는가?

주요 결과

• 진단 작업에서 FCL 모델은 7개 클라이언트에서 평균 AUC 97%를 기록하여 균형 잡힌 클래스 조건에서 베이스라인 FL 모델(95%)을 능가했다.
• 글리슨 등급 분류에서 FCL 모델은 6개 클라이언트에서 평균 카파 0.84를 달성하여 베이스라인 FL 모델의 0.74보다 유의미하게 높았다.
• FCL 모델은 공개(PANDA) 및 비공개(DiagSet-A, QHD) 외부 데이터셋 모두에서 강력한 성능을 유지하여 훈련 기관을 초월한 일반화 능력을 입증했다.
• 주의 히트맵은 WSI에서 암성 영역을 성공적으로 국소화했으며, 글리슨 패턴 형성과 같은 진단적으로 관련 있는 영역에 높은 주의 점수가 집중되어 있었다.
• FCL에 차별적 프라이버시를 통합함으로써 모델 유틸리티를 유지했으며, 가장 높은 기밀성 보장 설정(μ=0.1, z=0.1)에서도 0.7342 카파를 기록하여 기밀성-정확도 트레이드오���이 유리함을 보였다.
• 중앙 집중형 훈련 베이스라인은 카파 0.7171을 기록하여, FCL이 이 연합 환경에서 일관성과 강건성 향상 덕분에 중앙 집중형 학습을 능가함을 확인했다.