Oct 25, 2024
KoSAIM
Abstract
OSTEO
Enhancement of Osteoporosis Classification AI Model Using Federated Learning in Chest X-Ray
Minjun Kim
Object
연합학습을 활용해 데이터의 반출을 하지 않고, 두 병원의 데이터를 동시에 학습하여, 흉부 X선 영상(CXR) 기반 골다공증 스크리닝 인공지능 모델 개발을 하고 성능을 평가한다.
Methods
Figure 1. 연합학습에서 서버-클라이언트간 가중치를 업데이트하는 과정
국내 병원A의 54,793건의 학습데이터(골다공증 5.0%)와 7652건의 검증데이터(골다공증 4.2%), 국내 병원B의 8,819건의 학습데이터(골다공증 6.3%)와 1561건의 검증데이터(골다공증 6.4%)를 사용했다. 골다공증 이진 분류 모델 학습은 연합학습(Federated Learning; FL)을 활용해 두 병원의 각 서버에서 Fedavg[1] 방식의 연합 학습을 진행한다. 모델의 성능을 평가하기 위해서 병원A의 내부 데이터셋과 다른 기간에 수집된 병원A의 외부 검증데이터셋 1026건, 병원B 내부 데이터셋과 다른 기간에 수집된 외부 검증 데이터셋 3098건과 평균나이와 성별 분포가 다른 C병원 데이터셋 934건을 이용하였다. 가중치들의 합의 평균인 FedAvg[1], 서버에서 옵티마이저(optimizer)를 적용하는 방법인 Fedopt[2], 객관적 불일치를 제거하는 정규화 방법인 FedNova[3]를 집계 알고리즘(Aggregation algorithm)중 성능이 가장 좋은 FedAvg[1]을 사용하였다.
라운드가 시작되면, 클라이언트는 골다공증 분류를 위해 설계된 인셉션 네트워크를 기반으로 한 모델[4]로 학습을 진행한다. 학습이 종료되면 도출된 가중치를 암호화하여 서버로 전송한다. 서버는 암호화된 가중치를 다시 컴퓨터가 알아볼 수 있도록 해독하는 작업을 거친 후, 상기 집계(aggregation) 과정을 통해 새로운 가중치를 클라이언트에게 전달한다. 이러한 사이클을 라운드라고 지칭하며 사용자가 충분히 학습됬다고 판단할 때까지 진행한다. 마지막 라운드에서는 최종 집계가중치(aggregated weight)를 미세조정하는 과정을 거쳐 각 클라이언트에 최적화 되도록 설정 할수있고, 외부 검증을 진행할 때는 마지막 라운드의 서버가중치를 사용한다. 학습 성능 측정을 위해서 A 병원 데이터에서만 학습된 모델과 A, B병원에서 연합학습된 모델에 대해서 학습 커브와 성능을 비교하였다.(figure 2)
Figure 2. 학습, 검증 손실함수에 따른 학습 속도 비교 결과
Results
연합학습으로 학습된 모델은 수렴속도가 40% 이상 향상되었고, 일부 데이터셋에서는 민감도가 좋아지는 경향을 보였다. Hospital A의 내부 테스트셋(16.1% 남성; 평균 연령 58.7 ± 6.76)과 동일한 병원의 Hospital A의 외부 테스트셋(11.1% 남성; 평균 연령, 59.01 ± 6.65)은 외래, 검진센터 등 다른 기간에 수집되어진 데이터이고, Hospital B의 외부 테스트셋(44.5% 남성; 평균 연령, 59.38 ± 7.31), Hospital C의 외부 테스트셋(56.2% 남성; 평균 연령, 73.64 ± 6.74)에서 Table 1와 같이 검증하였다.
Hospital A internal | Hospital A external | Hospital B external | Hospital C external | ||
Origin | AUC | 0.9395 | 0.8823 | 0.9061 | 0.8770 |
민감도 | 70.78% | 67.61% | 56.93% | 75.97% | |
특이도 | 93.81% | 86.51% | 93.72% | 81.28% | |
ACC | 86.12% | 80.99% | 91.49% | 79.96% | |
F1 score | 0.773 | 0.675 | 0.4475 | 0.6531 | |
FL | AUC | 0.9421 | 0.8974 | 0.9079 | 0.8673 |
민감도 | 76.17% | 78.53% | 66.67% | 77.59% | |
특이도 | 91.10% | 83.33% | 91.66% | 78.92% | |
ACC | 86.12% | 81.87% | 91.06% | 80.79% | |
F1 score | 0.785 | 0.725 | 0.4615 | 0.635 |
Table 1. 단일 학습과 연합학습의 성능지표 비교 결과
Conclusions
우리는 두 개의 병원에서 연합학습을 실제로 시행해 봄으로써 데이터셋이 많을 수록 학습이 빠르게 수렴함을 알 수 있었고 일부 외부검증 데이터셋에서 민감도의 개선 효과는 보여주였다.
Reference
[1] McMahan, Brendan, et al. "Communication-efficient learning of deep networks from decentralized data." Artificial intelligence and statistics. PMLR, 2017.
[2] Reddi, Sashank, et al. "Adaptive federated optimization." arXiv preprint arXiv:2003.00295 (2020).
[3] Wang, Jianyu, et al. "Tackling the objective inconsistency problem in heterogeneous federated optimization." Advances in neural information processing systems 33 (2020): 7611-7623.
[4] Jang, Miso, et al. "Opportunistic osteoporosis screening using chest radiographs with deep learning: development and external validation with a cohort dataset." Journal of Bone and Mineral Research 37.2 (2020): 369-377.
*Keywords
Chest Radiology, Classification, Federated Learning, Osteoporosis
