연합 비전 트랜스포머와 적응형 초점 손실을 이용한 의료 영상 분류

초록

본 논문은 의료 영상 분류를 위해 Vision Transformer 기반 연합 학습 프레임워크를 제안한다. 각 클라이언트의 클래스 불균형을 실시간으로 반영하는 동적 적응형 초점 손실(DAFL)과 데이터 규모·특성을 고려한 가중 평균 집계 방식을 도입해 비IID 환경에서도 높은 정확도와 소수 클래스 성능을 달성한다. ISIC, Ocular Disease, RSNA‑ICH 데이터셋에서 기존 CNN·ViT·연합 모델들을 크게 능가함을 실험으로 입증한다.

상세 분석

이 연구는 의료 현장에서 데이터 프라이버시와 레이블 불균형이라는 두 가지 핵심 난제를 동시에 해결하고자 한다. 첫 번째로, Vision Transformer(ViT)를 백본으로 채택한 이유는 패치 단위의 전역 자기‑주의 메커니즘이 병변의 미세한 공간적 패턴을 효과적으로 포착할 수 있기 때문이다. 기존 CNN 기반 모델은 지역적 특징에 의존하는 반면, ViT는 전체 이미지 컨텍스트를 고려해 클래스 간 경계가 모호한 의료 영상에서도 강인한 표현력을 제공한다.

두 번째 핵심 기여는 동적 적응형 초점 손실(DAFL)이다. 전통적인 초점 손실은 고정된 감마(γ)와 클래스 가중치(α)를 사용해 어려운 샘플에 더 큰 페널티를 부여하지만, 연합 학습에서는 각 라운드마다 클라이언트별 데이터 분포가 변동한다. 저자들은 각 클라이언트 k의 현재 클래스 비율 p_{k,c}와 전체 라운드 t에서의 누적 비율을 이용해 동적 불균형 계수 β_{k,c}^{(t)}를 정의한다. 이 계수는 소수 클래스일수록 크게 증가하고, 학습이 진행됨에 따라 예측 난이도(예: 모델이 낮은 확신을 보이는 경우)와 결합돼 손실 함수 L_{k}^{(t)} = -∑{c} β{k,c}^{(t)} (1‑p_{k,c}^{pred})^{γ} log(p_{k,c}^{pred}) 로 구현된다. 따라서 클래스 가중치가 정적으로 고정되지 않아, 클라이언트 간 데이터 스키우가 심한 상황에서도 소수 클래스가 지속적으로 학습에 반영된다.

세 번째로 제안된 클라이언트‑인식 집계(weighted aggregation) 전략은 전통적인 FedAvg가 단순히 데이터 양에 비례해 가중치를 부여하는 한계를 극복한다. 저자들은 각 클라이언트의 불균형 계수 평균 β̄_{k}^{(t)}와 데이터 규모 N_k를 결합해 최종 집계 가중치 ω_k^{(t)} = (N_k·β̄_{k}^{(t)}) / ∑{j}(N_j·β̄{j}^{(t)}) 로 정의한다. 이 방식은 데이터가 많지만 내부적으로 편향된 클라이언트가 글로벌 모델에 과도히 영향을 미치는 것을 억제하고, 소규모이면서도 균형 잡힌 클라이언트의 기여도를 상승시킨다.

실험에서는 ISIC(피부 병변), Ocular Disease(안과 질환), RSNA‑ICH(뇌출혈) 세 가지 공개 데이터셋을 사용해 5~10개의 클라이언트 시뮬레이션을 수행하였다. 비교 대상에는 DenseNet121, ResNet50, ViT‑S/16, ViT‑L/32, Swin‑Transformer, CoAtNet, MixNet, 그리고 최신 연합 방법인 FedCLIP, FedProx, SCAFFOLD 등이 포함된다. 결과는 전반적으로 DAFL‑ViT가 정확도(ACC)와 F1‑score에서 0.98%~41.69% 향상을 보였으며, 특히 소수 클래스에 대한 재현율(recall)이 크게 개선되었다. Ablation study에서는 (1) 고정 초점 손실, (2) 단순 FedAvg, (3) β 계수 없이 DAFL만 적용한 경우와 비교해 제안된 두 요소가 모두 성능 향상에 기여함을 확인하였다.

추가적으로, 모델의 자기‑주의 맵을 시각화해 임상 전문가가 해석 가능한 영역을 강조했으며, 이는 단순 정확도 향상을 넘어 실제 진단 지원 시스템에 필요한 투명성을 제공한다는 점에서 의미가 크다. 한계점으로는 클라이언트 수가 급증하거나 통신 비용이 제한된 환경에서 β 계수 전송이 추가 오버헤드를 발생시킬 수 있다는 점, 그리고 현재는 이미지 분류에만 적용했으므로 세그멘테이션·검출 등 다른 의료 비전 태스크에 대한 확장성이 검증되지 않았다는 점을 언급한다. 향후 연구에서는 경량화된 β 전파 메커니즘과 멀티‑태스크 학습으로의 일반화를 목표로 할 수 있다.