[논문 리뷰] Real World Federated Learning with a Knowledge Distilled Transformer for Cardiac CT Imaging

Malte Tölle, Philipp Garthe외 23명|arXiv (Cornell University)|2024. 07. 10.

Advanced X-ray and CT Imaging인용 수 4

한 줄 요약

이 논문은 심장 CT 영상에서 부분적으로 레이블이 붙은 데이터를 효과적으로 활용하기 위해 임의의 병원 8곳에서의 개인정보 보호 제약과 레이블 이질성으로 인해 어려움을 겪는 강력한 딥러닝 모델을 훈련하는 데 도전하는 이중 단계의 준지도 학습 연합 학습 프레임워크를 제안한다. 이 방법은 작업별로 특화된 CNN에서 하나의 비전 트랜스포머로 지식을 흡수함으로써, 여러 심장 구조를 동시에 예측할 수 있는 통합된 트랜스포머를 훈련함으로써 정확도와 일반화 능력을 향상시킨다. 코드와 가중치는 오픈소스로 제공되어 향후 연구를 위한 기초 모델이 된다.

ABSTRACT

Federated learning is a renowned technique for utilizing decentralized data while preserving privacy. However, real-world applications often face challenges like partially labeled datasets, where only a few locations have certain expert annotations, leaving large portions of unlabeled data unused. Leveraging these could enhance transformer architectures ability in regimes with small and diversely annotated sets. We conduct the largest federated cardiac CT analysis to date (n=8,104) in a real-world setting across eight hospitals. Our two-step semi-supervised strategy distills knowledge from task-specific CNNs into a transformer. First, CNNs predict on unlabeled data per label type and then the transformer learns from these predictions with label-specific heads. This improves predictive accuracy and enables simultaneous learning of all partial labels across the federation, and outperforms UNet-based models in generalizability on downstream tasks. Code and model weights are made openly available for leveraging future cardiac CT analysis.

연구 동기 및 목표

• 개인정보 보호 제약과 레이블 이질성으로 인해 다수의 병원에서 부분적으로 레이블이 붙은 심장 CT 데이터셋에서 강력한 딥러닝 모델을 훈련하는 데 도전하기 위해.
• 임상 환경에서 일반적으로 사용되지 않는 방대한 양의 레이블이 없는 데이터를 활용하여 모델 성능을 향상시키기 위해.
• 레이블 커버리지가 상이한 기관들 간에 다수의 심장 구조(예: 석회화, 관상동맥 등)를 동시에 학습할 수 있는 연합 학습 프레임워크를 개발하기 위해.
• 다음 작업에서 표준 CNN보다 우수한 성능을 보이는 일반화 능력이 뛰어난 오픈소스 기초 모델을 심장 CT 영상에 대해 개발하기 위해.
• 중앙 집중식 데이터 저장을 피하고 로컬에서 모델을 훈련하며 오직 모델 가중치만을 집계함으로써 개인정보 보호 규정을 준수하기 위해.

제안 방법

• 이중 단계 준지도 학습 전략을 사용한다: 먼저 각 클라이언트 사이트에서 레이블이 있는 데이터를 바탕으로 작업별로 특화된 CNN(지점 검출 및 석회화 세그먼테이션용)을 훈련시킨다.
• 모든 클라이언트에서 레이블이 있는 데이터와 방대한 레이블이 없는 데이터를 활용하여 지식 흡수를 적용하여 이러한 CNN들에서 하나의 비전 트랜스포머 모델로 지식을 전이한다.
• 8개 병원 사이트에서 로컬 훈련 후 중앙 집중적으로 모델 가중치를 집계함으로써 데이터 프라이버시를 유지하면서 트랜스포머 모델을 연합 학습한다.
• 최종 트랜스포머 모델은 마지막 레이어만 훈련시켜 후행 작업인 관상동맥 세그먼테이션에 대해 피지컬 테스트를 수행함으로써 더 나은 특징 표현을 보여준다.
• FedBioMed 프레임워크는 TLS 암호화, 클라이언트 인증 및 IP 화이트리스트를 통해 안전한 통신을 보장하며, 단방향이고 안전한 연합 학습 파이프라인을 제공한다.
• 각 기관에서의 사전 처리, 훈련 및 검증을 위해 MONAI와 nnU-Net을 기반으로 한 커스터마이즈된 도커 오케스트레이션 및 파이토치 기반 파이프라인을 사용한다.

실험 결과

연구 질문

• RQ1 부분적으로 레이블이 붙은 데이터를 다수의 병원에서 사용하여 단일 비전 트랜스포머 모델이 다양한 심장 구조(예: 석회화, 관상동맥 등)를 효과적으로 학습할 수 있는가?
• RQ2 작업별로 특화된 CNN에서 트랜스포머로 지식 흡수를 적용함으로써 연합 환경에서 후행 세그먼테이션 작업의 성능이 향상되는가?
• RQ3 다양한 병원에서 온 레이블이 없는 데이터를 개인정보 보호를 지키는 연합 학습 환경에서 얼마나 효과적으로 활용하여 모델의 일반화 능력을 향상시킬 수 있는가?
• RQ4 트랜스포머 기반의 연합 모델은 특징 품질과 세그먼테이션 정확도 측면에서 CNN 기반 기준 모델과 비교해 어떻게 성능을 내는가?
• RQ5 중앙 집중식 데이터 수집 없이 이질적인 레이블 분포를 가진 기관들 간에 통합된 기초 모델을 훈련시킬 수 있는가?

주요 결과

• 제안된 연합 지식 흡수 트랜스포머 모델은 기준 UNet보다 심장 CT 영상 작업에서 뛰어난 성능을 보였다. 특히 후행 작업인 관상동맥 세그먼테이션에서 두드러진 성능 향상을 보였다.
• 지식 흡수를 통해 레이블이 없는 데이터를 활용함으로써, 클라이언트 당 레이블이 제한된 상황에서도 예측 정확도와 일반화 능력이 향상되었다.
• 이 방법은 연합 내에서 여러 부분 레이블(예: 석회화, 지점, 관상동맥 등)을 동시에 하나의 트랜스포머 아키텍처 내에서 학습시킬 수 있었다.
• 관상동맥 세그먼테이션을 위해 트랜스포머의 마지막 레이어만 미세조정함으로써 CNN 기준 모델보다 더 의미 있는 특징을 도출했으며, 이는 더 나은 표현 학습 능력을 시사한다.
• 이론적이고 실용적인 환경에서 이중 병원 8곳에서 이질적인 레이블 커버리지가 존재하는 상황에서도 최종 모델이 성공적으로 훈련되었으며, 실제 연합 환경에서의 확장성과 강건성을 입증했다.
• 코드, 모델 가중치, 데이터 처리 파이프라인은 FAIR 준수 프레임워크 하에 공개되어 있어 심장 영상 연구 분야에서의 재사용과 추가 개발이 가능하다.