[논문 리뷰] Application of Federated Learning in Building a Robust COVID-19 Chest X-ray Classification Model

Amartya Bhattacharya, Manish Gawali외 2명|arXiv (Cornell University)|2022. 04. 22.

COVID-19 diagnosis using AI인용 수 4

한 줄 요약

이 논문은 환자 데이터를 중앙 집중화하지 않고 흉부 X선 영상에서 코로나19를 분류하기 위한 강력한 딥러닝 모델을 훈련하기 위한 연합 학습(FL) 프레임워크를 제안한다. 비독립 동일분포(Non-IID) 데이터 소스 3개에서 훈련하고 FedAvg를 통해 국지적 모델 업데이트를 집계함으로써, 연합 학습 모델은 모든 데이터를 포함한 중앙 집중 모델과 동일한 성능을 달성했으며, 데이터 프라이버시를 유지하면서 다양한 데이터 분포 간의 일반화 능력을 향상시켰다.

ABSTRACT

While developing artificial intelligence (AI)-based algorithms to solve problems, the amount of data plays a pivotal role - large amount of data helps the researchers and engineers to develop robust AI algorithms. In the case of building AI-based models for problems related to medical imaging, these data need to be transferred from the medical institutions where they were acquired to the organizations developing the algorithms. This movement of data involves time-consuming formalities like complying with HIPAA, GDPR, etc.There is also a risk of patients' private data getting leaked, compromising their confidentiality. One solution to these problems is using the Federated Learning framework. Federated Learning (FL) helps AI models to generalize better and create a robust AI model by using data from different sources having different distributions and data characteristics without moving all the data to a central server. In our paper, we apply the FL framework for training a deep learning model to solve a binary classification problem of predicting the presence or absence of COVID-19. We took three different sources of data and trained individual models on each source. Then we trained an FL model on the complete data and compared all the model performances. We demonstrated that the FL model performs better than the individual models. Moreover, the FL model performed at par with the model trained on all the data combined at a central server. Thus Federated Learning leads to generalized AI models without the cost of data transfer and regulatory overhead.

연구 동기 및 목표

• 민감한 환자 영상 데이터의 중앙 집중화를 피하여 의료 AI에서의 데이터 프라이버시 및 규제 준수 문제를 해결한다.
• 분포 다양성이 제한된 단일 소스에서의 데이터 훈련으로 인한 모델 편향을 줄인다.
• 다양한 임상 데이터 분포에서의 일반화 능력을 향상시키기 위해 비독립 동일분포(Non-IID) 데이터 소스를 다수 활용한다.
• 연합 학습이 원시 데이터를 이동시키지 않으면서도 중앙 집중 훈련과 비교할 수 있는 성능을 달성할 수 있음을 입증한다.
• 개별, 연합, 통합 훈련 설정 간의 모델 성능 비교를 통해 연합 학습의 효과성을 검증한다.

제안 방법

• 다양한 출처에서 온 비독립 동일분포(Non-IID) 흉부 X선 영상 데이터셋 3개에서 각각 별도의 딥러닝 모델(CNN 기반)을 훈련한다.
• 기준 비교를 위해 모든 데이터를 통합하여 단일 서버에서 전체 데이터 접근 권한을 가진 중앙 집중식 '통합' 모델을 훈련한다.
• 각 클라이언트가 국지적으로 훈련하고 모델 업데이트만 공유하는 방식으로, FedAvg를 적용하여 글로벌 FL 모델을 3개의 클라이언트 간에 훈련한다.
• 5라운드 동안 총 20에포크를 반복하며, 매 라운드 후 글로벌 모델을 집계한다.
• 모델 평가를 위해 각 클라이언트의 테스트 세트에서 임계값에 의존하지 않는 지표—AUROC 및 AUPRC를 사용하여 공정한 비교를 확보한다.
• 모든 훈련 및 집계 과정을 국지적으로 안전하게 수행함으로써 원시 이미지를 중앙 서버로 이동시키지 않아 데이터 프라이버시를 보장한다.

실험 결과

연구 질문

• RQ1 연합 학습은 다양한 비독립 동일분포(Non-IID) 데이터 분포에서 코로나19 흉부 X선 영상 분류 모델의 일반화 능력을 향상시킬 수 있는가?
• RQ2 개별 데이터 소스에서 훈련된 모델과 모든 데이터를 통합한 중앙 집중 모델과 비교할 때 연합 모델의 성능은 어떠한가?
• RQ3 민감한 의료 영상 데이터를 이동시키지 않으면서도 연합 학습이 모델의 강건성을 유지할 수 있는 정도는 어느 정도인가?
• RQ4 단일 데이터 소스에서 훈련된 모델과 비교해 연합 학습 프레임워크는 데이터 편향을 줄일 수 있는가?
• RQ5 원시 데이터 이동 없이도 연합 학습이 환자 데이터 프라이버시를 훼손하지 않고 중앙 집중 훈련과 비교할 만한 성능을 달성할 수 있는가?

주요 결과

• FL 모델은 각 클라이언트의 테스트 세트에서 개별 클라이언트 모델보다 뛰어난 일반화 능력을 보였으며, 이는 뛰어난 일반화 능력을 입증한다.
• Client1의 테스트 세트에서 FL 모델은 AUROC 99.61%를 기록했으며, 개별 Client1 모델(99.96%)을 뛰어나며, Client2(37.85%)와 Client3(43.85%) 모델보다는 크게 뛰어났다.
• Client2의 테스트 세트에서 FL 모델은 AUROC 92.31%를 기록했으며, 개별 Client2 모델(93.94%)을 뛰어나며, Client1(37.85%)과 Client3(43.85%) 모델보다는 크게 향상되었다.
• Client3의 테스트 세트에서 FL 모델은 AUROC 88.31%를 기록했으며, 개별 Client3 모델(97.73%)을 뛰어나며, Client1(43.85%)과 Client2(37.85%) 모델보다는 크게 향상되었다.
• FL 모델은 각 클라이언트의 테스트 세트에서 AUROC 점수 99.61%, 92.31%, 88.31%를 각각 기록하여 다양한 데이터 분포 간에 일관된 높은 성능을 보였다.
• FL 모델의 성능은 중앙 집중식 '통합' 모델과 거의 동일했으며, 동일한 테스트 세트에서 AUROC 점수 99.86%, 99.92%, 93.51%를 기록하여, 데이터 이동 없이도 중앙 집중 훈련 수준의 성능을 달성할 수 있음을 확인했다.